KL-Proximal Splitting Scheme

• KL-Proximal Splitting is a family of convex optimization methods that decompose large-scale dynamic optimal transport problems into manageable subproblems.
• It employs advanced discretization such as staggered grids and leverages efficient proximal splitting techniques including Douglas–Rachford, ADMM, and primal–dual methods.
• The scheme extends to generalized cost functions and Riemannian geometries, delivering faster convergence and scalability in applications like imaging and physics.

The KL-Proximal Splitting Scheme is a family of first-order convex optimization algorithms tailored for efficiently solving large-scale discretized dynamic optimal transport (OT) problems. These approaches are formulated to minimize convex composite objectives by exploiting the structure of the discretized Benamou–Brenier formulation of OT and extend to generalized cost functions and Riemannian geometries. Proximal splitting leverages efficient computation of proximal operators and enables decomposition of complex problems into manageable subproblems, supporting scalability and flexibility in modeling.

1. Proximal Splitting Methodologies in Optimal Transport

Proximal splitting algorithms solve problems of the form $\min_{x} F(x) + G(x)$ by alternating between the computation of explicit proximal updates for each term. In the context of discretized dynamic OT, the composite functional encodes both the transport cost—often a convex, separable function—and the linear dynamical and boundary constraints that mass must obey.

Principal Algorithms

• Douglas–Rachford (DR) Splitting: Iteratively alternates between the proximal operators of two (or more) key components of the objective, allowing for separate enforcement of cost and constraints with tunable relaxation parameters.
• Alternating Direction Method of Multipliers (ADMM): Shown to be equivalent to DR in these settings and recovers classical methods (e.g., the Benamou–Brenier ALG2 algorithm) as special cases on centered grids.
• Primal–Dual (PD) Algorithms: Enable splitting involving additional linear operators by alternating updates to primal and dual variables, as in the Chambolle–Pock scheme.

The pivotal property is that the proximal operators for both the cost functional and the constraint indicator are either closed-form or efficiently solvable due to the problem's separable structure (1304.5784).

2. Discretization Strategies and Implementation

High-fidelity solution of the time-continuous optimal transport problem underlies the effectiveness of KL-Proximal Splitting methods. A key innovation in the referenced work is the introduction of a staggered grid discretization:

• Staggered grids allocate momentum variables on grid edges and density variables on grid centers, closely mirroring physical transport processes.
• This setup enables accurate, low-diffusion approximations to divergence and continuity equations, facilitating mass conservation and stability.
• Operators (proximal and projection) can be executed pointwise or locally, often decoupled across grid points, and major projection steps are implemented via FFTs or other fast solvers for Poisson-type equations.

This discretization yields a large but structured convex optimization problem, naturally suitable for decomposition by proximal splitting methods.

3. Extensions to General Costs and Manifolds

Beyond the canonical $L^2$-Wasserstein cost, the framework supports a family of convex costs of the form: $$j_\beta(m, f) = \begin{cases} \frac{\|m\|^2}{2f^\beta} & f > 0 \ 0 & (m, f) = (0,0) \ +\infty & \text{otherwise} \end{cases}$$ where $\beta \in [0,1]$, unifying the $L^2$-Wasserstein metric ($\beta=1$) and $H^{-1}$-type metrics ($\beta=0$). The approach accommodates:

• Spatially-varying weights ($w_k$): Allowing for modeling of OT on Riemannian manifolds or in domains with obstacles by setting $w_k = +\infty$ to exclude regions.
• Obstacles and boundary modifications: Arbitrary regions in spacetime may be declared impassable, supporting complex geometries in imaging and fluid contexts.

Proximal updates involving these generalized costs remain computationally tractable as they decompose into local cubic (or higher order) equations per grid cell.

4. Algorithmic Efficiency, Scalability, and Comparison

The KL-Proximal Splitting framework demonstrates significant improvements in computational efficiency and numerical stability, particularly in comparison to classic methods:

• Douglas–Rachford/ADMM on Staggered Grids: Achieve faster convergence and superior scaling properties compared to ALG2-centered grid approaches.
• Algorithmic flexibility: Parameters such as relaxation coefficients can be tuned for improved practical performance without compromising theoretical convergence guarantees.
• Scalability: The modularity of splitting, the local nature of the prox computations, and the use of FFT-based projections support deployment to high-resolution grids and large-scale imaging tasks.

A summary comparison is as follows:

Aspect	Benamou & Brenier ALG2	KL-Proximal Splitting (Staggered)
Grid	Centered	Staggered (more physical)
Main Algorithm	ADMM-type	DR, ADMM, PD (tunable, flexible)
Cost Functions	$\|m\|^2/2f$	General family $(\beta)$, obstacles
Riemannian Support	No	Yes (with spatial weights)
Scalability	Moderate	High
Applications	Fluid flows, interp.	Imaging, ML, physics, geoscience

5. Broad Applications and Modeling Implications

KL-Proximal Splitting algorithms for optimal transport are applied across diverse areas:

• Machine learning: Empower Wasserstein-based generative models, domain adaptation, and computation of barycenters for distributional clustering.
• Image processing and computer vision: Provide frameworks for color transfer, shape and texture interpolation, image registration, morphing, and tracking.
• Physical and computational sciences: Facilitate modeling in physics (e.g., incompressible flows), simulation in geophysics and oceanography, and general interpolation or regularization on manifolds and graphs.

The highly extensible architecture encourages their adoption as foundational blocks in modern computational toolkits for data analysis and modeling.

6. Practical Implementation Considerations

Successful deployment of KL-Proximal Splitting methods requires:

• Efficient solvers for local proximal steps and global projections; for mass-conserving projections, FFT-based Poisson solvers are standard.
• Memory and compute resources: As discretized problems are high-dimensional, memory management—especially when storing multiple time steps and variables—must be considered.
• Parameter tuning: Relaxation and stepsize parameters can often be set empirically; the convexity of the problem preserves robustness to moderate deviations.
• Boundary conditions and domain constraints: Handling boundaries or obstacles imposes additional logic but fits naturally within the indicator function and weighted cost framework.

• L. Benamou & J.D. Brenier, A computational fluid mechanics solution to the Monge-Kantorovich mass transfer problem, Numerische Mathematik, 2000.
• Chambolle & Pock, A First-Order Primal-Dual Algorithm for Convex Problems with Applications to Imaging, 2011.
• F. Santambrogio, Optimal Transport for Applied Mathematicians, 2015.

In summary, the KL-Proximal Splitting Scheme delivers an advanced, unified approach for large-scale optimal transport computations, leveraging proximal splitting theory, advanced discretization, and algorithmic flexibility to address modern needs in data science, computer vision, and scientific computing (1304.5784).