一眼概览
Diffusion Mamba (DiM) 是一种结合状态空间模型（SSM）和扩散模型的新架构，旨在实现高效的高分辨率图像生成。相比于基于 Transformer 的扩散模型，DiM 具有更优的计算效率，特别是在超高分辨率图像（1024×1024 及以上）生成任务中表现突出。
核心问题
目前的扩散模型主要依赖U-Net或Vision Transformer（ViT）作为骨干架构。然而，Transformer 的计算复杂度随 token 数量呈二次增长，使得高分辨率图像生成变得极为昂贵。如何在保证生成质量的前提下，提高推理效率，降低计算成本？ 这是 DiM 试图解决的核心问题。
技术亮点
1. Mamba 适配 2D 图像建模：
• 采用 多方向扫描，避免单向序列建模的局限性；
• 引入 可学习填充 token，在行列转换时保留空间连续性；
• 结合 轻量级局部特征增强模块，弥补 SSM 对局部结构的捕捉能力不足。
2. 高效训练策略：
• 采用 “弱到强”训练策略，先在低分辨率（256×256）上训练，再迁移至高分辨率（512×512），减少计算成本；
• 无训练超分辨率（training-free upsampling） 方案，使模型可直接生成 1024×1024 和 1536×1536 级别的图像。
3. 计算复杂度优化：
• 采用 Mamba 替代 Transformer 的自注意力机制，使计算复杂度由 O(n²) 降至 O(n)，在超高分辨率场景下更具优势；
• 在 1280×1280 及以上分辨率下，推理速度比 Transformer 快 2.2 倍。
方法框架
图片
DiM 采用 Mamba 作为扩散模型的骨干架构，其关键步骤如下：
1. 图像 Token 化：
• 输入噪声图像/潜变量，拆分为2D patch，并加入时间步、类别信息；
• 通过 3×3 深度可分卷积 进行局部特征增强。
2. Mamba 序列建模：
• 采用 四种扫描模式（行优先、列优先、反向行优先、反向列优先），使 token 拥有全局感受野；
• 在行末、列末插入可学习填充 token，保持空间连续性；
• 通过长跳跃连接（long skip connections）提升多尺度特征融合。
3. 训练与推理：
• 低分辨率预训练（256×256），再微调至 512×512；
• 采用 无训练超分辨率，可在 512×512 训练后，直接生成 1024×1024 及以上分辨率图像。
实验结果速览
1. 图像质量（FID 指标）
• CIFAR-10 数据集：DiM-Small 取得 FID = 2.92，优于 U-ViT-S（FID = 3.11）。
• ImageNet 256×256 训练：
a.DiM-Huge 在 319M 训练样本下 FID = 2.40；
b.进一步训练到 480M 样本后，超越 DiffuSSM-XL，取得 FID = 2.21。
• ImageNet 512×512 训练：
• 仅使用 15M 高分辨率训练样本，DiM-Huge 取得 FID = 3.94；
• 进一步训练至 110K 迭代，FID 降至 3.78。
2. 推理效率
• DiM 在 1280×1280 及以上分辨率，比 Transformer 快 2.2 倍；
• 比 Mamba 原生模型仅慢 1.4 倍，证明其设计在计算效率上几乎无额外损耗。
3. 超高分辨率生成
• 无训练超分辨率 方案可在 512×512 训练后直接生成 1024×1024 和 1536×1536 级别图像；
• 但仍存在细节塌缩、重复模式的问题，特别是人脸细节易失真。
实用价值与应用
DiM 作为一种高效的高分辨率图像生成模型，适用于多个场景：
• 艺术创作 & 设计：高质量图像合成，提高生产力；
• 自动驾驶 & 监控：高效解析高分辨率视觉数据；
• 生物医学影像：低成本生成高质量医学图像；
• 游戏 & 元宇宙：低延迟、高质量的虚拟场景生成。
开放问题
1. 如何改进无训练超分辨率策略，减少高分辨率图像的细节塌缩？
2. DiM 结构能否适用于视频生成，替代 Transformer 作为 backbone？
3. Mamba 的线性复杂度特性，是否可以应用于其他生成任务，如 3D 生成或文本生成？
