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

1. Problem Statement

The paper addresses the Dynamic Optimal Transport (OT) problem, specifically focusing on bridging two probability distributions, $\pi_0$ and $\pi_T$, over a time horizon $[0, T]$. The system evolves according to a controlled stochastic differential equation (SDE):
$$d\tilde{X}_t = U_t dt + \sqrt{2}\Sigma^{1/2} d\tilde{B}_t$$
where:

• $U_t$ is the control input.
• $\Sigma$ is the diffusion matrix.
• $\tilde{B}_t$ is a Brownian motion.

The objective is to minimize the expected control cost:
$$J(U) = \frac{1}{2} \int_0^T \mathbb{E}[\|U_t\|_R^2] dt$$
subject to the marginal constraints $\tilde{X}_0 \sim \pi_0$ and $\tilde{X}_T \sim \pi_T$.

Key Special Cases:

• Classical OT: $\Sigma = 0$ and $R = I$.
• Schrödinger Bridge Problem (SBP): $\Sigma = R = I$ (entropic regularization).
• Generalizations: The framework extends to underdamped Langevin dynamics and mean-field formulations like the Kalman–Wasserstein gradient flow.

2. Methodology: The McKean–Pontryagin Variational Approach

The author proposes a mean-field formulation based on an extension of the classical Pontryagin Maximum Principle (PMP), termed the McKean–Pontryagin Maximum Principle.

Core Concept

Instead of sampling stochastic paths directly, the method introduces a variational formulation using:

• States: $X_t(a) \in \mathbb{R}^{d_x}$, where $a$ represents a fixed "label" (initial condition) acting as an independent variable.
• Co-states: $P_t(a) \in \mathbb{R}^{d_x}$.
• Distribution: The law of $X_t(a)$ is denoted by $\rho_t$.
• Symmetry: The formulation accounts for "relabeling symmetry," ensuring the marginal densities match $\pi_0$ and $\pi_T$ regardless of the specific labeling of particles.

The Action Functional

The authors define an action functional $S$ involving states, co-states, controls, and auxiliary potentials ($\psi, \beta$):
$$S = \int_0^T \int_{\mathbb{R}^{dx}} \left( \langle P_t, \dot{X}_t - U_t \rangle - \Sigma : D^2_x \psi_t(X_t) + \frac{1}{2}\|U_t\|_R^2 \right) \rho_0(a) da dt + \text{Boundary Terms}$$
This leads to a Hamiltonian $H$ that incorporates a state-dependent running cost term $\Sigma : D^2_x \psi(X)$, which arises implicitly from the relation $P_t = \nabla_x \psi_t(X_t)$.

Derivation of Equations of Motion

By taking variations of the action functional with respect to $X_t, P_t, U_t, \beta_t,$ and $\psi_t$, the authors derive a system of constrained Hamiltonian equations:

1. State Evolution: $\dot{X}_t = U_t + \beta_t$
2. Co-state Evolution: $-\dot{P}_t = \nabla_x(\Sigma : D^2_x \psi_t(X_t)) - D^2_x \psi_t(X_t) \beta_t$
3. Constraints:
   • $P_t = \nabla_x \psi_t(X_t)$ (Linking co-state to potential)
   • $U_t = R \nabla_x \psi_t(X_t)$ (Optimal control law)
   • Continuity equation for $\rho_t$ involving $\beta_t$.

Key Insight on $\beta_t$: The auxiliary variable $\beta_t$ is not uniquely determined by the variational principle due to relabeling symmetry. The authors show that $\beta_t$ can be chosen arbitrarily (e.g., $\beta_t = -\Sigma \nabla_x \log \rho_t$) without affecting the optimal control law $u_t(x) = R \nabla_x \psi_t(x)$.

3. Key Contributions

1. Unified Framework: The paper provides a unified treatment for both deterministic ($\Sigma=0$) and stochastic ($\Sigma \neq 0$) optimal transport problems under a single mean-field variational principle.
2. Avoidance of Path Sampling: Unlike standard stochastic control approaches that rely on Forward-Backward Stochastic Differential Equations (FBSDEs) and Monte Carlo sampling, this approach formulates the problem as a system of ordinary differential equations (ODEs) in the mean-field limit. This avoids the noise inherent in stochastic path sampling.
3. Recovery of Known Results:
   • In the limit $\Sigma \to 0$, the formulation recovers the Benamou–Brenier dynamic OT formulation.
   • For $\Sigma = R = I$, it provides a deterministic mean-field formulation of the Schrödinger Bridge problem, recovering the Fokker–Planck and Hamilton–Jacobi–Bellman (HJB) equations.
4. Connection to FBSDEs: The paper demonstrates that by formally setting $\beta_t$ as a stochastic noise term (Stratonovich or Itô interpretation), the derived ODEs map directly to the standard FBSDE formulations used in stochastic optimal control, validating the consistency of the approach.
5. Extension to Control of Diffusion: The framework is extended to treat the diffusion matrix $\Sigma$ itself as a control variable, introducing a smoothness penalty on the potential $\psi_t$.

4. Results and Theoretical Findings

• Optimal Control Law: The optimal control is explicitly given by $u_t(x) = R \nabla_x \psi_t(x)$, where $\psi_t$ satisfies the HJB equation:
  $$-\partial_t \psi_t = \langle \nabla_x \psi_t, u_t \rangle + \Sigma : D^2_x \psi_t - \frac{1}{2}\|u_t\|_R^2$$
• Conservation of Hamiltonian: The Hamiltonian is shown to be conserved along the solutions, taking a constant value determined by the cost functional.
• Iterative Algorithm: An example is provided for transforming a distribution to a standard Gaussian. The paper outlines an iterative procedure (similar to Iterative Proportional Fitting) to update the boundary potential $\psi_T$ based on the mismatch between the current terminal distribution and the target $\pi_T$.
• Numerical Implications: The mean-field ODE formulation (12)-(13) is posited to yield less noisy approximations compared to stochastic counterparts (FBSDEs), as it does not require sampling Brownian paths.

5. Significance

This work is significant for the fields of Optimal Transport, Stochastic Control, and Machine Learning for several reasons:

• Computational Efficiency: By shifting from stochastic path sampling to a deterministic mean-field ODE system, the approach potentially offers more stable and efficient numerical algorithms for high-dimensional transport problems.
• Theoretical Unification: It bridges the gap between classical deterministic OT (Benamou–Brenier) and entropic-regularized stochastic OT (Schrödinger Bridge) through a single variational principle.
• Flexibility: The framework is highly extensible, accommodating complex dynamics (e.g., underdamped Langevin), state-dependent noise, and even the control of diffusion coefficients.
• New Perspective on SBP: It offers a fresh "McKean" perspective on the Schrödinger Bridge problem, treating the problem as a constrained variational problem on the space of labeled particles rather than just a probabilistic coupling problem.

In summary, Sebastian Reich presents a rigorous variational framework that generalizes the Pontryagin Maximum Principle to mean-field settings, offering a powerful alternative to FBSDE-based methods for solving entropic-regularized optimal transport problems.