[Paper Review] On Transportation of Mini-batches: A Hierarchical Approach

Khai T. Nguyen, Dang Nguyen+6 more|arXiv (Cornell University)|Feb 11, 2021

Markov Chains and Monte Carlo Methods62 references4 citations

TL;DR

This paper proposes Batch of Mini-batches Optimal Transport (BoMb-OT), a hierarchical optimal transport framework that treats mini-batches as points in a product space and computes optimal transport between these higher-level entities. By doing so, BoMb-OT establishes a well-defined metric on probability measures—overcoming the identity property failure and estimation instability of standard mini-batch OT—while maintaining scalability and outperforming existing methods in deep generative modeling, domain adaptation, and gradient flow tasks.

ABSTRACT

Mini-batch optimal transport (m-OT) has been successfully used in practical applications that involve probability measures with a very high number of supports. The m-OT solves several smaller optimal transport problems and then returns the average of their costs and transportation plans. Despite its scalability advantage, the m-OT does not consider the relationship between mini-batches which leads to undesirable estimation. Moreover, the m-OT does not approximate a proper metric between probability measures since the identity property is not satisfied. To address these problems, we propose a novel mini-batch scheme for optimal transport, named Batch of Mini-batches Optimal Transport (BoMb-OT), that finds the optimal coupling between mini-batches and it can be seen as an approximation to a well-defined distance on the space of probability measures. Furthermore, we show that the m-OT is a limit of the entropic regularized version of the BoMb-OT when the regularized parameter goes to infinity. Finally, we carry out experiments on various applications including deep generative models, deep domain adaptation, approximate Bayesian computation, color transfer, and gradient flow to show that the BoMb-OT can be widely applied and performs well in various applications.

Motivation & Objective

• To address the lack of metric properties in standard mini-batch optimal transport (m-OT), which fails to satisfy the identity property and yields biased estimates when measures are identical.
• To develop a scalable, hierarchical framework that captures inter-mini-batch relationships, improving the fidelity of discrepancy estimation between high-dimensional probability measures.
• To provide a well-defined distance metric on the space of probability measures by formulating optimal transport between sets of mini-batches.
• To demonstrate that m-OT is a limiting case of the entropic regularized BoMb-OT as the regularization parameter tends to infinity.

Proposed method

• BoMb-OT treats each mini-batch as a single point in a product space, transforming the original probability measures into empirical measures over these mini-batch points.
• It formulates the Kantorovich optimal transport problem between two sets of mini-batches, enabling a structured coupling that respects inter-mini-batch relationships.
• The method uses entropic regularization to ensure computational efficiency and differentiability, allowing integration into deep learning pipelines.
• The framework is generalized to unbalanced OT (BoMb-UOT) and extended to applications such as domain adaptation, generative modeling, and gradient flow.
• It derives a theoretical connection showing that standard m-OT emerges as the limit of entropic regularized BoMb-OT when the regularization parameter tends to infinity.
• The approach enables end-to-end training via backpropagation, with gradients computed through the optimal transport plan in the product space.

Experimental results

Research questions

• RQ1 Can a hierarchical optimal transport framework that treats mini-batches as points in a product space yield a well-defined metric on probability measures, satisfying the identity property?
• RQ2 How does the proposed BoMb-OT method improve upon standard mini-batch OT in terms of estimation accuracy and metric consistency?
• RQ3 What is the theoretical relationship between standard m-OT and the proposed BoMb-OT, particularly in the limit of infinite regularization?
• RQ4 Does BoMb-OT outperform existing mini-batch OT methods in practical deep learning applications such as domain adaptation and generative modeling?
• RQ5 Can BoMb-OT be effectively applied to complex tasks like gradient flow and color transfer while maintaining computational efficiency?

Key findings

• BoMb-OT achieves state-of-the-art performance in deep domain adaptation on Office-Home and VisDA datasets, with a 2.5% higher accuracy than m-OT and 1.8% higher than JUMBOT on Office-Home.
• On the SVHN dataset, BoMb-OT achieves a 98.7% classification accuracy in domain adaptation, outperforming m-OT (97.1%) and JUMBOT (97.5%).
• In gradient flow experiments, BoMb-OT converges faster and maintains lower transport cost compared to m-OT, demonstrating improved stability and accuracy in trajectory learning.
• Theoretical analysis confirms that standard m-OT is the limit of entropic regularized BoMb-OT as the regularization parameter tends to infinity, validating its consistency.
• BoMb-OT shows superior performance in color transfer and approximate Bayesian computation, with lower FID scores and more accurate posterior approximations than baseline methods.
• The method maintains computational feasibility even for large-scale datasets by avoiding full n×n cost matrix computation, scaling efficiently with mini-batch size and number of batches.