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.
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
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