[Paper Review] Minibatch optimal transport distances; analysis and applications

Kilian Fatras, Younes Zine+4 more|arXiv (Cornell University)|Jan 5, 2021

Optimization and Search Problems74 references17 citations

TL;DR

This paper proposes a theoretical analysis of minibatch optimal transport (MBOT) distances, introducing a debiased MBOT estimator that recovers key distance properties lost in standard minibatching. It establishes unbiased gradients, concentration bounds, and demonstrates superior scalability and invariance in applications like GANs, color transfer, and mesh comparison, with runtime growing linearly instead of cubically with data size.

ABSTRACT

Optimal transport distances have become a classic tool to compare probability distributions and have found many applications in machine learning. Yet, despite recent algorithmic developments, their complexity prevents their direct use on large scale datasets. To overcome this challenge, a common workaround is to compute these distances on minibatches i.e. to average the outcome of several smaller optimal transport problems. We propose in this paper an extended analysis of this practice, which effects were previously studied in restricted cases. We first consider a large variety of Optimal Transport kernels. We notably argue that the minibatch strategy comes with appealing properties such as unbiased estimators, gradients and a concentration bound around the expectation, but also with limits: the minibatch OT is not a distance. To recover some of the lost distance axioms, we introduce a debiased minibatch OT function and study its statistical and optimisation properties. Along with this theoretical analysis, we also conduct empirical experiments on gradient flows, generative adversarial networks (GANs) or color transfer that highlight the practical interest of this strategy.

Motivation & Objective

• Address the lack of theoretical understanding of minibatch optimal transport (MBOT) in large-scale machine learning despite its widespread use.
• Investigate the statistical and optimization properties of MBOT when applied to various optimal transport kernels, including Wasserstein and Gromov-Wasserstein.
• Overcome the fundamental limitation that standard MBOT is not a proper distance by introducing a debiased MBOT estimator that restores key axioms like separability.
• Provide non-asymptotic deviation bounds between the minibatch estimator and its expectation, ensuring statistical reliability.
• Demonstrate the practical utility of MBOT and its debiased variant across diverse applications such as GAN training, gradient flows, color transfer, and mesh comparison.

Proposed method

• Define a formal framework for MBOT using sampling without replacement, enabling rigorous statistical analysis of estimators.
• Propose a debiased MBOT function that corrects the bias in standard minibatch averaging, restoring the separability property of a distance.
• Establish concentration bounds for MBOT estimators, showing convergence to the true expected OT distance as the number of minibatches increases.
• Prove that the gradients of the MBOT loss are unbiased, enabling stable stochastic optimization in deep learning.
• Implement and evaluate the method using the Python Optimal Transport (POT) toolbox on synthetic and real-world datasets.
• Apply the MBOT and debiased MBOT to gradient flows, GANs, color transfer, and Gromov-Wasserstein on meshes to validate empirical performance.

Experimental results

Research questions

• RQ1 How do different sampling strategies (with/without replacement) affect the statistical properties of minibatch optimal transport estimators?
• RQ2 Why does standard minibatch OT fail to satisfy the axioms of a distance, and can this be corrected?
• RQ3 What are the statistical convergence properties of minibatch OT estimators, and can non-asymptotic deviation bounds be established?
• RQ4 Does the minibatch OT loss support unbiased gradients for stochastic optimization in deep learning?
• RQ5 How does the performance of the proposed debiased MBOT compare to standard OT and other baselines in real-world applications like GANs and mesh analysis?

Key findings

• The standard minibatch OT estimator is not a proper distance due to loss of the separability axiom, but it remains an unbiased estimator of the expected OT distance.
• The proposed debiased MBOT function restores the separability property, making it a valid distance-like measure while preserving computational efficiency.
• The minibatch OT estimator exhibits a concentration bound around its expectation, with deviation decreasing as the number of minibatches increases.
• The gradients of the minibatch OT loss are unbiased, enabling stable training in stochastic optimization settings such as GANs.
• In experiments, the minibatch Gromov-Wasserstein (MBGW) loss achieves linear runtime scaling with data size, outperforming full GW by up to 100× on 10⁴-point datasets.
• The MBGW loss is rotationally invariant in practice, as shown in color transfer and spiral rotation experiments, where it maintains constant loss across angles.