DiG

<div align="center"> <h1>DiG: 门控线性注意力扩散模型</h1> <h3>基于门控线性注意力的可扩展高效扩散模型</h3>

朱良辉<sup>1,2</sup>,黄子龙<sup>2 :email:</sup>,廖本成<sup>1</sup>,廖俊豪<sup>2</sup>, 严汉舒<sup>2</sup>, 冯佳时<sup>2</sup>, 王兴刚<sup>1 :email:</sup>

<sup>1</sup> 华中科技大学电子信息与通信学院, <sup>2</sup> 字节跳动

(<sup>:email:</sup>) 通讯作者.

ArXiv 预印本 (arXiv 2405.18428)

</div>

新闻

2024年5月28日: 我们在Arxiv上发布了论文。代码和模型即将发布，敬请期待！☕️

摘要

基于大规模预训练的扩散模型在视觉内容生成领域取得了重大成功，尤其以扩散变换器（DiT）为代表。然而，DiT模型在可扩展性和二次复杂度效率方面面临挑战。在本文中，我们旨在利用门控线性注意力（GLA）变换器的长序列建模能力，扩展其在扩散模型中的应用。我们提出了扩散门控线性注意力变换器（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倍。所有这些结果都证明了它在最新扩散模型中的卓越效率。

概述

训练环境

Python 3.9.2
- conda create -n your_env_name python=3.9.2
torch 2.1.1 + cu118
- pip3 install torch==2.2.0 torchvision==0.17.0 torchaudio==2.2.0 --index-url https://download.pytorch.org/whl/cu121

依赖项：

# triton
pip3 install triton

# GLA
git clone https://github.com/sustcsonglin/flash-linear-attention
git checkout 36743f3f14e47f23c1ad45cf5de727dbacb5600e
cd flash-linear-attention
pip3 install -e .

# 其他
pip3 install diffusers
pip3 install tensorboard
pip3 install timm
pip3 install transformers
pip3 install accelerate
pip3 install fvcore
pip3 install opt_einsum
pip3 install torchdiffeq
pip3 install ftfy
pip3 install PyAV

训练您的DiG

在train-multi-nodes.py中设置您的VAE路径。
在scripts/dig_s_d2_in1k_256_bs256_1node.sh中设置您的DATA_PATH。
运行bash DiG/scripts/dig_s_d2_in1k_256_bs256_1node.sh no_env_install。

致谢 :heart:

本项目基于GLA（论文，代码），flash-linear-attention（代码），DiT（论文，代码），DiS（论文，代码），OpenDiT（代码）。感谢他们出色的工作。

引用

如果您发现DiG在您的研究或应用中有用，请考虑给我们一个星标🌟并使用以下BibTeX条目进行引用。

@article{dig,
      title={DiG: Scalable and Efficient Diffusion Models with Gated Linear Attention}, 
      author={Lianghui Zhu and Zilong Huang and Bencheng Liao and Jun Hao Liew and Hanshu Yan and Jiashi Feng and Xinggang Wang},
      year={2024},
      eprint={2405.18428},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}