Adjoint Matching in Stochastic Control

• Adjoint Matching is a method that trains parameterized control fields by regressing on backward adjoint information from terminal objectives.
• It reformulates stochastic optimal control as a regression problem to fine-tune generative models, reinforcement learning models, and combinatorial processes.
• Variants like EAM, RAM, and discrete adjoint matching offer practical speedups and performance gains in complex simulation, optimization, and data assimilation tasks.

Adjoint Matching (AM) denotes a family of adjoint-based optimization methods in which a parameterized control, drift, or velocity field is trained by matching it to a backward-in-time adjoint, costate, or related target induced by a terminal objective. In the recent arXiv literature, the term is most strongly associated with stochastic optimal control formulations for flow and diffusion generative models, where AM converts reward fine-tuning into a regression problem with the same optimum as the underlying control problem (Domingo-Enrich et al., 2024). Closely related formulations now appear in reinforcement learning, discrete continuous-time Markov chains, combinatorial optimization, molecular generative modeling, data assimilation, and PDE-constrained registration. A separate optimal-transport paper explicitly notes that its “AM” stands for Action Matching, not Adjoint Matching, so acronym usage is not uniform across subfields (Kornilov et al., 31 Oct 2025).

1. Foundational formulation

The canonical AM formulation treats reward fine-tuning of an iterative generative model as a stochastic optimal control problem on a controlled SDE. In the unified continuous-time setup, the base process is

$dX_t = b(X_t,t)\,dt + \sigma(t)\,dB_t,$

and the controlled process is

$$dX_t^u = \big(b(X_t,t)+\sigma(t)u(X_t,t)\big)\,dt + \sigma(t)\,dB_t.$$

With terminal cost $g(X_1)$ and optional running cost $f(X_t,t)$, the control objective is

$$\min_u \;\mathbb{E}\!\left[\int_0^1 \left(\frac12\|u(X_t,t)\|^2 + f(X_t,t)\right)dt + g(X_1)\right].$$

In reward-tilting language, the target terminal law is $p^\star(x)\propto p_{\mathrm{base}}(x)\exp(r(x))$ when $g=-r$ (Domingo-Enrich et al., 2024).

A central theoretical point is the memoryless noise schedule. The base process is memoryless, meaning $X_0 \perp X_1$ under $p_{\mathrm{base}}$, if and only if

$\sigma(t)^2 = 2\eta_t + \chi(t),$

with

$$dX_t^u = \big(b(X_t,t)+\sigma(t)u(X_t,t)\big)\,dt + \sigma(t)\,dB_t.$$0

The simplest canonical choice is $$dX_t^u = \big(b(X_t,t)+\sigma(t)u(X_t,t)\big)\,dt + \sigma(t)\,dB_t.$$1. The original AM paper states that this schedule is necessary to avoid a value-function bias in the terminal marginal; under arbitrary $$dX_t^u = \big(b(X_t,t)+\sigma(t)u(X_t,t)\big)\,dt + \sigma(t)\,dB_t.$$2, the path distribution acquires an unwanted factor involving $$dX_t^u = \big(b(X_t,t)+\sigma(t)u(X_t,t)\big)\,dt + \sigma(t)\,dB_t.$$3, so fine-tuning no longer targets the intended reward-tilted law (Domingo-Enrich et al., 2024).

AM is built around an adjoint regression objective. The paper first defines a “full adjoint,” then introduces a lean adjoint that removes terms whose expectation vanishes at the optimum. The lean adjoint satisfies

$$dX_t^u = \big(b(X_t,t)+\sigma(t)u(X_t,t)\big)\,dt + \sigma(t)\,dB_t.$$4

and the corresponding Adjoint Matching loss is

$$dX_t^u = \big(b(X_t,t)+\sigma(t)u(X_t,t)\big)\,dt + \sigma(t)\,dB_t.$$5

The paper establishes that the AM objective has the optimal control $$dX_t^u = \big(b(X_t,t)+\sigma(t)u(X_t,t)\big)\,dt + \sigma(t)\,dB_t.$$6 as its unique critical point in expectation, and that the memoryless SOC formulation recovers the reward-tilted terminal law while permitting arbitrary sampling schedules at inference time (Domingo-Enrich et al., 2024).

This core construction already indicates the characteristic AM pattern that recurs in later work: a forward process supplies trajectories, a backward adjoint transports terminal information, and a local regression objective replaces unstable end-to-end differentiation. A plausible implication is that AM’s usefulness depends less on any single generative architecture than on whether the relevant dynamics admit a tractable backward sensitivity equation.

2. Variants for generative modeling

Subsequent work splits the original AM construction into several algorithmic families that preserve the adjoint-matching idea while altering the base dynamics, the estimator, or the computational pathway.

Variant	Distinguishing mechanism	Representative outcome
Efficient Adjoint Matching (EAM)	Linear memoryless base drift with corrected terminal cost	Converges up to 4x faster than AM (Shin et al., 12 May 2026)
Reinforce Adjoint Matching (RAM)	Reward-corrected consistency target with analytic noising	Reaches Flow-GRPO’s peak reward in up to 50x fewer training steps (Bergmeister et al., 11 May 2026)
Adjoint Sampling	Replay-buffer Reciprocal Adjoint Matching for energy-only sampling	About 0.002 energy evaluations per gradient update (Havens et al., 16 Apr 2025)
ASBM forward stage	AM inside non-memoryless Schrödinger bridge coupling construction	CIFAR-10 FID 3.16 at 100 NFE (Shin et al., 17 Feb 2026)
Deterministic/truncated AM	Terminal-window adjoint solve with general control penalties	345.5 s/iteration to 32.2 s/iteration on FLUX.2-Klein-4B (Guo et al., 7 May 2026)

“Efficient Adjoint Matching” reformulates the SOC problem with a linear base drift and a modified terminal cost. The redesigned base drift is

$$dX_t^u = \big(b(X_t,t)+\sigma(t)u(X_t,t)\big)\,dt + \sigma(t)\,dB_t.$$7

and the terminal correction becomes

$$dX_t^u = \big(b(X_t,t)+\sigma(t)u(X_t,t)\big)\,dt + \sigma(t)\,dB_t.$$8

Under this construction, the adjoint has the closed form

$$dX_t^u = \big(b(X_t,t)+\sigma(t)u(X_t,t)\big)\,dt + \sigma(t)\,dB_t.$$9

so EAM removes both full stochastic trajectory simulation and backward adjoint integration. The reported consequence is substantially lower training cost while matching or surpassing AM on PickScore, ImageReward, HPSv2.1, CLIPScore, and Aesthetics (Shin et al., 12 May 2026).

“Reinforce Adjoint Matching” takes a different route. It starts from the result that, under KL-regularized reward maximization, the optimal controlled process tilts the clean endpoint distribution while leaving the conditional noising law unchanged. Combined with a REINFORCE identity and the AM fixed-point condition, this yields a velocity-space regression target

$g(X_1)$0

with stop-gradient on the target. RAM therefore avoids SDE rollouts, backward adjoint sweeps, and reward gradients, and the paper reports that on Stable Diffusion 3.5M it reaches Flow-GRPO’s peak reward in up to $g(X_1)$1 fewer training steps (Bergmeister et al., 11 May 2026).

“Adjoint Sampling” specializes AM to the energy-based setting $g(X_1)$2. In the zero-drift, zero-stage-cost case, the lean adjoint becomes constant in time,

$g(X_1)$3

and the AM loss reduces to a regression on joint pairs $g(X_1)$4. The main algorithmic step is Reciprocal Adjoint Matching (RAM): sample $g(X_1)$5 on-policy, then sample $g(X_1)$6 analytically from the base conditional $g(X_1)$7 and reuse stored terminal gradients in many cheap inner-loop updates. On synthetic benchmarks, the paper reports roughly 500 gradient updates per energy evaluation and model sample (Havens et al., 16 Apr 2025).

“Adjoint Schrödinger Bridge Matching” places AM in a two-stage non-memoryless Schrödinger bridge pipeline. The forward stage learns a controlled coupling from data to an energy-defined prior using an AM target involving $g(X_1)$8 plus a terminal corrector $g(X_1)$9, and the backward stage then learns the generative dynamics under the induced optimal coupling. The paper argues that the non-memoryless regime yields straighter trajectories and reports, among other results, CIFAR-10 FID 3.16 at 100 NFE and FFHQ latent FID 6.47 at 250 steps (Shin et al., 17 Feb 2026).

A separate deterministic-control line reformulates AM directly over flow ODEs. In “Improved techniques for fine-tuning flow models via adjoint matching,” the control $f(X_t,t)$0 is chosen by minimizing a Hamiltonian

$f(X_t,t)$1

which gives the closed-form target

$f(X_t,t)$2

For polynomial regularization $f(X_t,t)$3, this becomes

$f(X_t,t)$4

The paper then truncates the adjoint computation to a terminal window, arguing that reward-relevant control intensity concentrates near the endpoint; on FLUX.2-Klein-4B, full adjoint AM takes 345.5 s/iteration, while one-step truncated AM takes 32.2 s/iteration (Guo et al., 7 May 2026).

3. Reinforcement learning formulations

In reinforcement learning, AM is used to optimize expressive diffusion or flow policies without backpropagating through long denoising or integration chains. “Q-learning with Adjoint Matching” formulates continuous-action policy optimization as a memoryless SOC problem driven by a learned critic $f(X_t,t)$5. The main policy parameterization is a state-conditioned flow field $f(X_t,t)$6 inside the SDE

$f(X_t,t)$7

The lean adjoint is propagated backward under the behavior flow $f(X_t,t)$8 rather than the current policy, with boundary

$f(X_t,t)$9

and the AM loss is

$$\min_u \;\mathbb{E}\!\left[\int_0^1 \left(\frac12\|u(X_t,t)\|^2 + f(X_t,t)\right)dt + g(X_1)\right].$$0

The paper states that this preserves the optimum of the SOC solution

$$\min_u \;\mathbb{E}\!\left[\int_0^1 \left(\frac12\|u(X_t,t)\|^2 + f(X_t,t)\right)dt + g(X_1)\right].$$1

while avoiding unstable backpropagation through the multi-step sampler (Li et al., 20 Jan 2026).

The empirical motivation is strong. On 50 OGBench offline RL tasks, QAM reports an aggregated score of 44, compared with 8 for FAWAC, 35 for BAM, 36 for FQL, and 42 for QSM; variants QAM-FQL and QAM-EDIT reach 45 and 46 respectively (Li et al., 20 Jan 2026). These numbers do not by themselves identify a universal ordering across all settings, but they do show that the adjoint-matching formulation can be competitive against value-only and distilled one-step alternatives.

“Scalable Maximum Entropy Reinforcement Learning for Diffusion Policies via Adjoint Matching” introduces AMDP, which casts MaxEnt RL for diffusion policies as a path-space reverse-KL problem and then trains with Reciprocal Adjoint Matching. Under a memoryless reference with $$\min_u \;\mathbb{E}\!\left[\int_0^1 \left(\frac12\|u(X_t,t)\|^2 + f(X_t,t)\right)dt + g(X_1)\right].$$2, the AM loss becomes

$$\min_u \;\mathbb{E}\!\left[\int_0^1 \left(\frac12\|u(X_t,t)\|^2 + f(X_t,t)\right)dt + g(X_1)\right].$$3

Because $$\min_u \;\mathbb{E}\!\left[\int_0^1 \left(\frac12\|u(X_t,t)\|^2 + f(X_t,t)\right)dt + g(X_1)\right].$$4 is sampled from a tractable conditional Gaussian bridge, the update is “simulation-free” in the sense used by the paper: there is no backpropagation through diffusion time. Runtime measurements show actor update overhead within about $$\min_u \;\mathbb{E}\!\left[\int_0^1 \left(\frac12\|u(X_t,t)\|^2 + f(X_t,t)\right)dt + g(X_1)\right].$$5 of fast Gaussian REPPO and independence from the number of diffusion steps, whereas reverse-KL chain backprop is reported as $$\min_u \;\mathbb{E}\!\left[\int_0^1 \left(\frac12\|u(X_t,t)\|^2 + f(X_t,t)\right)dt + g(X_1)\right].$$6 to $$\min_u \;\mathbb{E}\!\left[\int_0^1 \left(\frac12\|u(X_t,t)\|^2 + f(X_t,t)\right)dt + g(X_1)\right].$$7 slower in update time (Thilges et al., 21 Jun 2026).

Two later RL extensions modify the AM prior rather than the adjoint machinery itself. “Trust Region Q Adjoint Matching” introduces a trust-region parameter $$\min_u \;\mathbb{E}\!\left[\int_0^1 \left(\frac12\|u(X_t,t)\|^2 + f(X_t,t)\right)dt + g(X_1)\right].$$8 directly into the SOC dynamics and proves the path-space KL identity

$$\min_u \;\mathbb{E}\!\left[\int_0^1 \left(\frac12\|u(X_t,t)\|^2 + f(X_t,t)\right)dt + g(X_1)\right].$$9

Projected dual descent on $p^\star(x)\propto p_{\mathrm{base}}(x)\exp(r(x))$0 then controls the exact deviation from the pretrained flow policy. The paper reports an overall offline success rate of 68% on 50 OGBench tasks, compared with 46% for the strongest baseline and 35% for QAM (Dong et al., 26 May 2026).

“Entropy-Regularized Adjoint Matching for Offline RL” addresses what it calls popularity bias and support binding. It adds a Mirror Descent entropy term and a Mixture Behavior Prior

$p^\star(x)\propto p_{\mathrm{base}}(x)\exp(r(x))$1

and shows that the stationary policy takes the tempered form

$p^\star(x)\propto p_{\mathrm{base}}(x)\exp(r(x))$2

This replaces the standard behavior-weighted Boltzmann target with a flattened prior exponent $p^\star(x)\propto p_{\mathrm{base}}(x)\exp(r(x))$3 and an expanded geometric support (Ghanem et al., 7 May 2026).

4. Discrete-state and combinatorial extensions

AM has also been extended beyond differentiable continuous state spaces. “Discrete Adjoint Matching” models discrete generative processes as continuous-time Markov chains and derives a discrete adjoint from Dynkin’s formula rather than from continuous control calculus. For a base CTMC with rates $p^\star(x)\propto p_{\mathrm{base}}(x)\exp(r(x))$4 and terminal loss $p^\star(x)\propto p_{\mathrm{base}}(x)\exp(r(x))$5, the optimal controlled rate is

$p^\star(x)\propto p_{\mathrm{base}}(x)\exp(r(x))$6

where

$p^\star(x)\propto p_{\mathrm{base}}(x)\exp(r(x))$7

The lean discrete adjoint has terminal condition

$p^\star(x)\propto p_{\mathrm{base}}(x)\exp(r(x))$8

and backward equation

$p^\star(x)\propto p_{\mathrm{base}}(x)\exp(r(x))$9

DAM then matches nonnegative rates using a generalized KL divergence. On reasoning tasks with LLaDA-8B-Instruct, the paper reports accuracy improvements over the base model and strong results on Countdown and Sudoku, including 60.16/55.47 on Countdown and 89.20/88.12 on Sudoku at sequence lengths 128/256 (So et al., 6 Feb 2026).

The “Discrete Adjoint Schrödinger Bridge Sampler” pushes the discrete AM picture further by identifying the core AM mechanism as state-space agnostic provided the state space admits a cyclic group structure yielding additive reference kernels. In that setting, controller and corrector identities can be written in terms of translated ratios of Schrödinger potentials, and AM becomes a target-matching problem over local Hamming-1 moves. The paper uses a translation-invariant reference CTMC on $g=-r$0 and reports competitive sample quality with significant advantages in training efficiency and scalability on Ising and Potts benchmarks (Guo et al., 9 Feb 2026).

“Unsupervised Diffusion Solver for Combinatorial Optimization via Combinatorial Adjoint Matching” adapts the discrete adjoint idea to optimization over Hamming neighborhoods. The discrete adjoint is defined as a cost-to-go difference

$g=-r$1

and, after a low-temperature simplification, yields a lean recursion driven only by the observed forward path and its Hamming-1 neighbors. The resulting CAM loss uses a terminal perturbation estimator based on

$g=-r$2

so training becomes a low-variance matching of coordinate-wise flip rates to adjoint targets. The paper reports that CAM consistently outperforms existing unsupervised diffusion baselines and achieves performance competitive with strong supervised diffusion solvers and traditional solvers across MIS, MaxCut, TSP, and CVRP (Feng et al., 29 May 2026).

A plausible interpretation of these discrete papers is that AM does not fundamentally depend on Euclidean gradients; what matters is the existence of a backward quantity that represents future cost sensitivity and can be matched locally to the generator’s control parameters.

5. Scientific and engineering applications

One of the most detailed scientific instantiations is “FlowBack-Adjoint,” which uses AM to fine-tune a conditional flow-matching model for all-atom protein backmapping from C$g=-r$3 traces. The reward is defined from a CHARMM27 molecular mechanics energy,

$g=-r$4

so the adjoint boundary is

$g=-r$5

with $g=-r$6. The control correction is

$g=-r$7

and training minimizes the residual

$g=-r$8

Relative to FlowBack, FlowBack-Adjoint lowers single-point energies by a median of approximately 78 kcal/mol.residue, reduces bond-length errors by more than 92%, eliminates more than 98% of molecular clashes, maintains ensemble diversity, and generates configurations that can initialize stable all-atom molecular dynamics simulations without energy relaxation (Berlaga et al., 5 Aug 2025).

AM also appears in variational data assimilation. “Adjoint-Matching Neural Network Surrogates for Fast 4D-Var Data Assimilation” studies neural surrogates for the Lorenz-63 system in a strong-constraint 4D-Var setting. The surrogate forward model $g=-r$9 is trained either with the standard forward loss

$X_0 \perp X_1$0

or with adjoint-informed terms such as

$X_0 \perp X_1$1

In sequential 4D-Var, the exact solver reaches RMSE $X_0 \perp X_1$2, the full adjoint-matching surrogate reaches $X_0 \perp X_1$3, and the forward-only surrogate reaches $X_0 \perp X_1$4. Average wall time per 4D-Var solve drops from $X_0 \perp X_1$5 s for the exact model to $X_0 \perp X_1$6 s for the adjoint-matching surrogate (Chennault et al., 2021).

A different engineering use appears in “Real-time elastic partial shape matching using a neural network-based adjoint method.” There the control is a boundary force distribution $X_0 \perp X_1$7, the state is a hyper-elastic displacement field $X_0 \perp X_1$8, and the forward PDE solve is replaced by a feed-forward neural surrogate $X_0 \perp X_1$9. The registration objective is

$p_{\mathrm{base}}$0

with gradient

$p_{\mathrm{base}}$1

Backpropagation through the surrogate therefore plays the role of the adjoint. Reported timings are $p_{\mathrm{base}}$2 ms for beam registration and an approximately $p_{\mathrm{base}}$3 per-frame speedup for liver reconstruction relative to a Newton-based adjoint pipeline (Odot et al., 2023).

These applications broaden the meaning of AM beyond generative-model fine-tuning. In all three cases, however, the essential operation remains the same: a backward sensitivity is computed or approximated, and a local regression or optimization step matches a controllable field to that sensitivity.

6. Limitations, scope, and terminological distinctions

The original AM framework is mathematically tied to the memoryless schedule. The 2024 paper explicitly states that if fine-tuning is performed with an arbitrary $p_{\mathrm{base}}$4, the path distribution couples $p_{\mathrm{base}}$5 and $p_{\mathrm{base}}$6 and introduces a value-function bias, so the terminal marginal does not equal the intended reward tilt (Domingo-Enrich et al., 2024). This restriction is benign in some architectures, but it is a real scope condition rather than a cosmetic technicality.

Later variants retain their own assumptions. QAM depends on the quality of the critic gradient $p_{\mathrm{base}}$7 and on accurate vector-Jacobian products through the behavior flow; the paper lists critic quality, lean-adjoint accuracy, computational overhead, and support mismatch with the behavior policy as limitations (Li et al., 20 Jan 2026). AMDP relies on a memoryless reference with $p_{\mathrm{base}}$8 in its main setting and replaces exact policy entropy by a lower bound, although the paper proves contraction and fixed-point preservation under that surrogate (Thilges et al., 21 Jun 2026). EAM replaces AM’s original base drift by a linear one and estimates $p_{\mathrm{base}}$9 through a Tweedie-based approximation, so its speed gains come with an additional modeling assumption (Shin et al., 12 May 2026). Deterministic truncated AM requires tuning the truncation horizon and the order of the control penalty, and the paper’s strongest results use problem-dependent reward scales and polynomial regularizers (Guo et al., 7 May 2026). Discrete variants may require either tractable importance weights, a cyclic-group state space, or both (So et al., 6 Feb 2026, Guo et al., 9 Feb 2026).

A separate point concerns nomenclature. The paper “On the Equivalence of Optimal Transport Problem and Action Matching with Optimal Vector Fields” states unambiguously: “In this paper AM stands for Action Matching (not ‘Adjoint Matching’).” Its AM is a variational approach over scalar potentials $\sigma(t)^2 = 2\eta_t + \chi(t),$0 for reproducing a prescribed sequence of distributions,

$\sigma(t)^2 = 2\eta_t + \chi(t),$1

and the paper proves that, when restricted to optimal OT-type vector fields induced by a convex Brenier potential $\sigma(t)^2 = 2\eta_t + \chi(t),$2, this Action Matching loss reduces to the OT dual up to constants (Kornilov et al., 31 Oct 2025). This usage is conceptually separate from adjoint-based AM in generative modeling and RL, despite the acronym overlap.

Taken together, the literature supports a broad but technically precise view of Adjoint Matching. It is not a single algorithmic recipe; it is a family of optimization principles in which backward sensitivities are converted into local matching targets for a learnable dynamical law. In continuous generative modeling, this law is usually a drift or velocity field; in RL it is a policy flow or diffusion control; in discrete settings it is a CTMC rate; in scientific computing it may be a neural surrogate, a deformation control, or a physics-guided transport. The common thread is the adjoint itself.