Average the score or score the average

Let be a prior or reference measure on . The Kullback-Leibler term below is the regularizer. Bayesian and generalized Bayesian approaches can be written as

Standard Bayes is the special case with and the prior, while generalized Bayes allows other losses and scaling of (Bissiri, Holmes, and Walker 2016). A predictive posterior instead solves

where is the empirical distribution of the data and is a discrepancy between the predictive and the empirical distribution . The first term scores the mixture predictive directly, while the second term keeps close to a reference measure . The discrepancy may be chosen as the negative log likelihood, i.e.

but indeed it does not have to be. It can be any scoring rule and a few examples are given in the next section. The regularization term is important to avoid overfitting, especially when the model is misspecified and the mixture predictive can be very flexible. For a pointwise scoring rule , with smaller scores being better, the discrepancy can be written as

Note that the quantity can be chosen to depend on , and it can be tuned to optimize predictive performance, typically using a validation set or cross-validation.

Log score and mixture likelihood

With the log score and no KL penalty, Equation 3 becomes

This is the likelihood for a mixture model with unknown mixing distribution (Lindsay 1995). So in this case the predictive posterior is doing classical mixture estimation. If is restricted to a finite list of predictors, the same objective gives stacking or convex aggregation (Yao et al. 2018). Better prediction may simply mean that mixtures predict better than single components. The posterior-like part is the entropy penalty on .

Conditional prediction

For supervised data, write for covariates and for responses so that the observation is . Replace by a conditional distribution on and the mixture predictive is

If one uses a proper scoring rule as the discrepancy, the predictive posterior solves

where as usual the first term scores the mixture predictive and the second term keeps close to a reference measure . The log score is a special case, but other proper scoring rules can be used as well.

A latent-subpopulation example

Consider the Gaussian regression model where is known and is the regression coefficient. Now suppose the population has unobserved groups:

Bayes under the single-slope model learns one compromise slope, which can have arbitrarily poor predictive performance. A predictive posterior can put mass near , which gives a mixture predictive that can be much closer to the truth. If the groups are part of the scientific model, estimates population heterogeneity. If the same mixture is only approximating nonlinear regression, heteroscedastic noise, or missing covariates, is an ensemble representation.

A misspecified simulator example

The interpretation changes in mechanistic models. Consider a mechanistic simulator . For a deterministic simulator, this may be a point mass, or the simulator output plus observation noise. The parameter may contain rate constants, forcing terms, initial-state corrections, or boundary conditions. If the simulator omits a source of randomness or a reaction channel, ordinary Bayes can learn one best wrong simulator. The mixture can add predictive spread, as is empirically observed in ensemble weather forecasts. It can widen intervals, add heavier tails, or create several plausible predictive modes. Note that this predictive spread is not a discovery of the missing mechanism. It is just a predictive fix for a misspecified model. In other words, there is nothing magical about the mixture predictive: it is just a mixture of the original model components, and the mixture weights are chosen to optimize predictive performance. A broad may mean physical heterogeneity, missing process noise, numerical ensemble spread, or regularization.

What does Q mean?

The objective identifies , not and the map needs not be injective since different mixing distributions can give the same predictive distribution. Two nearly identical predictive distributions can come from very different mixing distributions. The objective directly supports claims about . On the other hand, claims about parameter uncertainty, latent heterogeneity, or scientific mechanism need more assumptions. They need identifiability of the mixture representation and a data-generating story in which has that meaning.

The regularization scale also matters. While a Bayes-scaled value keeps the posterior analogy closer, a validation-tuned value targets prediction. Large keeps close to . Small chases predictive fit. In the second case, the spread of is partly a tuning effect.

Computation

The optimization is over distributions on parameter space, so computation starts by choosing a representation of . A variational implementation chooses a parametric density on and optimizes . A particle implementation uses

and optimizes the locations of the particles. This leads to a nonconvex optimizaition problem that is often attacked with stochastic gradient descent. For MMD, CRPS, or energy scores, the objective can often be estimated from simulations. They only require simulation from , not evaluation of .

Where this appears

Several recent papers study this score-the-mixture update. Predictive variational inference, prediction-centric uncertainty quantification, and predictively oriented posteriors all fit the score-the-mixture pattern (Lai, Linero, and Yao 2024; Shen et al. 2024; McLatchie et al. 2025). Generalized Bayes fits the average-loss pattern (Bissiri, Holmes, and Walker 2016). MMD posterior bootstrap and MMD regression are nearby, but they remain parameter-projection methods unless they score the mixture predictive itself.

Setup: SOC on path space

Discretize the path into interior nodes: , with boundary conditions and fixed. The target density on the interior nodes is for some reference measure . As in adjoint sampling, the terminal cost absorbs the mismatch: , where is the Radon-Nikodym derivative of the OU marginal at time against the invariant measure . The log-RND term corrects for the mismatch between the OU marginal and the invariant measure, analogous to the path-space reweighting in the Jarzynski notes.

(Sign convention: is a cost, minimized; consistent with adjoint sampling. The adjoint matching note uses a reward .)

The reference process in adjoint sampling is Brownian motion. For paths, the natural reference is an Ornstein-Uhlenbeck (OU) process. Compared to Brownian motion, the OU reference adds a mean-reverting drift . Two reasons to prefer it: (a) the eigenbasis parameterization with Dirichlet eigenfunctions enforces boundary conditions exactly, and (b) the OU invariant measure concentrates on smooth paths, providing a better prior for path sampling than Brownian motion (which generates rough paths). The reference SDE on is where is the negative discrete Laplacian (positive-definite, eigenvalues ): the tridiagonal matrix with on the diagonal and on the off-diagonals, incorporating Dirichlet boundary conditions. Here is a -Wiener process with and for some .

The Laplacian has eigenvectors (discrete sine functions) with eigenvalues . Expanding , each mode evolves independently: where are independent scalar Brownian motions and is the noise amplitude from the covariance operator . Each coefficient follows a scalar OU process. In the eigenbasis, the problem decomposes into independent 1D SDEs, each a standard scalar SOC problem.

The controlled SDE adds a drift: Same structure as adjoint sampling, with base drift instead of . The optimal control is a Doob h-transform of the OU process, conditioning it to produce paths with low action.

The adjoint on path space

The lean adjoint from the adjoint matching notes satisfies with .

Here , so and the lean adjoint ODE becomes: This is a linear ODE with solution

Since is symmetric positive-definite with eigenvalues , the matrix exponential has eigenvalues . It is diagonal in the eigenbasis: just multiply each mode by the scalar , no matrix operations needed. In the eigenbasis: Low-frequency modes (small ) get weight close to 1, nearly identical to the Brownian case. High-frequency modes (large ) are exponentially : their adjoint contribution at early times is negligible.

Why is the adjoint no longer constant? For Brownian motion (), perturbing by shifts by the same amount (independent increments), so for all . For OU, the drift pulls perturbations back toward zero. A perturbation at time decays by before reaching . The terminal cost’s sensitivity to that perturbation is therefore , not . High-frequency modes are strongly damped by the OU reference, so the control has less leverage over them at early times: perturbing them costs energy but has minimal terminal effect.

The FAS loss is the adjoint matching loss with this non-constant adjoint: where the expectation is over trajectories of the controlled SDE Equation 3 (with stop-gradient). The factor appears because the control enters the SDE as ; the Girsanov change of measure produces , not , so the adjoint target picks up . The factor is the OU semigroup, which damps the terminal gradient according to how much the base dynamics attenuate perturbations.

Boundary conditions for free

The Dirichlet eigenfunctions of satisfy . Any linear combination automatically satisfies and .

To enforce and , write where is any fixed path connecting to (e.g., linear interpolation). The OU process evolves the coefficients; the boundary values are fixed by . No penalty terms, no projections. The controlled process respects boundary conditions by construction.

Discretization-invariant sampling. Train the control on a coarse grid (small ). At inference, sample at arbitrary resolution by adding more eigenmodes. The high-frequency modes () receive zero control and simply follow the OU reference, which already assigns them low variance ( decays with ). This is a standard property of Galerkin truncation for linear reference processes; FAS inherits it rather than inventing it.

Connection to existing framework

Set : the OU reference becomes Brownian motion, , , and the FAS loss Equation 6 reduces to the adjoint sampling loss . Standard adjoint sampling is the special case with no spatial structure in the base dynamics.

The reciprocal projection from adjoint sampling carries over directly. Replace the Brownian bridge with an OU bridge: given terminal state , sample from the OU bridge. Since each mode is an independent scalar OU process, the bridge decomposes into independent 1D Gaussian conditionals:

The buffer-based training from adjoint sampling transfers directly: simulate forward to fill a replay buffer, then sample from the OU bridge for cheap inner-loop gradient steps.

Setup

We work with a continuous-time Markov chain (CTMC) on a finite state space . Background on CTMCs, rate matrices, and time-reversal is in the discrete diffusion notes. The CTMC is characterized by its transition rate : the infinitesimal probability per unit time of jumping from state to state . A reference process with rate defines the base path measure . Adding a control means replacing by , producing a controlled path measure .

The path-space KL between and is (see the Girsanov notes for the continuous analogue):

Each summand is the generalized KL divergence between rates at a single pair . It is non-negative, equals zero when , and plays the role of from the continuous Girsanov penalty.

The discrete stochastic optimal control (SOC) problem (cf. the HJB notes) is:

where is a terminal cost. For sampling from , the terminal cost is (when the reference has uniform stationary distribution). By the Doob h-transform (discrete version), the optimal rate is

where is the discrete Doob -function. Writing (consistent with the HJB notes):

The structure here is multiplicative: the optimal rate tilts the reference rate by . In the continuous case, is an additive perturbation of the drift. The DAM paper defines as cost-to-go (); the formulas agree since both give .

Additive noise on cyclic groups

Continuous adjoint matching relies on translation invariance of Gaussian noise: the reference transition kernel depends only on the displacement .

For CTMCs on a finite state space, there is no obvious analogue. The key observation from (liu2025dasbs?; so2026dam?) is to equip the state space with a group structure. Take , the cyclic group of integers modulo in dimensions. Each state is with , and addition is component-wise modulo .

Choose a uniform reference rate:

where is the Hamming distance (differing in exactly one coordinate) and is a noise schedule. Restricting to single-coordinate jumps () keeps the rate matrix sparse and the computation per step . The resulting transition kernel depends only on (subtraction in ):

The argument is simple. The uniform rate is invariant under any translation on : the jump rate from to equals the rate from to . Starting from or and running the same Markov chain, the displacement has the same distribution. This is the discrete analogue of adding mean-zero Gaussian noise.

Discrete adjoint estimator

Matching loss

Dropping the memoryless condition: DASBS

The value function bias from adjoint matching and the SB decomposition from adjoint sampling apply here verbatim: without the memoryless condition, the SOC terminal marginal is not . The fix is the same: decompose the Schrodinger bridge into a controller (Doob h-transform via ) plus a corrector (backward potential ), giving

with modified terminal cost that absorbs the prior coupling.

Summary: continuous vs discrete

	Continuous	Discrete
State space
Base dynamics	SDE, Brownian motion	CTMC, uniform rate
Additive noise	Gaussian:	Cyclic group:
Optimal control	(additive)	(multiplicative)
Adjoint terminal	(gradient)	(exponential difference)
Adjoint ODE	Stochastic ( depends on )	Deterministic (admits closed form)
Matching loss	regression	Generalized KL divergence
Corrector	(score)	(ratio)

The discrete analogue of Nelson’s relation from the BMS notes is detailed balance: .

Can I learn without computing ?

The adjoint notes provide a tool for exactly this situation. For an ODE with running cost and terminal cost , the adjoint variable satisfies The full adjoint ODE for a controlled SDE with drift also contains terms involving : As described in the adjoint matching note, the -dependent terms ( and ) have zero conditional expectation at optimality. Dropping them does not change the fixed point but reduces gradient variance. The resulting lean adjoint satisfies which depends only on the base drift , not on the control . This is a practical win: no need to compute , which is expensive for neural network parameterizations.

Now set and : the base process is pure scaled Brownian motion, and the running cost is just the control penalty (which does not enter the lean adjoint). Both forcing terms vanish. The lean adjoint ODE becomes with terminal condition . So No backward ODE to solve. The regression target at every time is just the terminal gradient .

The intuition is clean. This decomposition holds for the base process (). Under the base process, the two pieces and are independent: Brownian increments do not remember the past. Holding the future noise realization fixed and perturbing by a small displacement , the terminal state shifts by the same amount: . The sensitivity of the terminal cost to the intermediate state is therefore along each base trajectory, confirming Equation 4. This argument uses the base process decomposition , which holds for . Along controlled trajectories, depends nonlinearly on through , and the full adjoint . The lean adjoint is constant because it is defined using the base dynamics (), not the controlled dynamics. Its virtue is that it avoids computing while having the same fixed point as the full adjoint. This is the simplest possible adjoint: the base dynamics offer no resistance to perturbations, so the terminal gradient propagates unchanged backward through time.

The adjoint matching loss with target Equation 4 is where : trajectories are simulated with the current (frozen) parameters, and gradients only flow through the term inside the square. Concretely, for each trajectory the computational graph is: simulate using (no gradient), compute (one energy gradient call), then for each saved along the path, evaluate and backpropagate through only. This regresses onto at each pair along the trajectory. The unique fixed point is (exchanging and , valid under regularity of ): the conditional expectation, under the optimally controlled process, of the terminal gradient. Comparing with , the adjoint target is a stochastic estimate of ; the regression averages many noisy terminal gradients to recover the smooth gradient of the value function.

This is just mean-squared-error regression. No backward ODE, no Hessians, no importance weights. For each trajectory, evaluate once at the endpoint, and you have training signal at every intermediate time. But there is still a computational bottleneck: each gradient step on Equation 5 requires simulating the controlled SDE forward (to produce trajectories), and evaluating the energy function at . Both operations are expensive. The next section removes the first cost almost entirely.

Can I avoid simulating the SDE to generate training data?

The loss Equation 5 only depends on the pair , not the full trajectory: Sampling these pairs still requires simulating the controlled SDE forward, which is the expensive part. The key practical innovation is to replace the joint with the product Keep the terminal marginal from the current control, but fill in the interior with the base Brownian bridge. This decouples the expensive step (generating ) from the cheap step (sampling given ).

The bridge formula. Under the base process, and where is independent of . The pair is jointly Gaussian with , , , , and (since ). Standard Gaussian conditioning (cf. the Brownian bridge in the Doob notes) gives Sampling from Equation 8 is a single Gaussian draw: with . For constant , we have and Equation 8 simplifies to , the classical Brownian bridge.

Substituting Equation 7 into Equation 6 gives the Reciprocal Adjoint Matching (RAM) loss: where is a time-weighting that normalizes the regression target magnitude across time (does not change the optimum). Sampling the training pair now proceeds as: (1) generate by running the controlled SDE forward (one expensive simulation), (2) for each , sample and draw from Equation 8 (one Gaussian draw). One energy evaluation of gives training signal for arbitrarily many pairs.

Why does this substitution work? At optimality, is the Schrodinger bridge, which factorizes as . Marginalizing over intermediate times, . At the fixed point, the replacement Equation 7 is exact. Away from optimality, the replacement defines a reciprocal projection : the Schrodinger bridge with the same terminal marginal as the current control.

After projection, RAM and AM coincide: is a SB with terminal marginal , so its path measure factorizes as , and the reciprocal projection is the identity. The two losses have the same fixed point and the same critical points on the projected space. In other words, we lose nothing by using the cheaper loss, and we gain a monotone improvement guarantee.

The algorithm. The full Adjoint Sampling algorithm (havens2025adjoint?) alternates between two phases. The outer loop (expensive) simulates the controlled SDE forward to produce terminal samples , evaluates , and stores the pairs in a replay buffer . The inner loop (cheap) draws from the buffer, samples and from the bridge Equation 8, and updates by gradient descent on the RAM loss Equation 9. Many gradient steps per energy evaluation, since the inner loop never touches .

This scheme has a clean fixed-point interpretation. Each outer step implicitly performs the reciprocal projection (by freezing the terminal samples), and each inner step minimizes the AM loss on the projected control. If denotes the current control and we fully converge the inner loop using samples from , the update satisfies The fixed point with is exactly the optimal control . In practice the inner loop is not converged fully: the buffer is refreshed when the control has drifted sufficiently from the stored samples.

The computational savings are significant. Previous SOC-based samplers require simulating the SDE and backpropagating through it at every gradient step, plus at least one energy evaluation per step. Adjoint Sampling decouples sampling from optimization: the outer loop produces terminal samples per buffer refresh (each requiring one SDE simulation + one evaluation), while the inner loop does gradient updates using stored buffer entries (each requiring one Gaussian draw from Equation 8 + one forward pass of the neural network). No SDE simulation, no energy call during the inner loop. The ratio of gradient updates to energy evaluations can be in the hundreds. For expensive energy functions like molecular force fields, this is the difference between feasibility and infeasibility.

The replay buffer also provides a form of experience replay: older samples from earlier stages of training remain in the buffer alongside fresh samples, smoothing the optimization. Buffer entries are weighted uniformly; more sophisticated prioritization schemes are possible but have not been explored.

What if the prior is not a Dirac?

The setup is restrictive. For molecular systems, a harmonic oscillator prior that starts particles near equilibrium bond lengths is a far better starting point: the transport cost is lower and convergence is faster. Even for non-molecular problems, starting from a Gaussian and using a moderate noise schedule is more natural than pumping all the stochasticity through the Brownian motion from a single point.

The base process now starts from (not ), and , , and the Brownian bridge all depend on this initial distribution. With a non-trivial prior, and are coupled under the base process, and the initial value function bias from the adjoint matching note reappears. The SOC optimal terminal marginal is , which is not proportional to because the factor cannot be pulled out of the integral. The memoryless condition ( under ) eliminates this bias by making the integral over collapse to a constant, but it also forces (or equivalently, requires a noise schedule that injects enough noise to erase all memory of by time 1, as discussed in the adjoint matching note).

The fix comes from the Schrodinger bridge formulation. Instead of solving a plain SOC (which only imposes via a terminal cost), solve the full SB problem that imposes both marginal constraints and simultaneously. The SB minimizing subject to these two constraints has path measure with time-dependent SB potentials and (in the SB notes, these are defined with endpoint potentials ; here corresponds to and to ). The forward potential is the Doob h-transform: the SB process has drift , steering trajectories toward the target. The backward potential propagates information about the prior forward in time.

The SB is also a controlled diffusion (its drift is the base drift plus ), so it must solve some SOC problem with an appropriate terminal cost. What terminal cost? The SOC optimal joint from the HJB notes is , while the SB joint from Equation 10 is . The SOC here optimizes both the drift and the initial distribution (the HJB formulation in Equation 1 jointly optimizes and ). The SB constraint pins the initial distribution, and enters through this constraint. Matching the terminal tilt requires . The SB marginal constraint at requires , giving . Substituting into identifies the modified terminal cost: Compared to the original , the base marginal is replaced by : a corrector that accounts for the coupling between and . This is the central insight from (liu2025adjoint?): every SB problem can be recast as an SOC problem with a modified terminal cost that absorbs the prior coupling.

Does this corrector actually remove the bias? Marginalize the SB joint over : The second step uses (the definition of combined with Bayes’ rule), and the last step uses the SB marginal constraint . This is a consistency check: IF the SB potentials satisfy the Sinkhorn equations, THEN the terminal marginal is . The existence of potentials satisfying these constraints follows from the Sinkhorn/IPFP theory in the SB notes (this requires for some satisfying the marginal constraints; see the SB notes for conditions). The corrector cancels the initial value function bias exactly. Every SB problem decomposes as: a Doob h-transform (the forward potential , which steers trajectories toward the target) plus a corrector (the backward potential , which absorbs the prior bias). The adjoint method handles the first piece; we now need a way to learn the second.

Corrector matching

The modified AM loss with terminal cost Equation 11 needs , which is unknown. Since , the lean adjoint is still constant (the same argument from Equation 4 applies), and the AM loss becomes where the expectation is over bridge pairs . The corrector score adds one extra term alongside the energy gradient . When , we have (no prior information to propagate), hence for all , and Equation 13 reduces to the plain AM loss Equation 5.

To derive a learning objective for , start from the definition . Converting from a backward conditional to a forward integral via Bayes’ rule (): Define to absorb the prior. Then Take of both sides (exchanging and , valid under standard regularity): Rewrite the numerator using : The ratio is a probability distribution over , and it is exactly the SB posterior .

Substituting back into Equation 15: The left side is , not just . To isolate the corrector score, subtract : At , the corrector enters the AM loss Equation 13 evaluated at . The base marginal score is known in closed form (it is the score of , namely ). So the corrector score at decomposes into a conditional expectation (which we learn by regression) minus a known Gaussian score (which we subtract analytically). Equation 16 is a Tweedie-type formula: the score of the product at is the conditional expectation of the transition score given . The structure is identical to the denoising score matching identity: the score of a mixture equals the expected score of the component, averaged over the posterior on the component index.

Since any conditional expectation minimizes a least-squares loss, Equation 16 at gives: This is the corrector matching (CM) loss: regress onto the base transition score at endpoint pairs. In practice, the expectation is over the current model’s endpoint pairs , not the unknown SB joint . Each CM step finds the best approximation to at the current iterate; alternation converges by the IPFP contraction argument. For the Brownian base, , and the transition score is just , known in closed form. The regression target is a Gaussian score evaluated at the endpoints; no energy evaluation needed. The minimizer of Equation 17 gives ; to recover we subtract the known base score .

Note the parallel with the AM loss. Both are mean-squared-error regressions. AM regresses the control onto the energy gradient at pairs from the bridge. CM regresses the corrector onto the transition score at endpoint pairs from the current process. Both are self-consistent: the expectation is over the model’s own samples. Both have a conditional expectation as their unique fixed point.

The CM loss Equation 17 should not be confused with the bridge matching objectives in data-driven SB methods like (de2021diffusion?), where must be drawn from the target . Here, both AM and CM use only on-policy samples from the current control . This is what makes the approach scalable: no target samples are needed.

Alternating AM + CM = IPFP

The AM loss Equation 13 needs . The CM loss Equation 17 needs samples from . Neither can be solved alone. The fix: alternate.

1. AM step: Solve Equation 13 with to get .
2. CM step: Solve Equation 17 with endpoint pairs from to get , then set (subtracting the known base score).

Initialize with . The first AM stage uses terminal cost (since ) and regresses the control onto : pure energy-guided transport with no corrector. This recovers exactly the Adjoint Sampling algorithm, which is the right starting point when no better guess for the corrector is available. The first CM stage then learns from the endpoint pairs produced by . Subsequent stages progressively refine both the control and the corrector.

The alternation has a variational interpretation. The AM step solves a forward half-bridge: , where is the backward SDE determined by the previous corrector . The CM step solves a backward half-bridge: , updating the corrector to absorb whatever bias remains. Alternating between forward and backward half-bridge projections is exactly IPFP/Sinkhorn on path-space (de2021diffusion?).

Recall from the SB notes that the static Sinkhorn algorithm alternates between projecting the coupling onto the set and the set . Here, we do the same thing on path-space: the forward half-bridge imposes but lets float; the backward half-bridge imposes but lets float. Each projection is itself a KL minimization over a convex constraint set, so each step reduces the KL to the SB solution. Convergence to follows from standard IPFP contraction (de2021diffusion?), provided each inner optimization is solved exactly (which never holds in practice with finite gradient steps; the practical version uses a few optimization steps per stage). In practice, 3-5 outer stages suffice; each stage consists of many AM or CM gradient steps.

Method	Prior	Terminal cost	Corrector
Adjoint Sampling (havens2025adjoint?)	(memoryless)		None ()
ASBS (liu2025adjoint?)	Arbitrary		Learned

ASBS reduces to Adjoint Sampling when : the corrector becomes constant and the CM step is trivially solved. A concrete example of a useful non-Dirac prior: for molecular conformer generation, one can use a harmonic oscillator prior where are equilibrium bond lengths read from the molecular graph. Particles start in a physically reasonable arrangement rather than at the origin, reducing the transport cost substantially.

Every piece was already in the earlier notes. The adjoint method gives the regression target ( is constant along trajectories when ). The Brownian bridge gives cheap training pairs (Gaussian conditioning, one draw per sample). The Schrodinger bridge gives the corrector (endpoint potentials that debias the prior). Adjoint sampling (havens2025adjoint?) assembles the first two; ASBS (liu2025adjoint?) adds the third. The reader who has followed the series could have assembled this themselves.

The overall picture: sampling from an unnormalized Boltzmann distribution reduces to solving a stochastic optimal control problem, which reduces to a sequence of mean-squared-error regressions on cheap Gaussian training pairs. The only expensive operation (energy evaluation at terminal samples) is amortized over many gradient updates. With a Dirac prior, a single matching objective suffices. With a non-trivial prior, two alternating objectives (AM + CM) converge to the Schrodinger bridge via IPFP.

The ideas generalize beyond the setting, though the simplifications are less dramatic. With a non-zero base drift (e.g. a pre-trained diffusion model as in the adjoint matching note), the lean adjoint ODE is no longer trivially constant; it must be solved backward along each trajectory. The reciprocal projection still applies in principle, but the bridge distribution is no longer a simple Gaussian, and must be approximated (e.g. by linearizing the base drift). The case is the sweet spot: everything is closed-form, the adjoint is constant, the bridge is Gaussian, and the only approximation is the neural network parameterization of .

A note on related methods. The RAM loss Equation 9 is closely related to the training objectives in PDDS (phillips2024particle?) and Target Score Matching (de2024target?), where the same squared-error expression appears inside the expectation. The difference is in what the expectation is taken over: PDDS and TSM sample from the target distribution (requiring importance sampling or SMC), while Adjoint Sampling samples from the current model . This makes the loss a moving target, which is harder to analyze but far more practical: no target samples needed, no importance weights, no resampling.

Nelson’s relation

Let denote the path measure of Equation 1 with time marginals . Write the Euler discretization with step :

The forward conditional mean is . Now compute the backward conditional mean using Bayes’ rule, exactly as in the reverse diffusions note. For :

Expanding and completing the square, the conditional mean is

(The second-order correction to affects the conditional variance at but not the conditional mean, which is all we need.)

Completing the square:

Drop multiplicative constants independent of . The exponent in the posterior is

Write . The quadratic piece is , with mean at . The linear piece shifts the Gaussian mean by (the standard “linear tilt of a Gaussian” identity: if , then the mean shifts from to , with and ). So , giving .

From the reverse diffusions note, the time-reversed process satisfies

The reversed drift is . Define by writing this reversed drift as , so that the reversed SDE takes the form . Then , i.e. . Rearranging:

This is Nelson’s relation:

As a sanity check, Equation 3 confirms this: the backward conditional mean identifies the reversed drift as , consistent with the time-reversal formula.

This holds for any Markov diffusion of the form Equation 1 with scalar noise schedule and marginals .

Reciprocal class and Markovian projection

Let denote the reference (uncontrolled) process . Write for the law of the reference process conditioned on . A path measure belongs to the reciprocal class if it has the form , where is an endpoint coupling and is the reference bridge (Brownian bridge for Brownian ); see the Schrodinger bridges note.

A reciprocal measure is generally non-Markovian: the bridge drift depends on . The Markovian projection finds a Markovian drift whose time marginals match those of (see (Brunick and Shreve 2013) for existence/uniqueness). (Existence and uniqueness of the Markovian projection requires the path-dependent drift to satisfy a linear growth condition; see Brunick and Shreve (2013).) This is an projection: if is the path-dependent drift of , then

Why? For any Markovian , expand and use the tower property:

The cross-term vanishes because by definition of . So minimizes the matching loss

Fixed-point iteration

All three methods (adjoint sampling with Dirac prior, adjoint sampling with corrector, BMS) follow the same template. Starting from some control :

1. Simulate the current SDE to generate endpoint pairs .
2. Reciprocal projection: form a coupling from the endpoints, define .
3. Markovianize: update by regressing onto the bridge drift via Equation 6.

If , then and , so is a fixed point. Convergence of this iteration is not guaranteed in general; the BMS paper treats it as empirically effective. What distinguishes the methods is the coupling in step 2, which determines the regression target .

Target score identity

To get a tractable , we need the score . Define the cumulative variance and . The bridge is Gaussian: for ,

The marginal density of at time is

Differentiate :

Since is Gaussian with mean ,

Integration by parts to swap gradients:

The identity relies on the Gaussian bridge having a mean that is affine in with endpoint-independent variance. Shifting by is equivalent to shifting by , which holds because the mean is . Beyond Gaussian references, the TSI still holds conceptually but the gradient swap takes a different form.

Integrate by parts in :

where we moved from onto (boundary terms vanish by decay of at infinity). (The minus from and the minus from integration by parts cancel, giving a positive coefficient .) Write for the conditional distribution of given under . Recognizing the conditional expectation:

The same argument with integration by parts in gives

Since both expressions equal the same score, any convex combination times the first plus times the second is also valid, for any .

This gives the generalized target score identity: for any ,

(The endpoints and are excluded because they would divide by or , which vanish at or .)

(The gradient-integral interchange and vanishing boundary terms require regularity of and decay of the integrand at infinity; these hold for sub-Gaussian couplings.)

General regression target

Now combine everything. For a reciprocal process , the bridge drift decomposes into forward and backward pieces: . The first piece points back toward ; the second points forward toward . The backward Markovian drift extracts the backward piece and averages it over :

Since is Gaussian , this is averaged over . This follows from the bridge drift decomposition above: the piece, when averaged over , gives the forward drift by the Doob h-transform. So is the remaining backward piece, averaged over .

Nelson’s relation Equation 4 was derived for with marginals . Here we apply it to , which has marginals (the Markovian projection preserves marginals). This gives the Markovian forward drift: . Both the score and the backward drift are conditional expectations (from Equation 9 and Equation 10 respectively). To get a regression target for the matching loss Equation 6, we need a non-Markovian drift whose conditional expectation equals . Substituting the integrands (before conditioning) from Equation 9 and Equation 10 into Nelson gives such a :

The first two terms are the TSI integrand (the quantity inside in Equation 9, which conditions on to give ). The third term is the backward drift integrand from Equation 10 (conditioning on gives ). Their combination, before conditioning, is a valid with . Both the TSI integrand (depending on ) and the backward drift integrand (depending on ) are expectations under conditioned on . Since is a marginal of , the two integrands live in the same conditional probability space and can be combined before taking .

Here are all determined by the bridge path: from the coupling, from Equation 7. The Markovianization step fits to via Equation 6.

The tractability of Equation 11 depends entirely on the coupling scores .

Three couplings, three algorithms

Half-bridge / adjoint sampling. Set (Dirac prior, memoryless condition). With and : since , we use only the branch of the TSI (the integration-by-parts is not needed when is deterministic), and Equation 11 reduces to Equation 12. The paper derives this in Prop 4 (Appendix C.2); the key steps are below.

Reduction from Equation 11 to Equation 12:

Start from Equation 11 with . Since , we use only the branch of the TSI (the integration-by-parts is not needed when is deterministic). With :

Since , its score at is . So

Since , its score is , so can be rewritten after Markovianization. Under the half-bridge coupling , the bridge from to gives . The conditional expectation (using ). So . After Markovianization:

since for . The non-Markovian that gives this after conditioning is simply , which depends only on . The -dependent bridge score cancels against the -dependent part of under the Markovian projection, leaving only a function of .

Alternatively: the paper shows directly that setting in the general SHB formula (Lemma C.2) causes all -dependent terms to cancel, leaving .

where (the terminal marginal of the reference). Simple, but requires Dirac prior and large for exploration.

Full Schrodinger bridge / adjoint sampling with corrector. Set , the Schrodinger bridge coupling. The drift becomes

where is the backward Schrodinger potential. This allows arbitrary priors, but is unknown and must be learned alongside , requiring alternating IPF-style updates.

Independent coupling / BMS. Set

Plug into Equation 11. The coupling scores factor trivially: and . The regression target becomes

Every term on the right is known: is the prior score (assumed known, e.g. Gaussian), is the target score (computable from the unnormalized density), and is the Gaussian transition score. No unknown potentials, no alternation.

The independent coupling: why it works

The independent coupling satisfies the boundary constraints by construction: marginalizing over gives , marginalizing over gives . The terminal marginal of is : since (the bridge is pinned at its endpoint),

The Markovian projection preserves time marginals, so : the controlled SDE with drift hits the target at time . Combined with (from the initial condition), is a fixed point.

The Schrodinger bridge coupling minimizes path-space KL (as shown in the SB notes). The independent coupling sacrifices this optimality for a fully tractable regression target.

At each iteration, the coupling is : independently resample and from their marginals under the current SDE. In practice, simulate trajectories, then randomly pair the initial and terminal samples.

Sampling the bridge is cheap: given , draw from Equation 7 and evaluate Equation 15. No full trajectory simulation needed during regression.

Damped iteration

The undamped iteration can overshoot in high dimensions. The damped version uses step size :

Setting , this solves

Deriving Equation 16 from Equation 17:

Apply the bias-variance decomposition (Pythagorean identity from the Markovian projection) to the first term:

where . The first piece is independent of (irreducible noise from the non-Markovian ). Dropping it, Equation 17 reduces to

Pointwise first-order condition: , giving . So . With : and , recovering Equation 16.

The penalizes deviation from the previous iterate. Each step balances fitting new bridge data against staying close to , preventing mode collapse from aggressive updates.

Summary

Method	Coupling	Regression target	Limitation
AS			Dirac prior
AS + corrector			Alternating opt.
BMS		Equation 15	None (single obj.)

All three converge to a fixed point transporting to . The matching loss Equation 6 is a forward KL objective: . This follows from the Girsanov KL decomposition: , so minimizing the matching loss over is equivalent to minimizing . Forward KL is mode-covering (it penalizes placing zero mass where has mass); since the Markovian projection preserves time marginals, mode coverage at the path level implies mode coverage at the terminal marginal level, which drives mode diversity in practice.

Setup

Consider a pre-trained generative model that produces samples by simulating the SDE

on . We restrict to scalar, state-independent throughout, matching the generative model setting. The drift encodes a learned score or velocity field. The path distribution of this process is and the terminal distribution is . Given a reward function , we want to fine-tune the model so that its output distribution becomes the tilted distribution

We assume is smooth and that all expectations below are well-defined. Introduce a control and the controlled SDE

The Girsanov theorem gives (the Girsanov notes derive with expectation under ; the two directions are equal for the control-affine case since has zero expectation under either measure), so the stochastic optimal control problem

maximizes the terminal reward while penalizing deviations from the base process. Equivalently, we maximize , which is the SOC problem from the HJB notes with and . Here we fix and optimize only over , unlike the joint optimization in the HJB notes.

Value function bias

The HJB notes show that the optimal control is where is the value function and

under the base process. Combining Ito’s formula on with the HJB equation and Girsanov, the bivariate marginal of the optimal path measure satisfies

The HJB notes derive (final displayed equation) the pathwise identity via Ito’s formula and the HJB equation. Plugging this into the Girsanov path-level RN derivative and taking expectations (which kills the martingale), the path-dependent terms reduce to endpoint functions, yielding Equation 5.

The factor is the conditional normalizing constant: given , the RN derivative of the optimal path conditioned on is .

Marginalizing the joint over :

When and are correlated under , the factor cannot be pulled out of the integral. The marginal is not proportional to . The SOC solution tilts both the terminal and the initial distribution, and the fine-tuned model does not sample from Equation 2.

As a concrete example, take (standard flow matching). The ODE is deterministic: fully determines for a diffeomorphism . Adding a control in Equation 3 has zero effect since kills both the noise and the control term (the control enters as , so it vanishes with ; this is specific to the control-affine parameterization). The optimal control is trivially since any nonzero adds cost () without affecting the trajectory. The bias is maximal.

Memoryless noise schedule

Look again at Equation 6. If under the base process, then and

The integral over collapses to a constant. The bias disappears.

A base process Equation 1 is memoryless when . For the family of generative SDEs that arise in flow matching and diffusion models, this condition pins down .

The generative SDE Equation 1 is built to have the same time marginals as the reference interpolation with independent of . Conditionally on , this reference has : the conditional variance is , which starts at (pure noise) and decays to (deterministic). The unified drift from (domingoenrich2024adjoint?) is

This drift is constructed so that the SDE has the same time marginals as the reference interpolation; the term handles the deterministic scaling and the score term corrects for the noise. Here and . For the SDE to produce independent and , it needs to inject enough noise to erase memory of by time 1. The base SDE with drift and diffusion has the same marginals as . For to hold, the noise must be large enough that the SDE “forgets” its initial condition by time 1. The memoryless noise level is the unique choice that makes the conditional law under the SDE match the reference Gaussian structure for all . The paper (Proposition in Section 12.1 of (domingoenrich2024adjoint?)) proves this by analyzing the time-reversed SDE and showing that the coefficient of in the explicit solution vanishes if and only if , with being the minimal such choice. This requires

Near , so : the process mixes aggressively, erasing memory of . Near , so : the process stabilizes around .

Why is the right condition (sketch):

The time-reversed SDE for is a linear SDE whose explicit solution expresses as a function of . The coefficient of in this expression involves . For , this coefficient must vanish, requiring . The threshold is (the paper proves this in Appendix 12.1). Below this threshold, correlation persists; above it, the process mixes too aggressively and distorts the time marginals.

(In practice, the time interval is truncated to to avoid the singularity at .)

For DDIM with , , the memoryless schedule gives , which is exactly the DDPM noise schedule (verified in the details block above for flow matching; for the DDPM verification, see Table 1 of the paper). The DDPM noise schedule is the memoryless schedule for DDIM.

The paper proves (Theorem 2) that within the family of SDEs sharing marginals with the reference flow, the memoryless schedule is not just sufficient but necessary for the fine-tuned velocity to be convertible to arbitrary noise schedules. The memoryless schedule is the unique choice that preserves the velocity-score relationship: a model trained under the memoryless schedule yields consistent velocity and score fields, and , enabling conversion to arbitrary noise schedules after fine-tuning.

The memoryless schedule is only needed during fine-tuning. After learning the control , the fine-tuned model (base drift plus control) can be converted back to a velocity field and sampled with any , including .

Adjoint Matching

With the memoryless schedule, we need to solve Equation 4. The adjoint method gives the gradient of the cost along a trajectory. The adjoint state is the sensitivity of the future cost to perturbations at time :

Applying the adjoint ODE from the adjoint notes to the controlled drift with running cost and terminal cost :

The conditional expectation is the gradient of the cost functional; heuristically, following the SDE extension described in the adjoint notes (the interchange of gradient and conditional expectation requires regularity of and integrability of ). The paper defines as cost-to-go (minimization); we use the HJB notes’ convention (reward-to-go, maximization). The two are negatives of each other: . Since and at optimality, we get

This is a fixed-point condition: is the unique control satisfying Equation 12.

Turn this into a regression. For a current control (parameters frozen), simulate , solve the adjoint ODE Equation 11 backward, and regress:

Here denotes stop-gradient: the trajectory and adjoint use frozen parameters; gradients only flow through the term. Expanding: . At (stop-gradient), the term produces the continuous adjoint gradient from Equation 11. So this basic version produces the same parameter updates as the standard adjoint method.

The lean adjoint

The full adjoint Equation 11 contains terms involving . At optimality, these terms have conditional expectation zero.

Removing the -dependent terms:

Both sides are column vectors in . The minimizer of the regression Equation 13 is the conditional expectation:

Right-multiply both sides by the Jacobian :

Since and are deterministic functions of , they can be pulled in and out of freely. Rearranging and applying the tower property:

The terms inside the expectation appear in the full adjoint ODE Equation 11 (they are the -dependent pieces: and ). They vanish in conditional expectation at optimality.

This gives the lean adjoint :

Compare with Equation 11: every -dependent term is gone. The lean adjoint depends only on the base drift , not on the control. No need to compute , which is expensive for neural networks. The lean adjoint also has smaller magnitude since the removed terms cancel in expectation but add gradient variance.

The lean adjoint is the adjoint ODE from the adjoint notes applied to the base dynamics with terminal cost .

The Adjoint Matching loss replaces the full adjoint with the lean one:

The unique critical point of is . The argument above shows that at optimality. This does NOT mean the lean adjoint equals pointwise; it means that replacing by in the regression loss produces the same gradient at . To see why: the gradient of with respect to is . At , this equals . Adding back the removed terms: . Since the removed terms have zero conditional expectation given and is a function of , the tower property gives . So the gradients of the lean and basic losses agree at , and by the uniqueness of the critical point of (Appendix 13.3 of the paper), is also the unique critical point of .

Unlike the basic version, the lean adjoint produces a different gradient than the standard adjoint method away from optimality (the removed terms have expectation zero only at , not elsewhere). The paper reports more stable convergence in practice.

After fine-tuning

The memoryless schedule is only used during training. After convergence, the fine-tuned drift can be converted to a velocity field and sampled with any noise schedule, including for deterministic generation.

Potts model

Consider the Potts model with colors on the complete graph with vertices. A configuration is with . Define the energy At inverse temperature , the Boltzmann distribution is . On the complete graph, the only relevant macroscopic variable is the empirical proportions vector

so that . A short computation rewrites the energy in terms of : up to an additive constant irrelevant for Gibbs weights. The number of configurations with a given is , where is the Shannon entropy. Putting energy and entropy together, typical samples from concentrate near minimizers of the mean-field free-energy functional:

constrained to the probability simplex. Everything that follows is geometry: how the minima of evolve as the inverse temperature varies.

Local minima: Two features make the analysis almost trivial. First, permuting colors leaves unchanged. Second, a stationary point under the simplex constraint satisfies a Lagrange-multiplier condition for . Hence each coordinate must solve the same scalar equation. From this, one finds that all local minima are necessarily of the following two types:

• Disordered point (uniform): always a stationary point.

• Ordered points (one dominant color, the rest equal): for some dominant proportion . There are such points by choosing which color is dominant.

The remaining work is algebra: determine, as a function of , when ordered stationary points exist, and which of the stationary points are local minima versus saddles. The details are routine and add little insight.

The phase diagram in

Assume . Then the model exhibits a first-order transition, and metastability appears on both sides. There are two key inverse temperatures:

• Spinodal threshold : the ordered stationary points appear as local minima.
• Coexistence threshold : the disordered minimum and the ordered minimum have equal free energy.

For the normalization above, the coexistence point is

The geometry of splits into three regimes.

Zonal spherical harmonics

Understanding the eigenfunctions of the spherical Laplacian is a central task. These eigenfunctions form the building blocks of harmonic analysis on spheres.

A particularly important situation arises when the function of interest depends only on the angular separation from a fixed axis.

Such zonal functions retain full rotational symmetry around that axis and reduce the spherical Laplacian to a one-dimensional operator. Its eigenfunctions turn out to be polynomials on the interval with a very specific weight. These polynomials are the Gegenbauer polynomials, which generalize several classical families of orthogonal polynomials, including the Legendre and Chebyshev polynomials.

Reminder on differential operators

Many problems require the ability to compute gradients, divergences, and Laplacians in arbitrary coordinate systems. Before specializing to spherical coordinates in higher dimensions, it is useful to recall the geometric meaning of these operators and the minimal formalism needed to manipulate them. Suppose that a point in space is described by a collection of coordinates If we make a small change , the physical displacement has a length denoted by . In an orthogonal coordinate system, each coordinate direction comes with a scale factor such that a change corresponds to a “physical” displacement of length along the unit vector . These functions encode how the coordinate system stretches or compresses distances along each axis. Once the scale factors are known, the rest of the differential operators follow directly.

The gradient: For a scalar function , the change produced by varying only is approximately . The physical distance traveled in this move is . The rate of increase of per unit distance in the direction is therefore . This is the th component of the gradient. Summing over all directions gives

The divergence describes the net outflow of a vector field from an infinitesimal volume around a point. Consider a vector field . Place a tiny coordinate-aligned box at a point. Its physical edge lengths are , so its infinitesimal volume is: The flux of through the pair of faces orthogonal to is multiplied by the physical area of those faces: . The divergence is the total outward flux divided by the infinitesimal volume:

This expression is simply “flux in direction ” minus “flux in direction ” across the opposite face, normalized by the physical volume.

The Laplacian is the divergence of the gradient. It describes how a scalar function curves around a point since it measures the net outflow of the gradient field. It also encodes how compares to its local averages over small spheres. In orthogonal coordinates, inserting the expression for the gradient into the divergence formula yields:

This is the expression we will apply to spherical coordinates in the next section.

Masked Discrete Diffusion

We consider a finite state space where is a special masked state and correspond to token values in a vocabulary of size . On the time interval , we define a continuous-time Markov chain with initial distribution and time-dependent infinitesimal rate matrix so that for any ,

If the total jump rate out of state is , then

Bayes’ rule implies that the time-reversal of this Markov chain is itself Markov, with infinitesimal rate matrix satisfying where is the marginal distribution of at time . We have:

We are interested in modeling a Markov chain that progressively masks the initial value into the masked state as time goes from to . Transitions are only allowed from any token to the masked state , and once in the process remains there. Thus, outside the diagonal, the only nonzero entries of are . As it will be useful later, we denote by the jump time to and we assume almost surely, so that with probability one, and that has a continuous distribution. In words: the process starts at some token value and at a random time jumps to the masked state , where it remains until time .

Extension to Sequences

We are interested in modeling sequences comprised of discrete tokens, eg: binary images, genomic sequences, chemical compounds, protein sequences, etc… Each token takes value in . We denote by the data distribution over such sequences. For this purpose, we consider independent copies of the above Markov chain, one per coordinate: each with rate matrix as defined previously. At time , the process reaches the fully masked sequence with probability one. Denote by the jump time of coordinate . Since the jump times are almost surely distinct, the infinitesimal rate matrix of the joint process is nonzero only when a single coordinate changes. If differ by a single coordinate , we have

As before, the time-reversal has infinitesimal rate matrix satisfying

where is the marginal distribution of . Since and differ at coordinate only, for to be non-zero, necessarily and . Let be the set of unmasked coordinates in . To observe configuration at time , the masked coordinates must have and the unmasked ones :

Similarly, and since differs from only at coordinate :

This shows that the time-reversal rate matrix becomes

with time dependent scalar . To simulate the reverse process that progressively unmasks a fully masked sequence, one only needs to model the conditional distribution of the data distribution . This is precisely the prediction task of masked language models such as BERT, which estimate token probabilities conditioned on visible context.

Training

To train the denoising model, Equation 1 shows that it is natural to parametrize the conditional distribution

for all sets with . Once done, one can define the rate matrix of the time-reversal process as:

If one denotes by the law of the forward noising process started from , and by the law of the time-reversal process started from the fully masked sequence at time and with learned denoising model , one can train the model by minimizing

Consider a trajectory of the forward noising process. The jump at time of the -th coordinate is denoted by . For simplifying the notations, we denote the reverse jump by . The log-likelihood ratio between the two processes is easily shown to be:

where and are the total jump rates of the forward and reverse processes respectively. Since in fact does not depend on , minimizing the KL divergence is equivalent to minimizing:

where is the set of unmasked coordinates at time . It is more convenient to rewrite this quantity as an integral over time so that one can sample a time uniformly in during training. With the Dirac delta function , we can rewrite this expectation as:

where so that . For training, it suffices to sample uniformly in , then choose , then sample the noised configuration according to the forward process, and finally obtain an unbiased estimate of the loss. The term is large for small , counter-balancing the fact that the reconstruction task is much easier when only a few tokens are masked. For standard denoising diffusion model, there is a similar “signal-to-noise” weighting term that balances the easy and hard denoising tasks.

Conclusion

Discrete diffusion models with one-way masking are mathematically almost identical to masked language models. Hence their similar behavior and performance on text generation tasks are not coincidental. The ideas summarized in these notes were developed in a very interesting stream of papers, including (Ou et al. 2024), (Sahoo et al. 2024), (Shi et al. 2024) and a number of more recent works. One of the potential drawbacks of such masked discrete diffusion models is that the support of the noising distribution is typically strictly smaller, and indeed often much smaller, than the whole state space. This means that when the denoising model is not perfect and wanders outside the support of the noising distribution, one can quickly end up in regions never seen during training. This can leads to poor sample quality and unstable behavior. Other discrete diffusion models such as the ones reaching the uniform distribution over all tokens at time are not as badly affected by this issue, although they do suffer from other important computational and modeling challenges. Exciting research directions remain to be explored in this area!

Masked Discrete Diffusion

We consider a finite state space where is a special masked state and correspond to token values in a vocabulary of size . On the time interval , we define a continuous-time Markov chain with initial distribution and time-dependent infinitesimal rate matrix so that for any ,

If the total jump rate out of state is , then

Bayes’ rule implies that the time-reversal of this Markov chain is itself Markov, with infinitesimal rate matrix satisfying where is the marginal distribution of at time . We have:

We are interested in modeling a Markov chain that progressively masks the initial value into the masked state as time goes from to . Transitions are only allowed from any token to the masked state , and once in the process remains there. Thus, outside the diagonal, the only nonzero entries of are . As it will be useful later, we denote by the jump time to and we assume almost surely, so that with probability one, and that has a continuous distribution. In words: the process starts at some token value and at a random time jumps to the masked state , where it remains until time .

Extension to Sequences

We are interested in modeling sequences comprised of discrete tokens, eg: binary images, genomic sequences, chemical compounds, protein sequences, etc… Each token takes value in . We denote by the data distribution over such sequences. For this purpose, we consider independent copies of the above Markov chain, one per coordinate: each with rate matrix as defined previously. At time , the process reaches the fully masked sequence with probability one. Denote by the jump time of coordinate . Since the jump times are almost surely distinct, the infinitesimal rate matrix of the joint process is nonzero only when a single coordinate changes. If differ by a single coordinate , we have

As before, the time-reversal has infinitesimal rate matrix satisfying

where is the marginal distribution of . Since and differ at coordinate only, for to be non-zero, necessarily and . Let be the set of unmasked coordinates in . To observe configuration at time , the masked coordinates must have and the unmasked ones :

Similarly, and since differs from only at coordinate :

This shows that the time-reversal rate matrix becomes

with time dependent scalar . To simulate the reverse process that progressively unmasks a fully masked sequence, one only needs to model the conditional distribution of the data distribution . This is precisely the prediction task of masked language models such as BERT, which estimate token probabilities conditioned on visible context.

Training

To train the denoising model, Equation 1 shows that it is natural to parametrize the conditional distribution

for all sets with . Once done, one can define the rate matrix of the time-reversal process as:

If one denotes by the law of the forward noising process started from , and by the law of the time-reversal process started from the fully masked sequence at time and with learned denoising model , one can train the model by minimizing

Consider a trajectory of the forward noising process. The jump at time of the -th coordinate is denoted by . For simplifying the notations, we denote the reverse jump by . The log-likelihood ratio between the two processes is easily shown to be:

where and are the total jump rates of the forward and reverse processes respectively. Since in fact does not depend on , minimizing the KL divergence is equivalent to minimizing:

where is the set of unmasked coordinates at time . It is more convenient to rewrite this quantity as an integral over time so that one can sample a time uniformly in during training. With the Dirac delta function , we can rewrite this expectation as:

where so that . For training, it suffices to sample uniformly in , then choose , then sample the noised configuration according to the forward process, and finally obtain an unbiased estimate of the loss. The term is large for small , counter-balancing the fact that the reconstruction task is much easier when only a few tokens are masked. For standard denoising diffusion model, there is a similar “signal-to-noise” weighting term that balances the easy and hard denoising tasks.

Conclusion

Discrete diffusion models with one-way masking are mathematically almost identical to masked language models. Hence their similar behavior and performance on text generation tasks are not coincidental. The ideas summarized in these notes were developed in a very interesting stream of papers, including (ou2024your?), (sahoo2024simple?), (shi2024simplified?) and a number of more recent works. One of the potential drawbacks of such masked discrete diffusion models is that the support of the noising distribution is typically strictly smaller, and indeed often much smaller, than the whole state space. This means that when the denoising model is not perfect and wanders outside the support of the noising distribution, one can quickly end up in regions never seen during training. This can leads to poor sample quality and unstable behavior. Other discrete diffusion models such as the ones reaching the uniform distribution over all tokens at time are not as badly affected by this issue, although they do suffer from other important computational and modeling challenges. Exciting research directions remain to be explored in this area!

Static Shrödinger Bridge Problem

A standard way to address these issues is to add an entropic regularization term to the objective function. The resulting problem is known as the Schrödinger bridge problem and can be formulated as follows. Consider a reference joint distribution over and find the coupling that minimizes the Kullback-Leibler divergence to while matching the marginals and :

Remark (invariance under separable rescaling). Only the “interaction structure” of matters: if one replaces by with positive , then the optimal coupling can still be written in the same form below, with the factors absorbed into the potentials. Equivalently, the solution depends on only through its kernel up to left/right diagonal scaling.

A common choice for the reference measure is a distribution of the form (optionally times separable factors). This choice encourages the coupling to put more mass on pairs that are close according to the distance , and the resulting optimization problem can be rewritten (up to an additive constant) as:

where is the entropy of the coupling . Note that since the marginals of are fixed, it is also equivalent (up to constants) to replacing the entropy term by the Kullback-Leibler divergence to the independent coupling ,

Writing the Lagrange dual formulation of the problem Equation 2 provides fast algorithms to solve it such as the iterative proportional fitting procedure (IPFP), also known as the Sinkhorn-Knopp algorithm in the optimal transport literature (Cuturi 2013). Importantly, it also almost immediately shows that the optimal coupling has a very simple form,

These potentials must satisfy the marginal constraints, equivalently the Schrödinger/Sinkhorn scaling equations:

Dynamic Schrödinger Bridge Problem

Now, suppose that the reference distribution is given as the two-time marginal of a Markov process with transition kernels and path measure on trajectories .

The dynamic Schrödinger bridge problem consists in finding a new path measure on trajectories such that the starting and ending marginals are and , while minimizing the Kullback-Leibler divergence to the reference path distribution :

When the reference is Markov, the chain rule for KL (disintegration with respect to ) shows this is equivalent to the static Schrödinger bridge problem for the two-time marginal of : one first solves for the optimal endpoint coupling , and then fills in intermediate times using the reference bridges :

1. Sample the endpoints .
   
2. Sample the intermediate points according to the conditional law of the reference process given the endpoints, i.e. .

Continuous-Time Stochastic Processes:

A typical setting is when the reference process is given as the solution of a stochastic differential equation of the form

started from some initial distribution (one may take it to be without loss of generality, since any mismatch can be absorbed into the endpoint tilt below). The solution to the dynamic Schrödinger bridge problem is given by the twisted path distribution:

for endpoint potentials and satisfying the marginal constraints at and .

The marginal density at intermediate time factorizes as:

where is the time- marginal density of the reference process, and where the time-dependent Schrödinger potentials are defined by:

(with and , up to normalization). Naturally, the dynamics of the Schrödinger bridge process can be described as a new stochastic differential equation obtained by applying a Doob h-transform to the reference SDE,

where, as explained in these previous notes, the function is the harmonic extension of the terminal potential defined above.

Discrete-Time Markov Chains:

It is often useful to state the Schrödinger bridge dynamics in a discrete setting. Let the reference process be a Markov chain with one-step kernels , i.e. Let be the terminal potential at time and define the backward (harmonic) potentials defined by and recursively as:

Then the Schrödinger bridge is Markov and its forward transition kernels are given by the discrete Doob -transform: This is the discrete-time counterpart of the continuous-time drift correction .

In general, Schrödinger bridge problems are difficult problems since the endpoint potentials and need to be found such that the marginal constraints are satisfied. These recent years have seen the development of many numerical methods to solve this problem approximately, especially in the machine learning community, eg: (Shi et al. 2023); (De Bortoli et al. 2021); (Vargas et al. 2021); (Chen, Liu, and Theodorou 2021).

Static Shrödinger Bridge Problem

A standard way to address these issues is to add an entropic regularization term to the objective function. The resulting problem is known as the Schrödinger bridge problem and can be formulated as follows. Consider a reference joint distribution over and find the coupling that minimizes the Kullback-Leibler divergence to while matching the marginals and :

Remark (invariance under separable rescaling). Only the “interaction structure” of matters: if one replaces by with positive , then the optimal coupling can still be written in the same form below, with the factors absorbed into the potentials. Equivalently, the solution depends on only through its kernel up to left/right diagonal scaling.

A common choice for the reference measure is a distribution of the form (optionally times separable factors). This choice encourages the coupling to put more mass on pairs that are close according to the distance , and the resulting optimization problem can be rewritten (up to an additive constant) as:

where is the entropy of the coupling . Note that since the marginals of are fixed, it is also equivalent (up to constants) to replacing the entropy term by the Kullback-Leibler divergence to the independent coupling ,

Writing the Lagrange dual formulation of the problem Equation 2 provides fast algorithms to solve it such as the iterative proportional fitting procedure (IPFP), also known as the Sinkhorn-Knopp algorithm in the optimal transport literature (cuturi2013sinkhorn?). Importantly, it also almost immediately shows that the optimal coupling has a very simple form,

These potentials must satisfy the marginal constraints, equivalently the Schrödinger/Sinkhorn scaling equations:

Dynamic Schrödinger Bridge Problem

Now, suppose that the reference distribution is given as the two-time marginal of a Markov process with transition kernels and path measure on trajectories .

The dynamic Schrödinger bridge problem consists in finding a new path measure on trajectories such that the starting and ending marginals are and , while minimizing the Kullback-Leibler divergence to the reference path distribution :

When the reference is Markov, the chain rule for KL (disintegration with respect to ) shows this is equivalent to the static Schrödinger bridge problem for the two-time marginal of : one first solves for the optimal endpoint coupling , and then fills in intermediate times using the reference bridges :

1. Sample the endpoints .
   
2. Sample the intermediate points according to the conditional law of the reference process given the endpoints, i.e. .

Continuous-Time Stochastic Processes:

A typical setting is when the reference process is given as the solution of a stochastic differential equation of the form

started from some initial distribution (one may take it to be without loss of generality, since any mismatch can be absorbed into the endpoint tilt below). The solution to the dynamic Schrödinger bridge problem is given by the twisted path distribution:

for endpoint potentials and satisfying the marginal constraints at and .

The marginal density at intermediate time factorizes as:

where is the time- marginal density of the reference process, and where the time-dependent Schrödinger potentials are defined by:

(with and , up to normalization). Naturally, the dynamics of the Schrödinger bridge process can be described as a new stochastic differential equation obtained by applying a Doob h-transform to the reference SDE,

where, as explained in these previous notes, the function is the harmonic extension of the terminal potential defined above.

Discrete-Time Markov Chains:

It is often useful to state the Schrödinger bridge dynamics in a discrete setting. Let the reference process be a Markov chain with one-step kernels , i.e. Let be the terminal potential at time and define the backward (harmonic) potentials defined by and recursively as:

Then the Schrödinger bridge is Markov and its forward transition kernels are given by the discrete Doob -transform: This is the discrete-time counterpart of the continuous-time drift correction .

In general, Schrödinger bridge problems are difficult problems since the endpoint potentials and need to be found such that the marginal constraints are satisfied. These recent years have seen the development of many numerical methods to solve this problem approximately, especially in the machine learning community, eg: (shi2023diffusion?); (de2021diffusion?); (vargas2021solving?); (chen2021likelihood?).

Some natural examples

One can then asks oneself what are some natural entire functions that vanish at some predetermined set of zeros. For example, what functions vanish on all the integers? Such a function cannot be of order less than one since otherwise it could be writtem as , but this product does not converge. Any entire function of order that has a simple zero at each integer is of the form

Checking the Taylor expansion of sine at zero, one finds that the celebrated formula:

and a similar interesting example by taking the derivatives of . Now, what about a function that vanishes on all the negative integers? Again, such a function needs to be of order at least one. Hadamard’s factorization theorem then tells us that such a function is, up to a multiplicative constant, of the form

Naturally, since such a function vanishes on all the negative integers, one of the first things one would like to try is to look at and relate it to itself. Since vanishes on , one knows that it can be expressed as so that is almost the same as . One can then do some algebra to choose the constant so that . One finds that the correct choice is the Euler-Mascheroni constant,

This gives the final expression for as

Furthermore, since as , it follows that , from which the identity gives that . In other words, is an analytic continuation on the whole complex plane of the inverse of the Gamma function,

This generalizes the usual definition valid for to the whole complex plane.

The connection to the sine function is also almost immediate. Indeed, the function vanishes on all the integers and the infinite product immediately shows that . But since , one obtains that , i.e.

This is the celebrated Euler reflection formula, from which it also follows that .

Zeta function

Let’s conclude these notes by deriving the analytic continuation of the Riemann zeta function , originnaly defined for , to the whole complex plane. I have always found this proof very elegant. A common trick that is used in many places is to do a change of variable in the definition of the Gamma function to obtains that

It is valid for . This allows one to obtain express the zeta function as:

From the previous discussion, one knows that is an entire function of order , and the integral converges for except for ; in other words, this already gives a meromorphic continuation of to the domain with a simple pole at . Since things get nasty near , the standard approach consist in splitting the integral into two parts. The part defines an entire function and one only needs to take care of the integral . This can be done by expressing , where are the Bernoulli numbers and integrating each term, but the way Bernhard Riemann did it is way more fun. The idea is to note use the boring but instead introduce the way more interesting function

where , defined for . Note that the function decreases exponentially rapidly to as . In other words, instead of Equation 1 one can just as easily write

This slight change of parametrization allows to write the zeta function as

One has not gained much doing this since there is indeed still an issue at where diverges. However, the Jacobi theta function enjoys some interesting symmetries. Crucially, the Poisson summation formula applied to the Gaussian function gives that satisfies modular inversion symmetry:

This means that splitting in Equation 2, using the change of variable to map to and finally use the modular inversion symmetry of the theta function leads after standard algebra:

First, one can note that since decreases exponentially quickly to as , the integral above defines and entire function. This means that the expression above defines a meromorphic continuation of to the whole complex plane with a simple pole at . There is no pole at since the simple zero of takes care of it and gives the value . This also shows that the function inherits from a simple zero at all the negative even integers, i.e.  for . And there are indeed a few other zeros, as the plot below shows… and they seem to be located on the critical line …

What is remarkable is that the term inside the curly brackets of Equation 3 is symmetric in and , i.e. symmetric with respect to the vertical line in the complex plane. This means that the function satisfies the functional equation

Fisher-Rao metric

Suppose we want to define a distance on the space of probability densities on . A natural but naive approach is to use an -type distance:

where are densities with respect to the Lebesgue measure. However, this definition has several shortcomings. For instance, if we change the base measure to for some positive density , and define the distance as

we obtain a different value. Perhaps more troubling, the distance is not invariant under reparametrizations. Let be a diffeomorphism of , and set . Then the transformed densities become , where is the Jacobian determinant. In general,

so the distance depends on the choice of coordinates. That is, measuring in Cartesian or polar coordinates yields different results—an undesirable feature. Ideally, we seek a distance that is invariant under reparametrizations and changes of base measure, such as the total variation distance,

One potential drawback of the total variation distance is that it is not differentiable, which can make it difficult to use in optimization problems. An alternative is to consider -divergences, defined as

where is a convex function with . These divergences are differentiable and invariant under reparametrizations and changes of base measure, although they are not symmetric and thus not true distances. Locally, however, all -divergences are equivalent, as a second-order Taylor expansion shows:

This means that all these divergences describe the same local geometry, defined by the Fisher-Rao information metric. Furthermore, it is relatively straightforward to derive the global geometry induced by the Fisher information metric. Consider the mapping , which maps a density to an element of the unit sphere of . Since

for , and we have just seen that any -divergence is locally equivalent to the Fisher-Rao metric, it follows that the geometry induced by the Fisher-Rao information information metric is the same as the geometry induced by the -norm on the unit sphere of . This implies that the geodesic distance between two densities and is given (up to an irrelevant constant) by the -geodesic distance between the points and . In other words, the geodesic distance, i.e. the Fisher-Rao distance, is given by

The geodesic path is a great circle, , where for . This shows, for example, that the Fisher-Rao geodesic between two Gaussian densities is composed of densities that are Gaussian mixtures; i.e., the geodesic is not composed of Gaussian densities in general. In other words, probability mass is not transported along the geodesic but reshaped, unlike the Wasserstein metric, which describes transport of probability mass. Note in passing that the Hellinger distance, defined as

is just just a slightly rescaled version of the Fisher-Rao distance since they are related by a deterministic function, i.e. . In this sense, the Hellinger distance is equivalent to the Fisher-Rao distance, and both describe a “correct” way to measure distances between probability densities for many applications.

Gradient flow

What do gradient flows look like in this Fisher-Rao geometry? For example, for a given distribution , the gradient flow of under the Wasserstein metric is given by the transport equation:

which describes the Langevin dynamics of the process , as informally described in these notes. So what does the gradient flow of look like in the Fisher-Rao geometry?

To answer this, one can consider the square-root mapping , express everything in terms of , compute the -gradient on the unit sphere (which is straightforward), and finally map back to the density using the inverse mapping . One readily finds that the gradient flow is described by:

This is quite intuitive: the flow tries to increase in regions where and decrease in regions where . Discretizing this flow can naturally be done using sampling-based methods. If is a system of weighted particles approximating , following the Fisher-Rao gradient flow corresponds to updating the weights as

for a small time-step. Indeed, it is because , where , can be discretized by updating the weights as . This is very much related to the resampling step in sequential Monte Carlo methods, and the recent article (Crucinio and Pathiraja 2025) make these connections explicit.

Table of Contents

• Linear Systems
• Adjoint Method
• PDE Inverse Problems
• Controlled Diffusions

The adjoint method is, at its core, the same idea as backpropagation or reverse-mode automatic differentiation. In practice, though, it’s often helpful to understand how it works under the hood. Basic implementations of backprop can be sub-optimal or impractical (eg. memory intensive), especially in settings like PDE-constrained optimization or stochastic optimal control.

Table of Contents

• Linear Systems
• Adjoint Method
• PDE Inverse Problems
• Controlled Diffusions

The adjoint method is, at its core, the same idea as backpropagation or reverse-mode automatic differentiation. In practice, though, it’s often helpful to understand how it works under the hood. Basic implementations of backprop can be sub-optimal or impractical (eg. memory intensive), especially in settings like PDE-constrained optimization or stochastic optimal control.

Nystrom approximation

Before describing the simple and most important case of additive Gaussian noise, let’s give a brief reminder on the Nystrom approximation. The distribution of is Gaussian with mean and covariance . This means that is distributed as the unconditional distribution . In particular:

where is the so-called Nystrom approximation of the covariance matrix based on the inducing variable . This shows that the Nystrom approximation simply consists in ignoring the conditional variance term , and is thus an underestimate of the covariance matrix . Furthermore, if is very informative of , then is small and the Nystrom approximation is accurate.

Observation with additive Gaussian noise

The case of additive Gaussian noise is particularly simple. Assume that

where the noise terms are independent. Since , algebra gives that

Using that and the matrix inversion lemma it quickly follows that optimal variational distribution is with

Indeed, these are approximations of the exact condition moments,

where the Nystrom approximation is used instead. One then finds that . With the optimal variational distribution , Equation 4 gives that the free energy is:

Furthermore, note that exact likelihood of the observations is so that the free energy can be expressed as

for pseudo-likelihood . This shows that is given by:

The term is just the sum of the conditional variances and can be thought of as a regularization term,

As the number of inducing variables increases, the pseudo-likelihood becomes more accurate , the conditional variances shrink to zero, and the KL divergence approaches zero.

The animation at the start of this note illustrates the effect of optimizing the location of the inducing points with a very simple gradient descent. A few experiments show that it is worth being careful with the initial choice of inducing points. Inducing points chosen very far from the data essentially remain fixed during the optimization (ie. the gradient is very small). Initializing with k-means++ clustering of the data points seems to be a robust strategy and give an almost optimal choice of inducing points.

The problems and notations

Recall that we are trying to estimate the connective constant of self-avoiding walks (SAW) in the 2D lattice . If denotes the number of SAWs of length , we have the following asymptotic behavior:

for some unknown constants , , and . The main objective of the assignment is to estimate , which can also be expressed as the limit of as . As of today, the best known estimate is , which required several tens of thousand hours of CPU time to compute. To estimate , one must approximate the number of SAWs of length starting at the origin for large values of if one hopes to get a good estimate.

Consider a sequence of distinct vertices in with and for all , i.e., a walk of length . For notational convenience, let us introduce the function that returns one if is a correct walk of length , and zero otherwise. In particular, this function returns zero if two consecutive vertices are the same, or if the walk does not start at zero. Similarly, introduce the function that returns one if is a SAW of length . One can define two important probability mass functions:

They describe the uniform distributions on all the walks of length and all the SAWs of length , respectively.

Importance sampling

One can approximate with naive Monte Carlo by estimating the proportion of walks that are SAWs,

This is an absolute disaster since the proportion of SAWs among all walks is extremely small. One can do significantly better using importance sampling. For this, consider a proposal distribution that starts at the origin and continues by choosing uniformly among the four neighbors of the last vertex that have not been visited yet. If there are no unvisited neighbors, the walk continues by standing still until length is reached: the resulting path is not even a valid walk, so as well as . The probability mass function of the proposal distribution is easy to compute, so estimating with importance sampling is straightforward. This is usually called the Rosenbluth method (Rosenbluth and Rosenbluth 1955). [Note to students: make it much clearer in your report that the Rosenbluth method is just importance sampling. Do note that even the “rejected” walks have to be taken into account!]

As one can see, the quality quickly deteriorates as increases. This is because the number of accepted walks is very small, and, among them, the importance weights are highly unequal.
[Note to students: you should explain this much more clearly, and possibly explore this more quantitatively. The reason it is failing is not only that the number of accepted walks is small]

Recursive formulation

We have just seen that importance sampling will not be able to estimate for large values of . This makes accurate estimates of difficult to obtain this way.

To make progress, one can exploit the recursive structure of the problem. Let us define the concatenation of two walks. Given a first walk and a second walk , one can define a new walk of length by starting at the origin, following the increments of the first walk, then the increments of the second. The concatenation of two SAWs is not always a SAW. However, it is not hard to prove the following. Define as the probability that, when sampling SAWs and independently and uniformly at random, their concatenation is still a SAW. Then:

[Note to students: it is OK for you to use this fact. It’s even better if you can prove it, but not absolutely necessary.]

Assuming one can generate SAWs of length uniformly at random ( a problem that will be discussed later), we can estimate in several ways:

1. For small values of , the number of SAWs is known exactly (e.g., , ). Suppose one can generate SAWs of length . One can then estimate empirically. Since , it follows that, for fixed and , Using the fact that , one can then estimate from the estimate of .

2. Alternatively, one can estimate for large recursively. For example, starting from , estimate to compute , then use to compute , and so on. Using this method and about hours of CPU time (see below for details) with SAWs of lengths , I obtained .

Generating SAWs

The previous discussion shows that, once we know how to generate uniform SAWs, we can estimate relatively easily. One of the most common methods is the pivot algorithm: see here for a nice visualization. The principle is simple: given a SAW, randomly select a pivot site and apply a symmetry operation (like rotation or reflection) to one part of the walk. If the resulting walk remains self-avoiding, accept it; otherwise, reject it. Repeating this process generates diverse, approximately uniform SAWs.
[Note to students: explain this much more clearly if you decide to use it]

In short, the pivot algorithm updates a SAW by applying a symmetry operation to a subpath. Given a SAW , one can obtain another SAW by applying to it the pivot algorithm a (large) number of times. To obtain a nearly independent SAW of length starting from , one typically need to apply about pivot steps. While it can be slow for large , it is far more efficient than naive importance sampling.
[Note to students: efficiently implementing the pivot algorithm is non-trivial, but LLM assistants can help a lot, and are actually quite useful for code optimization]

Sequential Monte Carlo

To estimate for large , one can use Sequential Monte Carlo (SMC). The idea is to grow a population of SAWs in parallel and estimate by recursively estimating the ratios . Suppose you have SAWs of length . Try to extend each SAW by choosing a neighbor of the last vertex that has not been visited yet. This is a form of importance sampling, giving new walks of length with associated weights (some of them being non-valid walks!). Then, resample times from this weighted set to get new SAWs of length (with possible duplicates). Apply the pivot algorithm to eliminate these duplicates and generate more diverse SAWs.
[Note to students: if you decide to use SMC, explain it much more clearly. It’s not entirely straightforward to understand or implement, but it is one of the most powerful and versatile Monte Carlo methods to this day. A good investment of your time if you decide to understand SMC]

Improving the estimation of

Suppose you have estimates of for various . Since

you can fit a linear regression to estimate , , and . I tried this approach using a naive and non-optimized SMC implementation with and , running for 10 hours on a free (and bad) online CPU, and obtained .
[Note to students: can you do much better?]

Running long simulations

The best known estimate of required several tens of thousands of CPU hours. While writing these notes, I was able to run simulations easily and for free using deepNote: it was my first time using it, and it was very user friendly. This allowed me to run simulations for 8 hours on a (free but slow) CPU without issue. Launch simulations in the evening and let them run overnight. [Note to students: for the more motivated ones, you can try writing GPU-friendly code to run simulations, possibly on Google Colab]

Some References:

• The original papers by Jarzynski (Jarzynski 1997) and Crooks (Crooks 1999).
• The book (Stoltz, Rousset, et al. 2010) is excellent!
• The two papers that prompted these notes: (Vargas et al. 2024) and (Albergo and Vanden-Eijnden 2024).