Learning Paths for Dynamic Measure Transport: A Control Perspective (2511.03797v1)

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.

• The paper introduces a control-theoretic framework to learn optimal paths for dynamic measure transport, mitigating artifacts in common geometric annealing paths.
• It employs Gaussian process PDE methods in RKHS to regularize both the drift field and tilting function, ensuring smoother and more efficient sampling.
• The learned interpolation significantly outperforms reference paths in multimodal scenarios by better aligning mass distribution and reducing error metrics.

Learning Paths for Dynamic Measure Transport: A Control Perspective

Introduction and Motivation

This paper addresses the problem of identifying optimal paths of probability measures for sampling via dynamic measure transport (DMT), with a particular focus on the control-theoretic perspective. DMT-based sampling methods, including neural ODEs, continuous normalizing flows, and diffusion models, rely on constructing a stochastic process that transforms samples from a reference distribution $\eta$ into samples from a target distribution $\pi$. The evolution of the process is governed by a drift (velocity) field $v$ and, optionally, a diffusion term. The choice of the path of intermediate measures $(\rho(t))_{t \in [0, T]}$ is critical for the tractability and efficiency of the sampling procedure.

The paper highlights that commonly used paths, such as the geometric annealing path $\mu(t) \propto \eta^{1-t}\pi^t$, may be suboptimal for certain $(\eta, \pi)$ pairs, leading to irregular or inefficient transport dynamics. The authors propose a flexible family of optimization problems for learning "tilted" paths of measures, advocating for explicit regularization terms that promote smoothness in the velocity field. The connection to mean-field games (MFGs) is established, providing a principled framework for path selection and regularization.

Issues with Canonical Paths and the Need for Learned Interpolations

The geometric annealing path is widely used due to its tractable log-derivative and connections to Fisher–Rao geometry. However, for multimodal targets or poorly aligned reference and target distributions, this path can induce undesirable transport phenomena, such as "teleportation of mass" between modes, which are difficult to capture with standard DMT algorithms. The paper demonstrates this with a mixture-of-Gaussians example, where the geometric path fails to allocate sufficient mass to the left mode of the target.

Figure 1: Geometric annealing path (top) and path resulting from solving the proposed control problem for a mixture-of-Gaussians example. Samples generated by the respective velocity fields are overlaid in red.

The authors propose to learn a tilting function $g(x, t)$ that perturbs the reference path, yielding a new path $\rho^g(x, t) \propto \mu(x, t) e^{g(x, t)}$. This approach is motivated by prior work using PINN-based optimization, but the paper provides a more principled control-theoretic formulation, connecting the problem to MFGs and explicit regularization.

Control-Theoretic Formulation and Regularization

The central contribution is the formulation of a control problem for path identification:

$$\inf_{v \in \mathcal{V}, g \in \mathcal{G}} \|v\|^2_{\mathcal{V}} + \lambda_g \|g\|^2_{\mathcal{G}} \quad \text{s.t.} \quad -\nabla \cdot (v \rho^g) = \rho^g (\partial_t \log \rho^g), \quad \rho^g \propto \rho^{\rm ref} e^{g}, \quad g(\cdot, 0) = g(\cdot, 1) \equiv 0$$

This framework allows for the use of general Banach or RKHS norms to regularize both the velocity and the tilting function, promoting smoothness and tractability. The constraints enforce exact matching of the terminal distribution, in contrast to MFGs or Schrödinger bridge approaches that use terminal costs. The flexibility of the regularization is argued to be crucial for practical DMT, especially in high-dimensional or multimodal settings.

Numerical Implementation via Gaussian Process PDEs

The paper presents a numerical algorithm for solving the control problem using Gaussian process (GP) methods for PDE-constrained optimization. Both the velocity potential $u$ (with $v = \nabla u$) and the tilting function $g$ are represented in RKHSs with separable kernels over space and time. The PDE constraints and boundary conditions are enforced at collocation points, and the optimization is performed via a penalized loss using Levenberg–Marquardt.

Figure 2: Space-time plots of the reference path $\mu(x, t)$, the tilting $e^{g(x, t)}$, and the learned path $\rho^g(x, t)$.

The learned path eliminates the teleportation artifact present in the geometric path, resulting in more efficient and smoother transport dynamics. The velocity field $v_g$ associated with the learned path is shown to be spatially smoother and better aligned with the target distribution.

Figure 3: Trajectories for DMT between $\eta$ and $\pi$ using the reference velocity, learned velocity, and McCann interpolant velocity.

Figure 4: Potentials $u_{\rm ref}$ and $u_g$ and their weighted versions, as well as the corresponding velocity fields.

Figure 5: Spatial RKHS norms of $u_g(\cdot, t)$ (blue) and $u_{\rm ref}(\cdot, t)$ (red) as a function of time.

The spatial RKHS norm of the reference velocity increases sharply during the teleportation phase, while the learned velocity maintains a relatively constant norm, indicating improved regularity.

Empirical Results and Performance Metrics

The paper provides quantitative metrics comparing the learned and reference interpolations:

Method	Fraction in Left Mode	Rel. Error Mean ↓	Rel. Error Var ↓	MMD ↓	$\|u\|_{\mathcal{H}}$
Reference Interpolation	0.005	1.80	0.96	0.743	770
Learned Interpolation	0.375	0.88	0.016	0.137	136
Ground Truth Samples	0.654	0.040	0.024	$7.21 \times 10^{-4}$	n/a

The learned interpolation dramatically improves sample quality, especially in allocating mass to the correct mode and reducing the RKHS norm of the velocity.

Theoretical and Practical Implications

The control-theoretic framework for path identification in DMT provides a principled approach to regularizing transport dynamics, with direct implications for the design of efficient samplers in Bayesian inference, generative modeling, and data assimilation. The explicit connection to MFGs and the flexibility in regularization norms enable adaptation to problem-specific requirements, such as spatial and temporal smoothness.

The numerical results suggest that learned paths can significantly outperform canonical choices, especially in challenging multimodal scenarios. The GP-PDE approach is scalable and can be extended to higher dimensions and more complex reference/target pairs.

Future Directions

The paper outlines several avenues for future research, including:

• Exploration of alternative regularization penalties (e.g., Bochner space norms) to further control spatial and temporal regularity.
• Extension to stochastic dynamics and SDE-based DMT.
• Application to high-dimensional problems and real-world Bayesian inference tasks.
• Investigation of the interplay between path regularity and sample efficiency in practical settings.

Conclusion

This work introduces a general control-based framework for learning paths of measures in dynamic measure transport, emphasizing the importance of regularization for smooth and efficient sampling. The proposed approach, grounded in mean-field game theory and implemented via GP-PDE methods, demonstrates clear advantages over standard geometric annealing paths, both theoretically and empirically. The framework is broadly applicable and sets the stage for further advances in principled sampler design and probabilistic inference.