DiG：使用门控线性注意力机制的高效可扩展 Diffusion Transformer

文章预览

关注 「 AIWalker 」 并 星标 从此AI不迷路 作者丨科技猛兽     编辑丨极市平台 极市导读   在相同的模型尺寸下，DiG-XL/2 比基于 Mamba 的扩散模型在 1024 的分辨率下快 4.2 倍，在 2048 的分辨率下比带有 CUDA 优化的 FlashAttention2 的 DiT 快 1.8 倍。这些结果都证明了其优越性能。  >> 加入极市CV技术交流群，走在计算机视觉的最前沿 太长不看版 Diffusion Transformer 模型面临的一个问题是计算复杂度与序列长度呈二次方关系，这不利于扩散模型的缩放。本文通过门控线性注意力机制 (Gated Linear Attention) 的长序列建模能力来应对这个问题，来提升扩散模型的适用性。 本文提出的模型称为 Diffusion Gated Linear Attention Transformers (DiG)，是一种基于门控线性注意力机制和 DiT[1]的简单高效的扩散 Transformer 模型。除了比 DiT 更好的性能外，DiG-S/2 的训练速度比 DiT-S/2 高 2.5 倍， ………………………………