论文题目Trainable Log-linear Sparse Attention for Efficient Diffusion Transformers用于高效扩散 Transformer 的可训练对数线性稀疏注意力会议CVPR 2026Highlight摘要扩散 TransformerDiT已经成为视觉生成中的先进架构但自注意力的二次计算开销从根本上限制了它向长 Token 序列扩展。近期的 Top-K 稀疏注意力通过把 Token 压缩为块级表示、再为每个 Query 选择少量相关 Key Block 来降低 DiT 的计算量但仍有两个问题压缩 Token 上的选择过程依旧具有二次复杂度随着序列变长为维持模型质量还需要不断增大 K。作者认为这种低效来自现有方法的单层设计——单一粗粒度层不足以描述长序列的全局结构。为此论文提出 Log-linear Sparse AttentionLLSA利用层次结构把选择和注意力计算从二次复杂度降到对数线性复杂度。LLSA 采用层次化 Top-K 选择从粗层得到的索引出发逐层向细粒度递归搜索同时提出 Hierarchical KV Enrichment把不同粒度的粗层 Key/Value 补充进最终注意力以更少的 Token 保留全局上下文。为了支持高效训练作者进一步实现了只依赖稀疏索引的 GPU 前向与反向计算不再构造稠密 Attention Mask。实验在不使用 VAE 编码、且高分辨率 Pixel DiT 不进行 Patchification 的设置下验证 LLSA在 256 × 256 Pixel Token 序列上注意力推理最高加速 28.27×DiT 训练加速 6.09×同时保持接近全注意力的生成质量。源码https://github.com/SingleZombie/LLSA一、研究背景与核心问题DiT 的瓶颈非常直接标准 Self-Attention 对长度为 N 的序列需要构造 Query-Key 两两关系计算复杂度是 O(N²)。在 Latent DiT 中Token 数往往已经通过 VAE 和 Patchification 被大幅压缩但如果希望直接在 Pixel Space 建模或者进一步扩展到长视频N 会迅速增大二次复杂度很快成为主要开销。已有 Top-K Block Sparse Attention 看起来已经把注意力“稀疏化”了先把 Q、K 压成块级 Token再在粗 Token 上计算相似度为每个 Query Block 选出 Top-K Key Block最后只对这些块做 Sparse Attention。问题在于它只稀疏了最后一步。粗 Token 之间仍然要做全量两两匹配所以 Selection Stage 仍有 O(N²) 主导项而且为了不丢掉远距离全局信息序列越长通常还要使用更大的 K。LLSA 的核心判断是真正需要改的不是“再少选几个 Key”而是 Top-K 的搜索结构本身。如果用多层层次结构表示全局信息那么最粗层负责全局定位细层只在上一级命中的局部候选中继续搜索就没有必要在每一层重新做全局两两比较。二、方法整体框架从单层 Top-K 变成层次搜索论文 Figure 1普通 Top-K Sparse Attention 与 LLSA 的整体对比左侧普通 Top-K 只有一次 Compression压缩后仍要在所有粗 Token 间做全局 Top-K。右侧 LLSA 则建立多层表示先在最粗层定位再沿已命中的索引逐层向细粒度递归最终 Attention 不只使用最细层 Top-K KV还加入高层选中的粗粒度 KV。整体流程可以概括为Hierarchical Compression → Coarse-to-Fine Top-K Selection → Hierarchical KV Enrichment → Sparse Attention前两步解决“怎么把选择成本降下来”KV Enrichment 解决“稀疏以后怎么保留全局信息”稀疏索引 GPU Kernel 则保证训练阶段也不会重新引入二次开销。三、Hierarchical Top-K把 Selection 从 O(N²) 降到 O(NK)LLSA 首先递归平均池化 Q、K、V。第 0 层是原始 Token第 l 层序列长度变成 N/Bˡ一个高层 Token 就是 B 个下一层 Token 的摘要。层数取 L ≈ log_B N因此越往上 Token 越少、感受野越大。搜索从最粗层开始。只有在这一层LLSA 计算完整的 QKᵀ 并得到 Top-K得到粗层索引后一个命中的粗 Key 在下一层只对应 B 个子 Key因此细一层的每个 Query 不再面对全部 Key而只需要在大约 KB 个候选中继续做 Top-K。这个过程不断向下递归直到得到最细层索引。对除最粗层外的各层选择成本可以写成关键就在最后一步几何级数收敛因此层数虽然是 O(log N)Selection Stage 并不会重新变成 O(N log N)而是保持 O(NK)。当 K 视为常数时就是关于序列长度 N 的线性复杂度。直观上这和“先在地图上确定城市再确定街区最后找门牌号”很像。普通 Top-K 每次都在全国范围重新搜索LLSA 则把上一级结果直接变成下一级的候选集合。四、Hierarchical KV Enrichment稀疏以后怎样保住全局上下文逐层向细粒度搜索会缩小 Query 的候选范围因此 LLSA 把搜索阶段得到的高层 KV 再利用起来最终每个 Query Block 同时使用最细层 Top-K KV 与多个层级的粗粒度 KV。近距离信息由细 Token 表达远距离上下文由粗 Token 提供。每层只补充 K 个左右的候选、层数为 O(log N)因此 Sparse Attention 为 O(NK log N)结合 O(NK) 的 Selection总体是 O(NK log N)K 固定时即 O(N log N)。粗 Token 是多个细 Token 的平均若与单个细 Token 等权会低估其信息量因此作者设置 KV Reweighting第 l 层权重 W⁽ˡ⁾ Bˡ。论文 Table 1LLSA 核心模块、Block Size 与 Top-K 的消融Table 1a 中两层 Top-K 的 FID 为 27.98加入 KV Enrichment 后改善到 25.31再加 Reweighting 后达到 24.37吞吐量仍为 436.40。这说明 Hierarchy 更偏向解决效率而 Enrichment 与 Reweighting 负责补回稀疏化带来的质量损失。Table 1c 进一步显示LLSA 用 K 8 就达到 FID 24.37 / 吞吐量 436.40而单层 Baseline 即使 K 32 也只有 25.88 / 357.95说明层次上下文比单纯增大 K 更有效。4.1 反向传播也必须保持稀疏一些 Sparse Attention 在 Forward 中只访问稀疏块但 Backward 会构造 T × T Binary Mask从而重新引入 O(T²)。LLSA 使用类似 CSR→CSC 的 Sparse Index Transpose先统计每个 Key 被哪些 Query 选中再通过 Prefix Sum 得到连续区间直接建立 Key-major 反向索引。这样 Forward 与 KV Backward 都只围绕真实命中的稀疏索引工作。4.2 Pixel DiT 的二维适配论文 Figure 22D Pixel Token 的 Index ReorderingRaster Order 会让层次池化把空间上不够相近的像素混在一起因此作者先重排索引使局部相邻像素在 1D 序列中也尽量相邻。高分辨率训练还使用 Noise Rescaling()对大于 64 × 64 的图像设置同时高分辨率模型从低分辨率 Checkpoint 初始化以加快收敛。五、实验结果与消融分析5.1 实验设置作者主要在 FFHQ 的 Pixel DiT-S 上验证不使用 VAE、且 Patch Size 1 × 1 的长 Pixel Token 建模并在 ImageNet-256 上把 LLSA 接入 PixelFlow-L。质量指标主要使用 10,000 个样本计算 FID效率看 H200 上的训练吞吐量。论文 Table 4FFHQ / ImageNet 不同分辨率训练配置Table 4 给出 Pretrained Model、SNR Rescale、Epoch、Batch Size 与 Learning RateFFHQ 采用 32→128→256→512 的逐级预训练策略Learning Rate 统一为 1 × 10⁻⁴。5.2 FFHQ 主实验论文 Table 2FFHQ-128 / FFHQ-256 主实验128 × 128 上LLSA 的 FID / 吞吐量为 24.37 / 436.40Full Attention 为 24.91 / 188.88VSA 和 SLA 的 FID 分别是 26.91、25.73。256 × 256 上Full Attention 的 FID 略好为 38.77LLSA 为 39.29但在稀疏方法中最好吞吐量达到 375.34而 Full Attention 只有 61.64对应约6.09× 的 DiT 训练加速。因此这里更准确的结论是长 Pixel Token 下LLSA 以很小的质量差距换来显著训练提速。5.3 ImageNet-256接入 PixelFlow 后仍然成立吗论文 Table 3PixelFlow ImageNet-256 上 VSA / SLA / LLSA 对比LLSA 的 FID 20.41、Inception Score 73.21、吞吐量 34.16 images/sVSA 为 23.59 / 64.07 / 32.30SLA 为 22.58 / 65.31 / 29.81。说明收益不只存在于 FFHQ 的轻量 Pixel DiT-S在更复杂的 PixelFlow 设置中也能同时改善质量与吞吐量。5.4 Kernel 效率理论复杂度是否真的变成速度论文 Figure 3不同 Sparse Attention 相对 FlashAttention2 的推理/训练加速比Figure 3 在不同序列长度与 B 16/64 下比较推理和训练速度。序列变长后VSA/SLA 的二次 Selection 与 Mask-based Backward 越来越明显而 LLSA 仍保持优势论文摘要报告 256 × 256 Pixel Token 上 Attention Inference 最高加速28.27×。论文 Figure 4Sparse KV Backward 吞吐量LLSA 的 Sparse Index Transpose CSC-style KV Backward 随序列长度增长保持近似稳定吞吐量而 Dense Mask Baseline 持续下降说明反向阶段隐藏的 O(N²) Mask 开销确实被去掉。5.5 附录消融与定性结果论文 Table 5Enrichment Level、512 × 512、SNR 与 Index Reordering 消融Enrichment Level 从 0→1→2 时FID 从 27.98→25.49→24.37体现了质量与有效 KV 数量之间的折中。512 × 512 上两层 LLSA 为 FID 39.26 / 吞吐量 292.66三层吞吐量进一步到 323.29单层版本受二次 Selection 限制。Noise Rescale 的 FID 最好为 29.46Index Reordering 也把 FID 从 31.19 改善到 29.46。论文 Figure 5–6低分辨率预训练与 ImageNet 训练曲线Figure 5 说明低分辨率预训练显著加快高分辨率收敛Figure 6 显示 ImageNet-256 前 4 个 Epoch 中LLSA 的训练曲线整体优于 VSA/SLA。论文 Figure 7–8FFHQ 多分辨率生成结果与 ImageNet 定性对比Figure 7 展示 FFHQ-128/256/512 样例其中 512 模型仅训练 2 个 EpochFigure 8 则对比 SLA、VSA、LLSA 的 ImageNet-256 样例与定量结果形成补充证据。六、总结与思考如果把这篇论文压缩成一句话LLSA 把“从所有 Key 中一次性挑 K 个”改造成多尺度、由粗到细的搜索再把搜索过程中产生的粗粒度 KV 变成全局上下文。最值得记住的是三个环节Hierarchical Top-K 把 Selection 降到 O(NK)KV Enrichment Reweighting 让小 K 仍能保留长距离信息使整体 Attention 成为 O(NK log N)Sparse Index Transpose 则把同样的稀疏性贯彻到 Backward。从实验范围看论文最充分的证据仍集中在 FFHQ 与 ImageNet 图像生成512 × 512 也只进行了较短训练。长视频是重要应用动机但本文没有直接给出对应训练实验因此其在更长视频序列上的实际收益仍需进一步验证。