Optimal Transport with Proximal Splitting

Nicolas Papadakis; Gabriel Peyré; Edouard Oudet

Source-linked AI summary

Optimal Transport with Proximal Splitting

Nicolas Papadakis, Gabriel Peyré, Edouard Oudet

arXiv:1304.5784v2Apr 21, 2013math.NA

TL;DR

The paper addresses scalable computation of discretized dynamic optimal transport, whose Eulerian formulation becomes a large convex optimization problem. It develops staggered-grid and proximal-splitting formulations, identifies the Benamou–Brenier method as a Douglas–Rachford instance, and extends the framework to generalized costs and Riemannian manifolds. Experiments include manifold transport, while algorithm comparisons are reported as specific-example results.

Problem
Discretized dynamic optimal transport requires solving a large convex problem, motivating methods suited to uniform-grid and imaging-scale computations.
Method
The paper combines staggered-grid Eulerian discretization with proximal splitting, including Douglas–Rachford and primal-dual formulations, generalized costs, and manifold extensions.
Results
The framework computes displacement interpolation on a 50×50×100 space-time grid for Gaussian inputs in a 2-D labyrinth, closely following the geodesic between their centers.
Takeaways & Limitations
The method provides a unified proximal-splitting framework that extends the original Benamou–Brenier algorithm to staggered grids, spatially varying costs, and Riemannian transport.
Takeaways & Limitations
The formulation imposes homogeneous Neumann velocity boundaries and fixed initial and final densities, with boundary handling dependent on the chosen conditions.

Abstract

from arXiv · show

This article reviews the use of first order convex optimization schemes to solve the discretized dynamic optimal transport problem, initially proposed by Benamou and Brenier. We develop a staggered grid discretization that is well adapted to the computation of the $L^2$ optimal transport geodesic between distributions defined on a uniform spatial grid. We show how proximal splitting schemes can be used to solve the resulting large scale convex optimization problem. A specific instantiation of this method on a centered grid corresponds to the initial algorithm developed by Benamou and Brenier. We also show how more general cost functions can be taken into account and how to extend the method to perform optimal transport on a Riemannian manifold.

1. Introduction.

The introduction situates computational optimal transport across discrete, PDE, dynamical, generalized-cost, and manifold formulations, motivating scalable first-order methods. The paper builds on convex reformulations and proximal splitting to extend dynamical transport computation.

Computational optimal transport has applications in computer vision, image processing, color transfer, shape retrieval, surface reconstruction, and computer graphics.
Discrete optimal transport.: Discrete transport commonly represents distributions as weighted Diracs and uses linear, network, auction, or assignment solvers, but finely discretized problems remain numerically intensive.
Dynamical optimal transport.: Dynamical optimal transport computes Wasserstein geodesics either by transporting discrete masses or by an Eulerian formulation on a fixed spatial grid.
Dynamical optimal transport.: The Benamou–Brenier change of variables converts the Eulerian problem into a convex optimization problem that can be solved iteratively.
First order and proximal methods.: Proximal splitting methods suit large-scale uniform-grid problems because they optimize sums of functionals with efficiently computable proximal operators.
The paper identifies the Benamou–Brenier method as a Douglas–Rachford instance and introduces staggered grids, generalized costs, and Riemannian-manifold extensions.

2. Dynamical Optimal Transport Formulation.

The paper formulates Wasserstein displacement interpolation through transport maps and then derives a convex fluid-mechanics problem in density and momentum. This formulation enforces mass conservation and endpoint densities while supporting recovery of the transport map from the computed flow.

2. Dynamical Optimal Transport Formulation.: A smooth transport map T pushes f 0(x)dx onto f 1(x)dx, subject to the density constraint induced by its differential.
2.1. Optimal Transport.: An optimal transport minimizes the assignment cost C(x,y), with squared Euclidean cost yielding the L2-Wasserstein distance.
2.1. Optimal Transport.: For the L2 ground cost, the geodesic density is generated by interpolating between the identity map and the optimal transport map.
2.2. Fluid Mechanics Formulation.: Benamou–Brenier’s fluid formulation minimizes kinetic energy over density f and velocity v subject to the continuity equation and endpoint conditions.
2.2. Fluid Mechanics Formulation.: The formulation uses homogeneous Neumann velocity boundaries and fixes the initial and final densities to f 0 and f 1.
2.2. Fluid Mechanics Formulation.: Existence analysis naturally relaxes the problem to Radon measures, where incompressibility is interpreted distributionally.
2.2. Fluid Mechanics Formulation.: After computing the optimal velocity field, integrating its flow in time recovers the transport map, with T obtained at t=1.
2.2. Fluid Mechanics Formulation.: The change of variables m=fv makes the optimization convex, with a perspective cost J(m,f) and linear constraints.

3. Discretized Dynamic Optimal Transport.

The discretization places density and momentum variables on a staggered space-time grid so divergence can be evaluated accurately, then expresses the problem as a constrained convex program. Proximal splitting targets this large-scale formulation.

The centered grid uses N+1 spatial and P+1 temporal points, with x_i=i/N and t_j=j/P.
3.2. Staggered Grid.: A staggered grid places vector-field components on different grids, enabling standard centered finite differences for each partial derivative and accurate divergence evaluation.
The staggered variables are connected to centered variables through a midpoint interpolation operator.
The space-time divergence is computed from momentum differences in space and density-flux differences in time, scaled by N and P.
Boundary extraction and imposed values encode zero spatial momentum at the boundary together with the discretized initial and final densities.
The finite-dimensional convex problem constrains variables to a set combining divergence-free flow and prescribed boundary values.
An equivalent formulation introduces auxiliary variables and rotated Lorentz-cone constraints, producing a second-order cone program.
The paper focuses on proximal splitting because it is more adapted to large-scale imaging problems than generic interior-point methods.

4. Proximal Splitting Algorithms.

The paper formulates discretized dynamic optimal transport as a convex problem addressed by proximal splitting. Different Douglas-Rachford formulations exploit simple proximal operators, staggered-grid constraints, and equivalences with the Benamou-Brenier algorithm.

Proximal operators and splitting: Proximal splitting replaces gradient steps with proximal operators for minimizing sums of simple convex functionals.The proximal operator is an implicit descent step, and Moreau’s identity transfers proximal computation between a functional and its Legendre-Fenchel transform.
Proximal operators and splitting: The constraint functional J has a closed-form proximal operator obtained from a positive root of a third-order polynomial, or (0, 0) otherwise.When a strictly positive root exists, it is unique and determines the proximal pair through m = µ(f).
Douglas-Rachford formulations: For 0 < α < 2 and γ > 0, Douglas-Rachford iterates converge to a solution of the split problem.The centered-grid projection associated with interpolation reduces to separable small linear systems whose inverse matrices can be precomputed.
Douglas-Rachford formulations: Four splitting algorithms map the optimal transport problem into a Douglas-Rachford formulation using staggered and centered variables, interpolation, and auxiliary variables.The formulations include asymmetric, exchanged asymmetric, symmetric, and related variants built from simple functionals and constraint sets.
Douglas-Rachford formulations: The A-DR iterates remain in the constraint set C, whereas A-DR’ iterates satisfy V^(ℓ) = I U^(ℓ) but generally violate C.For A-DR, the transported-geodesic approximation is retrieved from U^(ℓ) = (m^(ℓ), f^(ℓ)).
Relation to ALG2: On a centered grid, ALG2 produces the same iterates as Douglas-Rachford, while staggered grids require interpolation and an auxiliary variable V.An additional relaxation parameter α is introduced to accelerate convergence, and several Douglas-Rachford variants remain available.

5. Generalized Cost Functions.

The paper generalizes the transportation objective through exponent and spatial-weight parameters. These choices interpolate between Wasserstein and H^-1 geodesics and support Riemannian-manifold transport.

Generalized objective: The generalized functional uses spatial weights w_k satisfying c < w_k ≤ +∞, with infinite weights forbidding mass transport into selected cells.The parameter c is a positive small constant.
Geometric interpretations: When β = 1 with constant weights, the formulation corresponds to the L2-Wasserstein distance and discretizes displacement interpolation on a Riemannian manifold.The stated manifold interpretation uses squared geodesic ground cost and initially restricts attention to isotropic metrics.

6. Numerical Simulations.

The simulations compare proximal splitting algorithms on transport and demonstrate generalized interpolations, vanishing-density handling, and static or dynamic labyrinth constraints. They also illustrate applications with spatially varying costs, including an oceanographic setting.

6.1. Comparison of Proximal Schemes.: The experiments compare Douglas-Rachford, dual ADDM, and primal-dual algorithms, including asymmetric and symmetric Douglas-Rachford variants.The Gaussian experiment uses a 32×32×32 space-time discretization and an almost exact reference solution.
6.1. Comparison of Proximal Schemes.: The cost function reaches an almost minimal value after roughly 10^3 iterations, while the ℓ2 distance to the reference continues decreasing nearly linearly on a log-log scale.The cost alone therefore does not readily indicate approximation quality because the functional is very flat.
6.1. Comparison of Proximal Schemes.: A-DR and A-PD converge fastest to the reference solution, while S-DR performs as well as ADMM; PD iterations cost 0.13s versus 0.2s for DR or ADMM.The authors caution that these convergence results concern the specific Gaussian example but illustrate general algorithmic behavior.
6.1. Comparison of Proximal Schemes.: The PD method handles transport between characteristic functions whose geodesic is not composed of characteristic functions, demonstrating operation with vanishing densities.The experiment uses two characteristic functions of connected sets, one convex.
6.2. Interpolation Between L2-Wasserstein and H−1.: Varying β produces interpolations between linear density interpolation at β = 0 and displacement interpolation with transport at β = 1.The level-set experiment evaluates the estimated densities after 1000 PD iterations.
6.2. Interpolation Between L2-Wasserstein and H−1.: For β ∈]0, 1[, the functional converges faster than for β = 1, consistently with stronger convexity as β approaches 1/2; oscillations arise from approximate proximal computations.Figure 6.5 compares β = 1/2 and β = 3/4 using cost evolution and distance to a reference obtained after 10^5 PD iterations.
6.3. Riemannian Transportation.: On a static labyrinth, the displacement interpolation lies close to the geodesic path between Gaussian centers when the standard deviation is 0.04.The static example uses a 50×50×100 space-time discretization and constrains densities with labyrinth walls.
6.3. Riemannian Transportation.: Moving labyrinth walls alter the interpolation: when density is confined between t = 1/4 and t = 1/2, the mean’s propagation speed is no longer constant.A wall appears at t = 1/4 and another disappears at t = 1/2.