基于状态空间骨干的可扩展扩散模型
计算机视觉与模式识别
2024-03-29 v3 多媒体
摘要
本文提出了一种基于状态空间架构的新型扩散模型探索。我们致力于训练用于图像数据的扩散模型，其中传统的U-Net骨干被状态空间骨干取代，作用于原始块或潜在空间。鉴于其在处理长程依赖方面的显著有效性，扩散状态空间模型（DiS）通过将所有输入（包括时间、条件和噪声图像块）视为token而脱颖而出。我们对DiS的评估涵盖了无条件图像生成和类别条件图像生成场景，结果表明DiS在同等规模下表现出与基于CNN或Transformer的U-Net架构相当甚至更优的性能。此外，我们分析了DiS的可扩展性，以Gflops量化的前向传播复杂度衡量。通过增加深度/宽度或增加输入token实现更高Gflops的DiS模型，始终表现出更低的FID。除了展现出良好的可扩展性特征外，潜在空间中的DiS-H/2模型在256×256和512×512分辨率下的类别条件ImageNet基准测试中达到了与先前扩散模型相当的性能水平，同时显著降低了计算负担。代码和模型可在https://github.com/feizc/DiS获取。
引用
@article{arxiv.2402.05608,
title = {Scalable Diffusion Models with State Space Backbone},
author = {Zhengcong Fei and Mingyuan Fan and Changqian Yu and Junshi Huang},
journal= {arXiv preprint arXiv:2402.05608},
year = {2024}
}