頂刊TPAMI2025！北大、KAUST等提出可逆擴散模型賦能影像重建，程式碼已開源

論文資訊

論文標題：

Invertible Diffusion Models for Compressed Sensing

論文作者：

Bin Chen（陳斌），Zhenyu Zhang（張振宇），Weiqi Li（李瑋琦），Chen Zhao（趙琛），Jiwen Yu（餘濟聞），Shijie Zhao（趙世傑），Jie Chen（陳杰）and Jian Zhang（張健）

作者單位：

北京大學資訊工程學院、阿卜杜拉國王科技大學、字節跳動

發表刊物：

IEEE Transactions on Pattern Analysis and Machine Intelligence (TPAMI)

發表時間：

2025年2月5日

正式版本：

https://ieeexplore.ieee.org/document/10874182

ArXiv版本：

https://arxiv.org/abs/2403.17006

開原始碼：

https://github.com/Guaishou74851/IDM

任務背景

擴散模型作為當前非常知名且強大的生成模型之一，已在影像重建任務中展現出極大的潛力。擴散模型的基本實現方式是在訓練階段構建一個噪聲估計網路（通常是一個 UNet），並在推理階段透過迭代的去噪和加噪過程完成影像生成與重建。

然而，如何進一步提升擴散模型在影像重建中的效能與效率，仍然是業界探索的重點問題。

當我們將擴散模型應用於影像重建任務時，面臨兩個關鍵挑戰：

挑戰一：“噪聲估計”任務與“影像重建”任務之間的偏差。擴散模型中的深度神經網路主要針對“噪聲估計”任務（即，從當前變數中估計出噪聲）得到最最佳化，而非“影像重建”任務（即，從低質量的觀測資料中預測原始影像）本身。這可能導致其影像重建效能存在進一步提升的空間。
挑戰二：推理速度慢、效率低。儘管擴散模型能夠生成較為真實的影像，但其推理過程往往需要大量的迭代步驟，執行時間長，計算開銷大，不利於實際應用。

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

主要貢獻

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

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

傳統擴散模型在訓練階段的目標任務是“噪聲估計”，而實際的目標任務是“影像重建”。為了提升擴散模型的影像重建效能，我們將它的迭代取樣過程重新定義為一個整體的影像重建網路，對該網路進行端到端的訓練，突破了傳統噪聲估計學習正規化所帶來的侷限。

如圖所示，透過這種方式，模型的所有引數都針對“影像重建”任務進行了最最佳化，重建效能得到大幅提升。

實驗結果表明，基於 Stable Diffusion 的預訓練權重與這一端到端學習框架，在影像壓縮感知重建任務中，相比其他模型，我們的方法在 PSNR（峰值信噪比）指標上提升了 2dB，取樣步數從原本的 100 步降到了 3 步，推理速度提升了約 15 倍。

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

大型擴散模型（如 Stable Diffusion）取樣過程的端到端訓練需要佔用很大的 GPU 記憶體，這對於其實際應用來說是一個嚴重的瓶頸。為了減少記憶體開銷，我們提出了一種雙層可逆網路。可逆網路的核心思想是透過設計特殊的網路結構，讓網路每一層的輸出可以反向計算得到輸入。

在實踐中，我們將可逆網路應用到（1）所有擴散取樣步驟和（2）噪聲估計網路的內部，透過“佈線”技術將每個取樣步驟與其前後模組連線，形成一個雙層可逆網路。

這一設計使得整個訓練過程中，程式無需儲存完整的特徵圖資料，只需儲存較少的中間變數，顯著降低了訓練模型的 GPU 記憶體需求。最終，這使得我們可以在視訊記憶體有限的 GPU（如 1080Ti）上對該模型進行端到端訓練。

實驗結果

1. 影像壓縮感知重建

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

2. 影像補全與醫學成像

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

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

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

欲瞭解更多細節，請參考原論文。

作者簡介

陳斌：北京大學資訊工程學院博士生，主要研究方向是影像壓縮感知與超解析度。

張振宇：北京大學資訊工程學院碩士生，主要研究方向是影像重建。

李瑋琦：北京大學資訊工程學院博士生，主要研究方向是影像壓縮感知與超解析度。

趙琛：沙特阿卜杜拉國王科技大學（KAUST）的研究科學家，影像與影片理解實驗室（IVUL）影片分析課題組組長。她首次提出了可逆化預訓練神經網路方法，實現了大型預訓練模型的極低視訊記憶體微調。她在這一領域的代表工作包括 Re2TAL（CVPR 2023）、Dr2Net （CVPR 2024）等。個人主頁：https://zhao-chen.com/。

餘濟聞：北京大學資訊工程學院碩士生，主要研究方向是生成式擴散模型。

趙世傑：字節跳動多媒體實驗室研究員，負責影片處理與增強課題組。

陳杰：北京大學資訊工程學院副教授，主要研究方向是計算機視覺與模式識別和 AI4Science。個人主頁：https://aimia-pku.github.io/。

張健：北京大學資訊工程學院副教授，主要研究方向是視覺內容重建與生成、AIGC 內容鑑偽和版權保護。個人主頁：https://jianzhang.tech/cn/。

更多閱讀