O N T RANSPORTATION OF M INI - BATCHES : A H IERARCHICAL A PPROACH
Anonymous authors
Paper under double-blind review
A BSTRACT
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 satisﬁed.
To address these problems, we propose a novel mini-batching scheme for optimal transport, named Batch of Mini-batches Optimal Transport (BoMb-OT), that ﬁnds the optimal coupling between mini-batches and it can be seen as an approximation to a well-deﬁned 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 inﬁnity.
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 ﬂow, the BoMb-OT also yields either a lower quantitative result or a better qualitative result than the m-OT.
I NTRODUCTION
Optimal transport (OT) (Villani, 2021; 2008; Peyré et al., 2019) has emerged as an efﬁcient tool in dealing with problems involving probability measures.
Under the name of Wasserstein distance,
OT has been widely utilized to solve problems such as generative modeling (Arjovsky et al., 2017;
Tolstikhin et al., 2018; Salimans et al., 2018; Genevay et al., 2018; Liutkus et al., 2019), barycenter problem (Ho et al., 2017; Li et al., 2020), and approximate Bayesian computation (Bernton et al., 2019a; Nadjahi et al., 2020).
Furthermore, OT can also provide the most economical map of moving masses between probability measures, which is very useful in various tasks such as color transfer (Ferradans et al., 2014; Perrot et al., 2016), natural language procesing (Alvarez-Melis & Jaakkola, 2018), and domain adaptation (alignment) (Courty et al., 2016; Lee et al., 2019a; Xu et al., 2020), and graph processing (Titouan et al., 2019; Xu et al., 2019; Chen et al., 2020).
Although OT has attracted growing attention in recent years, a major barrier that prevents OT from being ubiquitous is its heavy computational cost.
When the two probability measures are discrete with n supports, solving the Wasserstein distance via the interior point methods has the complexity of O ( n 3 log n ) (Pele & Werman, 2009), which is extremely expensive when n is large.
There are two main lines of works that focus on easing this computational burden.
The ﬁrst approach is to ﬁnd a good enough approximation of the solution by adding an entropic regularized term on the objective function (Cuturi, 2013).
Several works (Altschuler et al., 2017; Lin et al., 2019) show that the entropic approach can produce a ε -approximated solution at the same time reduces the computational complexity to O ( n 2 /ε 2 ) .
The second line of works named the “slicing" approach is based on the closed-form of Wasserstein distance in one-dimensional space, which has the computational complexity of order O ( n log n ) .
There are various variants in this directions; i.e., (Bonneel et al., 2015; Deshpande et al., 2019; Kolouri et al., 2019; Nguyen et al., 2021a;b), these all belong to the family of sliced Wasserstein distances.
Recently, some methods are proposed to combine the dimensional reduction approach with the entropic solver of Wasserstein distance to inherit advantages from both directions (Paty & Cuturi, 2019; Muzellec & Cuturi, 2019; Lin et al., 2020a;b).
Although those works have reduced the computational cost of OT considerably, computing OT is nearly impossible in the big data settings where n could be as large as few millions.
In particular, solving OT requires to compute and store a n × n cost matrix that is impractical with current computational devices.
Especially in deep learning applications, both supports of empirical measures and the cost matrix must be stored in a same device (e.g. a GPU) for automatic differentiation.
This problem exists in all variants of OT including Wasserstein distance, entropic Wasserstein, and sliced Wasserstein distance.
Therefore, it leads to the development of the mini-batches method for OT (Genevay et al., 2018; Sommerfeld et al., 2019), which we refer to as mini-batch OT loss.
The main idea of the mini-batch method is to divide the original samples into multiple subsets (mini-batches), in the hope that each pair of subsets (mini-batches) could capture some structures of the two probability measures, meanwhile, the computing OT cost between two mini-batches is cheap due to a very small size of mini-batches.
Then the overall loss is deﬁned as the average of distances between pairs of mini-batches.
This scheme was applied for many forms of Wasserstein distances (Deshpande et al., 2018; Bhushan Damodaran et al., 2018; Kolouri et al., 2018; Salimans et al., 2018), and was theoretically studied in the works of (Bellemare et al., 2017; Bernton et al., 2019b; Nadjahi et al., 2019).
Recently, Fatras et al.
(Fatras et al., 2020; 2021b) formulated this approach by giving a formal deﬁnition of the mini-batch OT loss, studying its asymptotic behavior, and investigating its gradient estimation properties.
Despite being applied successfully, the current mini-batch OT loss does not consider the relationship between mini-batches and treats every pair of mini-batches the same.
This causes undesirable effects in measuring the discrepancy between probability measures.
First, the m-OT loss is shown to be an approximation of a discrepancy (the population m-OT) that does not preserve the metricity property, namely, this discrepancy is always positive even when two probability measures are identical.
Second, it is also unclear whether this discrepancy achieves the minimum value when the two probability measures are the same.
That naturally raises the question if we could propose a better mini-batching scheme to sort out these issues in order to improve the performance of the OT in practical applications.
Contribution:
In this paper, we propose a novel mini-batching scheme for optimal transport, which is named as Batch of Mini-batches Optimal Transport (BoMb-OT).
In particular, the BoMb-OT views every mini-batch as a point in the product space, then a set of mini-batches could be considered as an empirical measure.
We now could employ the Kantorovich formulation between these two empirical measures in the product space as a discrepancy between two sets of mini-batches.
In summary, our main contributions are two-fold:
First, the BoMb-OT could provide a more similar transportation plan to the original OT than the m-OT, which leads to a more meaningful discrepancy using mini-batches.
In particular, we prove that the BoMb-OT approximates a well-deﬁned 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.
Speciﬁcally, when the regularization parameter in the entropic population BoMb-OT goes to inﬁnity, its value approaches the value of the population m-OT. 2.
Second, we present the implementation strategy of the BoMb-OT and detailed algorithms in various applications in Appendix C. We then demonstrate the favorable performance of the BoMb-OT over the m-OT in two main applications that using optimal transport losses, namely, deep generative models and deep domain adaptation.
Moreover, we also compare BoMb-OT to m-OT in other applications, such as sample matching, approximate Bayesian computation, color transfer, and gradient ﬂow.
In all applications, we also provide a careful investigation of the effects of two hyper-parameters of the mini-batching scheme, which are the number of mini-batches and the size of mini-batches, on the performance of the BoMb-OT and the m-OT.
Organization: The remainder of the paper is organized as follows.
In Section 2, we provide backgrounds for optimal transport distances and the conventional mini-batching scheme (m-OT).
In Section 3, we deﬁne the new mini-batching scheme for optimal transport distances, Batch of mini-batches Optimal Transport, and derive some of its theoretical properties.
Section 4 benchmarks the proposed mini-batch scheme by extensive experiments on large-scale datasets, and followed by discussions in Section 5.
Finally, proofs of key results and extra materials are in the supplementary.
Notation:
For any probability measure µ on the Polish measurable space ( X , Σ) , we denote ⊗ m µ ( m ≥ 2) as the product measure on the product measurable space ( X m , Σ m ) .
For any p ≥ 1 , we deﬁne P p ( R N ) as the set of Borel probability measures with ﬁnite p -th moment deﬁned on a given metric space ( R N , (cid:107) .
(cid:107) ) .
To simplify the presentation, we abuse the notation by using both the notation X m for both the random vector ( x 1 , . . .
, x m ) ∈
X m and the set { x 1 , . . .
, x m } , and we deﬁne by P X m := 1 m (cid:80) mi =
δ x i the empirical measure (the mini-batch measure) associated with X m .
For any set X n := { x 1 , . . .
, x n } and m ≥ 1 , we denote by [ X n ] m the product set of X n taken m times and (cid:0) X n m (cid:1) the set of all m -element subsets of X n .
B ACKGROUND ON MINI - BATCH OPTIMAL TRANSPORT
In this section, we ﬁrst review the deﬁnitions of the Wasserstein distance, the entropic Wasserstein, and the sliced Wasserstein.
We then review the deﬁnition of the mini-batch optimal transport (m-OT).
2. W ASSERSTEIN DISTANCE AND ITS VARIANTS
We ﬁrst start with the deﬁnition of Wasserstein distance and its variants.
Let µ and ν be two probability measures on P p ( R N ) .
The Wasserstein p -distance between µ and ν is deﬁned as follows:
[Equation]
is the set of transportation plans between µ and ν .
The entropic regularized Wasserstein to approximate the OT solution (Altschuler et al., 2017;
Lin et al., 2019) between µ and ν is deﬁned as follows (Cuturi, 2013): W τp ( µ, ν ) :
= min π ∈ Π( µ,ν ) (cid:8)(cid:2) E π ( x,y ) (cid:107) x − y (cid:107) p (cid:3) 1 p + τ KL ( π | µ ⊗ ν ) (cid:9) , where τ > 0 is a chosen regularized parameter and KL denotes the Kullback-Leibler divergence.
Finally, the sliced Wasserstein (SW) (Bonnotte, 2013; Bonneel et al., 2015) is motivated by the closed-form of the Wasserstein distance in one-dimensional space.
The formal deﬁnition of SW is:
SW p ( µ, ν ) := (cid:2) E θ ∼U ( S N − 1 ) W pp ( θ(cid:93)µ, θ(cid:93)ν ) (cid:3) 1 p , where U ( S N − 1 ) denotes the uniform measure over the ( N − 1) -dimensional unit hypersphere and θ(cid:93) is the orthogonal projection operator on direction θ .
2. M INI - BATCH O PTIMAL T RANSPORT
In this section, we ﬁrst discuss the memory issue of large-scale optimal transport and challenges of dual solver.
Then, we revisit the mini-batch OT loss that has been used in training deep generative models, domain adaptations, color transfer, and approximate Bayesian computation (Bhushan Damodaran et al., 2018; Genevay et al., 2018;
Tolstikhin et al., 2018; Fatras et al., 2020; Bernton et al., 2019a).
To ease the presentation, we are given X n := (cid:8) x 1 , . . .
, x n (cid:9) , Y n := (cid:8) y 1 , . . .
, y n (cid:9) i.i.d.
samples from µ and ν in turn.
Let µ n := 1 n (cid:80) ni
=
δ x i and ν n := 1 n (cid:80) ni =
δ
y
i be two corresponding empirical measures from the whole data set.
Here, n is usually large (e.g., millions) and each support in X n ,
Y n can be a high dimensional data point (e.g. a high resolution image, video, etc).
Memory issue of optimal transport: Using an OT loss between µ n and ν n needs to compute and store a n × n cost matrix which has one trillion ﬂoat entries (about 4 terabytes) when n is about millions.
Moreover, when dealing with deep neural networks, both support points and the cost matrix are required to be stored in the same memory (e.g., a GPU with 8 gigabytes memory) as a part of the computational graph for automatic differentiation.
This issue applies to all variants of OT losses, such as Wasserstein distance, entropic Wasserstein, sliced Wasserstein distance.
Therefore, it is nearly impossible to compute OT and its variants in large-scale applications.
Challenges of stochastic dual solver : Using stochastic optimization to solve the Kantorovich dual form is a possible approach to deal with large-scale OT, i.e. Wasserstein GAN (Arjovsky et al., 2017; Leygonie et al., 2019).
However, the obtained distance has been shown to be very different from the original Wasserstein distance (Mallasto et al., 2019;
Stanczuk et al., 2021).
Using input convex neural networks is another choice to approximate the Brenier potential (Makkuva et al., 2020).
Nevertheless, recent work (Korotin et al., 2021) has indicated that input convex neural networks are not sufﬁcient (have limited power) in approximating the Brenier potential.
Furthermore, both mentioned approaches are restricted in the choice of ground metric.
In particular, L 1 norm is used in
Wasserstein GAN to make the constraint of dual form into the Lipchitz constraint and L 2 norm is for the existence of the Brenier potential.
Mini-batch solution: As a popular alternative approach for stable large-scale OT with ﬂexible choices of the ground metric, the mini-batch optimal transport is proposed (Genevay et al., 2018;
Sommerfeld et al., 2019)
and has been widely used in various applications (Arjovsky et al., 2017;
Deshpande et al., 2018; Sommerfeld et al., 2019; Bhushan Damodaran et al., 2018).
In this approach, the original n samples are divided into subsets (mini-batches) of size m , where m is often the largest number that the computer can process, then an alternative solution of the original OT problem is formed by aggregating these smaller OT solutions from mini-batches.
We now state an adapted deﬁnition of mini-batch OT (m-OT) scheme that was formulated in (Fatras et al., 2020), including its two key parts: its transportation cost and its transportation plan.
Deﬁnition 1 (Empirical m-OT) .
For p ≥ 1 and integers m ≥ 1 , and k ≥ 1 , let d : P p ( X ) ×
P p ( X ) → [0 , ∞ ) be a function, i.e., { W p , W τp , SW p } .
Then, the mini-batch OT (m-OT) loss and the transportation plan, a n × n matrix, between µ n and ν n are deﬁned as follows:
[Equation]
where X mi is sampled i.i.d from (cid:0) X n m (cid:1) , Y mi is sampled i.i.d from (cid:0) Y n m (cid:1) , and the transport plan π X mi ,Y mj , where its entries equal zero except those indexed of samples X mi × Y mj , is the transportation plan when d ( P X mi , P Y mj ) is an optimal transport metric.
The above deﬁnition was generalized to the two original measures as follows in (Fatras et al., 2020):
Deﬁnition 2 (Population m-OT) .
Assume that µ and ν are two probability measures on P p ( X ) for given positive integers p ≥ 1 , m ≥ 1 , and d : P p ( X ) × P p ( X ) → [0 , ∞ ) be a given function.
Then, the population mini-batch OT (m-OT) discrepancy between µ and ν is deﬁned as follows:
[Equation]
Issues of the m-OT: From Deﬁnition 1, the m-OT treats every pair of mini-batches the same by taking the average of the m-OT loss between any pair of them for both transportation loss and transportation plan.
This treatment has some issues.
First, a mini-batch is a sparse representation of the true distribution and two sparse representations of the same distribution could be very different from each other.
Hence, a mini-batch from X m would prefer to match to certain mini-batches of Y m , rather than treating every mini-batch of Y m equally.
For example, each mini-batch has one datum, then each term in the population m-OT now is the ground metric of the OT cost, the population m-OT degenerates to E [ d ( X m , Y m )] for independent X m and Y m .
This treatment also leads to an uninformative transportation plan shown in Figure 15, which is followed by a less meaningful transportation cost.
Second, although it has been proved that the population m-OT is symmetric and positive (Fatras et al., 2020), for the same reason it does not vanish when two measures are identical.
B ATCH OF M INI - BATCHES O PTIMAL T RANSPORT
To address the issues of m-OT, in this section, we propose a novel mini-batch scheme, named batch of mini-batches OT (BoMb-OT) .
We ﬁrst demonstrate the improvement of the BoMb-OT to the m-OT and discuss the practical usage of the BoMb-OT.
Then, we prove that the BoMb-OT loss is an approximation of a well-deﬁned metric in the space of Borel probability measures.
Finally, we discuss the entropic version of the population BoMb-OT (population eBoMb-OT) admits the population m-OT as a special case where the entropic regularization goes to inﬁnity in Appendix B. To simplify the presentation, we assume throughout this section that µ and ν are continuous probability measures in P p ( X ) for some given p ≥ 1 and X ⊂ R N where N ≥ 1 .
Furthermore, d : P p ( X ) × P p ( X ) → [0 , ∞ ) is a given divergence between probability measures in P p ( X ) .
3. D EFINITION OF B O M B -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
[Figure]
[Figure]
A A
B B
B B
A A
A A
B
B
[Figure]
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 n would prefer to match some certain mini-batches from Y n .
To address the limitation, we enhance the second step by considering the optimal coupling between mini-batches.
By doing so, we could deﬁne a new mini-batch scheme that is able to construct a good global mapping between samples that also leads to a meaningful objective loss.
Let X m 1 , X m 2 , . . .
, X mk be mini-batches that are sampled with replacement from [ X n ] m and
Y m 1 , Y m 2 , . .
, Y mk be mini-batches that are sampled with replacement from [ Y n ] m .
The batch of mini-batches optimal transport is deﬁned as follows:
Deﬁnition 3 (Batch of mini-batches) .
Assume that p ≥ 1 , m ≥ 1 , k ≥ 1 are positive integers and let d : P p ( X ) × P p ( X ) → [0 , ∞ ) be a given function (e.g., { W p , W τp , SW p } ).
The BoMb-OT loss and the BoMb-OT transportation plan between probability measures µ n and ν n are given by:
[Equation]
i = k (cid:88) j = γ ij d ( P X mi , P Y mj ); (cid:98) π mk ( µ n , ν n ) :=
k (cid:88) i =
k (cid:88)
j = ¯ γ ij π X mi ,Y mj , where ⊗ m µ k := 1 k (cid:80) ki =1 δ X mi and ⊗ m ν k := 1 k (cid:80) kj =1 δ Y mj are two empirical measures deﬁned on the product space via mini-batches (measures over mini-batches), ¯ γ is a k × k optimal transport plan between ⊗ m µ k and ⊗ m ν k , and π X mi ,Y mj is deﬁned as in Deﬁnition 1.
Compared to the m-OT, the BoMb-OT considers an additional OT between two measures on minibatches for combining mini-batch losses.
The improvement of the BoMb-OT over the m-OT is illustrated in Figure 1 in which we assume that we can only solve 2 × 2 OT problems and we need to deal with the OT problems with two empirical measures of 3 supports.
The optimal transportation plan is to map i -th square to i -th circle.
Assume that there are four mini-batches: A1 and A2 from X and B1 and B2 from Y .
The weight for the m-OT loss between A1 and B2 equals 14 , meanwhile, the map from A1 to B2 is unuseful.
In this toy example, A1 should be mapped to B1 and A2 is for B2.
It is also the solution that the BoMb-OT tries to obtain by re-weighting pairs of mini-batches through a transportation plan between two measures over mini-batches.
Practical usage of the BoMb-OT: There are three types of applications that the BoMb-OT can be utilized.
The ﬁrst one is gradient-based applications (deep learning) such as deep generative models, deep domain adaption, and gradient ﬂow.
In these applications, the goal is to estimate the gradient of a parameter of interest (a neural network) with an OT objective function.
The second one is mapping-based applications, namely, we aim to obtain a transportation map between samples.
Examples of mapping-based applications include color transfer, domain adaptation, and sample matching.
Finally, the third type of applications is discrepancy-based applications where we need to know the discrepancy between two measures, e.g., approximate Bayesian computation, and searching problems.
We now discuss the implementation of BoMb-OT in each scenario.
Gradient-based (deep learning) applications : For a better presentation, we assume that a system with a GPU (graphics processing unit) is used in the training process.
We further assume that the GPU has enough space to store the neural net, the 2 × m supports of two mini-batch measures, and the m × m cost matrix.
The algorithm of the BoMb-OT consists of three steps: Forwarding , Solving k × k OT , and Re-forwarding and Backwarding .
In the ﬁrst step, k 2 OT problems of size m × m are computed in turn on GPU to obtain the k × k cost matrix which indicates the transportation cost between each pair of mini-batch measures.
We recall that it is only enough memory for computing m × m OT on GPU at a time.
After that, an OT problem with the found k × k cost matrix is solved to ﬁnd the coupling between mini-batches.
Here, k can be much larger than m since the k × k OT can be solved on computer memory (RAM - Random-access memory), which is bigger than GPU’s memory.
We would like to emphasize that in the ﬁrst two steps, automatic differentiation is not required.
In the last step, we do re-forwarding and do backwarding (back-propagation) to obtain the gradients of each pair of mini-batches, then using the coupling in the second step to aggregate them into the ﬁnal gradient signal of the parameter of interest.
The detailed algorithms of the BoMb-OT for deep generative model and deep domain adaption are respectively described in Algorithm 1 in
Appendix C.1 and Algorithm 2 in Appendix C.2.
Compared to the m-OT, the BoMb-OT only costs an additional forwarding and an additional k × k OT problem that is not expensive since they do not require to create and store any computational graphs.
Mapping-based applications : In this type of application, the BoMb-OT also consists of three steps like in gradient-based applications.
However, in the last step, the transportation plan (the mapping) is aggregated instead of the gradient.
As an example, we present the BoMb-OT algorithm for color transfer in Algorithm 3 in Appendix C.3.
We would like to emphasize that the last step (reforwarding) can be removed by storing all k 2 mini-batch transportation plans of size m × m which can be represented by a sparse matrix of size n × n for saving memory.
Discrepancy-based applications
: The BoMb-OT needs only two steps in this scenario, namely, forwarding and solving k × k OT.
Since we do not need to re-estimate any statistics from minibatches, the ﬁnal value of the BoMb-OT can be evaluated at the end of the second step without the re-forwarding step.
Similar to the previous two types of applications, we present the BoMb-OT algorithm for approximate Bayesian computation in Algorithm 5 in Appendix C.4.
Compared to the m-OT, the BoMb-OT needs only an additional k × k OT here.
Choosing k and m : In practice, we want to have as large as possible values of k and m .
The mini-batch size m is often chosen to be the largest value of the computational memory that can be stored.
For choosing k , there are two practical cases.
Two (multiples) computational memoriesthis is the case of modern computational systems that have at least a CPU (central processing unit) with RAM and a GPU with its corresponding memory.
As discussed in the practical usage of the
BoMb-OT, OT problems between mini-batch measures of size m × m are solved on the GPU for high computational speed in estimating the parameter of interest or other statistics.
On the other hand, the k × k OT problems between measures over mini-batches can be computed on the CPU’s memory.
So, k can be chosen larger than m .
One (shared) computational memory - in this case, the largest value of k equals to m .
Note that, smaller values of k and m can still be used for faster computation.
We would like to recall that k is usually set to 1 in almost all recent applications of OT (Genevay et al., 2018; Bhushan Damodaran et al., 2018;
Tolstikhin et al., 2018; Bernton et al., 2019a).
In our experiments, we will demonstrate that increasing k could improve the performance.
Computational complexity of the BoMb-OT: The computational complexity of BoMb-OT depends on a particular choice of divergence d .
In the experiments, we speciﬁcally consider three choices of d ∈ { W p , W τp , SW p } where τ > 0 is a chosen regularized parameter.
Here, we only discuss the case when d = W p , the discussion of other choices of d is in Appendix B. When d is the entropic
Wasserstein distance, for each pair X mi and Y mj , the computational complexity of approximating d ( P X mi , P Y mj ) via the Sinkhorn algorithm is of order O ( m 2 /ε 2 ) where ε > 0 is some desired accuracy (Lin et al., 2019).
Hence, the computation cost of k 2 OT transport plans from k 2 pairs of mini-batches is of order O ( k 2 m 2 /ε 2 ) when using entropic regularization (see Appendix B for the deﬁnition).
With another OT cost within k 2 pairs, the total computational complexity of computing the BoMb-OT is O
( k 2 ( m 2 + 1) /ε 2 ) .
It means that the computational complexity of computing the
BoMb-OT is comparable to the m-OT, which is O ( k 2 m 2 /ε 2 ) .
The BoMb-OT’s transportation plan: We discuss in detail the sparsity of BoMb-OT’s transportation plan in Appendix B.
To visualize the BoMb-OT’s plan and compare it with the m-OT’s plan, we carry out a simple experiment on two measures of 10 supports in Figure 15.
The details of the experiment and discussion are given in Appendix D.6.
From this example, the BoMb-OT provides a more accurate transportation plan than the m-OT with various choices of m and k .
3. M ETRIC APPROXIMATION OF THE B O M B -OT LOSS
In this section, we show that the BoMb-OT is an approximation of a well-deﬁned metric, named population BoMb-OT , in the space of probability measures.
To ease the ensuing discussion, we ﬁrst deﬁne that metric as follows:
Deﬁnition 4 (Population BoMb-OT) .
Let µ and ν be two probability measures on P p ( X ) , for p ≥is a positive number.
The population BoMb-OT between probability measures µ and ν is deﬁned as:
[Equation]
Different from the m-OT in Deﬁnition 2, the optimal plan between two distributions on product spaces is crucial to guarantee that the population BoMb-OT is a well-deﬁned 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 X m and
the function d
( P X m , P Y m ) is continuous in terms of X m , Y m ∈
X m .
Then, the population BoMb-OT is a well-deﬁned metric in the space of probability measures.
The proof of Theorem 1 is in Appendix A.
The assumption of Theorem 1 is mild and satisﬁed when d is (sliced) Wasserstein metrics of order p , i.e., d ∈ { W p , SW p } .
Our next result shows the approximation error between the BoMb-OT loss and the population BoMbOT distance.
We provide the following result with the approximation error when d ∈ { W p , SW p } .
Theorem 2 (Population BoMb-OT) .
Assume that p ≥ 1 , X is a compact subset of R N , and all the possible mini-batches are considered.
Then, we have (i) (cid:12)(cid:12) (cid:98) D k,md [Equation]
O P (cid:0) m 1 − 1 p /n 1 N (cid:1) when d ≡ W p (ii) (cid:12)(cid:12) (cid:98) D k, m
d ( µ n , ν n ) − D md ( µ, ν ) (cid:12)(cid:12) =
O P (cid:0) m 1 − 1 p /n 12 (cid:1)
when d ≡
SW p and k is the number of mini-batches.
The proof of Theorem 2 is in Appendix A. We would like to remark that the dependency of the sample complexity of the BoMb-OT on N is necessary when d ≡ W p .
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 p .
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 is O ( k 2 m log m ) when d ≡ SW p .
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 − 12 (Mena &
Weed, 2019).
In the case of the entropic regularized BoMb-OT (eBoMb-OT) (see equation (4) in
Appendix B), when the regularized parameter λ is inﬁnity, namely, the m-OT, the result of (Fatras et al., 2020; 2021b) already established the sample complexity n − 12 .
However, for a general value of the regularized parameter, it is unclear whether we still have this sample complexity n − 12 of the eBoMb-OT.
We leave this question for the future work.
E XPERIMENTS
In this section, we demonstrate the favorable performance of BoMb-OT compared to m-OT in three discussed types of applications, namely, gradient-based, mapping-based, and value-based applications.
For gradient-based applications, we run experiments on deep generative model (Genevay et al., 2018;
Deshpande et al., 2018), deep domain adaptation (Bhushan Damodaran et al., 2018).
Experiments on color transfer (Rabin et al., 2014; Ferradans et al., 2014) are conducted as the example for mapping-based applications.
Lastly, we present results on approximate Bayesian computation (Bernton et al., 2019a) which is a value-based application.
In the main text of the paper, we only report and discuss the obtained experimental results, the details of applications and their algorithms are given in Appendix C. In Appendix D, we provide detailed experimental results of discussed applications including visualization and computational time.
Moreover, we also carry out experiments on gradient ﬂow (Santambrogio, 2017) and visualize mini-batch transportation matrices of the m-OT and the BoMb-OT.
From all experiments, we observe that the BoMb-OT performs better than m-OT consistently.
The detailed settings of our experiments including neural network architecture, hyperparameters, and evaluation metrics are in Appendix
E.
We would like to recall that k is the number of [Table] [Table] mini-batches and m is the mini-batch size.
In our experiments, we create mini-batches by sampling without replacement from supports of the original empirical measures.
4. D EEP GENERATIVE MODEL
We now show the deep generative model result on MNIST (LeCun et al., 1998), CIFAR (32x32) (Krizhevsky, 2009), and CelebA (64x64)
(Liu et al., 2015) using different mini-batch schemes with two OT losses for d : SW 2 (Deshpande et al., 2018) and W τ 2 (Genevay et al., 2018).
According to Table 1, the BoMb-OT always gives lower quantitative metrics than the m-OT on all datasets with all choices of the number of mini-batches k and mini-batch ground metric d .
Interestingly, the eBoMb-OT which is the entropic regularization version of BoMb-OT (deﬁned in Appendix B) is better than the BoMb-OT on CIFAR10, CelebA, and one setting on MNIST.
The explanation for this phenomenon is that the eBoMb-OT loss is smoother than the BoMb-OT due to the entropic regularization (Genevay et al., 2018).
From the table, we also observe that increasing the number of mini-batches k improves the result of generators that are trained with the same number of stochastic gradient updates.
It suggests that the gradient of an OT loss should not be estimated by only 1 pair of mini-batches like the way it has been implementing in practice.
Also, based on the experimental results, we recommend practitioners to consider using an additional transportation problem between measures over mini-batches when the number of mini-batches k > 1 .
4. D EEP DOMAIN ADAPTATION
In this section, we adapt two digits datasets SVHN (Netzer et al., 2011) and USPS (Hull, 1994) to
MNIST (LeCun et al., 1998).
Details about architectures of neural networks and hyper-parameters settings are given in Appendix E.2.
The classiﬁcation results are illustrated in Table 2.
The BoMb-OT always achieves better results than the m-OT 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 the
BoMb-OT, especially when dealing with a large number of mini-batches.
We also observe that the classiﬁcation accuracy on the target domain also increases as k changes from 8 to 32 and keeping the same number of stochastic gradient updates for deep neural networks.
Thus, we suggest that
DeepJDOT (Bhushan Damodaran et al., 2018) and its variants should be implemented with the
BoMb-OT’s strategy which is presented in Algorithm 2 in Appendix C.2. [Figure] [Figure]
4. C OLOR T RANSFER
Mini-batch OT has made color transfer to be able to transform color between two images that contain millions of pixels (Fatras et al., 2021b).
In this section, we show that our new mini-batch strategy can improve further the quality of the color transfer.
The details of the application and algorithm are given in Appendix C.3.
The full experimental results including qualitative images and their color palettes with different choices of k and m are presented in Appendix D.3.
Based on the experiment, we observe that the BoMb-OT provides a better barycentric mapping for color transfer than m-OT with every setting of k and m .
In this main paper, we show a color transfer result in Figure 2.
We can observe that the color of the transferred image of the BoMb-OT is more similar to the target image than one of the m-OT.
The color palettes in Appendix C.3 also reinforce this claim.
4. A PPROXIMATE B AYESIAN C OMPUTATION (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 k and m .
We compare them by plotting their approximated posteriors and computing the Wasserstein distance between approximated posteriors to the true posterior.
According to the experiments, the BoMb-OT always yield lower Wasserstein distances than the m-OT.
Here, we show some results in Figure 3.
From the ﬁgure, we can see that approximated posteriors from the
BoMb-OT are closer to the true posterior than the m-OT.
Moreover, we observe that increasing the value of k and m improves the quality of the approximated posteriors from both mini-batch schemes considerably.
Based on the fact that Wasserstein ABC is implemented with k = 1 (Bernton et al., 2019a), we suggest that the BoMb-OT should be used with k is greater than 1 in ABC.
C ONCLUSION
In the paper, we have presented a novel mini-batching scheme for optimal transport, named Batch of
Mini-batches Optimal Transport (BoMb-OT).
The idea of the BoMb-OT is to consider the optimal transport problem on the space of mini-batches with a Wasserstein-types ground metric.
We prove that the BoMb-OT is an approximation of a valid distance between probability measures and its entropic regularized version, eBoMb-OT, is the generalization of the conventional mini-batch optimal transport.
More importantly, we have shown that the BoMb-OT and the eBoMb-OT can be implemented efﬁciently and they have more favorable performance than the m-OT in various applications of optimal transport.
For future work, we could consider another extension of the BoMb-OT by changing the local OT to the Gromov-Wasseretein (GW) and unbalanced OT (UOT).
Reproducibility Statement: All datasets that we used in the paper are published and they are easy to ﬁnd on the Internet.
Source codes and instruction for our experiments are provided in the supplementary of the paper.
The details of experimental settings, computational infrastructure, and other used public libraries are given in Appendix E.
Ethics Statement: The paper investigates a fundamental practical problem of optimal transport in machine learning and deep learning.
Hence, we do not foresee any ethical issues of the paper that can be discussed.
Supplement to “On Transportation
of Mini-batches:
A Hierarchical Approach”
In this supplementary material, we collect several proofs 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 to hierarchical optimal transport, entropic regularization of BoMb-OT, sparsity of the BoMb-OT, and extension with non-OT mini-batch metrics in Appendix B. Furthermore, we provide a 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 in Appendix E.
A P ROOFS
In this appendix, we give detailed proofs of theorems that are stated in the main paper.
A. P ROOF OF T HEOREM
We ﬁrst prove that for any probability measures µ and ν ∈ P p (Θ) , there exists an optimal transportation plan γ ∗ such that
[Equation]
From the deﬁnition of the BoMb-OT, there is a sequence of transportation plans γ n ∈ Π( ⊗ m µ , ⊗ m ν ) such that
[Equation]
as n → ∞ .
Since Π( ⊗ m µ , ⊗ m ν ) is compact in the weak* topology (Villani, 2008), γ n weakly converges to some γ ∗ ∈ Π( ⊗ m µ , ⊗ m ν ) .
Since d ( P x m , P y m ) is continuous in terms of x m , y m ∈
X m , an application of Portmanteau theorem leads to
[Equation]
Putting the above results together, we obtain
[Equation]
Therefore, there exists an optimal transportation plan γ ∗ such that equation 3 holds.
We now proceed to prove that the BoMb-OT is a well-deﬁned metric in the space of Borel probability measures.
First, we demonstrate that D d ( µ, ν ) =
if and only if µ = ν .
In fact, from the deﬁnition of D d ( ., . ) , if we have two probability measures µ and ν such that D ( µ, ν ) = 0 , we ﬁnd that inf
[Equation]
Since d is a metric, it implies that there exists transportation plan γ ∗ ∈ Π( ⊗ m µ , ⊗ m ν ) such that
P X m =
P Y m γ ∗ -almost everywhere.
It demonstrates that for each x , there exists a permutation σ x of { 1 , . . .
, m } such that x i =
y σ x
( i ) for all i ∈ [ m ] .
Now, for any test function f : X → R , we have
[Equation]
Since (cid:82) X m (cid:81) mi =1 f ( x i ) d ⊗ m µ ( x ) = (cid:0)(cid:82)
X f ( x 1 ) dµ ( x 1 ) (cid:1) m and (cid:82)
X m (cid:81) mi =1 f ( y i )
d
[Equation]
y )
= (cid:0)(cid:82)
X f ( y 1 ) dν ( y 1 ) (cid:1) m , the above results show that
[Equation]
for any test function f : X → R .
It indicates that µ = ν .
On the other hand, if µ = ν , we have ⊗ m µ = ⊗ m ν .
We can construct the transportation plan ¯ γ ( x, y ) =
δ x ( y ) ⊗ m µ ( x ) for all ( x, y ) .
Then, we ﬁnd that
[Equation]
Therefore, we obtain the conclusion that D ( µ, ν ) =
if and only if µ = ν .
Now, we demonstrate the triangle inequality property of D d ( ., . ) , namely, we will show that for any probability measures µ 1 , µ 2 , and µ 3 , we have
[Equation]
The proof of this inequality is a direct application of gluing lemma (Berkes & Philipp, 1977;
De Acosta, 1982).
In particular, for any transportation plans γ 1 ∈ Π( ⊗ m µ 1 , ⊗ m µ 2 ) and γ 2 ∈ Π( ⊗ m µ 2 , ⊗ m µ 3 ) , we can construct a probability measure ξ on X m ×X m ×X m such that ξ ( ., .
, X m )
=
γ 1 ( ., . ) and ξ ( X m , ., . )
=
γ 2 ( ., . ) .
Therefore, we ﬁnd that
[Equation]
Taking the inﬁmum on the LHS of the above inequality with respect to γ 1 and γ 2 , we obtain the triangle inequality property of D ( ., . ) .
As a consequence, we reach the conclusion of the theorem.
A. P ROOF OF T HEOREM
To ease the presentation, for any product measures ⊗ m µ and ⊗ m ν , we denote the following loss between ⊗ m µ and ⊗ m ν as follows:
[Equation]
From the deﬁnitions of the BoMb-OT losses, we have D md ( µ, ν ) =
¯ D md ( ⊗ m µ , ⊗ m ν ) and
[Equation]
D k
[Equation]
m d
[Equation]
¯ D m d
[Equation]
⊗ m µ k , ⊗ m ν k ) .
Using the similar proof argument as that of Theorem 1, we can check that ¯ D md satisﬁes the triangle inequality, namely, for any product measures ⊗ m µ , ⊗ m ν , and ⊗ m η , we have ¯ D md ( ⊗ m µ , ⊗ m ν ) + ¯ D md ( ⊗ m ν , ⊗ m η ) ≥
¯ D md ( ⊗ m µ , ⊗ m η ) .
From the above deﬁnition and properties of ¯ D , we ﬁnd that
[Equation]
) − D md ( µ, ν ) (cid:12)(cid:12) = (cid:12)(cid:12) ¯ D md ( ⊗ m µ k , ⊗ m ν k )
− ¯ D md ( ⊗ m µ , ⊗ m ν ) (cid:12)(cid:12) ≤ (cid:12)(cid:12)
¯ D md ( ⊗ m µ k , ⊗ m ν k ) − ¯ D md ( ⊗ m µ , ⊗ m ν k ) (cid:12)(cid:12) + (cid:12)(cid:12)
¯ D md ( ⊗ m µ , ⊗ m ν k ) − ¯ D md ( ⊗ m µ , ⊗ m ν ) (cid:12)(cid:12) ≤ ¯ D md ( ⊗ m µ k , ⊗ m µ )
+ ¯ D md ( ⊗ m ν k , ⊗ m ν ) .
(i)
Since
d ≡ W p where p ≥ 1 , we have
d ( P X m , P Y m ) =
W p ( P X m , P Y m ) =
m 1 /p (cid:32) inf σ m (cid:88) i
= (cid:107) X
i − Y σ ( i ) (cid:107) p (cid:33) 1 /p ≤ m 1 /p m (cid:88)
i =
(cid:107) X
i − Y
i (cid:107) , where the inﬁmum is taken over all possible permutations σ of { 1 , 2 , . . .
, m } .
The above inequality indicates that
[Equation]
Since X is a compact subset of R N , the results of (Dudley, 1969; Dobri´c & Yukich, 1995; Fournier & Guillin, 2015) show that W 1 ( µ n , µ ) =
O P (
[Equation]
N ) .
Collecting these results together, we have
[Equation]
With similar argument, we also ﬁnd that ¯ D md
[Equation]
ν
[Equation]
= O P
[Equation]
m
[Equation] [Equation]
N
[Equation]
Putting all the above results together, we reach the conclusion of part (i) of the theorem.
(ii)
The proof of part (ii) follows directly from the argument of part (i) and the results that
SW p ( µ n , µ ) =
O P ( n − 1 / 2 ) and SW p ( ν n , ν ) =
O P ( n − 1 / 2 ) when X is a compact subset of
R N (Bobkov & Ledoux, 2019).
Therefore, we obtain the conclusion of part (ii) of the theorem.
B A DDITIONAL MATERIALS
Computational complexity: We now discuss the computational complexities of the BoMb-OT loss when d ≡ { W τp , SW p } .
When d ≡ W τp for some given regularized parameter τ > 0 , we can use the
Sinkhorn algorithm to compute d ( P X mi , P Y mj ) for any 1 ≤ i, j ≤ k .
The computational complexity of the Sinkhorn algorithm for computing it is of the order O ( m 2 ) .
Therefore, the computational complexity of computing the cost matrix from the BoMb-OT is of the order O ( m 2 k 2 ) .
Given the cost matrix, the BoMb-OT loss can be approximated with the complexity of the order O ( k 2 /ε 2 ) where ε stands for the desired accuracy.
As a consequence, the total computational complexity of computing the BoMb-OT is O ( m 2 k
+
k 2 /ε 2 ) .
When d ≡ SW p , the computational complexity for computing d ( P X mi , P Y mj ) is of the order
O ( m log m ) for any 1 ≤ i, j ≤ k .
It shows that the complexity of computing the cost matrix from the BoMb-OT loss is of the order
O ( m (log m ) k 2 ) .
Given the cost matrix, the complexity of approximating the BoMb-OT is at the order of O ( k 2 /ε 2 ) .
Hence, the total complexity is at the order of O ( m (log m ) k
+ k 2 /ε 2 ) .
Connection to hierarchical OT: At the ﬁrst sight, the BoMb-OT may look similar to hierarchical optimal transport (HOT).
However, we would like to specify the main difference between the BoMbOT and the HOT.
In particular, on the one hand, the HOT comes from the hierarchical structure of data.
For example, Yurochkin et al.
(Yurochkin et al., 2019) consider optimal transport problems on both the document level and the word level for document representation.
In the paper (Luo et al., 2020),
HOT is proposed to handle multi-view data which is collected from different sources.
Similarly, a hierarchical formulation of OT is proposed in (Lee et al., 2019b) to leverage cluster structure in data to improve alignment.
On the other hand, the BoMb-OT makes no assumption about the hierarchical structure of data.
We consider the optimal coupling of mini-batches as we want to improve the quality of mini-batch loss and its transportation plan.
Entropic regularized population BoMb-OT: We now consider an entropic regularization of the population BoMb-OT, which is particularly useful for reducing the computational complexity of the
BoMb-OT.
In particular, the entropic regularized population BoMb-OT (eBoMb-OT) between two probability measures µ and ν admits the following form:
[Equation]
where λ > 0 stands for a chosen regularized parameter.
From the above deﬁnition, the population eBoMb-OT is an interpolation between the population BoMb-OT and the population m-OT.
On the one hand, when λ → 0 , we have ED d ( µ, ν ) converges to D d ( µ, ν ) .
On the other hand, when λ → ∞ , the joint distribution γ approaches ⊗ m µ ⊗ ⊗ m ν and ED d ( µ, ν ) converges to U d ( µ, ν ) .
The transportation plan of the eBoMb-OT can be derived by simply replacing the coupling between mini-batches in Deﬁnition 3 by the coupling found by the Sinkhorn algorithm.
Sparsity of transportation plan from BoMb-OT: We assume that we are dealing with two empirical measures of n supports.
For d = W p , the optimal transportation plan contains n entries that are not zero.
In the m-OT case, with k mini-batches of size m , the m-OT’s transportation plan contains at most k 2 m non-zero entries.
On the other hand, with k mini-batches of size m , the BoMb-OT’s transportation plan has at most km non-zero entries.
For d = W τp ( τ > 0 ), the optimal transportation plans contains at most n 2 non-zero entries.
In this case, the m-OT provides transportation plans that have at most k 2 m 2 non-zero entries.
the BoMb-OT’s transportation plans contain at most kmnon-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 Deﬁnition 3 could be any discrepancy between empirical measures on mini-batches e.g. maximum mean discrepancy (MMD), etc.
In this case, we can see the outer optimal transport between measures over mini-batches as an additional layer to incorporate the OT property into the ﬁnal loss.
However, it is not easy to deﬁne the notion of a transportation plan in these cases.
C A PPLICATIONS AND B O M B -OT A LGORITHMS
In this section, we collect the details of applications that mini-batch optimal transport is used in practice including deep generative models, deep domain adaptation, color transfer, and approximate
Bayesian computation.
Moreover, we present corresponding algorithms of these applications with our new mini-batch scheme BoMb-OT.
[Figure]
[Equation]
[Figure]
G ENERATIVE M ODEL
Task description:
We ﬁrst consider the applications of the m-OT and the BoMb-OT into parametric generative modeling.
The goal of parametric generative modeling is to estimate a parameter of interest, says θ , which belongs to a parameter space Θ .
Each value of θ induces a model distribution µ θ over the data space, and we want to ﬁnd the optimal parameter θ ∗ which has µ θ ∗ as the closest distribution to the empirical data distribution ν n under a discrepancy (e.g. Wasserstein distance).
In deep learning setting, θ is the weight of a deep neural network that maps from a low dimensional manifold Z to the data space X , and the model distribution µ θ is a push-forward distribution G θ (cid:93)p ( z ) for p ( z ) is a white noise distribution on Z (e.g. N (0 , I ) ).
By using mini-batching schemes, we can estimate this parameter by minimizing the BoMb-OT loss in Deﬁnition 3:
[Equation]
where µ θ,n is the empirical distribution of µ θ .
The mentioned optimization is used directly in learning generative model on MNIST dataset with d ≡ ( SW 2 , W τ 2 ) in our experiments.
Algorithms: The algorithm for deep generative model with BoMb-OT is given in Algorithm 1.
This algorithm is used to train directly the generative model on MNIST dataset.
Metric learning:
Since L 2 distance is not a natural distance on the space of images such as
CelebA, metric learning was introduced by (Genevay et al., 2018; Deshpande et al., 2018) as a key step in the application of generative models.
The general idea is to learn a parametric ground metric cost c φ :
[Equation]
(6) where f φ : X → R h is a non-linear function that map from the data space to a feature space where
L 2 distance is meaningful.
The methodology to learn the function f φ depends on the choice of d .
For example, when d = W τ 2 , authors in (Genevay et al., 2018) seek for φ by solve following optimization:
[Equation]
Intialize C ∈ R k for i
1 to k do for j
1 to k do — off
Load X mi to GPU from I X mi
Load Z mj to GPU
Compute Y mj ← G θ ( Z mj )
Compute C ij ← L DGM ( P X mi , P Y mj ) end for end for —
Solve π ← OT ( u k , u k , C ) (entropic OT for eBoMb-OT) grad θ ←for i
1 to k do for j k do if then
Load X mi to GPU from I mi
Load Z mj to GPU
Compute Y mj ← G θ ( Z mj )
Compute C ij ← L DGM ( P X mi , P Y mj ) grad θ ← grad θ + π ij ∇ θ
C ij end
if end for end for θ end while
In (Deshpande et al., 2018), another metric learning technique is used for d = SW 2 :
[Equation] [Equation]
where g ψ : R h → [0 , 1] .
Learning the metric or f φ is carried out simultaneously with learning the generative model in practice, namely, after one gradient step for the parameter θ , φ will be updated by one gradient step.
C. D EEP D OMAIN A DAPTATION
We adapt the BoMb-OT into DeepJDOT (Bhushan Damodaran et al., 2018) which is a famous unsupervised domain adaptation method based on the m-OT.
In particular, we aim to learn an embedding function G θ : X → Z which maps data to the latent space; and a classiﬁer F φ : Z → Y which maps the latent space to the label space on the target domain.
For a given number of the mini-batches k and the size of mini-batches m , the goal is to minimize the following objective
Input: k , m , source domain ( S, Y ) , target domain T , chosen cost L DA in Equation 10.
Initialize G θ (parametrized by θ ), F φ (parametrized by φ ) while ( θ, φ ) do not converge do —
Sample indices uniformly I S m 1 ,Y m 1 , . .
, I S mk ,Y mk on ( S, Y )
Sample indices uniformly
I
T m 1 , . .
, I T mk on T
Initialize C ∈ R k × k for i = 1 to k do for j = 1 to k do —– On GPU, autograd off —–
Load ( S mi , Y mi ) to GPU from I S mi ,Y mi
Load T mj to GPU from I T mj
Compute mC ← L DA ( S mi , Y mi , T mj , G θ , F φ ) (Equation 10) u m ← (cid:0) 1 m , . .
, 1 m (cid:1)
Compute C ij ← OT
( u m , u m , mC ) end for end for —– On computer memory —– u k ← (cid:0) 1 k , . .
. , 1 k (cid:1)
Solve γ ← OT ( u k , u k , C )
(entropic OT for eBoMb-OT) grad θ ←grad φ ←for i = 1 to k do —
On
, autograd on —–
Load ( S mi , Y mi ) to GPU from I S mi ,Y mi grad θ ← grad θ + 1 k 1 m ∇ θ L s ( Y mi , F φ ( G θ ( S mi ))
(Equation 9) grad φ ← grad φ + 1 k 1 m ∇ φ L s ( Y mi , F φ ( G θ ( S mi ))
(Equation 9) for j = 1 to k do if γ ij (cid:54) =
then —
Load mj
Compute mC ← L DA
[Equation]
Compute C ij ← OT
( u m , u m , mC ) grad θ ← grad θ + γ ij ∇ θ C ij grad φ ← grad φ + γ ij ∇ φ C ij end if end for end for —– On GPU —– θ ← Adam ( θ, grad θ ) φ ← Adam ( θ, grad φ ) end while function:
[Equation]
where L s is the source loss function (e.g. classiﬁcation loss), S m 1 , . .
. , S mk are source mini-batches which are sampled with or without replacement from the source domain S m ∈
X m , Y m 1 , . .
, Y mk are corresponding labels of S m 1 , . . .
, S mk , with S mi
= { s i 1 , . . .
, s im } and Y m 1 :=
{ y i 1 , . . .
, y im } .
Similarly, T m 1 , . .
, T m k
( T m 1 := { t i 1 , . . .
, t
im } ) are target mini-batches which are sampled with or without replacement from the target domain T m ∈
X m .
The cost matrix C G,FS mi ,Y mi ,T mi denoted mC is deﬁned as follows:
[Equation]
where L t is the target loss function, α and λ t are hyper-parameters that control two terms.
[Equation]
In this appendix, we carry out more experiments on color transfer with various settings.
In detail, we compare the m-OT and the BoMb-OT in different domains of images, and we demonstrate visually the interpolation property of the eBoMb-OT by varying the regularization parameter.
Methodology:
In our experiment, we ﬁrst compress both the source image and the target image using K-means clustering with 3000 components.
After that, we look for the transportation plans between cluster centers of two images from different approaches, the m-OT, the BoMb-OT, and the full-OT.
Next, we change the values of all pixels in a cluster of the source image based on the value of the corresponding cluster centers of the target image with the previously found plans.
We present the algorithm that we use in color transfer with the BoMb-OT in Algorithm 3.
This algorithm is adapted from the algorithm that is used for the m-OT in (Fatras et al., 2020).
Algorithm 3 Color Transfer with BoMb-OT
Input: k, m, T source image X s ∈ R n × 3 , target image X t ∈ R n ×
Initialize Y s ∈ R n ×for t = 1 to T do
Initialize C ∈ R k × k
Sample indices I X m 1 , . .
, I X mk from X s Sample indices
I
Y m 1 , . .
, I Y mk from X t for i = 1 to k do for j = 1 to k do
Load X mi from I X mi
Load Y mi from I Y mi
Compute cost matrix M between X mi and Y mi π ← arg min π ∈ Π( u m , u m ) (cid:104) M, π (cid:105)
C ij ← (cid:104) M, π (cid:105) end for end for γ ← arg min γ ∈ Π( u k , u k ) (cid:104) C, γ (cid:105) for i = 1 to k do for j = 1 to k do if γ ij (cid:54) =
then
Load X mi from I X mi
Load Y mi from I Y mi
Compute cost matrix M between X mi and Y mi π ← arg min π ∈ Π( u m , u m ) (cid:104) M, π (cid:105)
Y s |
I
Xmi ← mγ ij π · X t | I mYj end if end for end for end for
Output: Y s
C.
A PPROXIMATE B AYESIAN C OMPUTATION (
ABC
In this appendix, details of Approximate Bayesian Computation (ABC) and the usage of mini-batches with BoMb-OT are discussed.
Input: Generative model p ( x, θ ) , observation { x i } ni =1 , number of iterations T , discrepancy measure D , tolerance threshold (cid:15) , number of particles m , and summary statistics s (optional).
for t = 1 to T do repeat
Sample until end for
Output
Review
on ABC
The central task of Bayesian inference is to estimate a distribution of parameter θ ∈ Θ given n data points { x i } ni =1 and a generative model p ( x, θ ) .
Using Bayes rule, the posterior can be written as p ( θ | x 1: n ) :
= p ( x 1: n | θ ) p ( θ ) p ( x 1: n ) .
Generally, the posterior is intractable since we cannot evaluate the normalizing constant p ( x 1: n ) (or the evidence).
It leads to the usage of approximate
Bayesian Inference, e.g., Markov Chain Monte Carlo and Variational Inference.
However, in some settings, the likelihood function p ( x 1: n | θ ) cannot be evaluated such as implicit generative models.
In these cases, Approximate Bayesian Computation (or likelihood-free inference) is a good framework to infer the posterior distribution since it only requires the samples from the likelihood function.
We present the algorithm of ABC in Algorithm 4 that is used to obtain posterior samples of θ .
The thing that sets ABC apart is that it can be implemented in distributed ways, and its posterior can be shown to have the desirable theoretical property of converging to the true posterior when (cid:15) → 0 .
However, the performance of ABC depends on the choice of the summary statistics s (e.g., empirical mean and empirical variance) and the discrepancy D .
In practice, constructing sufﬁcient statistics is not easy.
Thus, a discrepancy between empirical distributions is used to avoid this non-trivial task.
Currently, Wasserstein distances have drawn a lot of attention from ABC’s community because they provide a meaningful comparison between non-overlap probability distributions (Bernton et al., 2019a; Nadjahi et al., 2020).
When the number of particles m in Algorithm 4 is the largest OT problem that can be solved by the current computational resources, the mini-batch approach can be utilized to obtain more information about two measures, namely, more supports in sample space can be used.
In particular, we utilize the mini-batch OT losses as the discrepancy D in Algorithm 4.
The detail of the application of BoMb-OT in ABC is given in Algorithm 5.
D A DDITIONAL E XPERIMENTS
In this appendix, we provide additional experimental results that are not shown in the main paper.
In
Appendix D.1, we present random generated images on MNIST dataset, CIFAR10 dataset, CelebA dataset, and training time comparison between the m-OT and the BoMb-OT.
Furthermore, we give detailed result on deep domain adaptation in Appendix D.2 and we also provide the computational time.
In Appendix D.3 we carry out more experiments on color transformation with various settings on different images and their corresponding color palettes.
Moreover, details of setting of the ABC experiment and additional comparative experiments are presented in Appendix D.4.
After that, we compare the m-OT and the BoMb-OTs on gradient
ﬂow application in Appendix D.5.
Finally, we investigate the behavior of the m-OT and the (e)BoMb-OT in estimating the optimal transportation plan in Appendix D.6.
D. D EEP G ENERATIVE MODEL
Generated images:
We show the generated images on MNIST dataset in Figure 4, the generated images on CelebA in Figure 5, the generated images on CIFAR10 in Figure 6.
For both choices of d ( SW 2 , W τ 2 ), we can see that the images from (e)BoMb-OT is more realistic than the images from m-OT.
This qualitative result supports the quantitative in Table 1 in the main text.
Algorithm
Input: Generative model p ( x, θ ) , observation { x i } ni =1 , number of iterations T , optimal transport discrepancy measure D OT , tolerance threshold (cid:15) , number of particles m , number of mini-batches k , and summary statistics s (optional).
for t = 1 to T do repeat
Sample θ ∼ p ( θ )
Sample indices I X m 1 , . .
, I X mk from { x i } ni =
Sample { y i } ni
=1 ∼ p ( x | θ )
Sample indices I Y m 1 , . . .
, I Y mk from { y i } ni =
Initialize C ∈ R k × k for i = 1 to k do for j = 1 to k do
Load X mi , Y mj from their indices
I X mi , I Y mj
C ij ← D OT
( X mi , Y mj ) end for end for u k ← (cid:0) 1 k , . .
. , 1 k (cid:1)
D ← OT ( u k , u k , C ) until D ≤ (cid:15)
θ t = θ end for
Output: { θ t } Tt =
[Figure]
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 k does decrease the number of iterations per second in all experiments.
For both datasets, using the sliced Wasserstein distance ( d = SW 2 ) leads to higher iterations per second compared to the entropic Wasserstein distance ( d = W τ 2 ) on the high dimensional space.
This result is expected since the slicing approach aims to reduce the computational complexity.
On the MNIST dataset, we observe that the entropic regularized version of the BoMb-OT (eBoMb-OT) is the slowest for both choices of d .
The m-OT is the fastest approach with the number of iterations per second of 7 . 89 , 2 . 54 , and 0 .
85 for k = 2 , 4 , and 8 respectively.
In contrast, the BoMb-OT is faster than the m-OT for CIFAR10 and CelebA dataset because we set a high value for the regularized parameter (e.g. τ = 50 , 40 ).
For the CelebA dataset, although the eBoMb-OT is the slowest method when d = SW 2 , it becomes the fastest approach if we set d to W τ 2 .
The reason for this phenomenon is because of the sparsity of the (e)BoMb-OT’s transportation plan between measures over mini-batches.
In particular, pairs of mini-batches that are zero can be skipped
[Figure] [Figure]
Dataset
[Table]
MNIST 27.
CIFAR  22.
CelebA 6.
[Table]
(see in Algorithm 1.
So, the (e)BoMb-OT can avoid estimating gradient of the neural networks of a bad pair of mini-batches.
[Table] [Table]
[Table]
D. D EEP DOMAIN ADAPTATION
In this section
[Table]
we compare the performance of the BoMb-OT and the-OT on deep domain adaptation when changing the mini-batch size m .
Comparison between the m-OT and BoMb-OT:
As seen in Table 4, the BoMb-OT produces better classiﬁcation accuracy than the m-OT in all experiments.
In addition, it leads to a huge performance improvement (over 19% ) in comparison with the m-OT when the mini-batch size is small ( k = 10 for SVHN to MNIST and k = 5 for USPS to MNIST).
When the mini-batch size becomes larger, the performance of both methods also increases.
An explanation for such outcome is that a large mini-batch size makes mini-batch loss approach to its population version.
Computational speed:
The number of iterations per second of deep DA can be found in Table 5-6.
Similar to deep generative model, we observe a phenomenon that the speed of both the m-OT and
BoMb-OT decreases as k increases.
Interestingly, increasing the mini-batch size m when adapting from USPS to MNIST barely affects the running speed of both methods.
The run time of the m-OT nearly doubles that of the BoMb-OT.
Speciﬁcally, the BoMb-OT averagely runs 1 iteration in a second while an iteration of the m-OT consumes roughly 2 seconds.
Although having comparable time complexity, the BoMb-OT in practice runs faster than the m-OT in all deep DA experiments.
This is because of the sparsity of the BoMb-OT’s transportation plan between measures over mini-batches.
In particular, pairs of mini-batches that are zero can be skipped (see in Algorithm 2).
We would like to recall that, gradient estimation is the most time-consuming task in deep learning applications.
D. C OLOR T RANSFER
In this appendix, we compare the m-OT and the BoMb-OT in different domains of images including natural images and arts.
Comparison between the m-OT and the BoMb-OT:
Next, we illustrate the color-transferred images using the m-OT and the BoMb-OT with two values of the number of mini-batches and the size of mini-batches ( k, m ) =
(10 , 10) and ( k, m )
= (20 , 20) .
The number of incremental step T (see in
Algorithm 3) is set to 5000.
We show the source images, target images, and corresponding transferred images in Figure 7-Figure 10.
It is easy to see that the transferred images from the BoMb-OT look more realistic than the m-OT and the color is more similar to the target images.
The color palette of transferred images also reinforces the above claim when it is closer to the color palette of the target
[Figure]
images.
According to those ﬁgures, increasing the number of mini-batches and the mini-batch size improves the results considerably.
Computational speed:
For k = 10 , m = 10 , the m-OT has the speed of 103 iterations per second while the BoMb-OT has the speed of 100 iterations per second.
For k = 20 , m = 20 , the speed of m-OT is about 20 iterations per second and the speed of the BoMb-OT is also about 20 iterations per second.
It means that the additional k × k OT is not too expensive while it can improve the color palette considerably.
D. A PPROXIMATE B AYESIAN C OMPUTATION (ABC)
Settings: We use the same setup as in (Nadjahi et al., 2020): there are n = 100 observations { x i } ni =1 i.i.d from multivariate Gaussian N ( µ ∗ , σ 2 ∗
I N ) where N is the dimension, µ ∗ ∼ N (0 , I N ) and σ 2 ∗ = 4 .
The task is to estimate the posterior of σ 2 under the imaginary assumption that σfollows inverse gamma distribution IG (1 , 1) .
Under these assumptions, the posterior of σ 2 has the form IG (1 + n d 2 , 1 + 12 (cid:80) ni
=
|| x
i − ν ∗ || 2 ) .
For ABC, we use the m-OT and the BoMb-OT with the Wasserstein-2 ground metric for the acceptance-rejection sampling’s criteria in sequential
[Figure]
Monte Carlo ABC (Toni et al., 2009) (using implementation from pyABC (Klinger et al., 2018) with 100 particles and 10 iterations).
The mini-batches’ size is set in { 8 , 16 , 32 } and the number of mini-batches is in { 2 , 4 , 6 , 8 } .
Results on ABC with mini-batches:
After obtaining all samples, we estimate their densities by utilizing Gaussian kernel density estimation and then plot the approximated posteriors and the true posterior in Figure 11.
Here we run the algorithm with several values of ( k, m ) , namely, m ∈ { 8 , 16 , 32 } and k ∈ { 2 , 4 , 6 , 8 } .
From these results, we see that a bigger m usually returns a better posterior in both the m-OT and the BoMb-OT cases.
Similarly, increasing k also improves the performance of the two methods.
Moreover, these graphs strengthen the claim that the BoMb-OT outperforms the m-OT in every setting of ( k, m ) since its posteriors are always closer to the true posterior than those of the m-OT in both visual result and Wasserstein distance.
[Figure]
D. N ON - PARAMETRIC 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 ﬂow application
. Task description:
Gradient ﬂow is a non-parametric method to learn a generative model.
Like every generative model, the goal of gradient ﬂow is to mimic the data distribution ν by a distribution µ .
It leads to the functional optimization problem:
[Equation]
where D is a predeﬁned discrepancy between two probability measures.
So, a gradient ﬂows can be constructed:
[Equation]
[Figure]
We follow the Euler scheme to solve this equation as in (Feydy et al., 2019), starting from an initial distribution at time t = 0 .
In this paper, we choose D be (e)BoMb- W 2 and m- W 2 for the sake of comparison between them.
We ﬁrst consider the toy example as in (Feydy et al., 2019) and present our results in Figure 12.
The task is to move the colorful empirical measure to the "S-shape" measure.
Each measure has 1000 support points.
Here, we choose ( k, m ) =
(4 , 16) , the OT loss is Wasserstein-2, and we use the Wasserstein-2 score to evaluate the performance of the mini-batching scheme.
From Figure 12, BoMb-OT and eBoMb-OT provide better ﬂows than m-OT, namely, Wasserstein-2 scores of
BoMB-OT and eBoMb-OT are always lower than those of m-OT in every step.
In addition, we do an extra setup with a higher of mini-batches, ( k, m ) =
(16 , 16) to show the increasing the number of mini-batches improve the performance of both m-OT and (e)BoMb-OT.
The result is shown in
Figure 13.
In this setting, the BoMb-OT still shows its favorable performance compared to the m-OT, namely, its Wasserstein-2 scores are still lower than the m-OT in every step.
CelebA:
Let µ and ν denote the empirical measures deﬁned over 5000 female images andmale images in the CelebA dataset.
Thus, we can present the transformation from a male face to be a female one by creating a ﬂow from µ to ν .
Our setting experiment is the same as (Fatras et al., 2020).
We ﬁrst train an autoencoder on CelebA, then we compress two original measures to the
[Figure] [Figure]
low-dimensional measures on the latent space of the autoencoder.
After having the latent measures, we run the Euler scheme to get the transformed measure then we decode it back to the data space by the autoencoder.
The result is shown in Figure 14, we also show the closest female image (in sense of
L 2 distance) to the ﬁnal found image in each method and the corresponding L 2 distance between the middle step images and the nearest image.
As shown, (e)BoMb-OT provides a better ﬂow than m-OT does.
D. T RANSPORTATION PLANS
In this appendix, we investigate deeper the behavior of the m-OT and the (e)BoMb-OT in estimating the transportation plan.
We present a toy example to illustrate the transportation plans of the m-OT, the BoMb-OT, and the original OT (full-OT).
In particular, we sample two empirical distributions µ n and ν n , where n = 10 , from N (cid:0) (cid:20)
(cid:21) , (cid:20) 1 0
(cid:21) (cid:1) and N (cid:0) (cid:20) 44 (cid:21) , (cid:20)
[Figure]
8 −
[Figure] [Figure]
(cid:21) (cid:1) respectively.
[Figure]
[Figure]
Transportation plans:
We plot the graph of samples matching and transportation matrices in
Figure 15.
When the size of mini-batches equals 2 and the number of mini-batches is set to be 20, the m-OT approach produces messy OT matrices and disordered matching, meanwhile, the BoMb-OT can still concentrate the mass to meaningful entries of the transportation matrix, thus, its matching is acceptable.
Increasing the size of mini-batches to 8, m-OT’s performance improves signiﬁcantly however its matching and its matrices are visibly still not good solutions.
In contrast, the BoMb-OT is able to generate nearly optimal matchings and transportation matrices.
Next, we test the m-OT and the BoMb-OT in real applications in the case of an extremely small number of mini-batches.
When there are 8 mini-batches of size 2, the m-OT still performs poorly while the BoMb-OT creates a sparser transportation matrix that is closer to the optimal solution obtained by the full-OT.
Similarly, with 2 mini-batches of size 8, the BoMb-OT has only 4 wrong matchings while that number of the
O T o f s a m p l e s m-OT (n,m)=(10,2) n n 0 1 2 3 4
BoMb-OT (n,m)=(10,2) n n 0 1 2 3 4 5113 m-OT (n,m)=(10,8) n n 0 1 2 3 4
BoMb-OT (n,m)=(10,8) n n 0 1 2 3 4 5113 full-OT n n
O T m a t
r
i x                                        0 1 2 3 4
O T o f s a m p l e s m-OT (n,m)=(10,2) n n 0 1 2 3 4
BoMb-OT (n,m)=(10,2) n n 0 1 2 3 4 5113 m-OT (n,m)=(10,8) n n 0 1 2 3 4
BoMb-OT (n,m)=(10,8) n n 0 1 2 3 4 5113 full-OT n n
O T m a t
r
i x                                        m-OT is 10.
In conclusion, the BoMb-OT is the better version for mini-batching with any value of the number of mini-batches and mini-batches size.
Entropic transportation plan of the BoMb-OT:
We empirically show that the transportation plan of the eBoMb-OT, when the entropic regularization is sufﬁciently large, reverts into the m-OT’s plan.
The result is given in Figure 16.
When λ = 10 3 , the transportation plans of the eBoMb-OT are identical to plans of the m-OT with every choice of ( k, m ) .
However, when λ = 0 .
1 , despite not being identical to full-OT, the plans produced by the eBoMb-OT are still close to the true optimal plan.
Comparison with
stochastic averaged gradient (SAG):
We compare the BoMb-OT with the stochastic averaged gradient (SAG) (Aude et al., 2016) for computing OT.
The transportation matrices are given in Figure 17.
We would like to recall that SAG still need to store full cost matrix ( n × ) while the BoMb-OT only need to store smaller cost matrices ( m × m ).
We can see that the BoMb-OT gives better than transportation plan than SAG when m = 2 .
When m = 8 , the BoMb-OT still seems to be better.
E E XPERIMENT S ETTINGS
In this appendix, we collect some necessary experimental setups in the paper including generative model, gradient ﬂow.
We use POT (Flamary et al., 2021) for OT solvers, and the pyABC (Klinger et al., 2018) for the ABC experiments.
O T o f s a m p l e s m-OT (n,m)=(10,2) n n 0 1 2 3 4 5113 eBoMb-OT (n,m, )=(10,2,0.1) n n 0 1 2 3 4 5113 m-OT (n,m)=(10,8) n n 0 1 2 3 4 5113 eBoMb-OT (n,m, )=(10,8,0.1) n n 0 1 2 3 4 5113 full-OT n n
O T m a t
r
i x                                        0 1 2 3 4
O T o f s a m p l e s m-OT (n,m)=(10,2) n n 0 1 2 3 4 5113 eBoMb-OT (n,m, )=(10,2,10 3 ) n n 0 1 2 3 4 5113 m-OT (n,m)=(10,8) n n 0 1 2 3 4 5113 eBoMb-OT (n,m, )=(10,8,10 3 ) n n 0 1 2 3 4 5113 full-OT n n
O T m a t
r
i x
E. D EEP G ENERATIVE MODEL
Wasserstein-2 scores: We use empirical distribution with 11000 samples that are obtained by sampling from the generative model and the empirical test set distribution respectively, then we compute discrete Wasserstein-2 distance via linear programming.
FID scores:
We use 10000 samples from the generative model and all test set images to compute
FID score (Heusel et al., 2017).
Parameter settings:
We chose a learning rate equal to 0 .
0005 , batch size equal to 100 , number of epochs in MNIST equal to 100 , number of epochs in CelebA equal to 25 , number of epochs in CIFAR10 equal to 50 .
For d = SW 2 we use the number of projections L =on MNIST and L = 100 on CIFAR10 and CelebA. For d = W τ 2 , we use τ = 1 for MNIST, τ = 50 for CIFAR10, τ = 40 for CelebA. For the eBoMb-OT, we choose the best setting for λ ∈ { 1 , 2 , 3 , 4 , 5 , 10 , 20 , 30 , 40 , 50 , 60 , 70 , 80 }
Neural network architectures:
We use the MLP for the generative model on the MNIST dataset, while CNNs are used on the CelebA dataset.
Generator architecture was used for MNIST dataset:
[Equation]
Generator architecture was used for CelebA: z
[Equation]
O T o f s a m p l e s
SAG (n,m)=(10,2) n n 0 1 2 3 4 51135 BoMb-OT (n,m)=(10,2) n n 0 1 2 3 4
SAG (n,m)=(10,8) n n 0 1 2 3 4 51135 BoMb-OT (n,m)=(10,8) n n 0 1 2 3 4 5113 full-OT n n
O T m a t
r
i x
BatchNorm →
ReLU →
TransposeConv→
BatchNorm →
ReLU →
TransposeConv 128 → BatchNorm → ReLU → TransposeConv 64 → BatchNorm →
ReLU → TransposeConv 1 → Tanh
Metric learning neural network’s architecture was used for CelebA:
First part: x ∈ R 64 × 64 × 3 → Conv 64 → LeakyReLU 0 .
→ Conv 128 → BatchNorm → LearkyReLU 0 .
→ Conv 256 → BatchNorm → LearkyReLU 0 .
→ Conv 512 → BatchNorm → Tanh Second part: Conv 1 → Sigmoid
Generator architecture was used for CIFAR
: z ∈ R 128 → TransposeConv 256 → BatchNorm → ReLU → TransposeConv → BatchNorm →
ReLU →
TransposeConv 64 → BatchNorm → ReLU → TransposeConv 1 → Tanh
Metric learning neural network’s architecture was used for CIFAR:
First part: x ∈ R 32 × 32 × 3 → Conv 64 → LeakyReLU 0 .
→ Conv 128 → BatchNorm → LearkyReLU 0 .
→ Conv 256 → BatchNorm → Tanh Second part: Conv 1 → Sigmoid
E.
D EEP DOMAIN ADAPTATION
In this section we state the neural network architectures and hyper-parameters for deep domain adaptation.
Evaluation metric:
The classiﬁcation accuracy is utilized to evaluate the mini-batch methods.
Parameter settings:
The number of mini-batches k varies in { 8 , 16 , 32 } .
For the SVHN dataset, the mini-batch size m is set to 50 .
Because the USPS dataset has fewer samples than the SVHN dataset, m is set to 25 .
We train the network using Adam optimizer with an initial learning rate of 0 . 0002 .
The number of epochs is 80 for k = 8 .
As the number of mini-batches doubles, we double the number of epochs so that the number of iterations does not change.
Following DeepJDOT (Bhushan Damodaran et al., 2018) and JUMBOT (Fatras et al., 2021a), we stratify the data loaders so that each class has the same number of samples in the mini-batches.
The hyperparameters in equation 10 follow the settings in DeepJDOT:
α = 0 .
001 , λ t = 0 . 0001 .
For eBoMb, we choose the best value of (cid:15) in { 0 . 01 , 0 . 1 , 0 . 2 , 1 , 10 } .
Our neural networks are trained on the source domain during 10 epochs before applying either method.
Neural network architectures:
We use CNN for our generator and 1 FC layer for our classiﬁer in both adaptation scenarios.
Generator architecture was used for SVHN dataset: z ∈ R 32 × 32 × 3 → Conv 32 → BatchNorm → ReLU → Conv 32 → BatchNorm → ReLU →
MaxPool 2 D → Conv 64 → BatchNorm → ReLU
→ Conv 64 → BatchNorm → ReLU →
MaxPool 2 D → Conv 128 → BatchNorm → ReLU → Conv 128 → BatchNorm →
ReLU → MaxPool 2 D → Sigmoid → FC
Generator architecture was used for USPS dataset: z ∈ R 28 × 28 × 3 → Conv 32 → BatchNorm → ReLU → MaxPool 2 D → Conv 64 →
BatchNorm → ReLU → Conv 128 → BatchNorm → ReLU → MaxPool 2 D →
Sigmoid → FC
Classiﬁer archiecture was used for both SVHN and USPS datasets: z ∈ R
→ FC
E. G RADIENT FLOW
For the implementation of the gradient ﬂow, we use the geomloss library (Feydy et al., 2019).
The learning rate is set to 0.001.
For the autoencoder in CelebA experiments, we use the repo in https://github.com/rasbt/deeplearning-models/blob/master/ pytorch_ipynb/autoencoder/ae-conv-nneighbor-celeba.ipynb for the pretrained autoencoder.
E. C OMPUTATIONAL INFRASTRUCTURE
All deep learning experiments are done on a RTX 2080 Ti GPU and a GTX 1080 Ti.
Other experiments are done on a MacBook Pro 11inc M1.