跳转至

Diffusion Path Samplers via Sequential Monte Carlo

会议: NeurIPS 2026
arXiv: 2601.21951
论文: https://arxiv.org/abs/2601.21951
领域: 优化/理论
关键词: 扩散路径采样器, 序贯蒙特卡洛, 分数匹配, 控制变量, 随机插值

一句话总结

针对非归一化目标分布的采样难题,本文提出基于序贯蒙特卡洛的扩散路径采样框架(DPSMC),通过序列化演化条件后验辅助变量实现无偏边际密度与高效可并行的分数估计,并结合最优控制变量调度将分数方差压至最低。

研究背景与动机

从仅已知非归一化密度(已知能量函数或未归一化常数)的目标分布中进行采样,是贝叶斯统计推断、统计物理与机器学习中的核心基石问题。直接对高维、多峰或复杂几何流形上的目标分布进行采样往往极度困难,传统方法通常引入辅助分布路径在易采样的基底分布与目标分布间构建桥梁。然而,经典的几何退火路径(如退火重要性采样 AIS)在多峰和复杂流形下极易遭遇质量瞬移(mass teleportation)与对数索伯列夫常数(log-Sobolev constant)退化崩溃,导致采样粒子在不同峰之间迁移受阻、方差剧增。

扩散模型启发下的扩散路径凭借优良的数学正则性成为极具前景的替代方案——在温和条件下它具有一致有界的泛函不等式与有限作用量。然而,沿着时间反演扩散路径或扩散退火朗之万动力学(DALD)推进时,每一步都需要计算难以直接解析评估的边际分数函数 \(\nabla \log \mu_t(x)\)。现有蒙特卡洛分数估计方法陷入了两难困境:要么依赖 MCMC(如非调整朗之万 ULA/MALA)在每个离散时间步独立运行链条逼近后验分布,这不仅产生无法消除的串行收敛瓶颈,且随着时间推进需重复平衡;要么借助零阶拒绝采样,但其计算复杂度随问题维度指数级暴增;更关键的是,已有方法通常无法提供路径边际密度的无偏估计,无法天然契合分层 SMC 或正规化常数估计管线。

解决这一矛盾的关键契机在于:扩散路径的卷积结构天然诱导了关于干净辅助变量 \(Y\) 的后验条件分布 \(\varrho_{t,x}\),且沿着扩散时间轴推进时,该后验分布的演化是平滑连续的。本文的核心 idea:将扩散路径上的分数与边际密度估计转化为对辅助变量后验分布的序贯蒙特卡洛(SMC)追踪过程,结合严谨推导的降方差控制变量调度与动作泛函极小化路径设计,构建高度并行化且具有非渐近收敛保证的免训练采样器 DPSMC。

方法详解

整体框架

DPSMC 将高维扩散动力学粒子的空间轨迹推进与辅助变量集合的序贯蒙特卡洛(SMC)演化紧密解耦并协同推进。在时间区间 \(t \in [0, 1]\) 内,主粒子 \(X_k\) 沿着扩散路径的随机微分方程(SDE,涵盖 OU 反向过程、随机插值 SI 或扩散退火动力学 DALD)演化;而在每一个时间步 \(t_k\),算法利用上一时刻的辅助变量粒子系,通过设计的前向 MALA 转移核、重构的反向核与自适应重采样机制,逼近当前条件后验分布 \(\varrho_{t_k, X_k}\),从而闭式计算当前粒子所需的高精度分数估计量与无偏边际密度。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["初始状态<br/>主粒子采样与辅助变量初始化"] --> B["时间推进与主粒子积分<br/>OU / SI / DALD 离散步"]
    B --> C["SMC辅助变量演化<br/>MALA向前核传递与权重更新"]
    C --> D["自适应重采样<br/>ESS低于阈值时重构均匀粒子群"]
    D --> E["控制变量分数估计<br/>MCVSI / MMCVSI 最小化方差"]
    E -->|循环直至终点 t=1| F["目标分布粒子与归一化常数估计"]

关键设计

1. 基于 Feynman-Kac 路径测度的无偏边际密度与分数估计:打破 MCMC 逐时间步平衡瓶颈

针对传统扩散蒙特卡洛采样器(如 RDMC、SLIPS)在每个扩散时间步必须重新启动一条全新 MCMC 链条导致计算开销大且产生串行延迟的痛点,DPSMC 建立在时间连续的条件后验分布 \(\varrho_{t_k, X_k}(y) \propto \nu\left(\frac{X_k - \sqrt{\lambda_{t_k}}y}{\sqrt{1-\lambda_{t_k}}}\right) \pi(y)\) 基础之上。算法定义前向马尔可夫核 \(\mathsf{K}_k\) 与反向核 \(\mathsf{L}_{k-1}\),引入 Feynman-Kac 递推势函数: $\(G_k(y_{k-1}, y_k) = \frac{\tilde{\rho}_k(y_k)\mathsf{L}_{k-1}(y_k, y_{k-1})}{\tilde{\rho}_{k-1}(y_{k-1})\mathsf{K}_k(y_{k-1}, y_k)}\)$ 当选取反向核为前向核关于 \(\rho_k\) 的时间反演伴随核时,权重更新项精简为退火重要性采样形态 \(G_k(y_{k-1}, y_k) = \frac{\tilde{\rho}_k(y_{k-1})}{\tilde{\rho}_{k-1}(y_{k-1})}\)。由此,主粒子的时间变动分数函数与路径边际密度被严谨表示为关于前向路径测度的期望之比与期望积分: $\(\nabla \log \mu_{t_k}(X_k) = \frac{\mathbb{E}\left[\varphi_{k, X_k}(Y_k) \prod_{p=0}^k G_p(Y_{p-1}, Y_p)\right]}{\mathbb{E}\left[\prod_{p=0}^k G_p(Y_{p-1}, Y_p)\right]}, \quad \mu_{t_k}(X_k) = \mathbb{E}\left[\prod_{p=0}^k G_p(Y_{p-1}, Y_p)\right]\)$ 该设计使辅助变量粒子群可以在扩散时间步之间复用与迁移,不仅消除了内层重复预热收敛,还天然给出了路径边际密度的无偏估计量,允许嵌套进上层正规化常数计算。

2. 边际化控制变量分数恒等式(MCVSI / MMCVSI):最小化期望方差的自适应融合

去噪分数恒等式(DSI)与目标分数恒等式(TSI)在扩散两端存在固有的数值发散缺陷:DSI 在 \(t \to 1\)(\(\lambda_t \to 1\))时分母奇异发散,而 TSI 在 \(t \to 0\)(\(\lambda_t \to 0\))时方差急剧放大。本文引入广义矩阵权重参数族 \(\phi_{t,x}^{\mathbf{A}}(y) = \frac{1}{\sqrt{1-\lambda_t}}\mathbf{A} \nabla \log \nu\left(\frac{x-\sqrt{\lambda_t}y}{\sqrt{1-\lambda_t}}\right) + \frac{1}{\sqrt{\lambda_t}}(\mathbf{I}-\mathbf{A})\nabla \log \pi(y)\),并求解其在主粒子边际分布下的期望方差极小化问题: $\(\mathbf{A}_t^* = \arg\min_{\mathbf{A} \in \mathbb{R}^{d \times d}} \mathbb{E}_{X \sim \mu_t}\left[\mathrm{Var}_{Y \sim \varrho_{t,X}}\left[\phi_{t,X}^{\mathbf{A}}(Y)\right]\right] = \frac{1}{\lambda_t}\mathcal{I}_\pi \left(\frac{1}{1-\lambda_t}\mathcal{I}_\nu + \frac{1}{\lambda_t}\mathcal{I}_\pi\right)^{-1}\)$ 其中 \(\mathcal{I}_p = \mathbb{E}_p[\nabla \log p(X) \nabla \log p(X)^\top]\) 为分数协方差矩阵;若退化至标量调度 \(\mathbf{A}_t = \alpha_t \mathbf{I}\),极小解显式化为 \(\alpha_t^* = \frac{\frac{1}{\lambda_t}\mathrm{Var}_\pi[\nabla \log \pi(X)]}{\frac{1}{1-\lambda_t}\mathrm{Var}_\nu[\nabla \log \nu(X)] + \frac{1}{\lambda_t}\mathrm{Var}_\pi[\nabla \log \pi(X)]}\)。更关键的是,利用论文建立的 Stein 引理恒等式 \(\mathcal{I}_\pi = \mathbb{E}_{X \sim \mu_t}[\mathbb{E}_{Y \sim \varrho_{t,X}}[\nabla \log \pi(Y) \nabla \log \varrho_{t,X}(Y)^\top]]\),矩阵协方差可直接复用当前 SMC 粒子群在线估计,无需引入额外的真实能量梯度求值开销,在强各向异性目标分布(如 Funnel)上显著削减分数误差。

3. 基于测度曲线动作泛函极小化的最优扩散路径与非奇异随机插值(SI)

传统反向 OU 扩散依赖无限时间视界或人工设定的截断参数 \(T\),且起点无法严格精确置于标准高斯。为从动力学源头上提升粒子传输效率,本文以扩散路径在 Wassertein 空间上的测度曲线速度场为对象,考察其动作泛函上界 \(\mathcal{A}(\mu) = \int_0^T |\dot{\mu}|_t^2 \mathrm{d}t \le \int_0^T \left(\frac{\sigma^2 d}{4(1-\lambda_t)} + \frac{\mathbb{E}_\pi[\|X\|^2]}{4\lambda_t}\right) \dot{\lambda}_t^2 \mathrm{d}t\)。通过变分法求解欧拉-拉格朗日方程,证明当基底方差匹配目标二阶矩 \(\sigma^2 = \mathbb{E}_\pi[\|X\|^2]/d\) 时,使得泛函取极小的最优时间调度解析形式为余弦调度: $\(\lambda_t^* = \sin^2\left(\frac{\pi t}{2}\right)\)$ 该调度在起点缓慢加速帮助粒子探索与模式锁定,在终点平稳减速以细化局部几何。更进一步,结合 DSI 结构消除 \(t \to 0\) 处的漂移奇点,本文构造了扩散系数恒定为 \(\gamma_t = \sigma^2\) 的随机插值动力学(Stochastic Interpolant, SI),彻底摆脱了退火参数 \(\epsilon\) 或有限截断 \(T\) 等敏感超参数的困扰,实现了从先验高斯到目标分布的有限步精确动力学传输。

损失函数 / 训练策略

DPSMC 是一种免训练(training-free)的蒙特卡洛积分算法,不涉及神经网络参数反向传播。其实际算法流程(Algorithm 1)主要超参数包括:离散时间步数 \(K=1024\)、辅助变量粒子数 \(N_Y=128\)、有效样本容量(ESS)重采样阈值 \(\kappa=0.5\) 以及适应度 MALA 步长。初始辅助变量采用带有紧支撑半径估计的高斯包络 \(q_0(x) = \mathcal{N}(x; 0, \tilde{\sigma}^2 \min(1, (1-\lambda_0)/\lambda_0)\mathbf{I})\),在保证充分探索的同时避免权重坍塌。

实验关键数据

主实验

论文在多峰分布(GMM40,含 40 个随机高斯混合分量,测试 2 维与 50 维)、复杂几何分布(4 个同心圆 Rings \(d=2\)、强各向异性漏斗 Funnel \(d=10\))以及高维真实贝叶斯逻辑回归后验(Ionosphere \(d=35\) 与 Sonar \(d=61\))上全面对比了 DPSMC 与经典退火基线(AIS、SMC)及扩散蒙特卡洛基线(RDMC、SLIPS)。在相同能量函数评估预算下,DPSMC 依靠粒子并行化将批处理调用轮次减少了 32 至 131 倍。

算法 批处理评估轮次 (↓) GMM40 (\(d=2\)) \(\epsilon\text{-}W_2\) (↓) GMM40 (\(d=50\)) \(\epsilon\text{-}W_2\) (↓) Rings (\(d=2\)) \(\epsilon\text{-}W_2\) (↓) Funnel (\(d=10\)) KS 距离 (↓) Ionosphere (\(d=35\)) 测试对数似然 (↑) Sonar (\(d=61\)) 测试对数似然 (↑)
AIS 131K \(2.35 \pm 0.41\) \(115.13 \pm 2.98\) \(0.19 \pm 0.01\) \(0.037 \pm 0.003\) \(-87.79 \pm 0.10\) \(-110.94 \pm 0.10\)
SMC 131K \(5.50 \pm 0.22\) \(111.81 \pm 4.84\) \(\mathbf{0.18 \pm 0.01}\) \(0.035 \pm 0.003\) \(-87.79 \pm 0.10\) \(-110.94 \pm 0.10\)
RDMC 32K \(12.35 \pm 0.11\) \(79.37 \pm 0.10\) \(0.29 \pm 0.02\) \(0.082 \pm 0.003\) \(-109.97 \pm 0.23\) \(-129.82 \pm 0.18\)
SLIPS 32K \(4.07 \pm 0.40\) \(27.57 \pm 0.83\) \(\mathbf{0.18 \pm 0.01}\) \(0.073 \pm 0.003\) \(-87.88 \pm 0.10\) \(-110.88 \pm 0.06\)
DPSMC (DALD) 1K \(1.92 \pm 0.22\) \(29.05 \pm 1.51\) \(0.20 \pm 0.02\) \(0.062 \pm 0.003\) \(-83.77 \pm 0.07\) \(-108.45 \pm 0.05\)
DPSMC (DALD w/ LL) 1K \(1.57 \pm 0.33\) \(\mathbf{26.82 \pm 1.38}\) \(0.19 \pm 0.02\) \(0.046 \pm 0.004\) \(-85.22 \pm 0.11\) \(-109.34 \pm 0.08\)
DPSMC (OU) 1K \(\mathbf{1.14 \pm 0.24}\) \(27.27 \pm 1.06\) \(\mathbf{0.18 \pm 0.02}\) \(0.028 \pm 0.002\) \(\mathbf{-83.73 \pm 0.16}\) \(\mathbf{-108.23 \pm 0.07}\)
DPSMC (SI) 1K \(1.23 \pm 0.30\) \(28.78 \pm 1.78\) \(\mathbf{0.18 \pm 0.03}\) \(\mathbf{0.026 \pm 0.004}\) \(-87.54 \pm 0.11\) \(-111.06 \pm 0.08\)

消融与归一化常数估计实验

得益于 DPSMC 提出的 Feynman-Kac 公式给出的无偏路径边际密度估计,该算法能够无缝嵌入 RDSMC 嵌套重要性采样管线中,精确估计非归一化密度的对数配分函数(归一化常数 \(\log Z_\pi\))。下表展示了在 Rings (\(d=2\)) 与 Funnel (\(d=10\)) 上的正规化常数偏差及样本几何品质对比。

算法 Rings 半径 TVD (↓) Rings 角度 TVD (↓) Rings \(\log Z_\pi\) 偏差 (↓) Funnel KS 距离 (↓) Funnel \(\log Z_\pi\) 偏差 (↓)
AIS \(\mathbf{0.10 \pm 0.00}\) \(0.15 \pm 0.00\) \(0.05 \pm 0.00\) \(0.07 \pm 0.00\) \(0.28 \pm 0.01\)
SMC \(\mathbf{0.10 \pm 0.00}\) \(0.15 \pm 0.00\) \(0.05 \pm 0.00\) \(0.07 \pm 0.00\) \(0.28 \pm 0.01\)
RDSMC \(0.13 \pm 0.01\) \(0.21 \pm 0.01\) \(0.03 \pm 0.00\) \(0.11 \pm 0.00\) \(0.28 \pm 0.10\)
DPSMC (OU) \(0.31 \pm 0.01\) \(\mathbf{0.14 \pm 0.01}\) \(\mathbf{0.02 \pm 0.01}\) \(\mathbf{0.06 \pm 0.00}\) \(\mathbf{0.12 \pm 0.07}\)

关键发现

  • 在高维多峰混合高斯(GMM40, \(d=50\))中,传统几何退火方法(AIS 达到 115.13,SMC 达到 111.81)彻底失效发生粒子堆叠,而 DPSMC 系列方法稳定收敛在 26.82-29.05 区间内,验证了扩散路径在规避能量壁垒和质量瞬移方面的本质理论优势。
  • 在漏斗分布 Funnel 中,坐标轴方差的巨大尺度不匹配导致各向同性分数估计严重失真。引入基于全矩阵协方差的 MMCVSI 后,DPSMC (SI) 的 KS 误差大幅下降至 0.026,相比未采用矩阵控制变量的传统扩散方法降低了超过 60%。
  • 在批处理硬件场景下,DPSMC 将所需序列步数由传统 MCMC-based 扩散采样的 32K-131K 降低至仅需 1K,显著压缩了实际墙上时间(wall-clock time)。

亮点与洞察

  • 将扩散分数估计转化为辅助变量的序贯 Monte Carlo 问题:巧妙利用了扩散过程的时间邻近性,前一步的后验粒子天然构成了下一步后验的优质提议分布,避免了传统 MCMC 每次从头采样的重复开销。
  • 解析求解最优控制变量矩阵调度:不仅指出了 DSI 与 TSI 在扩散两端的互补性,更进一步从理论上给出了极小化边际期望方差的最优矩阵解 \(\mathbf{A}_t^*\),并通过 Stein 引理将其转化为可零额外开销计算的经验期望。
  • 动作泛函极小化导出余弦调度:为扩散路径的时间离散步长选择提供了严格的最优控制理论依托,证明了经验使用的余弦退火正是极小化 Wasserstein 测度加速度的解析最优解。

局限与展望

  • 终点狄拉克收缩的不稳定性:当时间 \(t \to 1\) 时,条件后验 \(\varrho_{t,X}\) 逐渐退化为狄拉克测度 \(\delta_X\),造成 MALA 转移核的拒绝率升高与粒子退化,未来需引入重参数化或高阶隐式流转移核加以平滑。
  • 嵌套常数估计的方差累积:在嵌入分层 RDSMC 进行归一化常数计算时,重要性权重的连乘在长序列步数下仍存在方差放大的现象,探索更加平滑的非高斯随机插值族是关键延伸方向。

相关工作与启发

  • vs RDMC (Huang et al., 2024) / SLIPS (Grenioux et al., 2024):两者均使用非调整朗之万算法(ULA)在每个时间步独立运行 MCMC 逼近后验并估计分数,存在严重的序列等待和累积偏差;DPSMC 改用 SMC 在时间轴上平滑延续后验样本,计算效率提升数倍且支持完全批处理。
  • vs RDSMC (Wu et al., 2025):RDSMC 在外层主粒子上做 SMC 并在内层每次通过 AIS 独立估计边际密度;DPSMC 聚焦于辅助变量粒子系演化,本身给出的边际密度是严格无偏的,二者具有高度的互补与嵌套兼容性。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次系统性将序贯蒙特卡洛与扩散路径分数估计深度融合,并推导出最优矩阵控制变量调度。
  • 实验充分度: ⭐⭐⭐⭐☆ 覆盖多峰高维合成基准、几何特异性分布、高维真实贝叶斯推断及归一化常数评估,对比详实。
  • 写作质量: ⭐⭐⭐⭐⭐ 数学推导严密精巧,从 Feynman-Kac 测度路径到非渐近误差界的理论链条完整清晰。
  • 价值: ⭐⭐⭐⭐⭐ 为免训练扩散采样器和高维复杂后验采样提供了高并行度、方差受控的新一代通用范式。