[Paper Review] Stochastic Optimization for Large-scale Optimal Transport

Aude Genevay, Marco Cuturi+2 more|arXiv (Cornell University)|May 27, 2016

Markov Chains and Monte Carlo MethodsMathematics20 references153 citations

TL;DR

The paper introduces stochastic optimization schemes to compute large-scale optimal transport distances across discrete, semi-discrete, and continuous settings, using dual formulations and entropic regularization to achieve provable convergence without discretization errors.

ABSTRACT

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.

Motivation & Objective

• Motivate efficient computation of optimal transport distances for large-scale distributions in machine learning.
• Develop stochastic optimization methods that operate via sampling from distributions, avoiding discretization.
• Provide algorithms with provable convergence for discrete, semi-discrete, and continuous OT settings.
• Show empirical comparisons where stochastic methods outperform traditional Sinkhorn-style solvers.

Proposed method

• Recast the dual OT problem as the maximization of an expectation to enable stochastic optimization (and semi-dual forms).
• Use entropic regularization to obtain a smooth dual allowing faster convergence (Sinkhorn-based techniques when appropriate).
• Propose SAG (stochastic averaged gradient) for the discrete OT setting to beat Sinkhorn in large-scale problems.
• Apply averaged SGD for semi-discrete OT to handle a discrete measure against a continuous one without discretizing the continuous density.
• For continuous-continuous OT, expand dual variables in an RKHS and apply kernel SGD, yielding convergence to the dual solution within RKHS.
• Provide algorithms with convergence guarantees and discuss practicalities like mini-batching, step sizes, and RKHS projections.

Experimental results

Research questions

• RQ1 Can stochastic optimization methods efficiently compute OT distances for large-scale discrete distributions and overcome the bottlenecks of Sinkhorn?
• RQ2 How can dual formulations and entropic regularization be leveraged to handle semi-discrete OT without discretization errors?
• RQ3 Is it feasible to solve OT distances between two continuous densities using stochastic optimization in an RKHS framework?
• RQ4 What are the convergence properties and practical performance (in terms of speed and accuracy) of SAG, SGD, and kernel SGD in the OT setting?
• RQ5 How do these stochastic methods compare empirically to state-of-the-art discrete OT solvers across discrete, semi-discrete, and continuous benchmarks?

Key findings

• Incremental stochastic optimization (SAG) can outperform Sinkhorn on large-scale discrete OT problems.
• Averaged SGD for semi-discrete OT yields convergence rates suitable for problems where one distribution is continuous and the other is discrete.
• Kernel SGD in RKHS provides a convergent approach to OT between two continuous densities, the first of its kind among practical methods besides finite-sample discretization.
• Entropic regularization enables smooth duals that facilitate stochastic optimization with provable convergence.
• Empirical tests on word embeddings and word mover’s distances show faster convergence than Sinkhorn in large-scale discrete settings, and the methods scale well on GPU hardware.