Low-Rank Sinkhorn Factorization

TL;DR

This paper introduces a novel low-rank factorization approach for optimal transport problems that imposes low-rank constraints directly on couplings, providing a flexible and efficient alternative to kernel approximation methods.

Contribution

It proposes a general algorithm for low-rank constrained optimal transport with arbitrary costs, using an explicit factorization and alternating updates, with proven convergence.

Findings

Algorithm converges non-asymptotically to a stationary point.

Demonstrates efficiency on benchmark experiments.

Applicable to general cost functions in OT problems.

Abstract

Several recent applications of optimal transport (OT) theory to machine learning have relied on regularization, notably entropy and the Sinkhorn algorithm. Because matrix-vector products are pervasive in the Sinkhorn algorithm, several works have proposed to \textit{approximate} kernel matrices appearing in its iterations using low-rank factors. Another route lies instead in imposing low-rank constraints on the feasible set of couplings considered in OT problems, with no approximations on cost nor kernel matrices. This route was first explored by Forrow et al., 2018, who proposed an algorithm tailored for the squared Euclidean ground cost, using a proxy objective that can be solved through the machinery of regularized 2-Wasserstein barycenters. Building on this, we introduce in this work a generic approach that aims at solving, in full generality, the OT problem under low-rank constraints with arbitrary costs. Our algorithm relies on an explicit factorization of low rank couplings as a product of \textit{sub-coupling} factors linked by a common marginal; similar to an NMF approach, we alternatively updates these factors. We prove the non-asymptotic stationary convergence of this algorithm and illustrate its efficiency on benchmark experiments.