Source-linked AI summary
Scaling Algorithms for Unbalanced Transport Problems
Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, François-Xavier Vialard
TL;DR
The paper addresses the limitation that classical optimal transport handles only unit-mass probability measures, while applications may require arbitrary positive measures and mass variation. It extends entropic regularization through iterative scaling algorithms, yielding a Sinkhorn generalization for unbalanced transport, barycenters, and gradient flows.
Problem
Classical optimal transport requires normalized unit-mass measures, limiting applications that need arbitrary positive measures or partial mass displacement.
Method
The paper develops iterative scaling algorithms that extend entropic regularization and Sinkhorn-style diagonal scaling to unbalanced transport problems.
Results
The methods are shown to solve unbalanced transport, barycenters, and gradient flows, with convergence established under stated continuous and discrete assumptions.
Takeaways & Limitations
The framework provides fast, parallelizable computational methods for transport problems involving mass variation and arbitrary positive measures.
Takeaways & Limitations
A sound theoretical analysis of the corresponding unbalanced gradient flows and their limit PDEs is not yet available, so the paper uses heuristic smooth-setting arguments.
Abstract
from arXiv · showhide
This article introduces a new class of fast algorithms to approximate variational problems involving unbalanced optimal transport. While classical optimal transport considers only normalized probability distributions, it is important for many applications to be able to compute some sort of relaxed transportation between arbitrary positive measures. A generic class of such "unbalanced" optimal transport problems has been recently proposed by several authors. In this paper, we show how to extend the, now classical, entropic regularization scheme to these unbalanced problems. This gives rise to fast, highly parallelizable algorithms that operate by performing only diagonal scaling (i.e. pointwise multiplications) of the transportation couplings. They are generalizations of the celebrated Sinkhorn algorithm. We show how these methods can be used to solve unbalanced transport, unbalanced gradient flows, and to compute unbalanced barycenters. We showcase applications to 2-D shape modification, color transfer, and growth models.
1. Introduction
Classical optimal transport assumes unit-mass probability measures, limiting applications that require arbitrary masses or partial transport. This paper extends entropic regularization into iterative scaling algorithms for unbalanced transport and related problems.
- Motivation: Classical optimal transport requires normalized unit-mass measures, whereas many applications need mass creation, destruction, or partial displacement.This motivates unbalanced formulations over arbitrary positive measures.
- Scope: The framework targets unbalanced transport, barycenters, and gradient flows, including flows over arbitrary positive measures with mass creation and destruction.Such flows are relevant to growth phenomena and other applications beyond fixed-mass transport.
- Numerical approach: Entropic regularization enables simple, scalable computations through diagonal scaling, complementing earlier methods that do not handle large problems with arbitrary transportation costs.The approach is connected to Sinkhorn iterations and uses pointwise scaling operations.
- Problem formulation: The paper situates unbalanced transport among formulations using approximate marginal constraints, semi-couplings, and dynamic source terms, selecting approximate marginals for numerical treatment.This formulation turns the transport problem into a convex problem involving ϕ-divergences.
- Contribution: The paper’s main contribution is a class of iterative scaling algorithms that solve entropic approximations of diverse unbalanced optimal transport problems.The algorithms are presented as direct generalizations of Sinkhorn’s method.
2. Unifying Formulation of Transport-like Problems
The paper unifies balanced and unbalanced transport-related problems through transport costs and convex functionals acting on coupling marginals. This framework covers divergences, soft-marginal transport, barycenters, and unbalanced gradient flows.
- 2.1. Divergence Functionals: Divergences measure how close positive measures are through pointwise ratios and are constructed from entropy functions.For entropy functions, the resulting divergences are jointly 1-homogeneous, convex, and weakly* lower semicontinuous.
- 2.1. Divergence Functionals: The Kullback–Leibler divergence is central, while total variation and equality or range constraints provide other divergence choices.Range constraints require α ν ⩽ µ ⩽ β ν, whereas equality constraints require µ = ν.
- 2.2. Balanced Optimal Transport: Classical optimal transport seeks the least-cost coupling between measures under fixed marginal constraints, but requires equal total mass.The formulation can be interpreted as finding the cheapest way to move mass from one distribution to another.
- 2.3. Unbalanced Optimal Transport: Unbalanced transport relaxes hard marginal constraints by penalizing deviations of coupling marginals from the input measures.The soft-marginal formulation includes balanced transport as a special case when the entropy domains enforce equality constraints.
- 2.3. Unbalanced Optimal Transport: WFR is a KL-based unbalanced transport distance, and it defines a distance when 0 ⩽ λ ⩽ 1, with the upper bound necessary on sufficiently large geodesic spaces.The construction uses a transport cost together with KL penalties on both marginals.
- 2.4. Barycenter Problem and Extensions: The unified framework also includes barycenters and minimizing-movement gradient flows over nonnegative measures, including mass creation and destruction.The barycenter minimizer can be obtained as a byproduct of the scaling algorithm, while WFR gradient-flow theory remains incomplete and is treated heuristically.
- 2.5.2. Minimization problem: The framework supports applications such as growth modeling because unbalanced metrics allow flows over arbitrary positive measures rather than only probability measures.The article explicitly notes that a sound theoretical analysis of WFR gradient flows and their limit PDEs is not yet available.
3. Entropic Regularization and Iterative Scaling Algorithm
The paper entropically regularizes a generic transport variational problem and solves the regularized problem with an iterative scaling algorithm. The regularizer yields strict convexity, while the algorithm generalizes Sinkhorn-style updates to marginal functionals.
- 3. Entropic Regularization and Iterative Scaling Algorithm: An iterative scaling algorithm solves a regularized version of the generic variational problem in a continuous setting.The method is presented as a direct computational procedure for the entropically regularized problem.
- 3.1. Entropic Regularization: The regularized objective adds ε H(γ) to J(γ), replacing the nonnegativity constraint with an entropy term based on reference measures dx and dy.The reference product measure dxdy supplies the densities used to define the entropy of the couplings.
- 3.1. Entropic Regularization: For finite spaces, as ε → 0, the unique regularized minimizer converges to the minimizer of J having minimal entropy.Strict convexity and lower semicontinuity ensure a unique regularized minimizer for each positive ε.
- 3.1.2. Interpretation 2.: The regularizer makes the objective strictly convex and can be interpreted as a KL proximal step.The proximal-point interpretation yields repeated updates γ^(ℓ+1) = prox_KL J/ε(γ^(ℓ)).
- 3.1. Entropic Regularization: The reference measure must contain the support of an optimizer, and its numerical choice corresponds to the discretization grid.Choosing normalized input measures as references can fail when divergences are not superlinear.
3.2. Reformulation using Densities and Duality.
The entropically regularized transport problem is reformulated over densities, yielding a primal–dual pair with strong duality and a unique primal minimizer. The marginal penalties enable an alternating dual maximization related to Sinkhorn iterations.
- Strong duality holds, and the entropically regularized primal problem has a unique minimizer.
- The dual formulation uses the marginal projections of the coupling through the operator A(u,v)(x,y)=u(x)+v(y).
- Because the marginal functionals act separately, Dykstra iterations reduce to alternating dual maximization related to Sinkhorn’s algorithm.
3.3. Scaling Algorithm.
The scaling algorithm alternates pointwise proximal updates of the two marginal functionals, implemented through multiplicative factors applied to the transport kernel. Under suitable regularity conditions, these updates correspond to alternating maximization of the dual problem.
- Scaling iterates update positive factors through kernel projections and proximal operators for the marginal functionals.
- The scaling iterations are equivalent to alternating maximization iterates of the dual problem whenever the iterates are well defined.
- Existence of the iterates is handled by assumptions ensuring the relevant functions and proximal operations are well defined.
- Each KL proximal update decomposes into pointwise optimization problems under admissible boundedness conditions.
3.4. Existence of the iterates for integral functionals.
For admissible integral marginal functionals, conjugation, subdifferentiation, and KL proximal operations can be treated pointwise. Positivity of the kernel and feasible positive points ensure uniquely defined, measurable scaling and dual iterates.
- The assumptions require nonnegative pointwise domains, finite feasibility, strictly positive feasible coordinates, and positive kernel values.
- Admissible integral functionals are convex and weakly lower semicontinuous, while their conjugates remain integral functionals.
- The KL proximal operator for an integral functional decomposes into pointwise minimization problems almost everywhere.
- Under admissible integral-function assumptions and a positive kernel, the scaling and dual iterates are uniquely well defined.
- Pointwise strict convexity and measurable minimizer selection support measurable dual updates and the relation between scaling and dual iterates.
3.5. Convergence Analysis.
The scaling formulation recovers the unique primal and dual solutions from a fixed point of the multiplicative factors. For KL marginal penalties with bounded logarithms and a positively lower-bounded kernel, the iterates converge linearly in the Thompson metric.
- The coupling r_k(x,y)=a_k(x)K_k(x,y)b_k(y) is the unique primal solution when the scaling factors form a suitable fixed point.
- The Thompson metric supplies a convergence proof whose rate does not depend on a bound on the transport cost.
- The metric is complete on each finite-distance equivalence class and contracts under suitable homogeneous order-preserving or order-reversing operators.
- The scaling iterates converge at a linear rate in the Thompson metric when log p and log q are bounded and K is lower bounded by a positive real number.
- For KL marginal penalties, the contraction factor is z1·z2<1, where z_i=λ_i/(λ_i+ε).
4. Algorithm for Discrete Measures
For finite discrete spaces, the paper turns unbalanced transport into implementable scaling iterations with convergence guarantees and numerical stabilization. The resulting algorithms use pointwise proximal updates and matrix-vector products, return approximate primal minimizers, and generalize to richer transport models.
- Convergence: Scaling iterations converge in finite dimension, with a pessimistic O(1/ℓ) guarantee while practice exhibits linear convergence.The convergence theorem concerns the scaling and associated dual iterates.
- Discrete formulation: The discrete formulation represents couplings as matrices, uses the Gibbs kernel K=e^(-C/ε), and evaluates scaling operations through weighted matrix-vector products.Vectors represent functions on X and Y, while matrices represent functions on their product space.
- Scaling algorithm: Algorithm 1 alternates proximal updates of a and b using K and K^T, then returns the primal optimizer as (a_iK_ijb_j)_ij.The required information about F1 and F2 is condensed into proxdiv functions, making the algorithm straightforward to implement.
- Numerical stabilization: For small ε, stabilized iterations absorb extreme scaling values into log-domain variables while retaining matrix-product computations.This redundant parametrization keeps auxiliary factors near one and mitigates overflow and numerical imprecision.
- Numerical stabilization: The stabilized method remains vulnerable because proxdiv may still involve the potentially extreme factor e^(-u/ε), although many cases avoid evaluating that exponential.The paper reports numerical stability in the small-ε limit for many practical proxdiv computations.
- Generalization: The scaling framework extends beyond two functionals and spaces, including partial transport and total-mass range constraints through an additional convex functional F3.When F3 enforces equality to a positive mass, the resulting algorithm solves partial optimal transport.
5. Applications
The paper applies stabilized scaling algorithms to unbalanced transport, color transfer, barycenters, gradient flows, and density evolution. These experiments show practical numerical behavior, including sharp low-regularization solutions, structure-preserving barycenters, and meaningful color selection.
- 5.1.2. Numerical examples for X = [0, 1].: At ε = 10^-7, stabilized scaling produces quasi-deterministic plans whose approximate supports expose optimizer structure across divergences.The displayed support contains entries greater than 10^-10; in the TV case, the minimal-entropy plan follows diagonal segments.
- 5.1.2. Numerical examples for X = [0, 1].: Unbalanced transport experiments show linear practical convergence of scaling iterations, despite a pessimistic general convergence-rate bound.The primal-dual gap is evaluated for Algorithm 1 on discretized quadratic-cost problems.
- 5.1.4. Color transfer.: In challenging color transfer, divergence choice selects the amount of target color mass, matching initial-histogram modes without quantitative image-quality measures.The experiment uses ε = 0.002, runs for 2000 iterations, and takes approximately 160 seconds.
- 5.2. Barycenters.: Relaxed marginal constraints preserve three-bump and other global input structures in unbalanced barycenter and related experiments.The reported comparisons contrast relaxed constraints with classical optimal transport.
- 5.3. Gradient Flows and Evolution of Densities.: In the growth model, a steady state is reached in finite time with the two densities summing to 1, while density-dependent growth can push one species spatially.The interaction is attributed to lower WFR effort for adding mass to regions of higher density.
Appendix A. Appendix
The appendix introduces convex conjugates for functions on topologically paired vector spaces.
- For a function f on E, the convex conjugate is defined on the paired space E∗ through the bilinear pairing.
A.1. Reminders on Convex Analysis.
This section recalls convex-analytic foundations, including conjugacy and a Fenchel–Rockafellar duality result with attainment conditions.
- Convex conjugates are defined for functions on paired locally convex Hausdorff spaces.
- The Fenchel–Rockafellar theorem applies to proper lower-semicontinuous convex functions linked by a continuous linear operator.Under a continuity condition at one feasible image, the minimum is attained.
- The appendix frames later divergence-functional results as function-space counterparts of measure-based divergences introduced earlier.
A.2. Properties of Divergence Functionals.
The appendix establishes convexity, lower semicontinuity, conjugacy, and subdifferential properties for divergence functionals on L1 spaces.
- Divergence functionals are positively 1-homogeneous, convex, and weakly lower semicontinuous as admissible integral functionals.
- The conjugate and subdifferential descriptions impose separate conditions where v is positive and where v vanishes.
- For fixed v, Dϕ(·|v) is proper, weakly lower semicontinuous, and convex on L1(X), with an explicit convex conjugate.
- The subdifferential is characterized pointwise through the entropy subdifferential and the condition ϕ′∞−a≥0.
- Normal-integrand and integral-functional results justify performing conjugation and subdifferentiation pointwise.
A.3. Proof of the Iterates for the Barycenter Problems.
The barycenter analysis derives the update expression for h by applying Proposition 5.3, reducing a case to a one-dimensional minimization problem.
- The expression for h in Table 2 is derived by applying Proposition 5.3 to the barycenter iterates.
- In the simple case, solving (5.6) reduces to a one-dimensional minimization over h.
A.3.1. Case Dϕ = ι{=}.
The section derives optimality conditions for the equality-indicator case by characterizing admissible subgradients and reducing feasibility to weighted conditions on the auxiliary variables. It also enumerates the resulting regimes according to the relative values of ˜s_k and h.
- KL case: The smooth KL case yields a singleton joint subdifferential for positive ˜s and h, while KL(0|h) has second subgradient {1} for h > 0.These subgradients produce the stated optimality system, including ˜s_k=0 when s_k=0 and the weighted condition P α_k(1−˜s_k/h)=0.
- Total-variation case: For total variation, the subgradient is obtained from a support-function representation, producing six regimes for the relative positions of ˜s_k and h, including h=0 boundary cases.The regimes specify values or ranges for a_k and b_k, such as a_k=1 with b_k≤−1 when ˜s_k>h=0.
- Piecewise characterization: For h>0 and ˜s_k>0, b_k is clipped to [−1,1] through the ratio involving ˜s_k and h, with the weighted constraint P α_kb_k=0.The four cases distinguish interior, boundary, and zero configurations through β_1h, β_2h, and ˜s_k.
- Feasibility regimes: If any s_k=0, feasibility forces h=0; otherwise h>0 and the condition P α_kb_k=0 determines h implicitly.The text identifies h=0 as the only feasible point when some s_k vanishes.