Generative Modeling by Value-Driven Transport

Abstract: We propose a new framework for generative modeling based on a discrete-time stochastic control formulation of measure transport. Adapting classic results from control theory, we formulate our problem as a linear program whose dual variables correspond to the \emph{optimal value function} of the control problem, which directly encodes the optimal control policy. Exploiting this LP formulation, we develop an efficient simulation-free primal-dual algorithm for computing approximately optimal value functions and the associated \emph{value-driven transport} (VDT) policies which approximate the true optimal policy. We show that well-trained VDT policies enjoy numerous favorable properties in comparison with other state-of-the-art methods based on flows, diffusions, or Schrödinger bridges: they lead to straight transport paths which can be simulated quickly and robustly, and can be enhanced in all the same ways as diffusion and flow-based models (e.g., conditional generation, classifier-free guidance, unpaired data-to-data translation are all easy to incorporate). We evaluate our methodology in a range of experiments, with results that indicate strong performance and good potential for scalability.

• The paper introduces a discrete-time, value-function-based paradigm that reformulates optimal transport as an MDP for generative modeling.
• It employs a simulation-free primal-dual algorithm leveraging neural network-based value functions to achieve few-step generation with near baseline fidelity.
• The method’s symmetric design enables robust paired/unpaired domain translation and conditional generation while reducing computational overhead.

Generative Modeling by Value-Driven Transport: A Technical Overview

Introduction and Theoretical Context

The paper "Generative Modeling by Value-Driven Transport" (2605.22507) formulates a new generative modeling paradigm rooted in the discrete-time stochastic control formulation of optimal transport (OT). Most state-of-the-art deep generative frameworks—such as diffusion, flow-matching, and Schrödinger bridge models—rely extensively on continuous-time mathematics, trajectory simulation, or score estimation. In contrast, this work proposes a discrete-time perspective, reframing transport as a Markov decision process (MDP) that is solved via linear programming (LP) techniques borrowed from control theory and reinforcement learning.

The approach builds upon the observation that the Monge-Kantorovich OT problem, classically expressed either as a coupling or a deterministic map, can be equivalently written as a dynamic stochastic control problem. In this setting, one transports the mass of a source distribution onto a target distribution using a control policy over a finite horizon, minimizing a quadratic cost functional. The dynamics of the control process are expressed as either smooth ODEs (continuous) or discrete sequences of mappings (discrete):

Figure 1: Continuous- and discrete-time transport plans, comparing solutions defined by ODEs to those defined by direct mappings per time step.

This discrete formulation lays the groundwork for leveraging a Bellman-like structure to represent optimal policies, which are fully characterized by their value functions—a key motivation for the value-driven transport (VDT) methodology.

Value-Driven Transport Framework and Algorithms

The framework hinges on parameterizing the dual variables of the associated linear program as value functions. These value functions encapsulate the OMDP's (Optimal Mass Transport MDP's) optimal transport cost and policy through backward-recursive optimality equations structurally reminiscent of the Bellman equations but adapted to distribution-matching constraints. The Bellman recursion for the value function $V^*$ and policy $\pi^*$ proceeds as:

$$V^*_h(x) = \min_{y} \left[ \frac{H+1}{2} \|x-y\|^2 + V^*_{h+1}(y) \right], \qquad \pi^*_h(x) = \arg\min_{y} \left[ \frac{H+1}{2} \|x-y\|^2 + V^*_{h+1}(y) \right].$$

Policy trajectories are guaranteed, under the squared-Euclidean cost, to follow linear interpolations between coupled source-target points, a property which greatly simplifies both analysis and algorithmic implementation.

The model's generative mechanism is realized via a primal-dual training algorithm:

• The dual (value function) is parameterized by a neural network $V_\theta$, trained by stochastic gradient ascent.
• The primal variables are occupancy measures discretized via particle clouds, updated by Wasserstein gradient descent.
• Crucially, only one neural network is needed—direct policy parameterization is eschewed in favor of extracting policies from the value function gradients.
• The algorithm achieves simulation-free primal-dual optimization, improving numerical stability and facilitating scalability.

The structure of value functions and their gradients, as well as the correspondence to straight paths in the state space, are visually elucidated in the following two-dimensional illustration:

Figure 2: Learned value functions and VDT policy directions in a 2D benchmark at varying time indices, with straight sample paths generated from a source point.

Practical Handling: Few-step Generation and Time-Scale Invariance

A central practical claim is the robustness of test-time time scales: since the learned value function relates to transport policies with arbitrary horizons, a model trained with horizon $H$ can reliably generate samples with a significantly reduced number of inference steps $K \ll H$. This is achieved by renormalizing the value gradient steps, leveraging the straight-path property of the optimal transport plan:

Figure 3: Few-step VDT generation for 2D data; fast convergence to the target with only a fraction of training steps used for interpolation paths.

With proper initialization (using OT maps between batch samples), numerical results show negligible degradation in sample quality even when reducing the number of sampling steps by an order of magnitude. This behavior is not matched by most diffusion or bridge-based counterparts, where fidelity is tightly correlated with the granularity of the sampling grid.

Advanced Applications: Data Translation and Conditional Generation

The fully symmetric formulation between source and target enables versatile applications, including:

• Unpaired domain translation: The model accommodates unpaired mappings (e.g., EMNIST letters to MNIST digits), exploiting the reversibility and distribution-agnostic nature of the VDT policy.
• Paired translation: Leveraging explicit couplings between sample pairs when available.
• Conditional and guided generation: Providing labels or auxiliary information as input to the value function enables conditional sampling and classifier-free guidance (CFG) methods analogous to those in diffusion models.

Empirical evidence for these capabilities is summarized with MNIST and EMNIST experimental results:

Figure 4: MNIST experiments—(a): paired deblurring, (b): conditional generation and CFG, (c): unpaired translation between EMNIST and MNIST via VDT with few-step predictions.

Notably, classifier-free guidance is realized by combining conditional and unconditional value functions, directly influencing sample sharpness and diversity, as typically done in CFG-based diffusion pipelines.

Comparative Evaluation and Claims

Empirical benchmarks on standard 2D OT tasks (e.g., moons, scurve, 8-gaussians) show that VDT achieves Wasserstein-2 distances and path energies near the best existing baselines—including simulation-free Schrödinger bridge and rectified flow approaches. The tenfold reduction in sampling steps, achieved without compromising accuracy, stands as a bold empirical claim:

• The model exhibits nearly identical sample quality (measured by W2 distance/path energy) at $K$-step prediction with $K = H/10$ as at full $H$-step generation.
• This few-step property is not observed in standard diffusion or SDE-based generative models, which suffer substantial degradation with coarse time grids.

Theoretical and Practical Implications

Theoretical implications include:

• Demonstration that discrete-time, control-theoretic formulations are not only mathematically equivalent but also practically competitive for large-scale generative modeling.
• Introduction of a value-function-centric approach—eschewing direct policy parameterization—that could generalize to other constrained control problems.

Practical implications are:

• Streamlined, simulation-free training reduces computational overhead, easing adoption in settings with limited resources.
• Robustness to horizon mismatches introduces flexibility in resource-constrained or real-time environments.
• The algorithm's symmetry in data translation tasks (paired and unpaired) makes it suitable for a broad range of multimodal generative tasks.

Future research should explore scaling the approach to high-dimensional datasets and more complex modalities. The value function-based approach may also find applications in robust or constrained RL settings, inverse modeling, and theoretical investigations into generalization properties of optimal-transport-based methods.

Conclusion

This work establishes a discrete-time, value-function-based optimal transport framework for generative modeling, introducing the VDT architecture and simulation-free primal-dual learning. Empirical evaluations on synthetic 2D OT problems and MNIST/EMNIST translation tasks validate both the effectiveness and efficiency of generation, particularly the ability to operate at drastically reduced sampling horizons without loss of generative fidelity. The symmetric, theoretically-grounded, and computationally practical structure suggests promising avenues for further research, both in advancing optimal transport generative models and in bridging connections to control and RL.