A McKean-Pontrygin maximum principle for entropic-regularized optimal transport

Abstract: This note outlines a mean-field approach to dynamic optimal transport problems based on the recently proposed McKean-Pontryagin maximum principle. Key aspects of the proposed methodology include i) avoidance of sampling over stochastic paths, ii) a fully variational approach leading to constrained Hamiltonian equations of motion, and iii) a unified treatment of deterministic and stochastic optimal transport problems. We also discuss connections to well-known dynamic formulations in terms of forward-backward stochastic differential equations and extensions beyond classical entropic-regularized transport problems.

• The paper formulates a mean-field variational principle that extends the McKean–Pontryagin maximum principle to model entropic-regularized optimal transport.
• It derives coupled Hamiltonian ODEs over probability measures to obtain optimal control laws for both deterministic and stochastic regularized transport problems.
• The framework bridges classical optimal transport with the Schrödinger bridge, offering numerically robust, ODE-based methods for high-dimensional stochastic control.

McKean–Pontryagin Maximum Principle for Entropic-Regularized Optimal Transport

Introduction

This work develops a mean-field, variational formulation for dynamic optimal transport, specifically addressing entropic-regularized scenarios such as the Schrödinger Bridge problem. The approach is rooted in an extension of the McKean–Pontryagin maximum principle, leading to fully coupled constrained Hamiltonian dynamics in the space of probability measures. The framework unifies deterministic optimal transport, Schrödinger bridge (stochastic regularized) problems, and more general stochastic and mean-field transport models, while maintaining structural connections to the classical Pontryagin principle and Hamilton–Jacobi–Bellman (HJB) theory.

Problem Formulation

The central focus is minimizing a quadratic control cost under the constraint that a controlled Itô diffusion steers an initial law $\pi_0$ to a target law $\pi_T$. The control variable is $U_t$, and the system dynamics are given by the SDE

$${\rm d}\tilde X_t = U_t {\rm d}t + \sqrt{2} \Sigma^{1/2} {\rm d}\tilde B_t,$$

where $\Sigma$ is a (possibly state-dependent) positive definite diffusion matrix. For $\Sigma = 0$, the problem reduces to the deterministic Benamou–Brenier formulation of optimal transport, while for $\Sigma = R = I$, one recovers the Schrödinger bridge. The cost functional is

$$\mathcal{J}(U) = \frac{1}{2} \int_0^T \mathbb{E}_{\tilde X_t}\|U_t \|^2_R \,{\rm d}t,$$

where the weighted norm is induced by $R$.

Crucially, the framework is flexible to accommodate both additive and multiplicative noise, degenerate and full rank diffusions, underdamped Langevin dynamics, and law-dependent coefficients (including mean-field dissipative flows).

McKean–Pontryagin Variational Principle

The main technical advance is the formulation of the stochastic optimal control problem using an adapted McKean–Pontryagin maximum principle in the space of probability measures. The states $X_t$ and co-states $\pi_T$0 evolve according to Hamiltonian ODEs derived from a global action functional, incorporating Lagrange multipliers for the density evolution and terminal constraints.

The action involves boundary potentials $\pi_T$1, $\pi_T$2 and a coupling variable $\pi_T$3, reflecting gauge freedom from the relabeling symmetry intrinsic to such mean-field functionals. Importantly, the constraints ensure that only the laws (not the trajectories) of the mean-field states are prescribed, which is crucial for computational tractability.

Taking variations leads to coupled evolution equations for $\pi_T$4:

• The state $\pi_T$5 evolves according to the controlled drift plus a gauge term,
• The co-state $\pi_T$6 evolves backwards, involving both the Hessian of the value function and the coupling to the density,
• The multipliers enforce that $\pi_T$7 and that the mean-field flow satisfies a continuity equation with additional divergence terms,
• The optimal control law emerges as $\pi_T$8.

Consistency of this framework with classical stochastic control is established via the HJB equation for $\pi_T$9 and the Fokker–Planck equation for $U_t$0.

Connections and Comparisons to Existing Methods

• Schrödinger bridge and forward-backward SDEs: By appropriate gauge fixing (randomizing the multiplier), the mean-field Hamiltonian formulation can be linked to standard FBSDE representations of the Schrödinger bridge, giving a unified perspective bridging variational, stochastic, and mean-field methods.
• Non-variational Lagrangian approach: The work makes contact with independently proposed Lagrangian HJB-based schemes, highlighting that these lack variational structure but still yield the correct control law and density evolution.
• Adjoint-matching, normalizing flows, and fluid dynamics: The framework generalizes to adjoint-matching methods and can be regularized or extended to include energy-like penalties and non-gradient drifts, encompassing a range of flows including normalizing flows, gradient flows, and fluid-dynamical models.

Numerical and Theoretical Implications

A significant feature of the mean-field McKean–Pontryagin approach is that it circumvents the necessity of sampling stochastic path space and works with deterministic coupled ODEs on the law (density) level. This is anticipated to yield numerically less noisy approximations compared to particle-based forward-backward SDE or path-space variational approaches, albeit at the cost of more involved optimization problems over functions rather than finite-dimensional variables.

By allowing the entropy regularization parameter ($U_t$1) to vanish, the method naturally recovers deterministic optimal transport, while regularization leads to smoother potential functions and interpolating flows.

Extensions to law-dependent drift and diffusion, as relevant for filtering (Kalman–Wasserstein flows) or mean-field games, are noted as direct applications. Treating the diffusion matrix $U_t$2 as a control variable is also discussed, leading to smoothness-promoting penalties on the transport potential.

Prospects for Future Research

• Generalization to interacting particle systems: Discretizations based on labels and time (e.g., Runge–Kutta methods), as well as connections to optimization over particle representations, promise tractable approximations for large-scale problems.
• Comparative numerical studies: Direct comparison of the mean-field variational approach with FBSDE and non-variational Lagrangian methods, especially regarding noise properties and convergence rates, is proposed as future work.
• Extensions to control of more complex diffusions: The framework is equipped to handle law-dependent drifts and diffusions, offering a natural approach to mean-field control and related inverse problems.

Conclusion

This work provides a rigorous, functional-analytical formulation of entropic-regularized optimal transport via the McKean–Pontryagin maximum principle. The resulting mean-field variational principle unifies previous stochastic control, Schrödinger bridge, and Benamou–Brenier approaches, admits a range of natural extensions (including to normalizing flows and fluid mechanics), and enables the development of numerically robust ODE-based methods for optimal transport under noise regularization. These advances have broad implications for computational optimal transport, stochastic control, and statistical inference, opening avenues for further research in mean-field and law-invariant control methodologies.