2025-11-18T05:16:13.529071

Local MAP Sampling for Diffusion Models

Zhang, Brekelmans, Steeg
Diffusion Posterior Sampling (DPS) provides a principled Bayesian approach to inverse problems by sampling from $p(x_0 \mid y)$. However, in practice, the goal of inverse problem solving is not to cover the posterior but to recover the most accurate reconstruction, where optimization-based diffusion solvers often excel despite lacking a clear probabilistic foundation. We introduce Local MAP Sampling (LMAPS), a new inference framework that iteratively solving local MAP subproblems along the diffusion trajectory. This perspective clarifies their connection to global MAP estimation and DPS, offering a unified probabilistic interpretation for optimization-based methods. Building on this foundation, we develop practical algorithms with a probabilistically interpretable covariance approximation, a reformulated objective for stability and interpretability, and a gradient approximation for non-differentiable operators. Across a broad set of image restoration and scientific tasks, LMAPS achieves state-of-the-art performance, including $\geq 2$ dB gains on motion deblurring, JPEG restoration, and quantization, and $>1.5$ dB improvements on inverse scattering benchmarks.
academic

Local MAP Sampling for Diffusion Models

基本信息

  • 论文ID: 2510.07343
  • 标题: Local MAP Sampling for Diffusion Models
  • 作者: Shaorong Zhang (UC Riverside), Rob Brekelmans (Vector Institute), Greg Ver Steeg (UC Riverside)
  • 分类: cs.GR cs.AI eess.IV
  • 发表时间/会议: Preprint (Under review)
  • 论文链接: https://arxiv.org/abs/2510.07343

摘要

Diffusion Posterior Sampling (DPS)通过从p(x0∣y)p(x_0 \mid y)中采样，为逆问题提供了一种有原则的贝叶斯方法。然而，在实践中，逆问题求解的目标不是覆盖后验分布，而是恢复最准确的重建结果，基于优化的扩散求解器通常在这方面表现出色，尽管缺乏清晰的概率基础。本文介绍了Local MAP Sampling (LMAPS)，这是一个新的推理框架，沿着扩散轨迹迭代求解局部MAP子问题。这种视角澄清了它们与全局MAP估计和DPS的联系，为基于优化的方法提供了统一的概率解释。基于这一基础，我们开发了具有概率可解释协方差近似、稳定性和可解释性的重新表述目标以及非可微算子的梯度近似的实用算法。

研究背景与动机

问题背景

扩散模型在逆问题求解中面临两个主要挑战：

  1. 目标不匹配：DPS旨在从后验分布p(x0∣y)p(x_0|y)中采样，但逆问题求解的实际目标是获得最准确的重建，而非多样性采样
  2. 理论基础缺失：基于优化的扩散求解器（如Resample、DiffPIR、DCDP等）在性能上表现出色，但缺乏清晰的概率理论基础

研究动机

  • 实用性导向：逆问题评估协议通常与单一真实参考进行比较，不为覆盖率或多样性提供奖励
  • 理论统一：需要为优化方法提供概率解释，澄清其与MAP估计和DPS的关系
  • 性能提升：在保持理论基础的同时，实现更好的重建性能

核心贡献

  1. 理论贡献：提出了Local MAP Sampling (LMAPS)框架，沿扩散轨迹迭代求解局部MAP子问题，分析了其与全局MAP和DPS的关系，统一了TMPD和基于优化的逆问题方法
  2. 方法贡献：
    • 提供概率可解释的协方差近似，替代现有求解器中的启发式选择
    • 引入目标重新表述以实现可解释参数和改善稳定性
    • 开发非可微算子的梯度近似策略
  3. 实验贡献：在10个图像恢复任务和3个科学逆问题上验证，在46/60个FFHQ/ImageNet案例中取得最佳结果，在运动去模糊、JPEG恢复和量化任务上获得≥2dB PSNR提升

方法详解

任务定义

逆问题旨在从先验分布π(x0)\pi(x_0)和噪声测量y∈Rmy \in \mathbb{R}^m中恢复未知图像或信号x0∈Rnx_0 \in \mathbb{R}^n： y=H(x0)+zy = H(x_0) + z 其中H(⋅):Rn→RmH(\cdot): \mathbb{R}^n \to \mathbb{R}^m是前向算子，z∼N(0,σy2I)z \sim \mathcal{N}(0, \sigma_y^2 I)是测量噪声。

核心理论框架

全局MAP vs 局部MAP

全局MAP直接优化完整后验： x0MAP:=arg⁡max⁡x0p(x0∣y)x_0^{\text{MAP}} := \arg\max_{x_0} p(x_0|y)

局部MAP在每个时间步tt求解条件优化问题： x0∗(t,xt,y):=arg⁡max⁡p(x0∣xt,y)x_0^*(t, x_t, y) := \arg\max p(x_0|x_t, y)xt−Δt=g(x0∗,xt,ϵ),ϵ∼N(0,I)x_{t-\Delta t} = g(x_0^*, x_t, \epsilon), \quad \epsilon \sim \mathcal{N}(0,I)

与DPS和DAPS的关系

  • DPS：使用条件均值E[x0∣xt,y]E[x_0|x_t, y]
  • DAPS：从p(x0∣xt,y)p(x_0|x_t, y)中采样
  • LMAPS：使用条件模式arg⁡max⁡p(x0∣xt,y)\arg\max p(x_0|x_t, y)

关键洞察：只有当p(x0∣xt,y)p(x_0|x_t, y)为高斯分布时，DPS和LMAPS才等价。

实用算法设计

协方差近似

采用各向同性近似： Σ0∣t≈kSNRI,SNR:=αt2σt2\Sigma_{0|t} \approx \frac{k}{\text{SNR}} I, \quad \text{SNR} := \frac{\alpha_t^2}{\sigma_t^2}

目标函数重新表述

将原始目标： x0∗=arg⁡min⁡{SNRk∥x0−m0∣t∥2+1σy2∥y−H(x0)∥2}x_0^* = \arg\min \left\{\frac{\text{SNR}}{k}\|x_0 - m_{0|t}\|^2 + \frac{1}{\sigma_y^2}\|y - H(x_0)\|^2\right\}

重新表述为： x0∗=arg⁡min⁡{(1−μt)12∥x0−m0∣t∥2+μtk2∥y−H(x0)∥2}x_0^* = \arg\min \left\{(1-\mu_t)\frac{1}{2}\|x_0 - m_{0|t}\|^2 + \mu_t k_2\|y - H(x_0)\|^2\right\}

其中μt=σt2σt2+k12∈(0,1)\mu_t = \frac{\sigma_t^2}{\sigma_t^2 + k_1^2} \in (0,1)，实现了：

  • 凸组合解释：权重为(1−μt)(1-\mu_t)和μt\mu_t
  • 自动退火：随着σt2\sigma_t^2减少，从测量驱动转向先验驱动
  • 数值稳定性：避免极端的SNR缩放

非可微算子处理

对于JPEG恢复、量化等非可微任务，使用代理梯度： ∇x0∥y−H(x0)∥2≈2JH′(x0)T(H(x0)−y)\nabla_{x_0}\|y - H(x_0)\|^2 \approx 2J_{H'}(x_0)^T(H(x_0) - y)

对于量化，采用H′(x0)=x0H'(x_0) = x_0，简化为： ∇x0∥y−H(x0)∥2≈2(H(x0)−y)\nabla_{x_0}\|y - H(x_0)\|^2 \approx 2(H(x_0) - y)

实验设置

数据集

  • 图像恢复：FFHQ 256×256和ImageNet 256×256，各使用100张测试图像
  • 科学逆问题：采用InverseBench数据集，包括荧光显微镜图像（线性逆散射）、GRMHD数据（黑洞成像）、fastMRI膝关节数据（压缩感知MRI）

评价指标

  • 图像恢复：PSNR、SSIM、LPIPS
  • 科学逆问题：主要使用PSNR，辅以任务特定指标

对比方法

包括DDNM、DDRM、ΠGDM、DPS、LGD、PnP-DM、FPS、MCG-diff、RedDiff、DAPS、DiffPIR、DCDP、DMPlug等12种方法

实现细节

  • 扩散步数：通常200步
  • 梯度更新步数：20-200步（取决于任务）
  • 学习率：0.01-1.0（任务相关）
  • 参数k1k_1：0-10，k2k_2：0.01-30000

实验结果

主要结果

图像恢复任务

在表1中，LMAPS在60个结果中的49个取得最佳性能：

  • 运动去模糊：FFHQ上32.62 dB vs DAPS的29.66 dB（+2.96 dB）
  • JPEG恢复：FFHQ上27.25 dB vs ΠGDM的25.04 dB（+2.21 dB）
  • 量化：FFHQ上29.51 dB vs ΠGDM的25.82 dB（+3.69 dB）

科学逆问题

在表2中，LMAPS在所有任务上都取得最佳PSNR：

  • 线性逆散射（NR=360）：38.07 dB vs RED-diff的36.56 dB（+1.51 dB）
  • 线性逆散射（NR=180）：37.19 dB vs RED-diff的35.41 dB（+1.78 dB）
  • 线性逆散射（NR=60）：30.75 dB vs RED-diff的27.07 dB（+3.68 dB）

消融实验

图4显示了优化步数与扩散步数的权衡：

  • 最佳性能通常在NFE=200-500时观察到
  • 增加每个扩散步的优化步数能显著提升性能
  • 相比SITCOM（600 NFEs），LMAPS以更少计算资源达到相似性能

计算效率

表3显示LMAPS在去模糊任务上的采样时间：

  • LMAPS（200步扩散，100步优化）：61秒/图像，30.88 dB
  • DAPS（200步扩散，100步优化）：110秒/图像，29.19 dB
  • SITCOM（600步）：73秒/图像，29.93 dB

相关工作

扩散后验采样

DPS及其变体通过直接从后验分布p(x0∣y)p(x_0|y)中采样来解决逆问题，包括TMPD、DDNM、ΠGDM等方法。

基于优化的方法

Resample、DiffPIR、DCDP、DMPlug等方法通过交替去噪、优化和重采样来解决逆问题，性能优异但缺乏理论基础。

MAP估计方法

最近的工作开始关注扩散先验下的MAP估计，但主要局限于线性逆问题。

结论与讨论

主要结论

  1. 理论统一：LMAPS为基于优化的扩散方法提供了统一的概率解释
  2. 性能提升：在多个任务上取得显著的PSNR提升，特别是在具有挑战性的非线性和非可微任务上
  3. 计算效率：相比现有方法具有更好的计算效率

局限性

  1. 收敛性：局部MAP序列不一定收敛到全局MAP
  2. 多样性：相比DPS，LMAPS可能产生较少的输出多样性
  3. 超参数敏感性：需要针对不同任务调整参数k1k_1和k2k_2

未来方向

论文指出全局MAP在贝叶斯推理中的关键作用被大大忽视，高效求解扩散先验下的全局MAP仍是开放挑战。虽然MAP可能通过集中在单一模式上减少输出多样性，但它提供更高的确定性和与观测数据的更好对齐。

深度评价

优点

  1. 理论贡献显著：首次为基于优化的扩散方法提供了清晰的概率解释
  2. 实验全面：涵盖10个图像恢复任务和3个科学逆问题，实验设置完整
  3. 性能提升明显：在多个具有挑战性的任务上取得2dB以上的显著提升
  4. 方法实用：提供了处理非可微算子的有效策略
  5. 写作清晰：理论分析和方法描述都很清楚

不足

  1. 理论分析深度：虽然提供了概率解释，但对收敛性和理论保证的分析相对有限
  2. 超参数复杂性：需要为每个任务调整多个超参数，可能影响方法的泛化性
  3. 计算开销：虽然相比某些方法更高效，但仍需要在每个时间步进行多次梯度更新
  4. 评估局限：主要关注重建质量，对不确定性量化的评估不足

影响力

  1. 学术价值：为扩散模型在逆问题中的应用提供了新的理论视角
  2. 实用价值：在图像恢复和科学计算领域具有直接应用价值
  3. 启发性：可能启发更多关于扩散模型MAP估计的研究

适用场景

  • 需要高质量重建而非多样性的逆问题
  • 涉及非可微前向算子的任务（如JPEG恢复、量化）
  • 科学计算中的逆问题求解
  • 对计算效率有一定要求的实时应用

参考文献

论文引用了扩散模型、逆问题求解、贝叶斯推理等领域的重要工作，包括DPS、DAPS、TMPD等关键方法的原始论文，为相关研究提供了良好的文献基础。