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

一、論文信息

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

二、任務(wù)背景

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

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

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

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

三、主要貢獻(xiàn)

我們的方法在圖像重建任務(wù)中帶來(lái)了兩個(gè)主要?jiǎng)?chuàng)新：

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

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

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

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

四、實(shí)驗(yàn)結(jié)果

圖像壓縮感知重建

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

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

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

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

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

欲了解更多細(xì)節(jié)，請(qǐng)參考原論文。

五、作者簡(jiǎn)介

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

• 個(gè)人主頁(yè)：https://zhao-chen.com/。

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

• 個(gè)人主頁(yè)：https://aimia-pku.github.io

• 張?。罕本┐髮W(xué)信息工程學(xué)院副教授，主要研究方向是視覺(jué)內(nèi)容重建與生成、AIGC內(nèi)容鑒偽和版權(quán)保護(hù)。

• 個(gè)人主頁(yè)：https://jianzhang.tech/cn