扩散模型作为一种强大的生成模型，已在图像重建任务中展现出巨大潜力。其基本原理是在训练阶段优化噪声估计网络（如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倍。
大型扩散模型（如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
“AI技术流”原创投稿计划
TechBeat是由将门创投建立的AI学习社区（
www.techbeat.net
）
。
社区上线600+期talk视频，3000+篇技术干货文章，方向覆盖CV/NLP/ML/Robotis等；每月定期举办顶会及其他线上交流活动，不定期举办技术人线下聚会交流活动。我们正在努力成为AI人才喜爱的高质量、知识型交流平台，希望为AI人才打造更专业的服务和体验，加速并陪伴其成长。
投稿内容
// 最新技术解读/系统性知识分享 //
// 前沿资讯解说/心得经历讲述 //
投稿须知
稿件需要为原创文章，并标明作者信息。
我们会选择部分在深度技术解析及科研心得方向，对用户启发更大的文章，做原创性内容奖励
投稿方式
发送邮件到
[email protected]
或添加
工作人员微信（
yellowsubbj
）
投稿，沟通投稿详情；还可以关注“将门创投”公众号，后台回复“
投稿
”二字，获得投稿说明。
