Pushing Probabilistic Machine Learning Out of Equilibrium

How out-of-equilibrium diffusion processes give us the most powerful generative models and inference algorithms in probabilistic ML today

Yuanqi Du

Microsoft Research

Joint work with

Jiajun He

Cambridge

Francisco Vargas

Cambridge · Xaira

Carles Domingo-Enrich

MSR New England

Eric Vanden-Eijnden

NYU · Capital Fund Mgmt

José Miguel Hernández-Lobato

Cambridge

Denis Blessing

KIT

Mouyang Cheng

MIT

and Zijing Ou, Yingzhen Li, Zongyu Guo, Yuanqing Wang, Paul Jeha, Peter Potaptchik, Leo Zhang, Saifuddin Syed, Dinghuai Zhang,
Botao Yu, Gerhard Neumann, Mingda Li, Maximilian Schebek, Jutta Rogal, and many others

Roadmap

Today

The problem

samples, estimations, and why densities are expensive

Path space & physics

time reversal, the path weight, the work

Inference & training

free energy, steering, learning from data or energy

Applications

ATLAS for glasses, JANUS for alloys and defects

Any state space

any Markov process: language models, whole paths

The problem · what we ask of models

We ask models for samples and estimations

Samples: a drug candidate, a conformation, a structure

Estimations: a free energy, a phase boundary, a rate

A learned model is always imperfect, so both need debiasing.

The problem · the fix

Debias by reweighting

$\mathbb E_\pi[f]=\mathbb E_{x\sim p_\theta}\!\Big[\underbrace{\textcolor{#E76F51}{\dfrac{\pi(x)}{p_\theta(x)}}}_{\textcolor{#E76F51}{\text{weight } w(x)}}\,f(x)\Big]$

Exact for any model; its quality only sets the variance.

With an exact density, all one line

Question	$w=\tilde\pi/p_\theta$
Expectations	$\mathbb E[wf]/\mathbb E[w]$
Free energy	$Z=\mathbb E[w]$
Samples	resample $\propto w$
Steering	weight by $h(x)$

Exact densities are valuable, but expensive for expressive models.

The problem · why it is hard

Why it is hard: expressive, cheap, or tractable, pick two

$\mathbb E_\pi[f]=\mathbb E_{x\sim p_\theta}\!\Big[\dfrac{\pi(x)}{\textcolor{#E76F51}{p_\theta(x)}}\,f(x)\Big]$

Reweighting needs the model's likelihood $p_\theta(x)$, for every sample. Which models can give it cheaply?

MCMC cost from time

Normalizing flows cost from architecture

Continuous flows cost from Jacobians

Diffusion models change the space

on paths, density costs as much as sampling, and an unbiased estimate is all reweighting needs

The problem · equilibrium methods

Equilibrium: the cost comes from time

The problem · the way out

Out of equilibrium: transport, not relaxation

Same path, but each sample is carried along by the exact transport, the process a perfect diffusion model learns.

MCMC

stuck in one well

Annealing

right only after a very long schedule

Transport

right in a few steps

Finite time can be exact if it transports mass instead of relaxing it.

Path space · diffusion models

Diffusion models: add noise, learn to reverse it

The generative models behind a decade of breakthroughs:

Images & video

Ho 2020 · Song 2021 · Rombach 2022

Protein design

RFdiffusion · Watson, Nature 2023

Biomolecular structure

AlphaFold 3 · Abramson, Nature 2024

Materials

MatterGen · Zeni, Nature 2025

Molecules & docking

DiffSBDD · Schneuing, Harris, Du et al., Nat. Comput. Sci. 2024 · DiffDock 2023

Weather

GenCast · Price, Nature 2025

Two processes through the same marginals $p_t$:

$\begin{aligned}\Pf\ \text{(generation, }0\to N)&:\quad \mathrm dX_t=\mu^+_t(X_t)\,\mathrm dt+\sigma_t\,\mathrm dW_t\\[4pt] \Pb\ \text{(noising, }N\to 0)&:\quad \mathrm dX_t=\mu^-_t(X_t)\,\mathrm dt+\sigma_t\,\mathrm d\overleftarrow{W}_t\end{aligned}$

Time reversal: their drifts differ by the score, the only unknown

$\mu^+_t-\mu^-_t=\sigma_t^2\,\textcolor{#E76F51}{\nabla\log p_t}$

Fokker–Planck: move the diffusion term into the drift

$\partial_tp_t=-\nabla\!\cdot(\mu^+_tp_t)+\tfrac12\sigma_t^2\Delta p_t=-\nabla\!\cdot\big((\mu^+_t-\tfrac12\sigma_t^2\textcolor{#E76F51}{\nabla\log p_t})\,p_t\big)$

so the ODE $\dot x=\mu^+_t-\tfrac12\sigma_t^2\textcolor{#E76F51}{\nabla\log p_t}$ has the same marginals: flows and diffusions are equivalent up to the score.

Learning a diffusion (or a flow) = learning the score at every noise level.

Sohl-Dickstein et al. 2015 · Ho et al. 2020 · Song et al. 2021 · Nelson 1967, Anderson 1982 · Lipman et al. 2023, Albergo et al. 2023 (flow matching, interpolants)

Path space · the path weight

Path weights: time reversal is Bayes' rule in time

$p_t(X_t)\,\Pf(X_{t+1}\mid X_t)=p_{t+1}(X_{t+1})\,\Pb(X_t\mid X_{t+1})$

⟹

$\dfrac{p_N(X_N)}{p_0(X_0)}=\dfrac{\prod_t\Pf(X_{t+1}\mid X_t)}{\prod_t\Pb(X_t\mid X_{t+1})}$

intractable marginals = a ratio of tractable kernels: one network call per step, no Jacobians

Path weight $\;w(X_{0:N})=\dfrac{\tilde\pi(X_N)\prod_t\Pb(X_t\mid X_{t+1})}{p_0(X_0)\prod_t\Pf(X_{t+1}\mid X_t)}\;$ computable for any model; $W=-\log w$ is the work.

He, Hernández-Lobato, Du†, Vargas†, ICLR 2026 (RNE) · Crooks 1999 · right: a real path of the toy sampler with its kernels

Path space · the physics

The path weight is the work

Jarzynski

$\mathbb E_{\Pf}\big[e^{-W}\big]=e^{-\Delta F}$

= annealed importance sampling

Crooks

$\dfrac{\dd\Pb}{\dd\Pf}=e^{-(W-\Delta F)}$

= time reversal

Second law

$\langle W\rangle\ge\Delta F$

= the ELBO

Zero variance $\iff$ reversible in finite time $\iff$ a perfect diffusion model.

Inference · the general recipe

One derivation, every estimator

what we want: an integral against $\tilde\pi=e^{-\beta U}$

$I[g]=\displaystyle\int g(x_N)\,\textcolor{#E76F51}{\tilde\pi(x_N)}\,\dd x_N$

insert the backward kernels; they integrate to one

$=\displaystyle\int g(x_N)\,\textcolor{#E76F51}{\tilde\pi(x_N)}\,\textcolor{#4C72B0}{\prod_t\Pb(x_t\mid x_{t+1})}\;\dd x_{0:N}$

multiply and divide by the forward path density

$=\displaystyle\int g(x_N)\,\underbrace{\frac{\textcolor{#E76F51}{\tilde\pi(x_N)}\,\textcolor{#4C72B0}{\prod_t\Pb}}{\textcolor{#1F7A6F}{p_0(x_0)\prod_t\Pf}}}_{w(x_{0:N})}\;\textcolor{#1F7A6F}{p_0(x_0)\prod_t\Pf}\;\dd x_{0:N}$

read it as an average over generated paths

exact for any model with full support

$=\mathbb E_{\Pf}\big[\,g(X_N)\,w(X_{0:N})\,\big]$

$g\equiv1$

$Z=\mathbb E_{\Pf}[w]$

free energy

$g=f$

$\mathbb E_\pi[f]=\mathbb E_{\Pf}[f\,w]\,/\,\mathbb E_{\Pf}[w]$

exact averages

$\tilde\pi\to h\,\tilde\pi$

$\mathbb E_q[f]=\mathbb E_{\Pf}[f\,h\,w]\,/\,\mathbb E_{\Pf}[h\,w]$

steering

Paths give the model's likelihood: noise $x_N$ back with $\Pb$ and average

$p_\theta(x_N)=\mathbb E_{X_{0:N-1}\sim\Pb(\cdot\mid x_N)}\Big[\dfrac{p_0(X_0)\prod_t\Pf(X_{t+1}\mid X_t)}{\prod_t\Pb(X_t\mid X_{t+1})}\Big]$

Du, blog: Pushing Probabilistic ML Out of Equilibrium, Part 5 · Jarzynski 1997 · Neal 2001 (AIS) · Du*, He* et al., NeurIPS 2025 (FEAT) · He, Hernández-Lobato, Du†, Vargas†, ICLR 2026 (RNE)

Inference · free energy

FEAT: free energy with learned transport

1. Learn a transport between the two states (stochastic interpolants).

2. Plug it into the path weight estimator, in both directions.

Consistent for any transport; learning only reduces the variance.

Du*, He*, Vargas, Wang, Gomes, Hernández-Lobato, Vanden-Eijnden, NeurIPS 2025 (FEAT) · He, Ou, Vargas, Li, Hernández-Lobato, Domingo-Enrich, Du, NeurIPS 2026

Inference · free energy results

Free energy: reference values at ~10× less compute

Method	Gaussian mixture 16 → 40 modes		Lennard-Jones free → interacting			Alanine dipeptide solvation · transition
	d = 40	d = 100	55×3	79×3	128×3	22×3	22×3
Reference	0	0	234.77	357.43	595.98	29.43	−4.25
Targeted FEP	0.09 ± 0.26	−17.96 ± 1.49	232.06 ± 0.03	*	*	29.47 ± 0.22	−4.78 ± 0.32
FEAT	0.04 ± 0.04	−5.34 ± 1.52	232.47 ± 0.15	356.74 ± 0.79	595.04 ± 0.52	29.38 ± 0.04	−4.56 ± 0.08

Peptide	ALA-4 d = 129	Chignolin d = 525
Reference	107.5	320.19
FEAT	109.91 ± 2.55	320.02 ± 0.70

chignolin

Du*, He* et al., NeurIPS 2025 (FEAT) · He, Ou, Vargas, Li, Hernández-Lobato, Domingo-Enrich, Du, NeurIPS 2026 · * not converged

Inference · steering

Steering: scaling sampling at inference time

Sample $q\propto h\,p$ instead of $p$, without retraining:

He, Hernández-Lobato, Du†, Vargas†, ICLR 2026 (RNE) · Wu et al. 2023 (TDS) · Singhal et al. 2025 (FK steering) · Skreta et al. 2025 (FKC)

Inference · inference-time scaling

One weight, two algorithms: SMC and replica exchange

Work of one step

$W_t=-\log\dfrac{q_{t+1}(x_{t+1})\,\Pb(x_t\mid x_{t+1})}{q_t(x_t)\,\Pf(x_{t+1}\mid x_t)}$

SMC: resample with

$w^i\propto e^{-\sum_{s\le t}W_s^i}$

Replica exchange: accept a swap with

$\min\big\{1,\ e^{-W_t^{\uparrow}+W_t^{\downarrow}}\big\}$

↑ replica moved up a level, ↓ moved down

More particles or rounds, more accuracy: inference-time scaling, with an unbiased $Z$ from both.

He, Jeha, Potaptchik, Zhang, Hernández-Lobato, Du, Syed, Vargas, ICLR 2026 (CREPE) · Del Moral, Doucet & Jasra 2006

Inference · steering results

Steering results

Annealing a sampler from 800 K to 300 K

trained at 800 K

steered to 300 K

reference 300 K

alanine dipeptide, Ramachandran plots $(\phi,\psi)$ · 10× less compute than sampling at 300 K

Composition

$p(x\mid\text{pocket}_1)\,p(x\mid\text{pocket}_2)$: molecules binding two targets

fraction of designs better than known binders

He, Hernández-Lobato, Du†, Vargas†, ICLR 2026 (RNE)

Inference · steering results

Steering composes, and works on discrete spaces

Composition + reward: arbitrary routes

Masked diffusion on ImageNet: reward tilting

random samples

tilted toward the prompt

“yellow ambulance”“orange balloon, red spots”“blue jeep”

He, Hernández-Lobato, Du†, Vargas†, ICLR 2026 (RNE)

Training · from data

With data: denoising score matching

Noise a data point with the backward kernel:

$X_t\sim\Pb(\cdot\mid X_N)=\mathcal N\big(\alpha_tX_N,\ s_t^2I\big),\qquad X_N\sim\pi\ \ (\text{data})$

Bayes' rule over the clean point gives the noised score:

$\nabla\log p_t(x_t)=\mathbb E_{\pi(x_N\mid x_t)}\big[\nabla_{x_t}\log\Pb(x_t\mid X_N)\big]=\mathbb E\Big[-\dfrac{X_t-\alpha_tX_N}{s_t^2}\,\Big|\,X_t=x_t\Big]$

Regress the score network onto that label, the direction back to the clean point:

$\mathcal L_{\rm DSM}(\theta)=\displaystyle\int_0^1\!\lambda_t\;\mathbb E\,\Big\|\,s_\theta(X_t,t)+\frac{X_t-\alpha_tX_N}{s_t^2}\Big\|^2\dd t$

Why this loss? It maximizes a bound on the likelihood (the ELBO):

$\mathrm{KL}(\pi\,\|\,p_\theta)\le\mathrm{KL}\big(\Pb\,\|\,\Pf_\theta\big)=\tfrac12\displaystyle\int_0^1\!\sigma_t^2\;\mathbb E\,\big\|s_\theta(X_t,t)-\nabla\log p_t(X_t)\big\|^2\dd t$

$-\mathbb E_\pi\log p_\theta=\mathrm{KL}(\pi\,\|\,p_\theta)+{\rm const}$; the path KL equals $\mathcal L_{\rm DSM}$ with $\lambda_t=\sigma_t^2/2$, up to a $\theta$-free constant

Its minimizer is the score $\nabla\log p_t$: plain regression, so it scales.

But every label needs a clean sample $X_N\sim\pi$: exactly what we lack for a Boltzmann distribution.

Hyvärinen 2005 · Vincent 2011 (denoising score identity) · Ho et al. 2020 · Song et al. 2021 · Song, Durkan et al. 2021 (likelihood) · left: exact score of the noised mixture ($\alpha_t=1$)

Training · from energy

Without data: target score matching

Same noising kernel; integrate by parts to move the gradient onto the target:

$\nabla\log p_t(x_t)=\dfrac{1}{\alpha_t}\,\mathbb E_{\pi(x_N\mid x_t)}\big[\nabla\log\tilde\pi(X_N)\big]=\dfrac{1}{\alpha_t}\,\mathbb E\big[\beta\mathbf F(X_N)\,\big|\,X_t=x_t\big]$

the noised score is the posterior-averaged force; $Z$ never appears

$\mathcal L_{\rm TSM}(\theta)=\displaystyle\int_0^1\!\lambda_t\;\mathbb E\,\Big\|\,s_\theta(X_t,t)-\frac{1}{\alpha_t}\,\beta\mathbf F(X_N)\Big\|^2\dd t$

① data → forces: physics enters through $\mathbf F=-\nabla U$, one force call per sample.

② but the expectation is still over $X_N\sim\pi$, the very thing we want to learn.

The identity changes the representation, from data samples to the target score (forces), but on its own it does not give us a learning algorithm.

De Bortoli et al. 2024 (target score matching) · Blessing et al., ICML 2026

Training · the fixed-point loop

From energy: the model proposes, the physics corrects

$s^{(k+1)}=\arg\min_s\displaystyle\int_0^1\!\lambda_t\;\mathbb E_{X_{0:N}\sim\Pf_{s^{(k)}}}\Big[\textcolor{#E76F51}{\bar w}\;\Big\|\,s(X'_t,t)-\frac{1}{\alpha_t}\beta\mathbf F(X_N)\Big\|^2\Big]\dd t,\qquad X'_t=\alpha_tX_N+s_t\varepsilon$

$\bar w$: the self-normalized path weight, the same weight we used for inference

Where the cost goes

① Generate: expensive, a full rollout of the model

② Label: expensive, one force call per sample, the only oracle calls

③ Interpolate: cheap, fresh noise, no oracle

④ Regress: cheap, plain regression

Why a buffer: pay ① and ② once per sample, then reuse every buffered sample for many cheap rounds of ③ and ④.

The buffer turns two expensive steps into a one-off cost per sample; training is cheap regression.

Blessing et al., ICML 2026 (Bridge Matching Sampler) · Havens et al. 2025 (Adjoint Sampling)

Training · why path weights

Forces find the modes; path weights fix their weights

$\pi$ is a fixed point; the path weight $\textcolor{#E76F51}{\bar w}$ fixes the mode weights.

schematic three-well landscape; left: Metropolis with small steps, run live · Blessing et al., ICML 2026

Applications · amorphous materials

Amorphous materials: properties come from ensembles

Why they matter: disorder gives tunable mechanical, transport and optical properties.
Why they are hard: below $T_g$, equilibration needs rare collective rearrangements.
No data: no database of equilibrated amorphous ensembles.

left: real 2D Kob–Andersen snapshot at T = 0.2 (glass) vs an ordered binary crystal

Applications · ATLAS

ATLAS: one model, many thermodynamic conditions

Path weights: $Z=\mathbb E_{\Pf}[w]$ gives the free energy $F$.

Glass at $T=0.2$: free-energy error $2\times10^{-3}$.

Cheng*, Blessing*, …, Domingo-Enrich†, Du†, ATLAS (2026), Fig. 2b,e · reference: converged PT-MCMC with ≈4.5×10⁸ evaluations

Applications · ATLAS for discovery

Discovery: an LLM agent queries ATLAS

Task: find glass compositions that are both stiff and ductile.

Result: the best trade-offs are found within 480 ATLAS queries.

Design rules: Co, Ni and Ti appear in every best alloy.

Cheng*, Blessing*, …, Domingo-Enrich†, Du†, ATLAS (2026), Fig. 4d · left: schematic of the asynchronous agent loop · candidates are model-guided hypotheses for higher-fidelity validation

Applications · chemical disorder

Chemical disorder: beyond the canonical ensemble

$p(x)\propto e^{-\beta U}$

$p(x,v)\propto e^{-\beta(U+Pv)}$

$p(a,x,v)\propto e^{-\beta(U+Pv-\mu\cdot N)}$

$p(a,x)\propto e^{-\beta(U-\mu\cdot N)}$

Each ensemble needs its own sampler today; JANUS covers all four with one model.

Blessing*, Cheng*, …, Domingo-Enrich†, Du†, JANUS (2026) · Schebek et al. 2024 (conditional Boltzmann generators for phase diagrams)

Applications · JANUS

JANUS: one model, many thermodynamic conditions

Each ensemble adds variables: volume (NPT), composition (semi-grand), particle number (grand canonical).
Each variable gets its own kernel; the path weight is still their product.
One network, conditioned on $(T,\Delta\mu)$: every temperature and composition from one model.

Soft label for the species

the analogue of the force

$\rho_i(b\mid\mathbf a_{-i},\mathbf x)=\dfrac{e^{-\beta U(\mathbf a_{i\to b},\,\mathbf x)}}{\sum_{b'}e^{-\beta U(\mathbf a_{i\to b'},\,\mathbf x)}}$

$\mathbb E\big[\,\rho_i(b\mid\mathbf a_{-i},\mathbf x)\,\big|\,\mathbf a_t,\mathbf x_t\big]=q^\star_i(b\mid\mathbf a_t,\mathbf x_t)$

in expectation, the target posterior

Austin et al. 2021 · Sahoo et al. 2024 · Shi et al. 2024 · Ou et al. 2025 (masked diffusion) · Blessing*, Cheng*, …, Domingo-Enrich†, Du†, JANUS (2026) · animation: toy lattice

Applications · JANUS for discovery

Discovery: defects for quantum technologies

Sampler: one grand-canonical model per host crystal; dopants and vacancies appear and disappear.

Screen: keep the defect motifs that appear more often than chance, then check them with DFT.

Found: known centres re-emerge (Si E centre, Ge–V); new: S–Ti (Si), B–O–O (diamond).

Blessing*, Cheng*, …, Domingo-Enrich†, Du†, JANUS (2026), Fig. 4 · left: schematic of the screening loop · candidates require excited-state validation (ZPL, Huang–Rhys, spin relaxation)

Any state space · Markov kernels

Any Markov process: the path weight only needs kernels

$w(X_{0:N})=\dfrac{\tilde\pi(X_N)\,\prod_t\Pb(X_t\mid X_{t+1})}{p_0(X_0)\,\prod_t\Pf(X_{t+1}\mid X_t)}$

Nothing here assumes $\mathbb R^d$: any pair of forward and backward kernels whose probabilities we can evaluate.

Continuous

one Gaussian step of an SDE

$\Pf(x_{t+1}\mid x_t)=\mathcal N\big(x_t+\mu_t\,\Delta t,\ \sigma_t^2\Delta t\big)$

molecules, materials

Discrete

reveal one masked site or pixel

$\Pf(a_{t+1}\mid a_t)$: a categorical

alloys, images

Autoregressive

append one token

$\Pf(x_{0:i}\mid x_{0:i-1})=p(x_i\mid x_{\lt i})$

language models

Whole paths

the state is a trajectory

$\dfrac{\dd\Pf^{\,*}}{\dd\Pf}\propto e^{\int r_t\,\dd t}$

rare events, transitions

Kernels compose: discrete species × continuous positions, as in JANUS

Conditionally independent: both parts of step $t{+}1$ are predicted from the same state at $t$.

$\Pf\big((a,x)_{t+1}\mid(a,x)_t\big)=\Pf(a_{t+1}\mid a_t,x_t)\;\Pf(x_{t+1}\mid a_t,x_t)\quad\Rightarrow\quad\log w=\log w_a+\log w_x$

Same weight, same estimators, same training: only the kernel changes.

Any state space · language models

Autoregressive language models are non-equilibrium processes too

Forward kernel generate

$\Pf_i(x_{0:i}\mid x_{0:i-1})=p(x_i\mid x_{\lt i})$

Backward kernel trivial

$\Pb_{i-1}=\dfrac{p_{i-1}\,\Pf_i}{p_i}=1$ (delete the last token)

So the path weight exact

$=\displaystyle\prod_i p(x_i\mid x_{\lt i})=p(x)$

Every tool here applies: steer to $p^\beta$, $p\,e^{r}$ or $p\,p'$ with SMC and replica exchange, and quality scales with inference compute, no reward model needed.

He, Guo, Hernández-Lobato, Du 2026 (Recipes for Steering and Scaling LLMs via Sampling)

Any state space · whole paths

Tilting whole paths: rare events and transitions

chignolin folding: a transition path sampled by stochastic optimal control

Tilt the entire path, not just its end:

$\dfrac{\dd\Pf^{\,*}}{\dd\Pf}\propto\exp\!\Big(\int r_t\,\dd t\Big)$

The tilted process adds a drift $\nabla\log h_t$: stochastic optimal control. For transitions, $h$ is the committor.

Endpoint tilts give thermodynamics; path tilts give kinetics.

Holdijk*, Du* et al., NeurIPS 2023 · Du*, Plainer*, Brekelmans* et al., NeurIPS 2024 · Du, He, Zhang, Vanden-Eijnden, Domingo-Enrich 2026

Outlook

What is still hard

variance

Variance

the variance of the importance weights can be high, especially when the model is imperfect

interpolant

The path matters

annealing or transport: which path from $p_0$ to $\pi$?

exploration

Exploration

without data, every basin must be discovered

Annealing

$q_\lambda\propto p_0^{\,1-\lambda}\,\pi^{\lambda}$

energy known at every step, but the mass of a mode can appear late: it must hop the barrier.

Transport

$q_\lambda=\mathrm{law}\big((1-\lambda)X_0+\lambda X_N\big)$

each sample keeps its mode, so mass moves smoothly, but $q_\lambda$ is unknown: it must be learned.

The path sets the work: we want known densities and smooth transport. Choosing it is open.

Takeaways

Push probabilistic ML out of equilibrium

Diffusion lifts the model to path spacetime reversal gives a cheap, unbiased path weight

One path weight: marginal likelihood, debiasing, steeringfree energy, exact averages, inference-time scaling

The path weight is the physical worknon-equilibrium thermodynamics: Jarzynski, Crooks, the second law

Training from data or from energylearn the time reversal: denoising, target score matching, the fixed-point loop

ApplicationsATLAS for glasses, JANUS for alloys and defects

Any state spaceany Markov process: language models; beyond endpoint marginals, whole paths

Equilibrium is exact only at infinite time. Out of equilibrium, every path carries its own correction.

Thank you!

Questions are very welcome.

Blog: Pushing Probabilistic Machine Learning Out of Equilibrium
RNE · FEAT · Free energy on any state space · CREPE · ATLAS · JANUS · LLM steering · Rare events via SOC
yuanqidu@microsoft.com · yuanqidu.github.io · @YuanqiD

Backup · training: related work

How do we train it? How the field got here

2018 · Boltzmann generators

A normalizing flow: exact likelihood $p_\theta(x)$, so reweighting to $\pi$ is exact.

But invertible layers limit expressiveness, and it is hard to scale to large, flexible systems.

2022–24 · a divergence on whole paths

$\mathrm{KL}\big(\Pf_\theta\,\big\|\,\Pb\big)=\mathbb E_{\textcolor{#E76F51}{X_{0:N}\sim\Pf_\theta}}\Big[\log\dfrac{\mathrm d\Pf_\theta}{\mathrm d\Pb}(X_{0:N})\Big]$

$\text{or}\ \ \mathrm{Var}_{X_{0:N}\sim\mathbb Q}\Big[\log\dfrac{\mathrm d\Pb}{\mathrm d\Pf_\theta}(X_{0:N})\Big]$

The reverse KL is evaluated on paths sampled from the model itself; the log ratio is the path weight (from earlier). But: backprop through the solver, no reuse of energy calls, mode collapse.

2024 · PDDS → iDEM: the right idea

Regress the score on the model's own samples; iDEM adds the replay buffer. Each step closer to the final loop.

But the debiasing sits inside the label: SMC weights (PDDS), an importance-sampled score $\nabla\log\frac1K\sum_k e^{-\beta U(x+s_t\varepsilon_k)}$ (iDEM). High variance → fails at scale.

2025–26 · fixed-point matching

Same loop, but the label is one force $\beta\mathbf F(X_N)$ at the model's own endpoint: no weights inside the regression, simulation-free, replay buffer.

Regression shapes each mode; it does not fix the mode weights.

ATLAS · JANUS = fixed-point matching + the 2022–24 path weight loss, back as a regularizer for the mode weights.

Noé et al. 2018 (Boltzmann generators) · Zhang & Chen 2022 (PIS) · Vargas et al. 2023 (DDS) · Berner et al. 2024 (DIS) · Richter & Berner 2024 (log-variance) · Vargas et al. 2024 (CMCD) · Phillips et al. 2024 (PDDS) · Akhound-Sadegh et al. 2024 (iDEM) · He*, Du* et al. 2025 (No trick, no treat) · Havens et al. 2025 (Adjoint Sampling) · Liu et al. 2025 (ASBS) · Blessing et al. 2026 (Bridge Matching Sampler) · Havens et al. 2026 (Flow Sampling)