On Transportation of Mini-batches: A Hierarchical Approach
Authors: Khai Nguyen
Dang Nguyen
Nguyen Dinh Quoc
Tung Pham
Hung Bui
Dinh Phung
Trung Le
Nhat Ho
Venue: ICLR 2022 Submitted (2022)
Keywords: Deep Generative Models
Deep Domain Adaptation
Color Transfer
Approximate Bayesian Computation
Gradient Flow
Optimal Transport

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.

Under review as a conference paper at ICLR 2022
ON TRANSPORTATION OF MINI-BATCHES: A HIERARCHICAL APPROACH
Anonymous authors Paper under double-blind review
ABSTRACT
Mini-batch optimal transport (m-OT) has been successfully used in practical ap- plications 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 bini-batches which leads undesirable estimation. Moreover, the m-OT does not approximate 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 opti- mal coupling between nini-batches and it can be well-defined distance on the space of probability measures. Furthermore, we show that the m-OT is 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 deep generative models and deep domain adaptation. From extensive experiments, we observe that the BoMb-O1 achieves 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 better qualitative result than the m-OT.
approximation to
INTRODUCTION
Optimal transport (OT) (Villani 2021 2008 Peyré 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 al., [017]]
Tolstikhin et al. 22188) Salimans et al. 20188 Genevay et al. 2018 Liutkus et al. (019)), barycenter problem (Ho et al.J [201 Li et al.) 2020), and approximate Bayesian computation (Bernton et al. 2019a Nadjahi al., (020). 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 al. 2014 Perrot et al., (2066), 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 al. 2019 Xu al. 2019 Chen et al.) 2020).
Although OT has attracted growing attention in recent years, major barrier that prevents OT from being ubiquitous is its heavy computational cost. When the two probability measures discrete with n supports, solving the Wasserstein distance via the interior point methods has the complexity of Onn log n) (Pele Werman, 2009). which is extremely expensive when 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 e-aappoximmtedd solution at the same time reduces the computational complexity to (nn//²²) 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 On l)).). There various variants in this directions; i.c., (Bonneel et al. 2015) Deshpande et al., 2019 Kolouri et al, 2019; Nguyen et al.) [021a[b), these all belong to the family of sliced Wasserstein distances. Recently, some methods are proposed to combine the

---

Under review as a conference paper at ICLR 2022
dimensional reduction approach with the entropic solver of Wasserstein distance 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. x 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 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-batchis method for OT (Genevay et al. 2018); Sommerteld et al. (019), which we refer to as mini- bani-bbtch 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 al.) 2018; Bhushan Damodaran et al., 2018 Kolouri et al. 2018 2 (allimans et al., (2018). and was theoretically studied in the works of (Bellemare et al. (21 [ernton et al. [019b) Nadjahi
ee al. 2019). Recently, Fatras et al. (Fatras et al.] 2020 2021b) formulated this approach by giving 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 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.
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 set of empirical measure. We now could employ the Kantorovich formulation between these two empirical measures in the product space main contributions are two-fold:
Contribution:
could be considered as an
discrepancy between two sets of mini-batches. In summary, our
1. First, the BoMb-OT could provide more similar transportation plan to the original OT than the m-OT, which leads to 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 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 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, 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
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.
2

---

Under review as a conference paper at ICLR 2022
Notation: For any probability measure j the Polish measurable space (X, denote ((mm
8m (m 2 2) as the product measure on the product measurable space (2" -.)) For any p 2 1, we define Py (IB= as the set of Borel probability measures with finite y-th moment defined on a given metric space (R-N II-II)- To simplify the presentation, we abuse the notation by using both the notation X for both the random vector ( Imm) E A and the set (I1,..... ,.m ), and we define by Pxm= := 1 -2 2i=1 82, the empirical measure (the mini-batch measure) associated with X". For any set X" := (II,...,Zn) Jn) and m 1, we denote by [X*]]* the product set of X taken m times and (X") the set of all m-element subsets of X"
m
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 ju and v be two probability measures on Pp (I.N). The Wasserstein p-distance between and v is defined as follows:
Wp(n, v) == milprel(p.v)) [E<(,,g) || y||²| } where 1I(,, v) == { Ind ʃ xdy 11} is the set of transportation plans between j and v.
The entropic regularized Wasserstein to approximate the OT solution (Altschuler et al. 2017
Lin et al. 2019) between j and v is defined as follows (Cuturi 2013):
W (1,
y/p] & TKL(T| @ v)}. where T 0 is a chosen regularized pa- min (a,))))
rameter and KL denotes the Kullback-Leible 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:
SW, (W,((,,)) == [Eg- (Oty,, 0tv)] , where (1()~~1) denotes the uniform measure over the (N 1) dimeensoonll unit hypersphere and 0c is the orthogonal projection operator on direction A.
2.2 MINC-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 mod- els. domain adaptations, color transfer, and approximate Bayesian computation (Bhushan Damodaran al. 2018: Genevay et al., 2018 Tolstikhin et al. 2018 Fatras et al. 2020 Bernton et al. 201 9a). To ease the presentation, we are given X" = 11....... Fn 4 - Un i.i.d. samples Σn 82, and v, 1 Σ 8,. be two corresponding empirical
from μ and in turn. Let (4, measures from the whole data set. Here, n is usually large (e.g., millions) and each support in X", y" can be high dimensional data point (e.g. a high resolution image, video, etc).
Memory issue of optimal transport: Using an OT loss between j4, and Un needs to compute and store 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., 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 solve the Kantorovich dual form is a possible approach to deal with large-scale OT, i.e. Wasserstein GAN (Arjovsky et al. 2017 Leygonic 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 restricted in the choice of ground metric. In particular, Lı norm is used in
3

---

Under review as a conference paper at ICLR 2022
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 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 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
We now state an adapted definition of mini-batch OT (m-OT) scheme that was formulated in (Fatras ct al. 2020), including its two key parts: its transportation cost and its transportation plan.
Definition 1 (Empirical m-OT). For p 2 1 and integers m 2 1, and k 2 1, let d : Pp(X) x Pp((X) -> [0, oo) be a function, i.e., (Wp, WT, SW,). Then, the mini-batch OT (m-OT) loss and the transportation plan, a
matrix, between Un, and vn are defined as follows:
k k d(pxp Py!"): (M,, Un) = R
k
k
Th:"((n: Un) =
1
ΣΣ j=1
1
ΣΣ i=1j=1
1.2
where X is sampled i..d from (x)) Yimh is sampled i..d from (²)). and the transport plan """,,
where its entries equal zero except those indexed of samples X x YJ", is the transportation plan when di (Xpx" Pxz") 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 (x(X) for given positive integers p 2 1, m 2 1. and d : Pp(X) x Pp(X) + [0, oo) be a given function. Then, the population mini-batch OT (m-OT) discrepancy between μ and is defined as follows:
U (n, v) E
am x dl Px Pym). (X),
(2)
Issues of the m-OT: From Definition I the m-OT treats every pair of mini-batche 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 X" would prefer to match to certain mini-batches of Y rather than treating every mini-batch of ym 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 Eld(XX", y"]] for independent X and Y 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 symmetric and positive (Fatras et al. 2020), for the same reason it does not vanish when two measures are identical.
3 BATCH OF MINI-BATCHES OPTIMAL TRANSPORT
To address the issues of m-OT, in this section, we propose a novel mini-batch scheme, named batch of mini-batches OT (T)) 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 well-defined metric 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 J and v are continuous probability measures in (X) for some given p 2 1 and X' C R. where N 2 1. Furthermore, d' (X) x Pp(X') [0, oo) is given divergence between probability measures in Pp())
3.1 DEFINITION OF BOMB-OT AND ITS PROPERTIES
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
4

---

Under review as a conference paper at ICLR 2022
m-OT
per deve
BoMb-OT
Figure 1: Visualization of the m-OT and the BoMb-OT in providing mapping between samples.
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 X" would prefer to match some certain mini-batches from Y". 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 good global mapping between samples that also leads a meaningful objective loss.
...... X be mini-batches that are sampled with replacement from [X"] and Y Y ..... Y be mini-batches that are sampled with replacement from [Y"]". The batch of mini-batches optimal transport is defined as follows:
Let
Definition (Batch of mini-batches). Assume that 2 m 2 k 1 are positive integers and let d: (x( (2)) + [0,oo] oo) be given function (e.g., (Wp, WT, SW, The BoMb-OT loss and the BoMb-OT transportation plan between probability measures and 1n are given by:
h ΣΣ Dí (pny Un)= = min
Ryd Pym));
WE (pms Un) == XXp
7E11 PA 2%
= E 6 are two empirical measures defined on the product space via batches (measures over mini-batches), 5 is a k x k optimal transport plan 8m 8m
60m where Mk
60en. : Et 0x.- and Vk-
between Pk and Vk, and i y." is defined as in Definition!
Compared to the m-OT, the BoMb-OT considers an additional OT between two measures mini- batches 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 x 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 B and B2 from y. The weight for the m-OT loss between Al and B2 equals 1, meanwhile, the map from Al to B2 is unuseful. In this toy example, Al should be mapped to B1 and A2 is for B2. It also the solution that the BoMb-OT tries to obtain by re- weightng pairs of mini-batchin through a transportation plan between 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 is gradient-baser applications (deep learning) such deep generative models, deep domain adaption, and gradient flow. In these applications, the goal is to estimate the gradient of parameter of interest neural network) with an OT objective function. The second one is mapping-basee 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-bate bbased 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 better presentation, we assume that a system with 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 x m supports of two mini-batch measures, and the X cost matrix. The algorithm of the BoMb-OT consists of three steps: Forwarding, Solving x k OT, and Re-forwarding and Backwarding. In the first step, k2 OT problems of size m x m computed in turn GPU to obtain the & 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 OT GPU at time. A After that, OT problem with the found k k cost matrix is solved
5

---

Under review as a conference paper at ICLR 2022
to find the coupling between mini- mini-batches. Here, & can be much larger than since the k x 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 rorworng and do backwarding prock-propagation) to obtain the gradients of each pair of nini-batches, then using the coupling in the second step aggregate them into the final gradient signal of the parameter of interest. The detailed algorithmss of the BoMb-OT for deep generative model and deep domain adaption are respectively described in Algorithm [ in Appendix C.I and Algorithm in Appendix C.2 Compared to the m- m-OT, the BoMb-OT only costs an additional forwarding and an additional l k OT problem that not expensive since they do not require to create and store any computational graphs.
Mapping-hare upplications: 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) 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 (re- forwarding) can be removed by storing all h mini-batch transportation plans of size m x m which can be represented by a sparse matrix of size x for saving memory.
Discrepancy-basee 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 mini- batches, the final value of the BoMb-OT can be evaluated at the end of the second step without the fe-roward step. Similar the previous algorithm for approximate Bayesian computation in Algorithm 5i Appendix C.4 Compared to the m-OT, the BoMb-OT needs only an additional k k OT here.
types of applications,
present the BoMb-OT
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 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 GPU with its corresponding memory. As discussed in the practical usage of the BoMb-OT, OT problems between mini-batch measures of size m X m are solved on the GPU for high computational speed in estimating the parameter of interest other statistics. On the other hand, the k x k OT problems between measures over mini-batches can be computed on the CPU's memory, So. i can be chosen larger than m. One (shared) computational memory- in this case, the largest value of l equals to m. Note that, smaller values of k and m can still be used for faster computation. We would like recall that l is usually set to 1 in almost all recent applications of OT (Genevay let al. 2018 Bhushan Damodaran et al. 2018; Tolstikhin et al. 2018 [Bernton et al.) 2019a). In our experiments, we will demonstrate that increasing Ie could improve the performance.
Computational complexity of the BoMb-OT: on particular choice of divergence d. In the experiments, we specifically consider three choices of d E (W,, WT. WyW where 7 > 0 is a chosen regularized parameter. Here, we only discuss the case when d "W. the discussion of other choices of d is in Appendix B, When d is the entropic Wasserstein distance, for each pair X." and YJm the computational complexity of approximating
The computational complexity of BoMb-OT depends
d(Pxp, Py:") via the Sinkhorn algorithm is of order (mm²/e²) where e O is some desired
accuracy (Lin et al., 2019). Hence, the computation cost of k22 OT transport plans from k.2 pairs of mini-batcher is of order O /E²) when using entropic regularization (see Appendix B for the definition). With another OT cost within h2 pairs, the total computational complexity of computing the BoMb-OT is O (m² + 1)/=?). It means that the computational complexity of computing the BoMb-OT is comparable to the m-OT, which is (k2m22/e3
The BoMb-OT's transportation plan:
We discuss in detail the sparsity of BoMb-OT's transporta- tion plan in Appendix B To visualize the BoMb-OT's plan and compare it with the m-OT's plan, we carry out simple experiment on two measures of 10 supports in Figure [5] 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.

---

Under review as a conference paper at ICLR 2022
3.2 METRIC APPROXIMATION OF THE BOMB-OT LOSS
In this section, we show that the BoMb-OT is an approximation of 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 v be two probability measures on (x))), for porp>> 1 is a positive number. The population BoMb-OT between probability measures μ. and v is defined as:
(u,
inf
y [x Pym)].
7El1(
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 R and the function d( Pxm, Pym) continuous in terms of X Y E x Then, the population BoMb-OT is well-defined metric in the space of probability measures.
The proof of Theorem T is in Appendix A The assumption of Theorem Tis mild and satisfied when d is (sliced) Wasserstein metries of order p. i.e., d E (Wp. p.
Our next result shows the approximation error between the BoMb-OT loss and the population BoMb- OT distance. We provide the following result with the approximation error when d E {Wps SW
Theorem 2 (Population BoMb-OT). Assume that p 2 1, X is compact subset of RR,, and all the possible mini-batches are considered. Then, we have (i) Di (Ma, Un) DJd DT vi]
Op(m²= / when d = Wpi (ū) [D"""" Un) - DDT (1,, v) = Op(m² / when d = SW, and 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 SW, Furthermore, the choice d = SW p 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 O(k2m log m) when d = SW, 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 Mena & Weed 2019). 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, However, for general value of the regularized parameter, it is unclear whether we still have this sample complexity n." & of the eBoMb-OT. We leave this question for the future work.
4 EXPERIMENTS
In this section, 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, et al. 2018 Deshpande et al. 2018), deep domain adaptation (Bhushan Damodaran et al. 201 2018). Experiments on color transfer (Rabin et al. 2014 Ferradans 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 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, consistently. The detailed settings of our experiments including neural network architecture, hyper- parameters, and evaluation metrics are in Appendix H We would like to recall that & is the number of
run experiments on deep generative model (Genevay
observe that the BoMb-OT OT performs better than m-OT
7

---

Under review as a conference paper at ICLR 2022
Table 1: Comparison between the BoMb-OT and the m-OT deep generative models. On the MNIST dataset, we evaluate the performances of generators by computing approximated Wasserstein-2 while we use FID score (Heusel al. 2017)
on CIFARIO and CelebA
Dataset
A m-OT(W3)
(13)
MINSTT
28.12
28.12
28.12
37.57
37.57
37.57
27.88
27.53
27.56
36.01
35,27
35,88
27.6
27.42
27.41
35.18
34,19
34.85
27.36
27.1
27.25
34.33
33.17
34,00
CIFARIO
78.34
78.34
78.34
80.51
80.51
80.51
76.20
75.59
74.25
86
62.80
76.01
74.60
74.12
62.30
62.11
58.78
75.22
74
73.33
59.68
58.94
53.44
CelebA
54.16
54.16
54.16
90.33
90.33
90.33
51.53
82.45
78.66
52.56
$1.71
50.55
73.06
72.37
72.19
51.92
51.18
71.95
69.3
Table 2: Comparison between two mini-batch schemes on deep domain adaptation on digits datasets.
We vary the number mini-batches k and report the classification accuracy on the target domain.
Sceuario
Number of epochs
m-OT
BoMb-OT
Improvement
m-UOT
<BoMb-UOT
Improvement
SVHN to MNIST
50
80
92
93.74
+1.74
98.76
98.81
+0.05
+1.26
98.83
+0.02
32
50
320
93.09
95.59
+2.40
98.83
98.90
+0.07
USPS to MNIST
25
95.86
96.16
+0.30
98.46
98.59
+0.13
16
25
160
96.04
96.48
+0.44
98.36
98.7
+0.34
32
25
320
96.22
96.71
+0.49
98.43
98.75
+0.32
mini-batcher and m is the mini-batch size. our experiments, we create mini-batches by sampling without replacement from supports of the original empirical measures.
4.1 DEEP GENERATIVE MODEL.
We now show the deep generative model result on MNIST LeCun et al. 1998), CIFARIO (32x32) (Krizhevsky 2009), and CelebA (64x64) (Liu et al. 2015) using different mini-batch schemes with two OT losses for .W2 (Deshpande et al. 2018) and W (Genevay et al. 2018). According Table [ the BoMb-OT always gives lower quantitative metrics than the m-OT all datasets with all choices of the number of mini-batches k and mini-batch ground metric d. In- terestingly, the eBoMb-OT which is the entropic regularization version of BoMb-OT (defined in Appendix B) is better than the BoMb-OT on CIFARIO, 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 al. 2018). From the table, we also observe that increasing the number of bini-batches 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 pair of mini-batches like the way has been implementing in practice. Also, based on the experimental results, recommend practitioners to consider using additional transportation problem between measures over mani-bachhes when the number of mini-batches k >>
4.2 DEEP DOMAIN ADAPTATION
In this section, we adapt two digits datasets SVHN (Netzer al., 2011) and USPS (Hull 1994) to MNIST (LeCun et al. 1998). Details about architectures of neural networks and arperrraramtterss settings are given in Appendix F.2 The classification results are illustrated 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 large number of mini-batches. We also observe that the classification accuracy on the target domain also increases as & changes from to 32 and keeping the same number of stochastic gradient updates for deep neural networks. Thus, 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
8

---

Under review as a conference paper at ICLR 2022
Source m-07
BoMb-07 OT
TuP-OT
Target
Figure 2: The transferred images of the m-OT and the BoMb-OT with k=10, m=10, and the full-OT from the most left source image to the most right larget image.
k- = 2, m 8
k= 6, m
k = 2, m = 16
h: ==
6, m 16
Figure 3: Illustration of approximated posteriors from ABC with the m-OT, the BoMb-OT, and the true posterior via kernel density estimation. The Wasserstein-2 distances from the approximated posteriors to the true posterior given next the label of corresponding mini-batch methods.
4.3 COLOR TRANSFER
Mini-batch OT has made color transfer 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 and m are presented in Appendix [D.3 Based on the experiment, we observe that the BoMb-OT provides better barycentric mapping for color transfer than m-OT with every setting of & and m. In this main paper, we show color transfer result in Figure 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.
4.4 APPROXIMATE BAYESIAN COMPUTATION (ABC)
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 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 = (Bernton et al. 2019a), we suggest that the BoMb-OT should be used with & is greater than in ABC.
5 CONCLUSION
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 to consider the optimal transport problem on the space of mini-batches with a Wasserstein-types ground metric. We prove that the BoMb-OT an approximation of a valid distance between probability measures and its entropic regularized version, (BoMb-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 (GW) and unbalanced OT (UOT).
changing the local OT to the

---

Under review as a conference paper at ICLR 2022
Reproducibility Statement: All datasets that used in the paper published and they 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 Appendix E
Ethics Statement: The paper investigates 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.
REFERENCES
Jason Altschuler, Jonathan Niles- Weed, and Philippe Rigollet. Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration. In Advances in neural information processing systems, pp. 1964-1974, 2017.
David Alvarez-Melis and Tommi Jaakkola. Gromov-wassers wassersteettt alignment of word embedding spaces. In Proceedings of the 2018 Conference Empirical Methods in Natural Language Processing, pp. 1881-1890, 2018.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International Conference on Machine Learning, pp. 214 214-223, 2017.
Genevay Aude, Marco Cuturi, Gabriel Peyré, and Francis Bach. Stochastic optimization for large- scale optimal transport. arXiv preprint / 2016.
Mare G Bellemare, Ivo Danihelka, Will Dabney, Shakir Mohamed, Balaji Lakshminarayanan. Stephan Hoyer, and Rémi Munos. The Cramer distance as solution to biased Wasserstein gradients. arXiv preprint arviv: 1005. 10743, 2017.
István Berkes and Walter Philipp. An almost sure invariance principle for the empirical distribution function of mixing random variables. Zeitschrift für Wahhsshheinlicchkettttheoo und verwandte Gebiete, 41(2):115-133 1977.
Espen Bernton, Pierre E Jacob, Mathieu Gerber, and Christian P Robert. Approximate Bayesian computation with the Wasserstein distance. Journal of the Royal Statistical Society, 2019a.
Espen Bernton, Pierre E Jacob, Mathieu Gerber, and Christian Robert. On parameter estimation with the Wasserstein distance. Information and Inference: A Journal of the IMA, 8(4):657-676.
2019b.
Bharath Bhushan Damodaran, Benjamin Kellenberger, Rémi Flamary, Devis Tuia, and Nicolas Courty. Deepjdot: Deep joint distribution optimal transport for unsupervised domain adaptation. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 447-463, 2018
S. Bobkov and M. Ledoux. 'One-dimensional empirical measures, order statistics, and Kantorovich transport distances. Memoirs of the American Mathematical Society, 261, 2019.
Nicolas Bonneel, Julien Rabin, Gabriel Peyré, and Hanspeter Pfister. Sliced and Radon Wasserstein barycenters of measures. Journal of Mathematical Imaging and Vision, 1(51):22-45, 2015.
Nicolas Bonnotte. Unidimensional and evolution methods for optimal transportation. PhD thesis, Paris 11, 2013.
Liqun Chen, Zhe Gan, Yu Cheng, Linjie Li, Lawrence Carin, and Jingjing Liu. Graph optimal transport for cross-domain alignment. In International Conference on Machine Learning, pp. 1542-1553. PMLR, 2020.
Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamonjy. Optimal transport for domain adaptation. IEFE transactions on pattern analysis and machine intelligence, 3959):1853-1866 2016.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in neural information processing systems, pp. 2292-2300, 2013.
10

---

Under review as a conference paper at ICLR 2022
Alejandro De Acosta. Invariance principles in probability for triangular arrays of b-valued random vectors and some applications. The Annals of Probability, pp. 346-373, 1982.
Ishan Deshpande, Ziyu Zhang, and Alexander G Schwing. Generative modeling using the sliced Wasserstein distance. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3483-3491, 2018.
Ishan Deshpande, Yuan-Ting Hu, Ruoyu Sun, Ayis Pyrros, Nasir Siddiqui, Sanmi Koyejo, Zhizhen Zhao, David Forsyth, and Alexander Schwing. Max-sliced Wasserstein distance and its use for GANs. In Proceedings of the IEEE Conference on Computer Vision and Patter Recognition, pp. 10648-10656, 2019.
V. Dobrić and J. E. Yukich. Asymptotics for transportation cost in high dimensions. Journal of Theoretical Probability, 8:97-118, 1995.
R. M. Dudley. The speed of mean Glivenko-Cante convergence. Annals of Mathematical Statistics, 40:40-50, 50, 1969.
Kilian Fatras, Younes Zine, Rémi Flamary, Rémi Gribonval, and Nicolas Courty, Learning with minibatch Wasserstein: asymptotic and gradient properties. In AISTATS 2020-23nd International Conference on Artificial Intelligence and Statistics, volume 108, pp. 1-20, 2020.
Kilian Fatras, Thibault Sejourne, Rémi Flamary, and Nicolas Courty, Unbalanced minibatch optimal transport; applications to domain adaptation. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 3186-3197, PMLR, 18-24 Jul 2021a. URL http:::/proceeeee/e
mlr. ww..ppres/vvee v 139 httr1/////attaala..oo html
Kilian Fatras, Younes Zine, Szymon Majewski, Rémi Flamary, Rémi Gribonval, and Nicolas Courty.
Minibatch optimal transport distances; analysis and applications. arXiv preprint arXiv:2101.0173 2021b.
Sira Ferradans, Nicolas Papadakis, Gabriel Peyré, and Jean-François Aujol. Regularized discrete optimal transport. SIAM Journal on Imaging Sciences, 1(3):1853-1882, 2014.
Jean Feydy, Thibault Séjourné, François-Xavies Vialard, Shun-ichi Amari, Alain Trouve, and Gabriel Peyré. Interpolating between optimal transport and mmd using Sinkhorn divergences. The 22nd International Conference Artificial Intelligence and Statistics, pp. 2681-2690, 2019.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Lactitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalic T.H. Gayraud. Hicham Janati, Alain Rakotomamonjy, Ievgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22((8):1-8, 2021. UtLbttp: /mm org/
html
N. Fournier and A. Guillin. On the rate of convergence in Wasserstein distance of the empirical measure. Probability Theory and Related Fields, 162:707-738, 2015.
Aude Genevay, Gabriel Peyre, and Marco Cuturi. Learning generative models with Sinkhorn divergences. In International Conference on Artificial Intelligence and Statistics, pp. 1608-1617, 2018.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by a two time-scale update rule converge to a local nash equilibrium. In Advances in Neural Information Processing Systems, pp. 6626-6637, 2017.
Nhat Ho, XuanLong Nguyen, Mikhail Yurochkin, Hung Hai Bui, Viet Huynh, and Dinh Phung. Multilevel clustering via Wasserstein means. In International Conference Machine Learning, pp. 1550- 2017.
Jonathan J. Hull. A database for handwritten text recognition research. IEEE Transactions on pattern analysis and machine intelligence, 16(5):550-555 1994.

---

Under review as a conference paper at ICLR 2022
Emmanuel Klinger, Dennis Rickert, and Jan Hasenauer. Pyabe: distributed, likelihood free inference. Bioinformatics, (1420):3591-599)) 2018.
Soheil Kolouri, Phillip E Pope. Charles Martin, and Gustavo K Rohde. Sliced Wasserstein
auto-encoders, In International Conference on Learning Representations, 2018.
Soheil Kolouri, Kimia Nadjahi, Umut Simsekli, Roland Badeau, and Gustavo Rohde. Generalized sliced Wasserstein distances. In Advances in Neural Information Processing Systems, pp. 261-272, 2019.
Alexander Korotin, Lingxiao Li, Aude Genevay, Justin Solomon, Alexander Filippov, and Evgeny Burnaev. Do neural optimal transport solvers work? a continuous Wasserstein-2 benchmark. arXiv preprint arXiv:2 (06.0/9544 2021.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86( 86(11):2278-23224 1998.
Chen-Yu Yu Lee, Tanmay Batra, Mohammad Haris Baig, and Daniel Ulbricht. Sliced wasserstein discrepancy for unsupervised domain adaptation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10285-10295, 2019a.
John Lee, Max Dabagia, Eva Dyer, and Christopher Rozell. Hierarchical optimal transport for multimodal distribution alignment. In Advances in Neural Information Processing Systems, pp. 13474-13484, 2019b.
Jacob Leygonie, Jennifer She, Amjad Almahairi, Sai Rajeswar, and Aaron Courville. Adversarial computation of optimal transport maps. arXiv preprint arXiv: (XX:::190600999) 2019.
Lingxiao Li, Aude Genevay, Mikhail Yurochkin, and Justin M Solomon. Continuous regularized Wasserstein barycenters. Advances in Neural Information Processing Systems, 33, 2020.
Tianyi Lin, Nhat Ho, and Michael Jordan. On efficient optimal transport: An analysis of greedy and accelerated mirror descent algorithms. In International Conference Machine Learning, pp. 3982-3991, 2019.
Tianyi Lin, Chenyou Fan, Nhat Ho, Marco Cuturi, and Michael Jordan. Projection robust Wasser- stein distance and Riemannian optimization. In Advances in neural information processing systems, 2020a.
Tianyi Lin, Zeyu Zheng, Elynn Y Chen, Marco Cuturi, and Michael Jordan. On projection robust op- timal transport: Sample complexity and model misspecification. arXiv preprint 1rX1v. 2020b.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference Computer Vision (ICCV), December 2015.
Antoine Liutkus, Umut Simsekli, Szymon Majewski, Alain Durmus, and Fabian-Robert Stöter.
Sliced-Wasseree Wasserstein flows: Nonparametric generative modeling via optimal transport and diffusions. In International Conference on Machine Learning, pp. 4104-4115 13. PMLR, 2019.
Dixin Luo, Hongteng Xu, and Lawrence Carin. Hierarchical optimal transport for robust multi-view learning. arXiv preprint arXiv:2006.03/66 2020.
Ashok Makkuva, Amirhossein Taghvaci, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. In International Conference Machine Learning, pp. 6672-6681 PMLR, 2020.
Anton Mallasto, Guido Montúfar, and Augusto Gerolin. How well do WGANs estimate the Wasser- 2019.
stein metric? arXiv preprint arXiv:
Gonzalo Mena and Jonathan Weed. Statistical bounds for entropic optimal transport: sample complexity and the central limit theorem. In Advances in Neural Information Processing Systems, 2019.
12

---

Under review as a conference paper at ICLR 2022
Boris Muzellec and Marco Cuturi. Subspace detours: Building transport plans that optimal on subspace projections. In Advances in Neural Information Processing Systems, pp. 6928-6928, 2019.
Kimia Nadjahi, Alain Durmus, Umut Simsekli, and Roland Badeau. Asymptotic guarantees for learning generative models with the sliced- Seiced-Wasserne Wasserstein distance. In Advances in Neural Information Processing Systems, pp. 250-260, 2019.
Kimia Nadjahi, Valentin De Bortoli, Alain Durmus, Roland Badeau, and Umut Şimgekli. Approxi- mate Bayesian computation with the sliced-Wassenee Wasserstein distance. ICASSP 2020-2020 IEEE International Conference Acoustics, Speech and Signal Processing (ICASSP), pp. 5470-5474. IEEE, 2020.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng. Reading digits in natural images with unsupervised feature learning. In NIPS Workshop Deep Learning and Unsupervised Feature Learning, 2011.
Khai Nguyen, Nhat Ho, Tung Pham, and Hung Bui. Distributional sliced- Wasserstein and applications to generative modeling. In International Conference Learning Representations, 2021a. URL
ltttps://peenzeiii
net/ 1o
Khai Nguyen, Son Nguyen, Nhat Ho, Tung Pham, and Hung Bui. Improving relational regular-
ized autoencoders with spherical sliced fused Gromov-Wassersten In International Confer- ence on Learning Representations, 2021b. URL httphhtt////perreeeee
net/ forum?id
DigD7FWL.233
François- Pierre Paty and Marco Cuturi. Subspace robust Wasserstein distances. In International Conference on Machine Learning, pp. 50755077255 5081, 2019.
O. Pele and M. Werman. Fast and robust earth mover's distance. In ICCV. IEEE, 2009.
Michaël Perrot, Nicolas Courty, Rémi Flamary, and Amaury Habrard. Mapping estimation for discrete optimal transport. Advances in Neural Information Processing Systems, 29:4197-4205, 2016.
Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications data science. Foundations and Trends@ in Machine Learning, 1 11(5-6)3355-677 2019
Julien Rabin, Sira Ferradans, and Nicolas Papadakis. Adaptive color transfer with relaxed optimal transport. In 2014 IEEE International Conference Image Processing (ICIP), pp. 4852 -855.. IEEE, 2014.
Tim Salimans, Han Zhang, Alec Radford, and Dimitris Metaxas. Improving GANs using optimal transport. In International Conference Learning Representations, 2018.
Filippo Santambrogio. (Euclidean, metric, and Wasserstein] Gradient Flows: An Overview. Bulletin of Mathematical Sciences, 7(1)877-154, 2017.
Max Sommerfeld, Jörn Schrieber, Yoav Zemel, and Axel Munk. Optimal transport: Fast probabilistic approximation with exact solvers. Journal of Machine Learning Research, 20(105): 2019.
Jan Stanczuk, Christian Etmann, Lisa Maria Kreusser, and Carola-Bibiane Schönlieb. Wasser- stein gans work because they fail (to approximate the Wasserstein distance). arXiv preprint arXiv:2103.01678, 2021.
Vayer Titouan, Nicolas Courty, Romain Tavenard, and Rémi Flamary. Optimal transport for structured data with application on graphs. In International Conference on Machine Learning, pp. 6275-6284, 2019.
Iyya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. Wasserstein auto-encoders. In International Conference on Learning Representations, 2018.
Tina Toni, David Welch, Natalja Strelkowa, Andreas Ipsen, and Michael PH Stumpf. Approximate Bayesian computation scheme for parameter inference and model selection in dynamical systems. Journal of the Royal Society Interface, ((71))187-204 2009.
13

---

Under review as a conference paper at ICLR 2022
Cédric Villani. Optimal transport: old and new, volume 338. Springer Science & Business Media, 2008.
Cédric Villani. Topics in optimal transportation, volume 58. American Mathematical Soc., 2021.
Hongteng Xu, Dixin Luo, and Lawrence Carin. Scalable Maamoo--asssesteinn learning for graph partitioning and matching. In Advances neural information processing systems, pp. 3052-3062, 2019.
Renjun Xu, Pelen Liu, Liyan Wang, Chao Chen, and Jindong Wang. Reliable weighted optimal transport for unsupervised domain adaptation. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 4394 4403, 2020.
Mikhail Yurochkin, Sebastian Claici, Edward Chien, Farzaneh Mirzazadeh, and Justin M Solomon. Hi- erarchical optimal transport for document representation. In H. Wallach, H. Larochelle, A. Beygelz- imer, F. d' d'Alché-Buc, Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates Inc., 2019, URL http:///......tooo/ooooc.e.oa.os http://2/11//11e//1//0/000 http:///wwZLL2aooaaaaaaaeeee3333
14

---

Under review as a conference paper at ICLR 2022
Supplement to "On Transportation of Mini-batches: A Hierarchical Approach"
In this supplementary material, we collect several proof' and remaining materials that were deferred from the main paper. In Appendix A we provide the proofs of the main results in the paper. We present additional materials including discussions about computation complexity of the BoMb-OT, the connection of the BoMb-OT hierarchical optimal transport, entropic regularization of BoMb-OT, sparsity of the BoMb--T,, and extension with non-OT mini-batch metrics in Appendix [8] Furthermore, we provide description of applications and the BoMb-OT algorithms in those applications in Appendix C Additional results of presented experiments in the main text in Appendix D. and their corresponding settings Appendix
A PROOFS
In this appendix, we give detailed proof's of theorems that stated in the main paper.
A.1 PROOF OF THEOREM
We first prove that for any probability measures and (P))), there exists optimal transporta- tion plan y such that
PT v) Exx" y"'"" [d (p Pym)
3m From the definition of the BoMb-OT, there is sequence of transportation plans "/n E II( /4, v)) such that
[xm ym)m,, [d( Pxm,, Pym)) )] - Da(u,, v)
as n c0. Since II( @m @m
μ 1 is compact the weak topology (Villani 2008), 7/n weakly converges
some 7thh II[ B0mm 60m Since d( ((2,,, Pym) is continuous in terms of xm nn E xm
application of Portmanteau theorem leads to
lim Exx",y")<nn [d(P ((xx,,, Pym)) 2 (mx [d(Pxm, Py.)]] 2-000
Putting the above results together, we obtain
DT (1, v) E(x", [d(Pxm Pym)].
Therefore, there exists optimal transportation plan 7" such that equation 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 (M v) 0 if and only if u. In fact, from the definition of Da(.,.), if we have two probability measures μ and such that D(p, v) 0, we find that
inf 7 7ETT (")
E(x, id(P Pym) ) = 0.
3m Since d is a metric, implics that there exists transportation plan x- E II M V' such that Pxm Py -lmoott everywhere. It demonstrates that for each there exists permutation a of (1. (1....,no m) such that 2; = (0) for all i E [m]. Now, for any test function f - R, we have
mm 1I f(x))d 6o (x) = f( xm XX
o xm x
II fax))dy =
II jml
f (v,, (0) (x.
f
TT
II fd (z. v)
pm x
II fludd
2m
15

---

Under review as a conference paper at ICLR 2022
Since Sxm II;-, (( = (Ix ((z)du((1)) and Sxm II f(uid (r)
(fx J(v ((v)) ((v))d((1)) "h the above results show that
Jx (on)dp(o))) J ((v))du(v)) ((in))ddon
for any test function f: X - R. It indicaasss that μ v. On the other hand, if /z v, have V. We can construct the transportation plan 7(x, y) (r) 1 (x) for all (r. v). Then, we find that
Dz (1,, v) S ((x,y)>> P x Py )]
Therefore, we obtain the conclusion that ((, v) = 0 if and only if p = V.
Now, we demonstrate the triangle inequality property of Da(,,.), namely, will show that for any probability measures μl, f2, and (ss, we have
(21, + DJ (H2: 2 Dj (H1,
The proof of this inequality is direct application of gluing lemma (Berkes & Philipp 1977,
De Acosta 1982). In particular, for any transportation plans ? E II 8 8m
1. (42) 2) and 7/2 E
II IH c0m 80m
1 1 3). we can construct probability measure on xm x xm x such that E(., Xm)
()...) and E(XT ...................... = -.(...) Therefore, we find that
Bx [d(Pxm, Py )] + ((x,2)~7al [d (P Pz Sccmps d(Pam, Pam) + d(P. m. pam) pm
2 d(Pam, P2m )d (sr, y. 2 D'd (s))
(Xm3a
Taking the infimum on the LHS of the above inequality with respect Y1 and 72, we obtain the
triangle inequality property of (..... As consequence, reach the conclusion of the theorem.
A.2 PROOF OF THEOREM
2
6m
To ease the presentation, for any product measures i and we denote the following loss between
@m and V as follows:
8m
6om 60m V)== Drr
inf Rm [d(Pxm. Pym YET J v")
From the definitions of the BoMb-OT losses, we have (TT = D 60m 60m 80m Dii (pn, Vx)) = Din [A ).). Using the similar proof argument as that of Theorem we can check that D satisfies the triangle inequality, namely, for any product measures we have Dim 20m oomn + Dar
V
and
3m &m @m and
com 60m
2 Dm
60m. c0m 1 17
From the above definition and properties of D, we find that
km (Pa, vx) - DTT (H,
= DD 3m 8m - Dmm
3m
< Dm (4k
sem som Dan
30m c0m
+
Dr
D mm. 60m
< tm @m @m
lk: #
D
(i) Since d = Wp where p 2 1, we have
1/p
m |X Yo(o 1P
d(Px (P Py
inf
m m/p
| Y.ll,
1/p i=1
16

---

Under review as a conference paper at ICLR 2022
where the infimum is taken over all possible permutations of {1, ........... The above inequality indicates that
m |
8m 3m Dd (PR P
inf B 3m )))
Yill ml-1/p. Wi Wi(Pns 11).
m
/p
yEII
Since X' is a compact subset of TR,, the results of (Dudley 1969 Dobrić & Yukich 1995
Fournier
& Guillin 2015) show that W1 (Hnnn u) = =O -1/N). Collecting these results together, we have
m 1/pp
50m P
&om
= Op n2/NN'
With similar argument, we also find that DT v = Op together, we reach the conclusion of part (i) of the theorem.
m /m-1//p a1/N
Putting all the above results
(ii) The proof of part (ii) follows directly from the argument of part (i) and the results that SW, (P). = Op(n 1/2) and SW = Op(n 1/2) when X' is compact subset of
R Bobkov & Ledoux 2019). Therefore, we obtain the conclusion of part (ii) of the theorem.
B ADDITIONAL MATERIALS
Computational complexity: We now discuss the computational complexities the BoMb-OT loss when d =
When WT for some given regularized parameter O, can the Sinkhorn algorithm compute d( d(Px,m, Py;") for any 1 S i, S k. The computational complexity of the order O(m²). Therefore, the computational
of the Sinkhorn algorithm for computing complexity of computing the cost matrix from the BoMb-OT is of the order (m.) Given the cost matrix, the BoMb-OT loss can be approximated with the complexity of the order O(k2/E²) where stands for the desired accuracy, As consequence, the total computational complexity of computing the BoMb-OT is Omm²k² - -k//=") /²?)
When d SW,,, p. the computational complexity for computing d(Pxç", (yz))) is of the order (m log m) for any 1 I i.j S k. It shows that the complexity of computing the cost matrix from the BoMb-OT loss is of the order (mm(lg m)) Given the cost matrix, the complexity of approximating the BoMb-OT is at the order of Ok e²2). Hence, the total complexity is at the order of O (m k2 + k222
Connection to hierarchical OT: At the first sight, the BoMb-OT may look similar to hierarchical optimal transport (HOT). However, we would like specify the main difference between the BoMb- OT and the HOT. In particular, on the one hand, the HOT comes from the hierarchical structure of data. For example, Yurochkin 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, hierarchical formulation of OT is proposed in (Lee et al.) (219b) 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 we want to improve the quality of mini-batch loss and its transportation plan.
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 (cBoMb-OT) between two probability measures is and v admits the following form:
Entropic regularized population BoMb-OT:
(H.v)== min 8m
Lom Rxxmym)jayl000 [d(Pxm, Py X & V
(4)
?EII
where À > 0 stands for 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 À + O, we have EDa(u, v) converges Pa(H, v). On the other hand, when 8m &m
oc, the joint distribution y approaches / & and EDa(p, v) converges to Ua(u, v). The
17

---

Under review as a conference paper at ICLR 2022
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 BoM---TT: We assume that we 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 & mini-batches of size m, the m-OT"s transportation plan contains at most mm non-zero entries. On the other hand, with & mini-batches of size m, the BoMb-OT's transportation plan has at most km non-zero entries, For W (r 0), the optimal transportation plans contains most n² non-zero entries. this case, the m-OT provides transportation plans that have at most m2 non-zero entries. the BoMb-OT's transportation plans contain at most km?? 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 B could be any discrepancy between empirical measures mini-batcher 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, is not easy define the notion of a transportation plan in these cases.
C APPLICATIONS AND BOMB-OT ALGORITHMS
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--T.
C.I DEEP GENERATIVE MODEL
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 0, which belongs to parameter space θ. Each value of 0 induces a model distribution the data space, and we want find the optimal parameter 0* which has jíg. the closest distribution to the empirical data distribution Vn, under a discrepancy (e.g. Wasserstein distance). In deep learning setting, f is the weight of deep neural network that maps from a low dimensional manifold Z to the data space X, and the model distribution μ4 is a push-forward distribution Gotp(z) for p(z) is white noise distribution on Z (e.g- (0, I)). By using haii-batching schemes, we can estimate this parameter by minimizing the BoMb-OT loss in Definition
Task description:
over
f- arg min D (M0,m; 0E
(5)
where n is the empirical distribution of ug. The mentioned optimization is used directly in learning generative model on MNIST dataset with = (SW2, W3) 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 natural distance on the space of images such as
CelebA, metric learning was introduced by (Genevay et al. 2018 Deshpande et al. 20018) as key step in the application of generative models. The general idea is learn parametric ground metric cost ca:
cs(x, v) (1f(()) (2)ll2
where fo : X + Rhth is non-linear function that map from the data space to a feature space where La distance is meaningful.
The methodology to learn the function fo depends on the choice of d. For example, when d = W authors in (Genevay al. 2018) seek for 6 by solve following optimization:
max D (f fetva))
80
18

---

Under review as a conference paper at ICLR 2022
Algorithm 1 Mini-batch Deep Generative Model with BoMb-OT
Input: k, m, data X. prior distribution p(z), chosen mini-batch loss LDGM E (W2, W SW2 Initialize Gg GPU;
while does not converge do
On computer memory
Sample indices uniformly Ixm...., Ixg on X'
Sample 7 ,....ZK 7 from p(=)
Intialize E R
for 1 to & do
for = to do
On GPU, autograd off
Load X to GPU from Lxm
Load 2 GPU
Compute yj (z))
Compute Cif t LOOM Pmm)
end for end for
computer memory-
Solve π OT Uko C*) (entropic OT for eBoMb-OT)
grad, f- 0
for 1 to k do
for to do
if Fij + () then
On GPU, autograd on
Load X." GPU from Ix!"
Load 2nnh to GPU
Compute ; + Go (Z
Compute Ci. + LOMM Px Pym)
grad, + grad, + Tay VoCig
end if
end for end for <- Adam (0, grad,)
On GPU
end while
In (Deshpande al. 2018), another metric learning technique is used for d = SW
max max E pp n. green, [log (x + log(1 (v(f(())))
(8)
where 96 R [0, 1].
Learning the metric or fø is carried out simultancously with learning the generative model in practice, namely, after gradient step for the parameter 0, á will be updated by gradient step.
C.2 DEEP DOMAIN ADAPTATION
We adapt the BoMb-OT into DeepJDOT (Bhushan Damodaran al. 2018 which is famous unsupervised domain adaptation method based on the m-OT. In particular, we aim to learn an embedding function Go X' + Z which maps data the latent space; and 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 and the size of mini-batches m, the goal is to minimize the following objective
19

---

Under review as a conference paper at ICLR 2022
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 Gg (parametrized by 0), Fo (parametrized by Φ)
while (0, Φ) do not converge do
On computer memory
Sample indices uniformly IsT y , Jm y on (S, Y) Sample indices uniformly r fr on T
Initialize R
for = 1 to k do
for j =
do
On GPU, autograd off
Load (S Y!n) to GPU from Is!"
Load T GPU from Ir
Compute mC &- LIA (5 Y TT Ga, F&) (Equation 10
Um + (
Compute Ci <- OT(um Um, mC)
end for
end for
On computer memory-
ux < .......,,,
Solve 7 t OT uks C²) (entropic OT for eBoMb-OT)
grad, 0
grad, d f 0
for 1 to do
On GPU. autograd on
Load , Y!") to GPU from Is y
grad, < grad, + (Y," F FS (Equation
grad - grad + k V VYL (S,)( (Equation
i
for
do
if %ij * 0 then
On GPU, autograd on
Load 1mth to GPU from 1Tm
Compute mC <- T.A [STA Y Tm, Go, F&) (Equation 10
Compute Cij OT
O0 (Um Umm mC)
grad, + grad, + 7ij CoCij
grad & grad Ye VoCi
end if
end for end for
GPU
0 < Adam (,, gragg))))))) Φ + Adam (0, grad
end while
function:
m
LDA Gg.
min
L. (Wy, F Fo(o((s)))
k m
min HEII MA MAL
min T (t6m Um
(C ym mmm.ymm πt)
where L, the source loss function (e.g. classification loss), STh S S are source batche which are sampled with or without replacement from the source domain S Xm y ..... pm are corresponding labels of STm ,...,Sk 5m. with S.m = {$11..... sim and Y == (ya..... Vim Similarly, TTm... ......., T (T = (li,.... "iu,,..,,any (am)) are target mini-batches which sampled with or
20

---

Under review as a conference paper at ICLR 2022
G.F without replacement from the target domain 7 E am The cost matrix S Mr" II." defined as follows:
denoted mC
mCigane = al|G(( G(( (|2 Le (z)(()))))).
(10)
where Lt is the target loss function, or and d are hyper-parameters that control two terms,
C.3 COLOR TRANSFER
In this appendix, carry compare the m-OT and the BoMb-OT different domains of images, and we demonstrate visually the interpolation property of the eBoMb-OT by varying the regularization parameter.
experiments color transfer with various settings. In detail,
Methodology: In our experiment, we first compress both the source image and the target image using K 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 cluster of the source image based on the value of the corresponding cluster centers the target image with the previously found plans. We present the algorithm that we use in color transfer with the BoMb-OT in Algorithm 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: m, T source image X, E R target image X, Rxx
Initialize Y, E R2=3
for to 7
Initialize C B
Sample indices Ixm,..., Ix from X, Iy from Xe Load X from 1x.
Sample indices Iym,....
for = 1 to & do
for to do
Load Y,M from lym
Compute cost matrix M between X!" and Y!"
< arg min un.) (M, T)
end for end for
7 arg min m (C,7)
for = 1 to k do
for =
do
if 7/j * 0 then
Load X," from Ix."
Load Y.T from Iym
Compute cost matrix M between X and Y
arg min me (M,
Ysli + mnj Xt Xx 1rgg.
end if end for
end for end for
Output: Y.
C.4 APPROXIMATE BAYESIAN COMPUTATION (ABC)
In this appendix, details of Approximate Bayesian Computation (ABC) and the usage of mini-bur batches with BoMb-OT are discussed.
21

---

Under review as a conference paper at ICLR 2022
Algorithm 4 Approximate Bayesian Computation
Input: Generative model p(x, 0), observation {z,)i=1, number of iterations T', discrepancy
measure D, tolerance threshold e, number of particles and summary statisties s (optional). for t = 1 to T do
repeat
Sample ~ p(0)
Sample (3 1= ())
until D(s U: (1 S
02 2==07
end for
Output: (0:)
The central task of Bayesian inference is to estimate distribution of parameter 0 E θ given n data points {z, |i= and generative model p(z, 0). Using Bayes rule, the posterior can be written as p(0 = p( (9)p(0) Generally, the posterior is intractable since we cannot p(2)).))
Review on ABC:
evaluate the normalizing constant p( (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(r (0) cannot be evaluated such as implicit generative models. In these cases, Approximate Bayesian Computation (or likelihood free inference) is 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 obtain posterior samples of 6. 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 -> However, the performance of ABC depends on the choice of the summary statistics (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 lot of attention from ABC's community because they provide 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 in Algorithm The detail of the application of BoMb-OT in ABC is given in Algorithm
D ADDITIONAL EXPERIMENTS
In this appendix, we provide additional experimental results that are not shown in the main paper. In Appendix D.I we present random generated images on MNIST dataset, CIFAR 10 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.3 Finally, we investigate the behavior of the m-OT and the Bo)-OTT in estimating the optimal transportation plan in Appendix [.6]
D.1 DEEP GENERATIVE MODEL.
Generated images: We show the generated images on MNIST dataset in Figure 4 the generated images CelebA in Figure the generated images CIFARIO in Figure For both choices of d (SW2, WT), ). we can see that the images from (c)BoMb- OT is more realistic than the images from m-OT. This qualitative result supports the quantitative in Table in the main text.
22

---

Under review as a conference paper at ICLR 2022
Algorithm 5 Approximate Bayesian Computation with BoMb-OT
Input: Generative model p(a, 0), observation {x,))=1 number iterations T, optimal transport discrepancy measure Dor, tolerance threshold number of particles m, number of mini-batches k. and summary statistics $ (optional). for t 1 to T do
repeat
Sample p(0)
Sample indices Tx.",..., Ixp- from {
Sample {yr J1=1 ~ p(x10)
Sample indices Iym,,..., by; from {wil1=1
Initialize E R
for 1 to k do
for j = to l: do
Load X,", Y from their indices Lxm Iym
Cj 4- Dor (x y
end for end for
u 1 1)
D + OTu Uk C)
until Se 0g
end for
Output: (0.)
7
0
7
9
8
7
2
0
1
4
8
7
3 U
1
4
3
f
3
3
7
6
6
7
6
1
3
7
9
6
7
2
/
i
3
7
6
7
7
/
1
0
0
7
7
8
7
3
3
0
3
2
7
9
1
3
3
6
7
7
9
7
3
0
6
9
&
2
0
9
7
0
7
Z
&
2
m-OT(SW,))
BoMb-OT (T(SW,)
eBoMb eBoMb-OT(SW-)
2
5
6
9
5
2
8
7
9 6
2
7
S
/ S
9 6
/
9
9
/
8
3
5
2
8
4
7
6
7
2
1
/
8
1
7
2
1
1
9
2
9
9
7
3
/
/
4
6
0
I
5
9
/
4
2
/
1
5
6
9
1
2
2
1
1
O
5
5
/
8
9
9
7
7
/
3
8
5
7
9
7
3
2
3
9
7
7
/
3
3
m-OT(W3)
BoMb-OT(W3)
eBoMb- (BoMb-OT(W:))
Figure 4: MNIST generated images from the m-OT and the (e)BoMb-OT for (k, m) (1, 100).
Computational speed: Table 3 details the computational speed of deep generative model when using mini-batch size m = 100. In general, the more the number of mini-batches is, the more complex the problem is. Increasing the number of mini-batches & does decrease the number of iterations per second in all experiments. For both datasets, using the sliced Wasserstein distance (d = SW,) leads higher iterations per second compared to the entropic Wasserstein distance (d W3) 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 = 4, and 8 respectively. In contrast, the BoMb-OT is faster than the m-OT for CIFARIO 10 and CelebA dataset because we set a high value for the regularized parameter (e.g. T = 50, 40). For the CelebA dataset, although the eBoMb-On OT is the slowest method when d SWa, 2, becomes the fastest approach if set d to WT. The reason for this phenomenon is because of the sparsity of the BoMbOOT's transportation plan between measures
batchbs In particular, pairs of mini- mini-batches that zero be skipped
23

---

Under review as a conference paper at ICLR 2022
m- OT(SW))
BoMb-OT(SIH)
2BoMb-SWa
m-OT(W3)
BoMb-OT(W3))
eBoMb-OT(W3))
Figure 5: CelebA generated images from the m- OT and the (e)BoMb OT for (k.0)=(4,100).
m-OT(SW22)
BoMb-OT(SW:)
eBoMb-SW2
m-OT(WE)
BoMb-OT(W))
eBoMb-OT(W3)
Figure 6: CIFAR10 generated images from the m-OT and the (e)BoMb-OT for (k,0)=(4,100).
Table 3: Number of iterations per second of deep generative models
Dataset
m-OT((W:)
BoMb-on!W)))
eBoMb-Or (W2))
m-OT(SWW)
BoMb-OI(SWW)
cHoMb-OT
MNIST
27.27
27.27
27.27
54.55
54.55
54.55
2
7.89
5.88
4.62
18.75
15.00
11.54
4
2.54
2.11
1.36
6.00
5.36
3,41
0.85
0.79
0.44
1.70
1.74
0.95
CIFAR10
22.73
22.73
22.73
25.00
25.00
25.00
2
6.94
5.81
7.35
9.62
10.00
7.58
4
2.05
2.31
2.16
3.21
4.17
2.40
8
0.53
0.78
0.59
0.97
1.54
0.71
CelebA
6.78
6.78
6.78
9.93
9.93
9.93
2
1.92
1.75
2.52
3.38
3.55
2.65
0.45
0.54
0.68
1.03
1.40
0.78
8
0.11
0.15
0.18
0.29
0.51
0.22
(see in Algorithm So, the (e)BoMb-OT can avoid estimating gradient of the neural networks of bad pair of mini-batches.
24

---

Under review as a conference paper at ICLR 2022
Table 4: Effect of changing the mini-batch size deep domain adaptation
Scenario
k
Number of epochs
m-OT
BoMb-OT
Improvement
SVHN to MNIST
10
32
64
69.41
88.49
+19.08
20
32
128
89.69
93.54
+3.85
50
32
320
93.09
95.59
+2.40
USPS to MNIST
50
73.92
93.15
+19.23
10
32
100
92.96
95.06
+2.10
20
32
200
95.85
96.15
+0.30
Table 5: Comparison between the two mini-batch schemes deep domain adaptation on Office-
Home datasets.
Methods
A2C
A2P
A2R
C2A
C2P
C2R
P2A
P2C
P2R
R2A
R2C
R2P
Avg
m-UOT
56.22
75.04
80.51
64.81
74.50
75.08
52.23
80.01
74.21
60.00
82.97
70.15
cRoMb-UOT
55.90
75.13
80.51
66.01
74.52
75.63
65.97
53.22
79.99
74.04
60.09
83.29
70.39
m-UOT
55.42
75.13
80.45
65.88
73.89
74,50
65.76
52.90
79.96
74.21
59.84
83.17
70.09
cBoMb-UOT
56.06
75.15
80.63
65.88
73.91
75.30
65.84
53.24
80.17
74.29
60.18
83,31
70.33
m-UOT
56.04
75.08
80.45
65.76
73.40
74.68
65.76
53.04
79.99
74.29
60.14
83.24
70.15
cBoMb-UOT
56.22
75.56
65.84
74.30
74.89
66.05
53.08
80.03
74.41
70.35
Table 6: Computational speed of deep DA
Table 7: Computational speed of deep DA when changing m
when changing
Scenario
k
m-OT
BoMb-OT
Scenario
m-OT
BoMb-OT
SVHN to MNIST
2.26
3,05
SVHN to MNIST
10
0.33
0.61
16
0.60
1.06
20
0.33
0.60
32
0.15
0.32
50
0.15
0.32
USPS to MNIST
6.00
9.00
USPS MNIST
5
0.49
1.00
1.80
2.57
10
0.49
1.00
32
0.45
0.75
20
0.46
1.00
D.2
DEEP DOMAIN ADAPTATION
In this section, we compare the performance of two mini- schemes digits and Office-Home datasets.
Comparison between the m-OT and BoMb-OT digits datasets: As seen in Table the BoMb-OT produces better classification accuracy than the m-OT in all experiments. In addition, it leads to 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 Jillutraass the performance of m- m-UOT and eBoMb-UOT on Office-Home dataset when changing the number of 8. When 4, eBoMb-UOT achieves the classification accuracy higher
mini-batches & from 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 Similar to deep generative model, we observe a phenomenon that the speed of both the m-OT and 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 iteration 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.
BoMb-OT decreases
25

---

Under review as a conference paper at ICLR 2022
D.3 COLOR TRANSFER
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-transferre 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 'I' (see in to 5000. We show the source images, target images, and corresponding transferred images in Figure 7 Figure TO It is easy to see that the transferred images from the BoMb- OT look more realistic than the m- m-OT and the color is more similar 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.
Algorithm 3)
Computational speed: For l = 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 ls x l OT is not too expensive while it can improve the color palette considerably.
D.4 APPROXIMATE BAYESIAN COMPUTATION (ABC)
Settings: We use the same setup in (Nadjahi et al. 2020): there {æ, 1a i.dd from multivariate Gaussian (pe, IN) where N is the dimension, μ4, ~ (O, Iv) and 0, = 4. The task is estimate the posterior of a² under the imaginary assumption that o² follows inverse gamma distribution IG(1, , 1). Under these assumptions, the posterior of g² has the form [ + n + 2 E ||=; v.113) For ABC, we use the m-OT and the BoMb-OT with the Wasserstein-2 ground metric for the acceptance-rejection sampling's 's criteria in sequential Monte Carlo ABC (Toni et al. 2009) (using implementation from py ABC (Klinger et al., 2018) with 100 particles and 10 iteration). The mini-batches' size is set in {8, 16, 32} and the number of mini-batcher is in {2, 4, 6, 8}.
100 observations
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 E {8, 16, 32} and & E 12 4. 6, 8}. From these results, we see that a bigger m usually returns better posterior in both the m- m-OT and the BoMb-OT cases. Similarly, increasing also improves the performance of the two methods. Moreover, these graphs strengthen the claim that the BoMb-OT outperforms the m-OT every setting of m) since its posteriors are always closer to the true posterior than those of the m-OT in both visual result and Wasserstein distance.
D.5 NON- NON-PARAMETRIA GENERATIVE MODEL VIA GRADIENT FLOW
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 distribution μ. It leads to the functional optimization problem:
min v),,
(11)
where a predefined discrepancy between two probability measures. So, gradient flows can be constructed:
0ip = -V, D(u,, v)
(12)
We follow the Euler scheme to solve this equation as in (Feydy al. 2019), starting from initial distribution at time t = O. In this paper, we choose D be (c)BoMb W2 and m- W2 for the sake of comparison between them.
26

---

Under review as a conference paper at ICLR 2022
m=10,
k=10, T=5000
Source
m-OT
BoMb-OT
Target
m=20, k= 20, T=5000
Source
m-oT
BoMb-OT
Target
Figure 7: Experimental results on color transfer for the m-OT, the BoMb-OT on natural images with (ki, m) = (10, 10), (k, m) (20, 20), and T = 5000. Color palettes are shown under corresponding images.
We first consider the toy example as in (Feydy et al., 2019) and present our results in Figure 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 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, 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.
Let ja and v 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 female one by creating 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 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 shown in Figure 14 also show the closest female image (in sense of
CelebA:
27

---

Under review as a conference paper at ICLR 2022
m=10,
k=10, T=5000
Source
m-OT
BoMo-OT
Target
m=20, 20 T=5000
Source
m-OT
BoMb-OT
Target
Figure 8: Experimental results color transfer for the m-OT, the BoMb-OT arts with (k, (10, 10).
(ks, m) = (20, 20), and T 5000. Color palettes are shown under corresponding images.
Le 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- -T provides better flow than m-OT does.
D.6 TRANSPORTATION PLANS
In this appendix, we investigate deeper the behavior of the m-OT and the (c)BoMb-OT in estimating the transportation plan. We present toy example illustrate the transportation plans of the m-OT, the BoMb-OT, and the original OT (full-OT). In particular, we sample two empirical distributions
and P,, where 10. from N N(4 |
and N [ [is -uN) respectively.
Transportation plans:
We plot the graph of samples matching and transportation matrices in Figure 15 When the size of mini-batche equals 2 and the number of mini-batches is set to be 20, the
28

---

Under review as a conference paper at ICLR 2022
m=10,
k=10, T=5000
Source
m-oT
BoMb-OT
Target
m=20, k=20, T=5000
Source
m-OT
BoMb-OT
Target
Figure 9: Experimental results on color transfer for the m-OT, the BoMb-OT natural images with (k, m) = (10, 10), (k, m) = (20, 20), and T = 5000. Color palettes are shown under corresponding images.
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 acceptable. Incrcasing the size of however its matching and its matrices are visibly still not good solutions. contrast, the BoMb-OT is able to generate nearly optimal matchings and transportation matrices. Next, 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 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 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-batcher size.
to 8, m-OT's performance improves significantly
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
29

---

Under review as a conference paper at ICLR 2022
m=10,
k=10, T=5000
Source
m.oT
BoMb-OT
Target
m=20, k=20, T=5000
BoMb-OT
Target
Figure 10: Experimental results on color transfer for the m-OT, the BoMb-OT on arts with (k, m) (10, 10),
(k) = (20, 20). and T 5000. Color palettes are shown under corresponding images.
identical to plans of the m-OT with every choice of (k, m). However, when 0.1, despite 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 given in Figure I7 We would like recall that SAG still need to store full cost matrix (πx) while the BoMb-OT only need to store smaller cost matrices (m x m). We can see that the BoMb-OT gives better than transportation plan than SAG when 2. When 8, the BoMb-OT still seems to be better.
E EXPERIMENT SETTINGS
In this appendix, collect some necessary experimental setups in the paper including generative model, gradient flow. We use POT (Flamary et al. 202I) for solvers, and the PyABC (Klinger al. 2018) for the ABC experiments.
30

---

Under review as a conference paper at ICLR 2022
8:10
#15
440000
Approximated posteriors from ABC with m-OT and BoMb-OT. The first row, the second row, and the last row have 8, 16, and 2, 4, and from left to right.
Figure 11:
32 respectively. In each row, the number of mini-batches &
Ws: 10
Wet 0
Wst O. 10
Wa: 1.1963xx10"
Ws: 0.1976x10
Wy: 10
Wy: -5.1100100
Wy: -
Wy: 0.1209x10-2
Ws: 0.1169x10-2
Ws: 0.1143x10-2
Wy: 0
Wy:
Ws:
We:
eBo.MbOT
steps
steps
Figure 12: Comparison between (e)Boob--T and m-OT gradient flow with 1000, m 16.
The entropic regularized parameter of cBoMb-OT À set to 0.01.
E. DEEP GENERATIVE MODEL.
We use empirical distribution with 1000 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.
Wasserstein-2 scores:
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, num- ber of epochs in MNIST equal 100, number of epochs in CelebA equal to 25, number of epochs in CIFAR1O 10 equal to 50. For d = SW2 we use the number of projections L = 1000 on MNIST and L 100 CIFAR10 and CelebA. For d Wy,
1 for MNIST,
31

---

Under review as a conference paper at ICLR 2022
W;: WA: Wg:
We:
Wy:
Wg: Wg:
Wg:
Wy:
Wy:
Wy:
CT
steps
steps
Figure 13: Visualization of gradient flow provided by m-OT, BoMb-OT and eBoMb-OT with corresponding Wasserstein-2 scores (n 1000, (ke, m) (16, 16)). The entropic regularized parameter eBoMb A is set to 0.01.
Lg
m-W;
3479
L:: 2359
1818
1332
L:: 1329
1144
BoMb-WW
Lg: 3479
Ly: 2406
Ly: 1682
Ly: 1421
Ly: 1225
Lx: 1154
G: 1126
eBooMie-ie W.
step 0
step 60
step 120
step 100
step 240
step 300
step 360
Transforming man face to woman face by using the gradient flow with m-OT and (e)-OT on CelebA dataset for (k:, m) (25, 100). The last column is the nearest image (in of La distance) to the final found female image, the corresponding La distance are also shown on the top middle-stage images.
Figure 14:
50 for CIFARIO, 40 for CelebA. For the eBoMb-OT, we choose the best setting for A E (1,2,3,4, 4,5, 10, 20, 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:
E R
FC ReLU + FC200 - ReLU + FC400 ReLU + FC784 + ReLU
Generator architecture used for CelebA: E R 128 TransposeConv512 BatchNorm + ReLU + TransposeContsaa + BatchNorm ReLU
TransposeConvvis BatchNorm ReLU TransposeConvia BatchNorm
32

---

Under review as a conference paper at ICLR 2022
(1mm)=(10.)
(100)
(10,)-(10.8)
BoMb-OT (n.m)=(10,8)
fuil-OT
07 of samples
Mn
m--TT (10.)=(10.2)
BoMb-OT (1..)
m-OT (1.0)=(10.8)
BoMb-OT (n_m)=(20.88
full-OT
07 matrix samples
M4
Figure 15: Visualization of transportation plans of m-OT and BoMb-OT OT between 2D empirical distributions
with 10 samples. The first two rows are results of k: 20. next two rows are (k, m) (8, 2) and (k, m) = (2, 8) turn.
ReLU - TransposeConvy + Tanh
Metric learning neural network's architecture was used for CelebA:
First part: I E R64x64x3 + Conu6s + LeakyReLUo.2 + Conv12s + BatchNorm LearkyReLUgz - Cont225 BatchNore - Tanh
BatchNorm LearkyReLUo. Cone512
Second part: Convi - Sigmoid
Generator architecture was used for CIFARIO: E 1228 + TransposeConv2ae BatchNorm
ReLU + TransposeCoonuis TransposeConuaa + BatchNorm - ReLU + Transpose + Tanh Metric learning neural network's s architecture was used for CIFAR:
BatchNorm ReLU
First part: I 3223223 Conves + LeakyReLUo.2 + Conw1ss + BatchNorm LearkyReeLop + Conv25e - ButchNoom + Tanh Second part: Cone + Sigmoid
E.2 DEEP DOMAIN ADAPTATION
In this section state the neural network architectures and hyper-parameters for deep domain adaptation.
Evaluation metrie: The classification accuracy is utilized evaluate the mini-hatch methods.
33

---

Under review as a conference paper at ICLR 2022
m-oT In.mi =110,21
eltMb--T 21
m-OT (ummm1100,0)
ettoNb-DT 21
hul-OT
ol samples
m-OT (10
<BoMb-OT (102000000.2.00
m-OT (1m))
eBoMb-OT (0m.0)010..0000
ful-OTT
07 matti
Figure 16: Visualization of the m-OT's transportation plan and eBoMb-OT's transportation plan between 2D empirical distributions provide the result for eBoMb-OT (A E (0.1, 10")), and k: = 40. The last row is the transportation plan with (ks m) = (8,2) 2) and (k, m) = (2, 8) respectively.
samples, and interpolation property eBoMb-OT. four rows
Parameter settings for digits datasets: The number of mini-batches & varies in {8, 16, 32}. For the SVHN dataset, the mini-batch size m set to 50. Because the USPS dataset has fewer samples than the SVHN dataset, is set to 25. We train the network using Adam optimizer with an initial learning rate of 0.0002. The number of epochs 80 for k 8. As the number of mini- batches doubles, double the number of epochs that the number of iterations does not change. The hyperparameters of m- m-OT and BoMb-OT in Equation [0 follow the settings in DeepJDOT: 0.001, = 0.0001. The hyperparameters of m- m-UOT and eBoMb-UOT are the same as JUMBOT: 0.1, X = 6. 0.1, = 1. For computing cBoMb-UOT, we choose the best value of X E (0.01, 0.1, 0.2, 1, 10] (entropic regularizer coefficient).
The number of mini mati-bbtches k varies in {2. 4, 8}. Following the settings in (Fatras et al. (021a)) we train models with mini-batch size 65 during 10000 iterations. The hyperparameters for computing the cost matrix follow the settings in JUMBOT: 0.01, X 0.5, r 0.5, e 0.01. For computing eBoMb-UOT, we choose the best value of A {0,0.01 1, 100} (entropic regularizer coefficient).
Parameter settings for Office-Home dataset:
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
trained from scratch, their learning rates 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 JUMBOT, The learning
to be 10 times that
34

---

Under review as a conference paper at ICLR 2022
SAG (10.)
BOMB-OT (1.
SAG (1.m)-(10.8)
BoMb-OT (n.m))(10,88
Tul-OT
01 of samples
Me
14
TTmmaaix
Visualization of the transportation plan of the BoMb-OT and the stochastic averaged gradient (SAG) method between two 2D empirical distributions with 10 samples. The number of mini-batch for the BoMb-OT is the learning rate for SAG is the best in {0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1,5, 10), the entropic regularization SAG is the best (0.01, 0.05, 0.1, 0.5, 1,5, 5,11. 103.
Figure 17:
rate at iteration is ?p (1+pg))* where is the training progress linearly changing from 0 to 1,
To = 0.01, 10, 0.75.
Neural network architectures: On digits datasets, we use CNN for our generator and FC layer for our classifier in both adaptation scenarios. For Office-Hoone dataset, our generator is ResNet50 pre-trained on ImageNet except for the last FC layer, which is our classifier.
Generator architecture was used for SVHN dataset: E 32x32x3 - Conuzz BatchNorm + ReLU + Conasz + BatchNorm - ReLU 4
Max + Convs4 + BatchN orm + ReLU - Conoss BatchNorv + ReLU
PoolDD + Conv12s + BatchNorm + ReLU + Convi2s + BatchNorm
ReLU + MoxPool2D + Sigmoid + FC128
Generator architecture used for USPS dataset: E 228x28x3 - Conuzz + BatchNorm + ReLU + MoxPool2D - Conuss
BatchNorm ReLU Conv12s BatchNorm
ReLU
Sigmoid + FC
Classifier archiecture was used for both SVHN and USPS datasets: E IR
FC10
E.3 GRADIENT FLOW
For the implementation of the gradient flow, we use the geomloss library (Feyyy et al. 2019). The learning rate is set to 0.001. For the autoencoder in CelebA experiments, we use the rеpo in http:://ww...aaooooooooo//oooooa
com
http:///tobb/o////////////////// .ipelebb..pynbb)) ipynb for the pre-
trained autoencoder.
E.4 COMPUTATIONAL INFRASTRUCTURE
All deep learning experiments are done RTX 2080 Ti GPU and a GTX 1080 Ti. Other experiments are done on a MacBook Pro line MI.
35


Reviews:
[{"review_id": "iKji3-TtxuR", "reviewer": "Reviewer_ScyT", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper describes correctly their new loss function as well as some weaknesses from the original minibatch formulation. They provide small toy examples to see the differences between the proposed and the original losses. They also discuss the computational complexity of their method and discuss some practical aspects of their contributions. \n\nEvaluation: My feeling is that the theoretical and experimental aspects are not studied enough. \n\nFor instance on the theoretical side:\n\n(i) While discussing the connections between samples from the transport plan, the authors do not discuss the transported mass. \n\n(ii) Also a study similar to [1] about the possible bias of stochastic gradient is important. Would minimizing your proposed loss function with SGD lead to the correct minimum ? \n\n\nOn the experimental side, there are many experiments, but some of them are not extensive (generative modelling and domain adaptation) which make the performances of their method unclear. I am also skeptical that the method is really competitive in practice (numerically and memory) with state of the art method due to the fact that it needs a bigger k for the loss to be computed.\n\n(i) I would encourage the authors to compelete their work on generative modelling and domain adaptation with comparison with recent methods, including methods which are also not based on optimal transport. \n\n(ii) On domain adaptation and generative modelling, I am not convinced that the proposed method outperforms other methods in practice due to its overall complexity (k>1). Other methods rely on a k equals to 1 (including deepjdot) which might make them faster to train and as efficient or more.\n\n(iii) On domain adaptation, there are more recent methods and I would like to see the comparison [2,3]. There is also this paper [4] that you cite in your supplementary and that you should compare to as this paper also tries to minimize the impact of undesirable connections from minibatches.\n\nRelated Work and Discussion: The related work is complete and well discussed. The bibliography style seems respected.\n\nClarity: The paper is clearly written except for the experimental part. A reader which do not know about generative modelling, domain adaptation, color transfer or bayesian computation can not understand the purpose of the different experiments as the explanations are in the supplementary materials. This gives the feeling that the authors just stacked experiments without discussing them properly and entirely. I think this is mostly due to the high number of different experiments. I suggest to focus on two different experiments, to develop the original problem, their methods to solve them and, finally to produce extensive experiments of different methods on several datasets. The remaining experiments could be added to the supplementary.\n\nQuestions and remarks:\n1. Does your loss transport all samples ? I think it does, but it should be proven. An overall theoretical study of the transport plan is lacking. \n2. Is your method an upper bound of OT and a lower bound of minibatch OT ?\n3. Time comparison between your method and original minibatch OT. This could be done on domain adaptation experiments for instance. \n4. Name of tables are unclear. Please add the meaning of the scores you five (Table 1 and 2).\n\n[1] Learning with minibatch Wasserstein: asymptotic and gradient properties, Fatras et al.\n[2] Sliced Wasserstein Discrepancy for Unsupervised Domain Adaptation, Lee et al.\n[3] Reliable Weighted Optimal Transport for Unsupervised Domain Adaptation, Xu et al.\n[4] Unbalanced minibatch Optimal Transport; applications to Domain Adaptation, Fatras et al.\n\n\n------------------------------\nRebuttal\n------------------------------\n\nThank you for your detailed answers to my questions. I have read the other reviews and the different answers. I still think that the idea is appealing, but the way the paper is written and the fact that the experiments are not extensive enough prevent a publication at that time in my opinion. I keep my score unchanged and I would encourage the authors to pursue the experimental evaluations or theoretical evaluations (distance of the plan marginals to the marginals ?) of their method.", "summary_of_the_paper": "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.", "summary_of_the_review": "The proposed formulation is appealing and I like the idea. However, in my opinion, some theoretical and experimental aspects are lacking, and that prevent a publication at the moment (see the different points above). There is a big discussion regarding where the method can applied and it could be replaced by a discussion on theoretical aspects which are not mentionned. Regarding experiments, I encourage the authors to focus on two experiments (Generative modelling and domaine adaptation) and to make extensive experiments with recent baselines. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "X4ARTWur5By", "reviewer": "Reviewer_9B4c", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper introduces BombOT, a hierarchical approach to combining mini-batches in order to enhance the approximation quality of mini-batch-based approaches to OT.\n\nI believe it is a nice contribution and the main strength is the comprehensive list of experiments performed, illustrating benefits on several tasks in statistics and machine learning.\n\nI have a few concerns, though\n\n1)The theoretical findings are not very strong. Showing that there is a distance being defined is a nice finding, and contrasting this with the lack of a metric structure in usual m-OT suggest that this construction is in the right direction. However, practitioners are more interested in convergence properties. In this regard, theorem 2 doesn't say much, as it corresponds to a comparison between batch and population versions of the proposed method. This doesn't say much about how does this relate to the original wasserstein/sliced/etc distance.\n\n2)Similarly, experiments mainly comparse m-OT and bomb-OT. Comparisons with true Wasserstein distance are mostly lacking. We should be able to at least empirically how large batches need to be so that we will get decent approximations. As a practitioner who tried used m-OT I believe the problem is super important, but the numbers provided don't say much if they are not put into perspective with respect to the true distance. The findings reporting in Fig 3 are honest, but a bit discouraging in that respect. What is the dimension here? how large need k and m need to be so we get a reasonable approximation? (related to (1)).  \n\n3)another suggestion (optional): I believe this method can be useful for aligning spaces, as in https://arxiv.org/abs/1809.00013. It would be good to have comparisons there as well, and to show whether this mini-batch approximation can be reasonable (i.e. compare against the no-mini batch case)", "summary_of_the_paper": "An enhancement over a mini-batch version of OT that provides better empirical evaluation over a large array of tasks", "summary_of_the_review": "Nice article. But requires a more thorough comparison with true wasserstein.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "TNKiQwur6bC", "reviewer": "Reviewer_C31a", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "strengths:\n\nThe optimal transport (OT) problem studied in this paper is important. Existing work and their advantages and disadvantages are discussed and the study is well motivated. \n\nTheorem 2 shows that the proposed BoMb-OT method is sound and approximating some well-defined metrics. On the other hand, the authors provided the implementation and demonstrated the usefulness of the proposed method using several different applications.\n\nweakness:\n\n1. As mentioned, the population m-OT has some problem of not preserving metric properties. As also noted the entropic regularized version of population BoMb-OT would recover m-OT in some cases. Does this mean the BoMb-OT will also recover the same problem that m-OT would have?\n\n2. Figure 1 is confusing to me. I get the authors wanted to provide some intuitive explanations for the advantages of using BoMb-OT. However, I do not understand why and how using BoMb-OT can achieve the right transportation in the figure while using m-OT cannot. I suggest if possible these examples can come with specific numbers such that we can calculate and verify the results, making them easier to understand and more convincing.\n\n3. The BoMb-OT has one additional OT comparing to m-OT as noted in \"Computational complexity of the BoMb-OT\" paragraph ($O(k^2(m^2+1)/\\epsilon^2)$ vs. $O(k^2 m^2 /\\epsilon^2)$). It looks contradiction to what is mentioned in the appendix, i.e., \"The run time of the m-OT nearly doubles that of the BoMb-OT\". Is this from a better implementation of BoMb-OT. Please clarify this.", "summary_of_the_paper": "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.\n\nThe 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.\n\nThe 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.\n\n\n", "summary_of_the_review": "Overall, I found the problems studied in this paper interesting and important, the proposed method reasonable, theoretically sound, and well supported by experiments and implementation details.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper propose a way to make minibatch Optimal transport (m-OT) more efficient by computing an optimal assignment (in the OT sens) and us this assignment to compute instead a hierarchical OT loss (bomb-OT ) that can be used instead of the m-OT loss. The authors discuss how the equivalent OT plan with bomb-OT is much more sparse, and how the proposed approach is actually not biased when the number of mini-batches $k\rightarrow \infty$ . Numerical experiments  show that the proposed method allows a gain in performances in applications such as generative modeling, domain adaptation, color transfer and approximate Bayesian computation.

The paper originally got borderline-negative scores from the reviewers. While the reviewers acknowledged that the idea is interesting, they had some concerns about the theoretical results strength, some missing baselines and discussions in the numerical experiments. The authors did a detailed reply that clarified some problems. the new numerical experiments with m-UOT were also greatly appreciated by the reviewers but they also raised some questions about the paper. Some concerns detailed below about the comparison with m-OT appeared during the reviewers discussion. Despite the new information,  the reviewers reached an agreement that this paper is interesting but needs more work and another round of reviews before acceptance. For theses reasons  the AC recommends a rejection for this paper.

More details and suggestions below:

- While it is clearly not the objective of the paper a discussion about the proximity of the average plan to the exact OT plan is interested. Also a short numerical experiments showing that the bomb-OT average plan is closer to the exact plan than m-OT would be a good illustration of the better performance of bomb-OT. This seems more important for the paper than the color transfer experiments that is kind of a toy problem.

- After checking the definition in the paper and discussion between reviewers it appeared that the comparison with m-OT is a bit unfair due to the reformulation of the problem in (1). indeed in the usual formulation, k pairs of independent  minibatches are used and the OT is done on those pairs (a sum of k OT) not on all the possible pairwise permutation as in definition of m-OT in equation (1). In other  words in m-OT the batches are supposed to be independent which is not the case  in the proposed formulation (it is equivalent in the population case though).  It means that in practical application, for the same computational complexity  (k^2 OT computed), m-OT actually uses $k^2m$ independent samples on each distribution  whereas the bomb-OD (and the m-OT defined in equation (1) ) use $km$ samples . By implementing m-OT as  in (1) they actually prevent m-OT to explore the dataset as its original  formulation does. This means that all the experiments should be done either with the original m-OT implementation of both the original and (1) in addition to bomb-OT. The proposed method will proably work better but the current experiment do not allow this fair comparison. 

- The theoretical result need more discussion and justification.  For instance  m-OT converges to its population value in 
 $O(m^{1/2}n^{-1/2}+k^{-1/2})$ that is independent from the dimensionality  $d$, but the authors prove the  concentration of bomb-OT in  of $O(m^{1/2}n^{-1/d})$  which is  clearly a problem for large $d$. Also the dependence on  $k$ of the convergence would be important since  bomb-OT is well defined is true only in the population case where $k$ is  large. Note that the claim that it is well defined and hence better is also a  bit dubious because it is well defined for $k=\infty$, which is also the case for m-OT when $m=\infty$. Both $m$ and $k$ large will lead to not practical optimization problems so they are comparable except that m-OT converged to the true OT plan when $m\rightarrow \infty$ which is not the case for bomb-OT.

- While the contribution of the paper in indeed a methodological method and does not require to be state of the art on all applications the numerical experiments should be improved. First as discussed above the comparison with m-OT is actually unfair an do not correspond to what in done in practice (where all mini batches are independent). m-OT should be implemented with  $k^2$ truly independent minibatches.

- Second , the authors use approximate W2 on two of  the GAN dataset and FID on the third. This is  problem because approximate W2 is not defined in the paper. FID is the standard performance measure and should be used for all dataset. 

- Third the novel experiments comparing also raises a lot of questions. m-UOT is far better than BoMb-OT suggesting that Unbalanced OT can compensate for the limits of m-OT far better than bomb-OT itself. Yes there is a slight increase in performance  for ebomb-UOT over m-UOT but is is so small (0.08 %) that it is hard to find them significant, especially since we have no variance. This result that is provided only for DA application actually  suggest that the competitor of bomb-OT is m-UOT and not m-OT so it should also be part of the comparison in the other experiments. The authors talk in their replay about the limits of m-UOT but stating that the experiences are not done in the paper is not an excuse for evaluating this clear competitor on other problems and showing numerically these limits.

- Finally in the current version of the paper puts a lot of things in the annex that make the paper clearly not self content. Some experiments could go in annex/supp for instance the color transfer to make place for more details in the main paper. 

Note that it is not one of those comments above that lead the the reject decision but the sum of them that clearly show that the paper needs more work.
Decision: Reject