Learning Paths for Dynamic Measure Transport: A Control Perspective

TL;DR

This paper reframes dynamic measure transport as a controllable path selection problem, arguing that naive paths can hinder efficient sampling due to phenomena like mass teleportation. It introduces a joint tilt-and-velocity control framework that regularizes path smoothness via RKHS/Gaussian-process PDEs and relates tilting to mean-field game theory, offering a principled alternative to geometric annealing. Numerically, the authors implement a GP-PDE-based solver for a PDE-constrained optimization over tilt $g$ and drift $v$, demonstrating in 1D that learned tilted paths produce smoother dynamics and better mass allocation than untitled or McCann interpolants. The approach promises improved DMT samplers for Bayesian inference and data assimilation by enabling explicit, smooth transport paths tailored to the target distribution.

Abstract

We bring a control perspective to the problem of identifying paths of measures for sampling via dynamic measure transport (DMT). We highlight the fact that commonly used paths may be poor choices for DMT and connect existing methods for learning alternate paths to mean-field games. Based on these connections we pose a flexible family of optimization problems for identifying tilted paths of measures for DMT and advocate for the use of objective terms which encourage smoothness of the corresponding velocities. We present a numerical algorithm for solving these problems based on recent Gaussian process methods for solution of partial differential equations and demonstrate the ability of our method to recover more efficient and smooth transport models compared to those which use an untilted reference path.

Figures (5)

• Figure 1: Geometric annealing path (top) and path resulting from solving our proposed control problem \ref{['eq:relaxed_problem']} (bottom) for the example $\eta = \mathcal{N}(0, 1)$ and $\pi = \frac{2}{3}\mathcal{N}(-8, 1) + \frac{1}{3}\mathcal{N}(4, 1)$. Samples generated by the respective velocity fields are plotted overtop in red.
• Figure 2: Space-time plots of the reference path $\mu(x, t) \propto \eta(x)^{1-t}\rho(x)^t$ (left), the tilting $e^{g(x, t)}$ (center), and the path $\rho^g(x, t) \propto \mu(x, t)e^{g(x, t)}$ resulting from \ref{['eq:relaxed_problem']} (right).
• Figure 3: Trajectories corresponding to three different velocity fields for DMT between $\eta$ and $\pi$: the reference velocity $v_{\rm ref} = \nabla u_{\rm ref}$ (left), the learned velocity $v_g = \nabla u_{g}$ (center), and the McCann interpolant velocity (right). The learned velocity $v_{g}$ places more mass in the left mode than the reference velocity $v_{\rm ref}$ and is spatially smoother than the McCann interpolant velocity.
• Figure 4: Potentials $u_{\rm ref}$ and $u_g$ and velocity fields $v_{\rm ref} = \nabla u_{\rm ref}$ and $v_{g} = \nabla u_g$ corresponding to the geometric path $\rho^{\rm ref} = \mu$ and the path $\rho^g$ obtained from \ref{['eq:relaxed_problem']}. In the first two columns of panels we show the absolute potentials/velocities, and in the second two columns we show the potentials/velocities weighted by their respective probability densities, which better capture how the mass is moving.
• Figure 5: Spatial RKHS norms of $u_g(\cdot, t)$ (blue) and $u_{\rm ref}(\cdot, t)$ (red) as a function of time.