Bridging Control, Inference, Transport, and Thermodynamics

Abstract

The last decade has seen the development of powerful methods for learning complex structure from high-dimensional data. These advances have brought to the foreground fundamental connections between subdisciplines of physics, applied mathematics, and machine learning. In this review, we bring together some of these ideas, often expressed in different languages, to highlight a conceptual thread that links five distinct fields: control theory, optimal transport, probabilistic inference, non-equilibrium thermodynamics, and machine learning. A common theme is the optimization of free-energy-like functionals under dynamical or statistical constraints. We offer a guided tour through this thread and present selected applications in reinforcement learning, variational inference, and generative modeling. The review does not assume prior familiarity with these topics, and begins with principles originating from physics.

INTRODUCTION

Recent decades have seen the development of remarkably powerful tools for learning complex structure from data. These methods have brought to the foreground certain key connections between subdisciplines of physics and applied mathematics. The conceptual ideas have appeared in different fields at different times, sometimes rediscovered and often expressed in different languages. In this review, we bring together some of these connections to highlight a fascinating conceptual thread that links five distinct fields: control, transport, inference, thermodynamics and learning. The visual table of contents, Fig. 1, lays out the topics we cover and the links between them. These connections have proved fruitful for the development of powerful practical algorithms for generative modeling and control. We offer a guided tour through this thread and present selected applications in reinforcement learning, variational inference, and generative modeling. The review is written in the language of physicists and does not assume prior familiarity with these topics. Ideas from mechanics and thermodynamics are used as familiar anchors to introduce concepts in other areas.

FIG. 1.

Graphical table of contents: topics of the review and the connections between them.

The conceptual thread begins at the level of paths (left column in Fig. 1) in Variational structure of mechanics, control, and inference. This chapter develops a common variational language that is reused throughout the review. We introduce optimal control, which seeks the best way to steer a dynamical system—for example, guiding a rocket to a target while minimizing fuel consumption. We illustrate two complementary structural approaches for solving control problems. Introducing Lagrange multipliers for imposing a constraint on the dynamics leads to Hamiltonian equations. Another technique, known as dynamic programming, reformulates the optimization as a recursive equation. This reformulation leads to the Hamilton–Jacobi equation and, for stochastic control, its more general Hamilton–Jacobi–Bellman (HJB) counterpart. Controls and rewards play roles analogous to velocities and the negative Lagrangian in classical mechanics. We then pass from optimizing individual paths to distributions (right column in Fig. 1). The Maximum Entropy Principle (MaxEnt) selects a probability distribution that remains as close as possible to a reference distribution, as measured by the Kullback–Leibler (KL) divergence, while remaining consistent with data. We then generalize these tools for probabilistic inference to distributions over paths. For controlled driftdiffusive processes, Girsanov’s theorem identifies this pathspace KL divergence with a quadratic control cost, thereby connecting inference and control.

The next chapter, Control and transport of densities, illustrates how the control of paths can be used to obtain a geometric description for the evolution of probability densities. The topic of transport brings together ideas from the study of (stochastic) paths and of distributions (center column in Fig. 1). The first step is to impose constraints on both the initial and final densities of a stochastic control problem. These constraints lead to the Schrödinger bridge problem, which finds the most likely ensemble of stochastic paths that connects the two densities. If one retains the endpoint constraints but describes only how initial and final positions are paired, the problem becomes entropic optimal transport. In general, optimal transport asks how to move one density into another at minimum cost, much like reshaping a pile of earth while minimizing the cost of moving it. The entropic version adds a term that favors spreading probability among different pairings. For a quadratic cost, the zero-noise limit defines the 2-Wasserstein distance, which has special mathematical properties. This distance equips the space of probability densities with a geometry and leads to Wasserstein gradient flows: steepest-descent dynamics in which the evolving object is a probability distribution.

The chapter on Thermodynamics grounds this geometry in physics. At equilibrium, statistical mechanics assigns probabilities to microscopic states while retaining only a small number of macroscopic constraints. Maximum-entropy inference reproduces the familiar Gibbs distributions from constraints on mean energy, making a heuristic connection between inference and equilibrium thermodynamics. Away from equilibrium, control and transport describe finite-time transformations between distributions. Onsager’s least-dissipation principle generates phenomenological field theories whose dissipation rate is bounded by the Wasserstein distance. For stochastic dynamics farther from equilibrium, the Wasserstein transport cost provides thermodynamic speed limits on the entropy produced during a finite-time transformation. Nonequilibrium work fluctuation relations provide a complementary link between fluctuations, free-energy differences, and the energetic cost of driving. Thermodynamics therefore supplies physical meaning for the variational structures introduced earlier, which in turn furnish quantitative limits on thermodynamic processes.

The fourth conceptual chapter, Sampling and control, presents sampling as a general approach to probabilistic inference, optimization, and control. Monte Carlo methods estimate averages from samples, while simulated annealing searches for low-energy states. Exploring possible trajectories and weighting them according to their rewards also links sampling to control and reinforcement learning; the Feynman–Kac formula makes this connection precise for a class of stochastic control problems. Sampling faces several distinct challenges, including high dimensionality, separated modes that impede efficient exploration, and observables whose averages are dominated by rare trajectories. When an unnormalized target density is available, Annealed Importance Sampling (AIS) replaces direct sampling by a sequence of intermediate distributions and corrects the resulting nonequilibrium process through trajectory-dependent weights. The control perspective shows how guiding the dynamics toward important trajectories can reduce fluctuations. Although this ideal control is generally unknown, it provides a target for practical algorithms that learn progressively better sampling dynamics. This chapter thus closes the conceptual thread by turning efficient sampling into a control problem.

These connections provide a framework and common language for Applications (bottom row in Fig. 1). The first application is Reinforcement learning. Stochastic optimal control assumes that the dynamics of the environment are known. Reinforcement learning instead seeks good controls—called actions—by interacting with an environment and observing the resulting states and rewards. A policy gives the probability of taking each action in a given state and therefore defines an ensemble of trajectories. The connection between inference and control leads to Maximum Entropy Reinforcement Learning (MaxEnt RL), in which high-reward trajectories are favored while the policy retains some randomness. Applying dynamic programming to this objective leads to soft Q-learning and Soft Actor-Critic (SAC), practical algorithms for learning stochastically controlled paths when the dynamics are not known in advance.

The second application concerns Wasserstein gradient flows. The geometry of probability distributions provides a variational method for integrating Fokker–Planck equations. The same viewpoint helps describe large systems of interacting particles. In a suitable mean-field limit, the individual particles are replaced by an evolving density, which can obey a Wasserstein gradient flow. This description has also been used for wide neural networks, where many parameters are treated as interacting particles. Finally, in variational inference, a difficult target distribution is approximated within a simpler family of distributions. The KL objective has the form of a free energy, and projecting its Wasserstein gradient flow onto that family yields a practical approximation scheme.

The final application, Generative modeling with flows and diffusion, treats targets represented by data rather than by a known density. Normalizing flows and diffusion models learn deterministic maps or stochastic dynamics that transform samples from a simple base distribution into samples from the target distribution. In this sense, generative modeling is a learned transport problem. Flow matching learns a velocity field that carries the base density toward the data, while diffusion models learn to reverse a gradual noising process. These methods need not find a minimum-cost transport; their goal is to learn a transformation that can be estimated from finite data and used efficiently for sampling. The earlier connections explain why ideas from transport, inference, control, and nonequilibrium physics repeatedly appear in their construction, and why additional guidance can be understood as a control that favors selected outcomes.

A brief note on notation.

We have sought to unify notation across the chapters to the extent possible. The notation tables at the end of this review are organized by the chapter in which a symbol is first introduced.

I. VARIATIONAL STRUCTURE OF MECHANICS, CONTROL AND INFERENCE

In this chapter, we develop a common variational language for mechanics, control, and inference. We introduce two complementary approaches for solving control problems. We then turn to inferring probability distributions from observations, and show how inference over stochastic paths leads back to a particular class of control problems.

A. Finding optimal paths in classical mechanics and control theory

We begin with a gentle introduction to control theory. Consider an illustrative example — suppose we would like to guide a rocket towards a target in a certain amount of time by controlling its thrust. A large thrust accelerates the rocket but consumes more fuel, whereas insufficient thrust may not be enough to reach the target on time. How does one choose a thrust protocol so that the rocket reaches the target while consuming the least amount of fuel? Different control strategies generate different paths. In optimal control, we assume that there is a reward function that assigns a value or utility to each path. This reward may favor reaching a target, penalize control costs, or balance competing goals. Control theory offers mathematical tools for finding a control strategy that maximizes the total reward along the path. Although its variational roots are much older, modern optimal control took shape in the 1950s. Bellman’s dynamic programming and the maximum principle of Pontryagin offer two powerful and complementary approaches for solving control problems [1–3]. We now describe these two solution principles, which will reoccur throughout the rest of the review. The principles are best illustrated by demonstrating their relationship with standard formulations of classical mechanics.

The optimal control problem.

In optimal control, a system is described by a (possibly multi-dimensional) state $x_t$ whose evolution is modulated by a control $u_t$. For the rocket, the state includes its position, velocity, and remaining fuel, while the control is the thrust. For convenience, we rescale time to the interval $0\le t\le 1$.

In stochastic optimal control, the state obeys¹

$${\dot{x}}_t=g_t\left(x_t,u_t\right)+\sqrt{\epsilon }{\eta }_t,$$ (1)

where $\epsilon$ sets the noise scale and ${\eta }_t$ is standard Gaussian white noise, with $⟨{\eta }_{t,i}{\eta }_{s,j}⟩={\delta }_{ij}\delta (t-s)$. A reward $r_t\left(x_t,u_t\right)$ assigns a value per unit time to each state-control pair, while a terminal reward $r_1\left(x_1\right)$ assigns value to the final state. Starting from a fixed $x_0$, the optimized objective is

$$V_0\left(x_0\right)=\underset{u_{0\le s<1}}{\text{max}}⟨r_1\left(x_1\right)+{\int }_0^1{r}_s\left(x_s,u_s\right)ds\mid x_0⟩,$$ (2)

where the expectation is taken over the noise histories $\left\{{\eta }_t\right\}$ conditioned on the initial state $x_0$ [4–6]. The controls may depend on the observed state, but not on future noise. A fuel cost, for example, enters with a minus sign in $r_t$, while a terminal reward favors arrival near the target.

Perhaps intuitively, classical mechanics can be described as a special case of this optimization problem. Consider a particle with coordinates $x_t$ and velocity ${\dot{x}}_t$ whose motion is described by a Lagrangian $L_t(x,\dot{x})$. Its action is

$$S\left[x\left(\cdot \right)\right]={\int }_0^1{L}_t\left(x_t,{\dot{x}}_t\right)dt,$$ (3)

where $x(\cdot )$ denotes the full path. The principle of stationary action states that the physical path $\left\{x_t^{*}\right\}$ makes $S$ stationary under variations that leave the endpoints $x_0$ and $x_1$ fixed. Imposing stationarity leads to the Euler–Lagrange equations. We now map action minimization onto a control problem by choosing

$$g_t(x,u)=u,\epsilon =0,$$ (4)

$$r_t(x,u)=-L_t(x,u),$$ (5)

and imposing that the terminal endpoint is $x_1$. Here the control plays the role of velocity, $u_t={\dot{x}}_t$. The endpoint constraint can equivalently be imposed through a terminal reward that is zero at $x_1$ and $-\infty$ elsewhere. Maximizing the reward then minimizes the action. Note that classical mechanics only requires stationarity, that is, physical trajectories need not be minima and can instead be saddles. The optimization therefore selects action-minimizing paths, while the variational equations also describe other stationary paths.

There are two structural approaches for solving the control problem. The first treats the dynamics in Eq. (1) as a constraint and introduces Lagrange multipliers, leading to a Hamiltonian formulation. The second recursively expresses the optimal solution in terms of optimal solutions to shorter remaining problems, leading to the Hamilton–Jacobi–Bellman equation. Relating the two will identify the Lagrange multiplier as the gradient of a value function.

Lagrange multipliers and Hamiltonian dynamics.

A Lagrange multiplier converts a constrained optimization problem into a stationary problem in an enlarged set of variables. Variation with respect to the multiplier enforces the constraint, while variation with respect to the original variables supplies the optimality equations. To see this, let’s first take the deterministic limit $\epsilon =0$. Introducing a time-dependent multiplier $p_t$ for the constraint ${\dot{x}}_t=g_t\left(x_t,u_t\right)$ gives the augmented objective

$$\tilde{V}(x,u,p)=r_1\left(x_1\right)+{\int }_0^1\left[r_t\left(x_t,u_t\right)\right.\left.+p_t\cdot \left(g_t\left(x_t,u_t\right)-{\dot{x}}_t\right)\right]dt,$$ (6)

with the initial endpoint $x_0$ held fixed. The sign of $p_t$ is conventional. Define the control Hamiltonian and its optimized form by

$${\mathscr{H}}_t\left(x,p,u\right)=p\cdot g_t\left(x,u\right)+r_t\left(x,u\right),$$ (7)

$$H_t(x,p)=\underset{u}{\text{max}}{\mathscr{H}}_t(x,p,u).$$ (8)

Variation with respect to $p_t$ and $x_t$ (after integrating by parts the term involving ${\dot{x}}_t$ in Eq. (6)), together with maximization over $u$, gives the necessary optimality conditions [2, 4]. When $H_t$ is differentiable, these take the form

$${\dot{x}}_t={\nabla }_p{H}_t\left(x_t,p_t\right),$$ (9)

$${\dot{p}}_t=-{\nabla }_x{H}_t\left(x_t,p_t\right),$$ (10)

$$u_t^{*}=\text{arg}{\text{max}}_u{\mathscr{H}}_t\left(x_t,p_t,u\right).$$ (11)

If the terminal state is free, the endpoint variation gives $p_1=\nabla r_1\left(x_1\right)$, whereas if it is fixed, the terminal variation vanishes. For a free terminal state, the multiplier equation is thus evolved backward from the terminal condition along the trajectory. We will see below that $p_t$ measures how the optimal future reward changes under a displacement of the state.

For the specialization $g_t(x,u)=u$ and $r_t(x,u)=-L_t(x,u)$, the optimized Hamiltonian becomes

$$H_t(x,p)=\underset{u}{\text{max}}\left\{p\cdot u-L_t(x,u)\right\},$$ (12)

which is the Legendre transform of the Lagrangian. The stationarity condition with respect to $u$ gives

$$p_t={\nabla }_u{L}_t\left(x_t,u_t\right)={\nabla }_{\dot{x}}{L}_t\left(x_t,{\dot{x}}_t\right),$$ (13)

so the multiplier is the canonical momentum, and the remaining conditions are Hamilton’s equations.

Dynamic programming and the Hamilton–Jacobi equation.

The second approach involves keeping track of the maximum reward that remains available from every possible intermediate state (‘reward-to-go’). For the rocket, suppose we already knew the maximum total future reward attainable from every state $x_{t+dt}$ at time $t+dt$. The value of choosing a particular thrust at time $t$ is its immediate reward, including the fuel cost, plus the expected best future reward from the state it produces. The optimal thrust is the one that maximizes this sum. In other words, if the expected best future reward is known, solving for the optimal thrust involves optimizing over a single step rather than over a full trajectory.

This method for solving a problem by breaking it into shorter subproblems and reusing their solutions is called dynamic programming [1]. In this setting, the key observation is that the control at the next state for an optimal strategy must itself be optimal. This gives a recursive relation between the best reward now and the best reward at the next time step, which we can solve backward from the terminal reward. The reuse of partial results is analogous to the transfer-matrix method in statistical mechanics, where a partial sum over configurations is updated one site at a time [7].

The central object in dynamic programming is the value function

$$V_t(x)=\underset{u_{t\le s<1}}{\text{max}}⟨{\int }_t^1{r}_s\left(x_s,u_s\right)ds+r_1\left(x_1\right)\mid x_t=x⟩,$$ (14)

with terminal condition $V_1(x)=r_1(x)$. Thus, $V_t(x)$ is the largest expected total reward that can still be collected from time $t$ onward, starting at $x$, including the terminal reward. It does not include rewards already received. This is the optimal reward-to-go, and $V_0\left(x_0\right)$ is the objective in Eq. (2).

The controls after time $t+dt$ have already been optimized inside $V_{t+dt}$, so only the control applied during the first short interval remains to be chosen. Splitting the total reward into the immediate and future contributions gives Bellman’s recursion, to first order in $dt$,

$$\begin{gathered} V_t(x)=\underset{u}{\text{max}}\left\{r_t(x,u)dt\right. \\ \left.+⟨V_{t+dt}\left(x+g_t(x,u)dt+\sqrt{\epsilon }d{W}_t\right)⟩\right\} \\ +o(dt). \end{gathered}$$ (15)

Here $d{W}_t$ is a Gaussian noise increment with zero mean and $⟨d{W}_{t,i}d{W}_{t,j}⟩={\delta }_{ij}dt$, and the expectation is over this increment. Assuming the value function is smooth, expanding $V_{t+dt}$ to first order in time and second order in the state increment, averaging over the noise, and retaining terms of order $dt$ gives the Hamilton–Jacobi–Bellman (HJB) equation [5, 6]

$${\partial }_t{V}_t(x)+\frac{\epsilon }{2}{\nabla }^2{V}_t(x)+\underset{u}{\text{max}}\left\{g_t(x,u)\cdot \nabla V_t(x)+r_t(x,u)\right\}=0.$$ (16)

The control chosen at state $x$ is the maximizer of the local expression,

$$u_t^{*}(x)=\text{arg}{\text{max}}_u\left\{g_t(x,u)\cdot \nabla V_t(x)+r_t(x,u)\right\}.$$ (17)

Thus, when $g_t$ and $r_t$ are known, the difficult optimization over an entire control history can be replaced by a backward equation for $V_t$ followed by a local maximization at each state and time. Finding this function over a high-dimensional state space can nevertheless be difficult.

In the deterministic limit, the Hamiltonian $H_t(x,p)$ introduced above allows us to write the HJB equation as

$${\partial }_t{V}_t(x)+H_t\left(x,\nabla V_t(x)\right)=0.$$ (18)

This equation also makes the connection between dynamic programming and the Lagrange-multiplier formulation explicit. Where $V_t$ and $H_t$ are smooth, set $p_t=\nabla V_t\left(x_t\right)$ along an optimal trajectory. The state evolves as

$${\dot{x}}_t=g_t\left(x_t,u_t^{*}\right)={\nabla }_p{H}_t\left(x_t,p_t\right).$$ (19)

The total time derivative of $p_t$ along this trajectory is

$${\dot{p}}_t={\partial }_t\nabla V_t\left(x_t\right)+\left({\dot{x}}_t\cdot \nabla \right)\nabla V_t\left(x_t\right).$$ (20)

Taking the spatial gradient of Eq. (18) gives

$${\partial }_t\nabla V_t+{\nabla }_x{H}_t+\left({\nabla }_p{H}_t\cdot \nabla \right)\nabla V_t=0,$$ (21)

where ${\nabla }_x{H}_t$ denotes differentiation with respect to the explicit state argument of $H_t(x,p)$ while holding $p$ fixed, and the Hamiltonian derivatives are evaluated at $p=\nabla V_t(x)$. Using ${\dot{x}}_t={\nabla }_p{H}_t$ and combining the above two equations, we get

$${\dot{p}}_t=-{\nabla }_x{H}_t\left(x_t,p_t\right).$$ (22)

Thus, the trajectories along which the Hamilton–Jacobi equation is solved, its characteristics, obey precisely the Hamilton equations obtained from the Lagrange-multiplier approach. The multiplier $p_t$ is the gradient of the value function because it measures the sensitivity of the optimal future reward to a displacement of the current state.

For action-minimizing trajectories in classical mechanics, the value function is the negative of the minimum action remaining to the fixed endpoint $x_1$,

$$V_t\left(x\right)=-\underset{\left\{x_s\right\}:x_t=x,}{\text{min}}{\int }_t^1{L}_s\left(x_s,{\dot{x}}_s\right)ds.$$ (23)

On a minimizing trajectory, variation of the initial endpoint gives $\delta V_t=p_t\cdot \delta x$, while variation with respect to the velocity gives $p_t={\nabla }_{\dot{x}}{L}_t$. Hence

$$p_t=\nabla V_t\left(x_t\right)={\nabla }_{\dot{x}}{L}_t,$$ (24)

so the multiplier of optimal control becomes the canonical momentum of mechanics.

The two approaches thus give complementary descriptions: the Lagrangian viewpoint follows individual trajectories, while the Eulerian viewpoint of dynamic programming finds a field $V_t(x)$ over all states. Note that we have assumed that the reward $r_t$ and dynamics $g_t$ are known. Reinforcement learning (RL) addresses the challenging problem of finding the optimal control when this information is not available in advance. In RL, one interacts with the environment to learn which controls, or actions, produce large cumulative reward from observed state transitions and rewards. The value function remains central because it provides an estimate of the long-term consequence of acting from each state. We return to the RL framework in Sec. V A of the Applications chapter.

B. Variational basis for probabilistic inference and control

We now turn to a different but related problem, wherein measurements constrain probabilities of random degrees of freedom, instead of control functions that steer individual paths. Since measurements rarely provide every microscopic detail of a system, we describe the system by a probability distribution that we seek to infer from the available observations. Bayesian inference is an example of such a capability: what is the distribution of $x$, given observations $A$ and a prior $p^0(x)$ (which encodes existing knowledge about the system [8])? By Bayes’ rule, the conditional distribution is $p(x\mid A)=p(A\mid x)p^0(x)/p^0(A)$, thereby recasting the problem as specifying a model for $A$ given $x$. More generally, the central task of probabilistic inference is to determine conditional probabilities. In this chapter, we show that this task can be formulated as constrained optimization—a variational problem under constraints arising from observations. This formulation highlights a conceptual connection to the optimization problems of mechanics and control (Sec. IA), and provides the theoretical basis for path-integral stochastic control in Sec. ID. Later in the review, variational probabilistic inference will serve as a foundation for practical inference algorithms: Chapter IV uses control theory to address inference problems, while Sec. VB 3 considers variational inference, where the system distribution is approximated using a parametrized ansatz.

Here, we will introduce maximum-entropy inference, a method for finding a probability distribution for the system without prior knowledge, constrained to be consistent with observations. The observations take the form of measurements of moments like the mean or variance, which, however, do not fully determine the distribution. There is thus a many-to-one correspondence between microscopic models and the measurements. How to choose one model out of the many? One well-motivated approach is to construct a model of behavior $p^A(x)$ that is closest to a reference distribution $p^0(x)$, but remains consistent with the observations $A$. As we will show below, a natural candidate for measuring closeness is the Kullback-Leibler divergence,

$${\text{D}}_{\text{KL}}(p\Vert q)=\int \text{d}xp(x)\text{log}\frac{p(x)}{q(x)}.$$ (25)

measuring how different (“the excess surprise”) the distribution $p$ is from $q$ [9]. We will furthermore see that using Lagrange multipliers that enforce consistency with observed paths has a natural interpretation as rewards. This choice will establish a conceptually and practically relevant connection between probabilistic inference and control.

1. Sanov’s theorem and the inference method

We first make a plausibility argument for the inference method through Sanov’s theorem. It equips ${\text{D}}_{\text{KL}}$ with a probabilistic interpretation as a measure of how a distribution estimated from samples deviates from the underlying true distribution. Namely, suppose that $N$ independent and identically distributed (IID) samples ${\left\{x_n\right\}}_{1,\dots ,N}$ are drawn from a base measure, $p^0(x)$. Having drawn these samples, to what extent does the empirical measure,

$$\hat{p}(x)=\frac{1}{N}\sum _{n=1}^N\delta \left(x-x_n\right),$$ (26)

depart from $p^0(x)?$ Sanov’s theorem [10, 11] states that the probability of the empirical measure, $\mathbb{P}\left[\hat{p}\right]$, deviating from the ground truth $p^0(x)$ is exponentially suppressed with an increasing number of samples $N$:

$$\mathbb{P}\left[\hat{p}\right]~e^{-N{\text{D}}_{\text{KL}}\left(\hat{p}\Vert p^0\right)},$$ (27)

Eq. (27) is achieved using large-deviations theory, which concerns finding the asymptotic decline in probability of a rare event, wherein the probability typically takes the form $~e^{-NI}$ with $I$ called the rate function. The rate function here is the ${\text{D}}_{\text{KL}}$, which possesses two appealing properties: ${\text{D}}_{\text{KL}}=0$ and minimal if and only if $\hat{p}=p^0$, and the spread itself gets smaller as $N\to \infty :\hat{p}-p^0~N^{-1/2}$. The rarer the event, the larger the KL divergence, and the more the probability is suppressed.

We will prove Eq. (27) via a path-integral calculation, exemplifying the path-integral machinery we will encounter later in the review. We compute the distribution of empirical measures $\mathbb{P}\left[\hat{p}\right]$ by definition, from the expectation of the functional delta function over the drawn samples,

$$\mathbb{P}\left[\hat{p}\right]=\prod _{m=1}^N\left[\int \text{d}{x}_m{p}^0\left(x_m\right)\right]\times \delta \left[\hat{p}(x)-\frac{1}{N}\sum _{n=1}^N\delta \left(x-x_n\right)\right].$$ (28)

Expressing the delta function via its Fourier transform, $\delta [f(x)]\propto \int \text{D}[\text{i}\mu (x)]e^{\int \text{d}x\mu (x)f(x)}$, and noting the identical and independent integrals over samples, we rewrite

$$\mathbb{P}\left[\hat{p}\right]\propto \int \text{D}[\text{i}\mu (x)]\text{exp}\left\{N\int \text{d}x\frac{\mu (x)}{N}\hat{p}(x)\right.\left.+N\text{log}\left[\int \text{d}{x}^{\prime }{p}^0\left(x^{\prime }\right)e^{-\mu \left(x^{\prime }\right)/N}\right]\right\}$$ (29)

Since $N\to \infty$, we note that only $\mu (x)~N$ contributes to the Fourier integral, overall implying that the argument of the exponential is $~N$. This justifies a saddle-point evaluation of the integral — the saddle point ${\mu }^{*}(x)$ satisfies

$$\frac{p^0(x)e^{-{\mu }^{*}(x)/N}}{\int \text{d}{x}^{\prime }{p}^0\left(x^{\prime }\right)e^{-{\mu }^{*}\left(x^{\prime }\right)/N}}=\hat{p}(x),$$ (30)

Inserting Eq. (30) in Eq. (29), we obtain Eq. (27).

Sanov’s theorem ensures that the deviations of the estimated measure from the ground truth are minimized with more samples. However, during probabilistic inference, the choice of a reference distribution is often that which can be sampled easily, meaning that there is no inherent reason why the underlying model is that of the base $p^0$. This suggests that if we observe data inconsistent with $p^0$ and seek the most likely alternative explanation, we should find the distribution $p^A$ that is closest to $p^0$ in KL divergence while remaining consistent with the observations $A$. By observations we mean, e.g., measurable moments or state occupancies of the system, dictated by the unknown underlying distribution. For instance, later we will be interested in inferring and controlling the distributions over paths. Via a large-deviation principle, we may motivate the fact that the distribution $p^A$ closest to $p^0$, consistent with observations $A$, satisfies

$$\begin{gathered} p^A=\text{arg}\underset{p}{\text{min}}{\text{D}}_{\text{KL}}\left(p\Vert p^0\right), \\ \text{s.t.}\int \text{d}xp(x)=1\text{and}\int \text{d}xp(x)a(x)=A. \end{gathered}$$ (31)

This will serve as the basis for probabilistic inference throughout this review.

To motivate this inference, we recall that not all $\hat{p}(x)\text{s}$ are possible if the value of a certain observable, $A$, is known and thus imposed as a constraint. The space of empirical distributions is thus reduced to those which satisfy $A=N^{-1}\sum _{n=1}^{N}a\left(x_n\right)$. Therefore, Eq. (28) should be rewritten instead as

$$\begin{gathered} \mathbb{P}[\hat{p}\mid A]=\prod _{m=1}^N\left[\int \text{d}{x}_m{p}^0\left(x_m\right)\right] \\ \times \delta \left[\hat{p}(x)-\frac{1}{N}\sum _{n=1}^N\delta \left(x-x_n\right)\right] \\ \times \delta \left[A-\frac{1}{N}\sum _{n=1}^{N}a\left(x_n\right)\right]. \end{gathered}$$ (32)

Even more so, note that Eq. (30) implies additionally that $\hat{p}(x)$ must be normalized ($\int \text{d}x\hat{p}(x)=1$). The rest of the procedure to obtain Eq. (27) remains unchanged other than the additionally imposed observables. Following the above saddle-point calculation, Eq. (27) can be rewritten as

$$\begin{gathered} \mathbb{P}[\hat{p}\mid A]~\text{exp}\left\{-N{\text{D}}_{\text{KL}}\left(\hat{p}\Vert p^0\right)\right. \\ +{\eta }^{*}\left[\int \text{d}x\hat{p}(x)-1\right] \\ \left.+{\lambda }^{*}\left[\int \text{d}x\hat{p}(x)a(x)-A\right]\right\}, \end{gathered}$$ (33)

where ${\eta }^{*}$ and ${\lambda }^{*}$ are such that the overall rate function (the argument of the exponential) is minimized. (This calculation should be regarded as qualitative, as the saddle point in $\eta$ and $\lambda$ may often not be justified.) Therefore, we conclude that Sanov’s theorem motivates turning the inference problem of $p^A(x)$ into a constrained optimization problem, Eq. (31).

2. Maximum-entropy inference

Sanov’s theorem is closely related to maximum-entropy inference. Suppose, for simplicity, that the prior is the uniform distribution, $p^0(x)=1/\int \text{d}{x}^{\prime }$. In this case, the KL divergence simplifies to the negative of the Shannon entropy of a distribution $p(x)$ [9],

$$H\left[p\right]=-\int \text{d}xp\left(x\right)\text{log}p\left(x\right).$$ (34)

Shannon argued that Eq. (34) is a functional that can quantify information [12]. When the only available information is a set of measured expectation values, the “least-committal” distribution consistent with them is the one that maximizes entropy [13, 14].

Mathematically put, suppose a set of observables was measured, $A=\int \text{d}xp(x)a(x)$. Combined with normalization, according to Eq. (31) (with entropy replacing the KL divergence) the inferred distribution is that which satisfies

$$\begin{gathered} p^A=\text{arg}\underset{p}{\text{max}}H\left[p\right],\text{s.t.}\int \text{d}xp(x)=1\text{and} \\ {\left\{\int \text{d}xp(x)a_m(x)=A_m\right\}}_{m=1,\dots ,M}. \end{gathered}$$ (35)

Using the method of Lagrange multipliers, one finds the solution

$$p^A(x)=\frac{1}{Z}\text{exp}\left[-\sum _{m=1}^M{\lambda }_m{a}_m(x)\right],$$ (36)

where

$$Z=\int \text{d}x\text{exp}\left[-\sum _{m=1}^M{\lambda }_m{a}_m(x)\right],$$ (37)

arose from the Lagrange multiplier enforcing normalization. ${\lambda }_m$ are the remaining Lagrange multipliers, determined from the implicit relation

$$A_m=-\frac{\partial \text{log}Z}{\partial {\lambda }_m}.$$ (38)

The maximized entropy reads

$$H\left[p^A\right]=\text{log}Z-\sum _{m=1}^M{\lambda }_m{A}_m.$$ (39)

The above is the general statement for an inference scheme based on the maximum-entropy principle. It can be readily generalized to the case where the prior $p^0(x)$ is not uniform, e.g., the inferred distribution will become the “exponential tilt” of the base distribution

$$p^A(x)\propto p^0(x)\text{exp}\left[-\sum _{m=1}^M{\lambda }_m{a}_m(x)\right].$$ (40)

In the above expressions, $a_m(x)$ can be polynomials, in which case $A$ gives the moments of the underlying distribution. However, a particularly relevant observable for later is that which checks if $x$ is in a region $𝒜,a(x)≔1_{𝒜}(x)$, so $A=\int a(x)p(x)dx=1$. The constraint is then $p^A(x\in 𝒜)≔p^A(𝒜)=1$. The maximum-entropy solution Eq. (35) is exactly the conditional probability distribution. At the same time, Bayes’ rule means $p(x\mid 𝒜)=1_{𝒜}(x)p^0(x)/p^0(𝒜)$. Thus, the constraint $p^A(𝒜)=1$ implies:

$$\begin{gathered} {\left.{\text{D}}_{\text{KL}}\left(p^A\Vert p^0\right)\right|}_{p^A(𝒜)=1}={\text{D}}_{\text{KL}}\left(1_{𝒜}\cdot p^A\Vert 1_{𝒜}\cdot p^0\right) \\ ={\text{D}}_{\text{KL}}\left(p^A\Vert p(\cdot \mid 𝒜)\right)-\text{log}{p}^0(𝒜) \end{gathered}$$ (41)

whose minimum is $p^A(x)=p(x\mid 𝒜)$. In this setting, maximum-entropy inference exactly reproduces Bayesian inference.

The minimum-KL-divergence formulation is closely related to Bayesian inference more generally. The “update” ratio $p^A(x)/p^0(x)=e^{\cdots }$ can be interpreted as the likelihood ratio, $p(A\mid x)$ by Bayes’ rule—the exponential factor reweighing the prior according to the information encoded by the imposed constraints. In this sense, the minimum-KL solution is thus a Bayesian-style update of the reference distribution, with the exponential tilt playing the role of an (unnormalized) likelihood. These general functional forms will appear repeatedly in the upcoming sections for specialized observables. We will furthermore see that the Lagrange multipliers will take the form of rewards, thereby generalizing ideas from Sec. IA to path-integral stochastic control. Furthermore, the ${\text{D}}_{\text{KL}}$ objective arises naturally in the context of Variational Inference (VI), the problem of approximating a complex distribution by a simpler—typically parametric—one, which will be discussed in Sec. V B 3.

Another noteworthy connection of maximum-entropy models is to exponential-family models [15, 16]. Maximum-entropy distributions subject to linear expectation constraints take the exponential-family form of Eq. (36), in which the Lagrange multipliers ${\lambda }_m$ enter linearly in the exponent and determine the partition function $Z$. Beyond their conceptual connection to maximum entropy, exponential-family models are useful because their parameters can be estimated through convex optimization, making them particularly attractive for statistical inference and applications such as image processing.

In addition to Sanov’s theorem and the Bayesian-inference connections, there is a subtle connection of maximum-entropy inference to equilibrium statistical mechanics. Sanov’s theorem provides concrete grounds for the maximum-entropy approach, as large deviations are exponentially suppressed as the number of samples increases. In the thermodynamic limit the dimensionality of $x$ (the number of particles), reaching as much as ~10²³, far exceeds any feasible number of samples that one may obtain. Instead, the second law of thermodynamics and statistical mechanics provide the theoretical basis for finding the microscopic statistical model of materials. Nevertheless, both mechanisms—equilibrium statistical mechanics and maximum-entropy modeling—provide the same mathematical structures [11]. We shall discuss this correspondence in Sec. III A.

C. Probabilistic inference over paths

In Sec. IA, we considered how to control a system with known dynamics to maximize a prescribed reward. In Sec. I B, we saw how inference problems can be formulated as optimization problems in which we seek a distribution that is consistent with observed data while remaining as close as possible to a prior distribution. We now turn to the application of the large-deviation principle to the inference of distributions over paths. We will exploit the tools of stochastic optimal control to solve these path-distribution inference problems.

In a stochastic system, repeated realizations of the noise generate different trajectories, so in the context of stochastic optimal control, a given control determines a probability distribution over paths. When investigating the properties of a distribution over paths, we will exploit that, in generic settings, there is a correspondence between stochastic differential equations (SDEs) and distributions over paths. In particular, an SDE determines a distribution over paths in which each realization of the noise generates a different trajectory. Under certain regularity conditions, there is a converse relationship where a distribution over paths can be represented by an SDE whose drift and diffusion coefficients are determined by the properties of the distribution; we will restrict the inference problems we consider to those for which this correspondence holds. We use uppercase $P$ for distributions over complete paths, reserving lowercase $p$ for state densities. If we have a prior distribution over paths, $P^0[x(\cdot )]$, and we observe some data about the system, the large-deviation principle tells us how to find the distribution over paths, $P^u[x(\cdot )]$, that is most similar to the prior while remaining consistent with the observations. Thus, because of the path distribution-SDE correspondence, the result of inference is a new SDE that describes the dynamics of the system under the inferred distribution over paths. Importantly, the inferred SDE will contain additional forces that will be determined by an optimization problem. To see this, consider a simple example of such a path-conditioning problem, namely, the Brownian bridge. In the Brownian bridge, we consider a particle known to start at a point $x_0$ at time $t=0$ and undergo Brownian motion, but we impose the condition that the particle is at a prescribed target $y$ at time $t=1$:

$$\begin{gathered} {\dot{x}}_t=\sqrt{\epsilon }{\eta }_t, \\ \text{conditioned}\text{on}{x}_{t=0}=x_0,x_{t=1}=y. \end{gathered}$$ (42)

Without this conditioning, the particle would hit $y$ at time $t=1$ formally with probability zero. Within the ensemble of such conditioned paths, there is apparently an additional force that drives the particle to $y$ at time $t=1$, which is not present in the unconditioned dynamics. Here, we will consider similar problems in which we condition on more general observations, and we will see that the large-deviation principle provides a natural way to infer the additional forces that emerge from this conditioning.

1. From inference to distributions over paths

Consider the prior distribution over paths, $P^0[x(\cdot )]$, induced by the passive dynamics (i.e., the dynamics without any additional forces),

$${\dot{x}}_t=f_t\left(x_t\right)+\sqrt{\epsilon }{\eta }_t.$$ (43)

We seek to infer how this ensemble changes when observations constrain its properties. We will represent the inferred ensemble through an additional control,

$${\dot{x}}_t=f_t\left(x_t\right)+u_t\left(x_t\right)+\sqrt{\epsilon }{\eta }_t,$$ (44)

and denote its path distribution by $P^u[x(\cdot )]$. The reference and controlled processes start at the same fixed $x_0$, as in Sec. IA. The marginal density at a single time is

$$p_t^u(x)={⟨\delta \left(x-x_t\right)⟩}_{P^u}.$$

We omit the superscript on a marginal when only one ensemble is under discussion. The large-deviation argument applies here with each sample being an entire trajectory. Among empirical ensembles compatible with the observations, it identifies those most likely to arise from repeated sampling of the passive dynamics. In the limit of many trajectories, the empirical path distribution concentrates on the corresponding inferred distribution.

Suppose the observations specify averages $A_m$ of trajectory observables $a_m[x(\cdot )]$, so that

$$\int 𝒟x{P}^u[x(\cdot )]a_m[x(\cdot )]=A_m.$$ (45)

Eq. (31) then states that the most likely path measure is found by minimizing over $P^u$ the KL divergence from the prior $P^0$ [17],

$${\text{D}}_{\text{KL}}\left(P^u\Vert P^0\right)=\int 𝒟x{P}^u[x(\cdot )]\text{log}\frac{P^u[x(\cdot )]}{P^0[x(\cdot )]}.$$ (46)

while satisfying these constraints. These constraints can be enforced with Lagrange multipliers $r_m$. We use the reward sign convention $r_m=-{\lambda }_m$, where ${\lambda }_m$ is the multiplier in the convention used in Sec. I B 2, so a positive reward appears with a positive sign in the exponential tilt. This gives the constrained objective:

$$ℒ={\text{D}}_{\text{KL}}\left(P^u\Vert P^0\right)-\sum _m{r}_m\left[\int 𝒟x{P}^u[x(\cdot )]a_m[x(\cdot )]-A_m\right].$$ (47)

Taking the variation with respect to the path distribution and setting it to zero gives the minimizing distribution. When represented by the optimal control $u^{*}$, it reads

$$P^{u^{*}}[x(\cdot )]\propto P^0[x(\cdot )]\text{exp}\left(\sum _m{r}_m{a}_m[x(\cdot )]\right),$$ (48)

meaning the inferred path measure is an exponential tilt of $P^0$ where the Lagrange multipliers $r_m$ enforce consistency with the observations. These Lagrange multipliers can be roughly interpreted as inferred rewards that the system, conditioned on the observations, is receiving. Because of the correspondence between path distributions and SDEs, the ‘tilts’ $r_m$ to the prior path distribution translate into apparent forces in the SDE describing paths drawn from $P^{u^{*}}$. Our goal now is to solve the constrained optimization problem introduced to get an expression for $u_t^{*}$ which is the apparent force that emerges when the system is conditioned on the observations.

2. Girsanov’s theorem

To do this, we exploit Girsanov’s theorem, which allows us to express the change of measure from $P^0$ to $P^u$ in terms of the control $u_t$ [18]. Girsanov’s theorem is a fundamental result in stochastic calculus because it provides a way to relate probability measures of stochastic processes with different drift terms. In the context of inference over paths, Girsanov’s theorem allows us to express the KL divergence between the prior and inferred path measures in terms of the added control $u_t$. This will allow us to formulate the inference problem as an optimization problem over the control, which can then be solved using techniques from stochastic optimal control.

Girsanov’s theorem states that the ratio of the path measures of two SDEs which are driven by (potentially inhomogeneous) Gaussian noise and have differing drift terms can be expressed in terms of the difference in the drift terms. To see this, note that the controlled dynamics can be made to look like the passive dynamics by relabeling the noise:

$$\begin{gathered} {\dot{x}}_t=f_t\left(x_t\right)+u_t\left(x_t\right)+\sqrt{\epsilon }{\eta }_t^u \\ =f_t\left(x_t\right)+\sqrt{\epsilon }\left[{\eta }_t^u+\frac{1}{\sqrt{\epsilon }}{u}_t\left(x_t\right)\right]. \end{gathered}$$ (49)

Here ${\eta }_t^u$ denotes the zero-mean noise under the controlled path measure. Thus, the same realized trajectory that requires the noise ${\eta }_t^u$ under the controlled dynamics would require the shifted noise ${\eta }_t^0={\eta }_t^u+\frac{1}{\sqrt{\epsilon }}{u}_t\left(x_t\right)$ under the passive dynamics. The ratio between the two measures can then be written as:²

$$\begin{gathered} \frac{P^u[x(\cdot )]}{P^0[x(\cdot )]}=\frac{\text{exp}\left[-\frac{1}{2}{\int }_0^1\text{d}t{‖{\eta }_t^u‖}^2\right]}{\text{exp}\left[-\frac{1}{2}{\int }_0^1\text{d}t{‖{\eta }_t^0‖}^2\right]} \\ =\text{exp}\left[\frac{1}{2\epsilon }{\int }_0^1\text{d}t{‖u_t\left(x_t\right)‖}^2+\frac{1}{\sqrt{\epsilon }}{\int }_0^1\text{d}t{u}_t\left(x_t\right)\cdot {\eta }_t^u\right]. \end{gathered}$$ (50)

The KL divergence is the expectation of the logarithm of this ratio under the controlled path measure:

$$\begin{gathered} {\text{D}}_{\text{KL}}\left(P^u\Vert P^0\right)=\frac{1}{2\epsilon }{⟨{\int }_0^1\text{d}t{‖u_t\left(x_t\right)‖}^2⟩}_{P^u} \\ +\frac{1}{\sqrt{\epsilon }}{⟨{\int }_0^1\text{d}t{u}_t\left(x_t\right)\cdot {\eta }_t^u⟩}_{P^u} \\ =\frac{1}{2\epsilon }{⟨{\int }_0^1\text{d}t{‖u_t\left(x_t\right)‖}^2⟩}_{P^u}. \end{gathered}$$ (51)

In the final equality, the term linear in the controlled noise has vanished because under the controlled path measure, the noise ${\eta }_t^u$ has mean zero. Eq. (50) is the standard statement of Girsanov’s theorem [5].

3. Probabilistic inference as a control problem

With Girsanov’s theorem, this means that we can write Eq. (31) more formally as an optimization problem:

$$\begin{gathered} u^{*}=\text{arg}{\text{max}}_u⟨-\frac{1}{2\epsilon }{\int }_0^1\text{d}t{‖u_t\left(x_t\right)‖}^2 \\ {+{\int }_0^1\text{d}t{r}_t\left(x_t\right)+r_1\left(x_1\right)⟩}_{P^u} \end{gathered}$$ (52)

where $r_t$ and $r_1$ are again Lagrange multipliers chosen to enforce the constraint of consistency with observations. Here the trajectory reward is additive in time, with a separate terminal term. In this special case, the running reward of Sec. IA is

$$\underset{\text{total}\text{running}\text{reward}}{\underbrace{r_t(x,u)}}=\underset{\text{state}\text{reward}}{\underbrace{r_t(x)}}-\frac{\Vert u{\Vert }^2}{2\epsilon }.$$

This now has the form of a control problem in which we are solving to find the optimal control $u_t^{*}\left(x_t\right)$ which maximizes the reward but has a cost from applying too strong a control [19–21]. We will see that the problem in Eq. (52) can be solved by converting it into a linear PDE for which Monte Carlo methods can be used to sample the optimal control.

In this context, the quadratic control cost has a probabilistic interpretation in that it measures how strongly the controlled ensemble departs from the passive dynamics. The same relation between noise and control cost that makes this inference interpretation possible also gives the associated HJB equation a particularly simple mathematical structure.

D. Mathematical structure of path integral stochastic control

Problems with objectives such as Eq. (52) are called path integral stochastic control (PISC) problems [22–26]. PISC problems are a special class of stochastic control problems in which the cost of control is quadratic in the control and the dynamics are linear in the control (i.e., in Eq. (1), $g_t(x,u)=\left.f_t(x)+u\right)$. Fig. 2 illustrates a solution to a PISC problem in which an ensemble of particles is controlled to achieve a distant terminal reward. These problems can be solved using a path integral perspective like in Eq. (48), but are most naturally solved using dynamic programming and the Hamilton–Jacobi–Bellman (HJB) equation as introduced in Sec. IA. PISC is a special class of stochastic control problems because the HJB equation takes a particularly simple form that can be transformed into a linear PDE.

FIG. 2.

Path integral stochastic control (PISC) for an ensemble of freely diffusing particles. Left: an ensemble of particles begins at a common initial condition and freely diffuses, failing to achieve some distant terminal reward. Middle: the value function which solves the PISC problem of the ensemble of particles achieving the maximum terminal reward subject to some control cost. Right: the resulting PISC policy leading a significant fraction of the particles to efficiently achieve the terminal reward.

1. The Cole–Hopf transformation of the HJB equation

To solve the inference problem, expressed as Eq. (52), we first define the value function to be the expected reward received if the optimal policy is taken starting at point $x$ at time $t$:

$$\begin{gathered} V_t(x)=\underset{\begin{array}{c}\left\{u_s\right\}\\ t\le s<1\end{array}}{\text{max}}⟨-\frac{1}{2\epsilon }{\int }_t^1\text{d}s{‖u_s\left(x_s\right)‖}^2 \\ +{\int }_t^1{r}_s\left(x_s\right)\text{d}s+r_1\left(x_1\right){\left|x_t=x\right.⟩}_{P^u} \end{gathered}$$ (53)

which using dynamic programming (Eq. (16)) must satisfy

$$\begin{gathered} 0=\underset{u}{\text{max}}\left\{r_t(x)+{\partial }_t{V}_t+\nabla V_t\cdot \left[f_t(x)+u\right]\right. \\ \left.+\frac{\epsilon }{2}{\nabla }^2{V}_t-\frac{1}{2\epsilon }\Vert u{\Vert }^2\right\} \end{gathered}$$ (54)

with the optimal control given by

$$u_t^{*}(x)=\epsilon \nabla V_t(x).$$ (55)

For PISC problems, the HJB equation is then the nonlinear PDE [27]

$$r_t(x)+{\partial }_t{V}_t+\nabla V_t\cdot f_t(x)+\frac{\epsilon }{2}{\nabla }^2{V}_t+\frac{\epsilon }{2}{‖\nabla V_t‖}^2=0$$ (56)

with the boundary condition $V_1(x)=r_1(x)$. Remarkably, this nonlinear PDE can be converted into a linear equation by the Cole–Hopf transformation [22, 23],

$$V_t(x)=\text{log}{\psi }_t(x),$$ (57)

which upon substitution gives:

$$\left[r_t(x)+{\partial }_t+f_t(x)\cdot \nabla +\frac{\epsilon }{2}{\nabla }^2\right]{\psi }_t(x)=0$$ (58)

which is a backward Kolmogorov equation with a potential term. Notably, this is also the Schrödinger equation in imaginary time (i.e., with $t\to it$) when $f_t=0$ with $\psi$ being the wavefunction, $\epsilon$ playing the role of $\hslash /m$, and $-r_t(x)$ playing the role of a potential energy term. In the small-$\epsilon$ limit, this parallels the passage from wave optics to geometrical optics, where the Hamilton–Jacobi equation becomes the eikonal equation. With this change of variables, the optimal control can be expressed as:

$$u_t^{*}\left(x\right)=\epsilon \nabla \text{log}{\psi }_t(x).$$ (59)

2. Feynman–Kac formula and Doob’s h-transform

The transformed equation is linear, but solving it over the full state space can still be difficult in high dimensions. An alternative is to express its solution as an average over trajectories of the passive dynamics, weighted by their accumulated rewards. This is the stochastic counterpart of the imaginary-time path-integral representation from quantum mechanics and provides a starting point for the sampling methods discussed in Chapter IV. From the same reasoning as in Eq. (48), the optimal path measure conditional on starting at $x$ at time $t$ is given below. In these conditional path laws, $x(\cdot )$ denotes the remaining trajectory on $t\le s\le 1$:

$$P^{u^{*}}\left[x(\cdot )\mid x_t=x\right]=\frac{1}{Z_t(x)}{P}^0\left[x(\cdot )\mid x_t=x\right]e^{R_t[x(\cdot )]}$$ (60)

where $R_t[x(\cdot )]={\int }_t^1{r}_s\left(x_s\right)\text{d}s+r_1\left(x_1\right)$ is the cumulative state and terminal reward from time $t$ to the final time 1, excluding control cost. We temporarily denote the normalization factor by $Z_t(x)$, which is a conditional partition function. Observe from Eq. (53) and Girsanov’s theorem that we can write the value function as

$$\begin{gathered} V_t(x)={⟨\left.R_t[x(\cdot )]-\frac{1}{2\epsilon }{\int }_t^1\text{d}s{‖u_s^{*}\left(x_s\right)‖}^2\right|x_t=x⟩}_{P^{u^{*}}} \\ ={⟨\left.R_t[x(\cdot )]-\text{log}\frac{P^{u^{*}}\left[x(\cdot )\mid x_t=x\right]}{P^0\left[x(\cdot )\mid x_t=x\right]}\right|x_t=x⟩}_{P^{u^{*}}} \\ ={⟨R_t[x(\cdot )]-\left(R_t[x(\cdot )]-\text{log}{Z}_t(x)\right)\mid x_t=x⟩}_{P^{u^{*}}} \\ =\text{log}{Z}_t\left(x\right), \end{gathered}$$ (61)

where the expectations above are all taken with respect to the optimal path measure $P^{u^{*}}\left[x(\cdot )\mid x_t=x\right]$. From Eq. (57), we then have ${\psi }_t(x)=Z_t(x)$. We henceforth use ${\psi }_t$ for this conditional partition function. It is a positive function of the current state, not a normalized state density, and can be expressed as a path integral over the passive dynamics:

$${\psi }_t(x)=\int 𝒟x{P}^0\left[x(\cdot )\mid x_t=x\right]\times \text{exp}\left({\int }_t^1{r}_s\left(x_s\right)\text{d}s+r_1\left(x_1\right)\right).$$ (62)

This is generalized to a result that a function ${\psi }_t(x)$ satisfying Eq. (58) can be expressed as a conditional expectation over the passive dynamics, which is known as the Feynman-Kac formula [5, 23]. Expressing ${\psi }_t$ in terms of a path integral is familiar from quantum mechanics, where the solution to the Schrödinger equation can be expressed as a path integral over all possible paths weighted by the exponential of the action, and the Feynman–Kac formula is the stochastic process analog of this result.

We can apply this result to the case where we are interested in inferring the dynamics of a system that is observed to be in a particular set $𝒜$ at time $t=1$. For a path measure $P$, we write $P(𝒜)\equiv P\left(x_1\in 𝒜\right)$ for the probability of this terminal event. In this conditioning example the running state reward is zero, $r_t(x)=0$. We are therefore interested in studying the case $P^u(𝒜)=1$, which requires $r_1\left(x_1\right)$ to be zero for $x_1\in 𝒜$ and $-\infty$ for $x_1\notin 𝒜$ (i.e., ${\psi }_1\left(x_1\right)=1_{𝒜}\left(x_1\right)$ where $1_{𝒜}$ is the indicator function for the set $𝒜$). From Eq. (62), we then have ${\psi }_t(x)=P^0\left(𝒜\mid x_t=x\right)$, which is the probability that the system is in the set $𝒜$ at time $t=1$ given that it is in state $x$ at time $t$ under the passive dynamics; this conditioning function is the function conventionally called $h$ in Doob’s $h$-transform [28]. We retain the symbol ${\psi }_t(x)$ here. Eq. (60) is then directly Bayes’ rule:

$$P^{u^{*}}\left[x(\cdot )\mid x_t=x\right]=P^0\left[x(\cdot )\mid x_t=x\right]\frac{{\mathbf{1}}_{𝒜}\left(x_1\right)}{P^0\left(𝒜\mid x_t=x\right)}$$ (63)

From Eq. (59), the inferred dynamics (known as Doob’s $h$-transform) that lead to the system being in the set $𝒜$ at time $t=1$ are then given by

$${\dot{x}}_t=f_t\left(x_t\right)+\epsilon \nabla \text{log}{\psi }_t\left(x_t\right)+\sqrt{\epsilon }{\eta }_t.$$ (64)

The logarithmic gradient $\nabla \text{log}{\psi }_t(x)$ is the conditioning or guidance term. It is closely related to the score $\nabla \text{log}{p}_t(x)$ used in generative modeling and diffusion models, but the two are generally distinct: ${\psi }_t(x)$ is a conditioning probability, whereas $p_t(x)$ is a marginal density. We return to scores and guidance in Sec. V C.

3. Example: Brownian bridge

The result in Eq. (64) allows us to now study the Brownian bridge (Eq. (42)) introduced in Sec. I C. In particular, we consider one-dimensional Brownian motion as the passive dynamics (i.e., $f_t(x)=0$) and the system is observed to be in a particular state $x_1=y$ at time $t=1$ while starting at $x_0=0$ at time $t=0$. To use Eq. (64) to solve this problem, we treat the target set $𝒜$ as a small ball of radius $\delta \to 0^{+}$ around the target state $y$, so that $𝒜=\{x:\Vert x-y\Vert <\delta \}$. Let $k_{t\mid s}(z\mid x)$ denote the transition density of the passive process from state $x$ at time $s$ to state $z$ at time $t>s$. Then:

$$\begin{gathered} P^0\left(𝒜\mid x_t=x\right)={\int }_{𝒜}\text{d}{x}_1{k}_{1\mid t}\left(x_1\mid x\right) \\ ={\int }_{𝒜}\text{d}{x}_1\frac{1}{\sqrt{2\pi \epsilon (1-t)}}\text{exp}\left(-\frac{{‖x_1-x‖}^2}{2\epsilon (1-t)}\right) \\ \approx \frac{1}{\sqrt{2\pi \epsilon (1-t)}}\text{exp}\left(-\frac{\Vert y-x{\Vert }^2}{2\epsilon (1-t)}\right){\int }_{𝒜}\text{d}{x}_1 \end{gathered}$$ (65)

The dynamics of a Brownian bridge are given by Eq. (64):

$${\dot{x}}_t=-\frac{x_t-y}{1-t}+\sqrt{\epsilon }{\eta }_t.$$ (66)

This means that in the conditioned dynamics, there is a drift term that pushes the system to move in a straight line from the starting point to the target point, with a strength that diverges as $t\to 1$.

Conditioning a stochastic system can therefore be understood dynamically. The information supplied by an observation is expressed as a control that makes the observed behavior typical. The Brownian bridge makes this equivalence concrete, with an endpoint constraint appearing as a drift toward the target.

II. CONTROL AND TRANSPORT OF DENSITIES

The variational viewpoint of Chapter I compares trajectories through their action or control cost. We now consider the problem of moving an entire distribution between prescribed initial and final distributions, rather than a particle between prescribed locations. An intuitive picture is moving a pile of earth into another shape: the total amount of earth moved is unchanged, but there are many ways to rearrange it. Choosing a rearrangement that minimizes a specified cost is an optimal transport problem. When the motion of mass is stochastic, one considers the most likely evolution consistent with the two distributions at the endpoints. We begin with a version of transport, called the Schrödinger bridge problem, formulated in terms of stochastic control. We then consider more general formulations of transport and the geometry it defines on probability distributions. This geometry allows us to describe the evolution of densities as steepest descent, connecting the present chapter to the discussions of nonequilibrium thermodynamics, sampling, and learning further below.

A. The Schrödinger bridge problem

The Schrödinger bridge problem (SBP) asks for the distribution over stochastic trajectories that matches two prescribed endpoint densities while differing as little as possible from a reference process [29–32]. Unlike the fixed-endpoint bridge considered in Sec. I D 3, the initial and final positions are now drawn from distributions ${\rho }_0$ and ${\rho }_1$. The control generates a curve of densities $t\mapsto {\rho }_t$ connecting these endpoints while accounting for all possible particle trajectories between them. Fig. 3 illustrates this passage from controlling trajectories to transporting densities.

FIG. 3.

The Schrödinger bridge problem (SBP) for an ensemble of freely diffusing particles. Left: an ensemble of particles begins distributed according to an initial density ${\rho }_0$ and freely diffuses to ${\rho }_1^0$, failing to achieve the target terminal density ${\rho }_1$. Middle: the value function $V_t(x)=\text{log}{\psi }_t(x)$ which solves the SBP of transporting the ensemble between the prescribed initial and terminal densities subject to a control cost. Right: the resulting control policy transports the particles so that the terminal density of the controlled ensemble matches the target density.

1. From trajectories to densities

As in Sec. I C, we consider non-interacting particles with controlled Langevin dynamics,

$${\dot{x}}_t=f_t\left(x_t\right)+u_t\left(x_t\right)+\sqrt{\epsilon }{\eta }_t,0\le t\le 1,$$ (67)

which is Eq. (1) with $g_t(x,u)=f_t(x)+u$ and feedback $u=u_t(x)$. Here ${\eta }_t$ is standard Gaussian white noise and $\epsilon >0$. The reference dynamics have $u_t=0$. Both the reference and controlled processes are initialized with $x_0~{\rho }_0$. We write $P^0[x(\cdot )]$ and $P^u[x(\cdot )]$ for the reference and controlled distributions over complete paths, respectively; these are the path distributions introduced in Sec. I C. In contrast, ${\rho }_t(x)$ is the density of the position at a single time under $P^u$, i.e., it is the marginal $p_t^u(x)$ of Chapter I. We use ${\rho }_t$ when emphasizing the transport of this density, and add a superscript when distinguishing different ensembles. For the independent particles considered here, this is both the single-particle marginal density and the density described by a large ensemble. The density generated by Eq. (67) obeys the Fokker–Planck equation, which expresses local conservation of probability:

$$\begin{gathered} {\partial }_t{\rho }_t=-\nabla \cdot \left[\left(f_t+u_t\right){\rho }_t\right]+\frac{\epsilon }{2}{\nabla }^2{\rho }_t, \\ {\rho }_{t=0}={\rho }_0,{\rho }_{t=1}={\rho }_1. \end{gathered}$$ (68)

It is useful to distinguish the added control $u_t$ from the velocity $v_t$ that transports the density. Eq. (68) can equivalently be written as

$$\begin{gathered} {\partial }_t{\rho }_t+\nabla \cdot \left({\rho }_t{v}_t\right)=0, \\ {v}_t=f_t+u_t-\frac{\epsilon }{2}\nabla \text{log}{\rho }_t. \end{gathered}$$

Here ${\rho }_t{v}_t$ is the probability current. In the deterministic, driftfree case, $v_t=u_t$. Note that even when intermediate densities ${\rho }_t$ are specified, we cannot identify a unique velocity field; adding a field $w_t$ with $\nabla \cdot \left({\rho }_t{w}_t\right)=0$ leaves the continuity equation unchanged. Specifying the two endpoint densities leaves still more freedom, since the intermediate densities are also unspecified. Thus, as in maximum-entropy inference, the constraints must be supplemented by a criterion for choosing among the compatible path distributions.

Suppose a large collection of independent particles evolving under $P^0$ is observed to begin with density ${\rho }_0$ and end with density ${\rho }_1$. The large-deviation argument of Eq. (31), now applied to paths, selects the distribution over paths that satisfies these observations and is closest to $P^0$ in KL divergence. This gives the SBP,

$$\begin{gathered} \underset{P^u}{\text{min}}{\text{D}}_{\text{KL}}\left(P^u\Vert P^0\right), \\ \text{s.t.}{\rho }_{t=0}={\rho }_0,{\rho }_{t=1}={\rho }_1, \end{gathered}$$ (69)

where ${\rho }_t$ is the marginal density under $P^u$. To relate this problem to control, we write $P^u[x(\cdot )]={\rho }_0\left(x_0\right)P^u\left[x(\cdot )\mid x_0\right]$, and similarly for $P^0$. The chain rule for KL divergence gives

$$\begin{gathered} {\text{D}}_{\text{KL}}\left(P^u\Vert P^0\right) \\ =\int \text{d}{x}_0{\rho }_0\left(x_0\right){\text{D}}_{\text{KL}}\left(P^u\left[\cdot \mid x_0\right]\Vert P^0\left[\cdot \mid x_0\right]\right) \\ =\frac{1}{2\epsilon }{\int }_0^1\text{d}t\int \text{d}x{\rho }_t\left(x\right){‖u_t\left(x\right)‖}^2, \end{gathered}$$ (70)

where the second equality follows from Girsanov’s theorem, Eq. (51), applied to paths with the same initial position. Consequently, minimizing Eq. (69) is equivalent to minimizing the expected quadratic control effort subject to Eq. (68) and its endpoint constraints. The SBP thus chooses the least surprising controlled ensemble relative to the reference process, or equivalently the one requiring the least expected control effort. This introduces the link between control, inference and transport, i.e., the control changes the path distribution and, through its single-time marginals, transports the density from ${\rho }_0$ to ${\rho }_1$.

2. The coupled forward–backward construction

We can impose the initial constraint directly by drawing $x_0$ from ${\rho }_0$, and impose the terminal constraint using a Lagrange multiplier function $r_1\left(x_1\right)$. As in Sec. I D, this multiplier acts as a terminal reward. For a fixed $r_1$, and after dropping terms independent of the control, the problem is

$$\begin{gathered} \underset{u}{\text{min}}{⟨\frac{1}{2\epsilon }{\int }_0^1\text{d}t{‖u_t\left(x_t\right)‖}^2-r_1\left(x_1\right)⟩}_{P^u}, \\ {x}_0~{\rho }_0. \end{gathered}$$ (71)

An initial reward $r_0\left(x_0\right)$ would have the fixed expectation $\int \text{d}x{\rho }_0(x)r_0(x)$ and would not change this minimization. Maximizing the negative of the objective in Eq. (71) is analogous to solving the PISC problem considered in Sec. I D, with no additional running reward. Its value function therefore admits the Cole–Hopf transformation $V_t=\text{log}{\psi }_t$, and Eq. (59) gives

$$\begin{gathered} u_t^{*}(x)=\epsilon \nabla \text{log}{\psi }_t\left(x\right),{\psi }_1\left(x\right)=e^{r_1\left(x\right)}, \\ {\partial }_t{\psi }_t+f_t\cdot \nabla {\psi }_t+\frac{\epsilon }{2}{\nabla }^2{\psi }_t=0. \end{gathered}$$ (72)

Unlike a prescribed terminal reward in PISC, $r_1$ must now be chosen so that the terminal density is ${\rho }_1$. Since ${\psi }_t$ satisfies the backward Kolmogorov equation, it is interpreted as a field that propagates information backwards from $t=1$.

To find the density under this control, define a second field ${\psi }_t^{†}(x)={\rho }_t(x)/{\psi }_t(x)$, so that

$${\rho }_t(x)={\psi }_t(x){\psi }_t^{†}(x).$$ (73)

The two factors need not be normalized densities. Schrödinger’s key insight was that ${\psi }_t^{†}$ obeys the forward Kolmogorov equation of the reference process,

$${\partial }_t{\psi }_t^{†}+\nabla \cdot \left(f_t{\psi }_t^{†}\right)-\frac{\epsilon }{2}{\nabla }^2{\psi }_t^{†}=0.$$ (74)

To see this, substitute ${\rho }_t={\psi }_t{\psi }_t^{†}$ and $u_t^{*}=\epsilon \nabla \text{log}{\psi }_t$ into Eq. (68). Expanding the derivatives gives

$$\begin{gathered} {\partial }_t\left({\psi }_t{\psi }_t^{†}\right)=-\nabla \cdot \left(f_t{\psi }_t{\psi }_t^{†}\right)-\epsilon \nabla \cdot \left({\psi }_t^{†}\nabla {\psi }_t\right)+\frac{\epsilon }{2}{\nabla }^2\left({\psi }_t{\psi }_t^{†}\right) \\ ={\psi }_t^{†}\left[-f_t\cdot \nabla {\psi }_t-\frac{\epsilon }{2}{\nabla }^2{\psi }_t\right]+{\psi }_t\left[-\nabla \cdot \left(f_t{\psi }_t^{†}\right)+\frac{\epsilon }{2}{\nabla }^2{\psi }_t^{†}\right]. \end{gathered}$$ (75)

The first bracket is ${\partial }_t{\psi }_t$ by Eq. (72). Cancelling ${\psi }_t^{†}{\partial }_t{\psi }_t$ from both sides and dividing by ${\psi }_t$ yields Eq. (74). Thus, ${\psi }_t$ propagates information backward from the final constraint, while ${\psi }_t^{†}$ propagates information forward from the initial constraint.

Let $k_{t\mid s}(x\mid y)$ denote the transition density of the reference process, namely, the conditional density of being at $x$ at time $t$ given position $y$ at time $s<t$, with $u=0$ in Eq. (67). It is the propagator of the reference Fokker–Planck equation. The solutions of the backward and forward equations can therefore be written as

$$\begin{gathered} {\psi }_t(x)=\int \text{d}y{k}_{1\mid t}\left(y∣x\right){\psi }_1\left(y\right), \\ {\psi }_t^{†}(x)=\int \text{d}y{k}_{t\mid 0}\left(x∣y\right){\psi }_0^{†}\left(y\right). \end{gathered}$$ (76)

The endpoint constraints couple these otherwise linear equations,

$$\begin{gathered} {\psi }_0\left(x\right){\psi }_0^{†}\left(x\right)={\rho }_0\left(x\right), \\ {\psi }_1\left(x\right){\psi }_1^{†}\left(x\right)={\rho }_1\left(x\right). \end{gathered}$$ (77)

Note that neither endpoint factor is known in advance; specifying ${\psi }_1$ determines ${\psi }_0$ by backward propagation, while specifying ${\psi }_0^{†}$ determines ${\psi }_1^{†}$ by forward propagation. One may solve this Schrödinger system by alternating between the two constraints. Starting from a positive guess for ${\psi }_1$, propagate it backward using the first line of Eq. (76) to obtain ${\psi }_0$. Set ${\psi }_0^{†}={\rho }_0/{\psi }_0$, propagate this field forward using the second line to obtain ${\psi }_1^{†}$, and update ${\psi }_1={\rho }_1/{\psi }_1^{†}$. Repeating these steps enforces the initial and final constraints in turn. This procedure is the so-called Sinkhorn algorithm [32, 33].

With the two fields determined, the optimal distribution over paths is an endpoint reweighting, or ‘tilt’, of the reference distribution [31, 32],

$$P^{u^{*}}[x(\cdot )]=P^0[x(\cdot )]\frac{{\psi }_1\left(x_1\right)}{{\psi }_0\left(x_0\right)}.$$ (78)

Eq. (78) has the same exponential-tilt structure encountered in maximum-entropy inference and extends the Doob $h$-transform to prescribed endpoint distributions. Instead of conditioning an endpoint to lie in a set, the endpoint factors are chosen together to reproduce two full densities. The central result is therefore not a new form of the optimal control, which remains $u_t^{*}=\epsilon \nabla \text{log}{\psi }_t$, but the coupled forward–backward construction that determines this control from the two distributional constraints.

B. Entropic optimal transport and the Wasserstein distance

The Schrödinger bridge specifies a complete distribution over paths, and thus solves a dynamic transport problem. More generally, optimal transport considers the least expensive rearrangement of mass from one particular configuration to another, given a cost $c\left(x_0,x_1\right)$ for moving a unit mass from $x_0$ to $x_1$. In Monge’s original formulation, one views optimal transport as the most efficient strategy for moving a pile of earth from one location to another. Each point $x_0$ is assigned a destination $T\left(x_0\right)$, and one minimizes $\int \text{d}{x}_0{\rho }_0\left(x_0\right)c\left(x_0,T\left(x_0\right)\right)$ while requiring $T\left(x_0\right)~{\rho }_1$ for $x_0~{\rho }_0$. However, a deterministic map cannot split mass, and therefore need not exist; for example, a point mass cannot be mapped deterministically into two separated point masses. Kantorovich’s formulation instead optimizes over a joint endpoint distribution $\pi \left(x_0,x_1\right)$ satisfying

$$\begin{gathered} \int \text{d}{x}_1\pi \left(x_0,x_1\right)={\rho }_0\left(x_0\right), \\ \int \text{d}{x}_0\pi \left(x_0,x_1\right)={\rho }_1\left(x_1\right). \end{gathered}$$ (79)

Any such joint distribution is called a coupling of ${\rho }_0$ and ${\rho }_1$. Minimizing the mean cost over couplings allows the probability at one $x_0$ to be distributed among several destinations [34].

To connect this static problem to the Schrödinger bridge, let $k\left(x_1\mid x_0\right)$ denote the transition density of the reference process from $t=0$ to $t=1$ (we have dropped the $1\mid 0$ subscript from our notation in Sec. II A). Consider the following optimization over couplings [31]:

$${\pi }^{*}=\text{arg}{\text{min}}_{\pi }\left\{\epsilon {\text{D}}_{\text{KL}}\left(\pi \Vert {\rho }_{0}k\right)\right\},$$ (80)

subject to Eq. (79), where $\left({\rho }_{0}k\right)\left(x_0,x_1\right)={\rho }_0\left(x_0\right)k\left(x_1\mid x_0\right)$. Now, le‘ts say $k\left(x_1\mid x_0\right)=\text{exp}\left(-c\left(x_0,x_1\right)/\epsilon \right)/Z\left(x_0\right)$. For example, if the reference process is Brownian motion, we have $c\left(x_0,x_1\right)={‖x_0-x_1‖}^2/2$. Here the same $\epsilon >0$ that sets the noise covariance in Eq. (67) sets the strength of the regularization (and is analogous to a temperature), and $Z\left(x_0\right)$ normalizes the kernel over $x_1$. Using the fixed marginal constraints, Eq. (80) is equivalent, up to constants, to

$${\pi }^{*}=\text{arg}{\text{min}}_{\pi }\left\{\int \text{d}{x}_0\text{d}{x}_1\pi \left(x_0,x_1\right)c\left(x_0,x_1\right)\right.\left.+\epsilon \int \text{d}{x}_0\text{d}{x}_1\pi \left(x_0,x_1\right)\text{log}\pi \left(x_0,x_1\right)\right\}.$$ (81)

This objective corresponds to entropic optimal transport [35]. The first term is the mean transport cost, while the second is minus the Shannon entropy of the coupling introduced in Sec. I B 2. Eq. (81) therefore has the form of a constrained free-energy minimization: the cost favors efficient transport, whereas the entropy spreads probability among several possible pairings. As $\epsilon \to 0$, the entropic term disappears and one obtains the Kantorovich formulation of optimal transport [34, 36].

Introducing Lagrange multipliers ${\phi }_0\left(x_0\right)$ and ${\phi }_1\left(x_1\right)$ for imposing the constraints in Eq. (79), and absorbing the overall normalization into them, gives

$$\begin{gathered} {\pi }^{*}\left(x_0,x_1\right)=\text{exp}\left(\frac{{\phi }_0\left(x_0\right)+{\phi }_1\left(x_1\right)-c\left(x_0,x_1\right)}{\epsilon }\right) \\ =Z\left(x_0\right)\text{exp}\left(\frac{{\phi }_0\left(x_0\right)}{\epsilon }\right)k_{1\mid 0}\left(x_1\mid x_0\right)\times \text{exp}\left(\frac{{\phi }_1\left(x_1\right)}{\epsilon }\right). \end{gathered}$$ (82)

These multipliers depend on $\epsilon$ and their zero-noise limits are known as the Kantorovich potentials. A connection to the Schrödinger potentials can be made by setting ${\psi }_0^{†}=Z\text{exp}\left({\phi }_0/\epsilon \right)$ and ${\psi }_1=\text{exp}\left({\phi }_1/\epsilon \right)$, which gives

$$\begin{gathered} {\pi }^{*}\left(x_0,x_1\right)={\psi }_0^{†}\left(x_0\right)k_{1\mid 0}\left(x_1\mid x_0\right){\psi }_1\left(x_1\right), \\ {\psi }_0\left(x_0\right)=\int \text{d}{x}_1{k}_{1\mid 0}\left(x_1\mid x_0\right){\psi }_1\left(x_1\right), \\ {\psi }_1^{†}\left(x_1\right)=\int \text{d}{x}_0{k}_{1\mid 0}\left(x_1\mid x_0\right){\psi }_0^{†}\left(x_0\right). \end{gathered}$$ (83)

together with ${\rho }_0={\psi }_0^{†}{\psi }_0$ and ${\rho }_1={\psi }_1^{†}{\psi }_1$. These are precisely the endpoint versions of Eqs. (76) and (77). Sinkhorn iteration, analogous to the Sinkhorn algorithm introduced previously, is the alternating rescaling of ${\psi }_0^{†}$ and ${\psi }_1$ needed to solve for these quantities, and thus the optimal coupling ${\pi }^{*}$ [35].

The quadratic cost $c\left(x_0,x_1\right)={‖x_0-x_1‖}^2/2$ has special geometric structure. For the zero-noise argument, write ${\pi }_{\epsilon }^{*}$ and ${\phi }_1^{\epsilon }$ to display their dependence on $\epsilon$. From Eq. (82), the conditional density of $x_1$ given $x_0$ is

$$\frac{{\pi }_{\epsilon }^{*}\left(x_0,x_1\right)}{{\rho }_0\left(x_0\right)}\propto \text{exp}\left[\frac{{\phi }_1^{\epsilon }\left(x_1\right)-\frac{1}{2}{‖x_0-x_1‖}^2}{\epsilon }\right],$$ (84)

where factors depending only on $x_0$ are absorbed into the normalization. Rewriting the exponent as

$$-\frac{{‖x_0‖}^2}{2\epsilon }+\frac{1}{\epsilon }\left[x_0\cdot x_1-\left(\frac{{‖x_1‖}^2}{2}-{\phi }_1^{\epsilon }\left(x_1\right)\right)\right],$$ (85)

shows by the Laplace principle that, as $\epsilon \to 0$, the conditional density concentrates at the maximizing endpoints. Assuming the potentials converge with a fixed choice of additive constants, denote the limiting potential by ${\phi }_1={\text{lim}}_{\epsilon \to 0}{\phi }_1^{\epsilon }$ and define the supremum

$$\begin{gathered} \zeta \left(x_0\right)=\underset{x_1}{\text{sup}}\left\{x_0\cdot x_1-{\bar{\phi }}_1\left(x_1\right)\right\}, \\ \text{where}{\bar{\phi }}_1\left(x_1\right)=\frac{1}{2}{‖x_1‖}^2-{\phi }_1\left(x_1\right). \end{gathered}$$ (86)

If the maximizer $x_1^{*}\left(x_0\right)$ is unique, differentiating the supremum with respect to $x_0$ gives $\nabla \zeta \left(x_0\right)=x_1^{*}\left(x_0\right)$, so the limiting optimal coupling is supported on the map

$$x_1=T\left(x_0\right)=\nabla \zeta \left(x_0\right).$$ (87)

Because $\zeta$ is a convex conjugate, it is convex. This result is known as Brenier’s theorem, which states that for a quadratic cost and absolutely continuous ${\rho }_0$, the optimal transport map is the gradient of such a convex function [37]. In one dimension, this reduces to the familiar method for mapping a uniform distribution to an arbitrary distribution using the inverse of its cumulative distribution. Specifically, defining the cumulative distribution functions $C_i(x)={\int }_{-\infty }^x\text{d}y{\rho }_i(y),C_0\left(x_0\right)$ is uniform on $(0,1)$ when $x_0~{\rho }_0$, and $x_1=C_1^{-1}\left(C_0\left(x_0\right)\right)$ has density ${\rho }_1$, with $C_i^{-1}$ denoting the quantile function. Hence,

$$T=C_1^{-1}\circ C_0$$ (88)

is the one-dimensional Brenier map.

These special properties motivate the 2-Wasserstein distance,

$$W_2^2\left({\rho }_0,{\rho }_1\right)=\underset{\pi }{\text{min}}\left\{\int \text{d}{x}_0\text{d}{x}_1{‖x_0-x_1‖}^2\pi \left(x_0,x_1\right)\right\},$$ (89)

where the minimization is over the couplings in Eq. (79). With the convention $c={‖x_0-x_1‖}^2/2$ used above, the unregularized minimum mean cost is $W_2^2/2$. In one dimension, the relationship to the cumulative distribution gives

$$W_2^2\left({\rho }_0,{\rho }_1\right)={\int }_0^1\text{d}z{\left|C_0^{-1}(z)-C_1^{-1}(z)\right|}^2.$$ (90)

The distinction between KL divergence and the 2-Wasserstein distance is illustrated in Fig. 4. In the chosen family of bimodal densities, changes in the relative weights and separation of the modes produce comparable KL divergences, whereas Wasserstein distance distinguishes them strongly because changing the relative weights requires transporting probability between the modes.

FIG. 4.

A comparison of the geometries induced by the KL divergence and the 2-Wasserstein distance on a two-parameter family of probability densities $p_{\theta }$ relative to the reference density $q$. The inset curves show representative densities at different points in parameter space. The ${\theta }_1$ parameter changes the balance of modes of a bi-modal distribution; the ${\theta }_2$ parameter shifts the distance between the modes. Left: the KL divergence assigns similar costs to changes along the two parameter directions, producing approximately isotropic contours. Right: the 2-Wasserstein distance assigns a greater cost to changes that transport probability between the separated modes, producing strongly anisotropic contours. The figure is adapted from Ref. [38], and the distribution’s parameterization is described in the Supplementary Information of that work.

Note that the static formulation specifies which initial and final points are paired, but not how mass moves between them. In this deterministic, drift-free setting, the transport velocity equals the control, $v_t=u_t$. For a particle following ${\dot{x}}_t=v_t\left(x_t\right)$, the Cauchy–Schwarz inequality gives ${‖x_1-x_0‖}^2\le {\int }_0^1\text{d}t{‖v_t\left(x_t\right)‖}^2$, with equality for straight motion at constant speed. For an entire density, the velocity becomes a field $v_t(x)$ and valid paths must conserve probability locally. The remarkable content of the Benamou–Brenier theorem is that minimizing this kinetic action over all density paths and velocity fields gives exactly the same quantity as minimizing the static endpoint cost over couplings [36, 39]:

$$W_2^2\left({\rho }_0,{\rho }_1\right)=\underset{{\rho }_t,v_t}{\text{min}}\left\{{\int }_0^1\text{d}t\int \text{d}x{\rho }_t(x){‖v_t(x)‖}^2\right\},{\partial }_t{\rho }_t+\nabla \cdot \left({\rho }_t{v}_t\right)=0,$$ (91)

with endpoint constraints on ${\rho }_0$ and ${\rho }_1$. For a Brownian reference with $f_t=0$, multiplying the objective in Eq. (69) by $2\epsilon$ does not change its minimizer. In the limit $\epsilon \to 0$, the diffusion term in Eq. (68) vanishes, so Eq. (91) can be interpreted as the zero-noise limit of the Schrödinger bridge, namely, the rescaled minimum $2\epsilon {\text{min}}_{P^u}{\text{D}}_{\text{KL}}\left(P^u\Vert P^0\right)$ approaches $W_2^2$ [40].

Once the optimal map $T$ is known, the minimizing density path is generated by linear interpolation of each transported particle,

$$x_t=(1-t)x_0+tT\left(x_0\right),x_0~{\rho }_0.$$ (92)

The resulting ${\rho }_t$ is called displacement interpolation, which is analogous to a straight line in the geometry of probability distributions [41]. A useful feature of the 2-Wasserstein distance is that the static map $T$ also gives us the transport map for all intermediate times $t$, which we will return to later when we discuss flow matching in Sec. V C 2. As expected, this picture is consistent with the Brownian bridge discussed previously (Eq. (42)). For a fixed endpoint pair, the mean at an intermediate time $t$ follows $(1-t)x_0+t{x}_1$, with covariance $\epsilon t(1-t)I$. As $\epsilon \to 0$, the fluctuations disappear and the bridge collapses onto that straight path.

Although we focus on the quadratic cost here, other choices can produce different optimal assignments between the same endpoints. Increasing the exponent in $c\left(x_0,x_1\right)={‖x_0-x_1‖}^p$ penalizes long displacements more strongly and changes how the source is partitioned among those targets. The optimal pairing also depends on the choice of transport cost. Fig. 5 shows an example with three target locations and compares linear, quadratic, and quartic distance costs for the same initial and final densities, showing how this choice changes the regions of the initial density assigned to each target mode.

FIG. 5.

Optimal transport maps between the same source density ${\rho }_0$ and three-mode target density ${\rho }_1$ under different choices of transport cost, $c$. In each panel, the three colored regions partition the initial density according to the correspondingly colored target mode, while one arrow shows a representative assignment from each region. Increasing the exponent penalizes long displacements more strongly and consequently changes how the initial density is partitioned among the target modes. The axes $x^{(1)},x^{(2)}$ denote spatial components, whereas the subscripts in $x_0,x_1$ denote initial and final endpoints.

C. Wasserstein geometry and Wasserstein gradient flows

As we've now seen, the 2-Wasserstein distance $W_2$ in Eq. (91) provides a notion of distance between probability densities in terms of the optimal flow field to transport one to another. This notion of distance provides the starting point for Wasserstein geometry, in which one treats the space of probability distributions like a Riemannian manifold in which the square root of the minimum in Eq. (91) is the distance between points [42]. This will allow us to generalize the notion of gradient flow for particles in a potential $\dot{x}=-\nabla F(x)$ to that of densities flowing down gradients of a functional. This geometric picture will provide the basis for elegant algorithms presented in the Applications Sec. V B.

In Wasserstein geometry, a probability density $\rho$ is regarded as a point, a time-dependent density ${\rho }_t$ as a curve, and its instantaneous rate of change $s_t(x)={\partial }_t{\rho }_t(x)$ as a tangent vector. Crucially, Wasserstein geometry only permits trajectories ${\rho }_t(x)$ that obey the continuity constraint, namely that they can be realized via a flow field $v_t$

$${\partial }_t{\rho }_t=-\nabla \cdot \left({\rho }_t{v}_t\right)$$ (93)

Some trajectories are not allowed because they cannot be realized by a smooth flow field, for instance ${\rho }_t={\rho }_{1}t+(1-t){\rho }_0$ when ${\rho }_0$ and ${\rho }_1$ have disjoint support. This space of probability measures with finite second moment, equipped with $W_2$, is denoted ${𝒲}_2$ [43, 44].

Due to the continuity constraint in Eq. (93), it is possible to set up a duality between each tangent vector $s(x)={\partial }_t{\rho }_t$ and the velocity field $v$ that generates it. The subtlety is that the velocity is not unique: adding a field whose probability current has zero divergence leaves $s$ unchanged. The canonical way to remove this redundancy is motivated by the metric $W_2$. Namely, given the tangent vector $s$ in the tangent space of the distribution $\rho$, choose the velocity field $v$ with minimum kinetic energy that generates $s$

$$v=\text{arg}\underset{v^{\prime }}{\text{min}}\left\{\int \text{d}x\rho {‖v^{\prime }‖}^2\right.\text{s.t.}\left.s+\nabla \cdot \left(\rho v^{\prime }\right)=0\right\}$$ (94)

This minimization problem can be solved with a Lagrange multiplier for the continuity constraint. Let’s say we choose its sign and normalization by adding $-2\int \text{d}x\chi (x)[s(x)+\nabla .(\rho v)]$ to the kinetic energy. Variation with respect to $v$ then gives $v(x)=-\nabla \chi (x)$. Thus, the minimum energy velocity is a potential gradient.

Having fixed $v$ in this way, the minimum $\Vert s{\Vert }_{\rho }^2=\int \text{d}x\rho \Vert v{\Vert }^2$ in Eq. (94) defines the squared norm of the tangent at $\rho$. Along a regular curve, this is the squared speed measured by $W_2$:

$${‖s_t‖}_{{\rho }_t}=\underset{\Delta t\to 0}{\text{lim}}\frac{W_2\left({\rho }_t,{\rho }_{t+\Delta t}\right)}{\left|\Delta t\right|},s_t={\partial }_t{\rho }_t.$$ (95)

Equivalently, $W_2^2\left({\rho }_t,{\rho }_{t+\Delta t}\right)=(\Delta t{)}^2{‖s_t‖}_{{\rho }_t}^2+o\left((\Delta t{)}^2\right)$. The inner product of two tangents $s,s^{\prime }$ follows from the polarization identity ${⟨s,s^{\prime }⟩}_{\rho }=\frac{1}{2}\left[{‖s+s^{\prime }‖}_{\rho }^2-\Vert s{\Vert }_{\rho }^2-{‖s^{\prime }‖}_{\rho }^2\right]$, giving

$${⟨s,s^{\prime }⟩}_{\rho }=\int \text{d}x\rho (x)v(x)\cdot v^{\prime }(x).$$ (96)

where $v(x)$ and $v^{\prime }(x)$ are the minimum-energy velocity fields associated with $s$ and $s^{\prime }$, respectively.

With the ingredients of Wasserstein geometry in hand, we can now define gradient descent for a functional $ℱ\left[\rho \right]$. Recall that for a function $F(x)$ in Euclidean space, its gradient is defined in such a way as to enforce the chain rule: $\nabla F$ is the vector with the property $\text{d}F\left(x_t\right)/\text{d}t={\dot{x}}_t\cdot \nabla F\left(x_t\right)$ for all trajectories $x_t$. In complete analogy, the Wasserstein gradient is defined to be the velocity field ${\nabla }_{\text{W}}ℱ\left[\rho \right]$, with associated tangent vector $s=-\nabla \cdot \left(\rho {\nabla }_{\text{W}}ℱ\left[\rho \right]\right)$, such that $\frac{\text{d}ℱ\left[{\rho }_t\right]}{\text{d}t}={⟨s,{\partial }_t{\rho }_t⟩}_{{\rho }_t}$ for all trajectories ${\rho }_t$. To derive an explicit expression for ${\nabla }_{\text{W}}ℱ\left[\rho \right]$, consider $\text{d}ℱ/\text{d}t$ over a flow generated by a potential $v=-\nabla \chi$. Writing $\delta ℱ/\delta \rho$ for the functional derivative, evaluated at the current density ${\rho }_t$, one obtains

$$\begin{gathered} \frac{\text{d}ℱ\left[{\rho }_t\right]}{\text{d}t}=\int \text{d}x\frac{\delta ℱ}{\delta \rho }{\partial }_t{\rho }_t=\int \text{d}x\frac{\delta ℱ}{\delta \rho }\nabla \cdot (\rho \nabla \chi ) \\ =-\int \text{d}x\rho \nabla \frac{\delta ℱ}{\delta \rho }\cdot \nabla \chi . \end{gathered}$$ (97)

Comparing Eq. (97) with the inner product in Eq. (96), the Wasserstein gradient velocity field is given by ${\nabla }_{\text{W}}ℱ\left[\rho \right]=\nabla (\delta ℱ/\delta \rho )$. This is the gradient in its velocity representation, whereas the associated density tangent is $-\nabla \cdot \left(\rho {\nabla }_{\text{W}}ℱ\left[\rho \right]\right)$. A Wasserstein gradient flow (WGF) follows the negative gradient, with descent velocity $v_t=-{\nabla }_{\text{W}}ℱ\left[{\rho }_t\right]$. The resulting density evolution is

$$\begin{gathered} {\partial }_t{\rho }_t=\nabla \cdot \left({\rho }_t{\nabla }_{\text{W}}ℱ\left[{\rho }_t\right]\right), \\ \frac{\text{d}ℱ\left[{\rho }_t\right]}{\text{d}t}=-\int \text{d}x{\rho }_t{‖{\nabla }_{\text{W}}ℱ\left[{\rho }_t\right]‖}^2\le 0. \end{gathered}$$ (98)

The WGF is a distributional analog of gradient descent $\dot{x}=-\nabla F(x)$ in Euclidean space. For example, consider the free energy

$$ℱ\left[\rho \right]=\int \text{d}x\rho (x)\Phi (x)+\int \text{d}x\rho (x)\text{log}\rho (x).$$ (99)

Eq. (98) becomes ${\partial }_t{\rho }_t=\nabla \cdot \left({\rho }_t\nabla \Phi \right)+{\nabla }^2{\rho }_t$, the familiar drift-diffusion Fokker–Planck equation [45]. The potential energy term drives probability down $\Phi$, while the entropy term produces diffusion. Thus, the relaxation of a diffusive system can be read as steepest descent of a free energy in Wasserstein geometry, a connection that will later be discussed in the context of non-equilibrium thermodynamics (Sec. III B) and applications (Sec. V B).

Once steepest descent has been defined, one can ask familiar questions from optimization, now in the space of densities. For example, it is useful to define a notion of convexity that rules out competing local minima and controls the convergence rate—these ideas will come up again in Sec. V B on applications of WGF, particularly in the setting of variational inference. Recall that in Euclidean space, a function $F$ is $\alpha$-convex if

$$\begin{gathered} F\left((1-t)x_0+t{x}_1\right) \\ \le (1-t)F\left(x_0\right)+tF\left(x_1\right) \\ -\frac{\alpha }{2}t\left(1-t\right){‖x_0-x_1‖}^2. \end{gathered}$$ (100)

Ordinary convexity corresponds to $\alpha =0$. Analogously, a functional $ℱ:{𝒲}_2\to \mathbb{R}$ is $\alpha$-geodesically convex if, along the constant-speed Wasserstein geodesic ${\rho }_t$ [41],

$$\begin{gathered} ℱ\left[{\rho }_t\right]\le (1-t)ℱ\left[{\rho }_0\right]+tℱ\left[{\rho }_1\right] \\ -\frac{\alpha }{2}t\left(1-t\right)W_2^2\left({\rho }_0,{\rho }_1\right). \end{gathered}$$ (101)

We emphasize that the comparison in Eq. (101) is made along geodesic interpolation, rather than along the point-wise mixture of two densities. For $\alpha >0$, under standard regularity assumptions, if ${\rho }_t$ follows the Wasserstein gradient flow of an $\alpha$-geodesically convex $ℱ$, then

$$ℱ\left[{\rho }_t\right]-\text{inf}ℱ\le e^{-2\alpha t}\left[ℱ\left[{\rho }_0\right]-\text{inf}ℱ\right].$$ (102)

assuming $ℱ$ is bounded from below [43]. Eq. (102) states that the functional gap decays exponentially. If a minimizer ${\rho }_{*}$ exists, positive geodesic convexity makes it unique and gives $ℱ\left[\rho \right]-ℱ\left[{\rho }_{*}\right]\ge \alpha W_2^2\left(\rho ,{\rho }_{*}\right)/2$, so the same estimate also implies exponential convergence in $W_2$. In addition to defining convex functionals, we can define convex sets. We say that a set of probability distributions is geodesically convex if it contains the Wasserstein geodesic between any two of its densities. Functionals that are $\alpha$-geodesically convex over ${𝒲}_2$ retain their convexity properties and convergence guarantees when restricted to geodesically convex subsets. In later sections, we will use this framework to describe irreversible thermodynamic relaxation, numerical integration of Fokker–Planck equations, variational inference, and the mean-field dynamics of interacting particles.

III. THERMODYNAMICS

As we elucidate in this section, the previously explored ideas in control and inference have a deep connection to the time-honored subject of thermodynamics. Upon specifying a state-space geometry, statistical inference aims to select probability distributions consistent with partial information, whereas control theory seeks protocols that transport systems between prescribed configurations while minimizing a cost. Thermodynamics simultaneously determines admissible equilibrium macroscopic states and identifies the least costly transformations among them, compatible with the imposed constraints such as fixed temperature, fixed volume, or adiabaticity. Here, we aim to make these connections concrete. In Sec. III A, we discuss the maximum-entropy interpretation of equilibrium statistical mechanics. Subsequently, in Sec. III B we demonstrate that the Onsager least-dissipation principle and its extension to stochastic processes amount to an optimal transport problem.

A. Equilibrium thermodynamics as maximum-entropy inference

Here, we discuss the extent to which statistical mechanics of materials in equilibrium can be interpreted as a maximumentropy inference. As materials are composed of ~ 10²³ degrees of freedom, it is infeasible to keep track of their individual dynamics. Instead, equilibrium statistical mechanics provides a means of describing microscopic particle configurations probabilistically. Using thermodynamic axioms (such as the thermodynamic entropy $S$ reaching a maximum in an isolated system at equilibrium, and the ergodic hypothesis wherein all states are equiprobable within a fixed energy shell), Boltzmann, Gibbs, and Maxwell obtained the expressions for the microstate distribution at equilibrium [46].

Boltzmann argued that by the ergodic hypothesis, all microstates $x$ with the same energy are equally likely in an isolated equilibrium system. He postulated that the microstate multiplicity $W$ compatible with the conserved quantities is maximized at equilibrium (by a reasoning akin to modern saddle-point approaches), and thus sets the thermodynamic entropy via $S=k_{\text{B}}\text{log}W$ (where the logarithm guarantees that entropy is additive). This is called the microcanonical ensemble and can be refined to accommodate additional conserved quantities, like particle number.

Boltzmann then considered a system whose energy fluctuates due to energy exchange with a large thermal bath, constituting the canonical ensemble. Suppose that the system has energy $E(x)$ and the bath $E_{\text{tot}}-E$, so that $E_{\text{tot}}$ is the conserved energy of the joint system (system and bath). The equilibrium probability distribution of the joint configuration $x_{\text{tot}}=\left(x,x_{\text{bath}}\right)$ is uniform, $p^{\text{eq}}\left(x_{\text{tot}}\right)=1/W_{\text{tot}}\left(E_{\text{tot}}\right)$. However, integrating the bath degrees of freedom for a fixed system configuration $x$ yields

$$\begin{gathered} p^{\text{eq}}(x)=\frac{p^{\text{eq}}\left(x,x_{\text{bath}}\right)}{p^{\text{eq}}\left(x_{\text{bath}}\mid x\right)}=\frac{W_{\text{bath}}\left(E_{\text{tot}}-E(x)\right)}{W_{\text{tot}}\left(E_{\text{tot}}\right)} \\ \propto e^{-k_{\text{B}}^{-1}{S}_{\text{bath}}\left(E_{\text{tot}}-E(x)\right)} \end{gathered}$$ (103)

(since $W_{\text{tot}}\left(E_{\text{tot}}\right)=\text{const}$). Therefore, physics naturally equipped us with a probability following a large-deviation principle, as $\text{log}{p}^{\text{eq}}$ is an extensive function. Since the bath is much larger than the subsystem $E(x)\ll E_{\text{tot}}$, we Taylor expand $S_{\text{bath}}\left(E_{\text{tot}}-E(x)\right)\simeq S_{\text{bath}}\left(E_{\text{tot}}\right)-E(x)/T$, where $1/T=\text{d}{S}_{\text{bath}}(E)/\text{d}E$ is the bath temperature. We write $\beta ={\left(k_{\text{B}}T\right)}^{-1}$ for the inverse thermal energy. Thus, the equilibrium distribution of the subsystem microstates $x$ reads:

$$p^{\text{eq}}(x)=\frac{1}{Z}{e}^{-\beta E(x)},$$ (104)

We thus arrive at the Boltzmann distribution from physical principles: energy conservation, the ergodic hypothesis, and a large-deviation principle arising from the thermodynamic limit. In the above, $Z$ is the partition function, and it is related to the free energy via ${ℱ}^{\text{eq}}=-k_{\text{B}}T\text{log}Z$. Similar fluctuation principles upon coupling to different reservoirs give rise to all other statistical-mechanical ensembles. Notably, the entropy across all ensembles can be written as

$$S=-k_{\text{B}}\int \text{d}x{p}^{\text{eq}}\left(x\right)\text{log}{p}^{\text{eq}}(x).$$ (105)

Note that the Boltzmann factor (Eq. (104)) belongs to the exponential family (Eq. (36)), and the thermodynamic entropy (Eq. (105)) takes the form of a Shannon entropy (Eq. (34)). This suggests a link between equilibrium statistical mechanics and maximum-entropy modelling (discussed in Sec. IB 2). Indeed, Jaynes proposed interpreting thermodynamics as enforcing a maximum-entropy inference [13] wherein, for example, the average energy of a system, $ℰ=\int dx{p}^{\text{eq}}(x)E(x)$, when coupled to a thermal bath, acts as a constraint. Upon maximization, the thermodynamic entropy is given by the Shannon entropy of the entropy-maximizing distribution,

$$S=k_{\text{B}}H\left[p^{\text{eq}}\right].$$ (106)

It is important to stress, however, that the conceptual justification behind maximum-entropy modelling and Boltzmann statistics is fundamentally different. The former is motivated by the principle that, given limited information, the least-committal distribution is the one that maximizes entropy, while the latter follows from physical laws governing equilibrium statistical mechanics. The connection between maximum-entropy modelling and statistical mechanics thus arises from the coincidence between the functional forms obtained in the two frameworks. In that regard, while Eq. (106) is conceptually fascinating, there is no indication that either this relation or the information-theoretic interpretation of statistical mechanics should hold away from equilibrium.

While thermodynamic entropy (like other thermodynamic state functions such as temperature [47] and pressure [48]) is an equilibrium property of matter, is it possible to extend thermodynamic notions far from equilibrium? Since Shannon entropy is well-defined for a system out of equilibrium given its instantaneous distribution $p(x)$, can it serve as an extension of entropy outside equilibrium, $S\to k_{\text{B}}H\left[p\right]$, even though this quantity is not generally maximized? The answer lies in whether the emergent relations obtained by presuming a particular entropy functional are consistent with thermodynamic theorems and whether they provide useful predictions, e.g., for the extractable work during thermodynamic transformations.

For example, in certain systems (most often those uncoupled from an external bath), Shannon entropy is explicitly an inadequate measure of thermodynamic entropy [49, 50]. For instance, it is a Lyapunov function of the Liouville equation describing isolated systems, indicating that it does not increase as the system approaches equilibrium (despite the system potentially being ergodic). Therefore, not only may there be no physical reason to invoke maximum-entropy inference in such systems, but Shannon entropy may not be an adequate functional in the first place.

On the other hand, in other cases (most often stochastic systems away from the thermodynamic limit, where the number of particles is small), Shannon entropy provides a useful generalization of thermodynamic entropy out of equilibrium. For example, as we show in Sec. III B 4, it recovers the second law of thermodynamics as well as additional useful fluctuation theorems concerning the extractable work from stochastic systems [51, 52]. Furthermore, returning to Shannon’s construction, Shannon entropy is a functional that adequately quantifies uncertainty [12]. Therefore, in the absence of any knowledge about a problem’s structure beyond a few known expectation values, maximizing entropy under constraints is not an unreasonable inference methodology [14]. Indeed, maximum-entropy inference, both static and dynamical (where $P$ denotes a distribution over trajectories), has proved useful across scales, including predicting protein structure [53], modeling neural communication [54], identifying bacterial swarming transitions [55], and characterizing alignment interactions during flocking [56]. Since thermodynamic variables are often uninformative in these systems, the imposed constraints are no longer temperature, pressure, and related quantities, but rather static and dynamic correlation functions [57, 58], compressed snapshot data [59, 60], transport coefficients [61, 62], and more, which nonetheless carry substantial information about the microstructure of the system. The success of such inference schemes depends not only on the adequacy of the maximum-entropy principle itself, but also on the judicious choice of the moments used to constrain the distribution [55]: moments that capture relevant aspects of the underlying dynamics can yield accurate predictions, whereas constraints on physically irrelevant observables may provide little useful information.

To conclude, there is a deep but subtle connection between equilibrium thermodynamics and information theory. Upon accepting the axioms of thermodynamics, Gibbsian statistical mechanics can be formulated as an inference problem on rigorous grounds. Since entropy is constructed to provide a consistent quantification of statistical uncertainty, exponential-family models can also be useful for inferring the steady-state microstructure of athermal systems, albeit with important limitations. The thermodynamic meaning of these models, however, remains subject to debate.

B. Nonequilibrium thermodynamics as a control problem

In Sec. III A, we recalled the statistical-mechanical characterization of systems at thermodynamic equilibrium, and discussed its information-theoretic interpretation. In this section, we will frame thermodynamic processes close to and far from equilibrium in the language of optimal transport. Thermodynamics provides a concrete notion of efficiency grounded in the second law,

$$T\Sigma =T\Delta S-Q=W-\Delta ℱ\ge 0,$$ (107)

where $W$ is the mechanical work along a process and $\Delta ℱ$ is the resulting change in free energy $ℱ$. Here $\Sigma$ is the entropy production, so that $T\Sigma$ is the dissipated heat—the energy irreversibly lost during the transformation. (Note that, confusingly, $\Sigma$ is often called entropy production; more often than not, entropy production $\Sigma$ and change in entropy $\Delta S$ are different objects.) In particular, an efficient process between an initial and final state is one which minimizes the entropy production $\Sigma$, in turn allowing the extractable work $-W$ to reach the thermodynamic upper bound $-\Delta ℱ$.

In this section, we will discuss useful approaches for estimating the extractable work $W$ and different computational strategies to minimize entropy production $\Sigma$ using optimal-transport and path-integral methods. We begin by studying thermodynamic relaxation for small departures from equilibrium, described by the Onsager least-dissipation theory. We then extend these ideas to stochastic processes far from the thermodynamic limit (few degrees of freedom instead of $N_{\text{p}}~{10}^{23}$ particles), for which we will derive useful thermodynamics-inspired bounds, a geometric interpretation for minimizing the entropy production, and nontrivial equalities for the extractable-work fluctuations.

1. Optimal transport bounds efficiency via Onsager’s least-dissipation principle

Let’s start with a modest question—what are the dynamics when we depart from equilibrium but remain close to it? In this setting, Onsager formulated an approach that is phenomenological in nature [63, 64], which has since permitted the calculation of bounds on dissipation in terms of the Wasserstein metric. The connection to Wasserstein geometry ultimately occurs in the context of field theories, but Onsager’s governing principles are perhaps easiest to state in a discrete setting. In brief, consider an isolated system described by the generalized variables ${\left\{X_m\right\}}_{m=1,\dots ,M}$ — for instance, concentrations of various compounds in a chemical mixture—that determine the system free energy, $ℱ\left(X_1,\dots ,X_M\right)$. Onsager postulated that, around equilibrium, the system responds linearly to gradients in the free energy:

$$\frac{\text{d}{X}_m}{\text{d}t}=-\sum _{n=1}^M{L}_{mn}\frac{\partial ℱ}{\partial X_n}.$$ (108)

where $L_{mn}$ is the Onsager matrix of phenomenological transport coefficients. Physically, the coefficients $L_{mn}$, which can be measured experimentally as well as predicted theoretically in some cases [65], describe how changing one variable affects all other variables in response. For instance, in the thermoelectric effect, the corresponding $L_{mn}$ is known as the thermoelectric coefficient and describes how changes in temperature drive changes in charge density. Eq. (108) serves as the basis for modern linear-irreversible thermodynamics [66], wherein a system relaxes to equilibrium linearly in the thermodynamic forces $\left(\partial ℱ/\partial X_m\right)$, but the forces may be arbitrarily nonlinear in the $X$’s.

It is worth noting that Onsager proved various nontrivial properties of $L_{mn}$ [64]. First, to satisfy the second law of thermodynamics (the nonnegativity of heat dissipation, Eq. (107)), $L_{mn}$ must be positive-definite. Namely, close to equilibrium, very generally [63] the dissipation functional is given by the thermodynamic dissipative force $-\sum _m{\left(L^{-1}\right)}_{nm}{\dot{X}}_m$ (friction coefficient times velocity) along a displacement ${\dot{X}}_n$,

$$T\dot{\Sigma }=\sum _{n,m}{\left(L^{-1}\right)}_{nm}{\dot{X}}_n{\dot{X}}_m,$$ (109)

which therefore requires a positive-definite $L_{mn}$ to satisfy $\dot{\Sigma }>0$. Moreover, the so-called Onsager reciprocity theorem states that $L_{mn}$ must be symmetric, $L_{mn}=L_{nm}$, when the microscopic equations of motion are time-reversal symmetric.

Useful for our purposes is that, for systems obeying Onsager reciprocity [64], Eq. (108) can be obtained from a variational principle. Namely, near equilibrium, the velocities $\left\{{\dot{X}}_m\right\}$ extremize the following function:

$$\begin{gathered} ℛ\left({\dot{X}}_1,\dots ,{\dot{X}}_M\right)=\frac{1}{2}T\dot{\Sigma }+\frac{\text{d}ℱ}{\text{d}t} \\ =\frac{1}{2}T\dot{\Sigma }+\sum _m\frac{\partial ℱ}{\partial X_m}{\dot{X}}_m \end{gathered}$$ (110)

where $ℛ$ is known as the Rayleighian, and in the second equality we used the chain rule for the free energy. Solving $\partial ℛ/\partial {\dot{X}}_m=0$ is equivalent to Eq. (108) with the heat dissipation of Eq. (109). In this formulation, it is natural to think of gradients of $ℱ$ as driving forces, and $L^{-1}$ as yielding a metric that defines steepest descent.

Onsager’s theory can be connected with optimal transport when Eq. (110) is extended to field theories. We follow the lines of Ref. [67]. In a continuum (hydrodynamic) approach, one models a fluid via a collection of conserved density fields (e.g., mass, momentum, or energy density fields [68]). For simplicity, we will concentrate on a single scalar field $n_t(x)$ representing the density of a conserved quantity, such as particle density. We denote the conserved particle number by $N_{\text{p}}=\int n_t(x)\text{d}x$; unlike a probability density, $n_t$ need not integrate to one. Here $n_t(x)$ plays the role of $X_m(t)$ in Eqs. (108)–(110), and the spatial coordinate $x$ plays the role of the index $m$. The field-theoretical generalization of Eq. (110) can be used as a prescription to derive the equations of motion for $n_t(x)$. First, because the field is conserved, it evolves according to the continuum equation

$$\frac{\partial n_t(x)}{\partial t}+\nabla \cdot \left[n_t(x)v_t(x)\right]=0,$$ (111)

where $v$ is the velocity field. Next, we assume that the free energy may be expressed as a local functional of $n$, namely $ℱ\left[n\right]$. Finally, we consider phenomenologically that the dissipation $\dot{\Sigma }$ is quadratic in $v_t(x)$ as in Eq. (109),

$$T\dot{\Sigma }=\int \text{d}x{n}_t(x)\Gamma \left(n_t(x)\right){\left|v_t(x)\right|}^2.$$ (112)

where the phenomenological dissipation coefficient $\Gamma (n)$ may be an arbitrary function of density $n$.

As a concrete example, one can have in mind a system of very dilute colloidal particles in a viscous fluid. In this case $n_t(x)$ represents the number density of the particles; neglecting hydrodynamic interactions, the drag is given by the Stokes drag $\Gamma =6\pi \eta r$ (where $r$ is the radius of the suspended particles, and $\eta$ is the viscosity); and, assuming the particles are non-interacting, the free energy is the negative of the entropy $ℱ=-TS=k_{\text{B}}T\int n_t(x)\text{log}{n}_t(x)\text{d}x$.

Upon extending the Onsager Rayleighian Eq. (110) to fields, we obtain

$$\begin{gathered} ℛ=\frac{1}{2}T\dot{\Sigma }+\frac{\text{d}ℱ}{\text{d}t} \\ =\frac{1}{2}T\dot{\Sigma }+\int \text{d}x{n}_t(x)v_t(x)\cdot \nabla \frac{\delta ℱ}{\delta n_t(x)}. \end{gathered}$$ (113)

with the entropy production $\dot{\Sigma }$ of Eq. (112), and we used the functional chain rule and Eq. (111) for evaluating the change in free energy. According to the Onsager variational principle, the equation of motion for $n_t(x)$ is determined by choosing the flow $v_t(x)$ that minimizes $ℛ$. Taking $\delta ℛ/\delta v=0$ yields

$$\Gamma \left(n_t(x)\right)v_t(x)=-\nabla \frac{\delta ℱ}{\delta n_t(x)}.$$ (114)

We write ${\mu }_t(x)=\delta ℱ/\delta n_t(x)$ for the local chemical potential. Eq. (114), when combined with Eq. (111), yields a closed dynamical equation for $n$:

$${\partial }_{t}n=\nabla \cdot \left(\frac{n}{\Gamma (n)}\nabla \frac{\delta ℱ}{\delta n}\right)$$ (115)

Returning to our example of a dilute colloidal suspension, one obtains

$$\frac{\partial n_t(x)}{\partial t}=D{\nabla }^2{n}_t(x),$$ (116)

where $D=k_{\text{B}}T/(6\pi \eta r)$ is the Stokes-Einstein diffusion constant. Indeed, many of the known dynamical theories in soft matter, ranging from diffusion to phase separation and viscoelasticity, can be derived from Onsager’s variational principle [67].

Onsager’s least-dissipation theorem is closely related to optimal transport—as we show, the dissipation associated with the WGF serves as a lower bound on the true dissipated heat. Note the close resemblance of Eq. (115) to Eq. (98). The Onsager Rayleighian, Eq. (113), can be compared per particle with probability transport. Defining ${\rho }_t(x)=n_t(x)/N_{\text{p}}$ gives

$$\begin{gathered} \frac{2}{\Gamma N_{\text{p}}}{\int }_0^1\text{d}t{ℛ}_t={\int }_0^1\text{d}t\int \text{d}x{\rho }_t(x){‖v_t(x)‖}^2 \\ +{\int }_0^1\text{d}t\int \text{d}x{\rho }_t\left(x\right)\left[\frac{2}{\Gamma }\nabla {\mu }_t\left(x\right)\right]\cdot v_t\left(x\right). \end{gathered}$$ (117)

where, for simplicity, we have taken $\Gamma =\text{const}$. At the same time, return to Eq. (91). Using the method of Lagrange multipliers, the minimum of the 2-Wasserstein distance subject to the conservation law and initial and final distributions is found from the saddle point of the following Lagrangian,

$$\begin{gathered} {\int }_0^1\text{d}t\int \text{d}x{\rho }_t(x){‖v_t(x)‖}^2 \\ +{\int }_0^1\text{d}t\int \text{d}x{\rho }_t(x)\left[2\nabla {\chi }_t(x)\right]\cdot v_t(x) \\ +(\text{terms}\text{independent}\text{of}v). \end{gathered}$$ (118)

Here ${\chi }_t(x)$ is the transport potential of Chapter II, introduced by adding $-2{\chi }_t\left[{\partial }_t{\rho }_t+\nabla \cdot \left({\rho }_t{v}_t\right)\right]$ to the integrand before integrating by parts. Variation with respect to $v_t$ gives $v_t=-\nabla {\chi }_t$, where terms independent of $v_t$ have been ignored.

Eqs. (117) and (118) have the same quadratic velocity term. The former describes the particle density $n_t=N_{\text{p}}{\rho }_t$, with dissipation expressed per particle, whereas the latter is the probability transport variational problem. Thermodynamics fixes the potential to ${\chi }_t={\mu }_t/\Gamma$, where ${\mu }_t=\delta ℱ/\delta n_t$ is the chemical potential, instead of optimizing over ${\chi }_t$. Thus, the Benamou-Brenier solution is a more efficient transportation mechanism;³ linear-irreversible thermodynamics is a Benamou-Brenier-type “partially” optimized transport, as the chemical potential is specified. We will return to the farreaching consequence of this realization—a bound, tighter than Eq. (107), for finite-time transformations—shortly. But until then—may a similar conclusion apply far from equilibrium?

2. Dynamics far from equilibrium

In constructing Onsager’s least-dissipation principle, we considered a coarse-grained hydrodynamic theory close to equilibrium. Yet, the same mathematical structures arise in many other types of dynamics, notably in stochastic processes far from equilibrium. In the following, we will be working within the framework of stochastic thermodynamics [51, 52], aiming to extend thermodynamic notions and theorems to fluctuation-dominated microscopic systems, such as driven colloidal suspensions and individual motor proteins. Since we are no longer in the thermodynamic limit, it is not obvious whether state functions such as temperature or pressure carry over to the microscale or whether the (fluctuating) mechanical work satisfies any relations akin to the second law (Eq. (107)). The immense success of stochastic thermodynamics, summarized in part here and in the coming sections, is the fact that many functions, including mechanical work, can indeed be used still; e.g., the fluctuating work is just any externally applied forces along a stochastic displacement. Upon selecting the surrounding bath properties (temperature, pressure, etc.) and adopting the Shannon entropy as the microscopic analogue of a nonequilibrium thermodynamic entropy, the second law of thermodynamics and many additional useful new fluctuation relations can be derived. The aim of this section is to explore the correspondence between the thermodynamics of stochastic systems and optimal control.

To demonstrate this point, we will consider the paradigmatic example of an overdamped colloid submerged in a thermal bath, which an agent drives via an external force $F_t(x)$,

$${\dot{x}}_t=D\beta F_t\left(x_t\right)+\sqrt{2D}{\eta }_t,$$ (119)

where $D$ is the particle diffusivity and ${\eta }_t$ is standard white noise. Thus, $\epsilon =2D$, the stochastic drift is $f_t(x)=D\beta F_t(x)$, and $F_t$ denotes a physical force. The time-dependent probability distribution to obtain a microstate, $p_t(x)$, satisfies the Fokker-Planck equation [69]—the continuity equation with the current velocity

$$v_t\left(x\right)=D\beta F_t\left(x\right)-D\nabla \text{log}{p}_t(x).$$ (120)

As preparation for later, we distinguish between two types of steady states: One is an equilibrium state (as in Sec. III A), where the external driving force is constant in time and conservative, $F=-\nabla E$, yielding the Boltzmann factor as the equilibrium distribution, Eq. (104), and a zero current, $v\to 0$. The other is a nonequilibrium steady state, where the current need not be zero ($v\ne 0$), but the flux is divergence-free $\nabla \cdot (pv)\to 0$ so $\partial p/\partial t=0$.

Since we may not be close to equilibrium, we may not rely on Onsager’s dissipation functional, and instead refer to the machinery of stochastic thermodynamics [52, 70]. Given our precaution in Sec. III A, we note that the Shannon entropy (Eq. (34)) is indeed an adequate choice for the thermodynamic entropy out of equilibrium in stochastic processes, $S\left[p_t\right]=-k_{\text{B}}\int \text{d}x{p}_t(x)\text{log}{p}_t(x)$, since it will provide insightful and thermodynamically consistent results shortly. The change in entropy over time is

$$\frac{\text{d}S}{\text{d}t}=-k_{\text{B}}\int \text{d}x{v}_t(x)\cdot \nabla p_t(x).$$ (121)

According to the first law of thermodynamics, in this setting, the heat exchanged with the bath is minus the external work — the average of the externally applied forces along a displacement,

$$\dot{Q}=-\dot{W}=-\int \text{d}x{p}_t(x)v_t(x)\cdot F_t(x).$$ (122)

Thus, the entropy production is given by

$${\dot{\Sigma }}_t=\frac{\text{d}S}{\text{d}t}-\frac{\dot{Q}}{T}=\frac{k_{\text{B}}}{D}\int \text{d}x{p}_t(x){\left|v_t(x)\right|}^2.$$ (123)

Indeed, it is zero at equilibrium $(v=0)$, otherwise positive, and functionally consistent with Onsager’s Eq. (112).

Upon further comparison with the Benamou-Brenier approach, Eq. (91), we obtain an optimal-transport-inspired bound on the entropy production [71]. The mere fact that one started and terminated the process at particular distributions during a finite time $\tau$ means that one must have produced entropy which is bounded by the Wasserstein metric between the distributions

$${\int }_0^{\tau }\text{d}t{\dot{\Sigma }}_t\ge \frac{k_{\text{B}}}{D\tau }{W}_2^2\left(p_0,p_{\tau }\right).$$ (124)

Since $W_2^2$ is itself nonnegative, Eq. (124) is a nontrivial bound on the total dissipated heat $T\Sigma$ that is tighter than the second law of thermodynamics, Eq. (107).

The minimization of Eq. (118) necessitates a gradient current velocity, $v_t=-\nabla {\chi }_t$. Therefore, unless $F_t$ is conservative ($F_t=k_{\text{B}}T\nabla \left[\text{log}{p}_t-{\chi }_t/D\right]$), there is no chance to saturate the inequality in Eq. (124), let alone approach $T\dot{\Sigma }=0$. Indeed, unless $F_t$ is given by the optimal value of ${\chi }_t$, even the Onsager least-dissipation principle will dissipate more heat than that prescribed by the Wasserstein-metric bound. To illustrate a related consequence of the above thermodynamic interpretations of Wasserstein geometry and its connection to the kinetics of stochastic processes, we outline in Sec. V B 1 the Jordan-Kinderlehrer-Otto scheme [45] in the language of optimal transport.

Interestingly, since $W_2^2$ is a property of the two distributions (but not time; cf. Eq. (89)), the bound of Eq. (124) is inversely proportional to the elapsed time. Indeed, the second law only saturates for quasistatic transformations; otherwise, Eq. (124) provides a lower bound on the minimal heat that must have been dissipated to enable a finite-time process. Conversely, Eq. (124) can be interpreted as a thermodynamic speed limit on how quickly one may transform from $p_0(x)$ to $p_{\tau }(x)$, giving rise to the lower bound $\tau \ge k_{\text{B}}{W}_2^2/(D\Sigma )$.

To conclude, the Benamou-Brenier optimal transport problem sets a bound on the thermodynamic cost of very general, finite-time, nonequilibrium transformations [71]. There are additional, fascinating extensions, stemming from combining optimal transport with stochastic mesoscopic systems. For instance, the above statements were shown to apply to discrete-state Markov systems as well [72]. Upon introducing feedback via measurements, the dissipated heat bound can be made tighter in terms of the acquired information [73]. It is important to stress that although the above are nontrivial connections for far-from-equilibrium systems, these connections only apply in cases where the Einstein relation is valid [71, 72]. Otherwise, the Wasserstein metric only yields a bound on the entropy production in athermal and active systems (which is still given by Eq. (123) [51, 74]), but not on the dissipated heat, as the connection between the latter two breaks in the absence of a scalar temperature [75].

Is it possible to saturate the inequalities provided by the second law of thermodynamics? We present here two strategies—one involving the expression of entropy production as a thermodynamic metric in protocol space [76, 77], and the other involving fluctuation properties of stochastic processes [78, 79].

3. Geometric interpretation of entropy production

In the previous sections we found nontrivial bounds on the entropy production using optimal transport. Here, our aim is more practical—is it possible to engineer efficient thermodynamic transformations? By efficient we mean that the entropy production $\Sigma$ is minimal so, according to Eq. (107), the extractable work is closest to the upper bound dictated by free energy. More concretely, while we did not specify $F_t(x)$ in Sec. III B 2, here we consider specifically a conservative force, $F_t(x)=-\nabla E_t(x)$. We furthermore note that the time dependence of the potential, in practice, arises from a “control knob” that the experimentalist varies over time (e.g., increasing the electric power of an optical trap, thereby increasing the radiation-pressure force on a colloid). These knobs are often called protocols, which we denote by $u_t$. Finding the most efficient thermodynamic transformation, therefore, entails optimizing the entropy production over the protocols $u_t$. For the purpose of this section, the protocol $u_t$ can be thought of as the controls in Sec. IA.

In trying to find the optimization framework for the entropy production, we mention several noteworthy properties of entropy production, Eq. (123): First, clearly the quasi-static transformations are most efficient—since, close to equilibrium $v(x)\to 0$ everywhere in $x$, the dissipation $\Sigma \to 0$, necessarily from above due to the quadratic form. Even more so, we see that work extraction in the quasi-static limit is indeed possible. Qualitatively, the slower the time scale over which the process occurs, $\tau \to \infty$, then $v_t(x)~{\tau }^{-1}$. Since changes in free energy, entropy (Eq. (121)), work, and heat (Eq. (122)) are proportional to $v$ but involve the integration over time, order-${\int }_0^{\tau }\text{d}tv~1$ energy extraction is possible while the dissipation is $~{\int }_0^{\tau }\text{d}t{v}^2~1/\tau \to 0$. Are all quasi-static transformations the same, or are some more efficient than others? For instance, due to critical slowing down, it may be counterproductive to move throughout Ising-model phase space through the critical point [77].

We restrict our attention to exploring the space of quasistatic transformations, as their entropy production is smallest. Since $v_t(x)$ is an involved function, requiring the knowledge of the instantaneous distribution in addition to the force, we seek to make the connection of quasistaticity to the protocol $u_t$ more concrete, as follows. For the process to be quasistatic, $p_t(x)$ must equilibrate on a faster timescale than the protocol evolves; namely all eigenmodes of the Fokker-Planck operator converge to equilibrium faster than the protocol varies. We first define the deviation $\delta p_t(x)$ of the instantaneous distribution $p_t(x)$ from the equilibrium Boltzmann factor corresponding to the instantaneous protocol value $p_{u_t}^{\text{eq}}(x)$:

$$\begin{gathered} p_t(x)≔p_{u_t}^{\text{eq}}\left(x\right)+\delta p_t\left(x\right), \\ {p}_u^{\text{eq}}(x)=\frac{1}{Z_u}\text{exp}\left[-\frac{E_u\left(x\right)}{k_{\text{B}}T}\right], \\ {Z}_u=\int \text{d}x\text{exp}\left[-\frac{E_u\left(x\right)}{k_{\text{B}}T}\right]. \end{gathered}$$ (125)

By inserting this solution into the Fokker-Planck equation,

$$\begin{gathered} \frac{\partial p_t(x)}{\partial t}=-{ℒ}_{u_t}{p}_t\left(x\right), \\ {ℒ}_u=-D\nabla \cdot \left[\frac{\nabla E_u\left(x\right)}{k_{\text{B}}T}+\nabla \right], \end{gathered}$$ (126)

where ${ℒ}_u{p}_u^{\text{eq}}=0$, we find the exact expression

$$\delta p_t(x)=-{\int }_{-\infty }^t\text{d}s{e}_{\text{T}}^{-{\int }_s^t\text{d}{s}^{\prime }{ℒ}_{u_{s^{\prime }}}}{\dot{u}}_s\cdot \frac{\partial p_{u_s}^{\text{eq}}(x)}{\partial u},$$ (127)

where $e_{\text{T}}$ is a time-ordered exponential. We will now use the assumption of a quasistatic protocol: Since all modes of ${ℒ}_u$ converge prior to $u_t$ changing considerably, $e_{\text{T}}^{-{\int }_s^t\text{d}{s}^{\prime }{ℒ}_{u_{s^{\prime }}}}$ suppresses integrand contributions in Eq. (127) where $u_s$ departed much from $u_t$. Thus, we approximately replace all protocol values $u_s$ with $u_t$,

$$\delta p_t(x)\simeq -{\int }_{-\infty }^t\text{d}s{e}^{-(t-s){ℒ}_{u_t}}{\dot{u}}_t\cdot \frac{\partial p_{u_t}^{\text{eq}}(x)}{\partial u}.$$ (128)

In preparation for later, we define the protocol score (a vector for a multicomponent protocol):

$$\begin{gathered} {\Theta }_u(x)≔\frac{1}{p_u^{\text{eq}}(x)}\frac{\partial p_u^{\text{eq}}(x)}{\partial u} \\ =-\frac{1}{k_{\text{B}}T}\left(\frac{\partial E_u(x)}{\partial u}-{⟨\frac{\partial E_u(x)}{\partial u}⟩}_u^{\text{eq}}\right) \end{gathered}$$ (129)

where $⟨a(x){⟩}_u^{\text{eq}}=\int \text{d}x{p}_u^{\text{eq}}(x)a(x)$.

Recall that the work is, by definition, the energy involved in shifting the potential in time [78],

$$W={\int }_0^1\text{d}t\int \text{d}x{p}_t(x){\dot{u}}_t\cdot \frac{\partial E_{u_t}(x)}{\partial u}.$$ (130)

and the change in free energy is $\Delta {ℱ}^{\text{eq}}=-k_{\text{B}}T\left[\text{log}{Z}_{u_1}-\text{log}{Z}_{u_0}\right]$. Using Eq. (107), we find the quasi-static expression for the entropy production:

$$\begin{gathered} \Sigma =\frac{1}{T}{\int }_0^1\text{d}t\int \text{d}x\delta p_t(x){\dot{u}}_t\cdot \frac{\partial E_{u_t}(x)}{\partial u} \\ =k_{\text{B}}{\int }_0^1\text{d}t{\dot{u}}_t\cdot {\zeta }_{u_t}\cdot {\dot{u}}_t, \end{gathered}$$ (131)

where the friction coefficient can be expressed as the integrated autocorrelation function of the thermodynamic force

$$\begin{gathered} {\zeta }_u={\int }_0^{\infty }\text{d}s\int \text{d}x{p}_u^{\text{eq}}(x){\Theta }_u(x)e^{-s{ℒ}_u^{†}}{\Theta }_u^{\top }(x) \\ ={\int }_0^{\infty }\text{d}s{⟨{\Theta }_u\left(x_s\right){\Theta }_u^{\top }\left(x_0\right)⟩}_u^{\text{eq}}, \end{gathered}$$ (132)

where the average is in the equilibrium problem, Eq. (119), where the protocol is fixed (with $x_0~p_u^{\text{eq}}$).

Thus, Eq. (131) is an exact expression for the entropy production in terms of a Riemannian geometry in the protocol space [76]. (This contrasts with Eq. (124), which provides an optimal-transport-inspired bound on the entropy production.) The thermodynamic metric, ${\zeta }_u$, describes the “friction” one would experience in a given position in phase space $u$. We now concretely see how entropy production decreases quadratically with slower driving $\dot{u}$. However, we may now further ask questions such as how to optimally transition between states quasistatically, e.g., how to vary temperature and magnetic field to flip the magnetization in the zero-temperature Ising model [77]. These trajectories would correspond to finding the geodesics on the curved surface described by ${\zeta }_u$, e.g., using geometric minimum-action methods [80] or by backpropagation [81]. Eq. (131) then allows us to investigate the efficiency and geometry of near-equilibrium thermodynamic transformations.

4. Thermodynamic fluctuation theorems

While the above has been a detailed, geometric treatise on close-to-equilibrium thermodynamics, one can find additional, surprising statistical characterizations of the work distributions far from equilibrium. Eq. (130) provides the mean work, and in the previous sections we derived second-law-like bounds $T\Sigma =W-\Delta {ℱ}^{\text{eq}}\ge 0$. Here we seek additional statistical properties of the work along a stochastic trajectory $\left\{x_t\right\}$,

$$w_t\left[x\left(\cdot \right)\right]={\int }_0^t\text{d}s{\dot{u}}_s\cdot \frac{\partial E_{u_s}\left(x_s\right)}{\partial u},W={⟨w_1⟩}_P,$$ (133)

where $P$ is the forward path distribution and $w_1$ is the work over the complete unit-time protocol. The trajectories start and end at equilibria ($x_0~p_{u_0}^{\text{eq}}$ and $x_1~p_{u_1}^{\text{eq}}$). For that purpose, we will be computing its moment-generating function (MGF), $⟨e^{\lambda w_1}⟩$. We will show that it satisfies two useful equalities: the Jarzynski equality [78, 82],

$$⟨e^{-w_1/\left(k_{\text{B}}T\right)}⟩=e^{-\Delta {ℱ}^{\text{eq}}/\left(k_{\text{B}}T\right)},$$ (134)

and the Crooks fluctuation theorem [79].

$$\frac{p\left(w_1\right)}{\bar{p}\left(-w_1\right)}=e^{\left(w_1-\Delta {ℱ}^{\text{eq}}\right)/\left(k_{\text{B}}T\right)},$$ (135)

where $\bar{P}$ is the reverse path distribution and $\bar{p}$ is its distribution of work along a trajectory that is being orchestrated by the reversed protocol, ${\bar{u}}_t≔u_{1-t}$, which starts from the corresponding (reversed) equilibrium initial condition, $p_{{\bar{u}}_0}^{\text{eq}}=p_{u_1}^{\text{eq}}$.

The Jarzynski equality has been proven using the Feynman-Kac formula in Ref. [83]. Inspired by their approach, we present a new method for deriving the Crooks fluctuation theorem, involving a Feynman-Kac calculation of the forward- and reverse-process MGFs. While Eq. (134) vividly involves the work MGF at $\lambda =-1/\left(k_{\text{B}}T\right)$, we convert Eq. (135) to involve the MGFs by multiplying it by $\bar{p}\left(-w_1\right)e^{\lambda w_1}$ and integrating over $w_1$, to find:

$${⟨e^{\lambda w_1}⟩}_P=e^{-\Delta {ℱ}^{\text{eq}}/\left(k_{\text{B}}T\right)}{⟨e^{-\left[\lambda +1/\left(k_{\text{B}}T\right)\right]w_1}⟩}_{\bar{P}}.$$ (136)

If the Crooks fluctuation theorem holds, Eq. (136) readily reproduces the Jarzynski relation, Eq. (134), by inserting $\lambda =-1/\left(k_{\text{B}}T\right)$. Therefore, we aim to derive the Crooks fluctuation theorem, involving both the forward- and backward-running controls, $u_t$ and ${\bar{u}}_t=u_{1-t}$.

Define the work-weighted endpoint density (similar to a moment-generating function) to appear at $x$ at time $t$ along the $u_t$-controlled trajectories:

$$M_t(x,\lambda )={⟨\text{exp}\left[\lambda {\int }_0^t\text{d}s{\dot{u}}_s\cdot \frac{\partial E_{u_s}\left(x_s\right)}{\partial u}\right]\delta \left(x-x_t\right)⟩}_P.$$ (137)

Using Itô’s lemma [69], one can show that it follows the initial-value partial-differential equation:

$$\begin{gathered} \frac{\partial M_t(x,\lambda )}{\partial t}=-{ℒ}_{u_t}{M}_t(x,\lambda )+\lambda {\dot{u}}_t\cdot \frac{\partial E_{u_t}(x)}{\partial u}{M}_t(x,\lambda ), \\ {M}_0(x,\lambda )=p_{u_0}^{\text{eq}}(x). \end{gathered}$$ (138)

We define the initially conditioned MGF along the reversed-protocol, ${\bar{u}}_t\equiv u_{1-t}$-controlled trajectories:

$${\bar{M}}_t(x,\lambda )={⟨\left.\text{exp}\left[\lambda {\int }_t^1\text{d}s{\dot{\bar{u}}}_s\cdot \frac{\partial E_{{\bar{u}}_s}\left(x_s\right)}{\partial \bar{u}}\right]\right|x_t=x⟩}_{\bar{P}}.$$ (139)

Using the Feynman-Kac formula, it satisfies the terminal-value partial-differential equation:

$$\begin{gathered} \frac{\partial {\bar{M}}_t(x,\lambda )}{\partial t}={ℒ}_{{\bar{u}}_t}^{†}{\bar{M}}_t\left(x,\lambda \right)-\lambda {\dot{\bar{u}}}_t\cdot \frac{\partial E_{{\bar{u}}_t}\left(x\right)}{\partial \bar{u}}{\bar{M}}_t\left(x,\lambda \right), \\ {\bar{M}}_1(x,\lambda )=1. \end{gathered}$$ (140)

In order to prove the Jarzynski and Crooks fluctuation theorems, we should find the connection between $M$ and $\bar{M}$, which will involve a time-reversal of the latter.

A useful, intermediate MGF we define is ${\tilde{M}}_t(x,\lambda )≔{\bar{M}}_{1-t}(x,-\lambda -\beta )$. Based on Eq. (140), one can show that $\tilde{M}$ satisfies the following equation:

$$\begin{gathered} \frac{\partial {\tilde{M}}_t(x,\lambda )}{\partial t}=-{ℒ}_{u_t}^{†}{\tilde{M}}_t(x,\lambda ) \\ +\left(\lambda +\frac{1}{k_{\text{B}}T}\right){\dot{u}}_t\cdot \frac{\partial E_{u_t}(x)}{\partial u}{\tilde{M}}_t(x,\lambda ), \end{gathered}$$ (141)

where the forward-time protocol has been restored, ${\bar{u}}_{1-t}=u_t$, and $\text{d}/\text{d}(1-t)=-\text{d}/\text{d}t$. Note the property $e^{-E_u(x)/\left(k_{\text{B}}T\right)}{ℒ}_u^{†}{e}^{E_u(x)/\left(k_{\text{B}}T\right)}={ℒ}_u$. With this in mind, we find that

$$M_t(x,\lambda )=\frac{1}{Z_{u_0}}{e}^{-E_{u_t}(x)/\left(k_{\text{B}}T\right)}{\tilde{M}}_t(x,\lambda ),$$ (142)

as by direct substitution, we see that the right-hand side of Eq. (142) indeed satisfies Eq. (138) along with its initial condition. Upon integrating Eq. (142) over $x$, evaluated at $t=1$, and recalling the definitions of the forward and backward MGFs, Eqs. (137) and (139), we recover the Crooks fluctuation theorem, Eq. (136),

$$\begin{gathered} {⟨e^{\lambda w_1}⟩}_P=\int \text{d}x{M}_1(x,\lambda ) \\ =\frac{Z_{u_1}}{Z_{u_0}}\int \text{d}x{p}_{u_1}^{\text{eq}}(x){\bar{M}}_0(x,-\lambda -\beta ) \\ =\frac{Z_{u_1}}{Z_{u_0}}{⟨e^{-\left[\lambda +1/\left(k_{\text{B}}T\right)\right]w_1}⟩}_{\bar{P}}. \end{gathered}$$ (143)

Upon replacing the partition functions with the free energy, $Z_u=e^{-{ℱ}_u^{\text{eq}}/\left(k_{\text{B}}T\right)}$, and inverting the Laplace transform, we find the original Crooks fluctuation theorem, Eq. (135).

It is remarkable that for any equilibrium-to-equilibrium transformation, independently of how fast or sub-optimal the controls were, how strong the stochasticity is in the microscopic system, and how far the system has been out of equilibrium during the transformation, equalities such as Eqs. (134) and (135) hold without any approximation. As a corollary, by the Jensen inequality, they reproduce the second law of thermodynamics, ${ℱ}_{u_1}^{\text{eq}}-{ℱ}_{u_0}^{\text{eq}}\le ⟨w_1⟩$, which is an inequality for the average work involved in a protocol-driven transformation of the stochastic system. Previous sections involved near-equilibrium expansions or optimal-transport-inspired bounds. Fluctuation relations such as these are routinely used to estimate free energy differences in microscopic transformations [83, 84]. In Sec. IV B, we show that the Jarzynski relation can also have utility for importance sampling.

IV. SAMPLING AND CONTROL

Sampling from complex, high-dimensional probability distributions is a fundamental task in scientific computing, statistics, and machine learning. The sampling problem asks for the production of independent samples distributed according to a target probability distribution $p_1$. In statistical mechanics, for instance, samples from the Boltzmann distribution allow estimation of physical observables, while in Bayesian inference, samples from the posterior distribution allow predictions based on observed data. This chapter describes the close connection between sampling, inference, control theory, and non-equilibrium thermodynamics, which forms the basis for modern machine-learning methods in reinforcement learning (Sec. V A) and generative modeling (Sec. V C). These links arise from two directions. On the one hand, sampling provides a general-purpose method for solving inference and control problems. On the other hand, control theory and non-equilibrium thermodynamics give rise to practical sampling algorithms. Indeed, such algorithms typically use a controlled, dynamic process that transforms samples from a simple base distribution $p_0$ into the more complex target $p_1$. Sampling thus becomes a control and transport problem of “guiding” base samples into regions of high target probability.

Sec. IV A begins by briefly reviewing the use cases and challenges of sampling in high dimension. Many classic algorithms generate samples using equilibrium physics. Sec. IV B introduces Annealed Importance Sampling (AIS), which allows harnessing non-equilibrium processes to generate samples, and can greatly outperform classical methods. AIS is closely connected to non-equilibrium thermodynamics (Sec. III B). Finally, Sec. IV C casts sampling as a control problem by optimizing the statistical efficiency of the sample-generation process. Vice versa, control problems can be solved by sampling trajectories and keeping count of the resulting reward (Sec. I D 2).

A. What is sampling and why is it hard?

We begin by fixing the setting. Our goal is to draw $N$ independent samples $x_n~p_1,n=1,\dots ,N$ from a complex target distribution $p_1$. We assume that we can evaluate an unnormalized density ${\tilde{p}}_1(x)$ for the target, with $p_1(x)={\tilde{p}}_1(x)/Z_1$. Throughout this chapter, $p$ denotes a normalized probability density and a tilde denotes its unnormalized form. Below, we use dimensionless “energies” and write $E_i(x)=-\text{log}{\tilde{p}}_i(x)$, so $p_i=e^{-E_i}/Z_i$ and $-\text{log}{p}_i=E_i+\text{log}{Z}_i$. As we explain next, the sampling task naturally arises in statistical mechanics and Bayesian inference.

1. The Monte Carlo method

A key problem in equilibrium statistical mechanics is to calculate expectation values of physical observables, denoted $a(x)$ for “average”, for example, the magnetization of a piece of metal [7]. This requires summing (or integrating) over all $d$-dimensional configurations $x$, weighted by the Boltzmann measure:

$$\begin{gathered} ⟨a⟩=\int a\left(x\right)p_1\left(x\right){\text{d}}^{d}x, \\ {p}_1(x)=\text{exp}(-\beta E(x))/Z \end{gathered}$$ (144)

where $\beta ={\left(k_{\text{B}}T\right)}^{-1},E$ is the energy, and $Z$ is the normalization constant. Computing this integral in large dimensions requires summing a number of configurations exponential in $d$ (the number of system degrees of freedom). This is generally intractable. A similar problem occurs in Bayesian inference [8], already encountered in Sec. I B. Here, the target distribution $p_1$ is the posterior of a statistical model. By Bayes’ rule, the posterior for parameters or hidden variables $x$, conditioned on observations $A$, reads:

$$\begin{gathered} p(x\mid A)=p(A\mid x)p(x)/p(A) \\ =\text{exp}\left(\text{log}p(A\mid x)+\text{log}{p}^0(x)\right)/Z \end{gathered}$$ (145)

The target is now $p_1(x)=p(x\mid A)$. The final form emphasizes that the prior $p(x)$, the statistical model $p(A\mid x)$, and the observations $A$ define an effective “energy landscape” for the parameters $x$. The $x$-independent $p(A)$ acts like a normalization constant. Like in statistical mechanics, in complex, high-dimensional models, directly computing expectation values, like the parameter means or normalized posterior probabilities, is intractable.

The Monte Carlo method relies on a set of $N$ (independent) samples from the target distribution $p_1$ to estimate expectation values:

$$\begin{gathered} ⟨a{⟩}_1\approx \frac{1}{N}\sum _{n=1}^{N}a\left(x_n\right)≕{\hat{a}}_N, \\ {x}_n~p_1,n=1,\dots ,N. \end{gathered}$$ (146)

We use $\hat{\cdot }$ for estimators, and simplify subscripts $⟨a{⟩}_1≔⟨a{⟩}_{p_1}$ where unambiguous. Since it uses a finite number of samples, the estimator Eq. (146) is not perfect (throughout, we use $\hat{\cdot }$ for estimators). It is correct on average, but has a nonzero mean squared error

$$⟨\left({\hat{a}}_n-⟨a⟩\right)⟩=0,⟨{\left({\hat{a}}_n-⟨a⟩\right)}^2⟩=\frac{\text{Var}(a)}{N}.$$ (147)

Estimating strongly fluctuating quantities requires many samples.

Computing statistical-mechanical observables and Bayesian inference thus become sampling tasks: Given an energy or an unnormalized probability density, draw independent samples to estimate observables or normalized probabilities.

2. Probability distributions in high dimension

Sampling in high dimension $d\gg 1$ is challenging. First, complicated probability distributions often feature multiple, disconnected modes. An example from physics is the Ising model in the ordered phase (Fig. 6, left). In the high-dimensional (thermodynamic) limit, the free energy barrier between the two magnetization states scales like $d$ and thus becomes very large [7]. In Bayesian statistics, different modes in the posterior are distinct “explanations” of the observed data. Capturing all modes is therefore essential.

FIG. 6.

(Left) Boltzmann distribution of the magnetization at different temperatures in the fully connected Ising model with energy $E=d^{-1}\sum _{i\ne j}{\sigma }_i{\sigma }_j$ where the ${\sigma }_i\in \{-1,1\}$ are binary “spins”. The magnetization is the average over its spins, $m(x)=d^{-1}\sum _i{\sigma }_i$. At high temperature, the distribution is simple and unimodal, but at low temperature, it has two disconnected modes. (Right) Magnetization of MCMC samples with $d=20$ spins in the low-temperature ferromagnetic phase $\beta =2$. The MCMC sampler carries out a random walk over spin configurations; jumps are barrier-crossing events between positive and negative magnetization. Convergence is very slow: it takes ~ 10⁶ steps to escape a local free energy minimum.

Second, the high-probability regions that sampling must locate and explore become needles in a high-dimensional haystack [85]. This issue is often called the curse of dimensionality. More precisely, high-dimensional probability measures tend to concentrate into very narrow sets (Fig. 7, right), which sampling algorithms must explore. This phenomenon, a generalization of the central limit theorem, is called concentration of measure [86]. In statistical mechanics, it underlies the equivalence of micro- and macrocanonical ensembles in the thermodynamic limit $d\to \infty$ : large energy fluctuations become unlikely.

FIG. 7.

Samples from a Gaussian with zero mean and variance ${\sigma }^2=1/d$ in dimension $d=2$ and $d=1000$, plotted in (projected) polar coordinates $r=\sqrt{\sum _i{x}_i^2},\theta =\text{arctan}\left(x_1/x_2\right)$. For large $d$, samples concentrate on a sphere of radius 1 around the mean.

How fluctuations grow with the number of dimensions is also of great importance in sampling. As a rule of thumb, estimating an extensive quantity like the energy is feasible. Let us consider a scalar observable $a$ that is approximately Gaussian and for which the mean and standard deviation are of the same order. If the magnitude is extensive, $\left|a\right|~d$, the estimator Eq. (146) requires $N\gtrsim d$ samples for an accurate estimate. However, estimating the exponential of an extensive quantity $e^a$ is very hard. Since $\text{Var}\left[e^a\right]~e^{2\text{Var}\left[a\right]}~e^{2d}$, it requires $N\gtrsim e^d$ samples and is thus infeasible. Controlling this variance is a key question in sampling algorithm design.

3. Basic sampling algorithms: importance and Langevin sampling

Most sampling algorithms begin with a base distribution $p_0$ from which one can sample easily (often, the standard $d$-dimensional Gaussian $𝒩(0,\mathbb{I})$) [87]. The game is then to convert samples from $p_0$ into samples of the target distribution $p_1$. The most straightforward implementation of this idea, transform sampling, applies a map to base samples. This idea underlies modern generative models and will be discussed in Application Sec. V C.

Importance sampling is based on reweighting. Assuming that the unnormalized density ${\tilde{p}}_1$ is known, one uses the base distribution $p_0$ to produce a set of proposals $x_n~p_0$ (in this context, $p_0$ is often called the proposal distribution). Weights ${\omega }_n$ compensate for the difference between $p_0$ and $p_1$:

$$\begin{gathered} ⟨a{⟩}_1=⟨a\omega {⟩}_0=\frac{⟨a\tilde{\omega }{⟩}_0}{⟨\tilde{\omega }{⟩}_0}\approx \frac{\sum _{n=1}^N{\tilde{\omega }}_{n}a\left(x_n\right)}{\sum _{n=1}^N{\tilde{\omega }}_n}, \\ \omega (x)=\frac{p_1(x)}{p_0(x)},{\tilde{\omega }}_n=\frac{{\tilde{p}}_1\left(x_n\right)}{{\tilde{p}}_0\left(x_n\right)},x_n~p_0. \end{gathered}$$ (148)

Crucially, the sample estimate only requires evaluating the unnormalized probability densities, since the normalization constants $Z_0,Z_1$ drop out.

Reweighting converts proposals into target samples. The closer $p_0$ is to $p_1$, the more efficiently the method works; if $p_0$ and $p_1$ are very different, most proposals land in regions of vanishing $p_1$-probability. This means that the finite-sample estimate will be dominated by a few samples with the largest weights. This phenomenon can be quantified by the effective sample size, the number $N_{\text{eff}}$ of samples from $p_1$ that yield the same error as $N$ importance samples from $p_0$, often approximated as [8]

$$N_{\text{eff}}\approx {\left(\sum _{n=1}^N{\tilde{\omega }}_n\right)}^2/\left(\sum _{n=1}^N{\tilde{\omega }}_n^2\right).$$ (149)

The effective sample size is thus set by the 2^nd moment of the weights. Let us see how this arises. By Eq. (147), $N_{\text{eff}}$ samples from $p_1$ yield an error $\text{Var}\left[a\right]/N_{\text{eff}}$. Instead of the “bare” observable $a$, importance sampling estimates the product $a\cdot \omega$, where $\omega =p_1/p_0=\text{exp}\left(\text{log}{p}_1-\text{log}{p}_0\right)$ is the weight. The error of the importance sampling estimator Eq. (148) is therefore $\text{Var}[a\cdot \omega ]/N~(\text{Var}[a]\cdot \text{Var}[\omega ])/N$. Comparing the two yields $N_{\text{eff}}=N/\text{Var}\left[e^{\text{log}{p}_1-\text{log}{p}_0}\right]$. In high dimensions, importance sampling thus becomes ineffective. The ratio $p_1/p_0$ generally fluctuates wildly since probability measures concentrate in narrow regions [88]. The “energies” $E_{0,1}=-\text{log}{\tilde{p}}_{0,1}$ are usually extensive, $E_{0,1}~d$, and so $\omega =p_1/p_0~e^d$. Indeed, constraining $p_1/p_0$ is a key part of importance-sampling-based machine learning algorithms like PPO [89], which we will discuss in Sec. V A.

Markov chain Monte Carlo (MCMC) uses an easy-to-simulate stochastic process that converges to the target distribution $p_1$ over time $t$. This generates samples incrementally and allows simulating complex, high-dimensional distributions. Its simplest variant, Langevin sampling, uses Brownian motion in a potential $E_1=-\text{log}{\tilde{p}}_1$:

$${\dot{x}}_t=-\nabla E_1+\sqrt{\epsilon }{\eta }_t,x_0~p_0$$ (150)

Here ${\eta }_t$ is the standard Gaussian white noise. We denote the marginal distribution of $x_t$, the sampler density, by $q_t(x)$. The stationary distribution of this process is precisely $p_1=\text{exp}\left(-E_1\right)/Z_1$, so after waiting long enough, $x_t$ can generate target samples. MCMC succeeds or fails based on how rapidly $q_t$ equilibrates to the target distribution. However, for complex, high-dimensional distributions, convergence can be painfully slow. For multi-modal distributions, Langevin dynamics can get stuck in a single energy minimum and thus perform poorly (Fig. 6, right). For instance, for a mixture of Gaussians with well-separated peaks ${\mu }_1,{\mu }_2$, trajectories must cross an energy barrier $\Delta E~{\left|{\mu }_1-{\mu }_2\right|}^2/{\sigma }^2$, which takes a time $~\text{exp}(\Delta E)$.

While Langevin sampling is conceptually based on equilibrium physics, the algorithms we discuss next use non-equilibrium physics. They transform the base into the target distribution over a fixed, finite time frame, akin to the thermodynamic protocols encountered in Sec. III B. The target distribution is no longer a stationary state. Intuitively, this allows proceeding at “full speed” where equilibrium approaches must “slow down” as the target is approached.

B. Annealed importance sampling

The method of simulated annealing, originally introduced in statistical physics [7, 90], overcomes this issue and can be seen as an antecedent of modern diffusion methods. Simulated annealing was originally an algorithm to optimize an energy $E(x)$. Because greedy optimization methods like gradient descent, ${\dot{x}}_t=-\nabla E$, can get stuck in local minima, one introduces noise (a finite temperature) to allow escape: ${\dot{x}}_t=-\nabla E+\sqrt{\epsilon }{\eta }_t$. However, if $\epsilon$ is too large, the state $x_t$ spends much time away from the energy minimum due to “thermal” fluctuations. Simulated annealing splits the difference by starting with a high temperature and gradually lowering it over time. The initial high-temperature distribution is supposed to act as a “bridge” connecting different energy minima.

1. Sampling with time-varying energies

By reparametrization, a changing noise level can be absorbed into a changing energy. In this section, we rescale so that $\epsilon =2$ to simplify notation. One thus generates samples by:

$${\dot{x}}_t=-\nabla E_t+\sqrt{2}{\eta }_t,x_0~p_0$$ (151)

The process generates proposal samples $x_t$, with actual marginal density $q_t(x)$. For independent runs, $x_t^{(n)}$ denotes the state in run $n$. By convention, time runs from $t=0$ to $t=1$. Thus, simulated annealing departs from conventional Langevin dynamics: samples are generated over finite time and using a time-dependent energy $E_t$. One simple strategy is linear interpolation between the energies:

$$\begin{gathered} E_t=\left(1-t\right)E_0+t{E}_1, \\ \text{where}{p}_0=e^{-E_0}/Z_0,p_1=e^{-E_1}/Z_1 \end{gathered}$$ (152)

Eq. (151) can be thought of as a thermodynamic protocol (Sec. III B). The “instantaneous equilibrium” distribution corresponding to Eq. (152) is $p_t^{\text{eq}}=e^{-E_t}/Z_t$, with $p_0^{\text{eq}}=p_0$ and $p_1^{\text{eq}}=p_1$.

However, this strategy is not yet perfect. The distribution $q_1$ of the trajectory endpoints $x_1$ from Eq. (151) is not the same as the equilibrium distribution $p_1$, our target. Physically, the transformation induced by the time-dependent energy is non-adiabatic. Mathematically, the Fokker-Planck equation for $q_t$ reads ${\partial }_t{q}_t=\nabla \cdot \left(q_t\nabla E_t\right)+{\nabla }^2{q}_t$. The $p_t^{\text{eq}}$ cannot be a solution since they are instantaneous equilibria, $\nabla \cdot \left(p_t^{\text{eq}}\nabla E_t\right)+{\nabla }^2{p}_t^{\text{eq}}=0\ne {\partial }_t{p}_t^{\text{eq}}$ [91]. Therefore, $x_1$ cannot be used directly as a sample for $p_1$.

Annealed importance sampling [92] (AIS) corrects the annealing process Eq. (151) to produce samples from $p_1$. Similarly to Girsanov’s theorem (Sec. I C 2), AIS computes how the marginal distribution of $x_t$ changes due to the additional force $\nabla E_t-\nabla E_0=t\nabla \left(E_1-E_0\right)$ from the time-dependent energy. AIS computes a weight ${\omega }_t$ for each realization $x_t$ that compensates for the difference between $q_t$ and $p_t^{\text{eq}}$ via importance sampling Eq. (148), i.e. ${\omega }_t=p_t^{\text{eq}}\left(x_t\right)/q_t\left(x_t\right)$. As we will derive shortly, this weight can be calculated “step-by-step” by integrating along the trajectory.

$$\begin{gathered} \frac{\text{d}}{\text{d}t}\text{log}{\omega }_t=\left({\partial }_t\text{log}{p}_t^{\text{eq}}\right)\left(x_t\right), \\ {\omega }_t=\frac{Z_0}{Z_t}{e}^{-w_t},w_t={\int }_0^t\left({\partial }_s{E}_s\right)\left(x_s\right)\text{d}s. \end{gathered}$$ (153)

where $w_t$ is the work performed on the trajectory particle due to the changing energy. In practice, the common factor $Z_0/Z_t$ need not be computed: when estimating averages, one simply normalizes by the total weight of all AIS samples. A special case of AIS is thermodynamic integration [93], in which the energy $E_t$ is changed quasi-statically.

In non-equilibrium thermodynamics, Eq. (153) is known as Jarzynski’s equality [78] (Sec. III B 4 deduced it from Crooks’ fluctuation theorem, another non-equilibrium thermodynamics result). Jarzynski’s equality relates the free energy difference ${ℱ}_1^{\text{eq}}-{ℱ}_0^{\text{eq}}=-\text{log}{Z}_1+\text{log}{Z}_0$ to the work $w_t$ on individual trajectories. It notably allows calculating the free energy difference between two states. This result from thermodynamics was, in fact, an important inspiration for AIS [92].

2. Proof of the annealed importance sampling identity

Jarzynski’s equality was proven in Sec. III B 4 using the Feynman-Kac formula. Here, we provide an alternative derivation of Eq. (153), effectively using a path-integral approach [78, 92]. The biggest challenge lies in notation: we already distinguished the instantaneous equilibria $p_t^{\text{eq}}$ from the distribution $q_t$ of the non-equilibrium process Eq. (151). In addition, we now introduce the time-reversed version of Eq. (151), in which the energy is changed from $E_1$ back to $E_0$:

$${\dot{\bar{x}}}_t=-\nabla E_{1-t}+\sqrt{2}{\eta }_t,{\bar{x}}_0~p_1$$ (154)

Note that the reverse process is initialized at $p_1$, our target: ${\bar{q}}_0=p_1$. We denote its marginal density by ${\bar{q}}_t$ and its path law by $\bar{P}$. We now show how $p_t^{\text{eq}}$ (equilibrium), $q_t$ (forward process), and ${\bar{q}}_t$ (reverse process) are related.

We discretize time into small intervals, $t=0,dt,2dt,\dots ,1$. Formally, the probability of a trajectory equals the product of the (infinitesimal) transition probabilities:

$$P[x(\cdot )]=p_0\left(x_0\right)\prod _t{k}_t\left(x_{t+dt}\mid x_t\right)$$ (155)

In the discretization, we alternate between taking trajectory steps $x_t\mapsto x_t+\left(-\nabla E_t+\sqrt{2}{\eta }_t\right)dt$ and updating the energy, $E_t\mapsto E_{t+dt}$. The $k_t(y\mid x)$ are one-step transition matrices. They are the equilibrium transition probabilities for an energy “frozen” at $E_t$, and, in particular, obey detailed balance. The transition probabilities of the reverse process are therefore⁴:

$$k_t\left(x_t\mid x_{t+dt}\right)=k_t\left(x_{t+dt}\mid x_t\right)\frac{p_t^{\text{eq}}\left(x_t\right)}{p_t^{\text{eq}}\left(x_{t+dt}\right)}$$ (156)

Hence, the difference in a trajectory $x_t$’s log-probability under the forward and reverse processes reads:

$$\begin{gathered} \text{log}\frac{\bar{P}[x(1-\cdot )]}{P[x(\cdot )]}=\sum _t\text{d}t\frac{\text{log}{p}_{t+dt}^{\text{eq}}\left(x_t\right)-\text{log}{p}_t^{\text{eq}}\left(x_t\right)}{\text{d}t} \\ =\int \text{d}t\left({\partial }_t\text{log}{p}_t^{\text{eq}}\right)\left(x_t\right)=\text{log}{\omega }_1. \end{gathered}$$ (157)

The reverse process starts at ${\bar{q}}_0=p_1$, so the weight ${\omega }_1$ corrects for the difference between $q_1$ and $p_1$. After reweighting with $\omega$, the trajectory endpoints $x_1^{(n)}$ can be used as importance samples.

In summary, AIS is an importance sampling scheme that uses time-dependent energy landscapes to generate samples from a target distribution. Because AIS leverages a non-equilibrium process, it can greatly outperform, say, Langevin sampling, when convergence to equilibrium is slow, for example due to energy barriers. The AIS importance weights are easy to calculate numerically⁵. An example is shown in Fig. 8 for the Ising model in the spontaneously magnetized phase. While Langevin sampling took 10⁶ steps to escape a local free energy minimum, AIS can generate independent samples in 100 steps by cooling down from the unmagnetized phase. While AIS is guaranteed to deliver unbiased importance samples, its practical performance greatly depends on the annealing schedule, i.e., on the choice of the time-dependent energy $E_t$. We discuss this question next.

FIG. 8.

(Left) Simulated annealing trajectories for the fully connected Ising model with $d=20$ spins, going from $\beta =0$ to $\beta =2$ in 100 steps. (Right) Distribution of magnetization at $\beta =2$ based on annealing runs. Without reweighting, the samples are not distributed correctly.

C. Sampling as a control problem

We now understand how to generate target samples using the AIS method. Does that mean that the sampling problem is “solved”? Further, the AIS scheme works for any time-dependent energy landscape. In which sense is one choice better than another? Just like thermodynamic protocols and the transport plans of Sec. II A can be optimized to minimize energy dissipation or transport costs, we will now discuss how sampling protocols can be optimized for statistical efficiency. As we noted above, importance sampling works best when the proposal and the target distributions are close. Otherwise, many samples land in regions of vanishing probability and are thus “lost”. By Eq. (149), the performance of importance sampling is quantified by the variance of the weights $\omega =\text{exp}\left(\text{log}{p}_1-\text{log}{p}_0\right)$. Our goal is thus to minimize weight variance to maximize statistical efficiency.

In the thermodynamic language of the Jarzynski equality (Sec. III B 4), minimal weight variance means that the work for moving the system between the two equilibria with free energies ${ℱ}_{0,1}^{\text{eq}}=-\text{log}{Z}_{0,1}$ is minimal. To see this, recall Eq. (153): the unnormalized AIS weights are ${\tilde{\omega }}_t=\text{exp}\left(-w_t\right)$, where $w_t$ is the work performed on a trajectory $x_t$. The Jarzynski equality Eq. (134) states that the free energy difference is ${ℱ}_1^{\text{eq}}-{ℱ}_0^{\text{eq}}=-\text{log}⟨e^{-w_1}⟩\le ⟨w_1⟩$. This bound is saturated when the work $w_1$ on all trajectories is equal, and thus the weight variance is zero (we caution that minimal work coincides with minimal variance, but away from the minimum, the two objectives differ). Physically, an efficient protocol is close to adiabatic: it ensures that samples at time $t$ rapidly thermalize to the next energy landscape $E_{t+dt}$, and avoids, for instance, regions of phase transitions [77].

In an ideal importance sampler, $\text{Var}\left[\omega \right]=0$ and thus $\omega (x)=\text{const}$. That is, the ideal proposals $x$ are already distributed according to the target $p_1$. To get around this somewhat circular conclusion, we will formulate the design of a good proposal distribution as an optimization problem. Indeed, importance sampling can be seen as an optimal control problem by considering the time-dependent energy gradient $-\nabla E_t$ in Eq. (151) as a control force $u_t$.

1. Sampler design as a control problem

To formalize sampler design as a Path Integral Stochastic Control (PISC) problem (Sec. ID), we consider a generalized setting. We aim to evaluate expectation values of the form $a\cdot e^{R_0}$ over trajectories of a stochastic process:

$$⟨a[x(\cdot )]e^{R[x(\cdot )]}⟩,{\dot{x}}_t=f_t\left(x_t\right)+\sqrt{\epsilon }{\eta }_t,$$ (158)

The brackets $[\cdot ]$ indicate that $a$ and $R$ may now be functionals of the entire trajectory $x(\cdot )$.

To map Eq. (158) to the importance sampling setting, choose:

$$\begin{gathered} R\left[x\left(\cdot \right)\right]=\text{log}\frac{p_1\left(x_1\right)}{p_0\left(x_1\right)}=E_0\left(x_1\right)-E_1\left(x_1\right)+\text{log}\frac{Z_0}{Z_1}, \\ {f}_t=-\nabla E_0,\epsilon =2,a[x(\cdot )]=a\left(x_1\right),x_0~p_0. \end{gathered}$$ (159)

The stochastic process produces Langevin samples from the base distribution; $R$ is the log importance weight, so $e^R$ reweights them to the target; the expectation on the left of Eq. (158) is then $⟨a{⟩}_{p_1}$. As discussed below Eq. (149), in high dimensions, the importance weight $e^R~e^d$ fluctuates wildly. This makes a finite-sample estimate of Eq. (158) unreliable.

To remedy this, we add a control force $u_t$ to drive trajectories toward higher rewards, and then reweigh the trajectories to account for the change in distribution [26, 94],

$${\dot{x}}_t=f_t\left(x_t\right)+u_t\left(x_t\right)+\sqrt{\epsilon }{\eta }_t.$$ (160)

We denote by $P^0[x(\cdot )]$ the probability of a path without control, and by $P^u[x(\cdot )]$ the probability with control. To compensate for the modified dynamics, we reweigh:

$${⟨a[x(\cdot )]e^{R[x(\cdot )]}⟩}_0={⟨a[x(\cdot )]e^{R[x(\cdot )]}\frac{P^0[x(\cdot )]}{P^u[x(\cdot )]}⟩}_u$$ (161)

We now need a suitable reward to choose the control $u$. We will aim to maximize $R$, balanced against a quadratic control cost:

$$u^{*}=\text{arg}\underset{u}{\text{min}}{⟨{\int }_0^1\frac{{‖u_t\left(x_t\right)‖}^2}{2\epsilon }\text{d}t-R[x(\cdot )]⟩}_u.$$ (162)

We thus arrive at a PISC problem. Recall from Eq. (60) that the optimal path measure is a tilted version of the uncontrolled measure:

$$P^{*}\left[x(\cdot )\mid x_0=x\right]=P^0\left[x(\cdot )\mid x_0=x\right]\cdot e^{R[x(\cdot )]}/Z$$ (163)

Indeed, Girsanov’s theorem from Sec. I C 2 shows that ${⟨{\int }_0^1{‖u_t\left(x_t\right)‖}^2/(2\epsilon )\text{d}t⟩}_u={\text{D}}_{\text{KL}}\left(P^u\Vert P^0\right)$. The control cost is thus an entropy, and as we saw in Sec. I B 2, entropy maximization generally leads to solutions that are exponential tilts. Plugging the optimal control into Eq. (161):

$$\begin{gathered} {⟨a[x(\cdot )]e^{R[x(\cdot )]}\mid x_0=x⟩}_0 \\ ={⟨\left.a[x(\cdot )]e^{R[x(\cdot )]}\frac{P^0[x(\cdot )]}{P^{*}[x(\cdot )]}\right|x_0=x⟩}_{u^{*}} \\ =Z{⟨a[x(\cdot )]\mid x_0=x⟩}_{u^{*}}. \end{gathered}$$ (164)

The weights $e^R$ in the expectation are now constant (conditional on the initial condition $x_0$ – unconditional expectation values must additionally average over the easy-to-sample $x_0~p_0$). In the case of importance sampling, we can compute $⟨a{⟩}_{p_1}$ purely from samples $x_t,x_0~p_0$ of the Langevin equation Eq. (160), with the control force found by optimizing Eq. (162). Since no reweighting is necessary, the effective sample size is maximal, $N_{\text{eff}}=N$: every proposal sample is worth one true sample.

2. Non-equilibrium systems and transition path sampling

The formulation Eq. (158) also applies to non-equilibrium systems, like energetically driven chemical reactions [95, 96]. Instead of a Boltzmann distribution, the state $x_t$ of the system at time $t$ follows a non-equilibrium Langevin equation, combining deterministic forces and random fluctuations. The dynamics in Eq. (158) thus become physically meaningful, rather than a sampling device. Evaluating expectation values over physical trajectories is of great importance, for example, to estimate chemical reaction rates. In this setting, so-called rare events can have outsize importance: the crossing of a potential barrier determines whether a reaction occurs or not. Mathematically, such events can be encoded in the reward $R$ in Eq. (158). One faces the same challenge as above: strong fluctuations in $e^R$ lead to poor finite-sample estimates (the target event may simply not occur in a finite sample). Variational Path Sampling [95] addresses it precisely as discussed in Sec. IV C 1. One introduces a control $u_t$ to encourage the target rare events, and reweighs trajectories to compensate for the fictitious control when computing physical observables (Eq. (161)).

3. Control as a sampling problem

In summary, the design of an optimal importance sampler is a control problem, where the reward is the “target” probability of the trajectories. This result enables systematic optimization of annealing samplers. But, in some sense, we simply kicked the can down the road: how do we solve the resulting control problem? The key to resolving this conundrum is that, vice versa, control problems can be solved by sampling. Indeed, we saw in Sec. I C 3 that control problems can be seen as probabilistic inference problems. Since sampling provides a general-purpose method to solve the latter, it should also be applicable to the former. This suggests an iterative approach in which the sampler is used to solve a control problem, whose solution in turn improves the sampler. This duality between sampling and control will play an important role in Sec. V A on reinforcement learning.

Let us make this precise. The notations “$\psi$” and “$R$” in Eq. (158) are no coincidence. Conditioned on the trajectory’s initial condition, the weighted expectation in Eq. (158) is the Cole–Hopf transform ${\psi }_s=e^{V_s}$ of the all-important value function of a control problem with accumulated reward $R[x(\cdot )]$ (Sec. I D 2). We write $R_s[x(\cdot )]$ for the accumulated reward from time $s$ onward, including the terminal reward and excluding control cost. Recall that the value function obeys the HJB equation (the recursive formulation of optimality), which can be Cole-Hopf transformed into a Schrödinger equation (in imaginary time). The Feynman-Kac formula Eq. (62) shows that the solution to this Schrödinger equation can be obtained by a Feynman path integral over stochastic trajectories:

$$\begin{gathered} {\psi }_s(x)={⟨e^{R_s[x(\cdot )]}\mid x_s=x⟩}_{P^0}, \\ {\dot{x}}_t=f_t\left(x_t\right)+\sqrt{\epsilon }{\eta }_t. \end{gathered}$$ (165)

The stochastic trajectories are now initialized at $x_s=x$ instead of at random. Up to this detail, the Feynman-Kac formula is precisely Eq. (158) with $a=1,f_t$ as the dynamics in the absence of control, and $R_s$ as the reward of the control problem. Intuitively, Eq. (165) computes the value function by “undirected exploration” (no control applied) and taking inventory of the rewards encountered. In the language of this chapter, the Feynman-Kac formula is a sampling-based approach to control problems: applying the Monte Carlo method to Eq. (165) yields an estimate of the value function from a finite number of stochastic trajectories. However, Monte Carlo estimation faces the same challenge as in importance sampling. In high dimensions, $e^{R_s}~e^d$ fluctuates wildly. This makes a finite-sample estimate of Eq. (165) unreliable. Uncontrolled trajectories are very unlikely to gather high reward $R_s$, while it is precisely high-reward trajectories that dominate Eq. (165).

Sec. IV C 1 explains the solution: add a control force to bias trajectories towards higher rewards. This suggests an iterative procedure, starting with a guess for the optimal control $u_t^{(0)}$:

1. Generate importance samples from the $u_t^{(k)}$-controlled process (Eq. (161))

2. Estimate $\text{log}{\widehat{\psi }}^{(k)}$ by Monte Carlo, using the reweighted trajectories (Eq. (165))

3. Recompute the optimal control from the value function by optimizing the expected reward at each timestep ($u_t^{(k+1)}(x)=\epsilon \nabla \text{log}{\widehat{\psi }}_t^{(k)}(x)$, Eq. (55))

and repeat until convergence. This procedure is closely related to the policy iteration algorithm for control problems [97], which iterates evaluation of the current control’s value function (steps 1–2) with improvement of the control using the value function (step 3). Here, evaluation happens through Monte Carlo estimation. In Sec. V A on reinforcement learning (RL), we will see that policy iteration forms the basis for state-of-the-art methods like the soft actor-critic algorithm [98], and that controlling the variance of Monte Carlo estimates is a crucial question in RL algorithm design [97]. RL generalizes control theory to the case where the dynamics $f_t$ and reward $r_t$ are not known in advance but must be learned through exploration. RL is often formulated in a somewhat different language, but makes use of the same principles that connect control, inference, and sampling problems.

V. APPLICATIONS

The Applications chapter will cover three topics. Sec. V A on reinforcement learning builds on control theory (Sec. I A), maximum-entropy methods (Sec. I C), and sampling (Chapter IV); Sec. V B on Wasserstein gradient flows draws on transport (Chapter II) and non-equilibrium thermodynamics (Sec. III B); and Sec. V C on generative models combines ideas from transport (Chapter II) and sampling (Chapter IV). A general theme is moving from settings where the dynamics, rewards, or distributional constraints are known, to settings where they have to be learned from data. The Applications chapter is thus centered on ideas from machine learning.

A. Reinforcement learning

Before introducing the reinforcement learning (RL) problem, we begin with a simple example to build intuition and illustrate what makes the problem hard. Take a simple problem setting of a robot in a new warehouse that it has never seen before (Fig. 9), where its goal is to reach a target as soon as possible. Importantly, the robot has no prior knowledge about the warehouse, and can only make decisions based on its collected experience and observed rewards. Let’s say that the robot executes random actions until it reaches the goal. This initial path is likely suboptimal and circuitous. Should the robot continue taking this path or explore to find a potentially shorter route? Thus, if the robot wishes to find the optimal route, it must balance between collecting information about its environment and refining the solutions found so far. This balance, known as the exploitation-exploration problem, is the core difficulty of the RL problem.

FIG. 9. The explore-exploit problem in reinforcement learning.

The robot wishes to reach a goal, and is given a reward of 1 at the goal (star), a reward of 0 in empty cells, and a slightly negative reward −0.01 in hazard cells (electricity). Furthermore, at any point in time, the robot might break down with a large independent probability 10%. Thus, if the robot wants to maximize its expected reward, it should heavily prioritize shorter routes to the goal. Assume that the robot’s initial state is fixed to the top left of the warehouse. At initialization, the robot’s control is random and chooses these two actions with equal probability. From this point, the agent can choose one of two controls: go right or go down. These two controls lead to two paths towards the goal, a suboptimal path that is long and circuitous with ample opportunity to break down but with no penalties, and an optimal path that routes directly to the goal with a small incurred penalty (e.g. an activation energy barrier).

The RL formalism and methods build on the stochastic optimal control (SOC) framework. As discussed in Sec. I A, SOC aims to find the optimal control $u_t(x)$ that maximizes the expected cumulative reward in a stochastic environment. Like the SOC problem, the RL problem also aims to solve for optimal controls that maximize a cumulative reward. However, general RL aims to solve for optimal actions from the perspective of a decision-making entity in a new environment where both environmental statistics and rewards are unknown [97], like the robot in a new warehouse. This entity is known as an agent. Recall that, given the dynamics of the environment, a reward function, and boundary conditions, the Hamilton–Jacobi–Bellman (HJB) equation determines the value function and hence the optimal control. Its underlying recursion decomposes the value into the immediate reward and the expected value at the next state. This recursion forms the theoretical backbone of modern RL methods. In RL, finding the optimal control not only involves solving HJB but also requires the agent to accumulate experience such that finding the optimal controls is tractable. For instance, while completely random controls are very suboptimal if an agent already has access to environment dynamics and reward functions (SOC setting), such a strategy may actually be somewhat sensible in a new environment (RL setting), where random controls can help the agent explore and build a better model of the world; more generally, the agent must carefully balance exploitation (e.g. greedily navigating to states which have already been shown to lead to high total rewards) and exploration. The desiderata to maximize reward and gather data to improve controls in the future often compete, in what is known as the exploitation-exploration problem. This balancing act is a function of the restricted, harder RL problem setting and makes the SOC and the RL problem fundamentally different, despite the shared overall objective of finding optimal controls.

The central tension between exploration and exploitation is illustrated by the so-called two-armed bandit problem. An agent repeatedly chooses between two slot machines—the “two arms of a bandit”— each offering one unit of reward with a fixed but unknown probability and zero otherwise. Only the outcome of the chosen arm is observed, and the aim is to maximize the accumulated reward. Playing the arm with the larger observed average reward exploits existing information, whereas exploring the other arm ensures that the arm that has been deemed optimal with finite observations is indeed optimal. The bandit problem admits a solution to the exploration-exploitation tradeoff in a minimal setting, known as the Lai–Robbins bound [99] (in the asymptotic limit of a large number of plays). This bound shows that for strategies that learn efficiently across possible winning probabilities, the expected reward lost relative to knowing the better arm from the outset grows at least logarithmically with the number of plays. The intuition comes from the fact that large deviations in the observed reward can make a sub-optimal arm appear optimal. Sanov’s theorem, Eq. (27), tells us that the probability of the better arm producing observations that make it seem sub-optimal decreases exponentially with the number of times it is sampled, at a rate determined by the KL divergence. Such misleading data could cause the agent to neglect the better arm, with a loss that grows with the duration of play. Suppressing these rare but costly errors therefore requires continued exploration. The competition between exponentially decreasing probabilities and accumulating losses explains the logarithmic scale in the bound. Thus, successful strategies for exploration and exploitation are ones that increasingly favor the apparently better arms (or, more generally, actions) while continuing to test alternatives that the data have not yet ruled out.

1. The Markov decision process

Now that we have some intuition about the structure and fundamental challenges in the RL problem, we introduce necessary terminology to discuss key algorithms. States $x$ are the RL vocabulary for coordinates in a particular state space $𝒳$. For example, in the state space of the Cartesian plane, one possible state an agent can occupy is the origin. An environment is a space of states that an agent can navigate subject to transition dynamics $k(\cdot \mid x,u)$, where the transitions are a (generally stochastic) function of the current state and current control. A reward function $r(x,u)$ is a scalar function that takes as input a state and control and outputs a reward. Reinforcement learning is defined on discrete time steps $n$; we will use the subscripts $x_n$ and $u_n$ to denote the states and controls at decision time $t_n$. For finite-horizon problems, we include the timestep $n$ in the state $x$ throughout this section; policies and value functions therefore need no separate time subscript. An agent is a decision-making entity that executes controls in the world based on its observed transitions and rewards [97]. The controls an agent executes are known as actions $u$ which lie within an action space $𝒰$. Reactive controls are also referred to as a policy in RL, where the policy $\pi (u\mid x)$ is a distribution over controls conditioned on the current state $x$.

If the transition dynamics are Markovian – that is, if the statistics of the next state depend only on the current state and action and not any prior states and actions – the transition dynamics, reward function, state space, and action space define what is called a Markov Decision Process (MDP). If the transition dynamics are non-Markovian and rely on additional history, this problem is referred to as a Partially Observed Markov Decision Process, or a POMDP [100]. We do not consider POMDPs here, but note that effective approaches include augmenting model-free MDP algorithms with recurrent neural networks to incorporate memory [101].

For the finite-horizon formulas, let $n=0,\dots ,N$ label decision times $t_n=n\Delta t$. The agent takes an action and receives a reward at the final step $N$. A policy induces a joint stateaction path law $P^{\pi }\left[x_{0:N},u_{0:N}\right]$, where $x_{0:N}=\left(x_0,\dots ,x_N\right)$ and similarly for $u_{0:N}$.

a. The Objective.

Let $\pi$ be an arbitrary policy. Denote the path measure of this policy as $P^{\pi }$. Then, for some reward function $r(x,u)$, the goal of RL is to maximize the objective $𝒥\left[\pi \right]$:

$$𝒥\left[\pi \right]\triangleq {\mathbb{E}}_{P^{\pi }}\left[\sum _{n=0}^{N}r\left(x_n,u_n\right)\right],\underset{\pi }{\text{max}}𝒥\left[\pi \right].$$

In RL, it is standard to consider a cumulative reward modified by a discount factor $0<\gamma <1$ that is taken to be close to 1. The discount factor weighs the reward and incentivizes accumulating reward sooner, rather than later:

$$𝒥\left[\pi \right]\triangleq {\mathbb{E}}_{P^{\pi }}\left[\sum _{n=0}^N{\gamma }^{n}r\left(x_n,u_n\right)\right],\underset{\pi }{\text{max}}𝒥\left[\pi \right].$$

Another way to understand the discount factor is that it imposes an independent $1-\gamma$ probability at every timestep that an agent permanently transitions to a “death” state with no reward (i.e. an absorbing state with no reward).

b. Value Functions

For the value definitions and Bellman equations in this subsection, we use the undiscounted convention $\gamma =1$. To solve the MDP, RL algorithms invoke theoretical tools from SOC. These methods rely on value estimation, where value denotes the expected cumulative future reward conditioned on a starting state $x$ at step $n$, for a particular policy $\pi$:

$$V^{\pi }(x)={\mathbb{E}}_{P^{\pi }}\left[\sum _{j=n}^{N}r\left(x_j,u_j\right)\mid x_n=x\right].$$ (166)

This value function is also known as a critic that critiques the “goodness” of a particular policy $\pi$ at state $x$. For $n<N$, the value function satisfies the Bellman equation for policy $\pi$:

$$V^{\pi }(x)={\mathbb{E}}_{u~\pi (\cdot \mid x)}\left[r(x,u)+{\mathbb{E}}_{x^{\prime }~k(\cdot \mid x,u)}\left[V^{\pi }\left(x^{\prime }\right)\right]\right].$$ (167)

Here $x^{\prime }$ denotes the next state. The optimal value function is

$$\begin{gathered} V^{*}(x)=\underset{\pi }{\text{max}}{V}^{\pi }(x) \\ =\underset{\pi }{\text{max}}{\mathbb{E}}_{P^{\pi }}\left[\sum _{j=n}^{N}r\left(x_j,u_j\right)\mid x_n=x\right], \end{gathered}$$ (168)

and measures the maximal cumulative expected reward achievable by a policy from a particular starting point. Like the value function in the continuous stochastic control setting, this discretized optimal value function similarly follows a recursive relationship:

$$V^{*}(x)=\underset{u}{\text{max}}\left[r(x,u)+{\mathbb{E}}_{x^{\prime }~k(\cdot \mid x,u)}\left[V^{*}\left(x^{\prime }\right)\right]\right].$$ (169)

This equation is known as the Bellman optimality equation, the discrete-time analogue of the Hamilton-Jacobi-Bellman equation. The Q function critiques the “goodness” of a particular policy $\pi$ for some state and control:

$$Q^{\pi }(x,u)={\mathbb{E}}_{P^{\pi }}\left[\sum _{j=n}^{N}r\left(x_j,u_j\right)\mid x_n=x,u_n=u\right].$$ (170)

In other words, the Q function measures the cumulative expected reward of a policy when fixing the control at step $n$ to $u$. For $n<N$, the Q function satisfies the Bellman equation for policy $\pi$:

$$Q^{\pi }(x,u)=r(x,u)+{\mathbb{E}}_{x^{\prime }~k(\cdot \mid x,u)}\left[{\mathbb{E}}_{u^{\prime }~\pi \left(\cdot \mid x^{\prime }\right)}\left[Q^{\pi }\left(x^{\prime },u^{\prime }\right)\right]\right].$$ (171)

Here $x^{\prime }$ and $u^{\prime }$ denote the next state and action, respectively. The Q function also shows how to improve the policy $\pi$ by scoring the values of all actions $u$. The simplest scheme for policy improvement is to choose, at each state $x$, an action that maximizes $Q^{\pi }(x,u)$. With exact Q values, the resulting policy has expected cumulative reward at least as high as that of the original policy. However, clearly, this greedy scheme can lead to failed exploration if policy improvement is performed using an estimate of the true Q function from prior experience. Later in this section, we discuss alternative policy formulations that naturally incentivize exploration.

The optimal Q function definition directly follows:

$$\begin{gathered} Q^{*}(x,u)=\underset{\pi }{\text{max}}{Q}^{\pi }(x,u) \\ =\underset{\pi }{\text{max}}{\mathbb{E}}_{P^{\pi }}\left[\sum _{j=n}^{N}r\left(x_j,u_j\right)\mid x_n=x,u_n=u\right]. \end{gathered}$$ (172)

Just like the previously introduced objects, the optimal Q function also satisfies a Bellman optimality equation:

$$Q^{*}(x,u)=r(x,u)+{\mathbb{E}}_{x^{\prime }~k(\cdot \mid x,u)}\left[\underset{u^{\prime }}{\text{max}}{Q}^{*}\left(x^{\prime },u^{\prime }\right)\right].$$ (173)

The starting point of RL methods is based on these Bellman equations, which are the discretized analogues of the Hamilton-Jacobi-Bellman equation.

2. An overview of online RL algorithms

We explicitly consider the RL setting where an agent has access to the environment, and can actively choose transitions to gather new data. This problem setting is referred to as online RL. Key methods in this category include Soft Actor Critic (SAC) [98] and Twin-Delayed Deep Deterministic Policy Gradient (TD3) [102]. In this section, we focus on building up to the SAC method, which, in addition to being an established method for continuous control on which newer algorithms build [103], is also a Maximum Entropy RL (Max-Ent RL) method with close connections to maximum entropy inference principles discussed throughout this review. To get there, we start from basic, early precursors to SAC that do not require deep learning at all, then go through the practical considerations that ultimately culminate in SAC.

a. Tabular Q Learning.

A starting point to solve the RL optimization problem is to utilize the theoretical tools from SOC. Recall from Sec. IA that recursion provides a PDE called the Hamilton-Jacobi-Bellman equation. This PDE is a functional of the environment dynamics and the true reward function, where the solution is the optimal value function. Then, one route to find optimal controls is to first find the optimal value function, now with the Bellman Equation instead of the HJB equation, and extract the corresponding optimal control. However, the RL problem setting makes this route difficult: without access to the environment dynamics or the reward, we cannot directly solve the Bellman equation for the true optimal value function. One possible solution is to, instead, utilize some of the intuition from sampling. In RL, the agent actively collects experience that it must use to find optimal controls. When the agent executes a control at timestep $n$, the agent observes a transition with reward $r_n$ from state $x_n$ to the next state $x_{n+1}$, which is a sample from the true environment dynamics and reward.

Recall the iterative connection between sampling and control developed in Sec. IV C: sampled trajectories update the value estimates, which guide changes to the controls, which change the next set of sampled trajectories. One route to update the value estimates and policy is to directly average the realized cumulative rewards of the current policy, improve the policy using these estimates, and then evaluate the updated policy. Evaluation and improvement thus iterate and build on each other, because changing the policy changes the expected future rewards, which is referred to as policy iteration. In the most naive approach, we can roll out $M$ independent trajectories each starting from state $x$ with initial control $u$ at time $t_n$, then follow policy $\pi$ for all subsequent controls. We now switch to a discounted reward formulation, with $0<\gamma <1$, so that rewards received later count less. Averaging the realized cumulative rewards of these rollouts gives:

$${\widehat{Q}}^{\pi }(x,u)=\frac{1}{M}\sum _{i=1}^M\sum _{j=n}^N{\gamma }^{j-n}r\left(x_j^{(i)},u_j^{(i)}\right).$$ (174)

Then, a simple way to find a good policy based on these value estimates is to take the argmax over these estimated Q values:

$$u\leftarrow \text{arg}{\text{max}}_u{\widehat{Q}}^{\pi }(x,u).$$ (175)

However, evaluating different actions from complete rollouts can have high variance. Another approach uses a running estimate of value to replace the remainder of a trajectory [104]. Instead of sampling and summing all realized rewards after a particular state, we can directly use a running value estimate for the future returns from the next state. This substitution can reduce sampling variance at the cost of potentially introducing error from utilizing an incorrect running value estimate. The observed immediate reward $r$ and the remaining estimated value from the next state $x^{\prime }$ provide a new prediction for the current Q value, referred to as the target $y$:

$$\begin{gathered} y=r+\gamma \underset{u^{\prime }}{\text{max}}{\widehat{Q}}^{*}\left(x^{\prime },u^{\prime }\right), \\ \text{for}\text{observation}\left(x,u,r,x^{\prime }\right). \end{gathered}$$ (176)

We can use these targets to update the current Q estimate, using a learning rate of $\xi$:

$${\widehat{Q}}^{*}(x,u)\leftarrow {\widehat{Q}}^{*}(x,u)+\xi \left[y-{\widehat{Q}}^{*}(x,u)\right].$$ (177)

Repeatedly updating the current Q estimate toward these target predictions converges to the optimal Q function under the assumption that every state and action is sampled infinitely often [97, 105]. Conceptually, this update can be thought of as a “backup” of value from the future value estimate to the present estimate. Because the future value is “easier” to estimate (e.g. requires aggregating rewards over fewer timesteps) and is trivial at the endpoint of a trajectory, where it becomes exactly the expected reward, this backup moves value backwards in time from the easier subproblem to harder subproblems. This intuition is also why we can refer to Hamilton-Jacobi-Bellman as a continuous form of dynamic programming and the Bellman Equation as discretized dynamic programming. Practically, a table holds these Q values, and the Bellman Backup indexes into this table and updates the value entries. As such, this algorithm is known as tabular Q-learning.

b. Deep Q-Networks.

The practical limitations of tabular Q Learning motivate the sequence of methods (Deep Q-Networks, Soft Q-Learning, and Soft Actor-Critic) we discuss below. First and foremost, storing and updating a separate Q value for every state and action in a table becomes infeasible in a large state space. Indeed, the tabular setup does not transfer to the continuous setting. Deep Q-networks (DQN) address this restriction by approximating the Q function with a neural network, allowing experience at one state to inform predictions at others [106]. We denote learned estimates of $Q^{\pi }$ and $Q^{*}$ by $Q_{\theta }^{\pi }$ and $Q_{\theta }^{*}$, where $\theta$ denotes the network weights.

DQN fits a neural-network estimate $Q_{\theta }^{*}$ of the optimal Q function using all of an agent’s prior observations – including transitions and observed rewards – aggregated within a replay buffer $𝒟$. Rather than do tabular updates of the value functions following the Bellman equation, DQN directly minimizes the “error” between the value function and the recursed form of the value function over this replay buffer of all past experiences, which is amenable to neural network parameterization:

$$\begin{gathered} L(\theta )=\frac{1}{2}{\mathbb{E}}_{\left(x,u,r,x^{\prime }\right)~𝒟}[\left.{\left(Q_{\theta }^{*}(x,u)-\left(r+\gamma \underset{u^{\prime }}{\text{max}}{Q}_{\bar{\theta }}^{*}\left(x^{\prime },u^{\prime }\right)\right)\right)}^2\right], \\ \theta \leftarrow \theta -\xi {\nabla }_{\theta }L(\theta ). \end{gathered}$$

The notation $\bar{\theta }$ means that gradients cannot flow through the parameters in the target. The reason for this “stop gradient” is to force value to flow backwards in time: the value at time $t$ should not be updated with the rewards that occurred before time $t$. While tabular Q learning provably converges, DQN does not provably converge without additional assumptions due to function approximation error [106, 107].

3. MaxEnt RL & Soft Q-learning

We now return to the undiscounted convention $\gamma =1$ for the MaxEnt derivation and the two algorithms below. Throughout the remainder of the RL section, $Q^{\pi }$ and $V^{\pi }$ denote entropy-regularized (“soft”) values. The entropy coefficient $\alpha$ has the same units as the per-step reward; it is distinct from the convexity parameter used in Wasserstein geometry.

Learning values alone does not ensure that the agent explores, as is shown in Fig. 9. Indeed, the tabular Q-learning algorithm can only guarantee convergence if the policies visit every state infinitely often. In large state spaces, this condition can be impossible to satisfy even with random actions. Furthermore, always choosing the action with the largest estimated Q value can ignore routes with activation energy. In RL, a common heuristic for exploration is to take a random action a small, epsilon proportion of the time. This strategy is called epsilon-greedy exploration where, typically, the probability epsilon goes to 0 over the course of training [97, 106]. This exploration mechanism can be thought of as a heuristic version of simulated annealing. Later works employ Maximum-entropy RL (MaxEnt RL) to control exploration in a more principled way, incorporating the policy Shannon entropy $H\left[\pi \left(\cdot \mid x_n\right)\right]$ in the objective itself:

$$𝒥\left[\pi \right]\triangleq {\mathbb{E}}_{P^{\pi }}\left[\sum _{n=0}^N\left[r\left(x_n,u_n\right)+\alpha H\left[\pi \left(\cdot \mid x_n\right)\right]\right]\right]$$ (178)

where $\alpha$ is a temperature that controls the strength of the entropy term. This approach encourages the policy to spread out actions while exponentially weighting those with high expected reward. Here, we discuss Soft Q-learning, one of the early examples of MaxEnt RL that utilizes entropy as a principled mechanism to encourage exploration [108].

The crux of MaxEnt RL is the idea that the policy should induce the maximum-entropy path measure, where the energy of the measure is the expected cumulative reward. Crucially, MaxEnt RL has close connections to control as inference principles (Sec. I C). In fact, one can derive the soft Q-learning objective to learn optimal policies directly from a maximumentropy inference problem over path measures [20]. We lay out this derivation here, connecting MaxEnt RL to the approaches and perspectives discussed previously in the review.

The inference-as-control problem begins with a definition of a path measure induced by a particular control policy. If this path measure exhibits mean observables (for example, the Schrödinger Bridge problem), and these moments are all that is known about the realized trajectories, Sanov’s theorem states that the most likely path measure is the maximum entropy distribution that satisfies these moments. Each of these observables corresponds to a Lagrange Multiplier $r$ that enforces this constraint, and the resulting max-ent path measure is:

$$P^{u^{*}}[x(\cdot )]\propto P^0[x(\cdot )]\text{exp}\left(\sum _m{r}_m{a}_m[x(\cdot )]\right),$$ (179)

Here $a_m[x(\cdot )]$ is the $m$-th observable measured along a trajectory. The Lagrange multiplier $r_m$ enforces a contraint on the averaged observable, as in Sec. IC. Thus, these approaches give controls $u$ that give a max-ent realization of these observables.

In the control setting, we instead wish for the path measure to satisfy particular observations – for example, that the path measure realizes a high expected cumulative $r$ for some fixed reward function. The rewards $r$ no longer directly result from enforcing specific observables (e.g. termination within a particular set). Therefore, despite MaxEnt RL optimizing path measures that look identical to the control-as-inference approach, maximum entropy is a postulate of the MaxEnt RL approach – there is not a fundamental physical principle like Sanov’s theorem for fixed observables that underlies the MaxEnt RL formulation.

To derive the MaxEnt RL objective, we now express the continuous-time path measures above using the discrete-time MDP notation introduced in Sec. V A 1. As an example, the controlled stochastic dynamics of Sec. IA

$${\dot{x}}_t=f_t\left(x_t\right)+u_t\left(x_t\right)+\sqrt{\epsilon }{\eta }_t$$

discretized by an Euler–Maruyama step of duration $\Delta t$ result in the Gaussian transition kernel:

$$k\left(x^{\prime }\mid x,u\right)=𝒩\left(x^{\prime };z+\left[f_{t_n}(x)+u\right]\Delta t,\epsilon \Delta tI\right).$$ (180)

where $𝒩$ denotes the Gaussian density for the physical coordinates, where $I$ is the identity matrix. Then, an agent executes controls and transitions according to this discretized transition kernel. When an agent visits the pair $(x,u)$, the agent receives a reward $r(x,u)$. Together, the kernel (or transition dynamics), the reward, and the state and action spaces define the finite-horizon Markov decision process (MDP), thus showing how the continuous control problem becomes the MDP.⁶ The policy $\pi \left(u_n\mid x_n\right)$ denotes the agent’s probability of choosing action $u_n$ in state $x_n$ at step $n$ and plays the role of the control law $u_t(x)$. Iteratively combining the transition kernel and the policy through time gives the discrete-time analog of the controlled path measure $P^u[x(\cdot )]$, which now induces a measure over discretized trajectories $\left\{x_n,u_n\right\}$:

$$P^{\pi }\left[x_{0:N},u_{0:N}\right]=p_0\left(x_0\right)\prod _{n=0}^N\pi \left(u_n\mid x_n\right)\prod _{n=0}^{N-1}k\left(x_{n+1}\mid x_n,u_n\right),$$ (181)

and the analog of the reference measure $P^0[x(\cdot )]$ is the measure of the passive dynamics with uniform action distribution ${\pi }^0\left(u_n\mid x_n\right)=c_0$:

$$P^0\left[x_{0:N},u_{0:N}\right]=p_0\left(x_0\right)\prod _{n=0}^N{\pi }^0\left(u_n\mid x_n\right)\prod _{n=0}^{N-1}k\left(x_{n+1}\mid x_n,u_n\right).$$ (182)

Say that we wish to find controls that lead to large rewards. One way to do so is to construct a target measure that upweights high-reward trajectories, then regress the control-induced path measure towards this target measure for the training objective. Following the maximum entropy postulate in the MaxEnt RL approach, one possible choice for this target measure is (a possibly unrealizable) maximum-entropy, reward-maximizing path measure with respect to a base measure $P^0\left[x_{0:N},u_{0:N}\right]$:

$$\begin{gathered} P^{*}\left[x_{0:N},u_{0:N}\right]\triangleq \text{arg}{\text{max}}_P\left[{\mathbb{E}}_P\left[\frac{1}{\alpha }\sum _{n=0}^{N}r\left(x_n,u_n\right)\right]-{\text{D}}_{\text{KL}}\left(P\Vert P^0\right)\right] \\ =\frac{1}{Z}{P}^0\left[x_{0:N},u_{0:N}\right]\text{exp}\left(\frac{1}{\alpha }\sum _{n=0}^{N}r\left(x_n,u_n\right)\right). \end{gathered}$$ (183)

Here the maximization over $P$ is over normalized path measures. $P^{*}$ is the reward-tilted target measure and $\alpha$ is a temperature that controls the strength of the tilt. Importantly, this $P^{*}$ is the theoretically optimal max-ent target distribution and need not be achievable by a particular policy. In practice, the rewards are only observed, as the agent does not have a priori access to the true reward function following standard RL assumptions. We seek a policy whose path measure approximates this target, where the limited policies can only determine the executed controls, not the dynamics of the system. The optimal policy for this inference objective therefore produces the closest realizable path measure, which is the minimizer of ${\text{D}}_{\text{KL}}\left(P^{\pi }\Vert P^{*}\right)$⁷:

$${\pi }^{*}\triangleq \text{arg}\underset{\pi }{\text{min}}{\text{D}}_{\text{KL}}\left(P^{\pi }\Vert P^{*}\right)$$ (184)

Note that the path distribution $P^{{\pi }^{*}}$ is the closest realizable distribution and need not equal $P^{*}$. Thus, we have derived an inference objective that should provide high-reward, maximum-entropy controls. What is the relationship between this inference objective and the MaxEnt RL objective? It turns out that in discrete time, the transition kernels cancel in the ratio of the two measures, which “simplifies” this inference into the MaxEnt RL objective. The result is shown in Eq. (185), Here $H\left[\pi \left(\cdot \mid x_n\right)\right]=-{\mathbb{E}}_{u~\pi \left(\cdot \mid x_n\right)}\left[\text{log}\pi \left(u\mid x_n\right)\right]$ denotes the Shannon entropy of the policy at state $x_n$ and step $n$ and the normalizing constant $Z$ is independent of $\pi$. Minimizing the KL divergence is therefore equivalent, up to the positive factor $1/\alpha$ and the additive constant $\text{log}Z-(N+1)\text{log}{c}_0$ (neither of which should change the fixed point), to maximizing the maximum-entropy RL objective:

$$𝒥\left[\pi \right]\triangleq {\mathbb{E}}_{P^{\pi }}\left[\sum _{n=0}^N\left[r\left(x_n,u_n\right)+\alpha H\left[\pi \left(\cdot \mid x_n\right)\right]\right]\right].$$ (186)

Thus, as desired, we have theoretically derived the maximum-entropy RL objective from MaxEnt inference principles. As a note, Eq. (186) has the same free-energy structure as entropic transport in Eq. (81), with negative reward playing the role of cost and temperature $\alpha$ the role of $\epsilon$. Maximum-entropy RL introduces the entropy bonus in Eq. (186) as a regularizer that aids robustness [109] and exploration [98]. In the above derivation, we show that the entropy bonus emerges directly from the inference problem following Ref. [20]: the entropy bonus is the KL cost of deviating from the base dynamics, now without any assumption of Brownian noise. With Brownian noise, the same divergence reduces to the quadratic control cost of Eq. (51).

FIG. 10.

Eq. (185). As the reader can see, RL-notation is no joke.

While the above derivation gives the MaxEnt RL objective, it is unclear how to practically solve this maximization. In continuous time, we solved the KL control problem of Eq. (52) by dynamic programming, arriving at the HJB equation linearized using the Cole–Hopf transformation $V_t^{*}(x)=\text{log}{\psi }_t(x)$ (Eq. (58)). The same transformation in discrete time leads to soft Q-learning [108], the MaxEnt RL variant of Q-learning [105]. The central object is the soft Q function, the action-conditioned, maximum-entropy counterpart of the value function with the state and the first action both held fixed:

$$Q^{\pi }(x,u)=r(x,u)+{\mathbb{E}}_{P^{\pi }}\left[\sum _{j=n+1}^N\left(r\left(x_j,u_j\right)-\alpha \text{log}\pi \left(u_j\mid x_j\right)\right)\mid x_n=x,u_n=u\right].$$ (187)

The soft Q and soft value functions are related by

$$V^{\pi }(x)={\mathbb{E}}_{u~\pi (\cdot \mid x)}\left[Q^{\pi }\left(x,u\right)-\alpha \text{log}\pi (u\mid x)\right].$$ (188)

Here $Q^{\pi }$ excludes the entropy contribution of the fixed current action, while $V^{\pi }$ includes the current entropy contribution. By construction, $Q^{\pi }$ satisfies the “soft” Bellman Eq. (189),

$$Q^{\pi }(x,u)=r(x,u)+{\mathbb{E}}_{x^{\prime }~k(\cdot \mid x,u)}\left[{\mathbb{E}}_{u^{\prime }~\pi \left(\cdot \mid x^{\prime }\right)}\left[Q^{\pi }\left(x^{\prime },u^{\prime }\right)-\alpha \text{log}\pi \left(u^{\prime }\mid x^{\prime }\right)\right]\right].$$ (189)

with the boundary condition $Q^{\pi }(x,u)=r(x,u)$ at the final step $N$. Then, the optimal policy maximizes the conditional expectation on the right-hand side of Eq. (189), the expected Q value plus the policy entropy, independently at the next states $x^{\prime }$. Setting the variation with respect to $\pi \left(u^{\prime }\mid x^{\prime }\right)$ to zero gives a Boltzmann distribution over actions,⁸

$${\pi }^{*}(u\mid x)=\frac{e^{Q^{*}(x,u)/\alpha }}{\sum _v{e}^{Q^{*}(x,v)/\alpha }}.$$ (190)

Substituting ${\pi }^{*}$ into Eq. (189) then reveals an alternative form of the soft Bellman equation, assuming that the policy $\pi$ is the true Boltzmann distribution:

$$Q^{*}(x,u)=r(x,u)+{\mathbb{E}}_{x^{\prime }~k(\cdot \mid x,u)}\left[\alpha \text{log}\sum _{u^{\prime }}{e}^{Q^{*}\left(x^{\prime },u^{\prime }\right)/\alpha }\right].$$ (191)

Thus, if the policy ${\pi }^{*}$ is known exactly, fixing $(x,u)$ leaves the environment transition $k\left(x^{\prime }\mid x,u\right)$ as the only remaining source of randomness in this backup.

Eq. (191) is the discrete counterpart of the max-ent HJB equation after the Cole–Hopf transformation, providing a consistency equation whose right-hand side is a conditional expectation under the path measure. In continuous time, we solved the consistency equation with explicit knowledge of the dynamics and the reward. In RL, neither is known, but every observed transition $\left(x,u,r,x^{\prime }\right)$ provides a one-sample Monte Carlo estimate of the right-hand side, giving an update rule:

$$\begin{gathered} y=r+\alpha \text{log}\sum _{u^{\prime }}{e}^{{\widehat{Q}}^{*}\left(x^{\prime },u^{\prime }\right)/\alpha }, \\ {\widehat{Q}}^{*}(x,u)\leftarrow {\widehat{Q}}^{*}(x,u)+\xi \left[y-{\widehat{Q}}^{*}(x,u)\right]. \end{gathered}$$ (192)

Here, ${\widehat{Q}}^{*}$ is the learned estimate of $Q^{*}$ and $\xi$ is the learning rate. This update rule is the soft Q-learning algorithm (Algorithm 1), where the log-sum-exp is a softened maximum. As the temperature $\alpha \to 0$, the softened maximum in the log-sum-exp becomes a hard maximum, recovering standard tabular Q-learning [105].

Soft Q-learning [108] has formed the basis for many popular RL algorithms, in which a neural network $Q_{\theta }^{*}$ approximates the soft Q function using the update rule in Eq. (192). Implementation details, such as the use of replay buffers and target networks, are important for getting these algorithms to work in practice.

Algorithm 1.

Soft Q-learning

Initialize $Q_{\theta }^{*}$ and replay buffer $𝒟$; fix temperature $\alpha >0$ and learning rate $\xi$

for each iteration do

Collect experience: Run ${\pi }_{\theta }(u\mid x)\propto e^{Q_{\theta }^{*}(x,u)/\alpha }$, and store $\left(x,u,r,x^{\prime }\right)$ in $𝒟$

Sample past experience: $\left(x,u,r,x^{\prime }\right)~𝒟$

Q target: $y=r+\alpha \text{log}\sum _{u^{\prime }}{e}^{Q_{\theta }^{*}\left(x^{\prime },u^{\prime }\right)/\alpha }$

Update Q: $\theta \leftarrow \theta -\xi {\nabla }_{\theta }{\left(Q_{\theta }^{*}(x,u)-y\right)}^2$, holding $y$ fixed

end for

4. Soft actor-critic

Each update must still evaluate the log-sum-exp over every action, which is cheap for a few discrete actions. With continuous actions, the sum becomes an integral, and computing the Q-conditioned partition function $Z_Q(x)=\int \text{d}u{e}^{Q(x,u)/\alpha }$ can become intractable and high-variance as the action dimension grows due to the curse of dimensionality. This curse makes it difficult to evaluate the normalized Boltzmann policy ${\pi }_Q(u\mid x)=e^{Q(x,u)/\alpha }/Z_Q(x)$ defined by the current critic $Q$, or the soft value $\alpha \text{log}{Z}_Q\left(x^{\prime }\right)$ on the right-hand side of the backup in Eq. (191).

This leads to Soft Actor-Critic (SAC), which learns a policy together with the value estimate [98]. The iteration in Algorithm 2 alternates between updating the critic’s estimate of value under the current policy and improving the policy using that estimate, with entropy included in both steps. The parameterized actor ${\pi }_{\theta }$ aims to fit the true Boltzmann target ${\pi }_Q$ defined by the current critic $Q$. Simultaneously, the learned critic $Q_{\theta }^{\pi }$ aims to fit the value induced by the current policy ${\pi }_{\theta }$.

The unknown reward function and transition dynamics are only available through samples. Each observed transition supplies $r$ and $x^{\prime }$. Thus, the integral in $Z_Q\left(x^{\prime }\right)$ can nevertheless remain intractable even when the current critic $Q$ is directly evaluable. To address this integration problem, consider the Boltzmann policy ${\pi }_Q$ defined above. Suppose we could draw the next action $u^{\prime }$ from ${\pi }_Q\left(u^{\prime }\mid x^{\prime }\right)$ and evaluate its log density. A single draw would then replace the log-sum-exp, because substituting $\text{log}{\pi }_Q\left(u^{\prime }\mid x^{\prime }\right)=Q\left(x^{\prime },u^{\prime }\right)/\alpha -\text{log}{Z}_Q\left(x^{\prime }\right)$ gives

$$\begin{gathered} Q\left(x^{\prime },u^{\prime }\right)-\alpha \text{log}{\pi }_Q\left(u^{\prime }\mid x^{\prime }\right) \\ =Q\left(x^{\prime },u^{\prime }\right) \\ -\alpha \left[\frac{Q\left(x^{\prime },u^{\prime }\right)}{\alpha }-\text{log}{Z}_Q\left(x^{\prime }\right)\right] \\ =\alpha \text{log}{Z}_Q\left(x^{\prime }\right) \end{gathered}$$ (193)

for every sampled $u^{\prime }$. Thus, with an exact ${\pi }_Q\left(u^{\prime }\mid x^{\prime }\right)$, this difference is independent of the sampled action $u^{\prime }$ and provides a zero-variance estimate of the soft value function. In practice, we do not have access to the true Boltzmann ${\pi }_Q$, so the estimator is not variance-free. Soft actor-critic [98, 110] keeps the combination $Q\left(x^{\prime },u^{\prime }\right)-\alpha \text{log}{\pi }_{\varphi }\left(u^{\prime }\mid x^{\prime }\right)$ as the estimate of the soft value. Evaluating the normalized density ${\pi }_Q$ itself requires the same intractable partition function $Z_Q(x)$. Soft actor-critic therefore draws $u^{\prime }$ from a second network, a policy ${\pi }_{\varphi }(u\mid x)$ with a tractable form, trained to stand in for ${\pi }_Q$. This substitution turns Eq. (191) into the practical target used in the critic update in Algorithm 2

$$\begin{gathered} y=r+Q_{\theta }^{\pi }\left(x^{\prime },u^{\prime }\right)-\alpha \text{log}{\pi }_{\varphi }\left(u^{\prime }\mid x^{\prime }\right), \\ \text{where}{u}^{\prime }~{\pi }_{\varphi }\left(\cdot \mid x^{\prime }\right). \end{gathered}$$ (194)

where, like Soft Q-learning, the learned Q function $Q_{\theta }^{\pi }$ is directly regressed towards the target. Then, training ${\pi }_{\varphi }$ to maximize

$$\underset{\varphi }{\text{max}}{\mathbb{E}}_{u~{\pi }_{\varphi }(\cdot \mid x)}\left[Q_{\theta }^{\pi }\left(x,u\right)-\alpha \text{log}{\pi }_{\varphi }(u\mid x)\right].$$ (195)

is equivalent to minimizing ${\text{D}}_{\text{KL}}\left({\pi }_{\varphi }\Vert {\pi }_Q\right)$, since the actor objective equals $-\alpha {\text{D}}_{\text{KL}}\left({\pi }_{\varphi }\Vert {\pi }_Q\right)+\alpha \text{log}{Z}_Q\left(x_k\right)$ at fixed $x_k$. Thus, for a fixed critic $Q$, the policy objective corresponds to projecting the learned policy ${\pi }_{\varphi }$ onto the Boltzmann policy ${\pi }_Q$. However, in practice, the critic estimate $Q_{\theta }^{\pi }$ can be noisy. Then, targets built from a maximum, or a soft maximum, of the estimates inherit an upward maximization bias. Practical implementations reduce overestimation by training two Q networks and using the smaller of the two targets, adding back in some pessimism [98].

Algorithm 2.

Soft Actor-Critic

Initialize critic $Q_{\theta }^{\pi }$, actor ${\pi }_{\varphi }$, and replay buffer $𝒟$; fix temperature $\alpha >0$

for each iteration do

Collect experience: Run ${\pi }_{\varphi }$, observe $\left(x,u,r,x^{\prime }\right)$, and store in $𝒟$

Sample past experience: $\left(x,u,r,x^{\prime }\right)~𝒟$

Sample $u^{\prime }~{\pi }_{\varphi }\left(\cdot \mid x^{\prime }\right)$

Critic target: $y=r+Q_{\theta }^{\pi }\left(x^{\prime },u^{\prime }\right)-\alpha \text{log}{\pi }_{\varphi }\left(u^{\prime }\mid x^{\prime }\right)$

Update critic: $\theta \leftarrow \theta -\xi {\nabla }_{\theta }{\left(Q_{\theta }^{\pi }(x,u)-y\right)}^2$, holding $y$ fixed

Update actor:

$\varphi \leftarrow \varphi +\xi {\nabla }_{\varphi }{\mathbb{E}}_{v~{\pi }_{\varphi }(\cdot \mid x)}[\left.Q_{\theta }^{\pi }\left(x,v\right)-\alpha \text{log}{\pi }_{\varphi }(v\mid x)\right]$

end for

Solving a fixed-point equation by stochastic approximation raises the question of convergence. When the algorithm stores one Q value per state-action pair, the update converges to the optimal Q function [105, 111], and similar guarantees exist for simple function classes such as linear MDPs [111, 112]. In continuous or high-dimensional settings, however, the algorithm cannot enumerate every state-action pair and must instead approximate $Q^{*}$ with a parametric function such as a neural network. Unfortunately, the convergence guarantee does not generally survive function approximation [107], and the failure is visible in practice. A standard benchmark is mountain car [113], an underpowered car on a sinusoidal hill that must accelerate back and forth to build enough energy to reach the top. Even in mountain car, learned RL solutions show a large gap from the known optimal control solution [114]. Thus, modern RL methods learn controls from samples alone and approximate intractable objects, at the cost of the exactness guarantees of the earlier control methods.

B. Wasserstein gradient flows

In this section, we examine three classes of applied problems for which Wasserstein geometry and Wasserstein gradient flows provide elegant solutions. The first problem is the numerical integration of Fokker-Planck equations, in which the Wasserstein geometry yields a powerful integration algorithm known as the Jordan-Kinderlehrer-Otto (JKO) scheme. Second, we discuss modeling the dynamics of many interacting particles, particularly the regime known as the mean-field limit in which each particle effectively interacts with a continuous density. The collective evolution of the particle density often takes the form of Wasserstein gradient flow, and this perspective has yielded insights into the training of wide neural networks, in which the parameters can be thought of as particles. Finally, we discuss the general problem of variational inference (VI), in which a complex distribution is approximated by a simpler one that can be used for sampling. The problem of finding the optimal parameters can be viewed as an instance of Wasserstein gradient flow, and this perspective yields physical intuition and stable algorithms. As we will see, variational inference as discussed in this section forms an “equilibrium” methodology, whose “non-equilibrium” counterparts are powerful, contemporary algorithms discussed in Sec. V C on flows and diffusion.

1. The Jordan-Kinderlehrer-Otto (JKO) scheme

The JKO scheme is a canonical example of using Wasserstein geometry to tackle numerical challenges, resulting in algorithms that are often more efficient, stable, and interpretable. Consider the general problem of numerically solving Fokker-Planck equations (FPE) driven by a conservative flow and constant diffusivity, which take the form:

$$\frac{\partial {\rho }_t(x)}{\partial t}=\nabla \cdot \left[{\rho }_t(x)\nabla \Phi (x)\right]+{\nabla }^2{\rho }_t(x)$$ (196)

The unit diffusion coefficient corresponds to $\epsilon =2$ in the convention of Chapter I. It may be tempting to numerically solve Eq. (196) by an explicit finite step update:

$${\rho }^{(k+1)}(x)={\rho }^{(k)}(x)+\Delta t\left\{\nabla \cdot \left[{\rho }^{(k)}(x)\nabla \Phi (x)\right]\right.\left.+{\nabla }^2{\rho }^{(k)}(x)\right\}$$ (197)

where ${\rho }^{(k)}(x)={\rho }_{k\Delta t}(x)$ is the numerical discretization and $\Delta t$ is the step size. While conceptually straightforward, this direct approach and those like it encounter certain challenges. First, under Eq. (197), ${\rho }^{(k)}$ is not guaranteed to be either positive or normalized, which is particularly problematic if ${\rho }^{(k)}$ represents a probability distribution. Second, one must explicitly compute derivatives of ${\rho }^{(k)}(x)$, which is costly in high dimensions and can introduce instabilities if ${\rho }^{(k)}(x)$ naturally concentrates onto singular manifolds such as points or lines. Third, the algorithm only converges for sufficiently small $\Delta t$. Fourth, the continuous-time FPE takes the form of a Wasserstein gradient flow $v=-{\nabla }_{\text{W}}ℱ$, yet under Eq. (197), $ℱ$ need not monotonically decrease at finite $\Delta t$.

Jordan, Kinderlehrer, and Otto proposed an integration scheme that overcomes these shortcomings [45]. It has two ingredients. First, the algorithm is variational: it breaks the numerical integration into a sequence of optimization problems. For intuition, consider integrating an ODE $\dot{x}=-\nabla \Phi (x)$. A variational approach takes the form

$$x^{(k+1)}=\text{arg}\underset{x}{\text{min}}\left\{\frac{{\left|x-x^{(k)}\right|}^2}{2\Delta t}+\Phi (x)\right\}$$ (198)

$$\approx \text{arg}\underset{x}{\text{min}}\left\{\frac{{\left|x-x^{(k)}\right|}^2}{2\Delta t}+\Phi \left(x^{(k)}\right)\right.\left.+\left(x-x^{(k)}\right)\cdot \nabla \Phi \left(x^{(k)}\right)\right\}$$ (199)

$$=x^{(k)}-\Delta t\nabla \Phi \left(x^{(k)}\right)$$ (200)

which recovers the expected Euler update as $\Delta t\to 0$. As another example, the diffusion equation ${\partial }_t\rho ={\nabla }^2\rho$, i.e. Eq. (196) with $\Phi =0$, can readily be put in variational form:

$${\rho }^{(k+1)}=\text{arg}\underset{\rho }{\text{min}}\left(\frac{{‖\rho -{\rho }^{(k)}‖}_{L^2}^2}{2\Delta t}+\int |\nabla \rho (x){|}^2\text{d}x\right)$$ (201)

where $\Vert f{\Vert }_{L^2}^2=\int \text{d}x|f(x){|}^2$ is the squared $L^2$ norm of $f$. While Eq. (201) is certainly variational, it still suffers some of the practical and conceptual limitations of Eq. (197). Practically, one still needs to compute derivatives of $\rho$, and Eq. (201) can only capture Eq. (196) with $\nabla \Phi =0$, i.e., there is no way to amend the action of Eq. (201) to account for the drift term $\nabla \cdot (\rho \nabla \Phi )$. Conceptually, there is no reason to think of $\rho$ as a probability distribution in Eq. (201): both the $L_2$ norm and the functional in the second term make sense even when $\rho$ is non-normalized or negative. In this sense, Eq. (201) emphasizes a particular interpretation of the diffusion equation — it punishes gradients. But we also know that diffusion has a probabilistic interpretation: it promotes randomness, thereby increasing entropy. Is there a version of Eq. (201) that emphasizes this interpretation?

These shortcomings are addressed by the second ingredient of the JKO scheme. Namely, one replaces the $L_2$ norm in Eq. (201) with one that is adapted to probability: the Wasserstein distance. In doing so, one also replaces the gradient penalty with a thermodynamically meaningful quantity: the free energy. The resulting update scheme takes the form

$${\rho }^{(k+1)}=\text{arg}\underset{\rho }{\text{min}}\left(\frac{W_2^2\left({\rho }^{(k)},\rho \right)}{2\Delta t}+ℱ[\rho ]\right)$$ (202)

Eq. (202) has several advantages. First, it accommodates both drift and diffusion within a single variational principle, and it has a clear probabilistic interpretation: descend $ℱ$ as fast as possible while not changing $\rho$ too much in the Wasserstein sense. This change in conceptual approach also resolves the practical issues arising in Eq. (197): ${\rho }^{(k)}$ is guaranteed to remain a valid probability distribution for all $k$, the integration is stable, and $ℱ\left[{\rho }^{(k)}\right]$ is guaranteed to decrease monotonically for all $\Delta t$. Moreover, Eq. (202) has the advantage of being derivative-free — the gradients and divergences of Eqs. (197)–(201) are replaced by an optimal transport problem [115]. In particular, the JKO scheme can be brought to the Benamou-Brenier form [116] (see Sec. II B), which allows one to rigorously take the continuum ($\Delta t\to 0$) limit, as well as to introduce efficient optimization methods [115] such as the Sinkhorn algorithm or point-mass approximations. The clear geometric motivation behind Eq. (202) also results in better interpretability; for instance, one can write down explicit formulas for the implicit bias (i.e., effective modifications of $ℱ$) introduced at finite $\Delta t$ [117].

As a practical matter, the JKO scheme tends to shine in the following contexts. In very high dimensions, explicit spatial derivatives or spectral methods become computationally expensive, whereas solving the optimal transport problem at each time step remains comparatively efficient using particle-based or other methods. Another context where JKO is helpful is when $\rho (x)$ has finite or irregular support, which is difficult to accommodate with spectral or other grid-based methods but is naturally handled by optimal transport optimization. Finally, JKO handles pattern-forming PDEs well, particularly when mass tends to aggregate onto singular points, lines, or manifolds. In short, JKO offers a practical advantage in high dimensions, with irregular boundaries, or in mass-concentrating situations.

2. The mean-field limit of interacting particles

Now we study how WGF helps us understand the dynamics of the interaction of many identical particles, particularly the so-called mean-field limit, in which the dynamics of the individual particles can be replaced by the dynamics of a coarse-grained density. This limit is of fundamental interest in mathematical physics (e.g., how do continuum theories like fluid mechanics rigorously emerge from interacting particles?), and in practical applications, such as the training of wide neural networks, in which the neurons are viewed as interacting particles. In certain situations, the coarse-grained density obeys WGF, which is typically simpler to analyze than a full many-body system and offers geometric intuition that underlies convergence theorems. For example, this approach allows one to prove under what conditions training a wide two-layer NN converges to global minima of its loss.

The mathematical setting of interest is a collection of $N_{\text{p}}$ interacting particles with coordinates $x_t^{(1)},\dots ,x_t^{\left(N_{\text{p}}\right)}$, obeying the following Langevin dynamics:

$$\text{d}{x}_t^{(i)}=-{\nabla }_i{\Phi }^{\left(N_{\text{p}}\right)}\left(x_t^{(1)},\dots ,x_t^{\left(N_{\text{p}}\right)}\right)\text{d}t+\sqrt{\epsilon }\text{d}{W}_t^{(i)}$$ (203)

where $\text{d}{W}_t^{(i)}$ are independent Wiener processes and ${\nabla }_i$ means that the gradient is evaluated with respect to argument $i$. To permit a mean-field limit, we assume that the particles are “identical” in the sense that their potential is only a function of their unlabeled density. That is, we can write ${\Phi }^{\left(N_{\text{p}}\right)}\left(x_t^{(1)},\dots ,x_t^{\left(N_{\text{p}}\right)}\right)=𝒱\left[{\widehat{\rho }}_t\right]$ where

$${\widehat{\rho }}_t(x)=\frac{1}{N_{\text{p}}}\sum _{i=1}^{N_{\text{p}}}\delta \left(x-x_t^{(i)}\right)$$ (204)

is the empirical probability measure, normalized to one. We reserve ${\rho }_t$ for the continuous mean-field density and $N_{\text{p}}$ is the particle count. A common class of such interactions consists of single- and two-particle potentials:

$$\begin{gathered} {\Phi }^{\left(N_{\text{p}}\right)}\left(x^{(1)},\dots x^{\left(N_{\text{p}}\right)}\right)=\sum _i\Phi \left(x^{(i)}\right)+\frac{1}{2N_{\text{p}}}\sum _{i,j}U\left(x^{(i)},x^{(j)}\right) \\ =N_{\text{p}}\int \Phi (x)\widehat{\rho }(x)\text{d}x\frac{N_{\text{p}}}{2}\int U(x,y)\widehat{\rho }(x)\widehat{\rho }(y)\text{d}x\text{d}y \\ \equiv 𝒱\left[\widehat{\rho }\right] \end{gathered}$$ (205)

where we assume $U(x,x)=0$ to exclude self interactions. Now Let’s assume that the initial positions of the particles are drawn independently from a continuous density ${\rho }_0$, namely $x_0^{(i)}\stackrel{\text{iid}}{~}{\rho }_0$. The mean-field limit boils down to the following approximation: instead of tracking all the individual particles, one tracks the evolution of the continuous density. More precisely, one assumes that at a later time $t$, the particles are once again independently and identically distributed, but now according to an updated density ${\rho }_t$. A representative particle $x_t$, for which ${\rho }_t$ is the distribution, evolves according to

$$\text{d}{x}_t=-{\nabla }_{\text{W}}𝒱\left[{\rho }_t\right]\left(x_t\right)\text{d}t+\sqrt{\epsilon }\text{d}{W}_t$$ (206)

where ${\nabla }_{\text{W}}𝒱=\nabla (\delta 𝒱/\delta \rho )$ is the Wasserstein gradient introduced in Sec. II C. Eq. (206) is often referred to as a McKean-Vlasov process [118], and is simply a rewriting of Eq. (203) in which the particle label has been dropped, and the potential is expressed in terms of the continuous density ${\rho }_t$ instead of the empirical density ${\widehat{\rho }}_t$. Qualitatively, Eq. (206) mathematizes the intuition that as $N_{\text{p}}\to \infty$, the typical particle $x_t$ effectively interacts with an average continuous field ${\rho }_t$ of particles. The corresponding Fokker-Planck equation for ${\rho }_t$ is given by:

$$\begin{gathered} {\partial }_t{\rho }_t=\nabla \cdot \left({\rho }_t{\nabla }_{\text{W}}ℱ\left[{\rho }_t\right]\right) \\ =\nabla \cdot \left({\rho }_t{\nabla }_{\text{W}}𝒱\left[{\rho }_t\right]\right)+\frac{\epsilon }{2}{\nabla }^2{\rho }_t. \end{gathered}$$ (207)

where $ℱ=𝒱-\frac{\epsilon }{2}H$ and $H\left[\rho \right]=-\int \rho \text{log}\rho \text{d}x$ is the entropy. Eq. (207) is sometimes referred to as a McKean-Vlasov equation⁹, and it precisely takes the form of WGF. Whether and for what values of $N_{\text{p}}$ and $t$ the mean-field approximation is a good approximation is a field of mathematical research known as propagation of chaos [119].

As a practical example, the WGF structure of Eq. (207) has been useful in analyzing the training dynamics of large machine learning (ML) models such as neural networks (NNs) [120–123]. Large ML models are often trained by defining a loss function and performing versions of gradient descent, e.g., Stochastic Gradient Descent (SGD) or Adaptive Moment Estimation (Adam), in their parameters. For a large but finite number of parameters, the loss landscape typically features many saddles, large manifolds of degenerate minima, and different minima-manifolds separated by shallow barriers. These features make the convergence of gradient descent difficult to analyze. However, when many of the parameters enter the loss identically, such as in a very wide NN, an $N_{\text{p}}\to \infty$ limit can be taken, allowing one to analyze the evolution of the density ${\rho }_t$ of the parameters over training time. The density ${\rho }_t$ generally obeys WGF and the functional $ℱ$ that it minimizes is generally much simpler than the loss landscape governing the evolution of a finite number $N_{\text{p}}$ of parameters. This simplification allows convergence results to be obtained. It should be noted that there is an older collection of results, sometimes referred to as Universal Approximation Theorems [124–126], that ask which classes of functions can be approximated by a given ML architecture. The question at stake here is perhaps even more relevant: under what conditions are these functions actually reached via training?

To illustrate typical ingredients of such an approach, consider a two-layer neural network

$$f^{N_{\text{p}}}\left(x;{\Theta }_t\right)=\frac{1}{N_{\text{p}}}\sum _{i=1}^{N_{\text{p}}}{w}_t^{(i)}\phi \left(x;z_t^{(i)}\right)$$ (208)

In Eq. (208), $f^{N_{\text{p}}}\left(x;{\Theta }_t\right)$ is the model to be trained. Here $\phi \left(x;z_t^{(i)}\right)$ is a nonlinear function, referred to as a “feature”, with weight $w_t^{(i)}$ and internal parameters $z_t^{(i)}$. All the parameters of $f^{N_{\text{p}}}$ may be collected in a single vector ${\Theta }_t=\left({\theta }_t^{(1)},\dots ,{\theta }_t^{\left(N_{\text{p}}\right)}\right)$ where ${\theta }_t^{(i)}=\left(w_t^{(i)},z_t^{(i)}\right)$. The parameters are given a time index $t$ because they will evolve over training. The function $f^{N_{\text{P}}}\left(x;{\Theta }_t\right)$ is trained to approximate a target function $y(x)$ with $x$ distributed according to a density $\nu (x)$. This is done by minimizing a loss function $L\left({\Theta }_t\right)$, a common choice for which is the mean-squared error (MSE):

$$\begin{gathered} L\left({\Theta }_t\right)=\frac{1}{2}\int {\left|y(x)-f^{N_{\text{p}}}\left(x;{\Theta }_t\right)\right|}^2\nu (x)\text{d}x \\ =\frac{1}{2}{\mathbb{E}}_{\nu }\left[{\left|y-f_t^{N_{\text{p}}}\right|}^2\right] \end{gathered}$$ (209)

where $f_t^{N_{\text{p}}}=f^{N_{\text{p}}}\left(\cdot ;{\Theta }_t\right)$ for brevity. The loss is minimized by using some form of gradient descent, typically with stochasticity. For instance

$$\text{d}{\theta }_t^{(i)}=-N_{\text{p}}{\nabla }_{i}L\text{d}t+\sqrt{\epsilon }\text{d}{W}_t^{(i)}$$ (210)

The factor $N_{\text{p}}$ compensates for the $1/N_{\text{p}}$ in the network output: differentiating $L$ with respect to one neuron’s parameters produces a factor $1/N_{\text{p}}$. Multiplying the learning rate by $N_{\text{p}}$ therefore keeps each neuron’s drift of order one as the network grows. In practice, more practical forms of gradient descent are used, for instance stochastic gradient descent in which stochasticity does not come from a Wiener process $\text{d}{W}_t^{(i)}$, but from using finite sample-size Monte Carlo estimates of the loss. Regardless of the precise method, the loss $L$ is a complicated, nonconvex function of the parameters ${\Theta }_t$, so proving (or understanding) anything about training may at first seem to be a daunting task.

To gain intuition, we can think of each neuron $i$ as being a particle with coordinate ${\theta }_t^{(i)}$, and (noisy) gradient descent, e.g. Eq. (210), as introducing interactions between the particles. For the two-layer neural network, the key insight is that Eq. (208) can be written as a function of neuron density only:

$$f^{N_{\text{p}}}\left(x;{\Theta }_t\right)=\int a(x;\theta ){\widehat{\rho }}_t(\theta )\text{d}\theta$$ (211)

where $\theta =(w,z),a(x;\theta )=w\phi (x;z)$, and ${\widehat{\rho }}_t(\theta )=\frac{1}{N_{\text{p}}}\sum _{i=1}^{N_{\text{p}}}\delta \left(\theta -{\theta }_t^{(i)}\right)$ is the empirical distribution of the parameters. Because $f^{N_{\text{p}}}$ only depends on the density of parameters, the loss $L\left({\Theta }_t\right)$ can also be expressed in terms of the density alone. For example, the MSE becomes:

$$\begin{gathered} L\left({\Theta }_t\right)=ℒ\left[{\widehat{\rho }}_t\right]=\int \Phi (\theta ){\widehat{\rho }}_t(\theta )\text{d}\theta \\ +\frac{1}{2}\int U\left(\theta ,{\theta }^{\prime }\right){\widehat{\rho }}_t(\theta ){\widehat{\rho }}_t\left({\theta }^{\prime }\right)\text{d}\theta \text{d}{\theta }^{\prime }+C_y \end{gathered}$$ (212)

where $\Phi (\theta )=-{\mathbb{E}}_{\nu }[y(x)a(x;\theta )]$ and $U\left(\theta ,{\theta }^{\prime }\right)={\mathbb{E}}_{\nu }\left[a(x;\theta )a\left(x;{\theta }^{\prime }\right)\right]$ are single- and two-particle potentials, and $C_y=\frac{1}{2}{\mathbb{E}}_{\nu }\left[y(x{)}^2\right]$ is independent of $\theta$. The fact that the interactions only depend on the density $\widehat{\rho }(\theta )$ makes the convergence of training dynamics amenable to mean-field-based proof techniques.

The exact convergence statements proven in the literature [120–123] require some mathematical formalism to state precisely, and depend on the exact technical hypotheses placed on the features $\phi$, the loss $L$, and the training algorithm. There are, however, some shared ingredients: one typically assumes that the network is “initialized” by drawing initial parameters, ${\theta }_0^{(i)}$, independently from a well-behaved distribution ${\rho }_0$, typically obeying some precise regularity conditions. One is then interested in the evolution of the coarsegrained distribution ${\rho }_t$, which obeys a McKean-Vlasov equation, whose precise form depends on the exact architecture, loss, or version of gradient descent under consideration. Typically, the McKean-Vlasov equation can be interpreted in terms of WGF of a functional $ℱ,{\partial }_t{\rho }_t=\nabla \cdot \left({\rho }_t{\nabla }_{\text{W}}ℱ\left[{\rho }_t\right]\right)$, where $ℱ$ consists of the loss $ℒ$ plus entropic terms arising from stochasticity. The dynamical problem of training the neural network is then mathematically transformed into studying the comparatively simple geometric and topological properties of $ℱ$. While the functional $ℱ$ is not necessarily geodesically convex in the Wasserstein sense, it often permits enough structure to make mathematical statements about convergence to global minima for ${\rho }_0$ obeying suitable conditions. The conditions placed on ${\rho }_0$ for convergence are of some practical use, because they can be interpreted as instructions on reasonable ways to initialize the NN for training. Intuitively, WGF is the right mathematical framework for analyzing these problems because of the underlying physical picture: the parameters ${\theta }_t^{(i)}$ are particles, but these particles cannot jump, or be created or destroyed. Therefore, continuous motion of a density, i.e. WGF, naturally emerges.

To give a sense of the results proven, Ref. [122] considers a quadratic loss and training via either stochastic gradient descent or noisy stochastic gradient descent; for the latter case, they prove that any absolutely continuous ${\rho }_t$ converges to a unique global minimum of $ℱ$ that controllably approximates the MSE. They also prove that $f_t^{N_{\text{p}}}$ converges to the $N_{\text{p}}\to \infty$ limit with bounds on the rate of convergence as a function of $N_{\text{p}}$ and training time $t$. Ref. [121] considers models of the form of Eq. (208) with technical hypotheses on the forms of the nonlinearity and data distribution, and identifies convergence to a global minimum, and studies the size of fluctuations for finite $N_{\text{p}}$ and $t$. Ref. [120] provides convergence proofs for exact gradient descent for a class of architectures and loss functions defined by the property that the functional $ℱ$ they give rise to obeys certain homogeneity conditions.

A couple of additional general remarks are in order. First, the mean-field approach based on a single density of neurons works for a two-layer neural network, but deeper networks require additional information about the connections between layers. Mean-field methods have been extended to these settings, with different assumptions and conclusions [127, 128]. Second, there is a common source of confusion: one observes that $ℒ$ in Eq. (212) is manifestly convex in the parameter density $\rho$, i.e. $ℒ\left[\lambda {\rho }^{\prime }+(1-\lambda )\rho \right]\le \lambda ℒ\left[{\rho }^{\prime }\right]+(1-\lambda )ℒ\left[\rho \right]$ for any two densities $\rho ,{\rho }^{\prime }$ and $\lambda \in [0,1]$. While true, this mixture-interpolation notion of convexity is not the useful one for training based on gradient descent. Under gradient descent, the parameters flow along Wasserstein gradients, not weighted mixtures, and they certainly do not teleport (as is sometimes required for interpolation). Therefore, convexity with respect to Wasserstein geodesics is the useful mathematical framework to use, and one cannot directly conclude that ${\rho }_t$ converges to a unique global minimum under training due to interpolation convexity. Finally, one might wonder: why not freeze the internal parameters $z_t^{(i)}$ and only optimize the feature weights $w_t^{(i)}$’s? After all, the MSE loss $L\left({\Theta }_t\right)$ is explicitly convex in $w_t^{(i)}$. This approach, often referred to as kernel regression, is certainly feasible, but it suffers from two drawbacks. First, a much larger $N_{\text{p}}$ is required for kernel regression to approximate the $N_{\text{p}}=\infty$ global minimum [120]. Second, kernel regression exhibits a different implicit bias: when both the $w$’s and the $z$’s are dynamic, it is observed that the network learns comparatively sparse distributions for the $w_t^{(i)}$’s, and the selected features $\phi \left(x;z_t^{(i)}\right)$ have been updated via $z_t^{(i)}$, and thus tend to be informative about the learned properties of the data. When only the $z_t^{(i)}$’s are frozen, the learned distribution over the $w$’s tends to be broader and, by construction, no individual feature $\phi \left(x;z_t^{(i)}\right)$ reveals properties about the data. Whether ingredients from the mean-field approach can be extended to other architectures remains an open question.

3. Variational inference

Our final example of applications of WGF is to variational inference (VI) [129]. VI addresses the following challenge: one is confronted with a probability density $p_1={\tilde{p}}_1/Z_1$ that is difficult to sample from and whose unnormalized form ${\tilde{p}}_1$ may be all that is available; however, we assume that we have access to a family of distributions $p_{\theta }$ defined by a finite number of parameters $\theta$, and that the $p_{\theta }$ are easy to normalize and sample from, e.g. the family of Gaussians. Our goal is to find the optimal parameters ${\theta }^{*}$ such that $p_{{\theta }^{*}}$ best approximates $p_1$. This arises, for example, in Bayesian settings, in which one has a prior $p^0(x)$ and observations $A$, and one would like to sample from the posterior $p_1(x)=p(x\mid A)=p(A\mid x)p^0(x)/p^0(A)$, where the normalization constant $p^0(A)=\int p(A\mid x)p^0(x)\text{d}x$ is often intractable. It is also thematically related to variational methods in quantum mechanics or statistical physics, in which the ground state of a model may be approximated by an ansatz featuring only a few parameters.

For the connection to WGF, we are interested in the setting where “best approximation” is defined as minimizing the Kullback-Leibler divergence

$${\theta }^{*}=\underset{\theta }{\text{arg}\text{min}}{\text{D}}_{\text{KL}}\left(p_{\theta }\Vert p_1\right)$$ (213)

As noted earlier in the review, the gradient of ${\text{D}}_{\text{KL}}\left(p_{\theta }\Vert p_1\right)$ with respect to $\theta$ is independent of $Z_1=\int {\tilde{p}}_1\text{d}x$, which is often infeasible to compute. One tempting option to minimize Eq. (213) is to simply perform gradient descent in the parameters, i.e.

$$\dot{\theta }=-{\nabla }_{\theta }L(\theta )$$ (214)

with $L(\theta )={\text{D}}_{\text{KL}}\left(p_{\theta }\Vert p_1\right)$. However, in practice, this naive parameter-centric approach is often problematic. The loss $L(\theta )$ may be poorly conditioned and nonconvex. Moreover, certain families of probability distributions are only well defined on a subset of parameters; for instance, a Gaussian distribution is only well defined if its covariance matrix is positive-definite. Unfortunately, Eq. (214) may lead to violations of these constraints. Finally, if one were to reparameterize the same family of probability distributions, Eq. (214) would yield a different trajectory through probability space—so Eq. (214) lacks parameterization invariance.

Wasserstein geometry is useful in addressing these issues. For convenience, write $p_1(x)=e^{-\Phi (x)}/Z_1$, with dimensionless potential $\Phi =-\text{log}{\tilde{p}}_1$, and write

$$\begin{gathered} ℒ\left[p_{\theta }\right]={\text{D}}_{\text{KL}}\left(p_{\theta }\Vert p_1\right)=\int \Phi (x)p_{\theta }(x)\text{d}x \\ +\int p_{\theta }(x)\text{log}{p}_{\theta }\left(x\right)\text{d}x+\text{log}{Z}_1. \end{gathered}$$ (215)

It seems that one approach to minimizing ${\text{D}}_{\text{KL}}\left(p_{\theta }\Vert p_1\right)$ would be to perform Wasserstein gradient flow down the functional $ℒ\left[p_{\theta }\right]$, namely update $p_{\theta }$ according to $\nabla \cdot p_{\theta }{\nabla }_{\text{W}}ℒ\left[p_{\theta }\right]$. Yet, there is a fundamental problem: there is no guarantee that one can find a trajectory of parameters ${\theta }_t$ such that $\dot{\theta }\cdot {\partial }_{\theta }{p}_{\theta }=\nabla \cdot \left(p_{\theta }{\nabla }_{\text{W}}ℒ\left[p_{\theta }\right]\right)$. In other words, the naive WGF may induce a flow field that flows out of the subspace of probability distributions defined by the parameters $\theta$.

The resolution is to project the Wasserstein descent velocity onto the tangent space of the parametric family. Recall from Sec. II C that a density tangent is represented by a minimum-energy velocity $v=-\nabla \chi$, so ${\partial }_t{p}_t=-\nabla \cdot \left(p_t{v}_t\right)$. The inner product of two such velocities at density $p$ is

$${⟨v,v^{\prime }⟩}_p=\int p\left(x\right)v\left(x\right)\cdot v^{\prime }\left(x\right)\text{d}x.$$ (216)

Consider a parametric family $ℳ=\left\{p_{\theta }:\theta \in \Omega \right\}$, where $\Omega$ is a parameter domain of dimension $d_{\theta }$. In local coordinates $\theta =\left({\theta }_1,\dots ,{\theta }_{d_{\theta }}\right)$, define velocities $v_i=-\nabla {\chi }_i$ by ${\partial }_{{\theta }_i}{p}_{\theta }=-\nabla \cdot \left(p_{\theta }{v}_i\right)$. They induce the metric $g_{ij}={⟨v_i,v_j⟩}_{p_{\theta }}$. The projection of a velocity field $v$ is

$${\text{proj}}_{ℳ}v=\sum _{i,j}{v}_i{\left(g^{-1}\right)}_{ij}{⟨v_j,v⟩}_{p_{\theta }}.$$ (217)

The projected descent velocity is $-{\text{proj}}_{ℳ}{\nabla }_{\text{W}}ℒ\left[p_{\theta }\right]$, and its density evolution is

$${\partial }_t{p}_{{\theta }_t}=\nabla \cdot \left(p_{{\theta }_t}{\text{proj}}_{ℳ}{\nabla }_{\text{W}}ℒ\left[p_{{\theta }_t}\right]\right).$$ (218)

Writing $L(\theta )=ℒ\left[p_{\theta }\right]$, the same evolution in parameter coordinates is

$${\dot{\theta }}_i=-\sum _{j=1}^{d_{\theta }}{\left(g^{-1}\right)}_{ij}{\partial }_{{\theta }_j}L(\theta ).$$ (219)

Comparing Eq. (219) to Eq. (214), the metric induced by Wasserstein geometry essentially acts to redirect the gradients with respect to $\theta$. Moreover, one can show that Eq. (219) is equivalent to a parametric version of the JKO style update introduced in Sec. V B 1

$${\theta }_{t+\Delta t}=\text{arg}\underset{\theta }{\text{min}}\left\{\frac{W_2^2\left(p_{{\theta }_t},p_{\theta }\right)}{2\Delta t}+ℒ\left[p_{\theta }\right]\right\}$$ (220)

in the limit $\Delta t\to 0$.

So what does all this formalism buy us? The minimization algorithm in Eq. (218), or equivalently Eq. (219), comes with elegant convergence guarantees. Recall from Sec. II C that the functional $ℒ$ can have a notion of convexity (namely $\alpha$-geodesic convexity) completely analogous to convexity in finite-dimensional optimization, and that this notion of convexity ensures that WGF converges to its global minimum. Recall also from Sec. II C that $ℳ$ is a geodesically convex subset of probability distributions if it contains the entire Wasserstein geodesic between any two of its elements. One can show that whenever $ℳ$ is geodesically convex and $ℒ$ is $\alpha$-geodesically convex, the projected WGF in Eq. (218) is guaranteed to reach its unique minimum on $ℳ$, and do so exponentially fast. Especially important in the context of VI is that the loss $ℒ$ in Eq. (215) is in fact $\alpha$-geodesically convex whenever $p_1$ is strongly log-concave, which means that it is written as $p_1\propto e^{-\Phi }$ for some $\alpha$-convex $\Phi$ (see Sec. II C for the definition of $\alpha$-convex). So in short, one can use the structure of Wasserstein geometry to turn the problematic optimization in Eq. (214) into a well-behaved one.

Let’s explore two pedagogical examples of parametric subspaces. First, take $ℳ$ to be the set of all non-degenerate Gaussians on ${\mathbb{R}}^d$, which goes by the name Bures-Wasserstein space $\text{BW}\left({\mathbb{R}}^d\right)$ [130]. Each distribution in this family can be represented by the mean $m\in {\mathbb{R}}^d$ and covariance matrix $\Sigma \in {\mathbb{S}}_{++}^d$, where ${\mathbb{S}}_{++}^d$ is the set of all symmetric, positive-definite $d\times d$ matrices. To see that the optimization problem is nontrivial, one can check that $L(m,\Sigma )$ defined by Eq. (215) is not necessarily convex in the parameters $(m,\Sigma )$ even when $\Phi (x)$ is $\alpha$-convex in $x$—for example, consider $\Phi (x)=x^4$ in $d=1$ dimension. Now Let’s apply the machinery of WGF and see what the dynamical optimization equations for $m$ and $\Sigma$ look like. First, Let’s check that $\text{BW}\left({\mathbb{R}}^d\right)$ is geodesically convex. One can show that the Wasserstein geodesic between any two Gaussians $p_0=𝒩\left(m_0,{\Sigma }_0\right)$ and $p_1=𝒩\left(m_1,{\Sigma }_1\right)$ is given by the transport map

$$\begin{gathered} x_t=(1-t)x_0+t\left[{\Sigma }_0^{-1/2}{\left({\Sigma }_0^{1/2}{\Sigma }_1{\Sigma }_0^{1/2}\right)}^{1/2}\right. \\ \left.{\Sigma }_0^{-1/2}\left(x_0-m_0\right)+m_1\right] \end{gathered}$$ (221)

Notably, Eq. (221) is an affine transformation for all $t$. Recall that any affine transformation of a normal distribution is a normal distribution, so the $p_t\in \text{BW}\left({\mathbb{R}}^d\right)$ for all $t$, confirming that $\text{BW}\left({\mathbb{R}}^d\right)$ is geodesically convex. Next, we would like to find an explicit expression for the projected Wasserstein gradient ${\nabla }_{\text{BW}}ℱ={\text{proj}}_{\text{BW}\left({\mathbb{R}}^d\right)}{\nabla }_{\text{W}}ℱ$ on a functional $ℱ$. Because all Gaussians are connected by affine maps, and all affine maps take Gaussians to Gaussians, it follows that the tangent space to $\text{BW}\left({\mathbb{R}}^d\right)$ corresponds to all affine maps. Using this fact and a bit of algebra, one finds that:

$$\left({\nabla }_{\text{BW}}ℱ\left[p\right]\right)(x)={⟨{\nabla }^2\frac{\delta ℱ}{\delta p}⟩}_p(x-m)+{⟨\nabla \frac{\delta ℱ}{\delta p}⟩}_p.$$ (222)

Here ${\nabla }^2$ in a matrix expression denotes the Hessian; the Laplacian is its trace. The parameter dynamics are

$${\dot{m}}_t=-{⟨\nabla \frac{\delta ℱ}{\delta p}⟩}_{p_t},$$ (223)

$${\dot{\Sigma }}_t=-{⟨{\nabla }^2\frac{\delta ℱ}{\delta p}⟩}_{p_t}{\Sigma }_t-{\Sigma }_t{⟨{\nabla }^2\frac{\delta ℱ}{\delta p}⟩}_{p_t}.$$ (224)

For the special case of VI, we take $ℱ={\text{D}}_{\text{KL}}\left(p\Vert p_1\right)$ and the flow equations become

$${\dot{m}}_t=-⟨\nabla \Phi {⟩}_{p_t},$$ (225)

$${\dot{\Sigma }}_t=-{⟨{\nabla }^2\Phi ⟩}_{p_t}{\Sigma }_t-{\Sigma }_t{⟨{\nabla }^2\Phi ⟩}_{p_t}+2I.$$ (226)

Eqs. (225)–(226) can be implemented straightforwardly because the expectations are taken with respect to the Gaussian measure $p_t=𝒩\left(m_t,{\Sigma }_t\right)$. In the naive approach, Eq. (214), of directly optimizing with respect to $m$ and $\Sigma$, the equations for the mean ${\dot{m}}_t$ are unchanged, while Eq. (226) for ${\dot{\Sigma }}_t$ differs in its dependence on ${\Sigma }_t$. With the standard Frobenius inner product for covariance matrices, ordinary gradient descent gives:

$${\dot{\Sigma }}_t=-\frac{1}{2}{⟨{\nabla }^2\Phi ⟩}_{p_t}+\frac{1}{2}{\Sigma }_t^{-1},$$ (227)

which can create situations in which $\Sigma$ becomes singular (loses its positive definiteness) during optimization. The geometric perspective of WGF avoids this pathology, makes the convex nature of the problem more transparent, and leads to other generalizations and principled convergence guarantees [131, 132].

As a second, theoretically rich, example, suppose $ℳ$ is the set of all discrete measures comprising $N_{\text{p}}$ point masses: ${\widehat{p}}_t(x)=\frac{1}{N_{\text{p}}}\sum _{i=1}^{N_{\text{p}}}\delta \left(x-x_t^{(i)}\right)$ [133, 134]. Here, the parameters $x_t^{(i)}$ can be thought of as point particles, whose positions need to be optimally situated to approximate the target distribution $p_1$. Upon inspection, this is exactly the setting of interacting particles considered in Sec. V B 2. Therefore, using the loss in Eq. (215), the particles update according to

$$\text{d}{x}_t^{(i)}=-\nabla \Phi \left(x_t^{(i)}\right)\text{d}t+\sqrt{2}\text{d}{W}_t^{(i)}$$ (228)

where the noise convention is $\epsilon =2$, as in Chapter IV. This is simply traditional MCMC sampling (see Chapter IV). However, there is a conceptually interesting point: Unlike Eqs. (225)–(226), or more generally, Eq. (219), the evolution in Eq. (228) is stochastic, not deterministic. This is ultimately connected to the fact that the discrete measure ${\widehat{p}}_t$ is singular, so gradients of the entropic contribution to Eq. (215) are ambiguous. While stochastic updates are not inherently problematic, this has led to a conceptual question: are there purely deterministic equations for the $x_t^{(i)}$ that would nevertheless give rise to an equivalent approximation in the $N_{\text{p}}\to \infty$ limit? In the language of Sec. V B 2, is there a deterministic version of Eq. (228) that gives rise to the same mean-field limit in Eq. (207)? An intuitive attempt would be to replace the delta functions $\delta \left(x-x_t^{(i)}\right)$ with a smoother function $k_h\left(x-x_t^{(i)}\right)$ of width $h$, which gives the evolution equations

$$\begin{gathered} {\dot{x}}_t^{(i)}=-\nabla \Phi \left(x_t^{(i)}\right) \\ -{\left.{\nabla }_x\text{log}\left[\frac{1}{N_{\text{p}}}\sum _j{k}_h\left(x-x_t^{(j)}\right)\right]\right|}_{x=x_t^{(i)}} \end{gathered}$$ (229)

where the density gradient differentiates the evaluation argument $x$ with particle centers held fixed, before setting $x=x_t^{(i)}$. The continuum limit $N_{\text{p}}\to \infty$ yields

$${\partial }_t{p}_t=\nabla \cdot \left[p_t\nabla \left(\Phi +\text{log}\left(k_h\star p_t\right)\right)\right]$$ (230)

where $\left(k_h\star p_t\right)(x)=\int k_h(x-y)p_t(y)\text{d}y$ denotes convolution. Notably, this approach is not entirely straightforward because $p_t$ in Eq. (230) does not converge to $p_1$ unless $h=0$, which raises a subtle question of limits for $N_{\text{p}}\to \infty$ and $h\to 0$. This subtlety motivates an alternative, more recent, approach known as Stein variational gradient descent [133, 134]. The basic idea is to put the regularization in a different location:

$${\partial }_t{p}_t=\nabla \cdot \left[p_t{𝒦}_h^{p_t}\left(\nabla \text{log}\frac{p_t}{p_1}\right)\right]$$ (231)

In Eq. (231), ${𝒦}_h^p$ implements convolution weighted by $p$. Namely, given $f(x)$, the smoothed version is

$$\begin{gathered} {𝒦}_h^p(f)(x)=\left(k_h\star (pf)\right)(x) \\ =\int f(y)p(y)k_h(x-y)\text{d}y \end{gathered}$$ (232)

Interestingly, the smoothing operator ${𝒦}_h^p$ can also be interpreted as a projection of the Wasserstein gradient ${\nabla }_{\text{W}}ℱ=\nabla \text{log}\left(p_t/p_1\right)$ [134]; however, this projection is distinct from ${\text{proj}}_{ℳ}$ in Eq. (218); it instead projects the velocity onto the reproducing kernel Hilbert space associated with the kernel $k_h$, see e.g. [134] for more detail. The discretized version of Eq. (231) takes the form:

$$\begin{gathered} {\dot{x}}_t^{(i)}=-\frac{1}{N_{\text{p}}}\sum _{j=1}^{N_{\text{p}}}\left[k_h\left(x_t^{(i)}-x_t^{(j)}\right)\nabla \Phi \left(x_t^{(j)}\right)\right. \\ \left.+\nabla k_h\left(x_t^{(i)}-x_t^{(j)}\right)\right] \end{gathered}$$ (233)

where $\nabla \Phi \left(x_t^{(j)}\right)$ differentiates the potential at particle $j$, while $\nabla k_h(z)$ differentiates the kernel argument $z=x_t^{(i)}-x_t^{(j)}$. Notice that Eq. (233) is straightforward to implement, and Eq. (231) genuinely permits $p_1$ as a stationary solution. Yet, questions surrounding its convergence remain a topic of investigation [135–139]. Nonetheless, Eq. (233) has a clear physical interpretation: the first term captures gradient descent, while the second is a particle repulsion encouraging exploration of the full potential.

In addition to these two examples, several other parametric spaces for VI have received attention. For example, the treatment of Bures-Wasserstein space above can be extended to Gaussian mixtures by considering measures over $\text{BW}\left({\mathbb{R}}^d\right)$ itself [131, 140, 141]¹⁰. Another powerful approach is to consider a mean-field approximation (not to be confused with the mean-field limit of interacting particles) in which $ℳ={\otimes }_{i=1}^d{𝒲}_2(\mathbb{R})\subset {𝒲}_2\left({\mathbb{R}}^d\right)$ is the space of products of one-dimensional distributions [43, 142, 143]. Across these examples, VI benefits from theoretical guarantees when $\Phi (x)$ is $\alpha$-convex. Much modern interest is in situations in which $\Phi (x)$ is not convex, but may be multi-modal and feature low-dimensional structure. Building models that are effective for sampling in this more advanced context motivates the next section on flow and diffusion (Sec. V C).

C. Generative modeling with flows and diffusion

In Chapter IV on Sampling, we explored the following problem: given a (potentially unnormalized) probability distribution $p(x)$, how to draw a set of $N$ independent samples ${\left\{x_n\right\}}_{n=1}^N$? In this section, we ask the inverse problem: given a set of samples ${\left\{x_n\right\}}_{n=1}^N$, how to infer the probability distribution $p(x)$ and generate a new independent sample $x_{N+1}$ from the estimated distribution? The former task is known as density estimation, and the latter task is generative modeling. Like many inverse problems, generative modeling is a somewhat ill-defined task: one aims to generate samples similar, but not identical to those from the training set (an issue referred to as “generalization versus memorization”). Generative modeling is an important problem: In many contexts, like protein structures [144] or natural images [145], training samples are abundant, but unlike in physics or chemistry, an explicit model for the data is not available. Over the last ten years, deep-learning-based generative models have proven astonishingly successful at producing samples from such complex, high-dimensional distributions. This section provides an overview of these new methods, their utility for scientific simulations, and their connection to the ideas around sampling, transport, and non-equilibrium physics discussed in Chapters I–IV. Note that we do not attempt a systematic literature review.

Like for sampling, the key difficulty in generative modeling is the high dimension $d$ of the sample space. A typical example is image generation, where $d\approx 200,000$ (for 256 × 256-pixel color images). Learning high-dimensional probability distributions is difficult, since the training samples only cover a vanishing fraction of high-dimensional space (the “curse of dimensionality”). Vice versa, a generative model needs to target this vanishing fraction to succeed. Put flippantly, randomly choosing pixel values has vanishing probability of producing an image of a cat (Fig. 11). In addition, complex probability distributions are typically multi-modal. As for sampling, modern designed models generate samples from a high-dimensional target distribution $p_1$ (e.g., the space of images) by dynamically transforming a simple base distribution $p_0$ (e.g., the standard Gaussian). We already encountered such transport problems in Chapter II on optimal transport. Like optimal transport, generative models are recipes for transforming probability distributions, now designed to make it easy to learn the transformation from training data. Diffusion models [145, 146] are perhaps the most well-known machine learning methods of this type. They learn a stochastic process that generates samples from a target distribution defined by a set of training samples. Diffusion models were originally inspired by the physical ideas of non-equilibrium thermodynamics and simulated annealing discussed in Chapter IV [146]. This section approaches generative modeling using the (effectively equivalent) frameworks of stochastic interpolation [147] and flow matching [148]. We believe that this choice will make the core idea behind diffusion models easier to understand and generalize.

FIG. 11.

Not a cat.

1. Variational transform sampling

Fundamentally, this section is based on the method of transform sampling, which applies a map $T$ to transform a sample $x~p_0$ into $y=T(x)$. Throughout this section, we use “$y$” to denote samples generated by a transform. The map $T$ can also be thought of as a “decoder”, and $x$ as a “hidden variable” [149]. The probability distribution of $y$ is called the pushforward $T_{*}{p}_0$. By changing variables, the probability density of $T_{*}{p}_0$ is:

$$\left(T_{*}{p}_0\right)(y)=p_0\left(T^{-1}(y)\right)\cdot \left|\text{det}{\nabla }_y{T}^{-1}\right|$$ (234)

For example, the linear map $y=M\cdot x$ transforms a standard Gaussian into one with covariance matrix $\Sigma =M{M}^{\top }$. A second example is provided by one-dimensional distributions, where the map $T$ is defined by the cumulative distribution functions of $p_0,p_1$ (Eq. (88)).

The algorithms in this section combine transform sampling with variational inference (VI, Sec. V B 3). VI approximates a target distribution $p_1$ by a member $q$ of a parametrized family. Here, the trial family is defined by a set of transformations $\widehat{T}$ via Eq. (234). VI is an intuitive match for modern ML techniques like expressive neural networks and largescale, gradient-based optimization. To bring these techniques to bear, one needs a loss function $ℒ$ to measure the difference between the estimated $q$ and the target distribution $p_1$. Here, our goal is to build loss functions that can be efficiently approximated using training samples, without explicit knowledge of $p_1$. Then, the transformation $\widehat{T}$ can be learned directly from data by minimizing $ℒ$.

2. Flow matching

The key idea of flow matching is to construct the transformation $\widehat{T}$ step-by-step from a velocity field $\widehat{v}$, parametrized by a neural network [150]. Flow matching generates target samples $y_1$ by integrating the ODE ${\dot{y}}_t={\widehat{v}}_t\left(y_t\right)$ starting from a random initial condition $y_0=x_0~p_0$ (say, an image of white noise). Flow matching defines a suitable loss function that allows learning the velocity $\widehat{v}$ from training samples. In contrast to diffusion models, the sample generation process is deterministic; Sec. V C 4 will introduce stochasticity.

Both flow matching and diffusion models are based on mass transport. Consider a time-dependent¹¹ vector field ${\widehat{v}}_t(x)$ on ${\mathbb{R}}^d$. Here and below, hats denote learned maps or fields. The marginal of the generated trajectories $y_t$ is $q_t$. By convention, time $t$ ranges from 0 to 1, so that $q_{t=0}=p_0$ is the base, and we want to adjust the generative process so that $q_1$ matches the target distribution $p_1$ (reader beware: the reverse convention is also often used in the literature). The flow map ${\widehat{T}}_t(x)$ transports particles along the integral curves of ${\widehat{v}}_t$, and is defined by solutions to the ODE

$${\dot{y}}_t={\widehat{v}}_t\left(y_t\right),y_{t=0}=x_0~p_0,{\widehat{T}}_t\left(x_0\right)≔y_t.$$ (235)

If the initial conditions are sampled $y_{t=0}~p_0$, the time-dependent distribution $q_t={\left({\widehat{T}}_t\right)}_{*}{p}_0$ is governed by conservation of mass [152]:

$${\partial }_t{q}_t(x)+\nabla \cdot \left(q_t(x)\cdot {\widehat{v}}_t(x)\right)=0$$ (236)

If a $q_t$ fulfills Eq. (236), it is said to be generated by the flow $\widehat{v}$. In fluid dynamics terms, Eq. (235) is the Lagrangian, and Eq. (236) the Eulerian formulation.

Note that many different velocity fields can generate the same $q_t$. Indeed, adding a divergence-free vector field, ${\widehat{v}}_t\mapsto {\widehat{v}}_t+v_t^{\prime }$, with $\nabla \cdot \left(q_t{v}_t^{\prime }\right)=0$, leaves Eq. (236) invariant. In Sec. II B, we saw how optimal transport singles out one of these vector fields by minimizing a transport cost. On the other hand, not all time-dependent distributions can be generated by a flow. For example, direct interpolation between two probability densities, $p_t(x)=(1-t)p_0(x)+t{p}_1(x)$ with disjoint support requires “mass teleportation” and cannot be generated by any flow.

Flow-generated distributions are very convenient: given $\widehat{v}$, one can generate samples of $q_1$ by integrating Eq. (235). How can one identify a velocity field $\widehat{v}$ so that $q_1={\left({\widehat{T}}_1\right)}_{*}{p}_0$ matches the target distribution $p_1$? A first idea is to directly optimize the transport map ${\widehat{T}}_1$, obtained by integrating Eq. (235). This approach is, however, very inefficient: evaluating the objective requires simulating the entire trajectory. Instead, one uses a stochastic interpolant [147, 148], a stochastic process $x_t$ with $x_0~p_0$ and $x_1~p_1$. Recall that in generative modeling, the target distribution $p_1$ is unknown. Instead, one has access to a set of training samples (for example, an image database). One constructs a stochastic interpolant from these training samples, for example:

$$x_t=(1-t)x_0+t{x}_1,x_0~p_0,x_1~p_1.$$ (237)

The endpoints $x_0,x_1$ are sampled independently¹². For instance, $x_1$ could be an image sampled from the database, and $x_0$ could be Gaussian noise, so that $x_t$ is a sequence of increasingly noisy images (Fig. 13). In this sense, Eq. (237) can be thought of as a version of the Brownian bridge from Sec. ID 3, conditioned to terminate in the target sample $x_1$. The interpolant Eq. (237) defines a time-dependent probability path $x_t~p_t$ from $p_0$ to $p_1$. The $p_t$ play a role similar to the instantaneous equilibria in annealed importance sampling (Sec. IV B). It is important to distinguish the interpolant $x_t$ from the generating process $y_t$; even if initialized at the same starting point $x_0~p_0$, the two processes lead to different trajectories.

FIG. 13.

A stochastic interpolant adds progressively larger amounts of uncorrelated Gaussian noise to an initial image.

Crucially, the probability path $p_t$ defined by Eq. (237) can be generated by a flow field (in the sense of mass conservation, Eq. (236)). As we show shortly, the velocity $v$ of the interpolant $p_t$ is the average of the trajectories’ speed:

$$v_t(x)=⟨{\dot{x}}_t\mid x_t=x⟩$$ (238)

Here and below, $⟨\cdot \mid \cdot ⟩$ denotes the conditional expectation. The average is over the randomly sampled endpoints $x_0,x_1$, which determine $x_t,{\dot{x}}_t$. To derive Eq. (238), integrate ${\partial }_t{p}_t$ against a test function $\phi (x)$, and use the law of total expectation:

$$\begin{gathered} \int \phi (x){\partial }_t{p}_t(x)d^{d}x=\frac{d}{dt}⟨\phi \left(x_t\right)⟩ \\ =⟨\nabla \phi \cdot \dot{x_t}⟩=⟨\nabla \phi \cdot ⟨\dot{x_t}\mid x_t⟩⟩ \\ =-\int \phi (x)\nabla \cdot \left(p_t⟨{\dot{x}}_t\mid x_t=x⟩\right)d^{d}x \end{gathered}$$ (239)

Fig. 12 shows a toy example ($p_1$ is a mixture of Gaussians), where the true flow field $v$ can be computed analytically.

FIG. 12.

(Left) Stochastic interpolation using Eq. (237) between $p_0=𝒩(0,I)$, a standard Gaussian, and $p_1=\sum _j{c}_j𝒩\left({\mu }_j,{\Sigma }_j\right)$, a Gaussian mixture with weights $c_j$ (in the 2d example shown, $c_1=0.4,c_2=0.6$). Since $x_t=(1-t)x_0+t{x}_1$ is a sum of Gaussian mixture variables, the intermediate distributions are also Gaussian mixtures, $p_t=\sum _j{c}_j𝒩\left(t{\mu }_j,(1-t{)}^{2}I+t^2{\Sigma }_j\right)$. The plot shows contour lines of $p_t$ for $t=0,1/2,1$. (Right) The flow map can be calculated from the intermediate $p_t$’s using Eq. (242). It transports samples from $p_0$ to samples from $p_1$ in unit time (black). Note that the integral curves are bent while the trajectories of the interpolant are straight lines.

The stochastic interpolant defines the objective: match the estimated flow field $\widehat{v}$ to the true one $v$. The resulting flow matching loss reads, using the law of total expectation:

$$\begin{gathered} {ℒ}_{\text{FM}}\left[\widehat{v}\right]={\int }_0^1⟨{‖{\widehat{v}}_t\left(x_t\right)-v_t\left(x_t\right)‖}^2⟩dt \\ ={\int }_0^1⟨{‖{\widehat{v}}_t\left(x_t\right)-⟨{\dot{x}}_t\mid x_t⟩‖}^2⟩dt \\ ={\int }_0^{1}dt⟨{‖{\widehat{v}}_t\left(x_t\right)‖}^2-2{\widehat{v}}_t\left(x_t\right)\cdot {\dot{x}}_t⟩ \\ +\text{const.}=⟨{‖{\widehat{v}}_t\left(x_t\right)-{\dot{x}}_t‖}^2⟩+\text{const.} \end{gathered}$$ (240)

The additive constant in Eq. (240) is independent of $\widehat{v}$, and thus does not affect training. Eq. (240) has three important properties. First, the unique minimizer is the true flow field. Thus, at the minimum, $q_t$ matches the interpolant distribution $p_t$, and $y_{t=1}$ is a sample from $p_1$ (note that $y_t$ and $x_t$ have the same marginals, but the trajectory laws can differ). Second, since $x_t=(1-t)x_0+t{x}_1$, the expectation value is easily estimated using independent samples $x_0~p_0,x_1~p_1$ from the training set. Third, the loss is simulation-free: it does not require simulating the entire trajectory $y_t$. Thus, one can efficiently train the neural network $\widehat{v}$ to generate target samples by stochastic gradient descent on ${ℒ}_{\text{FM}}$. In summary, the overall flow-matching workflow reads:

1. Design a stochastic interpolant $x_t$ whose distribution $p_t$ bridges the simple base distribution $p_0$ and the target $p_1$.

2. Using interpolant samples generated from a training set, train a velocity field $\widehat{v}$ by minimizing ${‖{\widehat{v}}_t\left(x_t\right)-{\dot{x}}_t‖}^2$.

3. Generate samples by integrating the ODE ${\dot{y}}_t={\widehat{v}}_t\left(y_t\right)$, with $y_{t=0}=x_0~p_0$.

One can see that flow matching bears a close resemblance to dynamic optimal transport (OT). Both aim to transform a base into a target distribution via “local”, mass-conserving dynamics, Eq. (236). Like in optimal transport, in Eq. (237), samples move along straight lines between endpoints. However, the choice of the endpoints is precisely the opposite: OT couples the endpoints deterministically via the transport map $T,x_1=T\left(x_0\right)$, while in Eq. (237), the endpoints are independent. Therefore, OT and stochastic interpolation are not identical. Sampling from the OT process is computationally hard, since it requires finding a very specific flow field that additionally minimizes a transport cost. In flow matching, the stochastic interpolant is instead designed to be easy to sample. These samples serve as a device to find a flow field $\widehat{v}$ that takes $p_0$ to $p_1$.

Flow matching can be viewed not only as data-driven transport, but also as an iterative denoising process. This perspective sheds light on how flow matching is able to learn multi-modal distributions. Multi-modality was already a challenge in the context of sampling from a known distribution (Chapter IV). Since $x_0$ is typically Gaussian white noise, the stochastic interpolant is a sequence of increasingly noisy images (Fig. 13). One can think of the velocity field as a denoiser, which estimates $x_{t+dt}$ from the (infinitesimally) more noisy version $x_t$ as $⟨x_{t+dt}\mid x_t⟩\approx x_t+v_t\left(x_t\right)dt$. The flow-matching loss Eq. (240) measures the squared error of the denoiser. Intuitively, flow matching thus breaks down a difficult sampling task (generating an image) into a sequence of simpler tasks (incrementally denoising an image).

This step-wise approach is essential to handle multi-modality. Indeed, suppose one were to add a large, finite amount of noise to a target sample and attempt to denoise it in one step. The denoiser $D_t$ at interpolation time $t$ has mean squared error $⟨{‖D_t\left(x_t\right)-x_1‖}^2⟩$ (smaller $t$ corresponds to more noise). Because the noise removes details, several different reconstructions are compatible with the noisy sample. The optimal denoiser (with minimum least-squares error) will return the average of all reconstructions, $D_t^{*}(x)=⟨x_1\mid x_t=x⟩$. If the distribution is multi-modal, this average is a blurred mix of different modes that does not resemble any sample of the target distribution. However, for an infinitesimal amount of noise, the reconstruction problem becomes uni-modal. This is the principle that flow matching exploits to learn multi-modal distributions.

3. The score function

Let us take a closer look at the distributions $p_t$ for a linear interpolant between a standard Gaussian and an arbitrary target. If $x_0$ is standard Gaussian and independent of $x_1$, the conditional density is also Gaussian, $p_t\left(\cdot \mid x_1\right)=𝒩\left(t{x}_1,(1-\right.\left.t{)}^{2}I\right)$. Taking the gradient, one obtains $\nabla \text{log}{p}_t\left(x_t\mid x_1\right)=-x_0/(1-t)$. Averaging over endpoints using the law of total expectation gives [153]:

$$\begin{gathered} \nabla \text{log}{p}_t(x)=⟨\nabla \text{log}{p}_t\left(x\mid x_1\right)\mid x_t=x⟩ \\ =-\frac{⟨x_0\mid x_t=x⟩}{1-t}. \end{gathered}$$ (241)

The quantity $\nabla \text{log}{p}_t$ is called the score. We already encountered the score in Sec. ID 2 in the context of optimal control, where it served as a control force that drives a particle towards a target. Here, the score plays a similar role. Eq. (241) shows that the score can separate the noise $x_0$ from the current $x_t$, and thus allows reconstructing the target $x_1$. After a little algebra:

$$\begin{gathered} v_t(x)=⟨{\dot{x}}_t\mid x_t=x⟩=\frac{x-⟨x_0\mid x_t=x⟩}{t} \\ =\frac{x+\left(1-t\right)\nabla \text{log}{p}_t(x)}{t},0<t<1. \end{gathered}$$ (242)

Flow matching with an initially Gaussian distribution is therefore closely related to score matching [154], a method that uses the score to fit statistical models. Flow and diffusion models can be viewed as a combination of simulated annealing (Sec. IV B) and score matching [155]. Eq. (242) demonstrates that the velocity field is a potential gradient: $\text{log}{p}_t$ acts as a time-dependent energy landscape that guides a sample from the initial to the target distribution (see the example in Fig. 12). Note that the relation between score and flow only holds for linear interpolation where one endpoint is Gaussian; in general, there is no simple relation between $p_t$ and $v$.

4. From flow matching to diffusion

In flow matching, sampling is a deterministic process, defined by the ODE Eq. (235). Diffusion models instead use a stochastic sampling process, which, however, has the same marginals. To see how, add $\nabla \text{log}{p}_t$ to and subtract it from the mass conservation equation for the interpolant (target) density $p_t$:

$$\begin{gathered} {\partial }_t{p}_t=-\nabla \cdot \left(\left(v_t+\frac{\epsilon }{2}\nabla \text{log}{q}_t-\frac{\epsilon }{2}\nabla \text{log}{p}_t\right)q_t\right) \\ =-\nabla \cdot \left(\left({\widehat{v}}_t+\frac{\epsilon }{2}\nabla \text{log}{p}_t\right)p_t\right)+\frac{\epsilon }{2}{\nabla }^2{p}_t \end{gathered}$$ (243)

Eq. (236) is thus reformulated as a Fokker-Planck equation with diffusion, compensated by an additional drift. This seemingly tautological manipulation allows defining a new sampling process. The Fokker-Planck Eq. (243) gives a Langevin equation for trajectories:

$$\begin{gathered} {\dot{y}}_t=v_t\left(y_t\right)+\frac{\epsilon }{2}\nabla \text{log}{p}_t\left(y_t\right)+\sqrt{\epsilon }{\eta }_t \\ =\frac{1}{t}\left(y_t+\left(1+\frac{\epsilon -2}{2}t\right)\nabla \text{log}{p}_t\left(y_t\right)\right)+\sqrt{\epsilon }{\eta }_t. \end{gathered}$$ (244)

In the second step, we used the relation between score and velocity. Because the density obeys Eq. (243), the stochastic process Eq. (244) still generates samples from $p_1$. Effectively, diffusion models add an additional stochastic term to the interpolant Eq. (237), such that there is stochasticity at intermediate $t$ even after fixing the endpoints. Stochastic sampling has practical advantages: it can improve stability so that even an imperfectly learned score produces reasonable samples [147].

With the above results, flow matching can be compared to the original formulation of denoising diffusion models [156]. Here, a forward diffusion process gradually adds noise to a target sample $x_1$. This forward process is the pendant of the stochastic interpolant and generates samples from which the score function is learned. Given the score function, target samples are generated by Eq. (244). In this context, Eq. (244) is called the reverse process, since it is the time-reversal of the forward diffusion process.

To see this, reparametrize time as $s=-\text{log}t$, so $s=0$ is the target distribution and $s=\infty$ is standard Gaussian noise. The forward process is initialized at the target distribution:

$${\dot{x}}_s=-x_s+\sqrt{2}{\eta }_s,x_{s=0}~p_1$$ (245)

Since we set $\epsilon =2$, as $s\to \infty$, the distribution $p_s(x)$ converges to a standard Gaussian (the base distribution). (Note that the intermediate distributions $p_s$ are not equivalent to those of the stochastic interpolant Eq. (237).) The forward process obeys the Fokker-Planck equation:

$$\begin{gathered} {\partial }_s{p}_s(x)=\nabla \cdot \left(x{p}_s\right)+\frac{1}{2}{\nabla }^2{p}_s \\ =\nabla \cdot \left(p_s\left(x+\frac{1}{2}\nabla \text{log}{p}_s\right)\right) \end{gathered}$$ (246)

To reverse the probability flow, one flips the sign of the effective flow field $\left(x+\frac{\epsilon }{2}\nabla \text{log}{p}_s\right)\mapsto -\left(x+\frac{\epsilon }{2}\nabla \text{log}{p}_s\right)$. One thus reverses time, starting from a large final time $S\gg 1$. The reverse process, which generates target samples starting from noise, reads [156, 157]

$${\dot{y}}_s=y_s+\frac{\epsilon }{2}\nabla \text{log}{p}_{S-s}\left(y_s\right)+\sqrt{\epsilon }{\eta }_s$$ (247)

This idea of time-reversing a process that starts at the target distribution $p_1$ is precisely what underlies the proof of annealed importance sampling and the Jarzynski identity in Sec. IV B 2.

Finally, Eq. (244) also allows revisiting the distinction between flow matching and Langevin sampling with energy $E=-\text{log}{p}_1$. Flow matching uses a non-equilibrium process to transform probability distributions in finite time, and the target $p_1$ is not a stationary state. The potential is time-dependent, and the intermediate-time distributions $p_t$ are explicitly distinct from the target and instead act as a bridge between $p_0$ and $p_1$. By contrast, Langevin dynamics only converge to its equilibrium distribution $p_1$ as $t\to \infty$.

5. One-step generation and normalizing flows

Once a flow-matching model is trained, generating samples requires integrating the flow field $\widehat{v}$ (or an SDE for a diffusion model). This can be slow and makes the generation process difficult to fine-tune for downstream applications. Consistency models [158] and rectifying flows [159] distill the flow field $\widehat{v}$ into a neural-network model for the flow map ${\widehat{T}}_1$, which allows generating samples in one step. Such one-step models parallel the transition from Benamou-Brenier dynamic optimal transport (Sec. II B, Eq. (91)) to the static Monge-Kantorovich problem (Sec. II B).

In this section, we will discuss another one-step generative method, normalizing flows [160], one of the inspirations for flow matching models. Despite their name, normalizing flows are not flows, but generate samples using a single-step transform. While flow matching excels at learning generative models from training data, normalizing flows can also learn from an energy function or unnormalized probability density. This makes them very useful for improving sampling algorithms in statistical mechanics simulations and Bayesian statistics (Chapter IV). In this setting, no training samples are available: producing them is the whole point.

As before, the goal is to transform a base $p_0$ into a target $p_1$ distribution via Eq. (234). Samples are generated as $y=\widehat{T}\left(x_0\right),x_0~p_0$ where the base distribution $p_0$ is often Gaussian (hence, “normalizing” flow). In contrast to flow matching, which constructs $\widehat{T}$ iteratively, normalizing flows directly parametrize $\widehat{T}$ by a neural network. However, not any neural network will do: Eq. (234) shows that the inverse ${\widehat{T}}^{-1}$ and the Jacobian determinant $\text{det}{\nabla }_x{\widehat{T}}^{-1}$ are needed to compute the transformed density $q={\widehat{T}}_{*}{p}_0$. One thus uses neural network architectures designed specifically to make these quantities easy to evaluate. A minimal example is the following neural network layer for a 2-component vector $x=\left(x^{(1)}{x}^{(2)}\right)$ [161]:

$$\left(x^{(1)}{x}^{(2)}\right)\mapsto \left(\widehat{a}\left(x^{(2)}\right)\cdot x^{(1)}+\widehat{b}\left(x^{(2)}\right)x^{(2)}\right)$$ (248)

where $\widehat{a}$ and $\widehat{b}$ are neural networks for scale and shift. Since Eq. (248) is an affine function in $x^{(1)}$ and copies $x^{(2)}$, it is easy to invert. More complex networks can be built by stacking layers.

Next, we need a procedure to adjust $q={\widehat{T}}_{*}{p}_0$. to match the target distribution $p_1$, paralleling the minimization of the flow matching loss in Eq. (240). One possibility is to use a training dataset $\left\{x_n\right\}$ of samples from $p_1$. In this case, one maximizes the $\text{log}$-likelihood $\sum _{n=1}^N\text{log}q\left(x_n\right)$ of the training data. In the setting of Bayesian inference or statistical mechanics simulations, one has access to the unnormalized density ${\tilde{p}}_1$ instead of training samples. In Sec. V B 3, we saw that variational inference uses $p_1$ to define a data-free objective:

$$\begin{gathered} {\text{D}}_{\text{KL}}\left(q\Vert p_1\right)={⟨\text{log}q-\text{log}{p}_1⟩}_q \\ \approx \frac{1}{N}\sum _{n=1}^N\left[\text{log}q\left(x^{\left(n\right)}\right)-\text{log}{p}_1\left(x^{\left(n\right)}\right)\right], \\ {x}^{(n)}=\widehat{T}\left(x_0^{(n)}\right),x_0^{(n)}~p_0 \end{gathered}$$ (249)

Here, $D_{\text{KL}}$ is the Kullback-Leibler (KL) divergence, a measure of the distance between two probability distributions encountered several times in previous chapters. Because of the logarithm, the normalization constant $Z_1$ drops out when optimizing Eq. (249). The expectation value in Eq. (249) is approximated using samples from $q$, so no target distribution samples are required (hence, “data-free”). Now, $\widehat{T}$ can be trained by minimizing ${\text{D}}_{\text{KL}}\left(q\Vert p_1\right)$ via gradient descent [162]. The data-free objective has certain drawbacks: it can encourage mode collapse, where the map $\widehat{T}$ always returns samples from a single mode and fails to explore the full distribution $p_1$ [163]. This is a well-known effect of the KL divergence. It strongly penalizes any probability density outside of the support of $p_1$ due to the term $\text{log}\left(q/p_1\right)$, but failing to sample from a part of $p_1$ receives a much milder penalty.

Once trained, one can generate samples $\widehat{T}\left(x_0\right)$. However, if the variational approximation is not perfect, these samples will be biased (not exactly distributed $~p_1$). This is a serious issue in scientific contexts. Fortunately, one can combine a normalizing flow with another sampling technique like importance sampling from Chapter IV to correct the bias [163]. As long as $q$ covers the support of $p_1$, the result is guaranteed to be unbiased. This achieves the “best of both worlds”: the conventional sampling algorithm provides mathematical guarantees, and the power of machine learning allows rapid identification of high-probability regions in $d$-dimensional space. Normalizing flows have been successfully used to accelerate statistical physics simulations, for instance in lattice QCD [164]. More generally, gradient-based optimization can systematically tune parameters of sampling algorithms that were traditionally tuned by hand, like the proposal distribution for an importance sampler.

6. Generalizations and ongoing work

In summary, machine learning has made striking progress on the classic problem of sampling from complex, high-dimensional distributions. These techniques transform a simple base distribution to the target via a transform that is fit variationally to data. The greatest success has come from methods based on mass transport, like diffusion and flow matching, which construct the transform map step by step. In flow matching, one first designs a dynamical process that interpolates between the base (often, uniform noise) and the target, and then learns the corresponding probability flow. In the domain of image and video generation, flows and diffusions are the state of the art and have eclipsed previous approaches, like Variational Autoencoders [149] and Generative Adversarial Networks [165]. On the other hand, for language and other “sequential datasets”, autoregressive models, which iteratively predict the next sequence element or “token”, dominate.

Flow matching can be generalized in several directions. First, the design of the stochastic interpolant offers significant flexibility. Eq. (237), linear interpolation, is only the simplest possibility. Second, there is an equal amount of freedom in choosing a sampling algorithm (stochastic or deterministic sampling), which can lead to important numerical efficiency gains. Lastly, different neural network architectures can be used to approximate the velocity field or score. This leads to a large zoo of methods [147, 148].

Next, in generative modeling, one is rarely interested in generating “random” samples from the target distribution $p_1$ (in contrast to statistical mechanics). After all, a random sample could simply be drawn from the training set. Instead, one is interested in a sample that fulfills some particular task: an image that matches a given caption, or a protein with a particular functionality. Guided diffusion uses additional “guidance” terms to nudge the sampling process $y_t$ towards a user-defined aim [166]. The population distribution $p_1$ of images/proteins/… serves as a conduit, a convenient fiction, towards the ultimate goal of a task-directed model. Guided diffusion is a form of sampling with a control $u_t$ which we discussed above. A running reward can be supplied by an ML model (like a classifier) to score whether the trajectory ${\text{x}}_t$ progresses towards a desired sample like a particular type of image, and a terminal reward can score the final sample.

Finally, flow matching can also be generalized from continuous distributions on ${\mathbb{R}}^d$, such as images, to discrete ones, such as text data. In this case, a sample is a sequence $x=\left\{x_1,\dots ,x_d\right\},x_i\in 𝒜$ from a finite alphabet $𝒜$. The stochastic interpolant becomes a Markov chain that converts a sequence into a sample from a simple base distribution $p_0$. For example, masked diffusion [167] continuously masks out sequence elements (so $p_0$ has only one element, the fully masked sequence). The goal is now to learn a Markov process that reverses the masking. The transition rates of this process play the role of the velocity field in the continuous setting, and one minimizes a loss similar to Eq. (240). Alternatively, the discrete data can be embedded into a continuous space. For language modeling, discrete diffusion has recently become competitive with the next-token-prediction paradigm that underlies most large language models [167]. In contrast to the latter, diffusion models can generate many tokens at once, and thus leverage more computing power.

Statistical mechanics of memorization and generalization.

Flow matching and diffusion models are optimized by using a set of training samples, but so far, we have not paid any attention to the limited size of this training set. The models are fit not to the true target $p_1$, but to the empirical distribution ${\widehat{p}}_{\text{data}}(x)=N^{-1}\sum _{n=1}^N\delta \left(x-x_n\right)$ of the finite training set. Therefore, the “optimal” flow map under Eq. (240) reproduces the empirical distribution, and generation always returns a sample from the training set. This evidently undesirable phenomenon is called memorization; the model fails to generalize to new samples. In practice, memorization is avoided by stopping the training process early, and systematic approaches to avoid memorization are a topic of current research. The phenomenon of benign overfitting, where a model perfectly interpolates the training data without impeding performance, is thus absent in flow models. A model with zero training flow matching loss always memorizes.

The dynamics of diffusion models starting from pure noise bear a close resemblance to a statistical mechanics system that is slowly cooled from an initially high temperature. Flow-matching time is thus analogous to temperature. Indeed, flow matching and diffusion models can be fruitfully analyzed using statistical mechanics, which has shed light on the generalization-memorization transition. Refs. [168, 169] argue that flow matching dynamics undergo a series of symmetry-breaking phase transitions as a function of time $t$. At a first “ferromagnetic” transition, samples $x_t$ commit to one of the modes of the target distribution, generating a novel sample. However, as the dynamics continue, a second “glass” transition traps $x_t$ near a sample from the training dataset (Fig. 14).

FIG. 14.

Memorization-generalization transition for a 1D mixture of Gaussians: interpolating distributions $p_t$ between a Gaussian base and the empirical distribution of a Gaussian mixture ($N=4$ samples). At the ferromagnetic transition, $p_t$ splits into the distinct modes of the target distribution. At the glass transition, trajectories become trapped near a training sample from the empirical distribution.

In Chapter IV, we saw how an analogy with physical annealing gave rise to powerful algorithms for sampling from a known distribution. An interesting avenue for future research is whether the analogy between annealing and flow matching could be used to design “better”, more statistically efficient generative models, now learned from training data.

TABLES OF NOTATION

The final column links to the section where each symbol first appears. Each chapter has a separate table, and we list notation in the order in which each symbol was first introduced. Familiar symbols may have different meanings in clearly identified contexts, for example, mechanical action $S[x(\cdot )]$ and thermodynamic entropy $S\left[p\right]$.

TABLE I:

Chapter I: Mechanics, Control, and Inference.

Symbol	Meaning or convention	Section
$x,x_t,x(\cdot )$	State argument, state at time $t$, and complete trajectory. In Chapter V, $x_n$ denotes the state at decision step $n$ and $y_t$ distinguishes a generated path from an interpolant.	IA
$u,u_t(x)$	Control value and feedback control; the applied control is $u_t\left(x_t\right)$. In Chapter III, $u_t$ also denotes a prescribed experimental protocol.	IA
$t,s$	Time; the control horizon is $0\le t\le 1$.	IA
$g_t(x,u)$	General controlled drift.	IA
$\epsilon$	Noise covariance per unit time; amplitude $\sqrt{\epsilon }$, diffusion coefficient $\epsilon /2$. Also the entropic transport regularization in Sec. II B. In Chapter III, $\epsilon =2D$.	IA
${\eta }_t,W_t$	Standard Gaussian white noise and standard Wiener process.	IA
$r_t(x,u)$	Running reward, including control costs. In Sec. V A, $r(x,u)$ is a per-step reward; its running-reward limit includes a factor $\Delta t$.	IA
$r_1(x)$	Terminal reward, separate from running reward. In Sec. II A it enforces the final density constraint.	IA
$V_t(x)$	Optimal expected reward-to-go, including control costs. In Sec. V A, $V^{\pi }$ evaluates a policy and $V^{*}$ is optimal.	IA
$L_t(x,\dot{x}),S[x(\cdot )]$	Mechanical Lagrangian and action.	IA
$p_t$	Costate; canonical momentum in the mechanics specialization.	IA
${\mathscr{H}}_t(x,p,u)$	Control Hamiltonian before maximization over $u$.	IA
$H_t(x,p)$	Optimized Hamiltonian, $H_t={\text{max}}_u{\mathscr{H}}_t$.	IA
$p(x),p_t(x)$	Probability density and time-dependent marginal density. In Chapter IV, $p_0,p_1$ denote the base and target densities.	IB–IC
$p^0,p^A$	Reference density and density inferred from observations $A$. Superscripts distinguish these from time subscripts.	IB
${\text{D}}_{\text{KL}}(p\Vert q)$	KL divergence, with averaging under its first argument; also used for path distributions.	IB
$⟨\cdot {⟩}_p$	Expectation under the indicated density; conditioning is written inside the brackets. Also written ${\mathbb{E}}_p[\cdot ]$ in Chapter V.	IB
$N$	Number of independent samples; in Sec. V A, $N$ instead denotes the final decision index and $M$ counts rollouts.	IB
$\hat{p}$	Empirical density estimated from samples; ${\widehat{p}}_{\text{data}}$ denotes the training-set density in Sec. V C. Hats indicate estimates, not samples.	IB
$I$	Large-deviation rate function.	IB
$a_m(x),a_m[x(\cdot )]$	State observable and, in Sec. I C, trajectory observable.	IB–IC
$A_m$	Prescribed mean of observable $a_m$.	IB
$H\left[p\right]$	Shannon entropy with natural logarithms. In Sec. V A, $H[\pi (\cdot \mid x)]$ is the policy entropy at state $x$.	IB
$M$	Number of imposed observable constraints; in Sec. V A, the number of independent rollouts.	IB
${\lambda }_m$	Multiplier enforcing the mean of observable $a_m$, with the negative sign in the exponential tilt.	IB
$Z$	Normalization constant or partition function. In Sec. II B, $Z\left(x_0\right)$ normalizes a transition kernel; in Sec. V A, $Z_Q(x)=\int e^{Q(x,u)/\alpha }du$ normalizes the Boltzmann policy.	IB
$𝒜,1_{𝒜}$	Conditioning set and its indicator.	IB
$P^0,P^u,P^{u^{*}}$	Reference, controlled, and optimally controlled path distributions. In Sec. V A, $P^{\pi }$ is the joint state-action path law induced by a policy, and $P^{{\pi }^{*}}$ is its optimal realizable form.	IC
$y$	Prescribed terminal point in the Brownian-bridge examples.	IC
$f_t(x)$	Passive drift; additive control gives $g_t(x,u)=f_t(x)+u$.	IC
$p_t^u(x)$	Marginal density under $P^u$; denoted ${\rho }_t(x)$ in Chapter II.	IC
$⟨\cdot {⟩}_{P^u}$	Expectation over the indicated path distribution.	IC
$r_m$	Reward coefficient enforcing a trajectory constraint, with $r_m=-{\lambda }_m$.	IC
$r_t(x)$	State reward, excluding the separately displayed quadratic control cost.	IC
${\psi }_t(x)$	Exponentiated PISC value, ${\psi }_t=e^{V_t}$. Also denoted $Z_t(x)$ during its derivation; the backward factor in Sec. II A.	ID
$R_t[x(\cdot )]$	Accumulated state and terminal rewards from $t$ onward, excluding control cost.	ID
$\nabla \text{log}{p}_t(x)$	Spatial score of a marginal density; distinct from the conditioning field $\nabla \text{log}{\psi }_t$ and the protocol score ${\Theta }_u$ of Chapter III.	ID
$k_{t\mid s}(z\mid x)$	Passive transition density from $x$ at $s$ to $z$ at $t>s;k=k_{1\mid 0}$ in Sec. II B. In Chapter IV, $k_t$ freezes the energy for one step; in Sec. V A, $k(x\prime \mid x,u)$ is the environment transition kernel.	ID

TABLE II:

Chapter II: Transport and Wasserstein Geometry.

Symbol	Meaning or convention	Section
${\rho }_t,{\rho }_0,{\rho }_1$	Evolving density and prescribed endpoint densities. In Sec. V B, ${\widehat{\rho }}_t=N_{\text{p}}^{-1}{\sum }_i\delta \left(x-x_t^{(i)}\right)$ is the empirical particle density.	II A
$v_t(x)$	Density-transport velocity: ${\partial }_t{\rho }_t+\nabla \cdot \left({\rho }_t{v}_t\right)=0$. Generally distinct from $u_t;{\widehat{v}}_t$ is a learned estimate in Sec. V C.	II A
${\psi }_t^{†}(x)$	Forward Schrödinger factor, with ${\rho }_t={\psi }_t{\psi }_t^{†}$. The dagger does not denote complex conjugation.	II A
$c\left(x_0,x_1\right)$	Cost of transporting unit mass between two endpoints.	II B
$T\left(x_0\right)$	Transport map assigning a final position to an initial position. In Sec. V C, $T_t$ is a flow map and ${\widehat{T}}_t$ a learned map, not necessarily optimal.	II B
$\pi \left(x_0,x_1\right),{\pi }^{*}$	Endpoint coupling and optimal coupling, with marginals ${\rho }_0,{\rho }_1$. Noise dependence is displayed as ${\pi }_{\epsilon }^{*}$ when taking a limit.	II B
${\phi }_0,{\phi }_1$	Endpoint dual potentials for entropic transport; ${\phi }_1=\epsilon r_1$ with compatible normalization. Their limits are Kantorovich potentials.	II B
$\zeta ,{\bar{\phi }}_1$	Convex conjugate and shifted limiting endpoint potential: ${\bar{\phi }}_1(x)=\Vert x{\Vert }^2/2-{\phi }_1(x),T=\nabla \zeta$.	II B
$C_0,C_1$	Cumulative distribution functions of the endpoint densities; $C_i^{-1}$ denotes the quantile function.	II B
$\theta ,p_{\theta }$	Parameters and corresponding density in the geometric illustration; reused for the parametric VI family in Sec. V B.	II B
$W_2\left({\rho }_0,{\rho }_1\right)$	2-Wasserstein distance; its square is the minimum mean squared displacement. The $\text{cost}c={‖x_0-x_1‖}^2/2$ has minimum $W_2^2/2$.	II B
$F(x)$	Ordinary scalar objective in the Euclidean gradient-descent analogy.	II C
$s(x),s_t(x)$	Density tangent; along a curve, $s_t={\partial }_t{\rho }_t$, with $\int s_{t}dx=0$.	II C
${𝒲}_2$	Space of probability measures with finite second moment, equipped with $W_2$.	II C
$\chi (x)$	Potential for the minimum-energy transport velocity, with $v=-\nabla \chi$. In Sec. V B, $v_i=-\nabla {\chi }_i$ represents the parameter tangent ${\partial }_{{\theta }_i}{p}_{\theta }$.	II C
$\Vert s{\Vert }_{\rho },{⟨s,s^{\prime }⟩}_{\rho }$	Wasserstein tangent norm and inner product at density $\rho$.	II C
$ℱ\left[\rho \right]$	Functional of a density, such as free energy.	II C
${\nabla }_{\text{W}}ℱ\left[\rho \right]$	Wasserstein gradient in its velocity representation; the descent velocity is its negative.	II C
$\delta ℱ/\delta \rho$	Functional derivative; ${\nabla }_{\text{W}}ℱ=\nabla (\delta ℱ/\delta \rho )$.	II C
$\Phi (x)$	Dimensionless potential in the illustrative free energy with unit diffusion coefficient; reused in Sec. V B, with $p_1=e^{-\Phi }/Z_1$ for VI.	II C
$\alpha$	Convexity or geodesic-convexity parameter; positive values indicate strong convexity.	II C

TABLE III:

Chapter III: Thermodynamics.

Symbol	Meaning or convention	Section
$S\left[p\right]$	Thermodynamic entropy, $S\left[p\right]=k_{\text{B}}H\left[p\right]$ in the statistical description used here; distinct from mechanical action $S[x(\cdot )]$.	III A
$T,k_{\text{B}},\beta$	Bath temperature, Boltzmann constant, and inverse thermal energy $\beta ={\left(k_{\text{B}}T\right)}^{-1}$.	III A
$p^{\text{eq}},p_u^{\text{eq}}$	Equilibrium density, including dependence on a fixed protocol $u$. In Chapter IV, $p_t^{\text{eq}}=e^{-E_t}/Z_t$ is the prescribed equilibrium bridge.	III A
$E(x),E_u(x),E_t(x)$	Microstate energy; $E_t=E_{u_t}$ for a protocol. General sampling algorithms in Chapter IV absorb $\beta$ into $E_t$ and use dimensionless energy.	III A
${ℱ}^{\text{eq}},{ℱ}_u^{\text{eq}}$	Equilibrium free energy, $-k_{\text{B}}T\text{log}{Z}_u;-\text{log}{Z}_u$ in Chapter IV’s thermal units.	III A
$ℰ$	Mean internal energy, $⟨E{⟩}_p$.	III A
$\Sigma ,{\dot{\Sigma }}_t$	Total entropy production and its rate; distinct from the change $\Delta S$ in system entropy.	III B
$Q,W$	Mean heat received by the system and mean work done on it.	III B
$N_{\text{p}}$	Particle number; distinct from $N$ statistical samples. In Sec. V B, also the number of neurons viewed as parameter particles.	III B
$X_m,L_{mn}$	Macroscopic variables ($m=1,\dots ,M$) and Onsager transport coefficients.	III B 1
$ℛ$	Rayleighian, $ℛ=\frac{1}{2}T\dot{\Sigma }+\dot{ℱ}$; distinct from accumulated reward $R_t$.	III B 1
$n_t(x)$	Conserved particle number density, with $\int n_{t}dx=N_{\text{p}}$. Its normalized form is ${\rho }_t=n_t/N_{\text{p}}$.	III B 1
$\Gamma (n)$	Dissipation or drag coefficient in the continuum Onsager description.	III B 1
$\eta ,r$	Fluid viscosity and colloid radius in the Stokes-drag example; $\eta$ here is distinct from white noise ${\eta }_t$.	III B 1
${\mu }_t(x)$	Chemical potential, $\delta ℱ/\delta n_t(x)$; for constant drag, relaxation has ${\chi }_t={\mu }_t/\Gamma$.	III B 1
$D$	Physical diffusion coefficient, $D=\epsilon /2$.	III B 1
$F_t(x)$	Physical force; the overdamped stochastic drift is $f_t=D\beta F_t$.	III B 2
$\tau$	Process duration when physical elapsed time is displayed explicitly.	III B 2
$\delta p_t(x)$	Deviation from instantaneous equilibrium: $p_t=p_{u_t}^{\text{eq}}+\delta p_t$.	III B 3
${ℒ}_u,{ℒ}_u^{†}$	Operators with convention ${\partial }_{t}p=-{ℒ}_{u}p;-{ℒ}_u^{†}$ is the backward generator. The dagger denotes the operator adjoint.	III B 3
${\Theta }_u(x)$	Protocol score, ${\partial }_u\text{log}{p}_u^{\text{eq}}(x)$; a vector for a multicomponent protocol.	III B 3
${\zeta }_u$	Integrated equilibrium score-correlation tensor defining the protocol-space metric. Its energy-unit friction normalization is $k_{\text{B}}T{\zeta }_u$.	III B 3
$w_t[x(\cdot )]$	Work accumulated along one trajectory up to time $t;W={⟨w_1⟩}_P$ for the unit-time protocol. Dimensionless in Chapter IV’s sampling formulas.	III B 4
$P,{\bar{u}}_t,\bar{P}$	Forward path law, reversed protocol ${\bar{u}}_t=u_{1-t}$, and reverse path law. In Chapter IV, $E_t=E_{1-t}$ and $x_t^{\text{rev}}=x_{1-t}$ denotes a reversed realization.	III B 4
$\lambda$	Work-MGF transform variable in $e^{\lambda w}$; inverse-energy units when $w$ is physical work.	III B 4
$M_t,{\bar{M}}_t,{\tilde{M}}_t$	Work-weighted endpoint density, reverse conditional MGF, and transformed reverse MGF used in the fluctuation-theorem proof.	III B 4

TABLE IV:

Chapter IV: Sampling and Control.

Symbol	Meaning or convention	Section
$\tilde{p},p=\tilde{p}/Z$	Unnormalized density and the corresponding normalized probability density.	IV A
$d$	Dimension of configuration space or number of model degrees of freedom.	IV A 1
${\hat{a}}_N$	Monte Carlo estimate of $⟨a⟩$ from $N$ independent samples $x_n$. For an ensemble of paths, $x_t^{(n)}$ labels time and run separately.	IV A 1
${\sigma }_i,m(x)$	Ising spin and configuration magnetization, $m(x)=d^{-1}{\sum }_i{\sigma }_i$.	IV A 2
$\omega ,\tilde{\omega }$	Normalized importance ratio and unnormalized weight. In AIS, ${\tilde{\omega }}_t=e^{-w_t}$ and ${\omega }_t=\left(Z_0/Z_t\right){\tilde{\omega }}_t$.	IV A 3
$N_{\text{eff}}$	Effective sample size estimated from the weights, ${\left({\sum }_n{\tilde{\omega }}_n\right)}^2/{\sum }_n{\tilde{\omega }}_n^2$.	IV A 3
$q_t,{\bar{q}}_t$	Actual forward and reverse sampler densities, distinct from prescribed densities. In AIS, $q_0=p_0,{\bar{q}}_0=p_1$; in Sec. V C, $q_t={\left({\widehat{T}}_t\right)}_{*}{p}_0$ is the learned model’s density.	IV A 3

TABLE V:

Chapter V: Applications in Learning.

Symbol	Meaning or convention	Section
$𝒳,𝒰$	RL state and action spaces.	V A 1
$\pi (u\mid x)$	Policy: conditional action distribution; time is included in the state for finite-horizon problems. Later, ${\pi }_Q$ denotes the Boltzmann policy defined by a critic.	V A 1
$n,N$	Decision step and final decision index; $n=0,\dots ,N$ and $t_n=n\Delta t$.	V A 1
$Q^{\pi }(x,u),Q^{*}(x,u)$	Exact policy-specific and optimal action values; hats denote learned tabular estimates. In MaxEnt RL, these denote entropy-regularized (“soft”) values.	V A 1
$\theta ,Q_{\theta }^{\pi },Q_{\theta }^{*}$	Neural-network estimates of policy-specific and optimal values; $\theta$ denotes the network weights.	V A 2
$i,M$	Rollout sample index and number of independent rollouts, $i=1,\dots ,M$.	V A 2
$𝒥\left[\pi \right]$	Expected RL return; the text specifies discounting and any entropy bonus.	V A 2
$\gamma$	Reward discount factor.	V A 2
$y$	Sampled Bellman target used to update a critic.	V A 2
$\xi$	Learning rate.	V A 2
$𝒟$	Replay buffer of observed transitions and rewards.	V A 2
$L(\theta ),ℒ\left[\rho \right]$	Loss as a function of parameters or a functional of a density, respectively.	V A 2
$\bar{\theta }$	Parameters held fixed when evaluating a critic target; target-network updates, if used, are specified separately.	V A 2
$\alpha$	RL entropy temperature in reward units; distinct from the convexity parameter in Wasserstein geometry.	V A 3
${\pi }^0(u\mid x),c_0$	Reference action distribution and its constant value in the uniform case.	V A 3
$P^{*}$	Unrestricted reward-tilted target path law; generally distinct from $P^{{\pi }^{*}}$.	V A 3
$\varphi ,{\pi }_{\varphi }$	Actor parameters and corresponding parameterized policy.	V A 4
${\Phi }^{\left(N_{\text{p}}\right)},𝒱\left[\rho \right]$	Many-particle potential and its density-functional representation, with the normalization displayed in Sec. V B 2.	V B 2
$U(x,y)$	Pair-interaction potential, also used for interactions between parameter particles.	V B 2
$f^{N_{\text{p}}}\left(x;{\Theta }_t\right)$	Finite-width neural-network predictor; distinct from the passive drift $f_t$.	V B 2
${\Theta }_t,{\theta }_t^{(i)},w_t^{(i)},z_t^{(i)}$	Full parameter collection, neuron parameters, feature weights and internal feature parameters; $\theta =(w,z)$.	V B 2
$\phi (x;z)$	Nonlinear feature; its weighted form is the observable $a(x;\theta )=w\phi (x;z)$.	V B 2
$y(x),\nu (x)$	Target regression function and input-data density; distinct from the Bellman target $y$.	V B 2
$ℳ,\Omega ,d_{\theta }$	Parametric density family, parameter domain and parameter dimension.	V B 3
$g_{ij}$	Parameter-space metric induced by Wasserstein geometry.	V B 3
${\text{proj}}_{ℳ}$	Orthogonal projection of a velocity onto the tangent space of the parametric density family.	V B 3
$\text{BW}\left({\mathbb{R}}^d\right),{\nabla }_{\text{BW}}$	Gaussian Bures–Wasserstein space and its projected gradient.	V B 3
$m,\Sigma ,{\mathbb{S}}_{++}^d$	Gaussian mean, covariance and space of symmetric positive-definite covariance matrices.	V B 3
$h,k_h,{𝒦}_h^p$	Smoothing bandwidth, spatial kernel and density-weighted smoothing operator.	V B 3
$T_{*}p$	Pushforward of a distribution through a map $T$.	V C 1
$\hat{v}$	Estimated probability flow field.	V C 2
$y_t$	Sample-generating process in flow and diffusion models.	V C 2
$x_t$	Stochastic interpolant defined from training samples.	V C 2
${ℒ}_{\text{FM}}\left[\widehat{v}\right]$	Flow-matching loss.	V C 2
$c_j$	Gaussian-mixture weights, with $c_j\ge 0$ and ${\sum }_j{c}_j=1$.	V C 2
$\widehat{a},\widehat{b}$	Learned scale and shift functions in an affine normalizing-flow layer.	V C 5