今日应用
今日话题
DiG：使用门控线性注意力机制的高效可扩展 Diffusion Transformer
重点标签 DiG模型、门控线性注意力、扩散模型、图像生成、性能提升
文章摘要
本文介绍了一种新型的扩散模型——DiG（Diffusion Gated Linear Attention Transformers），由华中科技大学和字节跳动的研究团队共同开发。DiG模型通过引入门控线性注意力机制，有效地提高了扩散模型在图像生成任务中的性能和效率。以下是对文章的详细总结：
DiG模型概述
DiG模型的核心创新在于其采用了门控线性注意力机制，这一机制在自然语言处理领域已被证明十分有效。研究团队将其应用于视觉内容生成，通过高效的线性注意力Transformer构建扩散模型的backbone。DiG模型在保持与DiT相似的建模能力的同时，在训练速度和GPU显存成本方面展现出更高的效率。
模型架构
DiG模型结合了空间重定向和增强模块（Spatial Reorient & Enhancement Module, SREM），用于分层扫描方向控制和局部感知。这一模块包含四个基本扫描模式，使每个Patch能够感知其他Patch，同时通过深度卷积（DWConv）为模型提供局部信息。
门控线性注意力Transformer
GLA结合了数据依赖的门控机制和线性注意力，通过计算Query、Key和Value向量，并利用门控矩阵实现高效的循环建模。这一机制在DiG中被用来处理图像的DDPM训练，以适应图像生成任务。
实验结果
在ImageNet数据集上的实验表明，DiG模型在不同分辨率下均展现出卓越的性能。特别是在高分辨率图像生成方面，DiG-XL/2模型在1024分辨率下比基于Mamba的扩散模型快4.2倍，在2048分辨率下比带有CUDA优化的FlashAttention2的DiT快1.8倍。此外，DiG模型在FID等指标上也表现出更好的可扩展性。
性能提升
DiG模型在训练速度上比DiT快2.5倍，并在1792×1792分辨率下节省了75.7%的GPU显存。这些结果证明了DiG在处理大规模长序列生成任务时的优越性能，有望成为下一代扩散模型的backbone。
代码和资源
文章提供了DiG模型的论文地址和代码链接，供读者进一步学习和研究。同时，作者还分享了在ImageNet上进行class-conditional图像生成任务的训练细节和实验结果，展示了DiG生成的图像具有正确的语义和精确的空间关系。
通过上述总结，我们可以看到DiG模型在图像生成领域的潜力和优势。其高效的门控线性注意力机制和扩散模型的结合，为高质量图像生成提供了新的可能性。
文章来源
原文地址: 点我阅读全文
原文作者: 极市平台
相关文章
