Diffusion models with large-scale pre-training have achieved significant success in the field of visual content generation, particularly exemplified by Diffusion Transformers (DiT). However, DiT models have faced challenges with scalability and quadratic complexity efficiency. In this paper, we aim to leverage the long sequence modeling capability of Gated Linear Attention (GLA) Transformers, expanding its applicability to diffusion models. We introduce Diffusion Gated Linear Attention Transformers (DiG), a simple, adoptable solution with minimal parameter overhead, following the DiT design, but offering superior efficiency and effectiveness. In addition to better performance than DiT, DiG-S/2 exhibits $2.5\times$ higher training speed than DiT-S/2 and saves $75.7\%$ GPU memory at a resolution of $1792 \times 1792$. Moreover, we analyze the scalability of DiG across a variety of computational complexity. DiG models, with increased depth/width or augmentation of input tokens, consistently exhibit decreasing FID. We further compare DiG with other subquadratic-time diffusion models. With the same model size, DiG-XL/2 is $4.2\times$ faster than the recent Mamba-based diffusion model at a $1024$ resolution, and is $1.8\times$ faster than DiT with CUDA-optimized FlashAttention-2 under the $2048$ resolution. All these results demonstrate its superior efficiency among the latest diffusion models. Code is released at https://github.com/hustvl/DiG.
翻译:大规模预训练扩散模型在视觉内容生成领域取得了显著成功,以扩散Transformer(DiT)为典型代表。然而,DiT模型在可扩展性和二次计算复杂度效率方面面临挑战。本文旨在利用门控线性注意力(GLA)Transformer的长序列建模能力,将其扩展应用于扩散模型。我们提出了扩散门控线性注意力Transformer(DiG),这是一种遵循DiT设计、参数开销极小的简洁可适配方案,同时具备更优的效率与性能。除性能优于DiT外,DiG-S/2在1792×1792分辨率下训练速度比DiT-S/2提升2.5倍,并节省75.7%的GPU显存。此外,我们系统分析了DiG在不同计算复杂度下的可扩展性。通过增加深度/宽度或扩展输入标记,DiG模型始终呈现FID指标持续下降的趋势。我们进一步将DiG与其他亚二次时间复杂度的扩散模型进行对比。在相同模型规模下,DiG-XL/2在1024分辨率下比近期基于Mamba的扩散模型快4.2倍,在2048分辨率下比采用CUDA优化FlashAttention-2的DiT快1.8倍。所有结果均表明其在当前前沿扩散模型中具有卓越的效率优势。代码已发布于https://github.com/hustvl/DiG。