提出分层稀疏注意力机制，让扩散模型高效处理超长序列图像生成。
Trainable Log-linear Sparse Attention for Efficient Diffusion Transformers
- 用分层Top-K选择减少冗余计算，实现对数级复杂度。
- 在256x256图像上推理提速28.27倍，训练提速6.09倍。
- 适合追求长序列生成效率的视觉生成研究者。
扩散变换器（DiTs）在视觉生成中表现卓越，但其自注意力计算呈二次方增长，限制了长序列扩展。现有Top-K稀疏注意力通过块化压缩和选取关键块降低计算量，但仍存在（i）压缩后令牌的二次选择开销，以及（ii）随序列增长需增大K值以维持质量的问题。我们发现根源在于单层设计无法有效表达全局结构。本文提出可训练的对数线性稀疏注意力（LLSA），通过分层结构将选择与注意力成本从二次降至对数级。LLSA采用逐层Top-K选择，利用上层索引引导下层稀疏选择，并引入分层键值增强机制，在减少不同粒度令牌使用的同时保留全局上下文。为支持高效训练，我们开发了仅用稀疏索引的高性能GPU实现，前向与反向均无需稠密注意力掩码。在不使用分块化或VAE编码的高分辨率像素空间图像生成任务中，LLSA在256x256 token序列上实现推理加速28.27倍、训练加速6.09倍，且生成质量保持不变。结果表明，LLSA为高效训练长序列DiTs提供了可行方向。代码已开源。
原文摘要 · Abstract (English)
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.27x and DiT training by 6.09x on 256x256 pixel token sequences, while maintaining generation quality. The results demonstrate that LLSA offers a promising direction for training long-sequence DiTs efficiently. Code is available at: https://github.com/SingleZombie/LLSA
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。