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

一、論文信息

• 論文標(biāo)題：Invertible Diffusion Models for Compressed Sensing
• 論文作者：Bin Chen（陳斌）, Zhenyu Zhang（張振宇）, Weiqi Li（李瑋琦）, Chen Zhao（趙?。? Jiwen Yu（余濟聞）, Shijie Zhao（趙世杰）, Jie Chen（陳杰） and Jian Zhang（張?。?/span>
• 作者單位：北京大學(xué)信息工程學(xué)院、阿卜杜拉國王科技大學(xué)、字節(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

二、任務(wù)背景

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

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

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

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

三、主要貢獻

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

1.端到端的擴散采樣圖像重建學(xué)習(xí)框架

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

2.雙層可逆網(wǎng)絡(luò)設(shè)計：減少內(nèi)存開銷

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

四、實驗結(jié)果

圖像壓縮感知重建

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

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

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

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

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

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

五、作者簡介

• 陳斌：北京大學(xué)信息工程學(xué)院博士生，主要研究方向是圖像壓縮感知與超分辨率。
• 張振宇：北京大學(xué)信息工程學(xué)院碩士生，主要研究方向是圖像重建。
• 李瑋琦：北京大學(xué)信息工程學(xué)院博士生，主要研究方向是圖像壓縮感知與超分辨率。
• 趙?。荷程匕⒉范爬瓏蹩萍即髮W(xué)（KAUST）的研究科學(xué)家，圖像與視頻理解實驗室（IVUL）視頻分析課題組組長。她首次提出了可逆化預(yù)訓(xùn)練神經(jīng)網(wǎng)絡(luò)方法，實現(xiàn)了大型預(yù)訓(xùn)練模型的極低顯存微調(diào)。她在這一領(lǐng)域的代表工作包括Re2TAL（CVPR 2023）、Dr2Net （CVPR 2024）等。

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

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

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

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

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