Overrelaxed Sinkhorn-Knopp Algorithm for Regularized Optimal Transport

This article describes a method for quickly computing the solution to the regularized optimal transport problem. It generalizes and improves upon the widely-used iterative Bregman projections algorithm (or Sinkhorn-Knopp algorithm). The idea is to overrelax the Bregman projection operators, allowing for faster convergence. In practice this corresponds to elevating the diagonal scaling factors to a given power, at each step of the algorithm. We propose a simple method for establishing global convergence by ensuring the decrease of a Lyapunov function at each step. An adaptive choice of overrelaxation parameter based on the Lyapunov function is constructed. We also suggest a heuristic to choose a suitable asymptotic overrelaxation parameter, based on a local convergence analysis. Our numerical experiments show a gain in convergence speed by an order of magnitude in certain regimes.

[1]  Marco Cuturi,et al.  Sinkhorn Distances: Lightspeed Computation of Optimal Transport , 2013, NIPS.

[2]  Jean-Luc Starck,et al.  Wasserstein Dictionary Learning: Optimal Transport-based unsupervised non-linear dictionary learning , 2017, SIAM J. Imaging Sci..

[3]  Louis A. Hageman,et al.  Iterative Solution of Large Linear Systems. , 1971 .

[4]  Philip A. Knight,et al.  The Sinkhorn-Knopp Algorithm: Convergence and Applications , 2008, SIAM J. Matrix Anal. Appl..

[5]  Alexandre d'Aspremont,et al.  Regularized nonlinear acceleration , 2016, Mathematical Programming.

[6]  Alexandre d'Aspremont,et al.  Online Regularized Nonlinear Acceleration , 2018 .

[7]  Gabriel Peyré,et al.  Iterative Bregman Projections for Regularized Transportation Problems , 2014, SIAM J. Sci. Comput..

[8]  Marco Cuturi,et al.  Principal Geodesic Analysis for Probability Measures under the Optimal Transport Metric , 2015, NIPS.

[9]  François-Xavier Vialard,et al.  Quantum Optimal Transport for Tensor Field Processing , 2016, ArXiv.

[10]  Leonidas J. Guibas,et al.  The Earth Mover's Distance as a Metric for Image Retrieval , 2000, International Journal of Computer Vision.

[11]  Klaus-Robert Müller,et al.  Wasserstein Training of Restricted Boltzmann Machines , 2016, NIPS.

[12]  Nicolas Courty,et al.  Optimal Transport for Domain Adaptation , 2014, IEEE Transactions on Pattern Analysis and Machine Intelligence.

[13]  Nicolas Papadakis,et al.  Regularized Optimal Transport and the Rot Mover's Distance , 2016, J. Mach. Learn. Res..

[14]  Jason Altschuler,et al.  Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration , 2017, NIPS.

[15]  A. Hadjidimos,et al.  On the Optimization of the Classical Iterative Schemes for the Solution of Complex Singular Linear Systems , 1985 .

[16]  S. K. Zavriev,et al.  Heavy-ball method in nonconvex optimization problems , 1993 .

[17]  Max-K. von Renesse,et al.  A note on overrelaxation in the Sinkhorn algorithm , 2020, Optimization Letters.

[18]  L. Richardson,et al.  On the Approximate Arithmetical Solution by Finite Differences of Physical Problems Involving Differential Equations, with an Application to the Stresses in a Masonry Dam , 1910 .

[19]  Julien Rabin,et al.  Non-convex Relaxation of Optimal Transport for Color Transfer Between Images , 2015, GSI.

[20]  Alain Rakotomamonjy,et al.  Screening Sinkhorn Algorithm for Regularized Optimal Transport , 2019, NeurIPS.

[21]  Boris Polyak Some methods of speeding up the convergence of iteration methods , 1964 .

[22]  Lénaïc Chizat,et al.  Scaling Algorithms for Unbalanced Transport Problems , 2016, 1607.05816.

[23]  Hossein Mobahi,et al.  Learning with a Wasserstein Loss , 2015, NIPS.

[24]  Gabriel Peyré,et al.  Fast Dictionary Learning with a Smoothed Wasserstein Loss , 2016, AISTATS.

[25]  Euhanna Ghadimi,et al.  Global convergence of the Heavy-ball method for convex optimization , 2014, 2015 European Control Conference (ECC).

[26]  Peter Ochs,et al.  Local Convergence of the Heavy-Ball Method and iPiano for Non-convex Optimization , 2016, J. Optim. Theory Appl..

[27]  Richard Sinkhorn A Relationship Between Arbitrary Positive Matrices and Doubly Stochastic Matrices , 1964 .

[28]  L. Bregman The relaxation method of finding the common point of convex sets and its application to the solution of problems in convex programming , 1967 .

[29]  Michael I. Jordan,et al.  On the Efficiency of the Sinkhorn and Greenkhorn Algorithms and Their Acceleration for Optimal Transport , 2019 .