Source-linked AI summary
Fast Computation of Wasserstein Barycenters
Marco Cuturi, Arnaud Doucet
TL;DR
Computing Wasserstein barycenters requires optimizing over empirical probability measures while repeatedly solving costly optimal transport problems. The paper introduces two subgradient algorithms and entropically smoothed transport gradients, then applies them to image visualization and constrained clustering. The methods support fixed or freely chosen barycenter support, with the paper demonstrating computations on large image and transport-problem workloads.
Problem
Direct subgradient computation for Wasserstein barycenters is too costly because it repeatedly solves large primal and dual optimal transport problems.
Method
The paper proposes two subgradient algorithms and replaces exact transport computations with entropically smoothed optimal transport problems whose gradients use matrix scaling.
Results
The algorithms visualize nested-ellipse and MNIST images and solve a constrained clustering problem, including a single MNIST barycenter computation reaching 100 iterations in about 2 hours on a GPU.
Takeaways & Limitations
The approach enables Wasserstein barycenter applications involving tens of thousands or tens of high-dimensional transport problems that network-flow solvers cannot realistically handle.
Takeaways & Limitations
A related sliced-Wasserstein approach may not work for d ≥4 and does not generalize to non-Euclidean metric spaces.
Abstract
from arXiv · showhide
We present new algorithms to compute the mean of a set of empirical probability measures under the optimal transport metric. This mean, known as the Wasserstein barycenter, is the measure that minimizes the sum of its Wasserstein distances to each element in that set. We propose two original algorithms to compute Wasserstein barycenters that build upon the subgradient method. A direct implementation of these algorithms is, however, too costly because it would require the repeated resolution of large primal and dual optimal transport problems to compute subgradients. Extending the work of Cuturi (2013), we propose to smooth the Wasserstein distance used in the definition of Wasserstein barycenters with an entropic regularizer and recover in doing so a strictly convex objective whose gradients can be computed for a considerably cheaper computational cost using matrix scaling algorithms. We use these algorithms to visualize a large family of images and to solve a constrained clustering problem.
1. Introduction
The paper targets efficient Wasserstein barycenter computation for empirical measures, using subgradient methods with smoothed optimal transport gradients. It motivates the approach through image means and applications requiring many large transport problems.
- Motivation: Comparing, summarizing, and reducing empirical probability measures are fundamental statistical and machine-learning tasks, but classic divergences have important limitations.For nonfinite spaces, these divergences may be ill-defined for empirical measures and may not incorporate geometry of the underlying space.
- Contribution: The paper computes Wasserstein barycenters by minimizing the sum of optimal transport distances from one variable measure to fixed target measures.The proposed gradients come from smoothed optimal transport problems.
- Illustration: The 2-Wasserstein mean captures the structure of 30 nested-ellipse images, whereas Euclidean, re-centered Euclidean, Jeffrey, and Gaussian RKHS means hardly make sense.The images are discrete measures on a 100 × 100 grid, and the Wasserstein result uses only the Euclidean ground distance.
- Algorithms: The paper develops two subgradient methods: one for known barycenter support and one for freely chosen support, while noting that direct computation is costly.The methods repeatedly require primal or dual optimal transport solutions, which the paper later bypasses with smoothed problems.
2. Background on Optimal Transport
The paper restricts attention to empirical measures and formulates discrete Wasserstein distance through a transportation problem. Pairwise costs form the objective, while transportation polytopes enforce prescribed marginals.
- Empirical measures: Empirical measures are represented as weighted sums of Dirac masses, with atom locations and weights lying in a probability simplex.The paper considers measures of the form µ = Σ_i a_iδ_xi.
- Support and weights: The paper allows barycenter measures supported on a predefined finite set or on up to k points, with weights constrained to a subset of the simplex.When weights are unconstrained, they may range over the full probability simplex.
- Discrete optimal transport: For two discrete measures, the Wasserstein distance is the pth root of a network-flow transportation optimum.The transportation problem uses pairwise distances raised to p as costs and couplings with specified row and column marginals.
- Discrete optimal transport: The squared Wasserstein distance can be written as the optimum of a parametric linear program over n × m variables.Its parameters are the two marginal weight vectors and the pairwise cost matrix.
3. Wasserstein Barycenters
A Wasserstein barycenter minimizes the aggregate Wasserstein distance to several measures over a chosen measure class. The formulation includes clustering and approximation problems, while existing computational approaches have dimensionality and geometry limits.
- Definition: A Wasserstein barycenter is a minimizer of the sum of Wasserstein distances from one measure to a collection of target measures over a prescribed set.The target measures may carry nonnegative distance weights, although the paper uses uniform weights for simpler notation.
- Special cases: The framework covers constrained k-means, histogram centroids, and Euclidean k-means as special cases of Wasserstein barycenter optimization.Uniform k-atom approximation corresponds to constrained k-means, while finite-space 1-Wasserstein distance is the Earth Mover’s Distance.
- Recent work: Sliced-Wasserstein barycenter computation may be effective in lower dimensions but may not work for d ≥4 and does not generalize to non-Euclidean metric spaces.This limitation is stated for the preceding sliced-Wasserstein approach to point clouds.
4. New Computational Approaches
The paper formulates Wasserstein barycenter computation over restricted supports and develops subgradient-based algorithms for fixed or freely chosen support. It derives convexity and subgradients from optimal transport dual or primal solutions, then combines projected or alternating updates while addressing computational cost through smoothing.
- Problem formulation: The search is restricted to barycenters supported on a fixed finite set or on at most k atoms, with weights constrained to a convex subset of the simplex.These are denoted P(X, Θ) and P_k(Ω, Θ), respectively.
- Fixed support: For fixed support X, the objective is convex in the weights, and optimal dual transport solutions provide subgradients for projected or proximal minimization.The proximal formulation uses a Bregman divergence; with the Kullback–Leibler divergence, the update becomes multiplicative and preserves the simplex.
- Free support: When support locations X are free in Euclidean space, the objective is nonconvex in X, but primal optimal transport solutions provide subgradients for location updates.The objective is a convex quadratic plus a piecewise linear concave term, so the resulting function is not convex with respect to X.
- Free support: The free-support method alternates weight optimization with Algorithm 1 and location optimization using Newton updates based on local quadratic approximations.This alternating procedure updates X and a until convergence.
- Connections: For one measure with unconstrained weights, Algorithm 2 reduces to k-means and is equivalent to Lloyd’s algorithm; constrained weights remain feasible through Algorithm 1.The same correspondence extends to the uniform-weight heuristic of Ng (2000).
- Computational cost: The proposed subgradient algorithms are computationally expensive and may converge slowly, motivating entropic smoothing to obtain strictly convex objectives with cheaper gradient computation.Algorithm 1 repeatedly solves N dual transport problems, while Algorithm 2 additionally solves N primal problems at each iteration.
5. Smoothed Dual and Primal Problems
The paper smooths primal and dual optimal transport problems with entropy, yielding strictly convex approximations whose solutions and gradients can be computed through Sinkhorn matrix scaling.
- Regularized primal and smoothed dual: Entropic smoothing replaces repeated exact primal and dual transport solves with strictly convex approximations of both problems.The approach extends Cuturi’s regularization to the barycenter objective.
- Use in barycenter algorithms: The smoothed optima substitute for the exact optima in the barycenter algorithms, providing gradients with a cheaper computational cost.The substitution applies to gradients with respect to both weights a and support locations X.
- Regularized primal and smoothed dual: The smoothed primal and dual optima are linked by positive scaling vectors u and v through the kernel K = e^−λM.The optimal transport has the matrix-scaled form diag(u)K diag(v).
- Matrix scaling computation: Sinkhorn’s matrix scaling algorithm computes the positive vectors u and v needed for the smoothed transport solutions.The algorithm alternates scaling operations, with linear convergence in Hilbert’s projective metric.
- Matrix scaling computation: For p = 2, K is the pairwise Gaussian kernel matrix between supports, formed from the squared-distance matrix M with bandwidth σ = 1/λ.The same smoothing parameter λ is used for the smoothed primal and dual transport problems.
6. Applications
The algorithms are applied to large-scale image visualization and constrained clustering. They produce intermediate digit barycenters and more balanced uniform-weight geographic clusters, while constraints increase Wasserstein distance and computation remains substantially cheaper than exact transport.
- Visualization of perturbed images: 50.000 perturbed MNIST images, approximately 5.000 per digit, are used to compute barycenters after 1, 10, and 60 gradient iterations.Images are randomly scaled and translated within a 50 × 50 grid.
- Visualization of perturbed images: A single MNIST barycenter takes about 2 hours to reach 100 iterations on a Quadro K5000 GPU with close to 1500 cores.Warm starts make iterations near the end typically less computationally intensive.
- Clustering with uniform centroids: Uniform weight constraints produce more balanced clustering than free assignments for income and population across 57.647 locations in the 48 contiguous states.Each Voronoi cell is constrained to hold the same aggregate wealth or population.
- Clustering with uniform centroids: Uniform constraints increase Wasserstein distance relative to unconstrained k-means, while Algorithm 2 converges at a comparable iteration speed with modest computational overhead.The larger distance is a consequence of restricting the centroid weights.
- Clustering with uniform centroids: 12.5 seconds for the constrained case and 1.55 seconds for regular k-means are required on one CPU core using Sinkhorn’s approximation.A regular transportation solver would require about 1 hour for one iteration on the 57.647-to-48 transport problem.