面向高效扩散 Transformer 的可训练对数线性稀疏注意力
Trainable Log-linear Sparse Attention for Efficient Diffusion Transformers
LLSA 把 DiT 自注意力的选择与计算复杂度从二次降到对数线性，并给出仅用稀疏索引的前向/反向 GPU 实现，可直接替换现有 DiT 中的 Top-K 块稀疏注意力，用于高分辨率像素空间长序列生成训练。
原文依据 LLSA accelerates attention inference by 28.27 xand DiT training by 6.09 xon 256 x256 pixel token sequences
摘要译文
由模型翻译 · 术语统一 · 请对照原文核对关键结论扩散 Transformer（DiT）在视觉生成领域达到了最先进水平，但其自注意力的二次方开销从根本上限制了向长词元序列的扩展。近期的 Top-K 稀疏注意力方法通过将词元压缩为块级表示并选取少量相关键块来降低 DiT 的计算量，但仍存在两个问题：(i) 在压缩词元上进行选择的开销仍是二次方；(ii) 随着序列增长，维持模型质量所需的 K 不断增大。我们发现其低效源于单层设计，因为单一粗粒度层级不足以表示全局结构。本文提出对数线性稀疏注意力（LLSA），一种面向极长词元序列的可训练稀疏注意力机制，它利用分层结构将选择开销与注意力开销都从二次方复杂度降为对数线性复杂度。LLSA 执行分层 Top-K 选择，逐级沿用上一层找到的索引进行稀疏 Top-K 选择，并引入分层 KV 增强机制，在注意力计算中使用不同粒度且数量更少的词元，同时保留全局上下文。为支持高效训练，我们开发了高性能 GPU 实现，其前向与反向传播均只使用稀疏索引，从而无需稠密注意力掩码。我们在不使用分块化（patchification）与变分自编码器编码的高分辨率像素空间图像生成上评估 LLSA。在 256×256 像素词元序列上，LLSA 将注意力推理加速 28.27 倍、DiT 训练加速 6.09 倍，同时保持生成质量。结果表明，LLSA 为高效训练长序列 DiT 提供了一个有前景的方向。
英文原文摘要 1768 字符
Diffusion Transformers (DiTs) set the state of the art in visual generation, yet their quadratic self-attention cost fundamentally limits scaling to long token sequences. Recent Top-K sparse attention approaches reduce the computation of DiTs by compressing tokens into block-wise representation and selecting a small set of relevant key blocks, but still suffer from (i) quadratic selection cost on compressed tokens and (ii) increasing K required to maintain model quality as sequences grow. We identify that their inefficiency is due to the single-level design, as a single coarse level is insufficient to represent the global structure.In this paper, we introduce Log-linear Sparse Attention (LLSA), a trainable sparse attention mechanism for extremely long token sequences that reduces both selection and attention costs from quadratic to log-linear complexity by utilizing a hierarchical structure. LLSA performs hierarchical Top-K selection, progressively adopting sparse Top-K selection with the indices found at the previous level,and introduces a Hierarchical KV Enrichment mechanism that preserves global context while using fewer tokens of different granularity during attention computation. To support efficient training, we develop a high-performance GPU implementation that uses only sparse indices for both the forward and backward passes, eliminating the need for dense attention masks.We evaluate LLSA on high-resolution pixel-space image generation without using patchification and VAE encoding. LLSA accelerates attention inference by 28.27 xand DiT training by 6.09 xon 256 x256 pixel token sequences, while maintaining generation quality. The results demonstrate that LLSA offers a promising direction for training long-sequence DiTs efficiently.