TPAMI 2025 | 北大、KAUST、字节联合提出“可逆扩散模型”赋能图像重建

扩散模型作为一种强大的生成模型，已在图像重建任务中展现出巨大潜力。其基本原理是在训练阶段优化噪声估计网络（如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的训练、推理效率和重建性能得到了显著提升。

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

llustration From IconScout By IconScout Store

-The End-

本周上新！

扫码观看！

“AI技术流”原创投稿计划

TechBeat是由将门创投建立的AI学习社区（

www.techbeat.net

投稿内容

// 最新技术解读/系统性知识分享 //

// 前沿资讯解说/心得经历讲述 //

投稿须知

稿件需要为原创文章，并标明作者信息。

我们会选择部分在深度技术解析及科研心得方向，对用户启发更大的文章，做原创性内容奖励

投稿方式

发送邮件到

melodybai@thejiangmen.com

或添加工作人员微信（yellowsubbj）投稿，沟通投稿详情；还可以关注“将门创投”公众号，后台回复“投稿”二字，获得投稿说明。

关于我“门”

将门是一家以专注于数智核心科技领域的新型创投机构，也是北京市标杆型孵化器。 公司致力于通过连接技术与商业，发掘和培育具有全球影响力的科技创新企业，推动企业创新发展与产业升级。

将门成立于2015年底，创始团队由微软创投在中国的创始团队原班人马构建而成，曾为微软优选和深度孵化了126家创新的技术型创业公司。

如果您是技术领域的初创企业，不仅想获得投资，还希望获得一系列持续性、有价值的投后服务，欢迎发送或者推荐项目给我“门”:

bp@thejiangmen.com

点击右上角，把文章分享到朋友圈