This paper presents a new exploration into a category of diffusion models
built upon state space architecture. We endeavor to train diffusion models for
image data, wherein the traditional U-Net backbone is supplanted by a state
space backbone, functioning on raw patches or latent space. Given its notable
efficacy in accommodating long-range dependencies, Diffusion State Space Models
(DiS) are distinguished by treating all inputs including time, condition, and
noisy image patches as tokens. Our assessment of DiS encompasses both
unconditional and class-conditional image generation scenarios, revealing that
DiS exhibits comparable, if not superior, performance to CNN-based or
Transformer-based U-Net architectures of commensurate size. Furthermore, we
analyze the scalability of DiS, gauged by the forward pass complexity
quantified in Gflops. DiS models with higher Gflops, achieved through
augmentation of depth/width or augmentation of input tokens, consistently
demonstrate lower FID. In addition to demonstrating commendable scalability
characteristics, DiS-H/2 models in latent space achieve performance levels akin
to prior diffusion models on class-conditional ImageNet benchmarks at the
resolution of 256$\times$256 and 512$\times$512, while significantly reducing
the computational burden. The code and models are available at:
this https URL

本文介绍了基于状态空间架构构建扩散模型的新探索，通过训练扩散模型用于图像数据，使用状态空间骨干替代传统的 U-Net 骨干，处理原始补丁或潜在空间的图像。通过对 DiS 的评估，包括无条件和类条件的图像生成场景，发现 DiS 在与同等大小的基于 CNN 或变压器的 U-Net 架构相比具有可比甚至更好的性能。此外，我们还分析了 DiS 的可扩展性，通过 Gflops 量化前向传递复杂性进行评估。通过增加深度 / 宽度或增加输入令牌的方式获得更高 Gflops 的 DiS 模型，始终表现出更低的 FID。在潜在空间中，DiS-H/2 模型在 256x256 和 512x512 分辨率下，达到了类条件 ImageNet 基准的性能水平，并显著减轻了计算负担。