Diffusion Mamba：用線性計算打造高效高分辨率圖像生成新范式
一眼概覽
Diffusion Mamba (DiM) 是一種結合狀態(tài)空間模型（SSM）和擴散模型的新架構，旨在實現(xiàn)高效的高分辨率圖像生成。相比于基于 Transformer 的擴散模型，DiM 具有更優(yōu)的計算效率，特別是在超高分辨率圖像（1024×1024 及以上）生成任務中表現(xiàn)突出。
核心問題
目前的擴散模型主要依賴U-Net或Vision Transformer（ViT）作為骨干架構。然而，Transformer 的計算復雜度隨 token 數(shù)量呈二次增長，使得高分辨率圖像生成變得極為昂貴。如何在保證生成質量的前提下，提高推理效率，降低計算成本？ 這是 DiM 試圖解決的核心問題。
技術亮點
1. Mamba 適配 2D 圖像建模：
? 采用 多方向掃描，避免單向序列建模的局限性；
? 引入 可學習填充 token，在行列轉換時保留空間連續(xù)性；
? 結合 輕量級局部特征增強模塊，彌補 SSM 對局部結構的捕捉能力不足。
2. 高效訓練策略：
? 采用 “弱到強”訓練策略，先在低分辨率（256×256）上訓練，再遷移至高分辨率（512×512），減少計算成本；
? 無訓練超分辨率（training-free upsampling） 方案，使模型可直接生成 1024×1024 和 1536×1536 級別的圖像。
3. 計算復雜度優(yōu)化：
? 采用 Mamba 替代 Transformer 的自注意力機制，使計算復雜度由 O(n2) 降至 O(n)，在超高分辨率場景下更具優(yōu)勢；
? 在 1280×1280 及以上分辨率下，推理速度比 Transformer 快 2.2 倍。
方法框架
圖片
DiM 采用 Mamba 作為擴散模型的骨干架構，其關鍵步驟如下：
1. 圖像 Token 化：
? 輸入噪聲圖像/潛變量，拆分為2D patch，并加入時間步、類別信息；
? 通過 3×3 深度可分卷積 進行局部特征增強。
2. Mamba 序列建模：
? 采用 四種掃描模式（行優(yōu)先、列優(yōu)先、反向行優(yōu)先、反向列優(yōu)先），使 token 擁有全局感受野；
? 在行末、列末插入可學習填充 token，保持空間連續(xù)性；
? 通過長跳躍連接（long skip connections）提升多尺度特征融合。
3. 訓練與推理：
? 低分辨率預訓練（256×256），再微調至 512×512；
? 采用 無訓練超分辨率，可在 512×512 訓練后，直接生成 1024×1024 及以上分辨率圖像。
實驗結果速覽
1. 圖像質量（FID 指標）
? CIFAR-10 數(shù)據(jù)集：DiM-Small 取得 FID = 2.92，優(yōu)于 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 迭代，F(xiàn)ID 降至 3.78。
2. 推理效率
? DiM 在 1280×1280 及以上分辨率，比 Transformer 快 2.2 倍；
? 比 Mamba 原生模型僅慢 1.4 倍，證明其設計在計算效率上幾乎無額外損耗。
3. 超高分辨率生成
? 無訓練超分辨率 方案可在 512×512 訓練后直接生成 1024×1024 和 1536×1536 級別圖像；
? 但仍存在細節(jié)塌縮、重復模式的問題，特別是人臉細節(jié)易失真。
實用價值與應用
DiM 作為一種高效的高分辨率圖像生成模型，適用于多個場景：
? 藝術創(chuàng)作 & 設計：高質量圖像合成，提高生產力；
? 自動駕駛 & 監(jiān)控：高效解析高分辨率視覺數(shù)據(jù)；
? 生物醫(yī)學影像：低成本生成高質量醫(yī)學圖像；
? 游戲 & 元宇宙：低延遲、高質量的虛擬場景生成。
開放問題
1. 如何改進無訓練超分辨率策略，減少高分辨率圖像的細節(jié)塌縮？
2. DiM 結構能否適用于視頻生成，替代 Transformer 作為 backbone？
3. Mamba 的線性復雜度特性，是否可以應用于其他生成任務，如 3D 生成或文本生成？
