Source-linked AI summary
Iterative Bregman Projections for Regularized Transportation Problems
Jean-David Benamou, Guillaume Carlier, Marco Cuturi, Luca Nenna, Gabriel Peyré
TL;DR
Optimal transport and its generalizations require efficient approximate solvers, especially as direct formulations become computationally expensive. The paper introduces entropic regularization and interprets the resulting problems as KL projections, solved by iterative Bregman or Dykstra methods. The framework is applied to several generalized OT problems and is most suitable when some smoothing is acceptable, since very small regularization causes instability and slow convergence.
Problem
Optimal transport and its generalizations lack sufficiently fast algorithms for many computationally demanding formulations.
Method
The paper introduces entropic regularization and solves the resulting KL projection problems using iterative Bregman projections or Dykstra’s algorithm.
Results
The framework provides numerical methods for entropic approximations of barycenters, tomographic reconstruction, multi-marginal, partial, and capacity-constrained transport.
Takeaways & Limitations
The approach is useful for fast approximations in applications where a small amount of smoothing is acceptable.
Takeaways & Limitations
The methods become numerically unstable for small ε, and iterative projection convergence slows significantly as ε approaches zero.
Abstract
from arXiv · showhide
This article details a general numerical framework to approximate so-lutions to linear programs related to optimal transport. The general idea is to introduce an entropic regularization of the initial linear program. This regularized problem corresponds to a Kullback-Leibler Bregman di-vergence projection of a vector (representing some initial joint distribu-tion) on the polytope of constraints. We show that for many problems related to optimal transport, the set of linear constraints can be split in an intersection of a few simple constraints, for which the projections can be computed in closed form. This allows us to make use of iterative Bregman projections (when there are only equality constraints) or more generally Bregman-Dykstra iterations (when inequality constraints are in-volved). We illustrate the usefulness of this approach to several variational problems related to optimal transport: barycenters for the optimal trans-port metric, tomographic reconstruction, multi-marginal optimal trans-port and in particular its application to Brenier's relaxed solutions of in-compressible Euler equations, partial un-balanced optimal transport and optimal transport with capacity constraints.
1 Introduction
The paper addresses computational bottlenecks in optimal transport by combining entropic regularization with KL-based projection methods. It develops a unified framework for generalized OT problems and applies it across several extensions.
- Computational optimal transport: Discrete optimal transport can be computationally demanding, motivating fast approximate algorithms for generalized OT problems.Matching methods have roughly O(N^3) complexity, while dynamic formulations add a time dimension and remain expensive.
- Entropic regularization: Entropic regularization replaces the original OT problem with a smoothed problem whose solution is unique and can be computed through diagonal scaling and iterative proportional fitting.The regularized solution is a diagonal scaling of e^(-C), and IPFP provides an efficient iterative scheme.
- Entropic regularization: Entropic penalization improves computational stability and enables KL projection methods, but it produces non-sparse couplings and smooths the transport solution.The regularized problem is strongly convex, unlike the initial OT problem, while its smoothing changes the support structure.
- Generalized OT problems: The framework treats OT barycenters, tomographic reconstruction, multi-marginal transport, and other generalized problems through suitable entropic formulations.The paper specifically presents applications to barycenters, tomographic reconstruction, multi-marginal transport, partial transport, and capacity-constrained transport.
- Computational speed: The proposed algorithms are competitive when moderate regularization is acceptable, whereas small ε causes numerical instability and slower iterative convergence.The authors report machine-precision issues for quantities of order e^(-1/ε) and significant degradation in convergence speed as ε approaches zero.
2 Iterative Bregman Projections and Dykstra Algorithm
The framework recasts regularized generalized OT problems as KL projections onto intersections of convex sets. Iterative KL projections solve affine-constraint cases, while Dykstra’s algorithm handles general convex constraints and converges to the desired projection.
- Problem formulation: Regularized generalized OT problems are formulated as KL projections of an input vector onto an intersection of closed convex sets.The constraint sets represent the structure of the regularized problem.
- Iterative Bregman projections: When all constraint sets are affine subspaces, iterative KL projections can solve the problem starting from γ(0) = ξ.The iterates are formed by repeatedly projecting onto the constraint sets.
- Iterative Bregman projections: The iterative KL projection sequence converges to the unique solution of the regularized projection problem.The convergence statement applies to the affine-subspace case.
- Dykstra’s algorithm: For non-affine convex sets, ordinary iterative Bregman projections need not converge to the KL projection on the intersection.This motivates the use of Dykstra’s algorithm in the general convex-constraint setting.
- Dykstra’s algorithm: Dykstra’s algorithm converges to the solution of the KL projection problem for general convex constraint sets.This extends the convergence guarantee beyond affine-subspace constraints.
- Dykstra’s algorithm: Dykstra’s algorithm augments successive KL projections with correction variables initialized alongside γ(0) = ξ.The correction terms are initialized to one before the iterative updates begin.
3 Entropic Regularization of Transport-like Problems
The paper formulates entropically regularized transport problems as KL-divergence projections and solves them through iterative Bregman projections or related closed-form projection schemes. This framework is applied to transport, barycenter, and partial Radon inversion problems.
- 3.1 Entropic Optimal Transport Regularization: Entropic regularization makes the transport solution smoother, strongly convex, and uniquely computable through elementary matrix and vector operations.As ε approaches zero, the regularized solution converges to an unregularized optimum; with multiple optima, it selects the one with largest entropy.
- 3.1 Entropic Optimal Transport Regularization: The regularized transport problem is a KL projection of ξ = e^−C onto affine marginal constraints.The two projections alternately normalize rows and columns to match the prescribed marginals p and q.
- 3.1 Entropic Optimal Transport Regularization: Sinkhorn iterations can be implemented using matrix-vector multiplications with a fixed matrix ξ, including parallel computation for transport problems sharing the same ground cost.The updates use u(n) = p/(ξv(n)) and v(n+1) = q/(ξ^T u(n)).
- 3.2 Optimal Transport Barycenters: The barycenter algorithm generalizes IPFP/Sinkhorn from two marginals to weighted barycenters by replacing KL with a weighted divergence and splitting the constraints into simple sets.The projections can be computed separately for each input marginal and in parallel.
- 3.3 Partial Radon Inversion with OT Fidelity: For partial Radon inversion, Wasserstein fidelity to measurements and a template is used to improve reconstructions when few projections produce artifacts.A larger λ1 yields a reconstruction closer to the template, whereas smaller λ introduces more geometric features of the target but also reconstruction artifacts.
4 Multi-marginal Optimal Transport
The paper extends KL-projection methods to multi-marginal optimal transport, enabling iterative projection algorithms for barycenters and generalized Euler flows. Exploiting the cost structure reduces high-dimensional computations to manageable transition-probability operations while retaining diffuse, non-deterministic flows.
- Multi-marginal formulation: Multi-marginal optimal transport is formulated as a KL projection onto the intersection of marginal-constraint sets.Each Bregman projection is a simple normalization, so iterative Bregman projections compute the regularized problem.
- Multi-marginal barycenters: The framework computes measure barycenters through a multi-marginal transport problem with weighted input measures and a quadratic Euclidean ground cost.The resulting barycenter can contain more than N Diracs located off the input discretization grid, unlike fixed-grid formulations.
- Multi-marginal barycenters: Figure 4 displays a three-marginal barycenter from annulus, diamond, and square inputs on a 60 × 60 grid with ε = 0.005.The displayed histogram aggregates the barycenter measure over discretization squares.
- Generalized Euler flows: For generalized Euler flows, discretizing time into K steps yields a multi-marginal transport problem whose transition probabilities represent particle evolution between times.Trajectories are non-deterministic, allowing particle mass to split and spread across the domain.
- Generalized Euler flows: The proposed Euler-flow method naturally represents non-diffeomorphic volume-preserving maps and avoids storing the full (N^d)^K coupling by using (N^d)^2 matrices.The cost structure permits reduction to two smaller coupling matrices and efficient matrix-product updates.
- Generalized Euler flows: With N = 200, ε = 10^-3, and K = 16, the computed transition maps agree with solutions produced by Brenier and capture mass spreading through two-marginal couplings.The experiments use uniform discretization in one dimension for three test cases.
5 Transport Problems with Inequality Constraints
The section extends KL-projection methods to transport problems with inequality constraints, including partial, capacity-constrained, and multi-marginal partial transport. These problems use convex-set projections and Dykstra iterations, with numerical tests illustrating active regions, symmetries, and higher-dimensional solutions.
- General framework: Inequality-constrained transport problems are formulated as KL projections onto intersections of convex sets rather than affine constraints.Because the constraint sets are non-affine, Dykstra iterations are used to converge to the regularized solution.
- 5.1 Partial Transport: Partial transport allows marginals with different total masses while transporting only a prescribed amount of mass.The active source and target regions identify where transported mass is concentrated, using a threshold η to detect active regions.
- 5.1 Partial Transport: The KL projection for partial transport with three convex constraint sets can be computed in closed form.The component-wise minimum appears in the projection formula, while Dykstra iterations handle the intersection of non-affine sets.
- 5.2 Capacity Constrained Transport: Capacity-constrained transport imposes an upper bound θ_i,j on the mass transferred from source i to target j.The regularized problem combines marginal constraints with the component-wise capacity constraint.
- 5.2 Capacity Constrained Transport: In capacity-constrained transport, optimal plans saturate the capacity constraint and can exhibit symmetries between minimizers with Hölder-conjugate capacities.The numerical tests reproduce expected symmetries in one- and two-dimensional settings, including N = 100 and N = 50 × 50 discretizations with ε = 10^-3.
Conclusion
The paper presents entropic regularization as a unifying approximation framework for OT-related linear programs. It supports fast KL-based iterative methods across applications, while introducing smoothing and losing competitiveness for highly accurate recovery of the original linear program.
- The paper presents a unifying framework for approximating solutions of various OT-related linear programs through entropic regularization.
- Entropic regularization enables simple iterative KL projection methods for these problems.
- The approach provides fast approximations for applications including machine learning, image processing, and economics, at the expense of extra smoothing.
Aknowledgements
The acknowledgements thank collaborators and funding organizations supporting the work. Support came through research grants, institutional programs, and donated computing hardware.
- The authors thank Yann Brenier and Brendan Pass for stimulating discussions.
- G. Peyré acknowledges support from the European Research Council through the SIGMA-Vision project.
- Benamou, Carlier, and Nenna acknowledge ANR and INRIA support, while Cuturi acknowledges JSPS support and a donated NVIDIA K40 card.