U-Shape Mamba: State Space Model for faster diffusion

📄 arXiv: 2504.13499v2 📥 PDF

作者: Alex Ergasti, Filippo Botti, Tomaso Fontanini, Claudio Ferrari, Massimo Bertozzi, Andrea Prati

分类: cs.CV

发布日期: 2025-04-18 (更新: 2025-04-26)

备注: Accepted at CVPR 2025 eLVM workshop. The code is here: https://github.com/ErgastiAlex/U-Shape-Mamba

DOI: 10.1109/CVPRW67362.2025.00307

💡 一句话要点

提出U型Mamba（USM），加速扩散模型并提升图像生成质量。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 扩散模型 Mamba模型 U-Net 图像生成 序列建模

📋 核心要点

  1. 扩散模型计算成本高昂，限制了其在资源受限环境中的应用。
  2. USM利用Mamba模型构建U-Net结构，通过序列长度的缩减和恢复，降低计算复杂度。
  3. 实验表明，USM在降低GFlops和内存需求的同时，显著提升了图像生成质量（FID指标）。

📝 摘要（中文）

扩散模型已成为高质量图像生成的主流方法，但其高计算成本仍然是一个重大挑战。为了解决这个问题，我们提出了一种新的扩散模型U型Mamba（USM），它在类似U-Net的分层结构中利用了基于Mamba的层。通过在编码器中逐步减少序列长度，并在解码器中通过Mamba块恢复序列长度，USM在保持强大生成能力的同时，显著降低了计算开销。与目前最有效的基于Mamba的扩散模型Zigma相比，实验结果表明，USM实现了三分之一的GFlops，需要更少的内存并且速度更快，同时在图像质量方面优于Zigma。在AFHQ、CelebAHQ和COCO数据集上，Frechet Inception Distance (FID) 分别提高了 15.3、0.84 和 2.7 个点。这些发现表明，USM 是一种高效且可扩展的扩散生成模型解决方案，使研究界更容易获得高质量的图像合成，同时降低了计算成本。

🔬 方法详解

问题定义：扩散模型在图像生成领域取得了显著成果，但其计算复杂度高，推理速度慢，阻碍了其更广泛的应用。现有基于Mamba的扩散模型虽然在一定程度上降低了计算量，但仍有改进空间。

核心思路：USM的核心思路是借鉴U-Net的结构，利用Mamba块作为基本构建单元，通过在编码器中逐步缩减序列长度，在解码器中逐步恢复序列长度，从而在保证生成质量的同时，显著降低计算复杂度。这种设计能够有效地捕捉图像的多尺度信息，并减少Mamba模型需要处理的序列长度。

技术框架：USM采用U-Net的编码器-解码器结构。编码器部分由多个Mamba块和下采样层组成，逐步提取图像特征并缩减序列长度。解码器部分由多个Mamba块和上采样层组成，逐步恢复图像分辨率并生成最终图像。跳跃连接（skip connection）被用于连接编码器和解码器中对应层级的特征图，以保留更多的细节信息。

关键创新：USM的关键创新在于将Mamba模型与U-Net结构相结合，并针对扩散模型进行了优化。通过在U-Net的编码器和解码器中使用Mamba块，USM能够有效地处理图像序列数据，并降低计算复杂度。此外，USM还采用了特定的训练策略和参数设置，以进一步提升生成质量。

关键设计：USM的具体实现细节包括Mamba块的配置、下采样和上采样的方式、跳跃连接的实现以及损失函数的选择。论文中可能详细描述了这些参数的具体数值和设计原则，例如Mamba块中状态变量的维度、下采样和上采样采用的具体算子、以及损失函数中各项的权重等（具体细节未知）。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

USM在图像生成任务上取得了显著的性能提升。与Zigma相比，USM实现了三分之一的GFlops，需要更少的内存并且速度更快，同时在图像质量方面优于Zigma。在AFHQ、CelebAHQ和COCO数据集上，Frechet Inception Distance (FID) 分别提高了 15.3、0.84 和 2.7 个点。这些结果表明USM在计算效率和生成质量方面都具有优势。

🎯 应用场景

U-Shape Mamba (USM) 有潜力广泛应用于图像生成、图像修复、图像编辑等领域。其高效的计算特性使其能够在资源受限的设备上运行，例如移动设备或嵌入式系统。此外，USM还可以用于生成高质量的合成数据，以训练其他机器学习模型，从而提高模型的泛化能力。未来，USM有望推动扩散模型在更多实际场景中的应用。

📄 摘要（原文）

Diffusion models have become the most popular approach for high-quality image generation, but their high computational cost still remains a significant challenge. To address this problem, we propose U-Shape Mamba (USM), a novel diffusion model that leverages Mamba-based layers within a U-Net-like hierarchical structure. By progressively reducing sequence length in the encoder and restoring it in the decoder through Mamba blocks, USM significantly lowers computational overhead while maintaining strong generative capabilities. Experimental results against Zigma, which is currently the most efficient Mamba-based diffusion model, demonstrate that USM achieves one-third the GFlops, requires less memory and is faster, while outperforming Zigma in image quality. Frechet Inception Distance (FID) is improved by 15.3, 0.84 and 2.7 points on AFHQ, CelebAHQ and COCO datasets, respectively. These findings highlight USM as a highly efficient and scalable solution for diffusion-based generative models, making high-quality image synthesis more accessible to the research community while reducing computational costs.