On transportation of mini-batches: a hierarchical approach

Khai Nguyen, Dang Nguyen, Quoc Nguyen, Tung Pham, Hung Bui, Dinh Phung, Trung Le, Nhat Ho

Research output: Chapter in Book/Report/Conference proceedingConference PaperResearchpeer-review


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-batching 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 present the new algorithms of the BoMb-OT in various applications, such as deep generative models and deep domain adaptation. From extensive experiments, we observe that the BoMb-OT achieves a favorable performance in deep learning models such as deep generative models and deep domain adaptation. In other applications such as approximate Bayesian computation, color transfer, and gradient flow, the BoMb-OT also yields either a lower quantitative result or a better qualitative result than the m-OT.
Original languageEnglish
Title of host publicationProceedings of the 39th International Conference on Machine Learning 2022
EditorsKamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, Sivan Sabato
Place of PublicationLondon UK
PublisherProceedings of Machine Learning Research (PMLR)
Number of pages34
Publication statusPublished - 2022
EventInternational Conference on Machine Learning 2022 - Baltimore, United States of America
Duration: 17 Jul 202223 Jul 2022
Conference number: 396th
https://icml.cc/virtual/2022/index.html (Website)
https://proceedings.mlr.press/v162/ (Proceedings)


ConferenceInternational Conference on Machine Learning 2022
Abbreviated titleICML 2022
Country/TerritoryUnited States of America
Internet address

Cite this