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

2026-03-28

arXiv.org

AI Summary

This work addresses the challenge of reward fine-tuning in diffusion models and Boltzmann distribution sampling by formulating generative model optimization as a stochastic optimal control problem governed by stochastic differential equations. Leveraging the Stochastic Maximum Principle (SMP), the paper rigorously derives, for the first time, a general Hamiltonian adjoint matching objective applicable to settings where both drift and diffusion coefficients depend on the control, and establishes its intrinsic connection to the Hamilton–Jacobi–Bellman (HJB) equation. By integrating the adjoint system with a continuous-time successive approximation algorithm, the method recovers the lightweight adjoint loss when the diffusion coefficient is state-independent, validates the necessity of higher-order terms in state-dependent cases, and provides a tractable iterative scheme based on SMP that circumvents 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.

Problem

• stochastic optimal control
• adjoint matching
• Stochastic Maximum Principle
• state-dependent diffusion
• Hamiltonian

Innovation

• Stochastic Maximum Principle
• Adjoint Matching
• Stochastic Optimal Control
• Hamiltonian Formalism
• Diffusion Models