LLM2D
DiG:具有门控线性注意力机制的可扩展高效扩散模型
DiG: Scalable and Efficient Diffusion Models with Gated Linear Attention
作者: Lianghui Zhu, Zilong Huang, Bencheng Liao, Jun Hao Liew, Hanshu Yan, Jiashi Feng, Xinggang Wang
发布日期: 11/28/2024
arXiv ID: oai:arXiv.org:2405.18428v2

摘要

基于大规模预训练的扩散模型在视觉内容生成领域取得了显著成功,特别是扩散Transformer(DiT)模型。然而,DiT模型面临着二次复杂度效率的挑战,尤其是在处理长序列时。本文旨在将门控线性注意力(GLA)的亚二次建模能力融入二维扩散骨干网络。具体来说,我们引入了扩散门控线性注意力Transformer(DiG),这是一种简单易用的解决方案,参数开销极小。我们提供了两种变体,即普通架构和U型架构,它们显示出优越的效率和具有竞争力的有效性。除了在256×256分辨率下比DiT和其他亚二次时间扩散模型具有更好的性能外,DiG从512分辨率开始就表现出更高的效率。具体来说,在1792分辨率下,DiG-S/2的速度是DiT-S/2的2.5倍,GPU内存节省了75.7%。此外,在1024分辨率下,DiG-XL/2的速度是基于Mamba的模型的4.2倍,在2048分辨率下,其速度是使用FlashAttention-2的DiT的1.8倍。我们将很快发布代码。代码已发布在https://github.com/hustvl/DiG。