顶刊TPAMI 2025 | 北大、KAUST、字节联合提出“可逆扩散模型”赋能图像重建，代码已开源！

前言 本篇分享 TPAMI 2025 论文Invertible Diffusion Models for Compressed Sensing，北大、KAUST、字节联合提出“可逆扩散模型”赋能图像重建，代码已开源！

欢迎关注公众号CV技术指南，专注于计算机视觉的技术总结、最新技术跟踪、经典论文解读、CV招聘信息。

本文转载自我爱计算机视觉

仅用于学术分享，若侵权请联系删除

CV方向的准研究生们，未来三年如何度过？

招聘高光谱图像、语义分割、diffusion等方向论文指导老师

• 论文作者：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的训练、推理效率和重建性能得到了显著提升。

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

 

欢迎关注公众号CV技术指南，专注于计算机视觉的技术总结、最新技术跟踪、经典论文解读、CV招聘信息。

计算机视觉入门1v3辅导班

【技术文档】《从零搭建pytorch模型教程》122页PDF下载

QQ交流群：470899183。群内有大佬负责解答大家的日常学习、科研、代码问题。

其它文章

分享一个CV知识库，上千篇文章、专栏，CV所有资料都在这了

明年毕业，还不知道怎么做毕设的请抓紧机会了

LSKA注意力 | 重新思考和设计大卷积核注意力，性能优于ConvNeXt、SWin、RepLKNet以及VAN

CVPR 2023 | TinyMIM：微软亚洲研究院用知识蒸馏改进小型ViT

ICCV2023|涨点神器！目标检测蒸馏学习新方法，浙大、海康威视等提出

ICCV 2023 Oral | 突破性图像融合与分割研究：全时多模态基准与多交互特征学习

听我说，Transformer它就是个支持向量机

HDRUNet | 深圳先进院董超团队提出带降噪与反量化功能的单帧HDR重建算法

南科大提出ORCTrack | 解决DeepSORT等跟踪方法的遮挡问题，即插即用真的很香

1800亿参数，世界顶级开源大模型Falcon官宣！碾压LLaMA 2，性能直逼GPT-4

SAM-Med2D：打破自然图像与医学图像的领域鸿沟，医疗版 SAM 开源了！

GhostSR|针对图像超分的特征冗余，华为诺亚&北大联合提出GhostSR

Meta推出像素级动作追踪模型，简易版在线可玩 | GitHub 1.4K星

CSUNet | 完美缝合Transformer和CNN，性能达到UNet家族的巅峰！​

AI最全资料汇总 | 基础入门、技术前沿、工业应用、部署框架、实战教程学习

计算机视觉入门1v3辅导班

计算机视觉交流群