Motivated by recent applications in generative modeling and sampling, we introduce a framework for optimal measure transport where cost captures the notion of neural network complexity. In transport-based generative models, samples from a reference distribution (e.g. Gaussian) are mapped to samples of a target distribution along ordinary or stochastic differential equations. These are implemented as deep residual networks when discretized in time, where each hidden layer approximates the associated instantaneous velocity. Thus, given a pair of target and reference measures, a natural question is to search for the most efficient neural representation that implements this transport. Our starting point is the kinetic formulation of OT, due to Benamou and Brenier. We replace the average kinetic L2 energy by the \emph{Barron} energy \cite{bach2017breaking, ma2022barron}, a natural norm which measures the complexity of representing a given vector field with a neural hidden layer, and which captures the adaptive properties of feature learning. This defines a metric on the space of probability measures, complementing existing Wasserstein and Stein geometries. In this work we examine the properties of this metric in the context of generative modeling. As a first application, we quantify the suboptimality of diffusion generative modeling in the Barron geometry by establishing super-polynomial score approximation lower bounds for data generated by neural network pushforwards of the Gaussian. We then investigate the benefit of adaptivity as a way to study alternative generative models. In a companion paper \cite{companionpaper} we leverage the Barron transport geometry for sampling applications, extending the scope of Stein variational gradient methods via feature adaptation.
Figures & tables
Figure 1: Diagram of the trajectories ρtπ and ρ^t connecting γd to π . The ideal reverse OU path ρtπ is depicted in red, while the effective path obtained by regressing the target scores with deep ResNet models is depicted in purple. The paths are matched at regular intervals, while the ResNet layers are free to implement arbitrary dynamics in between. Theorem 1 suggests that this freedom is useful for parameter efficiency.
Figure 2: Marginal densities of 2-dimensional slices (first and last coordinates) of ρt , along with histograms showing which coordinates are in motion, for the ResNet and reverse OU transports from γd to the lower bound instance πθ constructed in Theorem 3 (with ω=tanh , ψ=Id , and d=14 for simplicity). Starting from t=1 and looking backward, we see that the OU process adjusts the xd coordinate immediately, whereas the single-index ODE waits until t≈0.5 to do so (since it must first adjust all other coordinates).
Figure 3: Densities of 2-dimensional slices (the u -coordinate and one perpendicular to u ) for the standard Gaussian γd , the measure νε constructed in Example 2 with a ReLU pushforward, and the measure νε constructed in Example 3 with a smooth tanh flow; we take d=6 and ε=4 for visualization purposes. The ReLU pushforward has a clear nonsmooth transition, while the tanh flow almost looks like a Gaussian mixture.
Figure 4: Plots of the number of steps required so that KL(ρn∥π)≤41KL(ρ0∥π) as a function of d for the constructions π∈{π(1),π(2)} listed above. For each fixed dimension, we sweep over fixed learning rates (and bandwidth/smoothness in the SVGD cases) and plot the best performance; triangles indicate a failure to reach the threshold KL in an allotted number of steps. The hidden subspace is of dimension ⌊log2(d)/4⌋ , and the single-index tilt is constructed using a=0.1 and ψ(z)=0.256cos(z)+1.32cos(2z) .
We introduce, to our knowledge, the first deep generative modeling framework for probability distributions continuously supported on compact metric graphs. Given source and target measures on a metric graph, our method embeds the graph into a smooth ambient space, solves an entropic Kantorovich problem via a neural semidual parameterization, and projects generated samples back onto the original graph. We study two embedded geometries: an extrinsic Euclidean realization and the intrinsic tropical Abel--Jacobi embedding into the Jacobian torus. In both cases, the resulting generator is graph-supported by construction. We prove that, in the joint limit of increasing neural expressivity, the learned generator converges weakly to a valid transport coupling between the original graph measures. Empirically, across a range of geometrically distinct graphs, our method matches or improves upon heuristic transport baselines based on discrete graph OT, while scaling more favorably. Finally, we demonstrate scalability on real-world urban mobility data by training our model on one million Uber pickup locations in Manhattan, New York City.
Alessandro Micheli, Yueqi Cao, Anthea Monod +1
Imperial College London London, UK · KTH Royal Institute of Technology Stockholm, Sweden · Statens Serum Institut Copenhagen, Denmark +1
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.
Pablo Moreno-Muñoz, Adrian Müller, Gergely Neu
Universitat Pompeu Fabra · Barcelona, Spain · ETH Zürich +2
These notes recapitulate the high level mathematical principles behind different techniques for generative modeling. I show the connections between optimal transport and standard techniques such as Schr{ö}dinger bridge and flow matching.