Optimal Scheduling of Dynamic Transport

Flow-based methods for sampling and generative modeling use continuous-time dynamical systems to represent a transport map that pushes forward a source measure to a target measure. The introduction of a time axis provides considerable design freedom, and a central question is how to exploit this freedom. Though many popular methods seek straight line (i.e., zero acceleration) trajectories, we show here that a specific class of ``curved'' trajectories can significantly improve approximation and learning. In particular, we consider the unit-time interpolation of any given transport map T and seek the schedule τ: [0,1] → [0,1] that minimizes the spatial Lipschitz constant of the corresponding velocity field over all times t ∈ [0,1]. This quantity is crucial as it allows for control of the approximation error when the velocity field is learned from data. We show that, for a broad class of source/target measures and transport maps T, the optimal schedule can be computed in closed form, and that the resulting optimal Lipschitz constant is exponentially smaller than that induced by an identity schedule (corresponding to, for instance, the Wasserstein geodesic). Our proof technique relies on the calculus of variations and Gamma-convergence, allowing us to approximate the aforementioned degenerate objective by a family of smooth, tractable problems.