LearnIR通過訓練輕量網絡預測梯度校正分布，實現無需前向算子的擴散后驗采樣校正；并設計動態分辨率模塊，進一步抑制噪聲。

作者：vivo BlueImage Lab

01

論文主要是針對什么問題？

本文主要針對真實世界圖像復原問題，即從受到復雜退化（如霧霾、陰影、噪聲、運動模糊等）影響的圖像中恢復出高質量、高保真的清晰圖像。

問題出現的背景

• 真實世界中的成像環境往往引入多種異質退化，且這些退化常常同時出現并相互交織，使得圖像復原成為一個經典的病態逆問題。

• 現有基于擴散模型的圖像復原方法存在三類核心限制：

• 條件生成方法：難以在忠實復原和真實生成之間取得平衡；

• 基于反演的方法：將退化圖像反演到潛空間的過程中會累積誤差，導致與輸入明顯偏差，且效率低下；

• 后驗采樣方法（如DPS）：需要精確已知的前向測量算子 A（例如高斯模糊核、隨機掩碼等），但在真實場景中該算子通常不可獲得，嚴重限制了實際應用。

02

核心貢獻與效果概覽

2.1 核心貢獻

1. 可學習的擴散后驗采樣框架（LearnIR）：提出通過訓練輕量級網絡直接預測后驗采樣中的梯度校正項分布，無需已知前向退化算子即可實現擴散后驗采樣校正，從根本上突破了傳統DPS方法的關鍵限制。

2. 擴散后驗采樣校正（DPSC）：利用高斯分布的封閉性，證明了前向過程真實后驗與模型預測反向分布之間的偏差服從高斯分布，進而可以通過訓練一個輕量網絡來擬合該偏差的均值，作為即插即用的正則化項校正擴散軌跡，消除采樣過程中的結構偏差和色偏等不一致性。

3. 動態分辨率模塊（DRM）：設計了時間依賴的動態分辨率調度策略，在像素空間中實現"從粗到細"的采樣過程——高噪聲階段使用低分辨率捕獲全局上下文，低噪聲階段恢復高分辨率精修紋理細節，無需預訓練VAE即可簡化端到端流水線并降低計算開銷。

2.2 關鍵理論

論文的核心定理（Theorem 1）證明：在DRM潛空間中，DPS梯度正比于模型預測的反向分布與真實前向后驗之間的偏差：

利用高斯分布的封閉性，該偏差可以建模為：

其中均值 μ 和方差 σ2 均有解析閉式解。通過訓練網絡 μ_θ 去擬合解析均值 μ，可以有效地引導采樣軌跡與真實后驗對齊。

2.3 效果預覽

本文公式推導比較多，想了解細節的同學可以直接看原文附錄推導過程，先預覽下效果吧： 第一排是原始圖，第二排是對應處理后的圖：

在去霧和去陰影數據集上和一些其他模型的對比效果：

03

論文提出的方法是什么？

LearnIR的總體流程圖如上所示，整體框架由兩個互補模塊組成：

3.1 動態分辨率模塊（DRM）

• 定義時間依賴的縮放因子 s(t)，在不同擴散時間步將圖像映射到不同分辨率的潛空間：

• • 早期階段（ ）：對圖像進行大尺度下采樣（ ），聚焦全局結構建模；

  • 后期階段（ ）：恢復原始分辨率（ ），精修高頻紋理細節。

• 使用高效的非可訓練雙線性插值實現，無需預訓練VAE，顯著降低計算成本。
  

• 在標準去噪損失之外引入一致性正則化項：

• • 去噪損失：約束噪聲預測網絡 ε θ 準確估計殘差噪聲；

  • 一致性損失：約束校正網絡 μ θ 擬合前向-反向后驗偏差的解析均值。

• 總損失函數： λ

• 推理時，DPSC 作為即插即用模塊，在每個采樣步驟通過 μ θ 預測梯度校正，自適應修正擴散軌跡。

3.3 訓練與推理

• 訓練采用兩階段策略：

• Stage 1：固定分辨率訓練（DRM關閉），聚焦學習DPSC梯度校正；

• Stage 2：開啟DRM，以更小學習率在動態分辨率下微調。

• 推理基于殘差擴散的平滑等效變換確定穩定采樣起點 T'，僅需5步采樣即可生成高質量結果。

04

實驗結果展示

4.1 數據集與設置

實驗在5個數據集上進行：ISTD（陰影去除）、O-HAZE/HazyDet/REVIDE（去霧）以及新構建的FaceShadow數據集（人臉陰影去除，含30,000對合成數據 + 1,000對真實數據）。所有評估在單張A100 GPU上完成，采樣步數僅為5步。

4.2 陰影去除（ISTD數據集）

LearnIR在mask-based方法中取得最佳表現，與mask-free最優方法相比也具有競爭力。

4.3 去霧任務（O-HAZE / HazyDet / REVIDE）

LearnIR在三個去霧數據集上全面超越所有對比方法，在O-HAZE上PSNR提升 +2.27 dB，在HazyDet上PSNR提升 +1.65 dB 且SSIM提升 +0.124。

4.4 人臉陰影去除（自建的FaceShadow數據集）

LearnIR在合成和真實人臉陰影數據上均大幅領先，PSNR分別提升 +2.44 dB 和 +1.71 dB。

4.5 消融實驗

消融實驗驗證了DPSC和DRM兩個模塊的有效性：

• 去除DPSC導致PSNR下降 4.4 dB，說明后驗采樣校正對消除軌跡不一致至關重要；

• 去除DRM導致PSNR下降 1.27 dB，驗證了動態分辨率策略對全局結構保持的重要性；

• 同時去除兩者，性能大幅下降至22.86 dB，證明兩個模塊協同配合才能達到最優效果。

4.6 計算效率

DRM使用非可訓練的雙線性插值，計算開銷幾乎為零。完整模型僅需5步采樣，總推理時間約1.6秒。

05

關于作者：vivo BlueImage Lab

藍圖影像創新實驗室，主要負責移動影像算法創新，包括圖像/視頻處理、圖像/視頻交互、圖像/視頻增強、多模態理解大模型等方面的技術前沿探索。

致力于不斷提升vivo移動影像的算法能力，使用戶能夠拍攝出更加清晰、美觀的照片和視頻。同時積極探索增強現實、具身智能等新興技術領域的應用，努力為用戶提供更加豐富和便捷的影像體驗。

未經「AI科技評論」授權，嚴禁以任何方式在網頁、論壇、社區進行轉載！

公眾號轉載請先在「AI科技評論」后臺留言取得授權，轉載時需標注來源并插入本公眾號名片。