Dynamic Optimal Transport

• Dynamic optimal transport is a framework for time-dependent mass transfer that models evolving densities with minimal kinetic energy and strict mass conservation.
• It employs staggered grid discretization and a convex variational formulation to compute Wasserstein geodesics and displacement interpolations efficiently.
• Proximal splitting methods extend the approach to non-Euclidean metrics and weighted cost functions, enabling scalable solutions in imaging, geophysics, and beyond.

Dynamic optimal transport (dynamic OT) generalizes the classical optimal transport problem to time-dependent mass transfer, modeling the continuous evolution of a density as mass moves with minimal kinetic energy, subject to mass conservation. This framework, rooted in the Benamou–Brenier formulation, allows for the characterization of Wasserstein geodesics and provides a convex variational approach for computing displacement interpolations between probability distributions in both Euclidean and non-Euclidean domains. Dynamic OT underpins numerous theoretical and applied developments in fields ranging from imaging to traffic modeling and high-dimensional generative modeling.

1. Eulerian Dynamic OT Formulation and Staggered Grid Discretization

The Benamou–Brenier dynamic OT model expresses the squared L²-Wasserstein distance between densities $f^0$ and $f^1$ on $[0,1]^d$ as a convex optimization over a time-dependent density $f(x, t)$ and momentum $m(x, t)$: $$\min_{f, m} \int_0^1 \int_{[0,1]^d} \frac{\|m(x, t)\|^2}{2 f(x, t)} \, dx \, dt$$ subject to the continuity equation (mass conservation): $$\partial_t f + \nabla_x \cdot m = 0, \quad f(\cdot, 0) = f^0, \quad f(\cdot, 1) = f^1.$$

To ensure accurate numerical enforcement of the divergence constraint and boundary conditions, a staggered grid discretization is introduced. The fundamental idea is to discretize the primary variables $(m, f)$ on grids offset in both space and time:

• Centered grid: $(x_i, t_j) = (i/N, j/P)$ records the main variables for the entire domain.
• Spatial (staggered) grid: Offsets the $x$-coordinate by $1/2N$ to align flux evaluations at faces between cells.
• Temporal (staggered) grid: Offsets the $t$-coordinate by $1/2P$ to align density at midsteps.

The variables are expressed as $U = (\bar m, \bar f)$ on the staggered grid. Two key linear operators structure the discrete problem:

• Interpolation: Maps staggered variables to the centered grid,

$$m_{i, j} = \frac{\bar m_{i-1, j} + \bar m_{i, j}}{2}, \qquad f_{i, j} = \frac{\bar f_{i, j-1} + \bar f_{i, j}}{2}$$

• Divergence operator:

$$(\text{div}\,U)_{i,j} = N(\bar m_{i, j} - \bar m_{i-1, j}) + P(\bar f_{i, j} - \bar f_{i, j-1})$$

Boundary conditions and divergence constraints define the convex set $\mathcal{C}_{inc}$ in which the discrete solution must reside.

2. Proximal Splitting Schemes for Large-Scale Convex Optimization

The fully discretized dynamic OT problem combines smooth and nonsmooth terms, with the nonsmooth indicator enforcing the divergence and boundary constraints. The objective is written as: $$\min_{U\in\mathcal{C}_{inc}} \;J(\text{interp}(U)),$$ where $J$ is the sum over grid cells of $j(m, f) = \frac{\|m\|^2}{2f}$ if $f > 0$ (and $+\infty$ otherwise).

This structure is amenable to first-order convex optimization via proximal splitting:

• Proximal operator for $J$: Decomposes across grid cells. The update for a single cell with data $(\tilde m, \tilde f)$ is:

$$\mathrm{prox}_{\gamma j}(\tilde m, \tilde f) = \begin{cases} (\mu(f^*), f^*) & \text{if } f^* > 0,\ (0,0) & \text{otherwise} \end{cases}$$

with $\mu(f) = \frac{f\tilde m}{f+\gamma}$ and $f^*$ solves the cubic $$P(X) = (X-\tilde f)(X + \gamma)^2 - (\gamma/2)\|\tilde m\|^2 = 0$$.

• Proximal operator for $\iota_{\mathcal{C}_{inc}}$: The Euclidean projection onto the divergence/boundary constraint set, computed by solving a linear system (typically via an FFT-based Poisson solver).
• Splitting strategies: The problem is equivalently recast by introducing auxiliary coupling constraints (e.g., $V = \text{interp}(U)$). Proximal splitting can then be performed via, for example, a Douglas–Rachford (DR) scheme:
  • Alternate prox steps on $G_1$ (cost and incompressibility) and $G_2$ (coupling constraint),
  • Each update only requires an efficient projection and local nonlinear solve per cell.
• Primal–dual methods: Alternative splitting techniques (e.g., Chambolle–Pock) update both primal (grid variable) and dual (multiplier) variables, with steps satisfying $\sigma\tau\|\mathrm{interp}\|^2 < 1$.

3. Generalizations: Non-Euclidean and Weighted Cost Metrics

The convex framework is highly extensible. The local cost function is generalized from $j(m, f)$ to the family

$$j_\beta(m, f) = \begin{cases} \frac{\|m\|^2}{2f^\beta} & f > 0,\ 0 & (m, f) = (0, 0),\ +\infty & \text{otherwise} \end{cases}$$

for $\beta\in[0,1]$, interpolating between Wasserstein ($\beta=1$) and $H^{-1}$ ($\beta=0$) metrics. The discrete cost becomes $J_\beta^w(V) = \sum_k w_k j_\beta(m_k, f_k)$ for spatially varying weights $w_k>0$.

The corresponding cellwise proximal operator generalizes:

• Solution for $m$ is $m^*(f) = \frac{f^\beta \tilde m}{f^\beta+\gamma}$.
• $f^*$ is the unique positive root of

$$P_\beta(X) = X^{1-\beta}(X - \tilde f)(X^\beta+\gamma)^2 - \frac{\gamma}{2}\beta \|\tilde m\|^2$$

Spatially or temporally varying weights $w_k$ encode manifold geometry or obstacles (e.g., $w_k=+\infty$ in forbidden regions), thereby enabling the same proximal splitting machinery to address general Riemannian or constrained settings.

4. Numerical Implementation and Algorithmic Aspects

Efficient implementation of the scheme leverages the structure of the discrete problem:

• Cellwise proximal computations are independent and can be fully parallelized.
• The projection onto the constraint is a global linear operation with fast solvers available via FFT on uniform grids.
• The splitting approach enables decomposition into computationally tractable subproblems, each amenable to large-scale parallel and distributed computing.

Trade-offs include:

• Staggered grids are preferred for accurate discretization of divergence and mass conservation over centered grids, though the latter recovers earlier algorithms.
• Solving the cubic (or higher-degree) equations per cell is required at each iteration, but these are explicit and numerically stable.
• The global projection step may become the computational bottleneck for irregular domains or non-periodic geometries, but can be addressed by alternative linear solvers or preconditioners.

5. Applications, Extensions, and Flexibility of the Framework

The proximal splitting framework for dynamic OT is not restricted to $L^2$ OT but generalizes naturally to:

• Image interpolation and time-varying signal processing,
• Shape interpolation with obstacles or manifold constraints,
• Geophysical flows where domains include obstacles or variable terrain,
• Riemannian manifold interpolation by embedding geometry into the local weights $w_k$,
• Unbalanced and generalized OT formulations with suitable modification of the cost functional or constraint set.

The modularity of the approach allows for rapid adaptation to new cost structures, regularizations, and domain constraints, while maintaining scalability due to the cellwise decoupling and efficient projection.

6. Key Mathematical Formulas and Operators

Component	Mathematical Expression	Role
Centered grid	$\mathcal{G}_c = \{ (x_i = i/N, t_j = j/P) \}$	Main variables discretization
Staggered grid interpolation	$m_{i,j} = (\bar m_{i-1,j} + \bar m_{i,j})/2$, $f_{i,j} = (\bar f_{i,j-1} + \bar f_{i,j})/2$	Accurate divergence implementation
Divergence operator	$$(\mathrm{div}\,U)_{i,j} = N(\bar m_{i,j} - \bar m_{i-1,j}) + P(\bar f_{i,j} - \bar f_{i,j-1})$$	Mass conservation, incompressibility constraint
Proximal step for $j$	$\mathrm{prox}_{\gamma j}(\tilde m, \tilde f)$ with cubic $P(X)$ (see above)	Cellwise point update for the kinetic term
Generalized cost	$j_\beta(m,f) = \|m\|^2/(2f^\beta)$ (if $f>0$)	Riemannian/weighted OT, $H^{-1}$ metric
Constraint set	$$\mathcal{C}_{inc} = \{ U : \mathrm{div}\,U = 0, b(U) = b_0 \}$$	Linear divergence and boundary conditions

7. Summary and Impact

The proximal splitting approach to dynamic optimal transport is a unifying, flexible framework for efficiently solving large-scale, discretized OT problems in both Euclidean and Riemannian contexts. By combining a staggered grid discretization that respects mass conservation, explicit computation of cellwise proximal operators, and global projection via linear solvers, the scheme obtains accurate geodesic interpolations between distributions and accommodates a spectrum of cost structures and domain constraints. This methodology underlies scalable algorithms applicable to imaging, geometry processing, geophysical flows, and generic mass transport problems, and serves as a template for further generalization in modern applications.