On Transportation of Mini-batches: A Hierarchical Approach

1 INTRODUCTION . Optimal transport ( OT ) ( Villani , 2021 ; 2008 ; Peyré et al. , 2019 ) has emerged as an efficient tool in dealing with problems involving probability measures . Under the name of Wasserstein distance , OT has been widely utilized to solve problems such as generative modeling ( Arjovsky et al. , 2017 ; Tolstikhin et al. , 2018 ; Salimans et al. , 2018 ; Genevay et al. , 2018 ; Liutkus et al. , 2019 ) , barycenter problem ( Ho et al. , 2017 ; Li et al. , 2020 ) , and approximate Bayesian computation ( Bernton et al. , 2019a ; Nadjahi et al. , 2020 ) . Furthermore , OT can also provide the most economical map of moving masses between probability measures , which is very useful in various tasks such as color transfer ( Ferradans et al. , 2014 ; Perrot et al. , 2016 ) , natural language procesing ( Alvarez-Melis & Jaakkola , 2018 ) , and domain adaptation ( alignment ) ( Courty et al. , 2016 ; Lee et al. , 2019a ; Xu et al. , 2020 ) , and graph processing ( Titouan et al. , 2019 ; Xu et al. , 2019 ; Chen et al. , 2020 ) . Although OT has attracted growing attention in recent years , a major barrier that prevents OT from being ubiquitous is its heavy computational cost . When the two probability measures are discrete with n supports , solving the Wasserstein distance via the interior point methods has the complexity of O ( n3 log n ) ( Pele & Werman , 2009 ) , which is extremely expensive when n is large . There are two main lines of works that focus on easing this computational burden . The first approach is to find a good enough approximation of the solution by adding an entropic regularized term on the objective function ( Cuturi , 2013 ) . Several works ( Altschuler et al. , 2017 ; Lin et al. , 2019 ) show that the entropic approach can produce a ε-approximated solution at the same time reduces the computational complexity to O ( n2/ε2 ) . The second line of works named the “ slicing '' approach is based on the closed-form of Wasserstein distance in one-dimensional space , which has the computational complexity of order O ( n log n ) . There are various variants in this directions ; i.e. , ( Bonneel et al. , 2015 ; Deshpande et al. , 2019 ; Kolouri et al. , 2019 ; Nguyen et al. , 2021a ; b ) , these all belong to the family of sliced Wasserstein distances . Recently , some methods are proposed to combine the dimensional reduction approach with the entropic solver of Wasserstein distance to inherit advantages from both directions ( Paty & Cuturi , 2019 ; Muzellec & Cuturi , 2019 ; Lin et al. , 2020a ; b ) . Although those works have reduced the computational cost of OT considerably , computing OT is nearly impossible in the big data settings where n could be as large as few millions . In particular , solving OT requires to compute and store a n × n cost matrix that is impractical with current computational devices . Especially in deep learning applications , both supports of empirical measures and the cost matrix must be stored in a same device ( e.g . a GPU ) for automatic differentiation . This problem exists in all variants of OT including Wasserstein distance , entropic Wasserstein , and sliced Wasserstein distance . Therefore , it leads to the development of the mini-batches method for OT ( Genevay et al. , 2018 ; Sommerfeld et al. , 2019 ) , which we refer to as mini-batch OT loss . The main idea of the mini-batch method is to divide the original samples into multiple subsets ( mini-batches ) , in the hope that each pair of subsets ( mini-batches ) could capture some structures of the two probability measures , meanwhile , the computing OT cost between two mini-batches is cheap due to a very small size of mini-batches . Then the overall loss is defined as the average of distances between pairs of mini-batches . This scheme was applied for many forms of Wasserstein distances ( Deshpande et al. , 2018 ; Bhushan Damodaran et al. , 2018 ; Kolouri et al. , 2018 ; Salimans et al. , 2018 ) , and was theoretically studied in the works of ( Bellemare et al. , 2017 ; Bernton et al. , 2019b ; Nadjahi et al. , 2019 ) . Recently , Fatras et al . ( Fatras et al. , 2020 ; 2021b ) formulated this approach by giving a formal definition of the mini-batch OT loss , studying its asymptotic behavior , and investigating its gradient estimation properties . Despite being applied successfully , the current mini-batch OT loss does not consider the relationship between mini-batches and treats every pair of mini-batches the same . This causes undesirable effects in measuring the discrepancy between probability measures . First , the m-OT loss is shown to be an approximation of a discrepancy ( the population m-OT ) that does not preserve the metricity property , namely , this discrepancy is always positive even when two probability measures are identical . Second , it is also unclear whether this discrepancy achieves the minimum value when the two probability measures are the same . That naturally raises the question if we could propose a better mini-batching scheme to sort out these issues in order to improve the performance of the OT in practical applications . Contribution : In this paper , we propose a novel mini-batching scheme for optimal transport , which is named as Batch of Mini-batches Optimal Transport ( BoMb-OT ) . In particular , the BoMb-OT views every mini-batch as a point in the product space , then a set of mini-batches could be considered as an empirical measure . We now could employ the Kantorovich formulation between these two empirical measures in the product space as a discrepancy between two sets of mini-batches . In summary , our main contributions are two-fold : 1 . First , the BoMb-OT could provide a more similar transportation plan to the original OT than the m-OT , which leads to a more meaningful discrepancy using mini-batches . In particular , we prove that the BoMb-OT approximates a well-defined metric on the space of probability measures , named population BoMb-OT . Furthermore , the entropic regularization version of population BoMb-OT could be employed as a generalized version of the population m-OT . Specifically , when the regularization parameter in the entropic population BoMb-OT goes to infinity , its value approaches the value of the population m-OT . 2 . Second , we present the implementation strategy of the BoMb-OT and detailed algorithms in various applications in Appendix C. We then demonstrate the favorable performance of the BoMb-OT over the m-OT in two main applications that using optimal transport losses , namely , deep generative models and deep domain adaptation . Moreover , we also compare BoMb-OT to m-OT in other applications , such as sample matching , approximate Bayesian computation , color transfer , and gradient flow . In all applications , we also provide a careful investigation of the effects of two hyper-parameters of the mini-batching scheme , which are the number of mini-batches and the size of mini-batches , on the performance of the BoMb-OT and the m-OT . Organization : The remainder of the paper is organized as follows . In Section 2 , we provide backgrounds for optimal transport distances and the conventional mini-batching scheme ( m-OT ) . In Section 3 , we define the new mini-batching scheme for optimal transport distances , Batch of mini-batches Optimal Transport , and derive some of its theoretical properties . Section 4 benchmarks the proposed mini-batch scheme by extensive experiments on large-scale datasets , and followed by discussions in Section 5 . Finally , proofs of key results and extra materials are in the supplementary . Notation : For any probability measure µ on the Polish measurable space ( X , Σ ) , we denote ⊗m µ ( m ≥ 2 ) as the product measure on the product measurable space ( Xm , Σm ) . For any p ≥ 1 , we define Pp ( RN ) as the set of Borel probability measures with finite p-th moment defined on a given metric space ( RN , ‖.‖ ) . To simplify the presentation , we abuse the notation by using both the notation Xm for both the random vector ( x1 , . . . , xm ) ∈ Xm and the set { x1 , . . . , xm } , and we define by PXm : = 1m ∑m i=1 δxi the empirical measure ( the mini-batch measure ) associated with X m. For any set Xn : = { x1 , . . . , xn } and m ≥ 1 , we denote by [ Xn ] m the product set of Xn taken m times and ( Xn m ) the set of all m-element subsets of Xn . 2 BACKGROUND ON MINI-BATCH OPTIMAL TRANSPORT . In this section , we first review the definitions of the Wasserstein distance , the entropic Wasserstein , and the sliced Wasserstein . We then review the definition of the mini-batch optimal transport ( m-OT ) . 2.1 WASSERSTEIN DISTANCE AND ITS VARIANTS . We first start with the definition of Wasserstein distance and its variants . Let µ and ν be two probability measures on Pp ( RN ) . The Wasserstein p-distance between µ and ν is defined as follows : Wp ( µ , ν ) : = minπ∈Π ( µ , ν ) [ Eπ ( x , y ) ‖x− y‖p ] 1 p , where Π ( µ , ν ) : = { π : ∫ πdx = ν , ∫ πdy = µ } is the set of transportation plans between µ and ν . The entropic regularized Wasserstein to approximate the OT solution ( Altschuler et al. , 2017 ; Lin et al. , 2019 ) between µ and ν is defined as follows ( Cuturi , 2013 ) : W τp ( µ , ν ) : = minπ∈Π ( µ , ν ) { [ Eπ ( x , y ) ‖x − y‖p ] 1 p + τKL ( π|µ ⊗ ν ) } , where τ > 0 is a chosen regularized parameter and KL denotes the Kullback-Leibler divergence . Finally , the sliced Wasserstein ( SW ) ( Bonnotte , 2013 ; Bonneel et al. , 2015 ) is motivated by the closed-form of the Wasserstein distance in one-dimensional space . The formal definition of SW is : SWp ( µ , ν ) : = [ Eθ∼U ( SN−1 ) W pp ( θ ] µ , θ ] ν ) ] 1 p , where U ( SN−1 ) denotes the uniform measure over the ( N − 1 ) -dimensional unit hypersphere and θ ] is the orthogonal projection operator on direction θ . 2.2 MINI-BATCH OPTIMAL TRANSPORT . In this section , we first discuss the memory issue of large-scale optimal transport and challenges of dual solver . Then , we revisit the mini-batch OT loss that has been used in training deep generative models , domain adaptations , color transfer , and approximate Bayesian computation ( Bhushan Damodaran et al. , 2018 ; Genevay et al. , 2018 ; Tolstikhin et al. , 2018 ; Fatras et al. , 2020 ; Bernton et al. , 2019a ) . To ease the presentation , we are given Xn : = { x1 , . . . , xn } , Y n : = { y1 , . . . , yn } i.i.d . samples from µ and ν in turn . Let µn : = 1n ∑n i=1 δxi and νn : = 1 n ∑n i=1 δyi be two corresponding empirical measures from the whole data set . Here , n is usually large ( e.g. , millions ) and each support in Xn , Y n can be a high dimensional data point ( e.g . a high resolution image , video , etc ) . Memory issue of optimal transport : Using an OT loss between µn and νn needs to compute and store a n × n cost matrix which has one trillion float entries ( about 4 terabytes ) when n is about millions . Moreover , when dealing with deep neural networks , both support points and the cost matrix are required to be stored in the same memory ( e.g. , a GPU with 8 gigabytes memory ) as a part of the computational graph for automatic differentiation . This issue applies to all variants of OT losses , such as Wasserstein distance , entropic Wasserstein , sliced Wasserstein distance . Therefore , it is nearly impossible to compute OT and its variants in large-scale applications . Challenges of stochastic dual solver : Using stochastic optimization to solve the Kantorovich dual form is a possible approach to deal with large-scale OT , i.e . Wasserstein GAN ( Arjovsky et al. , 2017 ; Leygonie et al. , 2019 ) . However , the obtained distance has been shown to be very different from the original Wasserstein distance ( Mallasto et al. , 2019 ; Stanczuk et al. , 2021 ) . Using input convex neural networks is another choice to approximate the Brenier potential ( Makkuva et al. , 2020 ) . Nevertheless , recent work ( Korotin et al. , 2021 ) has indicated that input convex neural networks are not sufficient ( have limited power ) in approximating the Brenier potential . Furthermore , both mentioned approaches are restricted in the choice of ground metric . In particular , L1 norm is used in Wasserstein GAN to make the constraint of dual form into the Lipchitz constraint and L2 norm is for the existence of the Brenier potential . Mini-batch solution : As a popular alternative approach for stable large-scale OT with flexible choices of the ground metric , the mini-batch optimal transport is proposed ( Genevay et al. , 2018 ; Sommerfeld et al. , 2019 ) and has been widely used in various applications ( Arjovsky et al. , 2017 ; Deshpande et al. , 2018 ; Sommerfeld et al. , 2019 ; Bhushan Damodaran et al. , 2018 ) . In this approach , the original n samples are divided into subsets ( mini-batches ) of size m , where m is often the largest number that the computer can process , then an alternative solution of the original OT problem is formed by aggregating these smaller OT solutions from mini-batches . We now state an adapted definition of mini-batch OT ( m-OT ) scheme that was formulated in ( Fatras et al. , 2020 ) , including its two key parts : its transportation cost and its transportation plan . Definition 1 ( Empirical m-OT ) . For p ≥ 1 and integers m ≥ 1 , and k ≥ 1 , let d : Pp ( X ) × Pp ( X ) → [ 0 , ∞ ) be a function , i.e. , { Wp , W τp , SWp } . Then , the mini-batch OT ( m-OT ) loss and the transportation plan , a n× n matrix , between µn and νn are defined as follows : Ûk , md ( µn , νn ) : = 1 k2 k∑ i=1 k∑ j=1 d ( PXmi , PYmj ) ; π̂ m k ( µn , νn ) : = 1 k2 k∑ i=1 k∑ j=1 πXmi , Ymj , ( 1 ) whereXmi is sampled i.i.d from ( Xn m ) , Y mi is sampled i.i.d from ( Y n m ) , and the transport plan πXmi , Ymj , where its entries equal zero except those indexed of samples Xmi × Y mj , is the transportation plan when d ( PXmi , PYmj ) is an optimal transport metric . The above definition was generalized to the two original measures as follows in ( Fatras et al. , 2020 ) : Definition 2 ( Population m-OT ) . Assume that µ and ν are two probability measures on Pp ( X ) for given positive integers p ≥ 1 , m ≥ 1 , and d : Pp ( X ) × Pp ( X ) → [ 0 , ∞ ) be a given function . Then , the population mini-batch OT ( m-OT ) discrepancy between µ and ν is defined as follows : Umd ( µ , ν ) : = E ( Xm , Ym ) ∼⊗mµ ⊗⊗mν d ( PXm , PYm ) . ( 2 ) Issues of the m-OT : From Definition 1 , the m-OT treats every pair of mini-batches the same by taking the average of the m-OT loss between any pair of them for both transportation loss and transportation plan . This treatment has some issues . First , a mini-batch is a sparse representation of the true distribution and two sparse representations of the same distribution could be very different from each other . Hence , a mini-batch from Xm would prefer to match to certain mini-batches of Y m , rather than treating every mini-batch of Y m equally . For example , each mini-batch has one datum , then each term in the population m-OT now is the ground metric of the OT cost , the population m-OT degenerates to E [ d ( Xm , Y m ) ] for independent Xm and Y m. This treatment also leads to an uninformative transportation plan shown in Figure 15 , which is followed by a less meaningful transportation cost . Second , although it has been proved that the population m-OT is symmetric and positive ( Fatras et al. , 2020 ) , for the same reason it does not vanish when two measures are identical .

Summary: This paper proposed Batch of Mini-batches Optimal Transport (BoMb-OT) method, which finds the optimal coupling between mini-batches in mini-batch optimal transport (m-OT), which is achieved by solving another OT problem over the mini-batches.  The authors claimed that doing this will capture the relation between different mini-batches better. They firstly proved that BoMb-OT approximates the population BoMb-OT metric in probability measure space with and without entropic regularization.  The authors then implemented the proposed BoMb-OT method and applied it on deep generative models and deep domain adaptation, showing that BoMb-OT has favourable performance over m-OT.