頂刊TPAMI2025！北大&位元組等提出IDM：可逆擴散模型賦能影像重建

2025-04-23 03:14 CVer

點選下方卡片，關注“CVer”公眾號

AI/CV重磅乾貨，第一時間送達
點選進入—>【頂會/頂刊】投稿交流群

新增微訊號：CVer2233，小助手會拉你進群！
掃描下方二維碼，加入CVer學術星球！可以獲得最新頂會/頂刊上的論文idea和CV從入門到精通資料，及最前沿應用！發論文/搞科研/漲薪，強烈推薦！

轉載自：將門創投
擴散模型作為一種強大的生成模型，已在影像重建任務中展現出巨大潛力。其基本原理是在訓練階段最佳化噪聲估計網路（如UNet），並透過去噪和加噪過程完成影像生成。然而，如何提升擴散模型在影像重建中的效能與效率，仍是一個重要挑戰。
在影像重建任務中，擴散模型面臨兩個主要問題：首先，噪聲估計任務與影像重建任務之間存在偏差，限制了重建效果；其次，推理過程需要大量迭代步驟，導致速度慢、效率低。為解決這些問題，本文提出了可逆擴散模型（IDM），透過引入端到端訓練框架和可逆網路設計，顯著提升了影像重建的效能和效率。

論文標題： 
Invertible Diffusion Models for Compressed Sensing
論文連結：
https://arxiv.org/abs/2403.17006
開原始碼：
https://github.com/Guaishou74851/IDM

一、任務背景

擴散模型作為當前非常知名且強大的生成模型之一，已在影像重建任務中展現出極大的潛力。擴散模型的基本實現方式是在訓練階段構建一個噪聲估計網路（通常是一個UNet），並在推理階段透過迭代的去噪和加噪過程完成影像生成與重建。然而，如何進一步提升擴散模型在影像重建中的效能與效率，仍然是業界探索的重點問題。
當我們將擴散模型應用於影像重建任務時，面臨兩個關鍵挑戰：
  • 挑戰一：“噪聲估計”任務與“影像重建”任務之間的偏差。 擴散模型中的深度神經網路主要針對“噪聲估計”任務（即，從當前變數中估計出噪聲）得到最最佳化，而非“影像重建”任務（即，從低質量的觀測資料中預測原始影像）本身。這可能導致其影像重建效能存在進一步提升的空間。
  • 挑戰二：推理速度慢、效率低。 儘管擴散模型能夠生成較為真實的影像，但其推理過程往往需要大量的迭代步驟，執行時間長，計算開銷大，不利於實際應用。
針對這兩個挑戰，本文提出了一種可逆擴散模型（Invertible Diffusion Models，IDM）。這一方法透過引入（1）端到端的訓練框架與（2）可逆網路設計，有效提升了影像重建的效能與效率。

二、主要貢獻

我們的方法在影像重建任務中帶來了兩個主要創新：

2.1 端到端的擴散取樣影像重建學習框架

傳統擴散模型在訓練階段的目標任務是“噪聲估計”，而實際的目標任務是“影像重建”。為了提升擴散模型的影像重建效能，我們將它的迭代取樣過程重新定義為一個整體的影像重建網路，對該網路進行端到端的訓練，突破了傳統噪聲估計學習正規化所帶來的侷限。如圖所示，透過這種方式，模型的所有引數都針對“影像重建”任務進行了最最佳化，重建效能得到大幅提升。實驗結果表明，基於Stable Diffusion的預訓練權重與這一端到端學習框架，在影像壓縮感知重建任務中，相比其他模型，我們的方法在PSNR（峰值信噪比）指標上提升了2dB，取樣步數從原本的100步降到了3步，推理速度提升了約15倍。

2.2 雙層可逆網路設計：減少記憶體開銷

大型擴散模型（如Stable Diffusion）取樣過程的端到端訓練需要佔用很大的GPU記憶體，這對於其實際應用來說是一個嚴重的瓶頸。為了減少記憶體開銷，我們提出了一種雙層可逆網路。可逆網路的核心思想是透過設計特殊的網路結構，讓網路每一層的輸出可以反向計算得到輸入。
在實踐中，我們將可逆網路應用到（1）所有擴散取樣步驟和（2）噪聲估計網路的內部，透過“佈線”技術將每個取樣步驟與其前後模組連線，形成一個雙層可逆網路。這一設計使得整個訓練過程中，程式無需儲存完整的特徵圖資料，只需儲存較少的中間變數，顯著降低了訓練模型的GPU記憶體需求。最終，這使得我們可以在視訊記憶體有限的GPU（如1080Ti）上對該模型進行端到端訓練。

三、實驗結果

3.1 影像壓縮感知重建

在影像壓縮感知重建任務中，我們的方法IDM與現有基於端到端網路和擴散模型的重建方法進行了對比。實驗結果顯示，IDM在PSNR、SSIM、FID和LPIPS等指標上取得明顯提升。

3.2 影像補全與醫學成像

在掩位元速率90%的影像補全任務中，我們的方法能夠準確恢復出窗戶等複雜結構，而傳統的擴散模型（如DDNM）無法做到這一點。此外，我們還將該方法應用於醫學影像領域，包括核磁共振成像（MRI）和計算機斷層掃描（CT）成像，取得了良好的效果。

3.3 計算成本與推理時間的最佳化

基於傳統擴散模型的影像重建方法往往需要較長的推理時間和計算開銷，而我們的可逆擴散模型IDM顯著縮短了這一過程。在重建一張256×256大小的影像時，推理時間從9秒縮短至0.63秒，大幅降低了計算開銷。與現有方法DDNM相比，IDM的訓練、推理效率和重建效能得到了顯著提升。
欲瞭解更多細節，請參考原論文。

何愷明在MIT授課的課件PPT下載

在CVer公眾號後臺回覆：何愷明，即可下載本課程的所有566頁課件PPT！趕緊學起來！

CVPR 2025 論文和程式碼下載

在CVer公眾號後臺回覆：CVPR2025，即可下載CVPR 2025論文和程式碼開源的論文合集

ECCV 2024 論文和程式碼下載

在CVer公眾號後臺回覆：ECCV2024，即可下載ECCV 2024論文和程式碼開源的論文合集
CV垂直方向和論文投稿交流群成立
掃描下方二維碼，或者新增微訊號：CVer2233，即可新增CVer小助手微信，便可申請加入CVer-垂直方向和論文投稿微信交流群。另外其他垂直方向已涵蓋：目標檢測、影像分割、目標跟蹤、人臉檢測&識別、OCR、姿態估計、超解析度、SLAM、醫療影像、Re-ID、GAN、NAS、深度估計、自動駕駛、強化學習、車道線檢測、模型剪枝&壓縮、去噪、去霧、去雨、風格遷移、遙感影像、行為識別、影片理解、影像融合、影像檢索、論文投稿&交流、PyTorch、TensorFlow和Transformer、NeRF、3DGS、Mamba等。
一定要備註：研究方向+地點+學校/公司+暱稱（如Mamba、多模態學習或者論文投稿+上海+上交+卡卡），根據格式備註，可更快被透過且邀請進群

▲掃碼或加微訊號: CVer2233，進交流群
CVer計算機視覺（知識星球）來了！想要了解最新最快最好的CV/DL/AI論文速遞、優質實戰專案、AI行業前沿、從入門到精通學習教程等資料，歡迎掃描下方二維碼，加入CVer計算機視覺（知識星球），已彙集上萬人！

▲掃碼加入星球學習
▲點選上方卡片，關注CVer公眾號
整理不易，請點贊和在看


相關文章