Trainable Log-linear Sparse Attention for Efficient Diffusion Transformers

扩散变换器（DiTs）在视觉生成领域达到了最先进的水平，但其自注意力机制的二次方计算成本从根本上限制了模型向长令牌序列的扩展。近期一些基于Top-K稀疏注意力的方法通过将令牌压缩为分块表示，并选择少量相关的关键块来降低DiT的计算开销，但仍面临两个问题：（i）在压缩后的令牌上进行块选择时仍存在二次方级别的计算代价；（ii）随着序列长度增加，必须增大K值才能维持模型性能。我们发现这些问题的根源在于这些方法采用的是单层结构设计，而单一的粗粒度层次不足以有效表达全局结构信息。本文提出了一种名为对数线性稀疏注意力（Log-linear Sparse Attention, LLSA）的可训练稀疏注意力机制，专为处理极长令牌序列而设计。LLSA利用分层结构，将选择和注意力计算的复杂度从二次方降低至对数线性级别。LLSA采用分层的Top-K选择策略，逐级利用上一层确定的索引进行稀疏化选择；同时引入一种“分层键值增强”（Hierarchical KV Enrichment）机制，在注意力计算过程中使用不同粒度的较少令牌，同时保留全局上下文信息。为了支持高效的训练，我们开发了一种高性能GPU实现方案，在前向和反向传播中仅使用稀疏索引，无需生成稠密注意力掩码。我们在不使用图像分块（patchification）和变分自编码器（VAE）编码的情况下，于高分辨率像素空间图像生成任务中对LLSA进行了评估。实验结果表明，在256×256像素的令牌序列上，LLSA使注意力推理速度提升了28.27倍，DiT训练速度提升了6.09倍，同时保持了原有的生成质量。这些结果表明，LLSA为高效训练长序列DiT模型提供了一个极具前景的方向。代码地址：https://github.com/SingleZombie/LLSA

论文试图解决DiT（Diffusion Transformers）在处理长序列时因自注意力机制的平方计算复杂度而导致的可扩展性受限问题。随着图像分辨率提高，token序列增长，现有Top-K稀疏注意力方法仍面临选择过程中的二次计算开销以及需要不断增加K值以维持生成质量的问题。这是一个当前视觉生成模型向高分辨率、像素级建模扩展中的关键瓶颈问题。

提出Log-linear Sparse Attention (LLSA)，一种具有层次化结构的可训练稀疏注意力机制。其核心思想是采用分层Top-K选择策略，在多个粒度层级上逐步聚焦相关区域，并引入Hierarchical KV Enrichment机制，在减少参与计算的token数量的同时保留全局上下文信息。相比单层稀疏方法，LLSA将选择和注意力计算的复杂度从二次降至对数线性，显著提升了效率且更易于随序列增长而扩展。

LLSA在256x256像素的高分辨率图像生成任务中（无需patchification或VAE编码）实现了28.27倍的注意力推理加速和6.09倍的DiT训练加速，同时保持了生成质量。作者开发了高效的GPU实现，前向和反向传播均仅依赖稀疏索引，避免了稠密注意力掩码的使用，进一步优化性能。代码已开源：https://github.com/SingleZombie/LLSA。该工作为直接在像素空间训练长序列扩散Transformer提供了可行路径，值得在更高分辨率、视频建模等方向深入探索。

1. Swin Transformer: Hierarchical Vision Transformer using Shifted Windows 2. Sparse Transformers: Concentrated Attention Through Explicit Gating 3. Scaling Vision Transformers to Gigapixel Images via Hierarchical Attention 4. Elite Transformers: Eliminating Attention via Adaptive Token Selection 5. ToMe: Token Merging for Fast Stable Diffusion