Generative optimal transport via forward-backward HJB matching
Authors: Haiqian Yang, Vishaal Krishnan, Sumit Sinha, L. Mahadevan
Organizations: Department of Mechanical Engineering, MIT, Cambridge, MA 02139 · SEAS, Harvard University, Cambridge, MA 02138 · SEAS, Physics and OEB, Harvard University, Cambridge, MA 02138
Controlling the evolution of a many-body stochastic system from a disordered reference state to a structured target ensemble, characterized empirically through samples, arises naturally in non-equilibrium statistical mechanics and stochastic control. The natural relaxation of such a system - driven by diffusion - runs from the structured target toward the disordered reference. The natural question is then: what is the minimum-work stochastic process that reverses this relaxation, given a pathwise cost functional combining spatial penalties and control effort? Computing this optimal process requires knowledge of trajectories that already sample the target ensemble - precisely the object one is trying to construct. We resolve this by establishing a time-reversal duality: the value function governing the hard backward dynamics satisfies an equivalent forward-in-time HJB equation, whose solution can be read off directly from the tractable forward relaxation trajectories. Via the Cole-Hopf transformation and its associated Feynman-Kac representation, this forward potential is computed as a path-space free energy averaged over these forward trajectories - the same relaxation paths that are easy to simulate - without any backward simulation or knowledge of the target beyond samples. The resulting framework provides a physically interpretable description of stochastic transport in terms of path-space free energy, risk-sensitive control, and spatial cost geometry. We illustrate the theory with numerical examples that visualize the learned value function and the induced controlled diffusions, demonstrating how spatial cost fields shape transport geometry analogously to Fermat's Principle in inhomogeneous media. Our results establish a unifying connection between stochastic optimal control, Schrödinger bridge theory, and non-equilibrium statistical mechanics.
Figures & tables
Figure 1: Optimally connecting endpoint marginals via forward-backward HJB matching. We formulate generative modeling as a stochastic optimal control problem that transports samples from a reference distribution pref to a data distribution pdata , minimizing a trajectory-wise cost. The optimal control policy arises from the backward-time Hamilton–Jacobi–Bellman (HJB) equation for a value function U(s,x) , whose gradient defines the generative drift. By time-reversing this value function as W(s,x):=−U(1−s,x) , we obtain a forward HJB equation that is solved via a Feynman–Kac path integral, using uncontrolled forward trajectories sampled from pdata→pref . This establishes a dual connection between the forward (training) and backward (generation) dynamics, unified through the value function W and its governing HJB equation. The formulation incorporates a spatial cost function ν(x) that modulates the pathwise transport geometry.
Figure 2: HJB value function learning and generative transport on standard 2D benchmarks. We solve the stochastic optimal control problem (Eqn. 2 ): find the control u∗(t,x)=∇xW that transports the reference pref∼N(0,I) to the empirical target pdata while minimizing a pathwise cost combining spatial penalty ν(x) and quadratic control effort. The scalar value function W(t,x) , which solves the forward HJB equation (Eqn. 4 ), is parameterized as a 10-layer MLP (64 hidden units, GeLU activations, sinusoidal positional time embeddings) and trained via Feynman–Kac trajectory supervision using the total loss Ltotal=LFK+Ldual (Eqn. 14 , Algorithm 3 ). Forward trajectories are generated by an Ornstein–Uhlenbeck process with diffusion coefficient D=0.05 and risk parameter β=0.1 , discretized over T=128 steps. The model is trained for n=2000 epochs on N=4096 samples using the Adam optimizer (learning rate η=10−3 ), implemented in JAX. Results are shown on three standard 2D benchmarks used in the generative transport literature [ 10 , 23 ] : (a) 4 Gaussians, four isotropic clusters at (±2,0) and (0,±2) with σ=0.2 ; (b) 2 Moons, two interleaved half-moon arcs with additive Gaussian noise; (c) Swiss Roll, a planar spiral manifold. Top row: snapshots of the learned potential W(t,x) at t=0,0.5,1.0 , showing the emergence of structured basins aligned with the target geometry as transport progresses. Middle row: particle positions at each time step under the learned backward-time controlled diffusion (Algorithm 3 ), demonstrating successful aggregation onto pdata at terminal time t=1 . Bottom row: convergence of Ltotal over training epochs; monotonic decay across all three datasets confirms robust Feynman–Kac trajectory supervision without explicit score estimation or backward-time SDE simulation.
Figure 3: Geometric control of stochastic transport via spatial cost field ν(x) . We solve the same stochastic optimal control problem as in Fig. 2 (Eqn. 2 ), but with an inhomogeneous spatial cost field ν(x) that shapes the geometry of optimal transport paths (cf. Fermat’s principle). The source is a Gaussian p0=N((−1,0),0.01I) and the target is p1=N((1,0),0.01I) , with a Gaussian ν(x) profile (mean at origin, σν=0.1 ) placed at the midpoint between them. The value function W(t,x) is parameterized as a 10-layer MLP (64 hidden units, sinusoidal time embeddings) and trained via Feynman–Kac trajectory supervision (Algorithm 3 , Ltotal=LFK+Ldual , Eqn. 14 ). Forward trajectories use an Ornstein–Uhlenbeck process ( D=0.05 , β=0.1 ) with T=100 steps; training runs for n=1000 epochs on N=512 samples (Adam, η=10−3 , JAX). Trajectories (green) connect p0 (red) to p1 (blue). (a) Uncontrolled forward diffusion (no learned control), producing unstructured spreading. (b) Controlled transport with flat ν(x)=const , recovering straight-line baseline geometry. (c) Concave cost profile: low ν at the midpoint creates a potential well, concentrating optimal paths inward through the center. (d) Convex cost profile: high ν at the midpoint acts as an energetic barrier, deflecting optimal paths outward around the obstacle. (e) Convergence of Ltotal for the three controlled profiles, confirming robust training under different spatial cost fields.
Appendix figures & tables1 asset
Supplementary material from the paper’s appendix.
Appendix
Figure 4: HJB value function learning in high dimensions: transport from Gaussian noise to MNIST. We solve the same stochastic optimal control problem (Eqn. 2 ) as in the 2D setting, now in the high-dimensional pixel space of MNIST ( 28×28=784 dimensions). The reference pref=N(0,I) is transported to pdata (MNIST digit images) with a fixed uniform cost field ν(x)=1 . The value function Wθ(t,x) is parameterized as a convolutional U-Net (encoder channels [32,64,128,256] , GroupNorm, skip connections, Gaussian Fourier time embeddings with embed_dim=256 ) and trained via Feynman–Kac trajectory supervision using the total loss Ltotal=LFK+Ldual+Lspatial (Eqn. 14 , Algorithm 3 ). Forward trajectories are sampled from the closed-form conditional of an Ornstein–Uhlenbeck process ( D=0.01 , θ=5.0 , β=0.1 , T=100 steps), enabling memory-efficient training without rolling out full forward trajectories. Training runs for n=200 epochs on the full MNIST training set ( batch size=1024 ) using the Adam optimizer ( η=10−4 ), implemented in PyTorch. We visualize the scalar potential W(t,x(s)) evaluated along held-out test trajectories {x(s)} (unseen during training) at successive time steps. (a) At initialization, the potential is unstructured and shows no coherent alignment with the transport direction. (b) After 200 epochs, a clear propagating bump emerges: each trajectory exhibits a moving pulse in W aligned with the generative path, consistent with the learned cost-to-go structure observed in the 2D setting (Fig. 2 ). (c) Convergence of Ltotal over training epochs. The emergence of coherent potential structure on test trajectories confirms that the Feynman–Kac framework generalizes beyond the training ensemble to high-dimensional image data.
We establish the equivalence between the stochastic optimal control and path space formulations of the Schrödinger bridge problem (SBP) for the kinematic equation on a compact connected Lie group. Using the geometric concepts of horizontal lift and stochastic anti-development, we derive a Girsanov-type change-of-measure result, and show that the expected control energy equals the relative entropy of the controlled path law with respect to the reference Wiener measure. Thus, the SBP is equivalently a path space relative entropy minimization problem subject to prescribed endpoint marginals.Our result has three useful implications. From an analytic viewpoint, the shown equivalence helps prove the existence and uniqueness of the SB. From a probabilistic viewpoint, it helps interpret the SB as the most probable deviation of the uncontrolled stochastic dynamics consistent with the endpoint constraints. From a computational viewpoint, it allows using static Sinkhorn recursions to directly solve the relative entropy minimization problem and compute the optimal path measure. We illustrate the equivalence numerically on the torus T2. The code is publicly available at: https://github.com/gradslab/LargeDeviationSBP
Hamza Mahmood, Georgiy A. Bondar, Abhishek Halder +1
Dept. of Mechanical & Industrial Engineering at the New Jersey Institute of Technology, Newark, NJ 07102, USA · Department of Applied Mathematics at the University of California at Santa Cruz, CA, 95064, USA · Dept. of Aerospace Engineering at Iowa State University, Ames, IA 50011, USA
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, which avoids second-order terms and whose critical points coincide with the optimal control under mild uniqueness assumptions. Numerical experiments confirm that the extra terms it discards become necessary once the diffusion is state-dependent. 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.
Carles Domingo-Enrich, Jiequn Han
Microsoft Research New England · Flatiron Institute
We discuss three simple connections between generative diffusion and optimal transport. First, the joint law of the data and of their Ornstein-Uhlenbeck noising is, at every time, the solution of an entropic optimal-transport problem, at a temperature set by the noise level. Backward diffusion can consequently be read as an annealing which gradually decreases the entropic regularization. Second, we study the speciation transition--how the modes of the data distribution emerge during generation--and compare how it happens in diffusion and flow matching with respect to optimal transport. Using Gaussian mixtures as high-dimensional example, we show that the dynamical instability through which diffusion and flow matching separate the modes becomes, for optimal transport, an exponentially thin boundary layer at the beginning of generation, which is rounded by entropic regularization. Third, considering a generic forward dynamics with a general confining potential, we show that the backward process driven by the score also satisfies a drift-based variational principle whose minimal cost is the entropy difference between initial and final distributions.
Giulio Biroli, Marc Mézard
Laboratoire de Physique de l’Ecole Normale Supérieure, ENS, Université PSL, CNRS, Sorbonne Université, Université Paris-Diderot, Sorbonne Paris Cité, Paris, France · Department of Computing Sciences, Bocconi University, Milano, Italy