Adjoint Matching through the Lens of the Stochastic Maximum Principle in Optimal Control

TL;DR

This paper rigorously derives and generalizes Adjoint Matching for stochastic optimal control problems, connecting it to the Stochastic Maximum Principle and providing practical algorithms for control learning.

Contribution

It formulates a general Hamiltonian adjoint matching objective, shows its equivalence to the original SOC problem, and interprets it as a continuous-time approximation method.

Findings

Recovered the lean adjoint matching loss for state- and control-independent diffusion.

Demonstrated that critical points satisfy HJB stationarity conditions.

Provided a practical alternative to classical SMP algorithms avoiding intractable martingale terms.

Abstract

Reward fine-tuning of diffusion and flow models and sampling from tilted or Boltzmann distributions can both be formulated as stochastic optimal control (SOC) problems, where learning an optimal generative dynamics corresponds to optimizing a control under SDE constraints. In this work, we revisit and generalize Adjoint Matching, a recently proposed SOC-based method for learning optimal controls, and place it on a rigorous footing by deriving it from the Stochastic Maximum Principle (SMP). We formulate a general Hamiltonian adjoint matching objective for SOC problems with control-dependent drift and diffusion and convex running costs, and show that its expected value has the same first variation as the original SOC objective. As a consequence, critical points satisfy the Hamilton--Jacobi--Bellman (HJB) stationarity conditions. In the important practical case of state- and control-independent diffusion, we recover the lean adjoint matching loss previously introduced in adjoint matching, which avoids second-order terms and whose critical points coincide with the optimal control under mild uniqueness assumptions. Finally, we show that adjoint matching can be precisely interpreted as a continuous-time method of successive approximations induced by the SMP, yielding a practical and implementable alternative to classical SMP-based algorithms, which are obstructed by intractable martingale terms in the stochastic setting. These results are also of independent interest to the stochastic control community, providing new implementable objectives and a viable pathway for SMP-based iterations in stochastic problems.