本篇文章來自公眾號粉絲投稿，論文提出了一種可逆擴散模型（Invertible Diffusion Models，IDM）。這一方法通過引入（1）端到端的訓練框架與（2）可逆網絡設計，有效提升了圖像重建的性能與效率。

一、論文信息

• 論文標題：Invertible Diffusion Models for Compressed Sensing
• 論文作者：Bin Chen（陳斌）, Zhenyu Zhang（張振宇）, Weiqi Li（李瑋琦）, Chen Zhao（趙琛）, Jiwen Yu（余濟聞）, Shijie Zhao（趙世杰）, Jie Chen（陳杰） and Jian Zhang（張健）
• 作者單位：北京大學信息工程學院、阿卜杜拉國王科技大學、字節(jié)跳動
• 發(fā)表刊物：IEEE Transactions on Pattern Analysis and Machine Intelligence (TPAMI)
• 發(fā)表時間：2025年2月5日
• 正式版本：https://ieeexplore.ieee.org/document/10874182
• ArXiv版本：https://arxiv.org/abs/2403.17006
• 開源代碼：https://github.com/Guaishou74851/IDM

二、任務背景

擴散模型作為當前非常知名且強大的生成模型之一，已在圖像重建任務中展現(xiàn)出極大的潛力。擴散模型的基本實現(xiàn)方式是在訓練階段構建一個噪聲估計網絡（通常是一個UNet），并在推理階段通過迭代的去噪和加噪過程完成圖像生成與重建。然而，如何進一步提升擴散模型在圖像重建中的性能與效率，仍然是業(yè)界探索的重點問題。

當我們將擴散模型應用于圖像重建任務時，面臨兩個關鍵挑戰(zhàn)：

• 挑戰(zhàn)一：“噪聲估計”任務與“圖像重建”任務之間的偏差。擴散模型中的深度神經網絡主要針對“噪聲估計”任務（即，從當前變量中估計出噪聲）得到最優(yōu)化，而非“圖像重建”任務（即，從低質量的觀測數(shù)據(jù)中預測原始圖像）本身。這可能導致其圖像重建性能存在進一步提升的空間。
• 挑戰(zhàn)二：推理速度慢、效率低。盡管擴散模型能夠生成較為真實的圖像，但其推理過程往往需要大量的迭代步驟，運行時間長，計算開銷大，不利于實際應用。

針對這兩個挑戰(zhàn)，本文提出了一種可逆擴散模型（Invertible Diffusion Models，IDM）。這一方法通過引入（1）端到端的訓練框架與（2）可逆網絡設計，有效提升了圖像重建的性能與效率。

三、主要貢獻

我們的方法在圖像重建任務中帶來了兩個主要創(chuàng)新：

1.端到端的擴散采樣圖像重建學習框架

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

2.雙層可逆網絡設計：減少內存開銷

大型擴散模型（如Stable Diffusion）采樣過程的端到端訓練需要占用很大的GPU內存，這對于其實際應用來說是一個嚴重的瓶頸。為了減少內存開銷，我們提出了一種雙層可逆網絡。可逆網絡的核心思想是通過設計特殊的網絡結構，讓網絡每一層的輸出可以反向計算得到輸入。在實踐中，我們將可逆網絡應用到（1）所有擴散采樣步驟和（2）噪聲估計網絡的內部，通過“布線”技術將每個采樣步驟與其前后模塊連接，形成一個雙層可逆網絡。這一設計使得整個訓練過程中，程序無需存儲完整的特征圖數(shù)據(jù)，只需存儲較少的中間變量，顯著降低了訓練模型的GPU內存需求。最終，這使得我們可以在顯存有限的GPU（如1080Ti）上對該模型進行端到端訓練。

四、實驗結果

圖像壓縮感知重建

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

圖像補全與醫(yī)學成像

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

計算成本與推理時間的優(yōu)化

基于傳統(tǒng)擴散模型的圖像重建方法往往需要較長的推理時間和計算開銷，而我們的可逆擴散模型IDM顯著縮短了這一過程。在重建一張256×256大小的圖像時，推理時間從9秒縮短至0.63秒，大幅降低了計算開銷。與現(xiàn)有方法DDNM相比，IDM的訓練、推理效率和重建性能得到了顯著提升。

欲了解更多細節(jié)，請參考原論文。

五、作者簡介

• 陳斌：北京大學信息工程學院博士生，主要研究方向是圖像壓縮感知與超分辨率。
• 張振宇：北京大學信息工程學院碩士生，主要研究方向是圖像重建。
• 李瑋琦：北京大學信息工程學院博士生，主要研究方向是圖像壓縮感知與超分辨率。
• 趙琛：沙特阿卜杜拉國王科技大學（KAUST）的研究科學家，圖像與視頻理解實驗室（IVUL）視頻分析課題組組長。她首次提出了可逆化預訓練神經網絡方法，實現(xiàn)了大型預訓練模型的極低顯存微調。她在這一領域的代表工作包括Re2TAL（CVPR 2023）、Dr2Net （CVPR 2024）等。

• 個人主頁：https://zhao-chen.com/。

• 余濟聞：北京大學信息工程學院碩士生，主要研究方向是生成式擴散模型。
• 趙世杰：字節(jié)跳動多媒體實驗室研究員，負責視頻處理與增強課題組。
• 陳杰：北京大學信息工程學院副教授，主要研究方向是計算機視覺與模式識別和AI4Science。

• 個人主頁：https://aimia-pku.github.io

• 張健：北京大學信息工程學院副教授，主要研究方向是視覺內容重建與生成、AIGC內容鑒偽和版權保護。

• 個人主頁：https://jianzhang.tech/cn