DiG: Scalable and Efficient Diffusion Models with Gated Linear Attention

具有大规模预训练的扩散模型在视觉内容生成领域取得了显著的成功，尤其是Diffusion Transformers（DiT）模型。然而，DiT模型在可扩展性和二次复杂度效率方面面临挑战。本文旨在利用门控线性注意力（GLA）变压器的长序列建模能力，将其应用于扩散模型。我们引入了Diffusion Gated Linear Attention Transformers（DiG），这是一种简单、可采用的解决方案，参数开销最小，遵循DiT设计，但提供了更高的效率和有效性。除了比DiT表现更好外，DiG-S/2的训练速度比DiT-S/2快2.5倍，在分辨率为1792×1792时，节省了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。

本论文旨在解决视觉内容生成中的扩散模型的可扩展性和二次复杂度效率问题。

论文提出了一种采用门控线性注意力变换（GLA）转换器的扩散模型，称为Diffusion Gated Linear Attention Transformers（DiG），其具有比Diffusion Transformers（DiT）更高的效率和效果。

实验表明，相比于DiT，DiG-S/2的训练速度提高了2.5倍，且在分辨率为1792x1792时节省了75.7％的GPU内存。 DiG模型具有良好的可扩展性，增加深度/宽度或输入令牌的增强，会使FID不断降低。与其他次二次时间复杂度的扩散模型相比，DiG-XL/2在相同的模型大小下比最近的基于Mamba的扩散模型快4.2倍，在1024分辨率下比DiT快1.8倍。

近期相关研究包括Mamba-based diffusion model和CUDA-optimized FlashAttention-2。