顶刊TPAMI 2025！北大&字节等提出IDM：可逆扩散模型赋能图像重建

点击下方卡片，关注“CVer”公众号

AI/CV重磅干货，第一时间送达

点击进入—>【顶会/顶刊】投稿交流群

添加微信号：CVer2233，小助手会拉你进群！

扫描下方二维码，加入CVer学术星球！可以获得最新顶会/顶刊上的论文idea和CV从入门到精通资料，及最前沿应用！发论文/搞科研/涨薪，强烈推荐！

转载自：将门创投

扩散模型作为一种强大的生成模型，已在图像重建任务中展现出巨大潜力。其基本原理是在训练阶段优化噪声估计网络（如UNet），并通过去噪和加噪过程完成图像生成。然而，如何提升扩散模型在图像重建中的性能与效率，仍是一个重要挑战。

在图像重建任务中，扩散模型面临两个主要问题：首先，噪声估计任务与图像重建任务之间存在偏差，限制了重建效果；其次，推理过程需要大量迭代步骤，导致速度慢、效率低。为解决这些问题，本文提出了可逆扩散模型（IDM），通过引入端到端训练框架和可逆网络设计，显著提升了图像重建的性能和效率。

论文标题： 

Invertible Diffusion Models for Compressed Sensing

论文链接：

https://arxiv.org/abs/2403.17006

开源代码：

https://github.com/Guaishou74851/IDM

一、任务背景

扩散模型作为当前非常知名且强大的生成模型之一，已在图像重建任务中展现出极大的潜力。扩散模型的基本实现方式是在训练阶段构建一个噪声估计网络（通常是一个UNet），并在推理阶段通过迭代的去噪和加噪过程完成图像生成与重建。然而，如何进一步提升扩散模型在图像重建中的性能与效率，仍然是业界探索的重点问题。

当我们将扩散模型应用于图像重建任务时，面临两个关键挑战：

• 挑战一：“噪声估计”任务与“图像重建”任务之间的偏差。 扩散模型中的深度神经网络主要针对“噪声估计”任务（即，从当前变量中估计出噪声）得到最优化，而非“图像重建”任务（即，从低质量的观测数据中预测原始图像）本身。这可能导致其图像重建性能存在进一步提升的空间。

• 挑战二：推理速度慢、效率低。 尽管扩散模型能够生成较为真实的图像，但其推理过程往往需要大量的迭代步骤，运行时间长，计算开销大，不利于实际应用。

针对这两个挑战，本文提出了一种可逆扩散模型（Invertible Diffusion Models，IDM）。这一方法通过引入（1）端到端的训练框架与（2）可逆网络设计，有效提升了图像重建的性能与效率。

二、主要贡献

我们的方法在图像重建任务中带来了两个主要创新：

2.1 端到端的扩散采样图像重建学习框架

传统扩散模型在训练阶段的目标任务是“噪声估计”，而实际的目标任务是“图像重建”。为了提升扩散模型的图像重建性能，我们将它的迭代采样过程重新定义为一个整体的图像重建网络，对该网络进行端到端的训练，突破了传统噪声估计学习范式所带来的局限。如图所示，通过这种方式，模型的所有参数都针对“图像重建”任务进行了最优化，重建性能得到大幅提升。实验结果表明，基于Stable Diffusion的预训练权重与这一端到端学习框架，在图像压缩感知重建任务中，相比其他模型，我们的方法在PSNR（峰值信噪比）指标上提升了2dB，采样步数从原本的100步降到了3步，推理速度提升了约15倍。

2.2 双层可逆网络设计：减少内存开销

大型扩散模型（如Stable Diffusion）采样过程的端到端训练需要占用很大的GPU内存，这对于其实际应用来说是一个严重的瓶颈。为了减少内存开销，我们提出了一种双层可逆网络。可逆网络的核心思想是通过设计特殊的网络结构，让网络每一层的输出可以反向计算得到输入。

在实践中，我们将可逆网络应用到（1）所有扩散采样步骤和（2）噪声估计网络的内部，通过“布线”技术将每个采样步骤与其前后模块连接，形成一个双层可逆网络。这一设计使得整个训练过程中，程序无需存储完整的特征图数据，只需存储较少的中间变量，显著降低了训练模型的GPU内存需求。最终，这使得我们可以在显存有限的GPU（如1080Ti）上对该模型进行端到端训练。

三、实验结果

3.1 图像压缩感知重建

在图像压缩感知重建任务中，我们的方法IDM与现有基于端到端网络和扩散模型的重建方法进行了对比。实验结果显示，IDM在PSNR、SSIM、FID和LPIPS等指标上取得明显提升。

3.2 图像补全与医学成像

在掩码率90%的图像补全任务中，我们的方法能够准确恢复出窗户等复杂结构，而传统的扩散模型（如DDNM）无法做到这一点。此外，我们还将该方法应用于医学影像领域，包括核磁共振成像（MRI）和计算机断层扫描（CT）成像，取得了良好的效果。

3.3 计算成本与推理时间的优化

基于传统扩散模型的图像重建方法往往需要较长的推理时间和计算开销，而我们的可逆扩散模型IDM显著缩短了这一过程。在重建一张256×256大小的图像时，推理时间从9秒缩短至0.63秒，大幅降低了计算开销。与现有方法DDNM相比，IDM的训练、推理效率和重建性能得到了显著提升。

欲了解更多细节，请参考原论文。

何恺明在MIT授课的课件PPT下载

在CVer公众号后台回复：何恺明，即可下载本课程的所有566页课件PPT！赶紧学起来！

CVPR 2025 论文和代码下载

在CVer公众号后台回复：CVPR2025，即可下载CVPR 2025论文和代码开源的论文合集

ECCV 2024 论文和代码下载

在CVer公众号后台回复：ECCV2024，即可下载ECCV 2024论文和代码开源的论文合集

CV垂直方向和论文投稿交流群成立

扫描下方二维码，或者添加微信号：CVer2233，即可添加CVer小助手微信，便可申请加入CVer-垂直方向和论文投稿微信交流群。另外其他垂直方向已涵盖：目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch、TensorFlow和Transformer、NeRF、3DGS、Mamba等。
一定要备注：研究方向+地点+学校/公司+昵称（如Mamba、多模态学习或者论文投稿+上海+上交+卡卡），根据格式备注，可更快被通过且邀请进群

▲扫码或加微信号: CVer2233，进交流群
CVer计算机视觉（知识星球）来了！想要了解最新最快最好的CV/DL/AI论文速递、优质实战项目、AI行业前沿、从入门到精通学习教程等资料，欢迎扫描下方二维码，加入CVer计算机视觉（知识星球），已汇集上万人！

▲扫码加入星球学习

▲点击上方卡片，关注CVer公众号
整理不易，请点赞和在看