技术摘要：通过模式搜索损失解决扩散模型的逆问题

问题定义

本文研究了利用预训练的无条件扩散模型解决逆问题（inverse problems）的挑战。逆问题涉及在给定部分退化的测量值 $y$ 的情况下，估计原始数据样本 $x_0$，其过程由方程 $y = A(x_0) + \eta$ 描述，其中 $A$ 是已知的退化算子，$\eta$ 是高斯噪声。

虽然预训练的扩散模型可以以零样本（即插即用）的方式解决这些问题，但现有方法面临显著局限性：

1. 后验采样方法（Posterior Sampling Methods）： 这些方法旨在抽取与后验概率 $p(x_0|y)$ 成比例的样本。然而，它们通常依赖于建模近似（例如，对条件得分进行高斯近似）或顺序蒙特卡洛方法，这可能导致计算量巨大，并且在缺乏地面真值验证的情况下，可能无法严格遵循真实的后验分布。
2. MAP 估计方法（MAP Estimation Methods）： 这些方法旨在寻找最大后验（MAP）估计（即后验分布的众数/模态）。现有的求解器通常使用启发式近似，依赖于变量拆分（例如 HQS、ADMM），或需要一致性模型（consistency models）。这些方法通常涉及昂贵的梯度计算，且其建模假设对于所有线性算子可能并不成立。

作者认为，许多图像修复任务优先考虑生成与测量值一致的**合理（plausible）**图像，而非覆盖整个后验分布。这更接近于 MAP 估计的目标，从而为开发一种更具原则性、无近似的 MAP 策略提供了动力。

方法论：变分模式搜索损失 (VML)

核心贡献是变分模式搜索损失 (Variational Mode-Seeking Loss, VML)，这是一种旨在在反向扩散过程中引导生成样本趋向于 MAP 估计的新型目标函数。

理论基础

作者通过最小化每个反向扩散时间步 $t$ 时，扩散后验 $p(x_0|x_t)$ 与测量后验 $p(x_0|y)$ 之间的库尔贝克-莱布勒散度（KL 散度），从变分角度推导出了 VML：
$$\text{VML}(x_t, t) = D_{KL}(p(x_0|x_t) \parallel p(x_0|y))$$

直觉在于：随着 $t \to 0$，$p(x_0|x_t)$ 的分布变得越来越集中在 $x_t$ 周围。通过在每一步寻找使该 KL 散度最小的最优 $x_t^*$，该过程将样本引导至后验分布 $p(x_0|y)$ 的众数（modes）处。

针对线性算子的闭式解推导

一个关键的理论成就是，在不使用任何建模近似的情况下，推导出了线性逆问题（$A(x) = Hx$）的 VML 闭式表达式。完整的 VML 表达式包括：

1. 一个测量一致性项：$\|y - HD(x_t, t)\|^2 / (2\sigma_y^2)$。
2. 先验项：涉及去噪器 $D(x_t, t)$、噪声调度 $\sigma_t$ 以及后验协方差 $Cov[x_0|x_t]$。

简化与重权重的目标函数

完整的 VML 包含涉及协方差矩阵 $Cov[x_0|x_t]$ 的高阶项，这些项在计算和求导时非常昂贵。作者证明，当 $t \to 0$ 时，这些高阶项会一致地收敛到零（或常数）。因此，他们提出了一个简化版 VML ($VML_S$)，该版本舍弃了这些项，仅保留测量一致性和先验项。

为了进一步提高优化稳定性，他们引入了重权重的简化版 VML ($VML_S^\rho$)：
$$VML_S^\rho(x_t, t) = \frac{\|y - HD(x_t, t)\|^2}{2\sigma_y^2} - \rho_t \log p(x_t) - \rho_t \frac{\|D(x_t, t) - x_t\|^2}{2\sigma_t^2} + C$$
其中 $\rho_t$ 是一个先验权重调度，满足当 $t \to 0$ 时 $\rho_t \to 1$。该公式允许使用标准的自动微分进行高效的梯度计算。

算法

基于 VML 最小化，作者提出了以下算法：

1. VML-MAP：一种使用梯度下降法在每个反向扩散步骤中最小化 $VML_S^\rho$ 的算法。
2. VML-MAP$_{pre}$：针对病态（ill-conditioned）线性算子的预处理变体。它利用退化矩阵 $H$ 的奇异值分解（SVD）来构建预处理器 $P$，从而加速处理如去模糊（deblurring）等困难任务时的收敛。
3. VML-MAP$^\tau$：针对含噪声图像修复的策略。为了避免在低噪声水平下因优化挑战而产生的锐利伪影，该算法仅在达到扩散时间阈值 $\tau$ 之前最小化 VML，随后切换到标准 DDIM 采样以生成最终图像。
4. LatentVML-MAP：对潜在扩散模型（LDM）的扩展，在潜在空间中近似 VML。

核心贡献

1. 新型损失函数：引入了 VML，它通过最小化扩散后验与测量后验之间的 KL 散度，在理论上引导样本趋向于 MAP 估计。
2. 无近似推导：对于线性逆问题，VML 被推导为闭式解，且无需任何启发式的建模近似，为利用预训练无条件扩散模型进行 MAP 估计提供了正式的理论基础。
3. 收敛性分析：通过形式化证明，在温和的假设条件下，证明了简化版和重权重的 VML 目标函数在 $t \to 0$ 时收敛于负对数后验（即 MAP 目标）。
4. 实用算法：开发了 VML-MAP 及其预处理变体，为多种图像修复任务提供了经验有效的解决方案。
5. 广泛验证：在多个数据集（ImageNet, FFHQ, Celeba）和任务（图像补全、超分辨率、去模糊）上展示了其相对于现有基准方法（如 DDRM, $\Pi$GDM, MAPGA, DAPS）的卓越性能。

实验结果

作者使用在 ImageNet256, FFHQ256 和 CelebA 上预训练的模型，在图像修复任务上评估了其方法。

• 指标：论文强调使用感知指标（LPIPS 和 FID）而非像素级指标（PSNR, SSIM），理由是后者在处理病态逆问题中生成的具有合理性的重建图像时，往往无法捕捉到感知质量。
• 性能表现：
  • VML-MAP 在图像补全、超分辨率和去模糊任务中，在 LPIPS 和 FID 指标上一致优于后验采样方法（DDRM, $\Pi$GDM）和其他 MAP 估计器（MAPGA, DAPS）。
  • VML-MAP$_{pre}$ 对于病态线性算度任务（如去模糊）表现出比非预处理版本显著的提升，获得了更低的 LPIPS 和 FID 分数。
  • 计算效率：与需要许多反向扩散步骤的方法相比，VML-MAP 以更短的运行时间（更少的神经函数评估次数）实现了更好的感知质量。
• 噪声修复：VML-MAP$^\tau$ 策略成功处理了含噪声的测量值，在标准 VML-MAP 容易产生伪影的情况下，生成了高质量的重建图像。
• 潜在模型：LatentVML-MAP 在 CelebA 上展示了有效性，尽管作者指出由于解码器的非线性，在潜在空间中的优化仍然更具挑战性。

重要性与主张

本文声称，VML 代表了首次在仅利用预训练无条件扩散模型的情况下，对线性逆问题进行无近似建模的 MAP 估计方法，并得到了正式理论分析的支持。

作者强调，他们的方法与后验采样（旨在实现分布覆盖）以及现有的 MAP 求解器（通常依赖启发式或一致性模型）有着本质区别。通过最小化 VML，该方法有效地瞄准了后验分布的“众数”，作者认为这正是生成高质量、具有合理性的图像修复所需的行为。

作者对自身的局限性也进行了客观说明，承认：

• 其有效性依赖于是否存在能够高效处理 VML 目标的优化器。
• 对于病态或非线性算子，优化挑战依然存在，尽管预处理有所帮助。
• 未来需要开发更高阶的优化器，以在不产生过高计算成本的前提下进一步缓解优化问题。

文章总结道，VML-MAP 提供了一个有原则、有理论依据且在经验上有效的框架，用于利用扩散模型解决逆问题，弥合了理论 MAP 估计与实际图像修复之间的鸿沟。