Year: 2021

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-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.

heading: 1 INTRODUCTION
text: 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.

heading: 2 BACKGROUND ON MINI-BATCH OPTIMAL TRANSPORT
text: 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).

heading: 2.1 WASSERSTEIN DISTANCE AND ITS VARIANTS
text: 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 θ.

heading: 2.2 MINI-BATCH OPTIMAL TRANSPORT
text: 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.

heading: 3 BATCH OF MINI-BATCHES OPTIMAL TRANSPORT
text: To address the issues of m-OT, in this section, we propose a novel mini-batch scheme, named batch of mini-batches OT (BoMb-OT). We first demonstrate the improvement of the BoMb-OT to the m-OT and discuss the practical usage of the BoMb-OT. Then, we prove that the BoMb-OT loss is an approximation of a well-defined metric in the space of Borel probability measures. Finally, we discuss the entropic version of the population BoMb-OT (population eBoMb-OT) admits the population m-OT as a special case where the entropic regularization goes to infinity in Appendix B. To simplify the presentation, we assume throughout this section that µ and ν are continuous probability measures in Pp(X ) for some given p ≥ 1 and X ⊂ RN where N ≥ 1. Furthermore, d : Pp(X )× Pp(X )→ [0,∞) is a given divergence between probability measures in Pp(X ).

heading: 3.1 DEFINITION OF BOMB-OT AND ITS PROPERTIES
text: Intuitively, a mini-batch scheme in OT can be decomposed into two steps: (1) Solving local OT between every pair of mini-batches from samples of two original measures, (2) Combining local OT
solutions into a global solution. As discussed above, the main problem with the m-OT lies in the second step, all pairs of mini-batches play the same role in the total transportation loss, meanwhile, a mini-batch from Xn would prefer to match some certain mini-batches from Y n. To address the limitation, we enhance the second step by considering the optimal coupling between mini-batches. By doing so, we could define a new mini-batch scheme that is able to construct a good global mapping between samples that also leads to a meaningful objective loss.
Let Xm1 , X m 2 , . . . , X m k be mini-batches that are sampled with replacement from [X n]m and Y m1 , Y m 2 , . . . , Y m k be mini-batches that are sampled with replacement from [Y
n]m. The batch of mini-batches optimal transport is defined as follows: Definition 3 (Batch of mini-batches). Assume that p ≥ 1, m ≥ 1, k ≥ 1 are positive integers and let d : Pp(X )× Pp(X )→ [0,∞) be a given function (e.g., {Wp,W τp , SWp}). The BoMb-OT loss and the BoMb-OT transportation plan between probability measures µn and νn are given by:
D̂k,md (µn, νn) := min γ∈Π(⊗mµk , ⊗m νk ) k∑ i=1 k∑ j=1 γijd(PXmi , PYmj ); π̂ m k (µn, νn) := k∑ i=1 k∑ j=1 γ̄ijπXmi ,Ymj ,
where ⊗m µk := 1 k ∑k i=1 δXmi and ⊗m νk := 1 k ∑k j=1 δYmj are two empirical measures defined on the product space via mini-batches (measures over mini-batches), γ̄ is a k × k optimal transport plan between ⊗m µk and ⊗m νk , and πXmi ,Ymj is defined as in Definition 1.
Compared to the m-OT, the BoMb-OT considers an additional OT between two measures on minibatches for combining mini-batch losses. The improvement of the BoMb-OT over the m-OT is illustrated in Figure 1 in which we assume that we can only solve 2× 2 OT problems and we need to deal with the OT problems with two empirical measures of 3 supports. The optimal transportation plan is to map i-th square to i-th circle. Assume that there are four mini-batches: A1 and A2 from X and B1 and B2 from Y . The weight for the m-OT loss between A1 and B2 equals 14 , meanwhile, the map from A1 to B2 is unuseful. In this toy example, A1 should be mapped to B1 and A2 is for B2. It is also the solution that the BoMb-OT tries to obtain by re-weighting pairs of mini-batches through a transportation plan between two measures over mini-batches.
Practical usage of the BoMb-OT: There are three types of applications that the BoMb-OT can be utilized. The first one is gradient-based applications (deep learning) such as deep generative models, deep domain adaption, and gradient flow. In these applications, the goal is to estimate the gradient of a parameter of interest (a neural network) with an OT objective function. The second one is mapping-based applications, namely, we aim to obtain a transportation map between samples. Examples of mapping-based applications include color transfer, domain adaptation, and sample matching. Finally, the third type of applications is discrepancy-based applications where we need to know the discrepancy between two measures, e.g., approximate Bayesian computation, and searching problems. We now discuss the implementation of BoMb-OT in each scenario.
Gradient-based (deep learning) applications: For a better presentation, we assume that a system with a GPU (graphics processing unit) is used in the training process. We further assume that the GPU has enough space to store the neural net, the 2×m supports of two mini-batch measures, and the m ×m cost matrix. The algorithm of the BoMb-OT consists of three steps: Forwarding, Solving k × k OT, and Re-forwarding and Backwarding. In the first step, k2 OT problems of size m ×m are computed in turn on GPU to obtain the k × k cost matrix which indicates the transportation cost between each pair of mini-batch measures. We recall that it is only enough memory for computing m×m OT on GPU at a time. After that, an OT problem with the found k × k cost matrix is solved
to find the coupling between mini-batches. Here, k can be much larger than m since the k × k OT can be solved on computer memory (RAM - Random-access memory), which is bigger than GPU’s memory. We would like to emphasize that in the first two steps, automatic differentiation is not required. In the last step, we do re-forwarding and do backwarding (back-propagation) to obtain the gradients of each pair of mini-batches, then using the coupling in the second step to aggregate them into the final gradient signal of the parameter of interest. The detailed algorithms of the BoMb-OT for deep generative model and deep domain adaption are respectively described in Algorithm 1 in Appendix C.1 and Algorithm 2 in Appendix C.2. Compared to the m-OT, the BoMb-OT only costs an additional forwarding and an additional k × k OT problem that is not expensive since they do not require to create and store any computational graphs.
Mapping-based applications: In this type of application, the BoMb-OT also consists of three steps like in gradient-based applications. However, in the last step, the transportation plan (the mapping) is aggregated instead of the gradient. As an example, we present the BoMb-OT algorithm for color transfer in Algorithm 3 in Appendix C.3. We would like to emphasize that the last step (reforwarding) can be removed by storing all k2 mini-batch transportation plans of size m×m which can be represented by a sparse matrix of size n× n for saving memory. Discrepancy-based applications: The BoMb-OT needs only two steps in this scenario, namely, forwarding and solving k × k OT. Since we do not need to re-estimate any statistics from minibatches, the final value of the BoMb-OT can be evaluated at the end of the second step without the re-forwarding step. Similar to the previous two types of applications, we present the BoMb-OT algorithm for approximate Bayesian computation in Algorithm 5 in Appendix C.4. Compared to the m-OT, the BoMb-OT needs only an additional k × k OT here. Choosing k and m: In practice, we want to have as large as possible values of k and m. The mini-batch size m is often chosen to be the largest value of the computational memory that can be stored. For choosing k, there are two practical cases. Two (multiples) computational memories - this is the case of modern computational systems that have at least a CPU (central processing unit) with RAM and a GPU with its corresponding memory. As discussed in the practical usage of the BoMb-OT, OT problems between mini-batch measures of size m ×m are solved on the GPU for high computational speed in estimating the parameter of interest or other statistics. On the other hand, the k × k OT problems between measures over mini-batches can be computed on the CPU’s memory. So, k can be chosen larger than m. One (shared) computational memory - in this case, the largest value of k equals to m. Note that, smaller values of k and m can still be used for faster computation. We would like to recall that k is usually set to 1 in almost all recent applications of OT (Genevay et al., 2018; Bhushan Damodaran et al., 2018; Tolstikhin et al., 2018; Bernton et al., 2019a). In our experiments, we will demonstrate that increasing k could improve the performance.
Computational complexity of the BoMb-OT: The computational complexity of BoMb-OT depends on a particular choice of divergence d. In the experiments, we specifically consider three choices of d ∈ {Wp,W τp , SWp} where τ > 0 is a chosen regularized parameter. Here, we only discuss the case when d = Wp, the discussion of other choices of d is in Appendix B. When d is the entropic Wasserstein distance, for each pair Xmi and Y m j , the computational complexity of approximating d(PXmi , PYmj ) via the Sinkhorn algorithm is of order O(m 2/ε2) where ε > 0 is some desired accuracy (Lin et al., 2019). Hence, the computation cost of k2 OT transport plans from k2 pairs of mini-batches is of order O(k2m2/ε2) when using entropic regularization (see Appendix B for the definition). With another OT cost within k2 pairs, the total computational complexity of computing the BoMb-OT is O(k2(m2 + 1)/ε2). It means that the computational complexity of computing the BoMb-OT is comparable to the m-OT, which is O(k2m2/ε2). The BoMb-OT’s transportation plan: We discuss in detail the sparsity of BoMb-OT’s transportation plan in Appendix B. To visualize the BoMb-OT’s plan and compare it with the m-OT’s plan, we carry out a simple experiment on two measures of 10 supports in Figure 15. The details of the experiment and discussion are given in Appendix D.6. From this example, the BoMb-OT provides a more accurate transportation plan than the m-OT with various choices of m and k.

heading: 3.2 METRIC APPROXIMATION OF THE BOMB-OT LOSS
text: In this section, we show that the BoMb-OT is an approximation of a well-defined metric, named population BoMb-OT, in the space of probability measures. To ease the ensuing discussion, we first define that metric as follows: Definition 4 (Population BoMb-OT). Let µ and ν be two probability measures on Pp(X ), for p ≥ 1 is a positive number. The population BoMb-OT between probability measures µ and ν is defined as:
Dmd (µ, ν) := inf γ∈Π(⊗mµ ,⊗mν ) E(Xm,Ym)∼γ [d(PXm , PYm)].
Different from the m-OT in Definition 2, the optimal plan between two distributions on product spaces is crucial to guarantee that the population BoMb-OT is a well-defined metric in the space of probability measures. Theorem 1 (Metric property of population BoMb-OT). Assume that d is an invariant metric under permutation on Xm and the function d(PXm , PYm) is continuous in terms of Xm, Y m ∈ Xm. Then, the population BoMb-OT is a well-defined metric in the space of probability measures.
The proof of Theorem 1 is in Appendix A. The assumption of Theorem 1 is mild and satisfied when d is (sliced) Wasserstein metrics of order p, i.e., d ∈ {Wp, SWp}. Our next result shows the approximation error between the BoMb-OT loss and the population BoMbOT distance. We provide the following result with the approximation error when d ∈ {Wp, SWp}. Theorem 2 (Population BoMb-OT). Assume that p ≥ 1, X is a compact subset of RN , and all the possible mini-batches are considered. Then, we have (i)
∣∣D̂k,md (µn, νn) − Dmd (µ, ν)∣∣ = OP ( m1− 1 p /n 1 N ) when d ≡ Wp; (ii)
∣∣D̂k,md (µn, νn) − Dmd (µ, ν)∣∣ = OP (m1− 1p /n 12 ) when d ≡ SWp and k is the number of mini-batches.
The proof of Theorem 2 is in Appendix A. We would like to remark that the dependency of the sample complexity of the BoMb-OT on N is necessary when d ≡Wp. It is due to the inclusion of the additional optimal transport in the BoMb-OT. On the other hand, the curse of dimensionality of the BoMb-OT loss does not happen when d ≡ SWp. Furthermore, the choice d ≡ SWp improves not only the sample complexity but also the computational complexity of the BoMb-OT. In particular, in Appendix B, we demonstrate that the computational complexity of the BoMb-OT is O(k2m logm) when d ≡ SWp. Another potential scenario that the BoMb-OT does not have the curse of dimensionality is when we consider its entropic regularized version in equation 4. The entropic optimal transport had been shown to have sample complexity at the order n− 1 2 (Mena & Weed, 2019). In the case of the entropic regularized BoMb-OT (eBoMb-OT) (see equation (4) in Appendix B), when the regularized parameter λ is infinity, namely, the m-OT, the result of (Fatras et al., 2020; 2021b) already established the sample complexity n− 1 2 . However, for a general value of the regularized parameter, it is unclear whether we still have this sample complexity n− 1 2 of the eBoMb-OT. We leave this question for the future work.

heading: 4 EXPERIMENTS
text: In this section, we demonstrate the favorable performance of BoMb-OT compared to m-OT in three discussed types of applications, namely, gradient-based, mapping-based, and value-based applications. For gradient-based applications, we run experiments on deep generative model (Genevay et al., 2018; Deshpande et al., 2018), deep domain adaptation (Bhushan Damodaran et al., 2018). Experiments on color transfer (Rabin et al., 2014; Ferradans et al., 2014) are conducted as the example for mapping-based applications. Lastly, we present results on approximate Bayesian computation (Bernton et al., 2019a) which is a value-based application. In the main text of the paper, we only report and discuss the obtained experimental results, the details of applications and their algorithms are given in Appendix C. In Appendix D, we provide detailed experimental results of discussed applications including visualization and computational time. Moreover, we also carry out experiments on gradient flow (Santambrogio, 2017) and visualize mini-batch transportation matrices of the m-OT and the BoMb-OT. From all experiments, we observe that the BoMb-OT performs better than m-OT consistently. The detailed settings of our experiments including neural network architecture, hyperparameters, and evaluation metrics are in Appendix E. We would like to recall that k is the number of
mini-batches and m is the mini-batch size. In our experiments, we create mini-batches by sampling without replacement from supports of the original empirical measures.

heading: 4.1 DEEP GENERATIVE MODEL
text: We now show the deep generative model result on MNIST (LeCun et al., 1998), CIFAR10 (32x32) (Krizhevsky, 2009), and CelebA (64x64) (Liu et al., 2015) using different mini-batch schemes with two OT losses for d: SW2 (Deshpande et al., 2018) and W τ2 (Genevay et al., 2018). According to Table 1, the BoMb-OT always gives lower quantitative metrics than the m-OT on all datasets with all choices of the number of mini-batches k and mini-batch ground metric d. Interestingly, the eBoMb-OT which is the entropic regularization version of BoMb-OT (defined in Appendix B) is better than the BoMb-OT on CIFAR10, CelebA, and one setting on MNIST. The explanation for this phenomenon is that the eBoMb-OT loss is smoother than the BoMb-OT due to the entropic regularization (Genevay et al., 2018). From the table, we also observe that increasing the number of mini-batches k improves the result of generators that are trained with the same number of stochastic gradient updates. It suggests that the gradient of an OT loss should not be estimated by only 1 pair of mini-batches like the way it has been implementing in practice. Also, based on the experimental results, we recommend practitioners to consider using an additional transportation problem between measures over mini-batches when the number of mini-batches k > 1.

heading: 4.2 DEEP DOMAIN ADAPTATION
text: In this section, we adapt two digits datasets SVHN (Netzer et al., 2011) and USPS (Hull, 1994) to MNIST (LeCun et al., 1998). Details about architectures of neural networks and hyper-parameters settings are given in Appendix E.2. The classification results are illustrated in Table 2. The BoMb-OT and eBoMb-UOT always achieve better results than their counterparts for all choices of k. As k increases, the gap between two methods tends to become larger in both datasets. This demonstrates the effectiveness of our scheme, especially when dealing with a large number of mini-batches. We also observe that the classification accuracy on the target domain also increases as k changes from 8 to 32 and keeping the same number of stochastic gradient updates for deep neural networks. Thus, we suggest that DeepJDOT (Bhushan Damodaran et al., 2018) and its variants should be implemented with the BoMb-OT’s strategy which is presented in Algorithm 2 in Appendix C.2.
Source m-OT BoMb-OT full-OT Target

heading: 4.3 COLOR TRANSFER
text: Mini-batch OT has made color transfer to be able to transform color between two images that contain millions of pixels (Fatras et al., 2021b). In this section, we show that our new mini-batch strategy can improve further the quality of the color transfer. The details of the application and algorithm are given in Appendix C.3. The full experimental results including qualitative images and their color palettes with different choices of k and m are presented in Appendix D.3. Based on the experiment, we observe that the BoMb-OT provides a better barycentric mapping for color transfer than m-OT with every setting of k and m. In this main paper, we show a color transfer result in Figure 2. We can observe that the color of the transferred image of the BoMb-OT is more similar to the target image than one of the m-OT. The color palettes in Appendix C.3 also reinforce this claim.

heading: 4.4 APPROXIMATE BAYESIAN COMPUTATION (ABC)
text: As choices of sample acceptance rejection criteria, we compare the m-OT and the BoMb-OT in ABC. We present the detail of the application and the algorithm in Appendix C.4. The setting of the model and the full results are given in Appendix D.4. We compare the m-OT and the BoMb-OT with various choices of k and m. We compare them by plotting their approximated posteriors and computing the Wasserstein distance between approximated posteriors to the true posterior. According to the experiments, the BoMb-OT always yield lower Wasserstein distances than the m-OT. Here, we show some results in Figure 3. From the figure, we can see that approximated posteriors from the BoMb-OT are closer to the true posterior than the m-OT. Moreover, we observe that increasing the value of k and m improves the quality of the approximated posteriors from both mini-batch schemes considerably. Based on the fact that Wasserstein ABC is implemented with k = 1 (Bernton et al., 2019a), we suggest that the BoMb-OT should be used with k is greater than 1 in ABC.

heading: 5 CONCLUSION
text: In the paper, we have presented a novel mini-batching scheme for optimal transport, named Batch of Mini-batches Optimal Transport (BoMb-OT). The idea of the BoMb-OT is to consider the optimal transport problem on the space of mini-batches with a Wasserstein-types ground metric. We prove that the BoMb-OT is an approximation of a valid distance between probability measures and its entropic regularized version, eBoMb-OT, is the generalization of the conventional mini-batch optimal transport. More importantly, we have shown that the BoMb-OT and the eBoMb-OT can be implemented efficiently and they have more favorable performance than the m-OT in various applications of optimal transport. For future work, we could consider another extension of the BoMb-OT by changing the local OT to the Gromov-Wasseretein (GW) and unbalanced OT (UOT).
Reproducibility Statement: All datasets that we used in the paper are published and they are easy to find on the Internet. Source codes and instruction for our experiments are provided in the supplementary of the paper. The details of experimental settings, computational infrastructure, and other used public libraries are given in Appendix E.
Ethics Statement: The paper investigates a fundamental practical problem of optimal transport in machine learning and deep learning. Hence, we do not foresee any ethical issues of the paper that can be discussed.

heading: A PROOFS
text: In this appendix, we give detailed proofs of theorems that are stated in the main paper.

heading: A.1 PROOF OF THEOREM 1
text: We first prove that for any probability measures µ and ν ∈ Pp(Θ), there exists an optimal transportation plan γ∗ such that
Dmd (µ, ν) = E(Xm,Ym)∼γ∗ [d(PXm , PYm)]. (3)
From the definition of the BoMb-OT, there is a sequence of transportation plans γn ∈ Π( ⊗m µ , ⊗m ν ) such that
E(Xm,Ym)∼γn [d(PXm , PYm)]→ Dd(µ, ν)
as n→∞. Since Π(⊗mµ ,⊗mν ) is compact in the weak* topology (Villani, 2008), γn weakly converges to some γ∗ ∈ Π(⊗mµ ,⊗mν ). Since d(Pxm , Pym) is continuous in terms of xm, ym ∈ Xm, an application of Portmanteau theorem leads to
lim n→∞ E(Xm,Ym)∼γn [d(PXm , PYm)] ≥ E(Xm,Ym)∼γ∗ [d(PXm , PYm)].
Putting the above results together, we obtain
Dmd (µ, ν) = E(Xm,Ym)∼γ∗ [d(PXm , PYm)].
Therefore, there exists an optimal transportation plan γ∗ such that equation 3 holds.
We now proceed to prove that the BoMb-OT is a well-defined metric in the space of Borel probability measures. First, we demonstrate that Dd(µ, ν) = 0 if and only if µ = ν. In fact, from the definition of Dd(., .), if we have two probability measures µ and ν such that D(µ, ν) = 0, we find that
inf γ∈Π(⊗mµ ,⊗mν ) E(Xm,Ym)∼γ [d(PXm , PYm)] = 0.
Since d is a metric, it implies that there exists transportation plan γ∗ ∈ Π(⊗mµ ,⊗mν ) such that PXm = PYm γ
∗-almost everywhere. It demonstrates that for each x, there exists a permutation σx of {1, . . . ,m} such that xi = yσx(i) for all i ∈ [m]. Now, for any test function f : X → R, we have∫
Xm m∏ i=1 f(xi)d ⊗m µ (x) = ∫ Xm×Xm m∏ i=1 f(xi)dγ ∗(x, y) = ∫ Xm×Xm m∏ i=1 f(yσx(i))dγ ∗(x, y)
= ∫ Xm×Xm m∏ i=1 f(yi)dγ ∗(x, y)
= ∫ Xm m∏ i=1 f(yi)d ⊗m ν (y).
Since ∫ Xm ∏m i=1 f(xi)d ⊗m µ (x) = (∫ X f(x1)dµ(x1) )m and ∫ Xm ∏m i=1 f(yi)d ⊗m ν (y) =(∫
X f(y1)dν(y1) )m
, the above results show that∫ X f(x1)dµ(x1) = ∫ X f(y1)dν(y1)
for any test function f : X → R. It indicates that µ = ν. On the other hand, if µ = ν, we have ⊗m µ = ⊗m ν . We can construct the transportation plan γ̄(x, y) = δx(y) ⊗m µ (x) for all (x, y). Then, we find that
Dmd (µ, ν) ≤ E(X,Y )∼γ̄ [d(PXm , PYm)] = 0. Therefore, we obtain the conclusion that D(µ, ν) = 0 if and only if µ = ν. Now, we demonstrate the triangle inequality property of Dd(., .), namely, we will show that for any probability measures µ1, µ2, and µ3, we have
Dmd (µ1, µ2) +Dmd (µ2, µ3) ≥ Dmd (µ1, µ3). The proof of this inequality is a direct application of gluing lemma (Berkes & Philipp, 1977; De Acosta, 1982). In particular, for any transportation plans γ1 ∈ Π( ⊗m µ 1, ⊗m µ 2) and γ2 ∈ Π( ⊗m µ 2, ⊗m µ 3), we can construct a probability measure ξ on Xm×Xm×Xm such that ξ(., .,Xm) = γ1(., .) and ξ(Xm, ., .) = γ2(., .). Therefore, we find that E(X,Y )∼γ1 [d(PXm , PYm)] + E(Y,Z)∼γ2 [d(PYm , PZm)]
= ∫ (Xm)3 [ d(Pxm , Pym) + d(Pym , Pzm) ] dξ(x, y, z)
≥ ∫
(Xm)3 d(Pxm , Pzm)dξ(x, y, z) ≥ Dmd (µ1, µ3).
Taking the infimum on the LHS of the above inequality with respect to γ1 and γ2, we obtain the triangle inequality property of D(., .). As a consequence, we reach the conclusion of the theorem.

heading: A.2 PROOF OF THEOREM 2
text: To ease the presentation, for any product measures ⊗m µ and ⊗m ν , we denote the following loss between ⊗m µ and ⊗m ν as follows:
D̄md ( ⊗m µ , ⊗m ν ) := inf
γ∈Π(⊗mµ ,⊗mν ) E(Xm,Ym)∼γ [d(PXm , PYm)].
From the definitions of the BoMb-OT losses, we have Dmd (µ, ν) = D̄md ( ⊗m µ , ⊗m ν ) and D̂k,md (µn, νn) = D̄md ( ⊗m µk , ⊗m νk ). Using the similar proof argument as that of Theorem 1, we can check that D̄md satisfies the triangle inequality, namely, for any product measures ⊗m µ , ⊗m ν , and ⊗m η , we have D̄md ( ⊗m µ , ⊗m ν ) + D̄md ( ⊗m ν , ⊗m η ) ≥ D̄md ( ⊗m µ , ⊗m η ).
From the above definition and properties of D̄, we find that∣∣D̂k,md (µn, νn)−Dmd (µ, ν)∣∣ =∣∣D̄md (⊗mµk ,⊗mνk )− D̄md (⊗mµ ,⊗mν )∣∣ ≤ ∣∣D̄md (⊗mµk ,⊗mνk )− D̄md (⊗mµ ,⊗mνk )∣∣+∣∣D̄md (⊗mµ ,⊗mνk )− D̄md (⊗mµ ,⊗mν )∣∣ ≤ D̄md ( ⊗m µk , ⊗m µ ) + D̄md ( ⊗m νk , ⊗m ν ).
(i) Since d ≡Wp where p ≥ 1, we have
d(PXm , PYm) = Wp(PXm , PYm) = 1
m1/p ( inf σ m∑ i=1 ‖Xi − Yσ(i)‖p )1/p
≤ 1 m1/p m∑ i=1 ‖Xi − Yi‖,
where the infimum is taken over all possible permutations σ of {1, 2, . . . ,m}. The above inequality indicates that
D̄md ( ⊗m µk , ⊗m µ ) ≤ inf
γ∈Π(⊗mµk , ⊗m µ )
E(Xm,Ym)∼γ [ 1
m1/p m∑ i=1 ‖Xi − Yi‖ ] = m1−1/p ·W1(µn, µ).
Since X is a compact subset of RN , the results of (Dudley, 1969; Dobrić & Yukich, 1995; Fournier & Guillin, 2015) show that W1(µn, µ) = OP (n−1/N ). Collecting these results together, we have
D̄md ( ⊗m µk , ⊗m µ ) = OP
( m1−1/p
n1/N
) .
With similar argument, we also find that D̄md ( ⊗m νk , ⊗m ν ) = OP
( m1−1/p
n1/N
) . Putting all the above results
together, we reach the conclusion of part (i) of the theorem.
(ii) The proof of part (ii) follows directly from the argument of part (i) and the results that SWp(µn, µ) = OP (n−1/2) and SWp(νn, ν) = OP (n−1/2) when X is a compact subset of RN (Bobkov & Ledoux, 2019). Therefore, we obtain the conclusion of part (ii) of the theorem.

heading: B ADDITIONAL MATERIALS
text: Computational complexity: We now discuss the computational complexities of the BoMb-OT loss when d ≡ {W τp , SWp}. When d ≡W τp for some given regularized parameter τ > 0, we can use the Sinkhorn algorithm to compute d(PXmi , PYmj ) for any 1 ≤ i, j ≤ k. The computational complexity of the Sinkhorn algorithm for computing it is of the order O(m2). Therefore, the computational complexity of computing the cost matrix from the BoMb-OT is of the orderO(m2k2). Given the cost matrix, the BoMb-OT loss can be approximated with the complexity of the order O(k2/ε2) where ε stands for the desired accuracy. As a consequence, the total computational complexity of computing the BoMb-OT is O(m2k2 + k2/ε2). When d ≡ SWp, the computational complexity for computing d(PXmi , PYmj ) is of the order O(m logm) for any 1 ≤ i, j ≤ k. It shows that the complexity of computing the cost matrix from the BoMb-OT loss is of the order O(m(logm)k2). Given the cost matrix, the complexity of approximating the BoMb-OT is at the order of O(k2/ε2). Hence, the total complexity is at the order of O(m(logm)k2 + k2/ε2). Connection to hierarchical OT: At the first sight, the BoMb-OT may look similar to hierarchical optimal transport (HOT). However, we would like to specify the main difference between the BoMbOT and the HOT. In particular, on the one hand, the HOT comes from the hierarchical structure of data. For example, Yurochkin et al. (Yurochkin et al., 2019) consider optimal transport problems on both the document level and the word level for document representation. In the paper (Luo et al., 2020), HOT is proposed to handle multi-view data which is collected from different sources. Similarly, a hierarchical formulation of OT is proposed in (Lee et al., 2019b) to leverage cluster structure in data to improve alignment. On the other hand, the BoMb-OT makes no assumption about the hierarchical structure of data. We consider the optimal coupling of mini-batches as we want to improve the quality of mini-batch loss and its transportation plan.
Entropic regularized population BoMb-OT: We now consider an entropic regularization of the population BoMb-OT, which is particularly useful for reducing the computational complexity of the BoMb-OT. In particular, the entropic regularized population BoMb-OT (eBoMb-OT) between two probability measures µ and ν admits the following form:
EDmd (µ, ν) := min γ∈Π(⊗mµ ,⊗mν )
E(Xm,Ym)∼γ [d(PXm , PYm)] + λ · KL(γ| ⊗m µ ⊗ ⊗mν ), (4)
where λ > 0 stands for a chosen regularized parameter. From the above definition, the population eBoMb-OT is an interpolation between the population BoMb-OT and the population m-OT. On the one hand, when λ → 0, we have EDd(µ, ν) converges to Dd(µ, ν). On the other hand, when λ → ∞, the joint distribution γ approaches ⊗mµ ⊗ ⊗mν and EDd(µ, ν) converges to Ud(µ, ν). The
transportation plan of the eBoMb-OT can be derived by simply replacing the coupling between mini-batches in Definition 3 by the coupling found by the Sinkhorn algorithm.
Sparsity of transportation plan from BoMb-OT: We assume that we are dealing with two empirical measures of n supports. For d = Wp, the optimal transportation plan contains n entries that are not zero. In the m-OT case, with k mini-batches of size m, the m-OT’s transportation plan contains at most k2m non-zero entries. On the other hand, with k mini-batches of size m, the BoMb-OT’s transportation plan has at most km non-zero entries. For d = W τp (τ > 0), the optimal transportation plans contains at most n2 non-zero entries. In this case, the m-OT provides transportation plans that have at most k2m2 non-zero entries. the BoMb-OT’s transportation plans contain at most km2 non-zero entries. The sparsity is useful in the re-forward step of the BoMb-OT algorithm since we can skip pair of mini-batches that has zero mass to save computation.
Non-OT choice of d: We would like to recall that d in Definition 3 could be any discrepancy between empirical measures on mini-batches e.g. maximum mean discrepancy (MMD), etc. In this case, we can see the outer optimal transport between measures over mini-batches as an additional layer to incorporate the OT property into the final loss. However, it is not easy to define the notion of a transportation plan in these cases.

heading: C APPLICATIONS AND BOMB-OT ALGORITHMS
text: In this section, we collect the details of applications that mini-batch optimal transport is used in practice including deep generative models, deep domain adaptation, color transfer, and approximate Bayesian computation. Moreover, we present corresponding algorithms of these applications with our new mini-batch scheme BoMb-OT.

heading: C.1 DEEP GENERATIVE MODEL
text: Task description: We first consider the applications of the m-OT and the BoMb-OT into parametric generative modeling. The goal of parametric generative modeling is to estimate a parameter of interest, says θ, which belongs to a parameter space Θ. Each value of θ induces a model distribution µθ over the data space, and we want to find the optimal parameter θ∗ which has µθ∗ as the closest distribution to the empirical data distribution νn under a discrepancy (e.g. Wasserstein distance). In deep learning setting, θ is the weight of a deep neural network that maps from a low dimensional manifold Z to the data space X , and the model distribution µθ is a push-forward distribution Gθ]p(z) for p(z) is a white noise distribution on Z (e.g. N (0, I)). By using mini-batching schemes, we can estimate this parameter by minimizing the BoMb-OT loss in Definition 3:
θ∗ ← arg min θ∈Θ D̂k,md (µθ,n, νn), (5)
where µθ,n is the empirical distribution of µθ. The mentioned optimization is used directly in learning generative model on MNIST dataset with d ≡ (SW2,W τ2 ) in our experiments.
Algorithms: The algorithm for deep generative model with BoMb-OT is given in Algorithm 1. This algorithm is used to train directly the generative model on MNIST dataset.
Metric learning: Since L2 distance is not a natural distance on the space of images such as CelebA, metric learning was introduced by (Genevay et al., 2018; Deshpande et al., 2018) as a key step in the application of generative models. The general idea is to learn a parametric ground metric cost cφ:
cφ(x, y) = ‖fφ(x)− fφ(y)‖2, (6)
where fφ : X → Rh is a non-linear function that map from the data space to a feature space where L2 distance is meaningful.
The methodology to learn the function fφ depends on the choice of d. For example, when d = W τ2 , authors in (Genevay et al., 2018) seek for φ by solve following optimization:
max fφ
D̂k,md (fφ]µθ,n, fφ]νn), (7)
Algorithm 1 Mini-batch Deep Generative Model with BoMb-OT Input: k, m, data X , prior distribution p(z), chosen mini-batch loss LDGM ∈ {W2,W τ2 , SW2} Initialize Gθ on GPU; while θ does not converge do
—– On computer memory—– Sample indices uniformly IXm1 , . . . , IXmk on X Sample Zm1 , . . . , Z m k from p(z) Intialize C ∈ Rk×k for i = 1 to k do
for j = 1 to k do —– On GPU, autograd off —– Load Xmi to GPU from IXmi Load Zmj to GPU Compute Y mj ← Gθ(Zmj ) Compute Cij ← LDGM(PXmi , PYmj )
end for end for —– On computer memory—– Solve π ← OT(uk,uk, C) (entropic OT for eBoMb-OT) gradθ ← 0 for i = 1 to k do
for j = 1 to k do if πij 6= 0 then
—– On GPU, autograd on—– Load Xmi to GPU from IXmi Load Zmj to GPU Compute Y mj ← Gθ(Zmj ) Compute Cij ← LDGM(PXmi , PYmj ) gradθ ← gradθ + πij∇θCij
end if end for
end for —– On GPU—– θ ← Adam(θ, gradθ)
end while
In (Deshpande et al., 2018), another metric learning technique is used for d = SW2:
max fφ max gψ Ex∼µθ,n,y∼νn [log gψ(fφ(x)) + log(1− gψ(fφ(y))], (8)
where gψ : Rh → [0, 1]. Learning the metric or fφ is carried out simultaneously with learning the generative model in practice, namely, after one gradient step for the parameter θ, φ will be updated by one gradient step.

heading: C.2 DEEP DOMAIN ADAPTATION
text: We adapt the BoMb-OT into DeepJDOT (Bhushan Damodaran et al., 2018) which is a famous unsupervised domain adaptation method based on the m-OT. In particular, we aim to learn an embedding function Gθ : X → Z which maps data to the latent space; and a classifier Fφ : Z → Y which maps the latent space to the label space on the target domain. For a given number of the mini-batches k and the size of mini-batches m, the goal is to minimize the following objective
Algorithm 2 Mini-batch Deep Domain Adaptation with BoMb-OT Input: k, m, source domain (S, Y ), target domain T , chosen cost LDA in Equation 10. Initialize Gθ (parametrized by θ), Fφ (parametrized by φ) while (θ, φ) do not converge do
—– On computer memory—– Sample indices uniformly ISm1 ,Ym1 , . . . , ISmk ,Ymk on (S, Y ) Sample indices uniformly ITm1 , . . . , ITmk on T Initialize C ∈ Rk×k for i = 1 to k do
for j = 1 to k do —– On GPU, autograd off —– Load (Smi , Y m i ) to GPU from ISmi ,Ymi
Load Tmj to GPU from ITmj Compute mC ← LDA(Smi , Y mi , Tmj , Gθ, Fφ) (Equation 10) um ← ( 1 m , . . . , 1 m ) Compute Cij ← OT(um,um,mC)
end for end for —– On computer memory—– uk ← ( 1 k , . . . , 1 k
) Solve γ ← OT(uk,uk, C) (entropic OT for eBoMb-OT) gradθ ← 0 gradφ ← 0 for i = 1 to k do
—– On GPU, autograd on—– Load (Smi , Y m i ) to GPU from ISmi ,Ymi gradθ ← gradθ + 1k 1 m∇θLs(Y m i , Fφ(Gθ(S m i )) (Equation 9) gradφ ← gradφ + 1k 1 m∇φLs(Y m i , Fφ(Gθ(S m i )) (Equation 9) for j = 1 to k do if γij 6= 0 then
—– On GPU, autograd on—– Load Tmj to GPU from ITmj Compute mC ← LDA(Smi , Y mi , Tmj , Gθ, Fφ) (Equation 10) Compute Cij ← OT(um,um,mC) gradθ ← gradθ + γij∇θCij gradφ ← gradφ + γij∇φCij
end if end for
end for —– On GPU—– θ ← Adam(θ, gradθ) φ← Adam(θ, gradφ)
end while
function:
min Gθ,Fφ LDA = [1 k 1 m k∑ i=1 m∑ j=1 Ls(yij , Fφ(Gθ(sij)))
+ min γ∈Π(uk,uk) k∑ i=1 k∑ j=1 γij × min π∈Π(um,um) 〈CGθ,FφSmi ,Ymi ,Tmj , π〉 ] , (9)
where Ls is the source loss function (e.g. classification loss), Sm1 , . . . , S m k are source mini-batches which are sampled with or without replacement from the source domain Sm ∈ Xm, Y m1 , . . . , Y mk are corresponding labels of Sm1 , . . . , S m k , with S m i = {si1, . . . , sim} and Y m1 := {yi1, . . . , yim}. Similarly, Tm1 , . . . , T m k (T m 1 := {ti1, . . . , tim}) are target mini-batches which are sampled with or
without replacement from the target domain T m ∈ Xm. The cost matrix CG,FSmi ,Ymi ,Tmi denoted mC is defined as follows:
mC1≤z1,z2≤m = α||G(siz1)−G(tjz2)||2 + λtLt(yiz1 , F (G(tjz2))), (10)
where Lt is the target loss function, α and λt are hyper-parameters that control two terms.

heading: C.3 COLOR TRANSFER
text: In this appendix, we carry out more experiments on color transfer with various settings. In detail, we compare the m-OT and the BoMb-OT in different domains of images, and we demonstrate visually the interpolation property of the eBoMb-OT by varying the regularization parameter.
Methodology: In our experiment, we first compress both the source image and the target image using K-means clustering with 3000 components. After that, we look for the transportation plans between cluster centers of two images from different approaches, the m-OT, the BoMb-OT, and the full-OT. Next, we change the values of all pixels in a cluster of the source image based on the value of the corresponding cluster centers of the target image with the previously found plans. We present the algorithm that we use in color transfer with the BoMb-OT in Algorithm 3. This algorithm is adapted from the algorithm that is used for the m-OT in (Fatras et al., 2020).
Algorithm 3 Color Transfer with BoMb-OT Input: k,m, T source image Xs ∈ Rn×3, target image Xt ∈ Rn×3 Initialize Ys ∈ Rn×3 for t = 1 to T do
Initialize C ∈ Rk×k Sample indices IXm1 , . . . , IXmk from Xs Sample indices IYm1 , . . . , IYmk from Xt for i = 1 to k do
for j = 1 to k do Load Xmi from IXmi Load Y mi from IYmi Compute cost matrix M between Xmi and Y m i
π ← arg minπ∈Π(um,um)〈M,π〉 Cij ← 〈M,π〉
end for end for γ ← arg minγ∈Π(uk,uk)〈C, γ〉 for i = 1 to k do
for j = 1 to k do if γij 6= 0 then
Load Xmi from IXmi Load Y mi from IYmi Compute cost matrix M between Xmi and Y m i π ← arg minπ∈Π(um,um)〈M,π〉 Ys|IXm
i ← mγij π ·Xt|ImYj end if
end for end for
end for Output: Ys

heading: C.4 APPROXIMATE BAYESIAN COMPUTATION (ABC)
text: In this appendix, details of Approximate Bayesian Computation (ABC) and the usage of mini-batches with BoMb-OT are discussed.
Algorithm 4 Approximate Bayesian Computation Input: Generative model p(x, θ), observation {xi}ni=1, number of iterations T , discrepancy measure D, tolerance threshold , number of particles m, and summary statistics s (optional). for t = 1 to T do
repeat Sample θ ∼ p(θ) Sample {yi}mi=1 ∼ p(x|θ) until D(s({yi}mi=1), s({xi}ni=1)) ≤ θt = θ
end for Output: {θt}Tt=1
Review on ABC: The central task of Bayesian inference is to estimate a distribution of parameter θ ∈ Θ given n data points {xi}ni=1 and a generative model p(x, θ). Using Bayes rule, the posterior can be written as p(θ|x1:n) := p(x1:n|θ)p(θ)p(x1:n) . Generally, the posterior is intractable since we cannot evaluate the normalizing constant p(x1:n) (or the evidence). It leads to the usage of approximate Bayesian Inference, e.g., Markov Chain Monte Carlo and Variational Inference. However, in some settings, the likelihood function p(x1:n|θ) cannot be evaluated such as implicit generative models. In these cases, Approximate Bayesian Computation (or likelihood-free inference) is a good framework to infer the posterior distribution since it only requires the samples from the likelihood function.
We present the algorithm of ABC in Algorithm 4 that is used to obtain posterior samples of θ. The thing that sets ABC apart is that it can be implemented in distributed ways, and its posterior can be shown to have the desirable theoretical property of converging to the true posterior when → 0. However, the performance of ABC depends on the choice of the summary statistics s (e.g., empirical mean and empirical variance) and the discrepancy D. In practice, constructing sufficient statistics is not easy. Thus, a discrepancy between empirical distributions is used to avoid this non-trivial task. Currently, Wasserstein distances have drawn a lot of attention from ABC’s community because they provide a meaningful comparison between non-overlap probability distributions (Bernton et al., 2019a; Nadjahi et al., 2020). When the number of particles m in Algorithm 4 is the largest OT problem that can be solved by the current computational resources, the mini-batch approach can be utilized to obtain more information about two measures, namely, more supports in sample space can be used. In particular, we utilize the mini-batch OT losses as the discrepancy D in Algorithm 4. The detail of the application of BoMb-OT in ABC is given in Algorithm 5.

heading: D ADDITIONAL EXPERIMENTS
text: In this appendix, we provide additional experimental results that are not shown in the main paper. In Appendix D.1, we present random generated images on MNIST dataset, CIFAR10 dataset, CelebA dataset, and training time comparison between the m-OT and the BoMb-OT. Furthermore, we give detailed result on deep domain adaptation in Appendix D.2 and we also provide the computational time. In Appendix D.3 we carry out more experiments on color transformation with various settings on different images and their corresponding color palettes. Moreover, details of setting of the ABC experiment and additional comparative experiments are presented in Appendix D.4. After that, we compare the m-OT and the BoMb-OTs on gradient flow application in Appendix D.5. Finally, we investigate the behavior of the m-OT and the (e)BoMb-OT in estimating the optimal transportation plan in Appendix D.6.

heading: D.1 DEEP GENERATIVE MODEL
text: Generated images: We show the generated images on MNIST dataset in Figure 4, the generated images on CelebA in Figure 5, the generated images on CIFAR10 in Figure 6. For both choices of d (SW2, W τ2 ), we can see that the images from (e)BoMb-OT is more realistic than the images from m-OT. This qualitative result supports the quantitative in Table 1 in the main text.
Algorithm 5 Approximate Bayesian Computation with BoMb-OT Input: Generative model p(x, θ), observation {xi}ni=1, number of iterations T , optimal transport discrepancy measure DOT, tolerance threshold , number of particles m, number of mini-batches k, and summary statistics s (optional). for t = 1 to T do
repeat Sample θ ∼ p(θ) Sample indices IXm1 , . . . , IXmk from {xi} n i=1
Sample {yi}ni=1 ∼ p(x|θ) Sample indices IYm1 , . . . , IYmk from {yi} n i=1 Initialize C ∈ Rk×k for i = 1 to k do
Computational speed: Table 3 details the computational speed of deep generative model when using mini-batch sizem = 100. In general, the more the number of mini-batches is, the more complex the problem is. Increasing the number of mini-batches k does decrease the number of iterations per second in all experiments. For both datasets, using the sliced Wasserstein distance (d = SW2) leads to higher iterations per second compared to the entropic Wasserstein distance (d = W τ2 ) on the high dimensional space. This result is expected since the slicing approach aims to reduce the computational complexity. On the MNIST dataset, we observe that the entropic regularized version of the BoMb-OT (eBoMb-OT) is the slowest for both choices of d. The m-OT is the fastest approach with the number of iterations per second of 7.89, 2.54, and 0.85 for k = 2, 4, and 8 respectively. In contrast, the BoMb-OT is faster than the m-OT for CIFAR10 and CelebA dataset because we set a high value for the regularized parameter (e.g. τ = 50, 40). For the CelebA dataset, although the eBoMb-OT is the slowest method when d = SW2, it becomes the fastest approach if we set d to W τ2 . The reason for this phenomenon is because of the sparsity of the (e)BoMb-OT’s transportation plan between measures over mini-batches. In particular, pairs of mini-batches that are zero can be skipped
(see in Algorithm 1. So, the (e)BoMb-OT can avoid estimating gradient of the neural networks of a bad pair of mini-batches.

heading: D.2 DEEP DOMAIN ADAPTATION
text: In this section, we compare the performance of two mini-batch schemes on digits and Office-Home datasets.
Comparison between the m-OT and BoMb-OT on digits datasets: As seen in Table 4, the BoMb-OT produces better classification accuracy than the m-OT in all experiments. In addition, it leads to a huge performance improvement (over 19%) in comparison with the m-OT when the mini-batch size is small (k = 10 for SVHN to MNIST and k = 5 for USPS to MNIST). When the mini-batch size becomes larger, the performance of both methods also increases. An explanation for such outcome is that a large mini-batch size makes mini-batch loss approach to its population version.
Comparison between the m-UOT and eBoMb-UOT on Office-Home dataset: Table 5 illustrates the performance of m-UOT and eBoMb-UOT on Office-Home dataset when changing the number of mini-batches k from 2 to 8. When k = 4, eBoMb-UOT achieves the classification accuracy higher than m-UOT on 11 out of 12 scenarios, resulting in an improvement of 0.24 on average.
Computational speed: The number of iterations per second of deep DA can be found in Table 6-7. Similar to deep generative model, we observe a phenomenon that the speed of both the m-OT and BoMb-OT decreases as k increases. Interestingly, increasing the mini-batch size m when adapting from USPS to MNIST barely affects the running speed of both methods. The run time of the m-OT nearly doubles that of the BoMb-OT. Specifically, the BoMb-OT averagely runs 1 iteration in a second while an iteration of the m-OT consumes roughly 2 seconds. Although having comparable time complexity, the BoMb-OT in practice runs faster than the m-OT in all deep DA experiments. This is because of the sparsity of the BoMb-OT’s transportation plan between measures over mini-batches. In particular, pairs of mini-batches that are zero can be skipped (see in Algorithm 2). We would like to recall that, gradient estimation is the most time-consuming task in deep learning applications.

heading: D.3 COLOR TRANSFER
text: In this appendix, we compare the m-OT and the BoMb-OT in different domains of images including natural images and arts.
Comparison between the m-OT and the BoMb-OT: Next, we illustrate the color-transferred images using the m-OT and the BoMb-OT with two values of the number of mini-batches and the size of mini-batches (k,m) = (10, 10) and (k,m) = (20, 20). The number of incremental step T (see in Algorithm 3) is set to 5000. We show the source images, target images, and corresponding transferred images in Figure 7-Figure 10. It is easy to see that the transferred images from the BoMb-OT look more realistic than the m-OT and the color is more similar to the target images. The color palette of transferred images also reinforces the above claim when it is closer to the color palette of the target images. According to those figures, increasing the number of mini-batches and the mini-batch size improves the results considerably.
Computational speed: For k = 10,m = 10, the m-OT has the speed of 103 iterations per second while the BoMb-OT has the speed of 100 iterations per second. For k = 20,m = 20, the speed of m-OT is about 20 iterations per second and the speed of the BoMb-OT is also about 20 iterations per second. It means that the additional k × k OT is not too expensive while it can improve the color palette considerably.

heading: D.4 APPROXIMATE BAYESIAN COMPUTATION (ABC)
text: Settings: We use the same setup as in (Nadjahi et al., 2020): there are n = 100 observations {xi}ni=1 i.i.d from multivariate Gaussian N (µ∗, σ2∗IN ) where N is the dimension, µ∗ ∼ N (0, IN ) and σ2∗ = 4. The task is to estimate the posterior of σ
2 under the imaginary assumption that σ2 follows inverse gamma distribution IG(1, 1). Under these assumptions, the posterior of σ2 has the form IG(1 + nd2 , 1 + 1 2 ∑n i=1 ||xi − ν∗||2). For ABC, we use the m-OT and the BoMb-OT with the Wasserstein-2 ground metric for the acceptance-rejection sampling’s criteria in sequential Monte Carlo ABC (Toni et al., 2009) (using implementation from pyABC (Klinger et al., 2018) with 100 particles and 10 iterations). The mini-batches’ size is set in {8, 16, 32} and the number of mini-batches is in {2, 4, 6, 8}.
Results on ABC with mini-batches: After obtaining all samples, we estimate their densities by utilizing Gaussian kernel density estimation and then plot the approximated posteriors and the true posterior in Figure 11. Here we run the algorithm with several values of (k,m), namely, m ∈ {8, 16, 32} and k ∈ {2, 4, 6, 8}. From these results, we see that a bigger m usually returns a better posterior in both the m-OT and the BoMb-OT cases. Similarly, increasing k also improves the performance of the two methods. Moreover, these graphs strengthen the claim that the BoMb-OT outperforms the m-OT in every setting of (k,m) since its posteriors are always closer to the true posterior than those of the m-OT in both visual result and Wasserstein distance.

heading: D.5 NON-PARAMETRIC GENERATIVE MODEL VIA GRADIENT FLOW
text: In this appendix, we show the experiment that uses the m-OT and the (e)BoMb-OT in the gradient flow application.
Task description: Gradient flow is a non-parametric method to learn a generative model. Like every generative model, the goal of gradient flow is to mimic the data distribution ν by a distribution µ. It leads to the functional optimization problem:
min µ D(µ, ν), (11)
where D is a predefined discrepancy between two probability measures. So, a gradient flows can be constructed:
∂tµt = −∇µtD(µt, ν) (12) We follow the Euler scheme to solve this equation as in (Feydy et al., 2019), starting from an initial distribution at time t = 0. In this paper, we choose D be (e)BoMb-W2 and m-W2 for the sake of comparison between them.
We first consider the toy example as in (Feydy et al., 2019) and present our results in Figure 12. The task is to move the colorful empirical measure to the "S-shape" measure. Each measure has 1000 support points. Here, we choose (k,m) = (4, 16), the OT loss is Wasserstein-2, and we use the Wasserstein-2 score to evaluate the performance of the mini-batching scheme. From Figure 12, BoMb-OT and eBoMb-OT provide better flows than m-OT, namely, Wasserstein-2 scores of BoMB-OT and eBoMb-OT are always lower than those of m-OT in every step. In addition, we do an extra setup with a higher of mini-batches, (k,m) = (16, 16) to show the increasing the number of mini-batches improve the performance of both m-OT and (e)BoMb-OT. The result is shown in Figure 13. In this setting, the BoMb-OT still shows its favorable performance compared to the m-OT, namely, its Wasserstein-2 scores are still lower than the m-OT in every step.
CelebA: Let µ and ν denote the empirical measures defined over 5000 female images and 5000 male images in the CelebA dataset. Thus, we can present the transformation from a male face to be a female one by creating a flow from µ to ν. Our setting experiment is the same as (Fatras et al., 2020). We first train an autoencoder on CelebA, then we compress two original measures to the low-dimensional measures on the latent space of the autoencoder. After having the latent measures, we run the Euler scheme to get the transformed measure then we decode it back to the data space by the autoencoder. The result is shown in Figure 14, we also show the closest female image (in sense of
L2 distance) to the final found image in each method and the corresponding L2 distance between the middle step images and the nearest image. As shown, (e)BoMb-OT provides a better flow than m-OT does.

heading: D.6 TRANSPORTATION PLANS
text: In this appendix, we investigate deeper the behavior of the m-OT and the (e)BoMb-OT in estimating the transportation plan. We present a toy example to illustrate the transportation plans of the m-OT, the BoMb-OT, and the original OT (full-OT). In particular, we sample two empirical distributions µn and νn, where n = 10, from N ([ 0 0 ] , [ 1 0 0 1 ] ) and N ( [ 4 4 ] , [ 1 −0.8 −0.8 1 ] ) respectively.
Transportation plans: We plot the graph of samples matching and transportation matrices in Figure 15. When the size of mini-batches equals 2 and the number of mini-batches is set to be 20, the
m-OT approach produces messy OT matrices and disordered matching, meanwhile, the BoMb-OT can still concentrate the mass to meaningful entries of the transportation matrix, thus, its matching is acceptable. Increasing the size of mini-batches to 8, m-OT’s performance improves significantly however its matching and its matrices are visibly still not good solutions. In contrast, the BoMb-OT is able to generate nearly optimal matchings and transportation matrices. Next, we test the m-OT and the BoMb-OT in real applications in the case of an extremely small number of mini-batches. When there are 8 mini-batches of size 2, the m-OT still performs poorly while the BoMb-OT creates a sparser transportation matrix that is closer to the optimal solution obtained by the full-OT. Similarly, with 2 mini-batches of size 8, the BoMb-OT has only 4 wrong matchings while that number of the m-OT is 10. In conclusion, the BoMb-OT is the better version for mini-batching with any value of the number of mini-batches and mini-batches size.
Entropic transportation plan of the BoMb-OT: We empirically show that the transportation plan of the eBoMb-OT, when the entropic regularization is sufficiently large, reverts into the m-OT’s plan. The result is given in Figure 16. When λ = 103, the transportation plans of the eBoMb-OT
are identical to plans of the m-OT with every choice of (k,m). However, when λ = 0.1, despite not being identical to full-OT, the plans produced by the eBoMb-OT are still close to the true optimal plan.
Comparison with stochastic averaged gradient (SAG): We compare the BoMb-OT with the stochastic averaged gradient (SAG) (Aude et al., 2016) for computing OT. The transportation matrices are given in Figure 17. We would like to recall that SAG still need to store full cost matrix (n×) while the BoMb-OT only need to store smaller cost matrices (m×m). We can see that the BoMb-OT gives better than transportation plan than SAG when m = 2. When m = 8, the BoMb-OT still seems to be better.

heading: E EXPERIMENT SETTINGS
text: In this appendix, we collect some necessary experimental setups in the paper including generative model, gradient flow. We use POT (Flamary et al., 2021) for OT solvers, and the pyABC (Klinger et al., 2018) for the ABC experiments.

heading: E.1 DEEP GENERATIVE MODEL
text: Wasserstein-2 scores: We use empirical distribution with 11000 samples that are obtained by sampling from the generative model and the empirical test set distribution respectively, then we compute discrete Wasserstein-2 distance via linear programming.
FID scores: We use 10000 samples from the generative model and all test set images to compute FID score (Heusel et al., 2017).
Parameter settings: We chose a learning rate equal to 0.0005, batch size equal to 100, number of epochs in MNIST equal to 100, number of epochs in CelebA equal to 25, number of epochs in CIFAR10 equal to 50. For d = SW2 we use the number of projections L = 1000 on MNIST and L = 100 on CIFAR10 and CelebA. For d = W τ2 , we use τ = 1 for MNIST,
τ = 50 for CIFAR10, τ = 40 for CelebA. For the eBoMb-OT, we choose the best setting for λ ∈ {1, 2, 3, 4, 5, 10, 20, 30, 40, 50, 60, 70, 80}
Neural network architectures: We use the MLP for the generative model on the MNIST dataset, while CNNs are used on the CelebA dataset.
Generator architecture was used for MNIST dataset: z ∈ R32 → FC100 → ReLU → FC200 → ReLU → FC400 → ReLU → FC784 → ReLU
Generator architecture was used for CelebA: z ∈ R128 → TransposeConv512 → BatchNorm → ReLU → TransposeConv256 → BatchNorm → ReLU → TransposeConv128 → BatchNorm → ReLU → TransposeConv64 → BatchNorm →
ReLU → TransposeConv1 → Tanh Metric learning neural network’s architecture was used for CelebA:
First part: x ∈ R64×64×3 → Conv64 → LeakyReLU0.2 → Conv128 → BatchNorm → LearkyReLU0.2 → Conv256 → BatchNorm → LearkyReLU0.2 → Conv512 → BatchNorm→ Tanh Second part: Conv1 → Sigmoid Generator architecture was used for CIFAR10: z ∈ R128 → TransposeConv256 → BatchNorm → ReLU → TransposeConv128 → BatchNorm → ReLU → TransposeConv64 → BatchNorm→ ReLU → TransposeConv1 → Tanh Metric learning neural network’s architecture was used for CIFAR:
First part: x ∈ R32×32×3 → Conv64 → LeakyReLU0.2 → Conv128 → BatchNorm → LearkyReLU0.2 → Conv256 → BatchNorm→ Tanh Second part: Conv1 → Sigmoid

heading: E.2 DEEP DOMAIN ADAPTATION
text: In this section we state the neural network architectures and hyper-parameters for deep domain adaptation.
Evaluation metric: The classification accuracy is utilized to evaluate the mini-batch methods.
Parameter settings for digits datasets: The number of mini-batches k varies in {8, 16, 32}. For the SVHN dataset, the mini-batch size m is set to 50. Because the USPS dataset has fewer samples than the SVHN dataset, m is set to 25. We train the network using Adam optimizer with an initial learning rate of 0.0002. The number of epochs is 80 for k = 8. As the number of minibatches doubles, we double the number of epochs so that the number of iterations does not change. The hyperparameters of m-OT and BoMb-OT in Equation 10 follow the settings in DeepJDOT: α = 0.001, λt = 0.0001. The hyperparameters of m-UOT and eBoMb-UOT are the same as JUMBOT: α = 0.1, λt = 0.1, = 0.1, τ = 1. For computing eBoMb-UOT, we choose the best value of λ ∈ {0.01, 0.1, 0.2, 1, 10} (entropic regularizer coefficient). Parameter settings for Office-Home dataset: The number of mini-batches k varies in {2, 4, 8}. Following the settings in (Fatras et al., 2021a), we train models with a mini-batch size m = 65 during 10000 iterations. The hyperparameters for computing the cost matrix follow the settings in JUMBOT: α = 0.01, λt = 0.5, τ = 0.5, = 0.01. For computing eBoMb-UOT, we choose the best value of λ ∈ {0, 0.01, 1, 100} (entropic regularizer coefficient). Training details: Similar to both DeepJDOT and JUMBOT, we stratify the data loaders so that each class has the same number of samples in the mini-batches. For digits datasets, we also train our neural network on the source domain during 10 epochs before applying our method. For Office-home dataset, because the classifiers are trained from scratch, their learning rates are set to be 10 times that of the generator. We optimize the models using an SGD optimizer with momentum = 0.9 and weight decay = 0.0005. We schedule the learning rate with the same strategy used in JUMBOT. The learning
rate at iteration p is ηp = η0(1+µq)ν , where q is the training progress linearly changing from 0 to 1, η0 = 0.01, µ = 10, ν = 0.75.
Neural network architectures: On digits datasets, we use CNN for our generator and 1 FC layer for our classifier in both adaptation scenarios. For Office-Home dataset, our generator is a ResNet50 pre-trained on ImageNet except for the last FC layer, which is our classifier.
Generator architecture was used for SVHN dataset: z ∈ R32×32×3 → Conv32 → BatchNorm → ReLU → Conv32 → BatchNorm → ReLU → MaxPool2D → Conv64 → BatchNorm → ReLU → Conv64 → BatchNorm → ReLU → MaxPool2D → Conv128 → BatchNorm → ReLU → Conv128 → BatchNorm → ReLU →MaxPool2D → Sigmoid→ FC128
Generator architecture was used for USPS dataset: z ∈ R28×28×3 → Conv32 → BatchNorm → ReLU → MaxPool2D → Conv64 → BatchNorm → ReLU → Conv128 → BatchNorm → ReLU → MaxPool2D → Sigmoid→ FC128
Classifier archiecture was used for both SVHN and USPS datasets: z ∈ R128 → FC10

heading: E.3 GRADIENT FLOW
text: For the implementation of the gradient flow, we use the geomloss library (Feydy et al., 2019). The learning rate is set to 0.001. For the autoencoder in CelebA experiments, we use the repo in https://github.com/rasbt/deeplearning-models/blob/master/ pytorch_ipynb/autoencoder/ae-conv-nneighbor-celeba.ipynb for the pretrained autoencoder.

heading: E.4 COMPUTATIONAL INFRASTRUCTURE
text: All deep learning experiments are done on a RTX 2080 Ti GPU and a GTX 1080 Ti. Other experiments are done on a MacBook Pro 11inc M1.

Review summaries:
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.

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.