[论文解读] Path Integral Sampler: a stochastic control approach for sampling

Qinsheng Zhang, Yongxin Chen|arXiv (Cornell University)|Nov 30, 2021

Model Reduction and Neural Networks参考文献 55被引用 6

一句话总结

路径积分采样器（PIS）是一种新颖的基于随机控制的采样算法，通过神经网络控制的随机微分方程（SDE）求解薛定谔桥问题，从非归一化目标密度中生成样本。该方法支持自由形式架构的端到端训练，在次优控制条件下提供理论上的Wasserstein距离保证，并利用基于路径积分的重要度权重纠正偏差，其在样本质量和效率方面优于现有方法。

ABSTRACT

We present Path Integral Sampler~(PIS), a novel algorithm to draw samples from unnormalized probability density functions. The PIS is built on the Schrödinger bridge problem which aims to recover the most likely evolution of a diffusion process given its initial distribution and terminal distribution. The PIS draws samples from the initial distribution and then propagates the samples through the Schrödinger bridge to reach the terminal distribution. Applying the Girsanov theorem, with a simple prior diffusion, we formulate the PIS as a stochastic optimal control problem whose running cost is the control energy and terminal cost is chosen according to the target distribution. By modeling the control as a neural network, we establish a sampling algorithm that can be trained end-to-end. We provide theoretical justification of the sampling quality of PIS in terms of Wasserstein distance when sub-optimal control is used. Moreover, the path integrals theory is used to compute importance weights of the samples to compensate for the bias induced by the sub-optimality of the controller and time-discretization. We experimentally demonstrate the advantages of PIS compared with other start-of-the-art sampling methods on a variety of tasks.

研究动机与目标

• 开发一种灵活的、可端到端训练的采样方法，用于非归一化概率密度，而无需显式密度模型。
• 利用薛定谔桥框架，将采样问题形式化为与目标分布相关的终端代价的随机最优控制问题。
• 在少于传统MCMC方法的步数下实现高质量采样，同时提供理论性能边界。
• 通过路径积分理论计算重要度权重，校正次优控制和时间离散化带来的偏差。
• 通过隐式、可扩展且可微分的采样框架，克服变分推断（显式密度约束）和MCMC（混合缓慢、难训练的提议）的局限性。

提出的方法

• 将采样问题形式化为随机最优控制问题，其中运行代价为控制能量，终端代价源自目标密度。
• 使用神经网络建模控制策略，以支持自由形式架构设计和端到端训练。
• 利用Girsanov定理重加权SDE动力学，通过受控扩散将初始维纳过程与目标分布关联。
• 应用路径积分理论计算样本的重要度权重，补偿控制策略中的次优性和离散化误差。
• 通过受控SDE的终端分布与目标测度之间的KL散度目标训练控制器。
• 采用Euler-Maruyama SDE积分方法，无需伴随方法以增强稳定性，尤其在低维情况下；在内存受限时使用可逆Heun方法。

实验结果

研究问题

• RQ1 能否基于薛定谔桥问题的随机控制框架，设计一种灵活的、可端到端训练的非归一化密度采样算法？
• RQ2 如何校正次优控制策略，以确保渐近无偏采样？
• RQ3 在使用次优控制器时，可为采样质量提供哪些理论性能保证？
• RQ4 时间跨度T和时间步长Δt的选择如何影响训练稳定性和样本质量？
• RQ5 PIS能否在多样化的任务中，于样本质量和效率方面超越MCMC和归一化流等最先进方法？

主要发现

• 即使使用次优控制策略，PIS在Wasserstein距离方面仍提供采样质量的理论保证。
• 基于路径积分的重要度权重有效校正了次优控制和时间离散化带来的偏差。
• 较大的时间跨度T可提升训练稳定性和模式覆盖能力，而较小的T会导致高Lipschitz常数的控制策略，更难训练。
• 较小的时间步长Δt可减少离散化误差并提升样本质量，尤其在结合更高积分步数时效果更明显。
• 在高斯混合、漏斗分布以及MNIST和CIFAR-10等真实数据集任务中，PIS在样本质量和效率方面均优于基线方法。
• 使用批量归一化和梯度裁剪可提升训练稳定性，且实践中Euler-Maruyama方法（无需伴随）比基于伴随的方法更可靠。