Technical Summary: Generative Modeling by Value-Driven Transport

Problem Formulation

The paper addresses the problem of generative modeling, framed as a measure transport task: transforming samples from a source distribution $\nu_{src}$ into a target distribution $\nu_{tgt}$. While optimal transport (OT) provides a solid mathematical foundation for this, the authors argue that standard continuous-time approaches (e.g., diffusion models, Schrödinger bridges) rely on stochastic differential equations (SDEs) and score functions.

Instead, this work formulates the problem as a discrete-time stochastic control problem. The goal is to steer an initial state distribution through a sequence of $H$ decision rounds to match the target distribution at the final step while minimizing the cumulative movement cost. The cost function is defined as the sum of squared Euclidean distances between consecutive states, scaled by the horizon. A key challenge is the hard constraint that the final state distribution must exactly match $\nu_{tgt}$, which renders standard dynamic programming inapplicable.

Methodology

1. Linear Programming Formulation

The core theoretical contribution is reformulating the discrete-time dynamic OT problem as a Linear Program (LP).

• Primal Variables: The problem is cast in terms of occupancy measures $\mu_h$, representing the joint distribution of state pairs $(X_h, X_{h+1})$ at each time step. The constraints enforce temporal consistency (flow constraints) and boundary conditions (source and target constraints).
• Dual Variables: The dual variables of this LP correspond to a sequence of functions $V = \{V_0, \dots, V_{H+1}\}$, interpreted as value functions.
• Duality: The authors establish strong duality, showing that the optimal value of the LP equals the squared Wasserstein-2 distance ($W_2^2$). Crucially, the optimal dual variables (value functions) directly encode the optimal control policy.

2. Value-Driven Transport (VDT) Policies

A key insight derived from the optimality conditions is that the optimal policy $\pi^*$ is "greedy" with respect to the optimal value function $V^*$. Specifically, the optimal policy moves a state $x$ in the direction of the negative gradient of the value function:
$$\pi^*_h(x) = x - \frac{1}{H+1} \nabla_x V^*_h(x)$$
This structure implies that optimal transport paths are linear interpolations between optimally coupled source and target points. The authors define Value-Driven Transport (VDT) policies as those that follow this gradient field, allowing for explicit computation of the transport path once the value function is known.

3. Primal-Dual Algorithm

To compute the value functions without simulating trajectories from the learned model, the authors propose a simulation-free primal-dual algorithm:

• Dual Update: The value function $V_\theta$ is parametrized by a neural network. The algorithm performs stochastic gradient ascent on the dual function using minibatches of source and target samples.
• Primal Update: To evaluate the gradients of the dual function, the algorithm must minimize the Lagrangian with respect to the primal variables (the occupancy measures). Instead of optimizing over distributions directly, the authors use Wasserstein gradient descent on a system of particles.
  • Initialization: Particles are initialized as linear interpolations between source and target samples (using an OT coupling for the minibatch).
  • Update: Particles are updated via gradient descent steps to satisfy the flow constraints and minimize the cost.
• Efficiency: This approach requires only a single neural network for the value function and avoids the need for simulating forward trajectories during training.

4. Extensions

The framework naturally supports:

• Time Scaling: A value function trained at a specific horizon $H$ can be used for generation at different time scales (fewer steps) by adjusting the step size, leveraging the invariance of the optimal value function under time discretization.
• Conditional Generation: Class labels or embeddings can be concatenated to the input of the value function.
• Classifier-Free Guidance (CFG): The method supports CFG by combining gradients from conditional and unconditional value functions, similar to diffusion models.
• Reverse Generation: Due to the symmetry of the formulation, the learned value function can be used to transport from target to source by reversing the gradient direction.

Key Results

The authors evaluate VDT on 2D synthetic datasets and the MNIST dataset, comparing it against state-of-the-art methods like Flow Matching (FM), Conditional Flow Matching (CFM), Rectified Flow (RF), and Schrödinger Bridge (SB) variants.

• 2D Experiments: VDT achieves competitive Wasserstein-2 distances to the target. Notably, using an OT coupling for particle initialization (VDT+) significantly improves performance, allowing for high-quality generation even with a single step.
• Few-Step Generation: A major finding is that VDT policies maintain high quality even when the number of generation steps ($H_{test}$) is drastically reduced compared to the training horizon ($H$). For example, a model trained with $H=100$ steps can generate high-quality samples in $H_{test}=10$ steps with negligible loss in quality, offering a 10x speedup.
• MNIST Tasks:
  • Paired Deblurring: Successfully reconstructs clean images from downsampled/blurred inputs.
  • Unpaired Translation: Translates handwritten letters (EMNIST) to digits (MNIST) without explicit pairing, demonstrating the ability to learn non-trivial domain mappings.
  • Conditional Generation: Generates digits conditioned on labels, with CFG improving image quality at the cost of diversity, mirroring behavior in diffusion models.
• Path Energy: The path energy of VDT trajectories closely approximates the theoretical $W_2^2$ distance, confirming the straightness of the transport paths.

Significance and Claims

The paper positions VDT as a distinct alternative to the dominant continuous-time, score-based, or flow-based generative modeling paradigms.

• Conceptual Shift: It roots generative modeling in discrete-time stochastic optimal control (a reinforcement learning setting) rather than continuous-time SDEs.
• Straight Paths: The method inherently learns policies that generate straight transport paths, enabling efficient few-step generation without the complex training schedules often required by diffusion models to achieve similar speeds.
• Primal-Dual Efficiency: By solving the LP via a primal-dual method with particle-based primal updates, the method avoids the computational overhead of simulating trajectories during training.
• Modest Limitations: The authors acknowledge that the method is currently limited to small-scale benchmarks (2D and MNIST) and has not yet been scaled to high-resolution image generation (e.g., ImageNet) due to the lack of established best practices for this specific primal-dual architecture. They view the current work as a proof of concept that the LP framework can serve as a basis for competitive large-scale algorithms, inviting future work to refine the training methodology for more challenging tasks.