顶刊TPAMI 2025！北大、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/。

（文：PaperWeekly）

欢迎分享

发表评论