Source-linked AI summary
Stochastic Optimization for Large-scale Optimal Transport
Genevay Aude, Marco Cuturi, Gabriel Peyré, Francis Bach
TL;DR
Optimal transport captures geometry between probability distributions but is computationally burdensome at large scale. This paper develops sampling-based stochastic optimization methods for discrete, semi-discrete, and continuous OT, reporting improved scalability and the first converging approach for two arbitrary continuous densities under an RKHS assumption.
Problem
Existing OT solvers face computational burdens, cannot directly handle continuous densities, and lack a method for distances between two continuous densities.
Method
The paper recasts dual OT objectives as expectations and applies stochastic optimization, including incremental methods, averaged SGD, and RKHS expansions.
Results
The methods cover discrete, semi-discrete, and continuous OT; discrete incremental methods can surpass Sinkhorn in efficiency, and the continuous method converges under an RKHS assumption.
Takeaways & Limitations
Sampling-based optimization extends OT to high-dimensional settings without requiring discretization of arbitrary input densities.
Takeaways & Limitations
The continuous method assumes that both optimal potentials belong to reproducing kernel Hilbert spaces, and the algorithms require both input measures to be sampleable.
Abstract
from arXiv · showhide
Optimal transport (OT) defines a powerful framework to compare probability distributions in a geometrically faithful way. However, the practical impact of OT is still limited because of its computational burden. We propose a new class of stochastic optimization algorithms to cope with large-scale problems routinely encountered in machine learning applications. These methods are able to manipulate arbitrary distributions (either discrete or continuous) by simply requiring to be able to draw samples from them, which is the typical setup in high-dimensional learning problems. This alleviates the need to discretize these densities, while giving access to provably convergent methods that output the correct distance without discretization error. These algorithms rely on two main ideas: (a) the dual OT problem can be re-cast as the maximization of an expectation ; (b) entropic regularization of the primal OT problem results in a smooth dual optimization optimization which can be addressed with algorithms that have a provably faster convergence. We instantiate these ideas in three different setups: (i) when comparing a discrete distribution to another, we show that incremental stochastic optimization schemes can beat Sinkhorn's algorithm, the current state-of-the-art finite dimensional OT solver; (ii) when comparing a discrete distribution to a continuous density, a semi-discrete reformulation of the dual program is amenable to averaged stochastic gradient descent, leading to better performance than approximately solving the problem by discretization ; (iii) when dealing with two continuous densities, we propose a stochastic gradient descent over a reproducing kernel Hilbert space (RKHS). This is currently the only known method to solve this problem, apart from computing OT on finite samples. We backup these claims on a set of discrete, semi-discrete and continuous benchmark problems.
1 Introduction
Optimal transport lifts known feature geometry into a geometrically faithful distance between probability distributions, but its computational cost limits large-scale use. The paper introduces stochastic methods covering discrete, semi-discrete, and continuous OT settings.
- Motivation: Optimal transport uses feature geometry to define Wasserstein or earth mover’s distances between probability distributions.Applications include visual words, colors, shapes, words, and labels.
- Prior methods: Existing discrete OT reduces to a large-scale linear program, while entropic regularization enables efficient approximate computation with Sinkhorn’s algorithm.These approaches remain primarily finite-dimensional.
- Limitations: Current methods cannot handle continuous densities directly; semi-discrete solvers are restricted to squared Euclidean costs and low dimensions, while continuous-continuous OT remains open.Semi-discrete methods compare a discrete distribution with a continuous density.
- Contributions: The paper introduces stochastic optimization methods for discrete, semi-discrete, and continuous OT, including classical and entropically regularized formulations.The methods require sampling from the input distributions rather than discretizing arbitrary densities.
- Contributions: The discrete method uses incremental algorithms such as SAG, requiring N operations per iteration and scaling better than Sinkhorn on large problems.Its stated convergence rate is O(1/k).
- Contributions: Averaged SGD addresses semi-discrete OT, while an RKHS expansion supplies a converging method for OT between two arbitrary densities under an RKHS assumption on the potentials.The paper presents this as the first known converging algorithm for the continuous-continuous case.
2 Optimal Transport: Primal, Dual and Semi-dual Formulations
The paper formulates OT and entropic regularization through primal, dual, and semi-dual optimization problems over measures and potentials. Recasting the dual objectives as expectations enables stochastic optimization, with finite-dimensional semi-dual problems for discrete targets and RKHS methods for continuous targets.
- Primal formulation: The Kantorovich formulation minimizes transport cost over joint measures with prescribed marginals, and the ground cost c(x, y) represents the cost of moving unit mass.The framework allows discrete and continuous measures.
- Primal formulation: For ε > 0, entropic regularization makes the primal problem strongly convex, yielding a unique optimal transport plan and enabling Sinkhorn computations.The regularization uses KL divergence relative to the product measure.
- Dual and semi-dual formulations: The dual and semi-dual are concave maximization problems over Kantorovich potentials, with the semi-dual obtained by eliminating one potential through its smoothed c-transform.The dual variables are nonunique up to scalar translation.
- Dual and semi-dual formulations: When ν is discrete, the semi-dual becomes a finite-dimensional concave maximization problem suitable for stochastic programming.This advantage does not require μ to be discrete.
- Stochastic formulation: Both dual problems can be rewritten as expectations over independent samples X ~ μ and Y ~ ν, providing the basis for stochastic optimization.At ε = 0, the dual retains a potential constraint that prevents unconstrained expectation maximization.
3 Discrete Optimal Transport
The discrete OT formulation is finite-dimensional, but batch Sinkhorn iterations require costly matrix-vector products. The paper instead uses incremental stochastic optimization, with SAG providing faster convergence in large discrete problems.
- Discrete OT represents two measures as finite weighted sums of Dirac masses and rewrites its primal, dual, and semi-dual problems using a cost matrix.
- Sinkhorn solves regularized discrete OT by block-coordinate maximization, but each iteration uses a matrix-vector multiplication with O(N^2) complexity.
- Incremental stochastic gradient methods sample individual points instead of using the full batch, reducing the computational burden of discrete optimization.
- SAG stores previously computed sample gradients and updates toward the average gradient, improving convergence at the cost of gradient memory.
- In a 35-corpus word-embedding benchmark, SAG was more than twice as fast as Sinkhorn on average across tolerance thresholds while retaining GPU parallelization.
4 Semi-Discrete Optimal Transport
Semi-discrete OT compares an arbitrary, sampleable measure with a discrete measure through a finite-dimensional semi-dual problem. Averaged SGD avoids discretizing the continuous measure and converges at rate O(1/k), whereas empirical discretization can introduce bias.
- Semi-discrete OT assumes an arbitrary measure μ and a discrete measure ν, making the semi-dual a finite-dimensional expectation-based maximization problem.
- Approximating the continuous measure by an empirical measure changes the problem and introduces discretization noise in the resulting solution.
- Averaged SGD operates directly on the arbitrary measure without requiring it to be discrete and converges at rate O(1/k).
- In Gaussian-mixture experiments, regularized SGD iterates converged faster but did not converge to the correct unregularized solution.
- Compared with SAG on sampled empirical measures, SGD converged to the true semi-discrete solution, while SAG remained biased because of discretization error.
5 Continuous optimal transport using RKHS
For two continuous densities, the paper uses kernel stochastic gradient descent to optimize the regularized OT dual in RKHSs, yielding a convergent algorithm under an RKHS assumption. The method is demonstrated on a Gaussian-to-mixture benchmark, where convergence is slow but noticeable.
- Method: Continuous OT is infinite-dimensional, so the paper optimizes its regularized dual using expansions of both potentials in two RKHSs.The approach applies when ε > 0 because the unregularized dual cannot be cast as expectation maximization.
- Method: Kernel SGD updates the potentials from i.i.d. samples drawn from µ ⊗ ν and represents each iterate as a finite kernel expansion.Projection onto a centered ball ensures convergence when the dual solutions lie in the chosen RKHSs and the radius is sufficiently large.
- Implementation: Both potentials are approximated by linear combinations of kernel functions, but evaluating an iterate has complexity O(k^2).Random Fourier features and incremental incomplete Cholesky decomposition are cited as possible ways to reduce kernel-method runtime.
- Kernel choice: Universal kernels can approximate arbitrary potentials, while Gaussian kernels require careful bandwidth tuning for good convergence.For Euclidean domains, the paper gives a Gaussian kernel with bandwidth σ as a natural universal-kernel choice.
- Regularization: Because kernel SGD operates on the dual, the primal regularizer can be replaced by divergences such as χ2 rather than only KL regularization.This flexibility does not apply to the semi-discrete methods, for which entropy regularization is instrumental.
- Numerical illustrations: In a one-dimensional Gaussian-to-Gaussian-mixture experiment, RKHS iterates converge slowly but noticeably toward a semi-discrete proxy solution.The iterates converge faster where µ has more mass, where the potential contributes most to the objective.
Conclusion
The paper argues that stochastic optimization substantially reduces the computational burden of regularized OT and enables methods for semi-discrete and continuous settings. It reports improved discrete efficiency and convergent approaches for the other two settings.
- Conclusion: Stochastic optimization can considerably alleviate the computations required for regularized optimal transport or enable them in otherwise challenging settings.The conclusion covers discrete, semi-discrete, and continuous OT.
- Conclusion: In discrete OT, incremental gradient methods can surpass Sinkhorn in efficiency when the constant stepsize is correctly selected.The stated condition is that the stepsize selection should be feasible in practical applications.
- Conclusion: The paper presents the first known methods addressing semi-discrete and continuous OT, opening perspectives for applying OT to high-dimensional problems.This consequence is stated across the three settings considered by the paper.
A Convergence of (Sε) as ε →0
The paper studies how the semi-dual solution behaves as entropic regularization vanishes when the second measure is discrete. Under integrability assumptions, normalized regularized solutions are bounded, and every convergent subsequence solves the unregularized semi-dual problem.
- Novelty: The result fills a stated gap because the behavior of (Sε) had not previously been studied, unlike convergence results for (Pε) and discrete (Dε).The proposed theorem targets the discrete-ν setting where the semi-dual formulation is most advantageous.
- Proposition A.1: For discrete ν and integrable costs c(·, y), every convergent subsequence of normalized solutions vε as ε →0 solves the unregularized problem (S0).The normalization fixes vε(x0) = 0, and the proposition also establishes boundedness of the family.
- Proof strategy: The proof establishes pointwise convergence Hε → H0 and uses boundedness to extract a convergent subsequence of regularized solutions.The integrability condition is applied to obtain the required convergence and bounds.
- Normalization: Regularized potentials are unique up to an additive constant, so the proof fixes one potential value to zero before taking limits.This normalization supports boundedness and subsequence extraction.
- Proof strategy: The proof compares the regularized objective at vε with values at an optimal unregularized candidate v⋆ to show convergence of objective values.The argument concludes Hε(vε) → H0(v̄), where v̄ is a limit point.