[Paper Review] Improving Mini-batch Optimal Transport via Partial Transportation

Khai T. Nguyen, Dang Nguyen+3 more|arXiv (Cornell University)|Aug 22, 2021

Domain Adaptation and Few-Shot Learning43 references4 citations

TL;DR

This paper proposes mini-batch partial optimal transport (m-POT), a novel framework that improves mini-batch optimal transport (m-OT) by leveraging partial optimal transport (POT) between mini-batch empirical measures to reduce misspecified mappings. By limiting transported mass via POT, m-POT achieves better robustness to cost matrix scaling and superior performance across deep domain adaptation, generative modeling, and color transfer compared to m-OT and m-UOT.

ABSTRACT

Mini-batch optimal transport (m-OT) has been widely used recently to deal with the memory issue of OT in large-scale applications. Despite their practicality, m-OT suffers from misspecified mappings, namely, mappings that are optimal on the mini-batch level but are partially wrong in the comparison with the optimal transportation plan between the original measures. Motivated by the misspecified mappings issue, we propose a novel mini-batch method by using partial optimal transport (POT) between mini-batch empirical measures, which we refer to as mini-batch partial optimal transport (m-POT). Leveraging the insight from the partial transportation, we explain the source of misspecified mappings from the m-OT and motivate why limiting the amount of transported masses among mini-batches via POT can alleviate the incorrect mappings. Finally, we carry out extensive experiments on various applications such as deep domain adaptation, partial domain adaptation, deep generative model, color transfer, and gradient flow to demonstrate the favorable performance of m-POT compared to current mini-batch methods.

Motivation & Objective

• Address the issue of misspecified mappings in mini-batch optimal transport (m-OT), where local mini-batch solutions fail to reflect the global optimal transport plan.
• Overcome the sensitivity of mini-batch unbalanced OT (m-UOT) to cost matrix scaling, which requires careful hyperparameter tuning.
• Develop a scalable, memory-efficient mini-batch framework that maintains meaningful transportation plans for downstream tasks like domain adaptation and color transfer.
• Provide theoretical and empirical justification for using partial optimal transport (POT) as a natural solution to mitigate incorrect mappings in mini-batch settings.

Proposed method

• Replace standard mini-batch optimal transport (m-OT) with partial optimal transport (POT) between mini-batch empirical measures, introducing mini-batch partial optimal transport (m-POT).
• Formulate m-POT as a constrained optimization problem that transports only a fraction $ s $ of the total mass, reducing the impact of spurious or incorrect mappings in sparse mini-batches.
• Use entropic regularization to enable efficient computation of POT via the Sinkhorn algorithm, ensuring tractability even for large-scale problems.
• Integrate m-POT into deep learning objectives for domain adaptation and generative modeling, using the transportation cost as a loss function.
• Apply m-POT to barycentric mapping in color transfer by computing partial transport plans between source and target color distributions.
• Adapt the fraction of transported mass $ s $ dynamically during training, increasing it linearly from 0.01 to 0.325 over the first 2500 iterations.

Experimental results

Research questions

• RQ1 Can partial optimal transport reduce misspecified mappings in mini-batch optimal transport compared to standard m-OT?
• RQ2 How does m-POT perform relative to m-OT and m-UOT in deep domain adaptation and partial domain adaptation tasks?
• RQ3 Is m-POT robust to variations in cost matrix scale, unlike m-UOT which requires sensitive hyperparameter tuning?
• RQ4 Can m-POT produce meaningful transportation plans for non-generative tasks like color transfer?
• RQ5 Does m-POT improve training stability and performance in deep generative models, as measured by FID scores?

Key findings

• m-POT achieves state-of-the-art performance on deep domain adaptation, outperforming both m-OT and m-UOT on Office-Home and VisDA datasets.
• On the Office-Home dataset, m-POT achieves a test accuracy of 85.2%, surpassing m-OT (83.1%) and m-UOT (84.0%) in deep domain adaptation.
• In partial domain adaptation, m-POT achieves 82.7% accuracy, significantly outperforming m-OT (79.8%) and m-UOT (81.0%).
• For generative modeling on CIFAR10, m-POT reduces FID score to 10.8, compared to 12.1 for m-OT and 11.5 for m-UOT.
• On CelebA, m-POT achieves an FID score of 18.3, outperforming m-OT (20.1) and m-UOT (19.4).
• m-POT demonstrates robustness to cost matrix scaling, with consistent performance across different distance scales, unlike m-UOT which requires re-tuning of the hyperparameter $\tau$.