Path Integral Sampler: a stochastic control approach for sampling (2111.15141v2)

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\"odinger 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\"odinger 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.

Path Integral Sampler: A Stochastic Control Approach for Sampling

The paper presents the Path Integral Sampler (PIS), a novel algorithm designed to sample from unnormalized probability density functions using stochastic control methodology. The foundation of PIS is the Schrödinger bridge problem, which aims to ascertain the most probable path of a diffusion process given its initial and terminal distributions. The algorithm leverages this concept to conduct sampling by propagating samples from the initial distribution through the Schrödinger bridge to reach the terminal distribution.

Theoretical Foundations

The formulation of PIS involves solving a stochastic optimal control problem where the cost is defined as the control energy while imposing a terminal cost based on the target distribution. This is achieved using the Girsanov theorem with a straightforward prior diffusion. The optimal control policy that dictates the sample trajectory is derived by solving a Hamilton-Jacobi-BeLLMan (HJB) equation, which is transformed into a linear partial differential equation through a well-known transformation. The optimal control can represent paths accurately and efficiently, increasing the sampling quality.

Implementation via Neural Networks

The control mechanism within PIS is modeled using a neural network, allowing end-to-end training based on empirical samples. This setup not only facilitates learning the optimal path efficiently but also provides flexibility in network architecture design, which is often constrained in explicit density models used in traditional Variational Inference (VI). The authors propose using additional gradient information from the target density to enhance the neural network's performance, a technique that empirically accelerates convergence and aids in overcoming local optima issues.

Performance and Evaluation

The paper provides a theoretical guarantee of the sampling quality achieved by PIS, particularly by evaluating the Wasserstein distance between PIS-generated samples and the target distribution. Furthermore, it addresses suboptimal control policies by incorporating the path integral theory to compute importance weights, thereby correcting biases introduced by suboptimality or time discretization errors.

PIS's performance is compared empirically against other state-of-the-art methods across different tasks, demonstrating its efficiency and ability to generate high-quality samples with fewer steps. Specifically, PIS outperforms typical Monte Carlo methods that often suffer from slow convergence, especially in high-dimensional spaces or in distributions with well-separated modes.

Contributions and Future Work

The contributions of PIS include introducing a generic sampler built upon a stochastic control problem, offering a training framework that quantifies sampling performance via a set evaluation metric, and the ability to provide unbiased samples by recalibrating using importance weights.

Challenges remain in optimizing the initialization and scaling for high-dimensional applications, and future research could explore more elaborate model structures or incorporate domain knowledge to improve the initial propose distribution.

In summary, the Path Integral Sampler is a viable and competitive method for probabilistic sampling in machine learning applications, driven by a foundation grounded in optimal control theory and efficient network-based implementation.