Generative optimal transport via forward-backward HJB matching

This paper proposes a generative optimal transport framework that solves the minimum-work control problem for steering stochastic systems from a disordered reference to a structured target by leveraging time-reversal duality to compute the required backward dynamics via a tractable forward HJB equation derived from path-space free energy.

The Gist

Imagine you have a jar full of marbles scattered randomly on a table (the disordered reference state). Your goal is to magically arrange them into a perfect, intricate pattern, like a smiley face or a star (the structured target).

In the real world, nature does the opposite. If you leave a neat pattern alone, it naturally falls apart into randomness due to heat and chaos (diffusion). The big question this paper asks is: "What is the most energy-efficient way to push those random marbles back into that perfect pattern?"

Usually, figuring out the perfect path to do this is a nightmare. You'd need to know exactly where every marble will be in the future to plan the push now. But you don't know the future yet! That's the catch-22.

This paper solves that puzzle with a clever trick: Time Reversal and a "Cost Map."

Here is the breakdown using simple analogies:

1. The Problem: The "Backwards" Trap

Imagine you are trying to walk from a messy room to a clean room. To find the perfect path, you might think, "I need to see the clean room first to know where to step." But you can't see the clean room until you get there.

• Old methods tried to guess the path by simulating the walk backwards from the clean room, but they didn't have a map of the clean room to start with.
• This paper says: "Let's walk forward from the messy room to the clean room first, just to learn the terrain. Then, we'll use what we learned to walk backwards perfectly."

2. The Solution: The "Feynman-Kac" Magic Trick

The authors use a mathematical magic trick called the Feynman-Kac formula. Think of it like this:

Instead of trying to calculate the perfect path directly (which is hard), they calculate the "Free Energy" of all possible paths.

• Imagine every possible path a marble could take has a "price tag" attached to it.
• Some paths go through mud (high cost).
• Some paths go through smooth ice (low cost).
• The "Free Energy" is like an average of all these prices, but it heavily penalizes the expensive, muddy paths.

By simulating a simple, random walk from the messy room to the clean room (forward in time), they can calculate this "average price" for every spot on the floor. This creates a Value Function—a giant, invisible 3D landscape where low points are "good" paths and high points are "bad" paths.

3. The "Refractive Index" (Fermat's Principle)

This is the coolest part. The paper introduces a Spatial Cost Field (let's call it $\nu$).

• Think of this like the refractive index in optics (how light bends when it passes through water vs. air).
• If you want the marbles to avoid a specific area (like a hot stove), you make that area "expensive" (high cost).
• If you want them to go through a specific tunnel, you make that area "cheap" (low cost).

Just like light bends to take the fastest route through different materials (Fermat's Principle), the marbles will naturally bend their paths to avoid the "expensive" zones and hug the "cheap" zones. The math automatically figures out the curved path that minimizes the total "effort" and "cost."

4. How It Works in Practice (The Two-Step Dance)

The method works in two distinct phases:

Phase A: The Forward Scouting Run (Training)

1. Start with random marbles (the messy room).
2. Let them drift naturally toward the target shape, but keep a record of how much "effort" it would take to stay on different paths.
3. Use a neural network (a smart computer brain) to learn the Value Function (the 3D cost map) based on these forward drifts.
   • Analogy: It's like a hiker walking up a mountain in the fog, marking the steepness of the ground, so they can build a map of the easiest route down.

Phase B: The Backward Generation (Creation)

1. Now, take a new set of random marbles.
2. Use the map (the Value Function) learned in Phase A.
3. Push the marbles backwards (from the target shape back to the start, or rather, generate the target from the start by reversing the logic).
4. The marbles follow the gradient of the map, sliding down the "valleys" of low cost, perfectly arranging themselves into the target shape with minimal wasted energy.

Why Is This a Big Deal?

• No Guessing: You don't need to know the final answer to start the process. You just need to simulate the "messy" forward process, which is easy.
• Physical Meaning: It's not just a black-box AI trick. It's grounded in physics (thermodynamics and fluid dynamics). It treats the generation of data like moving a fluid with the least amount of energy.
• Control: You can literally "paint" the cost map. Want the generated images to avoid certain shapes? Just make those areas "expensive" on the map. The AI will learn to route around them automatically.

Summary

The paper teaches us how to reverse-engineer chaos into order without needing a crystal ball. By simulating the easy "forward" drift and calculating the "cost" of every path, they build a map that guides the system to the perfect target state with the least amount of work, bending the paths of particles just like a lens bends light.

Technical Summary

1. Problem Statement

The paper addresses the challenge of controlling the evolution of a many-body stochastic system from a disordered reference state ($p_{ref}$) to a structured target ensemble ($p_{data}$), characterized empirically by samples. This is a fundamental problem in non-equilibrium statistical mechanics and stochastic control.

• The Core Difficulty: The natural relaxation of such a system (driven by diffusion) moves from the structured target to the disordered reference. Reversing this process to generate samples requires finding the minimum-work stochastic process.
• The Circular Dependency: Computing the optimal control process typically requires knowledge of trajectories that already sample the target ensemble. However, the goal is to construct these trajectories. Existing methods (like score-matching or backward SDE simulation) often struggle with this circularity or lack a direct link to physical optimality principles (trajectory-level costs).
• Objective: To find a stochastic control policy that minimizes a pathwise cost functional combining spatial penalties (e.g., avoiding invalid states) and control effort, without requiring prior knowledge of the generative process or backward-time simulation.

2. Methodology: Forward-Backward HJB Matching

The authors propose a framework rooted in Stochastic Optimal Control (SOC) and Dynamic Optimal Transport (OT). The core innovation is a time-reversal duality that transforms an intractable backward control problem into a tractable forward learning problem.

A. The Control Formulation

The problem is formulated as minimizing the expected total cost over trajectories $x_t$:
$$\min_{u_t} \mathbb{E}_{P_u} \left[ \int_0^1 \nu(x_t) dt + \frac{\gamma}{2} \int_0^1 \|u_t\|^2 dt \right]$$
subject to the controlled SDE $dx_t = u_t dt + \sqrt{2D} dB_t$, with fixed marginals $x_0 \sim p_{ref}$ and $x_1 \sim p_{data}$.

• $\nu(x)$: A spatial cost function (acting as a refractive index).
• $\gamma$: Weight for control effort.
• $u_t$: Control input.

B. The Duality Theorem (Key Insight)

Instead of solving the backward-time Hamilton-Jacobi-Bellman (HJB) equation directly (which requires target-to-reference samples), the authors establish a duality:

1. Backward Problem: The optimal control for the generative process is the gradient of a value function $U(t, x)$ satisfying a backward HJB equation.
2. Forward Transformation: By defining a forward potential $W(s, x) := -U(1-s, x)$, the backward HJB transforms into a forward HJB equation:
   $$\frac{\partial W}{\partial s} - D\Delta W - \frac{1}{2\gamma} \|\nabla W\|^2 + \nu(x) = 0$$
3. Cole-Hopf Linearization: Using the transformation $W = \frac{1}{\beta} \log Z$ (where $\beta = 1/2D\gamma$), the nonlinear HJB becomes a linear parabolic PDE:
   $$\frac{\partial Z}{\partial t} = D\Delta Z - \beta \nu Z$$
4. Feynman-Kac Representation: The solution $Z$ can be computed as a path-space expectation (free energy) over forward-time trajectories sampled from the data distribution $p_{data}$ evolving toward the reference $p_{ref}$ (e.g., via Langevin dynamics or Ornstein-Uhlenbeck processes).

C. Learning Algorithm

The framework learns a scalar potential $W_\theta(s, x)$ (parameterized by a neural network) using forward-time supervision:

• Training Data: Trajectories generated by a forward diffusion process (e.g., OU process) moving from $p_{data} \to p_{ref}$.
• Loss Function: A combination of:
  1. Feynman-Kac Loss ($L_{FK}$): Enforces consistency between the network output and the path-integral estimate of the cost-to-go.
  2. Local Loss ($L_{FK-local}$): Ensures one-step temporal consistency (semigroup property).
  3. Dual Loss ($L_{dual}$): Enforces boundary conditions derived from the Kantorovich dual problem.
• Generation: Once trained, samples are generated by reversing time: starting from $p_{ref}$ and simulating a controlled SDE where the drift is defined by the gradient of the learned potential $\nabla W$.

3. Key Contributions

1. Forward-Backward HJB Duality: Theorem 2.2 establishes a rigorous connection between the generative (backward) control problem and a forward stochastic control problem. This allows learning the optimal transport policy using only forward-time samples from the data, eliminating the need for backward SDE simulation or explicit score estimation.
2. Spatial Cost Geometry (Fermat's Principle): The introduction of the spatial cost function $\nu(x)$ acts as a "refractive index" on the path space. It shapes the transport geometry analogously to Fermat's Principle in optics, allowing trajectories to be deflected, focused, or confined based on physical constraints without modifying the underlying model architecture.
3. Risk-Sensitive Control: The formulation naturally incorporates variance control. The parameter $\gamma$ modulates the trade-off between expected cost and path variance (risk), emerging naturally from the stochastic control formulation rather than requiring ad-hoc regularization.
4. Unified Framework: The work unifies Stochastic Optimal Control, Schrödinger Bridge theory, and Non-equilibrium Statistical Mechanics, interpreting the value function as a path-space free energy.

4. Results

The authors validate the theory through numerical experiments:

• 2D Benchmarks: On datasets like "4 Gaussians," "2 Moons," and "Swiss Roll," the learned value function $W(t, x)$ successfully develops structured basins (potential wells) that align with the target geometry. The model generates samples that aggregate onto the target manifold with monotonic loss convergence.
• Geometric Control (Refraction): In a controlled experiment moving a Gaussian source to a Gaussian target, the authors demonstrated that modifying $\nu(x)$ acts as an optical lens:
  • Convex $\nu$ (High cost barrier): Deflects trajectories around the obstacle.
  • Concave $\nu$ (Low cost well): Focuses trajectories through the center.
  • This confirms that $\nu(x)$ provides a powerful, interpretable mechanism for guiding stochastic transport.
• High-Dimensional Scalability: The method was successfully applied to the MNIST dataset (784 dimensions). Using a U-Net architecture and conditional sampling from the OU process, the model learned a coherent propagating potential structure that guided noise to digit images, demonstrating scalability beyond 2D.

5. Significance

• Physical Interpretability: Unlike black-box generative models, this approach provides a physically interpretable description of transport in terms of free energy, risk-sensitive control, and spatial cost geometry.
• Efficiency and Stability: By avoiding backward-time simulation and explicit score matching, the method sidesteps the instability and data requirements associated with those techniques. It leverages the tractability of forward diffusion for training.
• Constraint Handling: The spatial cost field $\nu(x)$ offers a principled way to enforce physical constraints (e.g., safety regions, energy barriers) in generative modeling, making it highly relevant for physics-informed machine learning and engineering applications.
• Theoretical Bridge: It provides a unifying mathematical link between the macroscopic evolution of probability distributions (OT) and microscopic trajectory optimization (SOC), offering new tools for controlling complex non-equilibrium systems.