[Paper Review] Unbalanced minibatch Optimal Transport; applications to Domain Adaptation

Kilian Fatras, Thibault Séjourné+2 more|arXiv (Cornell University)|Mar 5, 2021

Domain Adaptation and Few-Shot LearningComputer Science72 references34 citations

TL;DR

The paper proposes unbalanced minibatch optimal transport (MBOT) to improve robustness of minibatch OT, develops its theoretical properties, and demonstrates superior domain adaptation performance over standard OT and related baselines.

ABSTRACT

Optimal transport distances have found many applications in machine learning for their capacity to compare non-parametric probability distributions. Yet their algorithmic complexity generally prevents their direct use on large scale datasets. Among the possible strategies to alleviate this issue, practitioners can rely on computing estimates of these distances over subsets of data, {\em i.e.} minibatches. While computationally appealing, we highlight in this paper some limits of this strategy, arguing it can lead to undesirable smoothing effects. As an alternative, we suggest that the same minibatch strategy coupled with unbalanced optimal transport can yield more robust behavior. We discuss the associated theoretical properties, such as unbiased estimators, existence of gradients and concentration bounds. Our experimental study shows that in challenging problems associated to domain adaptation, the use of unbalanced optimal transport leads to significantly better results, competing with or surpassing recent baselines.

Motivation & Objective

• Motivate and address limitations of standard minibatch OT in large-scale learning and domain adaptation.
• Introduce an unbalanced OT formulation at the minibatch level to reduce spurious sample pairings.
• Provide theoretical guarantees: existence of gradients, unbiased estimators, and concentration bounds.
• Show empirically that Unbalanced MBOT yields robust domain adaptation performance against challenging datasets.

Proposed method

• Define minibatch unbalanced OT (UOT) as a relaxation of marginal constraints using Csiszár divergences.
• Adopt a minibatch estimator that averages UOT costs over all minibatch pairs, plus an incomplete (k-sample) estimator to reduce computation.
• Prove boundedness of the UOT cost and compactness of optimal transports, enabling concentration bounds.
• Establish unbiased gradient properties via Clarke regularity to enable SGD optimization.
• Demonstrate that MBOT mitigates undesirable sample couplings in minibatches and closely approximates unbalanced OT with scalable cost.
• Apply the MBOT framework to a domain adaptation method (jumbot) that jointly optimizes embeddings and label mappings under an unbalanced minibatch OT loss.

Experimental results

Research questions

• RQ1 How does the standard minibatch OT behave under sampling and outliers, and can relaxing marginal constraints improve robustness?
• RQ2 What are the statistical and optimization properties of minibatch unbalanced OT (existence, unbiased gradients, concentration)?
• RQ3 Can unbalanced MBOT provide improved domain adaptation performance compared to balanced OT and other OT variants?
• RQ4 What is the practical impact of MBOT on gradient-based learning and large-scale neural network training for domain adaptation?
• RQ5 How does unbalanced MBOT perform on benchmark DA datasets (digits, Office-Home, VisDA) relative to state-of-the-art methods?

Key findings

• Unbalanced MBOT yields a more robust transport plan at the minibatch level, reducing spurious cross-class pairings seen with standard MBOT.
• Theoretical results establish finite UOT costs, bounded optimal plans, and concentration bounds for estimators, independent of data dimension.
• MBOT provides unbiased gradient estimators under Clarke regularity, enabling SGD optimization for learning models.
• Empirical studies on domain adaptation problems show that UOT-based approaches can compete with or surpass recent baselines, including state-of-the-art OT-based methods.
• Gradients derived from MBOT do not introduce bias in the optimization, supporting reliable training of neural network embeddings and classifiers.
• The unbalanced formulation mitigates the sensitivity of OT to outliers and sampling, improving stability in gradient flows and DA tasks.