Diffusion posterior sampling for simulation-based inference in tall data settings

1. 研究背景与问题 (Problem)

核心任务：
在科学计算领域，经常需要识别非线性模型的参数，以解释观测数据。当模型依赖复杂的模拟器（如基于非线性随机微分方程）时，似然函数 $p(x|\theta)$ 通常是**不可计算（intractable）**的。这使得传统的贝叶斯推断方法（如 MCMC）难以直接应用。

现有方案与局限：

• 基于模拟的推断 (SBI)： 通过训练深度生成模型来近似后验分布 $p(\theta|x)$，从而绕过似然计算。
• 高维数据 (Tall Data) 挑战： 在实际应用中，我们往往拥有多个独立同分布（i.i.d.）的观测值 $x^*_{1:n} = (x^*_1, ..., x^*_n)$。利用这些额外信息可以得到更锐利的后验分布，提高参数可识别性。
  • 后验公式为：$p(\theta | x^*_{1:n}) \propto \lambda(\theta)^{1-n} \prod_{j=1}^n p(\theta | x^*_j)$。
• 现有方法的不足：
  • NPE/NLE/NRE 的扩展： 直接训练处理 $n$ 个观测值的模型需要巨大的模拟预算（$N_s \times n$），且缺乏灵活性（观测数变化需重新训练）。
  • F-NPSE (Factorized NPSE)： 最近提出的方法，通过组合单个观测的得分（Score）来处理高维数据，无需重新训练。
    • 致命缺陷： 由于组合后的得分并不对应任何已知的正向扩散过程，无法使用确定性采样器（如 DDIM）。F-NPSE 必须依赖退火朗之万动力学 (Annealed Langevin Dynamics) 进行采样。
    • 后果： 朗之万采样计算成本高、收敛慢、对超参数敏感且数值不稳定，成为整个流程的瓶颈。

本文目标：
提出一种新的算法，在保留 F-NPSE 的“摊销（amortized）”和“组合（compositional）”优势的同时，消除对朗之万采样的依赖，实现更快、更稳定的高维数据后验采样。

---

2. 方法论 (Methodology)

本文的核心思想是显式近似高维数据后验的扩散过程，从而推导出一个可计算的得分函数，使其能够直接用于确定性采样器（如 DDIM）。

2.1 理论基础

• 得分分解： 根据贝叶斯公式，高维后验的对数得分可以分解为先验得分和单个观测后验得分的加权和，但包含一个修正项（Correction Term）：
  $$\nabla_{\theta_t} \log p_t(\theta_t | x^*_{1:n}) = (1-n)\nabla \log p^\lambda_t + \sum \nabla \log p_t(\theta_t | x^*_j) + \nabla \log L_\lambda$$
  其中 $L_\lambda$ 涉及先验和各个后验的逆向核（backward kernels）的积分，通常不可计算。F-NPSE 忽略了这一项，导致需要朗之万修正步。

2.2 二阶近似 (Second-order Approximation)

为了计算修正项，作者引入了Tweedie 框架，假设逆向扩散核可以用高斯分布近似：
$$\hat{q}_{0|t}(\theta_0 | \theta_t, x) \approx \mathcal{N}(\theta_0; \mu_t, \Sigma_t)$$
其中均值 $\mu_t$ 和协方差 $\Sigma_t$ 可以通过得分函数及其导数计算得出。

在此假设下，修正项 $L_\lambda$ 被近似为高斯分布的乘积，其对数形式转化为得分函数的线性组合。最终导出了高维后验得分的近似公式（公式 16）：
$$\nabla_{\theta_t} \log p_t(\theta_t | x^*_{1:n}) \approx \Lambda(\theta_t)^{-1} \left( \sum_{j=1}^n \Sigma_{t,j}^{-1} \nabla \log p_t(\theta_t | x^*_j) + (1-n)\Sigma_{t,\lambda}^{-1} \nabla \log p^\lambda_t(\theta_t) \right) + F$$
其中 $\Lambda$ 是精度矩阵（Precision Matrix）的和，$F$ 是残差项。

2.3 两种实现算法

为了实际计算上述公式中的协方差矩阵 $\Sigma_{t,j}$ 并消除残差项 $F$（通过假设协方差为常数），作者提出了两种算法：

1. JAC (Jacobian Approximation):

   • 利用雅可比矩阵近似协方差：$\Sigma_{t,j} \propto \nabla_{\theta_t} \mu_{t,j}$。
   • 优点： 理论精度高。
   • 缺点： 需要计算 $m \times m$ 的雅可比矩阵，计算量大；且对得分网络的输入求导不稳定，容易受噪声影响。

2. GAUSS (Gaussian Approximation):

   • 核心策略： 假设协方差矩阵是常数。
   • 实现： 对于每个观测 $x^*_j$，先运行少量步骤的 DDIM 采样，计算样本的经验协方差 $\hat{\Sigma}_j$，将其作为 $\Sigma_{t,j}$ 的初始估计。
   • 优点： 避免了雅可比计算，数值极其稳定，且 $F$ 项自动为零。
   • 适用性： 特别适用于得分函数是学习得到的（存在误差）的场景。

采样流程：
将上述近似得分直接输入到 DDIM (Denoising Diffusion Implicit Models) 采样器中，从噪声逐步生成后验样本。这完全取代了 F-NPSE 中的朗之万步骤。

---

3. 主要贡献 (Key Contributions)

1. 理论突破： 推导了高维数据后验扩散过程的显式近似得分公式，填补了组合得分与确定性扩散采样之间的理论空白。
2. 算法创新： 提出了 JAC 和 GAUSS 两种算法，成功消除了对昂贵且不稳定的朗之万动力学（Langevin Dynamics）的依赖。
3. 性能提升：
   • 速度： 相比 F-NPSE（朗之万），GAUSS 快约 2.5 倍，JAC 快约 1.8 倍。
   • 稳定性： 在噪声存在或得分函数不完美（学习得到）的情况下，GAUSS 表现出极强的鲁棒性，而朗之万方法容易发散。
4. 可扩展性验证： 在合成数据、标准 SBI 基准测试（SLCP, SIR, Lotka-Volterra）以及真实的**计算神经科学模型（Jansen-Rit Neural Mass Model）**上进行了验证，证明了其在处理复杂、高维真实世界问题时的有效性。

---

4. 实验结果 (Results)

• 高斯玩具模型 (Gaussian Toy Models)：
  • 在解析解已知的情况下，GAUSS 在所有设置下均优于朗之万采样器，且随着观测数 $n$ 增加，优势更明显。
  • JAC 在无噪声时精度极高，但在有噪声时迅速变得不稳定。
• SBI 基准测试 (SBIBench)：
  • 在 SLCP（多模态）、SIR（单模态）、Lotka-Volterra（高维）任务中，GAUSS 在切片 Wasserstein 距离 (sW)、最大均值差异 (MMD) 和分类器双样本检验 (C2ST) 等指标上均显著优于 F-NPSE。
  • 随着观测数 $n$ 增加，F-NPSE（朗之万）的性能急剧下降甚至发散，而 GAUSS 保持稳健。
• 真实世界应用 (Jansen-Rit Neural Mass Model)：
  • 应用于脑电（EEG）信号模拟器的参数推断。
  • 结果显示，随着观测数据量增加，后验分布能够正确收敛到真实参数，并能揭示参数间的耦合不确定性（如 $g$ 与 $(\mu, \sigma)$ 的耦合），证明了该方法在处理病态推断问题上的能力。
• 对比其他确定性采样器：
  • 与 Geffner et al. (2023) 提出的另一种确定性采样器（基于各向同性方差假设）相比，GAUSS 在保持速度的同时，通过保留协方差结构，在精度和稳定性上更优，特别是在非各向同性后验分布中。

---

5. 意义与影响 (Significance)

1. 解决 SBI 中的采样瓶颈： 成功将高维数据推断从“昂贵的 MCMC/朗之万采样”转变为“高效的确定性扩散采样”，极大地降低了计算成本。
2. 提升参数可识别性： 证明了利用多个观测值（Tall Data）结合组合得分方法，可以有效解决复杂模型中的参数不可识别问题，且无需重新训练模型。
3. 推动科学发现： 该方法为计算神经科学、天体物理等依赖复杂模拟器的领域提供了一种可扩展、稳健的工具，使得从大量观测数据中反推复杂物理/生物机制成为可能。
4. 方法论启示： 展示了通过显式建模扩散过程的逆向动力学（而非仅仅依赖启发式修正），可以显著提高生成式模型在科学推断任务中的可靠性。

总结： 本文提出了一种基于扩散模型的高维数据后验采样新范式，通过数学近似消除了朗之万采样的必要性，在速度、稳定性和准确性上全面超越了当前的 SOTA 方法（F-NPSE），为科学计算中的复杂参数推断提供了强有力的工具。