PaperlayerView exact revision on arXiv

Source-linked AI summary

Path Integral Sampler: a stochastic control approach for sampling

Qinsheng Zhang, Yongxin Chen

arXiv:2111.15141v2cs.LG

TL;DR

Sampling from unnormalized densities remains challenging, motivating methods that avoid restrictive explicit-density parameterizations and difficult MCMC behavior. PIS uses a Schrödinger bridge formulated as stochastic optimal control, with a neural-network policy and path-integral importance weights. The paper provides Wasserstein-based quality guarantees for sub-optimal control and reports advantages across varied sampling tasks.

  • Problem

    Generating asymptotically unbiased samples from target distributions known only up to a normalizing constant remains challenging, while existing VI and MCMC methods have important restrictions or difficulties.

  • Method

    PIS transports samples from an initial distribution to the target through a Schrödinger bridge, formulated as stochastic optimal control with a neural-network policy and importance weighting.

  • Results

    PIS provides Wasserstein-based justification for sampling quality under sub-optimal control and uses importance weights to compensate for policy and time-discretization bias.

  • Takeaways & Limitations

    PIS offers an implicit sampling model whose calibration can preserve sampling quality with sub-optimal policies and supports estimation of normalization constants.

  • Takeaways & Limitations

    PIS requires training neural networks for each given distribution, adding computational overhead, and finite-step quality depends on policy optimality and hyperparameter choices.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper introduces Path Integral Sampler (PIS), which formulates sampling from unnormalized densities as a trainable stochastic-control problem. PIS uses an implicit model and importance weighting to address limitations of existing variational and Monte Carlo methods.

  • 1 INTRODUCTION: Variational inference requires an explicit proposal density, restricting parameterizations such as normalizing flows through bijectivity and tractable Jacobians.PIS instead uses an implicit model.
  • 1 INTRODUCTION: MCMC methods can have long mixing times, become stuck in local minima, and be difficult to evaluate or train when augmented with neural networks.Existing neural-network MCMC work often focuses on data-aware proposals.
  • 1 INTRODUCTION: PIS generates samples by simulating a stochastic differential equation in finite steps.Its control policy is modeled with a neural network and trained gradually.
  • 1 INTRODUCTION: The method connects sampling with the Schrödinger bridge problem, whose controlled diffusion can reach a target density through an appropriate terminal cost.The bridge describes the most likely diffusion evolution between endpoint distributions.
  • 1 INTRODUCTION: PIS can use path-integral importance weights to compensate for bias from sub-optimal control policies and time discretization.The method also uses the discrepancy from the optimal policy as an evaluation metric.

2 SAMPLING AND STOCHASTIC CONTROL PROBLEMS

The paper frames sampling as transporting particles from an easy initial distribution to a target distribution using stochastic processes. It evaluates these processes through path-space divergence and formulates controlled SDE dynamics whose optimal policy solves the sampling objective.

  • 2.1 SAMPLING PROBLEMS: Sampling constructs a stochastic process that drives particles from an easily sampled initial distribution ν to a target distribution µ.The target may be specified only up to a normalization constant.
  • 2.1 SAMPLING PROBLEMS: Variational inference uses a parameterized stochastic process whose terminal marginal µQ approximates the target µ.The path-space measure induced by the process is compared with the target-process measure.
  • 2.1 SAMPLING PROBLEMS: A small D_KL(Q∥P) value serves as a performance criterion for the sampler.The divergence is defined over path measures and bounds the divergence between terminal marginals.
  • 2.2 STOCHASTIC CONTROL: The stochastic-control model follows d x_t = u_t d t + d w_t with x_0 ∼ ν, where u_t is the control and w_t is standard Brownian motion.The objective is to find a feedback control strategy minimizing a specified cost.
  • 2.2 STOCHASTIC CONTROL: The cost considered has a terminal-cost component Ψ, and dynamic programming solves the corresponding optimal-control problem through the Hamilton-Jacobi-Bellman equation.The value function V_t(x) is the cost-to-go function used to compute the optimal policy.

3 PATH INTEGRAL SAMPLER

PIS formulates sampling as stochastic optimal control through a Schrödinger bridge, then parameterizes the control with neural networks for end-to-end training. Importance weights and Wasserstein analysis address bias from sub-optimal control and time discretization.

  • Sampling as stochastic optimal control problem: The stochastic control formulation uses a terminal loss tied to the target density, and its optimal policy transports samples from ν to µ.The optimal controlled process is also interpreted as the posterior of uncontrolled dynamics conditioned on the target distribution.
  • Sampling as stochastic optimal control problem: Path Integral control evaluates the optimal value function using expectations over uncontrolled trajectories, avoiding direct knowledge of the optimal policy.A logarithmic transformation and the Feynman-Kac formula convert the nonlinear HJB equation into a linear PDE representation.
  • Optimal control policy and sampler: For general targets, Monte Carlo value-function estimation can have intolerable variance and exponentially increasing sample requirements with dimension.These limitations motivate replacing empirical value estimation with a neural-network control policy.
  • Optimal control policy and sampler: PIS parameterizes control with neural networks and can incorporate the target score ∇log µ(x), which empirically accelerates convergence and reduces D_KL(Q_u∥Q*).The score-informed policy is described as modulated Langevin dynamics reaching µ within finite time T.
  • Optimal control policy and sampler: The learned sampler is biased by non-optimal neural control and SDE discretization, so the paper provides a Wasserstein-distance bound for finite sampling steps.The training procedure uses backpropagation through Neural SDE trajectories.
  • Importance sampling: Importance weights compensate for mismatch between trained and optimal policies, while PIS also supports unnormalized targets and normalization-constant estimation.When the policy is optimal, all particles share the same importance weight; the unnormalized terminal cost differs by the constant −log Z.

4 EXPERIMENTS

Experiments evaluate PIS against MCMC, SMC, AFT, and variational methods across synthetic, molecular, and VAE-based sampling tasks. PIS benefits from gradient guidance and importance weighting, with advantages especially visible for separated multimodal targets.

  • Experimental setup: PIS is evaluated against gradient-guided MCMC, annealed SMC methods, and variational approaches across Bayesian inference, normalization estimation, molecular sampling, and VAE posterior tasks.The benchmark includes HMC, NUTS, SMC, AFT, VI-NF, and SNF.
  • PIS-Grad versus PIS-NN: Gradient guidance improves mode coverage, with PIS-Grad outperforming PIS-NN more clearly when target modes are well separated.PIS-NN and VI-NF can cover only part of a separated ring, whereas gradient information provides exploring directions.
  • Synthetic distributions: PIS produces samples visually indistinguishable from a separated two-dimensional Gaussian-mixture target when the modes are narrow and far apart.All methods work well when Gaussian modes are close; the advantage of PIS becomes clearer as variances decrease and modes separate.
  • Benchmarking datasets: Importance weighting improves log-normalization estimates, and PIS-Grad outperforms PIS-NN most clearly on multimodal MG and LGCP datasets.The difference is less pronounced for the unimodal Funnel distribution.
  • Benchmarking datasets: PISRW-Grad is better than AFT and SMC in all reported cases, while PISRW-NN can outperform annealed SMC and HMC on the Funnel distribution.These comparisons concern estimation of log normalization constants.
  • Alanine dipeptide: PIS achieves lower KL divergence than the compared variational approaches for overall atom states and five multimodal Alanine-dipeptide torsion angles.The comparison uses unweighted samples and methods without gradient information.
  • VAE posterior sampling: For VAE posterior normalization estimation on binary MNIST, PIS has lower bias and variance than the evaluated alternatives.The reference normalization constant is estimated from 10 long-run SMC runs with 1000 temperature levels.

5 CONCLUSION

The paper concludes that PIS connects sampling with stochastic control and can achieve target distributions under optimal control, while importance weights calibrate sub-optimal sampling. It also acknowledges training, finite-step, and numerical limitations.

  • Conclusion: PIS uses stochastic control to drive particles from a simple initial distribution to a target density when the control policy is optimal.The terminal cost depends on the target distribution.
  • Conclusion: Importance weights compensate for sampling-quality errors caused by sub-optimal policies.This calibration is presented as preserving sampling quality despite controller sub-optimality.
  • Limitations: PIS requires neural-network training for each given distribution, adding computational overhead that amortization can mitigate.This is contrasted with popular non-learnable MCMC algorithms.
  • Limitations: Finite-step sampling quality depends on network optimality, while unsuitable hyperparameters can cause numerical issues and failure modes.The paper discusses these failure modes in appendix G.1.

A PROOF OF THEOREM 1

The proof establishes identities for the optimally controlled diffusion and uses them to show that the associated terminal distribution matches the target under the stated construction.

  • Proof setup: The transition density associated with the optimal control policy is introduced as the central lemma for the theorem proof.The lemma is attributed to Dai Pra and Pavon.
  • Notation: The notation s,t(x, y) denotes the transition probability from state x at time s to state y at time t.This transition-density notation supports the preceding lemma.
  • Proof derivation: The proof uses the relation V_t(x) = −log φ_t(x) together with earlier equations to derive the required path-measure identities.The derivation begins from the initial Dirac distribution.
  • Proof conclusion: The expected optimal control cost plus terminal cost equals V_0(x̄_0) = 0 in the displayed result.This equality is linked to equation (13).

B PROOF OF IMPORTANCE WEIGHTS

The importance-weight analysis assumes regularity of the optimal policy and quantifies finite-step sampling error under a time discretization with step size Δt.

  • Proof strategy: The derivation proceeds by substituting earlier equations into the importance-weight expression.This substitution is stated immediately before the theorem-related calculations.
  • Assumptions: The analysis assumes the initial state is a Dirac distribution and imposes Lipschitz conditions on the optimal policy.These conditions ensure existence of a unique strong solution whose terminal distribution is the target μ.
  • Finite-step analysis: Under Condition 1 and sampling step size Δt, Theorem 2 provides a bound for the discretized process generated by the optimal policy.The formal theorem statement begins in the supplied passage.
  • Finite-step analysis: The proof introduces a lemma bounding E[∥x_t2 − x_t1∥^2] using constants C_0 and d and the time separation |t_2 − t_1|.The bound applies for t_1, t_2 ∈ [0, T].

C.2 PROOF OF THEOREM 2

The proof bounds the discrepancy between the discretized process under a sub-optimal policy and the optimally controlled trajectory. Choosing parameters proportional to the step size yields an O(dT(∆t + ϵ)) bound.

  • Proof strategy: The argument extends the discrete-time process under sub-optimal control to continuous time and bounds its deviation from the optimal-policy trajectory.The proof uses triangle, Cauchy–Schwarz, and quadratic inequalities to control the resulting terms.
  • Proof strategy: The proof successively substitutes intermediate bounds into earlier inequalities to obtain the final discretization-and-control-error estimate.These substitutions are described as plugging equations (30), (29), and (31) into preceding bounds.
  • Bound: With fixed step size ∆t, the parameter choice α = C1∆t, β = η = 1 leads to the stated asymptotic bound.The resulting expression is O(dT(∆t + ϵ)).

D PROOF OF THEOREM 3

The proof analyzes importance weights under a feasible control policy and derives variance bounds for correcting sampling bias. It also identifies the optimal policy through trajectory-distribution KL divergence.

  • Importance weights: The proof defines a path-dependent random variable and uses importance sampling with a density ratio to analyze the estimator.The importance weight is explicitly treated as a random variable.
  • Variance bound: The variance of the importance weight is bounded through an upper bound on its second moment.The proof uses Var(wu) + 1 = EQu[(wu)2] to obtain the final bound.

E PROOF OF THEOREM 4

The section connects PIS estimation to trajectory-based importance sampling and discusses practical sampling, training, and comparison settings. It also describes extensions of the prior SDE and relationships to existing samplers.

  • Normalization estimation: PIS estimates the normalization constant using trajectory distributions, with an unbiased estimator available under a sub-optimal policy through importance sampling.The derivation first estimates the constant and then applies importance sampling under policy u.
  • Comparisons: PIS explores mode-separated targets competitively without annealing and resampling, while SVGD remains less efficient and lower quality than PIS-Grad in the cited comparison.The Figure 4 comparison uses 100 steps, and the SVGD comparison reports degradation in efficiency as active particles increase.
  • Efficiency: PIS can be trained once and reused for sampling, allowing its training cost to be amortized when generating many samples.The training time depends strongly on NeuralSDE efficiency.
  • Efficiency: The adjoint SDE solver avoids caching the full computational graph, with reported memory use of around 1.5GB for toy and Funnel examples and 5GB for LGCP.The required memory is described as approximately one forward and backward pass of the control network.
  • Experimental settings: PIS uses 100 uniform SDE discretization steps in the reported experiments, while HMC and NUTS use separate trajectories rather than samples from one trajectory.The HMC setup uses 100 leapfrog steps, and the trajectory choice reduces concentration on a single mode.
  • SDE design: A more general SDE can encode prior knowledge through its drift and diffusion functions, with training easier when the prior terminal distribution is close to the target.PIS fits the optimal controlled process using a parameterized controlled process.
  • Related methods: The paper relates normalization estimation in SMC and AFT to incremental importance weights and treats HMC and NUTS as two-temperature SMC cases.The initial distribution for these SMC-based comparisons is standard Gaussian.
  • Related methods: VI-NF uses importance sampling for unbiased normalization-constant estimation, while PIS uses its ELBO and an unbiased estimator for PISRW.The normalizing-flow distribution q is normalized.

F.8 BENCHMARKING DATASETS

The benchmark section specifies datasets, comparison procedures, and sampling configurations across normalization-constant, molecular, and image experiments. It also reports long-run and slice-sampler comparisons.

  • Normalization-constant benchmarks: Long-run MCMC comparisons vary annealing and leapfrog depth across 10, 20, 30, 50, and 100 layers.The table reports how AFT and SMC performance changes as the number of steps increases.
  • Normalization-constant benchmarks: For the mode-separated Gaussian mixture, nine grid centers use variance 0.3, with 2000 samples estimating the log normalization constant.A standard MLP parameterizes the control drift, with Fourier features for time.
  • Normalization-constant benchmarks: The Funnel benchmark uses 6000 particles and compares HMC-based settings with a slice-sampler configuration using 1000 slice-sampling steps per temperature.The HMC comparison keeps the number of integral steps aligned across methods.
  • Normalization-constant benchmarks: The Cox-process experiment uses 2000 samples and finds that gradient clipping stabilizes and accelerates PIS-Grad training.The experimental setup follows prior work and uses repeated SMC estimates as ground-truth normalization constants.
  • Normalization-constant benchmarks: With no annealing and only 100 steps, PIS is reported as competitive against AFT and SMC with slice-sampler kernels.This comparison is summarized in Table 7.
  • Normalization-constant benchmarks: The paper also evaluates PIS with a large number of integral steps and compares it against an AFT configuration with 100 annealing layers and 100 leapfrog steps per layer.The total-step suffix identifies the integral-step count.
  • Molecular benchmark: The molecular benchmark uses coordinate transformations and dataset normalization before evaluating unweighted samples under the comparison setup.The target distribution is constructed from alanine-dipeptide simulations.
  • Image benchmark: The image experiment trains a vanilla VAE on binary MNIST with a 50-dimensional diagonal Normal latent distribution and reconstructs images with a multilayer perceptron decoder.The reported reconstruction comparison notes smoother reconstructions than the original images.

G.1 TIPS

PIS training and sampling quality depend strongly on the diffusion horizon, discretization, and controller optimization. The reported failure cases are empirical and may not generalize across target densities.

  • Scope: The authors characterize these observations as experimental failure cases whose correctness and generality across target densities are uncertain.
  • Training stability: Small T can produce large control Lipschitz constants, making neural-network approximation and training more difficult.Increasing T or constraining the control magnitude is recommended when numerical instability occurs.
  • Discretization and weighting: 12?
  • Optimization choices: Training with Tweedie’s formula yields more stable and smaller training loss without obvious improvements in sampler testing or normalization-constant estimation.
  • Failure modes: PIS can miss modes when initialization is improper and T is extremely small.The stated factors include an initial Gaussian far from target modes, difficult policies with large Lipschitz constants, and limited policy representation under KL training.
  • Training stability: Large T improves training stability and mode coverage, but large δt causes discretization error that degrades sample quality.Increasing the number of integration steps can reduce this degradation.
Loading 2111.15141v2…