Scalable Diffusion Models with State Space Backbone

本文提出了一种基于状态空间结构的扩散模型类别的新探索。我们试图训练适用于图像数据的扩散模型，其中传统的U-Net骨干被状态空间骨干所替代，该骨干可以在原始补丁或潜在空间上运行。由于其在适应长程依赖方面的显著效果，扩散状态空间模型（DiS）通过将包括时间、条件和嘈杂图像补丁在内的所有输入视为令牌而得以区分。我们对DiS的评估包括无条件和类条件图像生成场景，揭示了DiS表现出与相应大小的基于CNN或Transformer的U-Net架构相当，如果不是更好的性能。此外，我们分析了DiS的可扩展性，通过Gflops量化前向传递复杂性来衡量。通过增加深度/宽度或增加输入令牌，具有更高Gflops的DiS模型始终表现出更低的FID。除了展示值得称赞的可扩展特性外，潜在空间中的DiS-H/2模型在256×256和512×512的分辨率下，达到了与以前的扩散模型相似的类条件ImageNet基准性能水平，同时显著降低了计算负担。代码和模型可在以下网址中获得：https://github.com/feizc/DiS。

本文尝试使用基于状态空间架构的扩散模型来训练图像数据，以解决图像生成中的长程依赖问题。

文章中提出了一种新的扩散状态空间模型（DiS），将传统的U-Net骨干网络替换为状态空间骨干网络，将所有输入包括时间、条件和噪声图像补丁视为令牌，以实现无条件和类条件图像生成，并且在与同等大小的CNN或Transformer U-Net架构相比表现出相当甚至更好的性能。

实验结果表明，DiS模型在256x256和512x512分辨率下，以较低的计算负担实现了与以前扩散模型相当的性能水平。此外，文章还分析了DiS的可扩展性，通过增加深度/宽度或增加输入令牌的数量，DiS模型的Gflops增加，FID值不断降低。作者开源了代码和模型。

最近的相关研究包括：《Diffusion Models Beat GANs on Image Synthesis》、《Improved Precision and Recall Metric for Assessing Generative Models》等。