Pushing Probabilistic Machine Learning Out of Equilibrium

Since the 1990s, statistical physics has built a theory of driven non-equilibrium processes: systems pushed from one state to another in finite time, never given the chance to relax. About a decade ago, machine learning picked up this idea in the form of diffusion models, starting with the part that makes training work: destroy structure with noise, then learn to reverse it. The rest of the theory, which says how to do exact inference with a process that never reaches equilibrium, is where the idea is now leading. This blog is a summary of where it leads, based on a line of work with many wonderful collaborators.

PART 1The problem we actually care about

Most of the excitement around generative models is about samples: images, molecules, protein structures, materials. Often a sample is exactly the deliverable: a candidate drug to synthesize, a conformation to seed a simulation, a crystal structure to test. Just as often the deliverable is an estimation computed from the distribution: a population, a free energy, a phase boundary, the probability of a rare event. A probabilistic model should serve both kinds of question, and the two share one requirement that is easy to overlook: a sample is only as good as the distribution it comes from. Figure 1 shows four places where that requirement decides the answer.

Sampling and estimation are therefore not a goal and a means. They are two queries against the same object, and they fail in the same way.

So let us state the problem in its full form. There is a target distribution

and we reach it in one of two ways:

The first row is the familiar learning from data setting. The other three are sampling from an unnormalized density, and mathematically they are the same problem. The first row also becomes one of them the moment we want to condition or steer a model trained on data, because $p_\theta(x)\,h(x)$ is again an unnormalized density.

Given $\pi$, there are five things we want to do:

1. Sample $x\sim\pi$.
2. Evaluate the density (likelihood) $\pi(x)$ or $p_\theta(x)$.
3. Estimate expectations $\mathbb{E}_\pi[f]$ (observables, posterior means).
4. Estimate the normalizing constant $Z$, which is the free energy $F=-\beta^{-1}\log Z$ or the model evidence.
5. Condition and steer: sample and estimate under $q \propto h\cdot\pi$ for some new factor $h$.

One constraint shapes everything that follows: the model will be imperfect. Whatever we learn, $p_\theta\neq\pi$. The useful question is therefore not “can we learn $p_\theta = \pi$?” but rather:

The only general way to do that is to correct the model. In its simplest form the correction is importance sampling:

Here $p_\theta$ is easy to sample and $\pi$ is hard to sample. The identity is exact for any $p_\theta$ whose support covers $\pi$. The quality of the model affects only the variance, never the correctness. That is the property we want. The same weights also correct samples: resampling the draws in proportion to $\pi/p_\theta$ (or accepting and rejecting them) turns samples from $p_\theta$ into samples from $\pi$. One mechanism serves both kinds of query. The catch is that it requires evaluating $p_\theta(x)$. This single requirement, a usable density for the proposal, drives the rest of the post.

PART 2Why parameterizing a probabilistic model is computationally hard

A probabilistic model has to deliver three things at once:

• Expressivity: it can represent complicated, multimodal, high-dimensional, symmetric distributions.
• Cheap sampling: drawing $x\sim p_\theta$ costs a modest, fixed amount of computation.
• Tractable density: $p_\theta(x)$, normalized, can be evaluated, or at least the density ratios that inference needs can be.

Each classical model family gives up at least one of these, and in each case the cost comes from a different place.

The equilibrium approach: the cost comes from time

The oldest answer skips the learned model: simulate Langevin dynamics $\dd X_t = \nabla\log\tilde\pi(X_t)\,\dd t + \sqrt{2}\,\dd W_t$, which leaves $\pi$ invariant and never needs $Z$. But exactness holds only as $t\to\infty$, and the time to cross a barrier of height $\Delta$ grows like $e^{\beta\Delta}$. The three demos below run on the same hard double well, so they can be compared directly.

Annealing is the standard remedy: start from a broad, easy distribution and let the barrier grow gradually.

Annealing helps, but only slowly. On this problem it needs about a million steps to get the populations right, because it is exact only in the quasi-static limit $N\to\infty$, where the system stays in equilibrium at every step. In the equilibrium paradigm, the cost of exactness comes from time.

Now keep the annealing just as fast, but change how the chains move. Instead of relaxing with local MCMC steps, carry each chain along with the changing distribution, using the exact transport from one $\pi_\beta$ to the next. This is the analytically optimal diffusion process, the one a perfectly trained diffusion model learns:

Same target, same path, a handful of steps, and the chains are in equilibrium at every $\beta$. A finite-time process can be exact if it transports mass instead of relaxing it. Learning that transport is what diffusion models do. Parts 3 and 4 show how to stay exact when the learned transport is imperfect, using a path weight that reweights each trajectory.

Explicit densities: the cost comes from architecture

The opposite answer is to design $p_\theta$ so that the density is exact. Autoregressive models do this with an ordering, which molecules and materials do not have. Normalizing flows push a Gaussian through an invertible map $x=T(z)$, giving $p_\theta(x) = p_0(z)\,|\det\nabla T(z)|^{-1}$. This is exactly what reweighting needs, and it powered Boltzmann generators [Noé et al. 2019]. The price is the architecture: invertibility with a cheap determinant limits expressivity and scaling.

Continuous flows: the cost comes from Jacobians

Continuous-time flows $\dot x_t = b_t(x_t)$ remove the architectural constraint and train easily with flow matching. Their density still exists:

but the divergence costs $d$ backward passes per solver step, where $d$ is the dimension, for every sample. For a protein with thousands of atoms, evaluating the density becomes far more expensive than generating the sample. The exact density is valuable, but here it is usually too expensive to use.

The pivot: when the exact density is too expensive, an unbiased estimate is enough

When an exact density is cheap, use it; the box above is all you need. When it is not, look again at what inference actually uses. Importance sampling, self-normalized estimators, sequential Monte Carlo and $Z$-estimation all consume weights. A standard fact from computational statistics, the pseudo-marginal principle [Andrieu & Roberts 2009], says that a weight can be replaced by any non-negative unbiased estimate of it without losing exactness. The reason is that we are then doing exact importance sampling on an extended space, the original variable plus the auxiliary randomness of the estimator, whose marginal is the one we care about.

PART 3Diffusion: lifting the model to path space

Everything is a Markov kernel

Before any equations, one point that the rest of the post depends on. Every derivation below uses only two things: Markov transition kernels $\Pf,\Pb$ whose densities we can evaluate, and the marginals they connect. We never use that the kernels are Gaussian, that the state is continuous, or that there is an SDE behind them. A Gaussian denoising step is one choice of kernel. Others include an MCMC move, unmasking a token, appending a token to a sentence, or swapping an atom's chemical species. Every identity, estimator and algorithm in Parts 3–6 therefore applies to all of them. We return to this in Part 7. For now, read $\Pf$ and $\Pb$ as any Markov kernels, with diffusion as the running example.

Two processes, one set of marginals

In continuous time, the running example is a generative SDE on $[0,1]$:

with marginals $p_t$ obeying the Fokker–Planck equation $\partial_t p_t + \nabla\!\cdot(\mu^+_t p_t) = \tfrac12\sigma_t^2\Delta p_t$. The same marginals can be traversed backward in time by a process with backward drift $\mu^-_t(x) = \lim_{h\downarrow0}\mathbb{E}[X_t-X_{t-h}\mid X_t=x]/h$. The two drifts are linked by a remarkably simple relation [Nelson 1967; Anderson 1982]:

The difference between the two directions of time is the score. That is why diffusion models learn scores. Averaging the two drifts gives the probability-flow ODE $\dot x = \mu^+_t - \tfrac12\sigma_t^2\nabla\log p_t$, which has the same marginals. Flows and diffusions are therefore “the same model” up to how well the score is estimated. As we will see, they are very much not the same for inference.

Discretizing with step $\delta$ gives Gaussian kernels, each with a closed-form density:

The key computational fact is easy to miss: evaluating the density of a whole path costs exactly as much as simulating it. That is $N$ network calls and $N$ kernel log-densities, with no Jacobian, trace or determinant.

Time reversal is Bayes' rule in time

A perfectly trained diffusion model satisfies, for every $t$,

This is Bayes' rule written across time: the joint of two consecutive states factorizes either way. Chaining it along the path gives the joint over all states:

This is the joint form of Bayes' rule over the whole path. In measure notation, which carries over unchanged to continuous time:

The central identity

Rearrange the joint form and divide both sides by $p_0(X_0)\prod_t\Pb(X_t\mid X_{t+1})$:

Real models are imperfect, so this pathwise equality becomes an estimator, the Radon–Nikodym Estimator (RNE). The deeper fact, which turns the identity into an inference engine, holds for any pair of kernels, perfect or not. It is the subject of Part 5. Here is the simplest instance, and the pattern to watch for. Let $p_N$ be the actual marginal of the forward chain, meaning the distribution of the samples you really generate. Then, for any backward kernels,

The backward kernels serve as an importance proposal over the paths that end at $X_N$. Nothing requires them to be the exact reversal. If they are exact, the estimator has zero variance. If they are not, it is still unbiased.

PART 4Reading the path weight: non-equilibrium thermodynamics

This is the right place to bring in physics. Not at the start as motivation, and not at the end as decoration, but here, at the moment we hold a path weight. That is because physics has studied exactly this object since the late 1990s, and its results are precisely the estimators we need.

Equilibrium vs. out of equilibrium

MCMC is an equilibrium method: it is correct once the system has relaxed, and the animation in Part 2 shows how long that can be. Generative models are the opposite. They drive a system from a prior to a target in finite time, $0\lt N\ll\infty$, and at intermediate times the system is not in equilibrium with any intermediate Hamiltonian, by design. Physicists call this a driven non-equilibrium process, and the question they asked about it is ours as well: if a system is pushed out of equilibrium, can we still recover equilibrium quantities exactly? The answer, a theme of 1990s statistical physics, is yes, through the fluctuation theorems.

Jarzynski's equality is our normalizing-constant estimator

Part 5 shows that the normalizing constant is the average path weight, $Z_N = \mathbb{E}_{\Pf}\big[\tilde\pi_N(X_N)\prod\Pb\,/\,p_0(X_0)\prod\Pf\big]$. Now make the choice physicists make. Let $U_t$ interpolate between $U_0$ (prior, $p_0 = e^{-U_0}/Z_0$) and $U_N$ (target, $\tilde\pi_N=e^{-U_N}$). Let $\Pf(\cdot\mid X_t)$ be any MCMC kernel that leaves $e^{-U_{t+1}}$ invariant, and let $\Pb$ be its reversal. The path weight then collapses to the work:

$W$ is the physical work: the energy injected when the Hamiltonian is switched while the configuration is held fixed. This is Jarzynski's equality, $\mathbb{E}[e^{-\beta W}] = e^{-\beta\Delta F}$ [Jarzynski 1997]. It is also, symbol for symbol, Neal's annealed importance sampling [Neal 2001], which Neal derived with an explicit nod to Jarzynski. The SMC-sampler framework of Del Moral, Doucet & Jasra (2006) generalized the “backward kernel” $\Pb$ into a free design choice, which is exactly the freedom a learned diffusion model fills.

Crooks' theorem is time reversal; the ELBO is the second law

For general (learned) kernels, normalize $p_0$ and define the work as the negative log path weight, $W := -\log w$. The backward path law built from the normalized target $\pi_N=\tilde\pi_N/Z_N$ is $\pi_N(X_N)\prod_t\Pb$, so

This is Crooks' fluctuation theorem [Crooks 1999]: the probability of a trajectory relative to its time reverse is the exponential of the dissipated work $W_{\rm diss}=W-\Delta F$. Compare it with the time-reversibility box above. A perfect model is exactly the case where this ratio equals $1$. Three consequences, each a familiar ML statement in disguise:

• Second law ⇔ KL ≥ 0. $\mathbb{E}_{\Pf}[W] - \Delta F = \mathrm{KL}(\Pf\,\|\,\Pb) \ge 0$. The mean dissipated work is a KL divergence between path measures.
• ELBO ⇔ free-energy bound. $\log Z_N \ge \mathbb{E}_{\Pf}[\log w]$ is the statement $\Delta F\le\langle W\rangle$. Evidence lower bounds are the second law.
• Perfect model ⇔ reversible process. Zero dissipation, $W\equiv\Delta F$ on every path, is the same thing as exact time reversal, which is the same thing as a constant path weight. A perfectly trained diffusion model is a reversible, finite-time protocol, something physics normally reaches only in the infinitely slow limit.

Crooks already tells us how to reach zero variance

Crooks' theorem also says what an ideal estimator looks like. The path weight is $e^{-W}$, and its variance is zero exactly when $W=\Delta F$ on every path, that is, when the forward and backward path laws coincide. In that case a single trajectory returns $Z$ exactly. So the recipe for a perfect estimator is already written in the theorem: build a process that is reversible in finite time.

Plain switching with MCMC kernels cannot do this. MCMC kernels only re-equilibrate particles where they are; they do not move them along with the changing distribution, so a fast protocol always lags behind and dissipates. The fix is to add a transport that carries particles along as the distribution changes. Physicists discovered exactly this as escorted free-energy simulations [Vaikuntanathan & Jarzynski 2008]: the closer the escort is to the true transport, the less work is dissipated, and a perfect escort gives zero variance. A learned diffusion or flow is such a transport. Learning is dissipation reduction, and a perfectly trained model is the zero-variance estimator that Crooks' theorem describes.

There is a historical irony here. The paper that started diffusion models was titled “Deep Unsupervised Learning using Nonequilibrium Thermodynamics” [Sohl-Dickstein et al. 2015]. The community kept the thermodynamics needed to train: destroy information with a known process, learn to reverse it. It mostly set aside the half needed for inference: fluctuation theorems that turn a finite-time, imperfect, out-of-equilibrium process into exact equilibrium answers. My central claim is that the second half is at least as important as the first.

PART 5One weight, four inferences

All of inference now comes from one template, the same three moves used in Part 3: write the quantity as an integral over paths, multiply and divide by the proposal path density, read it as an expectation. The four inferences differ only in which endpoint density appears and which kernels play which role.

5.1 Marginal likelihood, with the ELBO as a corollary

The derivation at the end of Part 3 is already the marginal-likelihood estimator for any trained diffusion model:

Applying the log and Jensen's inequality gives the diffusion ELBO:

Averaging $K$ path weights inside the log gives a tighter, IWAE-style bound that becomes exact as $K\to\infty$. The estimator refers to the model you actually sample from, in discrete time, and needs only forward passes: no divergence, no adjoint.

5.2 Imperfect model, exact statistics

Now suppose the target is an energy, $\tilde\pi_N = e^{-U}$, and we have a sampler $\Pf$ that reaches it only approximately. Expectations under the target are ratios of path-weighted averages:

The self-normalized estimator $\sum_i w_i f(X_N^{(i)})/\sum_i w_i$ is consistent whatever the quality of the model. This is the Boltzmann-generator recipe of Noé et al., with the flow replaced by a diffusion and the Jacobian replaced by a kernel ratio. The effective sample size $\mathrm{ESS}=(\sum w)^2/\sum w^2$ is an honest, built-in diagnostic: it measures how far from reversible the learned process still is.

5.3 Normalizing constants and free energies

The denominator above is itself the normalizing constant:

Free energies are among the most consequential quantities in computational science: binding affinities, solvation, phase stability, conformational equilibria. Classical estimators differ only in the choice of kernels in this one formula:

FEAT (Free-energy Estimators with Adaptive Transport) [Du*, He*, Vargas, Wang, Gomes, Hernández-Lobato, Vanden-Eijnden, NeurIPS 2025] turns this table into a recipe. Given samples from both end states and their energies:

1. Learn a transport between the two states with stochastic interpolants. This is simulation-free and uses the same machinery as training a diffusion model.
2. Plug the learned kernels into the estimator: escorted Jarzynski and controlled Crooks give consistent estimators for any learned transport.
3. Run both directions and combine them to reduce variance. Variational bounds from both sides sandwich the answer.

The message is short: when you can learn a good diffusion model, you get the free energy almost for free. The learning reduces variance and the fluctuation theorem guarantees correctness. We validated this on very different systems: high-dimensional Gaussian mixtures, Lennard-Jones clusters against non-interacting references, coarse-grained monatomic ice against an LJ solid, biomolecular solvation (vacuum against implicit solvent), the transition free energy between metastable states of alanine dipeptide, and the cold and hot phases of $\varphi^4$ lattice field theory. In each case the estimates recovered reference values at roughly an order of magnitude less compute.

The follow-up, Free-Energy Estimation on Any State Space [He, Ou, Vargas, Li, Hernández-Lobato, Domingo-Enrich, Du, NeurIPS 2026], has two parts. It formulates transport through a general Markov generator, so the estimator applies to discrete, multimodal and autoregressive spaces. It also shows that, at optimality, the commonly used estimators above are special cases. In practice this pushed learned free-energy estimation to biomolecules with more than 1,000 atoms.

5.4 Steering: sample from what you did not train on

Often we have a good model $p_N$ and want something nearby: $q_N(X_N) \propto h(X_N)\,p_N(X_N)$. A surprising range of tasks has exactly this form:

Instead of sampling from $p_N$, we want to actively sample from $q_N$. So we introduce a new pair of processes $\Qf,\Qb$ that push mass toward high-$q_N$ regions, and apply the same three moves:

With $q_0=p_0$ and $q_N\propto h\,p_N$, the weight contains $h(X_N)\,p_N(X_N)/p_0(X_0)$. This marginal ratio is the only intractable quantity, and it is exactly what the central identity supplies [RNE]:

Applying one weight to the whole path lets its variance compound over $N$ steps. The second design freedom fixes this: a curve of intermediate targets $q_t$, with $q_N$ fixed by the task. The same derivation on any segment $[t',t'']$ gives

so variance can be controlled one segment at a time. The one-step version is the building block of the next section.

Examples from our work: annealing an alanine-dipeptide model from 800 K to 300 K with roughly 10× compute savings; composing conditional generators to design molecules that bind two targets, with affinities better than known binders; reward tilting of discrete (masked) diffusion on ImageNet-256 for prompt alignment; and inpainting-style conditioning for structure-based drug design with equivariant diffusion.

5.5 SMC and replica exchange: two ways to spend the same weights

Define the incremental weight of one step from level $t$ to level $t+1$:

The product telescopes back to the full path weight, $\prod_{t=0}^{N-1}G_t = \dfrac{q_N(X_N)\prod_t\Qb}{q_0(X_0)\prod_t\Qf}$, because every intermediate $q_t$ appears once in a numerator and once in a denominator. Two classical algorithms use this one quantity in two different ways.

Sequential Monte Carlo moves a population of particles forward together. At each level it multiplies every particle's weight by $G_t$, and it resamples (duplicating heavy particles and dropping light ones) whenever the ESS drops. The genealogy collapses where the proposal is wrong, and the population arrives at $t=N$ as a weighted batch.

Replica exchange (parallel tempering along diffusion time) keeps one replica per level. Neighbouring levels $(t,t+1)$ propose to exchange information through the kernels: push the lower state up, $y_{t+1}\sim\Qf(\cdot\mid x_t)$, and pull the upper state down, $y_t\sim\Qb(\cdot\mid x_{t+1})$. The swap is accepted with a Metropolis ratio that is a forward incremental weight divided by a backward one:

Fresh prior samples enter at $t=0$, travel up the ladder through accepted swaps, and stream out at $t=N$ one per round. CREPE [He, Jeha, Potaptchik, Zhang, Hernández-Lobato, Du, Syed, Vargas, ICLR 2026] develops this view for controlling diffusion models. Unlike SMC, it generates samples sequentially, keeps diversity high after burn-in, and allows online refinement or early stopping.

5.6 A compute knob for correctness

This leads to inference-time scaling with a guarantee. More particles, more segments and more exchange rounds improve accuracy, and the estimators are consistent, so extra compute brings you closer to the true answer rather than to the model's bias. A model's quality is no longer a ceiling on accuracy. It is a rate: how much compute exactness costs.

PART 6Where training enters: learning the time reversal

Why does training appear only now? Because we can now say precisely what training is for. Everything in Part 5 is correct for any pair of processes; training only controls the variance of the path weight. The ideal is the time-reversibility identity of Part 3: the forward and backward path laws coincide, the path weight equals $Z$ on every path, and every estimator has zero variance. Training is therefore learning the time reversal: fitting the generative process so that it reverses the backward one.

Two settings differ in one thing only: whether we can sample the backward path measure.

From data: the ELBO is the path KL

With data, the backward path measure $\pi_N(X_N)\prod_t\Pb(X_t\mid X_{t+1})$ is something we can sample exactly. Draw a data point, then noise it with a fixed process chosen so that its conditional marginals $\Pb(X_t\mid X_N)$ have closed form. Fit the generative path measure to it:

where $\mathrm{ELBO}_\theta$ is exactly the bound of Section 5.1. Maximizing the ELBO is minimizing the KL between path measures, i.e. the dissipated work.

From energy: target score identity and fixed-point matching

With only $\tilde\pi_N = e^{-U}$, there is no data to noise. We can still write the marginal score of the noised target, and remarkably it needs only forces. Keep the same noising kernel $\Pb(x_t\mid x_N)=\mathcal N(x_t;\alpha_t x_N, s_t^2 I)$. Then the noised score satisfies the target score identity:

This suggests regressing a score network onto forces, target score matching:

whose minimizer is the conditional expectation above, i.e. the exact score. (One can also estimate the identity directly by self-normalized importance sampling over the posterior $\pi_N(x_N\mid x_t)\propto e^{-U(x_N)}\,\mathcal N(x_N; x_t/\alpha_t, s_t^2/\alpha_t^2)$.) The catch is visible: the loss asks for $X_N\sim\pi_N$, which is what we are trying to learn. The resolution is a fixed point. Let the current model supply the samples, and let its own path weights correct them toward $\pi_N$:

where $X_t' = \alpha_tX_N + s_t\varepsilon$ and $\bar w$ is the self-normalized path weight from Part 5. The target $\pi_N$ is a fixed point of this map: if the sampler already produces $\pi_N$, the regression target is the exact score and the update returns the same model. The path weights keep each iteration consistent rather than self-reinforcing.

The energy-trained foundation samplers we built, ATLAS and JANUS, run this loop at scale: target score matching with fixed-point iteration on a replay buffer, plus a path-weight regularizer. The regularizer is what gets the mode weights right: score regression shapes each mode well, while the path weights correct the relative populations between modes. At deployment the same weights reweight to the exact Boltzmann distribution and estimate free energies. ATLAS [Cheng, Blessing, Yu, Neumann, Li, Domingo-Enrich, Du 2026] reproduces molecular-dynamics and Monte-Carlo distributions of 2D amorphous systems with free-energy errors below 0.2% while using about 500× fewer energy evaluations. It recovers short-range order in Cu–Zr and Cr–Co–Ni metallic glasses, and it amortizes over system size, temperature and composition, so inverse design costs hundreds of times less.

Beyond NVT: generalized ensembles

Real materials are not simulated only at fixed particle number, volume and temperature. The physically relevant ensembles let other variables fluctuate:

Each step down the table adds a variable of a different type, and the state space becomes a product of continuous and discrete parts. JANUS [Blessing, Cheng, Schebek, Rogal, Li, Domingo-Enrich, Du 2026] couples continuous diffusion over positions with masked discrete diffusion over chemical identities, through one equivariant network trained directly from energy evaluations. It uses no pre-generated equilibrium data. It reproduces Monte Carlo observables and free energies with about three orders of magnitude fewer energy evaluations. It allows chemical control and inverse design for multicomponent alloys, and for semiconductors it explores vacancies and dopants across 15 elements, recovering known defects and proposing new candidates validated with DFT. The step from NVT to $\mu$VT forces the theory off $\mathbb{R}^d$. Because everything was written with Markov kernels, nothing has to be re-derived.

PART 7Extending the picture: any state space, any path

What the identities actually need

Look back at Parts 3–6. Every derivation used the same three ingredients:

1. a reference measure on the state space, so kernels have densities;
2. Markov kernels $\Pf,\Pb$ whose densities we can evaluate;
3. the time-reversal relation (Bayes' rule in time) linking them through the marginals.

So the whole machinery (marginal likelihood, exact statistics, free energies, steering with SMC or replica exchange, training by learning the time reversal) instantiates on any process built from kernels:

In continuous time, (2) and (3) are statements about a Markov generator $\mathcal L_t$ and its time reversal, and the path weight is the Radon–Nikodym derivative between the two path laws. For diffusions this is Girsanov's theorem. For continuous-time Markov chains with jump rates $\lambda^P_t(x\to y)$ and $\lambda^Q_t(x\to y)$, and total exit rates $\lambda_t(x)=\sum_{y\neq x}\lambda_t(x\to y)$, it is

On a product space the generator is a sum, and so is the log path weight. Free Energy Estimation on Any State Space [He et al., NeurIPS 2026] generalizes these results to any state space by writing them in terms of the generator of the Markov process.

Autoregressive models are non-equilibrium processes too

This is the point I most want readers to keep from Part 7. Following [He, Guo, Hernández-Lobato, Du 2026], let $\mathcal V$ be the vocabulary and $\tilde{\mathcal V}:=\mathcal V\cup\{\varnothing\}$ add an empty padding token, so that sequences of different lengths live in one state space:

with padding absorbing: $p(x_i=\varnothing\mid x_{1:i-1})=1$ whenever $x_{i-1}=\varnothing$. The state at step $i$ is the prefix $x_{0:i}$. Generation appends one token:

The backward kernel follows from Bayes' rule in time, exactly as in Part 3. Since $p_i(x_{0:i})=p_{i-1}(x_{0:i-1})\,p(x_i\mid x_{\lt i})$,

it deterministically deletes the last token. Plug this into the central identity:

An autoregressive model is a non-equilibrium process with a trivial backward kernel. Only one path leads to each endpoint, so the path weight has zero variance and equals the exact likelihood. Every section of this post therefore applies to language models without modification:

• Marginal likelihood (5.1): exact, the familiar product of conditionals, as a special case rather than a separate fact.
• Exact statistics and normalizing constants (5.2–5.3): expectations under, and the partition function of, targets such as the powered $\pi(x_{0:N})\propto p(x_{0:N})^{\beta}$, the tilted $\pi(x_{0:N})\propto p(x_{0:N})\,e^{r(x_{0:N})}$, or the product $\pi(x_{0:N})\propto p(x_{0:N})\,p'(x_{0:N})$. These are free-energy problems on sequence space.
• Steering (5.4): the prefix marginals $p_i(x_{0:i})$ are exact for an autoregressive model, so intermediate targets such as $p_i(x_{0:i})^{\beta}$ need no RNE approximation.
• SMC and replica exchange (5.5): particles become prefixes, and levels become prefix lengths. With a proposal model $q$, the incremental weight for the powered target is simply $$G(x_{0:i},x_{0:i-1}) = \frac{p(x_i\mid x_{\lt i})^{\beta}}{q(x_i\mid x_{\lt i})},$$ and the same weight drives both algorithms.

This is the basis of Recipes for Steering and Scaling LLMs via Sampling [He, Guo, Hernández-Lobato, Du 2026]. SMC and replica exchange steer generation toward powered, product and tilted versions of the base distribution, and quality scales with inference compute without external supervision or a learned reward model, improving on best-of-$N$ and standard MCMC baselines. Masked diffusion sits between the two extremes: its backward kernel is a random masking order, so its path weight averages over orders.

Tilting whole paths: rare events and control

So far we have only tilted the endpoint. Many scientific questions concern the path itself: how a protein folds, how a reaction crosses its barrier, how often a rare transition happens. The generalization is a Feynman–Kac tilt of the entire path measure:

The tilted process is again Markov. It is a Doob $h$-transform whose extra drift $\sigma_t\sigma_t^\top\nabla\log h_t$ solves a stochastic optimal control problem with exactly this objective. For transition paths between metastable states $A$ and $B$, the relevant $h$ is the committor, the probability of reaching $B$ before $A$. The optimal control is proportional to its log-gradient, and it drives trajectories through the reactive channel that unbiased simulation almost never visits.

This is the oldest part of my own stack, and it is the same picture. In order: stochastic optimal control for collective-variable-free transition path sampling [Holdijk*, Du*, et al., NeurIPS 2023]; Doob's Lagrangian, a sample-efficient variational approach to transition path sampling [Du*, Plainer*, Brekelmans*, et al., NeurIPS 2024]; and rare-event analysis via stochastic optimal control [Du, He, Zhang, Vanden-Eijnden, Domingo-Enrich 2026]. The last casts committor estimation as a hitting-time control problem with an off-policy value-matching loss, which yields reaction rates and equilibrium constants. The marginal path weight gives thermodynamics. The tilted path measure gives kinetics.

The whole picture in one sentence

PART 8What is still hard

An honest account should say where the picture strains.

Variance is extensive

Importance sampling needs roughly $e^{\mathrm{KL}}$ samples [Chatterjee & Diaconis 2018], and the dissipated work, like free energy, grows with system size. A model that is 0.01 nats per atom from reversible is many nats from reversible for $10^4$ atoms. Splitting the path (SMC, replica exchange), local or conditional structure, and amortization are necessary, not optional. Weight degeneracy in very large systems is the central open problem.

The interpolant matters

Every method in this post walks along a curve of intermediate distributions $q_t$, and the choice of curve decides how much dissipation is unavoidable before any learning happens. Today there are two families, and neither is what we ultimately want:

• Annealing-based curves are defined through energies, for example the geometric path $\tilde\pi_t\propto\tilde\pi_0^{1-\beta_t}\tilde\pi_N^{\beta_t}$ or temperature ladders. Their great advantage is that every intermediate unnormalized density is known, so MCMC kernels, Jarzynski weights and replica exchange are available at every level. Their weakness is the geometry of the path. Mass does not move along a geometric mixture; it vanishes in one mode and reappears in another. Separated modes are weighted wrongly at intermediate $\beta$, and physical systems can cross first-order phase transitions along the path, where any finite-time protocol dissipates heavily.
• Transport-based curves are defined through samples, as in stochastic interpolants $X_t=\alpha_tX_0+\beta_tX_1+\gamma_tZ$. They are smooth, mass moves continuously, and the velocity or score is learnable by simple regression. That is why FEAT and diffusion training are stable. Their weakness is that the intermediate densities are not known in closed form, only up to learning error. They also need samples at the endpoints, which in the energy setting must be bootstrapped (Section 6).

We need better: curves that keep the tractable intermediate energies of annealing and the smooth mass transport of interpolants, that avoid phase transitions along the way, and that are chosen to minimize dissipation for the system at hand rather than fixed in advance. Near equilibrium, excess work is governed by a friction metric, and minimum-dissipation protocols are geodesics in thermodynamic length [Sivak & Crooks 2012]. Learning the interpolant jointly with the transport, guided by such a metric, looks like the next frontier.

Exploration in energy-based training

Without data, the backward process has to be discovered, and a sampler that never visits a basin cannot learn it. The fixed-point loop of Section 6 converges to $\pi$ only if the buffer eventually covers $\pi$'s mass. Path weights reveal a missing mode (through a collapsing ESS) but do not supply it. Coupling learned samplers with tempering, replica exchange and enhanced-sampling biases, so that exploration and learning reinforce each other, is still underexplored.

TAKEAWAYSWhat I hope you remember

1. Samples and estimations both need debiasing.A learned model is always imperfect. Whether you want samples or estimations (likelihoods, expectations, free energies, conditionals), the general fix is the same: reweight by the ratio of target to model density.
2. Densities are valuable, but exact ones are expensive.With an exact density, every inference problem is a one-line reweighting formula. Expressive models rarely afford one: for MCMC the cost comes from time, for normalizing flows from architecture, for continuous flows from Jacobians.
3. Transport, not relaxation, makes finite time exact.On a hard double well, MCMC stays stuck and annealing freezes in the wrong proportions, while the analytically optimal diffusion process stays in equilibrium along the whole path with only a handful of steps.
4. Diffusion lifts the model to path space.Time reversal is Bayes' rule in time, $p_0\,\dd\Pf/p_N\,\dd\Pb = 1$. It turns the intractable marginal density into the conditional expectation of a cheap path weight, unbiased for any model.
5. The path weight is the work.Time reversal is Crooks' theorem, normalizing constants come from Jarzynski's equality, and the ELBO is the second law. Crooks also says what a zero-variance estimator is: a process that is reversible in finite time.
6. One weight, every estimation, with inference-time scaling.Marginal likelihoods, debiased expectations, free energies (FEAT) and steering (RNE) are all reweighting of paths. SMC and replica exchange (CREPE) use the same incremental weight, and accuracy improves with the compute spent at inference time.
7. Training means learning the time reversal.From data: denoising score matching, or the ELBO, which is the path KL. From an energy: target score matching in a fixed-point loop, with path-weight regularization to get the mode weights right (ATLAS, JANUS).
8. Everything is a Markov kernel, so it works on any state space.Written through the generator of a Markov process, the same identities hold on discrete, multimodal and mixed spaces, including generalized ensembles such as $NPT$ and $\mu VT$.
9. Autoregressive models are non-equilibrium processes too.Their backward kernel deletes the last token, so the path weight is the exact likelihood, and every tool here (free energies, SMC, replica exchange, steering) applies to language models unchanged.
10. Push probabilistic ML out of equilibrium.Equilibrium methods are exact only in the infinite-time limit. Non-equilibrium methods are exact at any finite time, because each path carries its own correction, and learned processes keep that correction small.

Acknowledgements

This line of work is deeply collaborative. Special thanks to Jiajun He, Francisco Vargas, Carles Domingo-Enrich, Eric Vanden-Eijnden, José Miguel Hernández-Lobato, Denis Blessing, Mouyang Cheng, and many others who shaped these ideas.

REFERENCESSelected references

From this line of work

1. J. He, J. M. Hernández-Lobato, Y. Du†, F. Vargas†. RNE: plug-and-play diffusion inference-time control and energy-based training. ICLR 2026. arXiv:2506.05668
2. Y. Du*, J. He*, F. Vargas, Y. Wang, C. P. Gomes, J. M. Hernández-Lobato, E. Vanden-Eijnden. FEAT: Free energy Estimators with Adaptive Transport. NeurIPS 2025. arXiv:2504.11516
3. J. He, Z. Ou, F. Vargas, Y. Li, J. M. Hernández-Lobato, C. Domingo-Enrich, Y. Du. Free Energy Estimation on Any State Space. NeurIPS 2026. arXiv:2605.31063
4. J. He, P. Jeha, P. Potaptchik, L. Zhang, J. M. Hernández-Lobato, Y. Du, S. Syed, F. Vargas. CREPE: Controlling Diffusion with Replica Exchange. ICLR 2026. arXiv:2509.23265
5. M. Cheng, D. Blessing, B. Yu, G. Neumann, M. Li, C. Domingo-Enrich, Y. Du. ATLAS: A Foundation Neural Sampler for Amorphous Materials. 2026. arXiv:2607.19198
6. D. Blessing, M. Cheng, M. Schebek, J. Rogal, M. Li, C. Domingo-Enrich, Y. Du. JANUS: A Multi-modal Foundation Neural Sampler for Disordered Materials. 2026. arXiv:2608.19116
7. J. He, Z. Guo, J. M. Hernández-Lobato, Y. Du. Recipes for Steering and Scaling LLMs via Sampling. 2026. arXiv:2608.26120
8. Y. Du, J. He, D. Zhang, E. Vanden-Eijnden, C. Domingo-Enrich. Rare Event Analysis via Stochastic Optimal Control. 2026. arXiv:2604.13213
9. Y. Du*, M. Plainer*, R. Brekelmans*, C. Duan, F. Noé, C. P. Gomes, A. Aspuru-Guzik, K. Neklyudov. Doob's Lagrangian: A Sample-Efficient Variational Approach to Transition Path Sampling. NeurIPS 2024. arXiv:2410.07974
10. L. Holdijk*, Y. Du*, F. Hooft, P. Jaini, B. Ensing, M. Welling. Stochastic Optimal Control for Collective Variable Free Sampling of Molecular Transition Paths. NeurIPS 2023. arXiv:2207.02149

Selected related work

Inference-time steering of diffusion models

1. L. Wu, B. L. Trippe, C. A. Naesseth, D. M. Blei, J. P. Cunningham. Practical and Asymptotically Exact Conditional Sampling in Diffusion Models (TDS). NeurIPS 2023. arXiv:2306.17775
2. R. Singhal, Z. Horvitz, R. Teehan, M. Ren, Z. Yu, K. McKeown, R. Ranganath. A General Framework for Inference-time Scaling and Steering of Diffusion Models (FK steering). ICML 2025. arXiv:2501.06848
3. M. Skreta et al. Feynman-Kac Correctors in Diffusion: Annealing, Guidance, and Product of Experts (FKC). ICML 2025. arXiv:2503.02819

Neural samplers

1. F. Vargas, S. Padhy, D. Blessing, N. Nüsken. Transport meets Variational Inference: Controlled Monte Carlo Diffusions (CMCD). ICLR 2024. arXiv:2307.01050
2. M. S. Albergo, E. Vanden-Eijnden. NETS: A Non-Equilibrium Transport Sampler. 2024. arXiv:2410.02711
3. D. Blessing, L. Richter, J. Berner, E. Malitskiy, G. Neumann. Bridge Matching Sampler: Scalable Sampling via Generalized Fixed-Point Diffusion Matching (BMS). 2026. arXiv:2603.00530
4. A. Havens, B. Karrer, N. Shaul. Flow Sampling: Learning to Sample from Unnormalized Densities via Denoising Conditional Processes. ICML 2026. arXiv:2605.03984

Foundations

1. C. Jarzynski. Nonequilibrium equality for free energy differences. Phys. Rev. Lett. 1997.
2. G. E. Crooks. Entropy production fluctuation theorem and the nonequilibrium work relation for free energy differences. Phys. Rev. E 1999.
3. S. Vaikuntanathan, C. Jarzynski. Escorted free energy simulations: improving convergence by reducing dissipation. Phys. Rev. Lett. 2008.
4. D. A. Sivak, G. E. Crooks. Thermodynamic metrics and optimal paths. Phys. Rev. Lett. 2012.
5. E. Nelson. Dynamical Theories of Brownian Motion. 1967. B. D. O. Anderson. Reverse-time diffusion equation models. Stoch. Proc. Appl. 1982.
6. R. M. Neal. Annealed importance sampling. Statistics and Computing 2001.
7. P. Del Moral, A. Doucet, A. Jasra. Sequential Monte Carlo samplers. JRSS-B 2006.
8. R. H. Swendsen, J.-S. Wang. Replica Monte Carlo simulation of spin-glasses. Phys. Rev. Lett. 1986.
9. C. Andrieu, G. O. Roberts. The pseudo-marginal approach for efficient Monte Carlo computations. Ann. Stat. 2009.
10. S. Chatterjee, P. Diaconis. The sample size required in importance sampling. Ann. Appl. Probab. 2018.
11. J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, S. Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. ICML 2015.
12. J. Ho, A. Jain, P. Abbeel. Denoising diffusion probabilistic models. NeurIPS 2020. Y. Song et al. Score-based generative modeling through SDEs. ICLR 2021.
13. V. De Bortoli et al. Target score matching. 2024. T. Akhound-Sadegh et al. Iterated denoising energy matching. ICML 2024.
14. F. Noé, S. Olsson, J. Köhler, H. Wu. Boltzmann generators. Science 2019.
15. R. T. Q. Chen et al. Neural ODEs. NeurIPS 2018. W. Grathwohl et al. FFJORD. ICLR 2019.
16. M. S. Albergo, N. M. Boffi, E. Vanden-Eijnden. Stochastic interpolants: a unifying framework for flows and diffusions. 2023. Y. Lipman et al. Flow matching. ICLR 2023.