Unbiased Diffusion Variational Inversion via Principled Posterior Matching

以下是论文《基于原则性后验匹配的无偏扩散变分反演》的详细技术总结。

问题陈述

计算成像任务涉及从含噪的间接测量值 $y = A(x) + \eta$ 中恢复未知信号 $x$，其中 $A$ 是一个病态前向算子。这些逆问题中的一个关键挑战，不仅在于获得高质量的点估计，还在于通过恢复完整的贝叶斯后验分布 $p(x|y)$ 来进行可靠的不确定性量化。

现有的基于扩散的逆问题求解方法通常分为三类，它们在优化 Kullback-Leibler (KL) 散度方面均存在显著的理论局限性：

1. 梯度引导的蒙特卡洛方法（如 DPS）： 这些方法近似难以处理的时间相关似然项，往往导致采样轨迹不稳定和后验估计有偏。
2. 基于优化的变分推断（如 RED-Diff, SDS）： 这些方法通常假设退化的变分分布（狄拉克 $\delta$ 函数），有效地丢弃了变分目标中的熵项。这将问题坍缩为最大后验（MAP）估计，导致严重的模式崩溃，并无法捕捉后验的多模态结构。
3. 摊销推断（如 DAVI）： 这些方法训练前馈网络直接预测后验。然而，它们依赖于积分 KL (IKL) 散度，这是一种启发式近似而非严格的变分下界。这种偏差，加上对配对真实数据进行监督的依赖，限制了它们进行零样本泛化和恢复精确贝叶斯后验的能力。

方法论：原则性后验匹配 (PPM)

作者提出了原则性后验匹配 (PPM)，这是一个回归变分推断基础以实现无偏优化精确 KL 散度的框架。

1. 理论基础：费舍尔散度积分
PPM 不依赖启发式近似或退化分布假设，而是通过沿扩散过程积分费舍尔散度 (Fisher divergence)，严格构建了精确优化 KL 散度的形式。作者利用了一个经典理论结果，即两个分布之间的 KL 散度可以展开为沿扩散轨迹的费舍尔散度的积分。

• 关键洞察： 最小化费舍尔散度的积分本质上鼓励覆盖质量 (mass-covering) 行为，从理论上缓解了先前方法中观察到的模式崩溃。

2. 梯度等价定理
为了使费舍尔散度优化可行，作者推导了一个新颖的梯度等价定理。该定理提供了一个精确、无偏的变分参数梯度估计量，无需通过冻结的预训练得分网络进行反向传播。

• 该梯度是利用“停止梯度 (stop-gradient)"算子和辅助得分网络推导得出的，允许通过随机梯度下降进行高效优化。

3. 统一框架
PPM 在单一优化目标下统一了两个互补的范式：

• 变分推断 (VI)： 变分后验被建模为图像粒子 $\{\mu_k\}$ 的集合。这允许针对特定观测值进行高保真的、基于粒子的后验恢复。
• 摊销推断 (AI)： 变分后验被建模为重建网络 $x = g_\phi(z)$。这通过学习从观测值到后验样本的全局映射，实现了快速的单步推理。
• 架构： 该框架利用预训练的扩散得分网络 $s_\theta$ 作为先验，并利用辅助得分网络 $s_\phi$（通过低秩适应进行适配）来近似变分分布 $q_\phi$ 的得分。优化过程在更新辅助得分网络和更新变分参数（粒子或网络权重）之间交替进行。

4. 泛化性
该公式通过将费舍尔散度的积分推广，自然地扩展到更广泛的散度度量族，将 PPM 定位为基于分数的变分推断的通用基础。

主要贡献

• 无偏优化： PPM 是首个严格优化基于扩散的逆问题精确 KL 散度的框架，无需诉诸现有方法中普遍存在的有偏近似（如狄拉克 $\delta$ 假设或 IKL 启发式）。
• 理论统一： 论文证明，通过费舍尔积分最小化精确 KL 散度，将现有的基于分数的方法统一为特例，并为先前方法遭受模式崩溃的原因（“近似差距”）提供了理论解释。
• 双范式支持： 它提供了一种统一的训练逻辑，既支持用于高保真后验恢复的基于粒子的 VI，也支持用于高效无监督生成的摊销推断。
• 鲁棒的不确定性量化： 通过避免模式崩溃，PPM 产生校准良好的不确定性估计，真实地反映后验的多模态性质。

实验结果

作者在包括自然图像和科学成像任务在内的全面计算成像基准测试中验证了 PPM。

1. 二维后验估计（玩具示例）

• 在二维高斯混合先验上，PPM 准确恢复了所有后验模式。相比之下，RED-Diff 表现出严重的模式崩溃（收敛到单一点），而 RLSD 则显示出测量不一致性或先验遵循度差。

2. 计算摄影 (FFHQ & ImageNet)

• 任务： 图像修复、运动去模糊、超分辨率和高斯去噪。
• 性能： 在重建保真度 (PSNR, SSIM) 和样本多样性方面，PPM 优于最先进基线 (DPS, RED-Diff, RLSD, DAVI)。
• 定性： 在修复任务中，PPM 生成多样且合理的补全内容（例如不同的发型），而基线方法产生几乎相同的输出。
• 不确定性： PPM 的标准差图准确反映了位置不确定性（例如在物体边缘），而 RED-Diff 的图指示了崩溃的模式，RLSD 的图则包含伪影。

3. 科学成像应用

• 超分辨率荧光显微镜： 应用于亚细胞结构（微管、内质网）。PPM 实现了比 RED-Diff 更高的保真度，并产生与成像物理一致的不确定性图（结构中心不确定性较低，模糊边缘处较高）。
• 射电干涉黑洞成像： 应用于合成事件视界望远镜 (EHT) 数据。PPM 忠实地恢复了黑洞阴影的环状结构和明亮新月。与发生崩溃的 RED-Diff 和模糊的 DPS 不同，PPM 提供了具有物理意义的不确定性分布，捕捉了结构方差。

意义与主张

论文声称，PPM 通过解决当前基于扩散的逆求解器的根本理论局限性，建立了一个用于科学成像的鲁棒框架。

• 可靠性： 通过消除导致模式崩溃的近似偏差，PPM 实现了忠实的多模态后验恢复和校准良好的不确定性估计，这对于科学和医学应用至关重要，因为在这些领域中，理解完整的解空间与重建本身同样重要。
• 零样本泛化： 与需要配对训练数据的摊销方法不同，PPM 以无监督方式运行，使其适用于缺乏真实数据的科学领域。
• 理论根基： 该工作在费舍尔散度与 KL 最小化之间提供了严格的理论桥梁，为长期主导该领域的启发式近似提供了一种原则性的替代方案。

总之，PPM 表明，通过费舍尔散度积分促进的变分目标精确优化，相较于现有的有偏近似，能产生更优越的重建质量、多样性和不确定性量化。