Learning few-step posterior samplers by unfolding and distillation of diffusion models

📄 arXiv: 2507.02686v2 📥 PDF

作者: Charlesquin Kemajou Mbakam, Jonathan Spence, Marcelo Pereyra

分类: cs.CV, cs.LG

发布日期: 2025-07-03 (更新: 2025-11-18)

备注: 34 pages, 18 figures, 11 tables

💡 一句话要点

通过扩散模型展开与蒸馏学习少量步骤的后验采样器

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 扩散模型 深度展开 模型蒸馏 贝叶斯计算成像 后验采样

📋 核心要点

  1. 现有基于扩散模型的贝叶斯计算成像方法，如即插即用方法，依赖近似，而条件扩散模型缺乏泛化能力。
  2. 本文提出一种结合深度展开和模型蒸馏的框架，将扩散模型先验转化为少量步骤的条件模型，用于后验采样。
  3. 实验表明，该方法在保持灵活性的同时，实现了优异的精度和计算效率，优于现有技术水平。

📝 摘要（中文）

扩散模型(DMs)已成为贝叶斯计算成像中强大的图像先验。目前有两种主要策略利用DMs：即插即用方法，该方法是零样本且高度灵活的，但依赖于近似；以及专门的条件DMs，通过监督训练，针对特定任务实现更高的准确性和更快的推理速度。本文介绍了一种新颖的框架，该框架集成了深度展开和模型蒸馏，将DM图像先验转换为用于后验采样的少量步骤条件模型。我们方法的一个核心创新是马尔可夫链蒙特卡洛(MCMC)算法的展开——特别是最近提出的LATINO Langevin采样器(Spagnoletti et al., 2025)——代表了深度展开应用于蒙特卡洛采样方案的第一个已知实例。通过大量的实验以及与最先进技术的比较，我们展示了我们提出的展开和蒸馏采样器，它们实现了出色的准确性和计算效率，同时保留了适应推理时前向模型变化的能力。

🔬 方法详解

问题定义：论文旨在解决贝叶斯计算成像中，如何高效且灵活地利用扩散模型作为图像先验进行后验采样的问题。现有方法，如即插即用方法，虽然灵活，但依赖于近似，导致精度受限。而专门训练的条件扩散模型虽然精度高，但缺乏泛化能力，难以适应不同的前向模型。

核心思路：论文的核心思路是将扩散模型先验通过深度展开和模型蒸馏，转化为一个少量步骤的条件模型。具体来说，是将一个MCMC采样算法（LATINO Langevin采样器）展开成一个深度网络，然后利用模型蒸馏技术，将该网络的知识转移到一个更小的网络中，从而实现高效的后验采样。

技术框架：整体框架包含以下几个主要步骤：1) 选择一个扩散模型作为图像先验；2) 选择一个MCMC采样算法（LATINO Langevin采样器）；3) 将MCMC采样算法展开成一个深度网络，网络的每一层对应MCMC算法的一个迭代步骤；4) 使用模型蒸馏技术，将展开后的网络的知识转移到一个更小的网络中，得到最终的少量步骤后验采样器。

关键创新：最重要的创新点在于将深度展开技术应用于蒙特卡洛采样方案，这是首次尝试。通过展开MCMC算法，可以将扩散模型先验的信息有效地融入到后验采样过程中，从而提高采样效率和精度。此外，结合模型蒸馏技术，可以进一步压缩模型大小，提高推理速度。

关键设计：论文使用了LATINO Langevin采样器作为展开的基础，该采样器具有良好的理论性质。在模型蒸馏过程中，使用了合适的损失函数来保证蒸馏后的模型能够保留原始模型的性能。具体的网络结构和参数设置在论文中有详细描述，但摘要中未提及具体细节。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文通过实验验证了所提出的展开和蒸馏采样器的有效性。实验结果表明，该方法在保持灵活性的同时，实现了出色的准确性和计算效率，优于现有的即插即用方法和条件扩散模型。具体的性能数据和对比基线在论文中给出，但摘要中未提及具体数值。

🎯 应用场景

该研究成果可广泛应用于贝叶斯计算成像领域，例如医学图像重建、遥感图像处理、计算机断层扫描(CT)等。通过高效的后验采样，可以提高图像重建的质量和速度，从而提升诊断效率和准确性。此外，该方法具有较强的泛化能力，可以适应不同的成像模型，具有重要的实际应用价值。

📄 摘要（原文）

Diffusion models (DMs) have emerged as powerful image priors in Bayesian computational imaging. Two primary strategies have been proposed for leveraging DMs in this context: Plug-and-Play methods, which are zero-shot and highly flexible but rely on approximations; and specialized conditional DMs, which achieve higher accuracy and faster inference for specific tasks through supervised training. In this work, we introduce a novel framework that integrates deep unfolding and model distillation to transform a DM image prior into a few-step conditional model for posterior sampling. A central innovation of our approach is the unfolding of a Markov chain Monte Carlo (MCMC) algorithm - specifically, the recently proposed LATINO Langevin sampler (Spagnoletti et al., 2025) - representing the first known instance of deep unfolding applied to a Monte Carlo sampling scheme. We demonstrate our proposed unfolded and distilled samplers through extensive experiments and comparisons with the state of the art, where they achieve excellent accuracy and computational efficiency, while retaining the flexibility to adapt to variations in the forward model at inference time.