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: Goals: This paper introduces a new optimal transport loss based on a minibatch computation in order to alleviate some weaknesses from the original minibatch OT formulation. The formulation treats minibatches as data and seek to transport the minibatches from the source distribution to the minibatches from the target distribution. By doing so, their method prevents some undesirable connections between data.