Diffusion Mamba:用线性计算打造高效高分辨率图像生成新范式

人工智能
目前的扩散模型主要依赖U-Net或Vision Transformer(ViT)作为骨干架构。然而,Transformer 的计算复杂度随 token 数量呈二次增长,使得高分辨率图像生成变得极为昂贵。如何在保证生成质量的前提下,提高推理效率,降低计算成本? 这是 DiM 试图解决的核心问题。

一眼概览

Diffusion Mamba (DiM) 是一种结合状态空间模型(SSM)和扩散模型的新架构,旨在实现高效的高分辨率图像生成。相比于基于 Transformer 的扩散模型,DiM 具有更优的计算效率,特别是在超高分辨率图像(1024×1024 及以上)生成任务中表现突出。

核心问题

目前的扩散模型主要依赖U-NetVision 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 生成或文本生成?

责任编辑:武晓燕 来源: 萍哥学AI
相关推荐

2012-01-05 16:08:57

佳能激光打印机

2023-02-16 07:30:55

图形编辑器屏幕开发

2011-11-14 13:30:27

惠普扫描仪

2020-09-21 16:11:04

Adobe 软件平台

2022-12-05 15:27:52

VR技术

2012-11-02 13:18:11

笔记本

2024-06-05 09:26:50

2023-10-27 08:30:52

传输接口刷新率

2025-02-26 13:23:08

2011-10-31 17:12:42

激光打印机评测

2012-01-17 10:35:31

惠普扫描仪

2011-10-28 15:52:22

激光打印机评测

2018-11-16 09:50:46

Windows 10高分辨率屏幕截图

2020-06-30 14:08:57

算法模型deepfake

2022-03-23 10:44:02

图像训练框架

2018-07-19 15:00:41

卫星图像

2022-11-01 08:00:00

2025-01-21 09:50:00

模型生成AI

2011-04-28 13:09:49

2025-02-26 11:14:10

点赞
收藏

51CTO技术栈公众号