[Paper Review] Scaling Algorithms for Unbalanced Transport Problems

Lénaïc Chizat, Gabriel Peyré+2 more|arXiv (Cornell University)|Jul 20, 2016

3D Shape Modeling and AnalysisEngineering54 references80 citations

TL;DR

This paper introduces scalable, entropic regularization-based algorithms for unbalanced optimal transport problems, extending the Sinkhorn algorithm to handle arbitrary positive measures with mass creation/destruction. The method enables fast, parallelizable computation of unbalanced transport, barycenters, and gradient flows, with applications in shape modification, color transfer, and growth modeling.

ABSTRACT

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.

Motivation & Objective

• Address the limitation of classical optimal transport, which requires normalized probability measures, by enabling computation on arbitrary positive measures with mass variation.
• Develop efficient numerical methods for unbalanced optimal transport problems that are scalable and parallelizable, overcoming the computational intractability of traditional linear programming approaches.
• Extend entropic regularization and the Sinkhorn algorithm framework to unbalanced settings, enabling practical computation of unbalanced transport, barycenters, and gradient flows.
• Provide a unified computational framework for diverse applications in image processing, computer graphics, and machine learning that require flexible mass handling.

Proposed method

• Apply entropic regularization to unbalanced optimal transport problems using Bregman divergences, particularly the Kullback-Leibler divergence, to ensure convexity and numerical stability.
• Derive iterative algorithms based on diagonal scaling (pointwise multiplication) of transport couplings, generalizing the Sinkhorn algorithm to unbalanced settings.
• Formulate the unbalanced transport problem as a convex optimization problem with a strictly convex regularization term, enabling closed-form updates at each iteration.
• Use Bregman projections and subdifferential calculus to derive optimality conditions and update rules for transport plans and barycenters.
• Implement scalable solvers by leveraging the structure of the regularized problem, allowing for efficient GPU-accelerated and distributed computation.
• Extend the framework to compute unbalanced Wasserstein barycenters and unbalanced gradient flows by adapting the iterative scaling procedure to these higher-level problems.

Experimental results

Research questions

• RQ1 How can entropic regularization be generalized to handle unbalanced optimal transport problems involving arbitrary positive measures with mass creation or destruction?
• RQ2 What is the structure of the iterative algorithms derived from entropic regularization in unbalanced settings, and how do they maintain computational efficiency and parallelizability?
• RQ3 Can the extended Sinkhorn-style algorithms be effectively applied to compute unbalanced barycenters and unbalanced gradient flows in practical applications?
• RQ4 What are the theoretical properties of the regularized unbalanced transport problem, particularly in terms of convergence and consistency as regularization vanishes?
• RQ5 How do the proposed algorithms compare in performance and accuracy to classical linear programming or combinatorial methods on large-scale unbalanced transport problems?

Key findings

• The proposed entropic regularization framework successfully generalizes the Sinkhorn algorithm to unbalanced optimal transport, enabling fast and scalable computation on arbitrary positive measures.
• The iterative algorithms rely solely on diagonal scaling operations (pointwise multiplications), making them highly parallelizable and suitable for GPU acceleration.
• The method achieves accurate approximations of unbalanced transport couplings, barycenters, and gradient flows with convergence rates comparable to classical entropic OT.
• Applications in 2D shape modification, color transfer, and growth modeling demonstrate the method’s practical utility and robustness on real-world data.
• Theoretical analysis confirms that the regularized problem converges to the classical unbalanced transport solution as the regularization parameter tends to zero.
• The framework supports various divergences (KL, TV, robust gamma), enabling flexibility in modeling different types of mass variation and noise.