论文题目：DiM: Diffusion Mamba for Efficient High-Resolution Image Synthesis

论文题目：DiM: Diffusion Mamba for Efficient High-Resolution Image Synthesis

• 论文地址：https://arxiv.org/pdf/2405.14224
• 项目地址：http://github.com/tyshiwo1/DiM-DiffusionMamba/

摘要

扩散模型在图像生成方面取得了巨大成功，其骨干从 U-Net 发展到视觉变压器。然而，变压器的计算成本与令牌数量成二次方关系，在处理高分辨率图像时带来了显著挑战。在这项工作中，本文提出了 Diffusion Mamba (DiM)，它结合了基于状态空间模型 (SSM) 的序列模型 Mamba 的效率和扩散模型的表现力，用于高效的高分辨率图像合成。为了解决 Mamba 无法泛化到二维信号的问题，本实验进行了多项架构设计，包括多方向扫描、在每行和每列末尾的可学习填充令牌以及轻量级局部特征增强。我们的DiM 架构在高分辨率图像方面实现了推理时间效率。此外，为了进一步提高DiM 在高分辨率图像生成中的训练效率，本文研究了“弱到强”的训练策略，即先在低分辨率图像 (256 × 256) 上预训练DiM，然后在高分辨率图像 (512 × 512) 上进行微调。本实验还探索了无训练的上采样策略，使模型能够在不进一步微调的情况下生成更高分辨率的图像（例如 1024 × 1024 和 1536 × 1536）。实验表明，DiM 具有有效性和效率。

Part.01 研究贡献

1. 本文提出了一种新的基于 Mamba 的扩散模型DiM，用于高效的高分辨率图像生成。本文提出了几种有效的设计，使 Mamba 能够处理二维图像，而 Mamba 最初是为处理一维信号而设计的
2. 为了解决高分辨率图像训练成本高的问题，本文研究了对预训练在低分辨率图像上的DiM 进行微调的策略，以实现高分辨率图像生成。此外，本实验探索了无训练的上采样方案，使模型能够在不进一步微调的情况下生成更高分辨率的图像。
3. 在ImageNet 和CIFAR 上的实验表明，DiM 在高分辨率图像生成中的训练效率、推理效率和有效性。

Part.02 模型结构

DiM 框架可以处理有噪声的二维 (2D) 输入，比如图像或者 latent 的特征，同时需要输入 time step 和 class condition。这种噪声输入可以被视为由对应于输入时间步长的特定高斯噪声级别扰动的干净信号。噪声输入首先被分成 2D Patches，每个 Patch 可以通过全连接层转换为高维特征向量。接下来，这些 Patches 被送入 3×3 Depth-Wise Convolution 层，其中局部信息被注入到 Patches 中。Patches 也在行和列的末尾用可学习的 tokens 填充，允许模型在一维顺序扫描期间感知二维空间结构。然后，使用图1所示的四个扫描模式之一，将 Patch tokens 展平为 Patch 序列。time step 和 class condition 也通过全连接层转换为 tokens，然后附加到序列中。随后，序列被送入 Mamba Blocks 进行扫描。此外，本文还在浅层和深层之间添加了长跳跃连接，以将低级信息传播到高级特征，这也被证明有利于扩散模型中的像素级预测目标。

图1 DiM模型结构

Part.03 数据集

本文使用 FID-50K 作为所有数据集的评估指标。每个数据集的具体设置如下：

1. CIFAR：模型以 128 的批量大小进行无条件图像生成训练。
2. ImageNet：本文训练模型进行条件图像生成。本实验还使用无分类器指导进行评估，计算 FID 的指导权重与中的相同。在 ImageNet 256×256 上进行预训练时，本文将 DiM-Large 和 DiM-Huge 的批量大小分别设置为 1024 和 768。在 ImageNet 512×512 上微调 DiM-Huge 时，本文将批量大小设置为 240，并进行梯度累积。

Part.04 实验结果

定量实验

本文将 DiM 与其他基于Transformer和基于 SSM 的扩散模型在表1中进行了比较。在对 3.19 亿张图像样本进行训练后，DiM-Huge 在 FID-50K 上可以达到 2.40 的得分。在使用 U-ViT63% 的训练数据（3.19 亿对比 5 亿）的情况下，本文的模型性能与其他基于变压器的扩散模型相当，即在 FID-50K 上仅差约 0.1。当本文使用 4.8 亿张图像样本训练模型时，本文的模型可以超越其他模型，在 FID-50K 上达到 2.21 的得分。此外，与 DiffuSSM-XL 相比，本文基于 Mamba 的扩散模型的 Gflops 要小得多，即 DiM 需要更少的推理资源。

表1 在 ImageNet 上以 256 × 256 的分辨率进行预训练

消融实验

本文在 CIFAR-10 数据集上进行了消融研究。按照 U-ViT的方法，本文的模型在无条件生成任务中进行训练，采用 VP 调度器。本文在表2中报告了 FID 分数，其中第一行包含了性能最佳模型的结果，其他行的性能对应于没有某些组件的模型。根据结果，当比较第一行和最后两行时，本文发现多扫描方向对性能贡献最大，显示了全局感受野的重要性。本文还发现长跳跃连接有利于训练收敛，这与中的发现一致。此外，本文的两个卷积层和填充令牌也有助于性能提升。

表2 在CIFAR-10上的架构消融研究

Part.05 研究展望

本文提出了一种新的基于 Mamba 的扩散模型骨干 Diffusion Mamba (DiM)，用于高效的高分辨率图像生成。在本文的框架中，序列模型 Mamba 被用于处理二维噪声输入的补丁特征。为了使 Mamba 适应二维数据，本文提出了几种方法，包括扫描模式切换、可学习的填充令牌和轻量级局部特征增强。然后，为了高效地训练高分辨率图像样本的模型，本文提出了使用“弱到强”的训练和微调策略。实验表明，本文的模型在高分辨率图像生成方面可以达到与其他基于Transformer的扩散模型相当的性能。本文还探索了 DiM 的无训练上采样，以在不进一步微调的情况下生成更高分辨率的图像。