Sparsity-Constrained Optimal Transport

by   Tianlin Liu, et al.

Regularized optimal transport (OT) is now increasingly used as a loss or as a matching layer in neural networks. Entropy-regularized OT can be computed using the Sinkhorn algorithm but it leads to fully-dense transportation plans, meaning that all sources are (fractionally) matched with all targets. To address this issue, several works have investigated quadratic regularization instead. This regularization preserves sparsity and leads to unconstrained and smooth (semi) dual objectives, that can be solved with off-the-shelf gradient methods. Unfortunately, quadratic regularization does not give direct control over the cardinality (number of nonzeros) of the transportation plan. We propose in this paper a new approach for OT with explicit cardinality constraints on the transportation plan. Our work is motivated by an application to sparse mixture of experts, where OT can be used to match input tokens such as image patches with expert models such as neural networks. Cardinality constraints ensure that at most k tokens are matched with an expert, which is crucial for computational performance reasons. Despite the nonconvexity of cardinality constraints, we show that the corresponding (semi) dual problems are tractable and can be solved with first-order gradient methods. Our method can be thought as a middle ground between unregularized OT (recovered in the limit case k=1) and quadratically-regularized OT (recovered when k is large enough). The smoothness of the objectives increases as k increases, giving rise to a trade-off between convergence speed and sparsity of the optimal plan.


page 16

page 17


Smooth and Sparse Optimal Transport

Entropic regularization is quickly emerging as a new standard in optimal...

Quadratically-Regularized Optimal Transport on Graphs

Optimal transportation provides a means of lifting distances between poi...

Informative GANs via Structured Regularization of Optimal Transport

We tackle the challenge of disentangled representation learning in gener...

On the Convergence of Gradient Extrapolation Methods for Unbalanced Optimal Transport

We study the Unbalanced Optimal Transport (UOT) between two measures of ...

Regularized Discrete Optimal Transport

This article introduces a generalization of the discrete optimal transpo...

Algorithms for Weak Optimal Transport with an Application to Economics

The theory of weak optimal transport (WOT), introduced by [Gozlan et al....

Rethinking Initialization of the Sinkhorn Algorithm

Computing an optimal transport (OT) coupling between distributions plays...

Please sign up or login with your details

Forgot password? Click here to reset