DiM: Diffusion Mamba for Efficient High-Resolution Image Synthesis

扩散模型在图像生成方面取得了巨大成功，其骨干从U-Net发展到视觉Transformer。然而，Transformer的计算成本与令牌数量的平方成正比，处理高分辨率图像时面临重大挑战。在本文中，我们提出了Diffusion Mamba（DiM），将基于状态空间模型（SSM）的序列模型Mamba的效率与扩散模型的表达能力相结合，实现了高效的高分辨率图像合成。为了解决Mamba无法推广到二维信号的挑战，我们进行了多方向扫描、在每行和每列的末尾使用可学习的填充令牌以及轻量级局部特征增强等多种架构设计。我们的DiM架构实现了高分辨率图像的推理时间效率。此外，为了进一步提高使用DiM进行高分辨率图像生成的训练效率，我们研究了“弱到强”的训练策略，即在低分辨率图像（$256\times 256$）上预训练DiM，然后在高分辨率图像（$512 \times 512$）上进行微调。我们进一步探索了无需训练的上采样策略，以使模型能够生成更高分辨率的图像（例如$1024\times 1024$和$1536\times 1536$），而无需进一步微调。实验证明了我们的DiM的有效性和效率。

本论文旨在解决高分辨率图像生成中Transformers模型计算成本高的问题，提出了结合State Space Models（SSM）和Diffusion Models的DiM架构，以提高高分辨率图像生成的推理效率。

DiM架构结合了SSM和Diffusion Models的优势，通过多方向扫描、可学习填充令牌、轻量级局部特征增强等设计解决了SSM无法推广到二维信号的问题，并探索了“弱到强”的训练策略和无需微调的上采样策略，提高了训练和生成效率。

实验结果表明DiM架构在高分辨率图像生成的推理时间和生成效率方面具有优势，同时还探索了“弱到强”的训练策略和无需微调的上采样策略。论文使用了多个数据集进行实验，并开源了代码。

相关研究包括U-Net和Vision Transformers等模型在图像生成中的应用，以及其他一些结合SSM和Diffusion Models的图像生成模型，如SSDGM和D-VAE等。