Generative optimal transport via forward-backward HJB matching

TL;DR

This paper introduces a novel method for controlling stochastic systems by leveraging a time-reversal duality, enabling the computation of optimal transport processes using forward trajectories and path-space free energy, bridging control and statistical mechanics.

Contribution

It establishes a forward-backward HJB duality that allows efficient computation of optimal stochastic transport without backward simulations, connecting control, Schrödinger bridges, and statistical mechanics.

Findings

The value function satisfies an equivalent forward HJB equation.

Forward relaxation trajectories suffice to compute the value function.

Numerical examples visualize how cost fields influence transport geometry.

Abstract

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\"odinger bridge theory, and non-equilibrium statistical mechanics.