先分组后智能注意：扩散变换器无训练加速方法

扩散变换器展现出惊人的生成能力，但其高计算成本阻碍了实际部署，例如在 A100 GPU 上生成 $8192\times8192$ 图像需超过一小时。本文提出 GRAT（\textbf{GR}ouping first, \textbf{AT}tending smartly），一种无需训练的注意力加速策略，用于快速生成图像和视频，同时不牺牲输出质量。关键思想是利用预训练扩散变换器中学习到的注意力图的内在稀疏性（其倾向于局部聚焦），并借助更好的 GPU 并行度。具体而言，GRAT 首先将相邻的 token 分组为非重叠的组，这既符合 GPU 执行模式，也符合预训练生成式变换器中学习到的局部注意力结构。随后，通过让同一组内的所有查询 token 共享一组公共可注意的键和值 token 来加速注意力计算。这些键和值 token 进一步限制在结构化区域（如周围块或十字交叉区域）内，显著降低计算开销（例如在生成 $8192\times8192$ 图像时实现 \textbf{35.8$\times$} 的加速），同时保留essential注意力模式和长程语境。我们在预训练的 Flux 和 HunyuanVideo 上验证了 GRAT，分别用于图像和视频生成。在两种情况下，GRAT 在无需微调的情况下实现了显著更快的推理速度，同时保持了与完整注意力的性能。我们希望 GRAT 将激发未来研究者加速扩散变换器以实现可扩展视觉生成的热情。