Path Integral Diffusion

Path Integral Diffusion (PID) is a linearly-solvable formulation of the Schrödinger bridge problem, in which the optimal drift of a controlled diffusion transporting a prescribed initial distribution to a prescribed terminal distribution is expressed as an explicit integral over forward and backward Green functions of two linear, mutually adjoint partial differential equations Behjoo & Chertkov (2025). PID generalises the path-integral approach to stochastic optimal control of Kappen (2005), extended to an arbitrary gauge and a density–current calculus of variations by Chernyak, Chertkov, Bierkens & Kappen (2014), and has a discrete-state counterpart in the linearly-solvable Markov decision problems of Todorov (2007). Within PID the harmonic path integral diffusion algorithm is the exactly closed-form, quadratic-potential specialisation, and the mean-field extension of PID recovers a bridge formulation of score-based generative modelling Song et al. (2021).

The stochastic optimal control setup

PID formalises generative sampling from a target density \(p^{\mathrm{tar}}\) on \(\mathbb{R}^d\) as a constrained stochastic optimal control problem. Given a reference Itô diffusion with drift \(f\) and unit diffusion coefficient, one looks for an additional control \(u_t(x)\) such that

\[ dx_t = \bigl(f_t(x_t) + u_t(x_t)\bigr)\,dt + dW_t,\qquad x_0 = x^{\mathrm{init}},\qquad \mathrm{Law}(x_1) = p^{\mathrm{tar}}, \]

and the expected cost

\[ C[u] = \mathbb{E}\!\left[\int_0^1 \Bigl(\tfrac{1}{2}\|u_t\|^2 + V_t(x_t)\Bigr) dt\right] \]

is minimised subject to the terminal-law constraint Kappen (2005); Behjoo & Chertkov (2025). The state potential \(V_t\) shapes the path distribution on the interior of \([0,1]\), while the fixed endpoints make the problem a Schrödinger bridge rather than an unconstrained optimal control problem. The Hamilton–Jacobi–Bellman (HJB) equation for the optimal cost-to-go \(J_t(x)\) is

\[ -\partial_t J_t = -\tfrac{1}{2}\|\nabla J_t\|^2 + f_t \cdot \nabla J_t + \tfrac{1}{2}\Delta J_t + V_t, \]

with terminal condition \(J_1(x) = +\infty \cdot \mathbb{1}[\,\cdot\,]\) encoding the hard marginal constraint Kappen (2005).

Hopf–Cole linearisation

The HJB equation is quadratic in \(\nabla J_t\) and therefore nonlinear. The Hopf–Cole substitution \(J_t = -\log \psi_t\) converts it into a linear backward equation

\[ -\partial_t \psi_t = f_t \cdot \nabla \psi_t + \tfrac{1}{2}\Delta \psi_t - V_t\,\psi_t, \]

i.e. an imaginary-time Schrödinger equation for \(\psi_t\) Kappen (2005). The optimal drift is then

\[ u^*_t(x) = \nabla_x \log \psi_t(x), \]

and \(\psi_t\) admits a Feynman–Kac path-integral representation integrating the cost \(V_t\) along sample paths of the reference diffusion Kappen (2005); this representation is what gives the framework its name Behjoo & Chertkov (2025). A companion forward equation for a conjugate density \(\varphi_t\) arises in the density–current formulation Chernyak, Chertkov, Bierkens & Kappen (2014), and the controlled marginal density satisfies \(p_t(x) = \varphi_t(x)\,\psi_t(x)\). The pair \((\psi_t,\varphi_t)\) thus plays the role of a backward and a forward Green function, and the drift admits the explicit integral representation

\[ u^*_t(x) = \nabla_x \log \int dy\; p^{\mathrm{tar}}(y)\,G_-(t,x;1,y)\big/G_+(1,y;0,x^{\mathrm{init}}), \]

where \(G_-, G_+\) are the backward and forward heat kernels of the linearised operator Behjoo & Chertkov (2025).

Linearly-solvable Markov decision problems

The discrete-state, discrete-time analogue of PID is the class of linearly-solvable Markov decision problems of Todorov (2007). Here the control cost penalises the Kullback–Leibler divergence of the controlled transition kernel from an uncontrolled passive dynamics, and the Bellman equation for the exponentiated cost-to-go \(z = \exp(-J)\) becomes linear in \(z\). The optimal policy is recovered by a single matrix–vector product against a Perron–Frobenius-type operator, without iterative dynamic programming Todorov (2007). PID is the continuous-state, continuous-time instance of the same phenomenon: the nonlinear HJB equation linearises under Hopf–Cole, and the nonlinear Bellman equation linearises under exponentiation Todorov (2007); Kappen (2005).

Gauge invariance and the density–current calculus

Chernyak, Chertkov, Bierkens & Kappen (2014) recast the path-integral control problem as a calculus of variations over a pair of conjugate fields — a density \(\rho_t(x)\) and a current \(j_t(x)\) — subject to the continuity equation. In this formulation an additional vector potential \(A_t(x)\) enters the linearised equations through the minimal-coupling prescription \(\nabla \mapsto \nabla - A_t\), which enlarges PID beyond the scalar-potential setting of Kappen (2005). Gauge transformations of \(A_t\) and \(V_t\) that leave observable currents invariant are analogues of electromagnetic gauge symmetry for stochastic control, and the linear structure of the Hopf–Cole picture is preserved under them Chernyak, Chertkov, Bierkens & Kappen (2014). This gauge extension is what allows PID to cover, within a single linear framework, the general Schrödinger bridge with non-reversible reference dynamics.

Three levels of integrability

Behjoo & Chertkov (2025) organise the tractability of the PID family into three nested levels:

  1. General bounded potentials, general forcing, general gauge. The Hopf–Cole reduction gives two linear conjugate PDEs for \(\psi_t, \varphi_t\); the optimal drift is a formal Green-function integral that can be evaluated by numerical path integration or operator-splitting.
  2. Quadratic, positive-definite potentials. The linearised operator is the Hamiltonian of a multi-dimensional quantum harmonic oscillator; the Green functions are Gaussian with hyperbolic time-dependence, and \(u^*_t\) becomes an affine map in \(x\). This is the harmonic path integral diffusion regime.
  3. Uniform quadratic potential with zero forcing and zero gauge. The drift further reduces to convolutions of the target with Gaussian kernels, yielding the closed-form sampler that can be run on a standard CPU Behjoo & Chertkov (2025).

The mean-field extension of PID, in which a Vlasov-type coupling replaces the external potential \(V_t\) by its convolution with the controlled density, preserves the linear structure of the quadratic case and is treated in a companion article on mean-field PID Chertkov (2025); Chertkov & Behjoo (2025).

Relation to score-based generative models

In score-based generative modelling Song et al. (2021) a reverse-time denoising SDE is driven by the gradient of the log-density \(\nabla_x \log p_t\), the so-called score, which is approximated by a neural network trained on samples from the target. PID realises the same generative object — a controlled diffusion whose terminal law is \(p^{\mathrm{tar}}\) — but with the drift \(u^*_t = \nabla_x \log \psi_t\) given by an integral over Green functions rather than by a trained neural network Behjoo & Chertkov (2025). In the linearly-solvable regimes above, \(u^*_t\) is available in closed form, which turns PID into a ground-truth drift against which learned scores can be benchmarked. Conversely, any score-based diffusion whose forward process is driven by a quadratic potential and whose target is a Gaussian mixture admits a PID representation, so score-based modelling and PID can be regarded as two computational implementations of the same underlying stochastic optimal transport problem Kappen (2005); Behjoo & Chertkov (2025).

See also

References