Generating Synthetic Datasets by Interpolating along Generalized Geodesics
Abstract
Data for pretraining machine learning models often consists of collections of heterogeneous datasets.
Although training on their union is reasonable in agnostic settings, it might be suboptimal when the target domain —where the model will ultimately be used— is known in advance.
In that case, one would ideally pretrain only on the dataset(s) most similar to the target one.
Instead of limiting this choice to those datasets already present in the pretraining collection, here we explore extending this search to all datasets that can be synthesized as ‘combinations’ of them.
We define such combinations as multi-dataset interpolations, formalized through the notion of generalized geodesics from optimal transport (OT) theory.
We compute these geodesics using a recent notion of distance between labeled datasets, and derive alternative interpolation schemes based on it: using either barycentric projections or optimal transport maps, the latter computed using recent neural OT methods.
These methods are scalable, efficient, and —notably— can be used to interpolate even between datasets with distinct and unrelated label sets.
Through various experiments in transfer learning, we demonstrate this promising new approach to targeted on-demand dataset synthesis.
Introduction
Recent progress in machine learning has been characterized by the rapid adoption of large pretrained models as a fundamental building block (Brown et al., 2020).
These models are typically pretrained on large amounts of general purpose data and then adapted (e g., fine-tuned) to a specific task of interest.
Such pretraining datasets usually draw from multiple heterogeneous data sources, e g., from arising from different domains or sources.
Traditionally, all available datasets are used in their entirety during pretraining, for example by pooling them together into a single dataset (when they all share the same label sets) or by training in all of them sequentially one by one.
These strategies, however, come with important disadvantages.
Training on the union of multiple datasets might be prohibitive or too time consuming, and it might even be detrimental.
Indeed, there is a growing line of research showing evidence that removing pretraining data sometimes helps transfer performance (Jain et al., 2022).
On the other hand, sequential learning (e.g., consuming datasets one by one) is infamously prone to catastrophic forgetting (McCloskey & Cohen, 1989; Kirkpatrick et al., 2017): the information from earlier datasets gradually vanishing as the model is trained on new datasets.
While all of these suggest that training on only some subset of the pretraining datasets, how to choose these is unclear.
However, when the target dataset on which the model is to be used is known in advance, the answer is much easier: intuitively, one would train only on those relevant to the target one: e.g., those most similar to it.
Indeed, recent work has shown that selecting pretraining datasets based on the distance to the target is a successful strategy (Alvarez-Melis & Fusi, 2020; Gao & Chaudhari, 2021).
However, such methods are limited to selecting (only) among individual datasets already present in the collection.
In this work, we propose a novel approach to generate synthetic pretraining datasets as combinations of existing ones.
In particular, this method searches among all possible continuous combinations of the available datasets, and thus is not limited to select one of them necessarily.
When given access to the target dataset of interest, we seek among all such combinations the one closest (in terms of a metric between datasets) to the target.
By characterizing datasets as sampled from a underlying probability distribution, this problem can be understood as a generalization (from Euclidean to probability space) of the problem of finding among the convex hull of a set of reference points, that closest to a query point.
While this problem has a simple closed-form solution in Euclidean space (via an orthogonal projection), solving it probability space is, as we shall see here, much more challenging.
We tackle this problem from the perspective of interpolation.
Formally, we model the combination of datasets as an interpolation between their distributions, formalized through the notion of geodesics in probability space endowed with the Wasserstein metric
(Ambrosio et al., 2008; Villani, 2008).
In particular, we rely on generalized geodesics (Craig, 2016;
Ambrosio et al., 2008), constant-speed curves connecting a pair (or more) distributions parametrized with respect to a ‘base’ distribution, whose role is played by the target dataset in our setting.
Computing such geodesics requires access to either an optimal transport coupling or map between the base distribution and every other reference distribution.
The former can be computed very efficiently with off-the-shelf OT solvers, but are limited to generate only as many samples as the problem is originally solved on.
In contrast, OT maps allow for on-demand out-of-sample mapping, and can be estimated using recent advances in neural OT methods (Fan et al., 2020;
Korotin et al., 2022b; Makkuva et al., 2020).
However, most existing OT methods assume unlabeled (feature-only) distributions, but our goal here is to interpolate between classification (i.e., labeled) datasets.
Therefore, we leverage a recent generalization of OT to labeled datasets to compute couplings (Alvarez-Melis & Fusi, 2020), and adapt and generalize neural
OT methods to the labeled setting to estimate OT maps.
In summary, the contributions of this paper are: (i) a novel approach to generate new synthetic classification datasets from existing ones by using geodesic interpolations, applicable even if they have disjoint label sets, (ii) two efficient methods to compute generalize geodesics, which might be of independent interest, (iii) empirical validation of the method in a transfer learning setting.
Background 2.1 Distributional interpolation with OT
Consider P ( X ) the space of probability distribution finite second moments over some Euclidean space X .
Given µ, ν ∈ P ( X ) , the Monge formulation optimal transport problem seeks a map
T : X → X that transforms µ into ν at minimal cost.
Formally, the objective of this problem is min T : T♯µ =
ν (cid:82) R d ∥ x − T ( x ) ∥
22 d ν ( x ) , where the minimization is over all the maps that pushforward distribution µ into distribution ν .
While a solution to this problem might not exist, a relaxation due to Kantorovich is guaranteed to have one.
This modified version yields the 2-Wasserstein distance:
W 22 ( µ, ν ) =
min π ∈ Π( µ,ν ) (cid:82) R d ∥ x − x ′ ∥
22 d π ( x, x ′ ) , where now the constraint set Π( µ, ν ) = { π ∈
P ( X 2 ) | P 0 ♯ π = µ, P 1 ♯π = ν } contains all couplings with marginals µ and ν .
The optimal such coupling is known as the OT plan.
A celebrated result by Brenier (1991) states that whenever P has density with respect to Lebesgue measure, the optimal T ∗ exists and is unique.
In that case, the
Kantorovich and Monge formulations coincide and their solutions are linked by π ∗
= ( Id , T ∗ )
♯ µ where Id is the identity map.
The Wasserstein-2 distance enjoys many desirable geometrical properties compared to other distances for distributions (Ambrosio et al., 2008).
One such property is the characterization of geodesics in probability space (Agueh & Carlier, 2011; Santambrogio, 2015).
When P ( X ) is equipped with metric W p , the unique minimal geodesic between any two distributions µ 0 and µ 1 is fully determined by π , the optimal transport plan between them, through the relation:
[Equation]
known as displacement interpolation .
If the Monge map exists, the geodesic can also be written as
[Equation]
and is known as McCann’s interpolation (McCann, 1997).
It easy to see that ρ M 0 = µ 1 and ρ M 1 = µ 2 .
Such interpolations are only defined between two distributions.
When there are m ≥ 2 marginal distributions { µ 1 , . . .
, µ m } , the Wasserstein barycenter
ρ Ba :=
arg min ρ (cid:80) mi =
a i W 22 ( ρ, µ i ) , a ∈ ∆ m ⊂ R m generalizes McCann’s interpolation (Agueh & Carlier, 2011).
Intuitively, the interpolation parameters a =
[ a 1 , . .
, a m ] determine the ‘mixture proportions’ of each dataset in the combination, akin to a convex combination of points in Euclidean space.
In particular, when a is a one-hot vector with a i = 1 , then ρ Ba = µ i , i.e., the barycenter is simply the i -th distribution.
Barycenters have attracted great attention in machine learning recently (Srivastava et al., 2018; Korotin et al., 2021), but they remain challenging to compute in high dimension (Fan et al., 2020; Korotin et al., 2022a).
Another limitation of these interpolation notions is the non-convexity of W 22 along them.
In Euclidean space, given three points x 1 , x 2 , y ∈ R d , the function t (cid:55)→ ∥ x t − y ∥ 22 , where x t is the interpolation x t = (1 − t ) x 1 + tx 2 , is convex.
In contrast, in Wasserstein space, neither the function t (cid:55)→
W 22 ( ρ Mt , ν ) or (cid:55)→ W 22 ( ρ Ba , ν ) are guaranteed to be convex (Santambrogio, 2017, Sec. 4.4).
This complicates theoretical analysis, such as in gradient flows.
To circumvent this issue, Ambrosio et al.
(2008) introduced the generalized geodesic of { µ 1 , . . . , µ m } with base ν and defined it as
[Equation]
∆ m , where T ∗ i is the optimal map from ν to µ i .
Lemma 1.
The functional µ (cid:55)→ W 22 ( µ, ν ) is convex along the generalized geodesics, and
W
[Equation]
Thus, unlike the barycenter the generalized geodesic does yield a notion of convexity satisfied by the Wasserstein distance, and is also easier to compute.
For these reasons, we adopt this notion of interpolation for our approach.
It remains to discuss how to apply it on (labeled) datasets.
2.2 Dataset distance
Consider a dataset D P
= { z ( i ) } Ni
=
= { x ( i ) , y ( i ) } Ni =1 i.i.d.
∼ P ( x, y ) .
The Optimal Transport Dataset
Distance (OTDD) (Alvarez-Melis & Fusi, 2020) measures its distance to another dataset D Q as:
[Equation]
which defines a proper metric between datasets.
Here, α y , α y ′ are class-conditional measures corresponding to P ( x | y ) and Q ( x | y ′ ) .
This distance is strongly correlated with transfer learning performance, i.e., the accuracy achieved when training a model on D P and then fine-tuning and evaluating on D P .
Therefore, it can be used to select pretraining datasets for a given target domain.
Henceforth we abuse the notation P to represent both a dataset and its underlying distribution for simplicity.
To avoid confusion, we use ν and µ to represent distributions in the feature space, which is Euclidean space, and use P, Q to represent distributions in the product space of features and labels.
Method and algorithm
Our method consists of two steps: estimating optimal transport maps between the target dataset and all training datasets (Sec. 3.1), and using them to generate a convex combinations of these datasets by interpolating along generalized geodesics (Sec. 3.2).
For some downstream applications we will additionally project the target dataset into the ‘convex hull’ of the training datasets (Sec. 3.3).
3.1 Solving optimal map between labelled datasets
The OTDD is a special case of Wasserstein distance, so it is natural to consider the alternative Monge (map-based) formulation to (2).
We propose two methods to approximate the OTDD map, one using the entropy-regularized OT and another one based on neural OT.
OTDD barycentric projection.
Barycentric projections (Ambrosio et al., 2008; Pooladian &
Niles-Weed, 2021) can be efficiently computed for entropic regularized OT using the Sinkhorn algorithm (Sinkhorn, 1967).
Assume that we have i.i.d.
samples X ν = ( x (1) ν , . .
, x ( N ν ) ν ) , X µ = ( x (1) µ , . . .
, x ( N µ ) µ ) from two distributions ν and µ separately.
After solving the optimal coupling π ∗ :
= min π ∈ Π( ν,µ ) (cid:82) ∥ x − x ′ ∥
22 d π ( x, x ′ ) , the barycentric projection can be expressed as T B ( X ν ) =
N ν π ∗ X µ .
We extend the method to two datasets Z Q = { X Q , Y Q } , Z P = { X P , Y P } , where we have additional i.i.d.
label data Y Q =
( y (1) Q , . . .
, y ( N Q ) Q ) , Y P =
( y (1) P , . . .
, y ( N P ) P ) .
We first solve the optimal coupling π ∗ for OTDD (2) following the regularized scheme in Alvarez-Melis & Fusi (2020), and represent labels as one-hot vectors y ∈ R C .
The barycentric projection is divided into two parts:
[Equation]
However, this approach has two important limitations: it can not naturally map out-of-sample data and it does not scale well to large datasets (due to the quadratic dependency on sample size).
OTDD neural map.
Inspired by recent approaches to estimate Monge maps using neural networks (Rout et al., 2022; Fan et al., 2021), we design a similar framework for the OTDD setting.
Fan et al.
Gazdieva et al.
(2022) approach the Monge OT problem with general cost functions by solving its max-min dual problem sup f inf T (cid:82)
[ c ( x, T ( x )) − f ( T ( x ))]
d ν ( x ) + (cid:82) f
( x ′ )
d µ ( x ′ ) .
We extend this method to the distributions involving labels by introducing an additional classifier in the map.
Given two datasets P, Q , we parameterize the map T N : R d × R C Q → R d × R C P as
[Equation]
where G ( · ) : R d → R d is the pushforward feature map, and the ℓ ( · ) : R d → R C P is a frozen classifier that is pre-trained on the dataset P .
Notice that, with the cost c ( z, T ( z )) =
∥ x − G ( z ) ∥
22 + W 22 ( α y , α ¯ y ) , the Monge formulation of OTDD (2) reads inf T♯Q
= P (cid:82) ∥ x − G ( z ) ∥
+
W 22 ( α y , α ¯ y )
d Q ( z ) .
We therefore propose to solve the max-min dual problem
[Equation]
Implementation details are provided in Appendix D. Compared to previous conditional Monge map solvers (Bunne et al., 2022a; Asadulaev et al., 2022), the two methods proposed here: (i) do not assume class overlap across datasets, allowing for maps between datasets with different label sets; (ii) are invariant to class permutation and re-labeling; (iii) do not force no one-to-one class alignments, e.g., samples can be mapped across similar classes.
3.2 Convex combination in dataset space
Computing generalized geodesics requires constructing convex combinations of data points from different datasets.
Given a weight vector a ∈ R m , features can be naturally combined as x a = (cid:80) mi =1 a i x i .
But combining labels is not as simple because: (i) we allow for datasets with different number of labels, so adding them directly is not possible; (ii) we do not assume different datasets have the same label sets, e.g. MNIST (digits) vs CIFAR10 (objects).
Our solution is to represent all labels in the same dimensional space by padding them with zeros in all entries corresponding to other datasets.
As an example, consider three datasets with 2 , 3 , and 4 classes respectively.
Given a label vector y 1 ∈ R 2 for the first one, we embed it into R 9 as ˜ y
= [ y 1 ; 0 3 ; 0 4 ] ⊤ .
Defining ˜ y 2 , ˜ yanalogously, we compute their combination as y a = a 1 ˜ y 1 + a 2 ˜ y
+ a 3 ˜ y 3 .
This representation is loss-less and preserves the distinction of labels across datasets.
3.3 Projection onto generalized geodesic of datasets
We finally put together the components in Sec 3.1 and 3.2 to construct generalized geodesics between datasets in two steps.
First, we compute OTDD maps T ∗ i between Q and all other datasets
P i , i = 1 , . . .
, m using the discrete or neural OT approaches.
Then, for any interpolation vector a ∈ ∆ m we identify a dataset along the generalized geodesic via P a :=
( (cid:80) mi =
a i T ∗ i ) ♯Q .
By using the convex combination method in Sec.
3.2 for labeled data, we can efficiently sample from P a .
We next consider locating the dataset P ∗ a that minimizes the distance between P a and Q , i.e. the projection of Q onto the generalized geodesic.
We firstly approach this problem from a Euclidean viewpoint.
Suppose there are several distributions { µ i } mi =1 and an additional distribution ν on Euclidean space R d , Lemma 1 guarantees there exists a unique parameter a ∗ that minimizes W 22 ( ρ Ga , ν ) .
However, it is not straightforward to locate a ∗ because there is no closed-form formula of the map a (cid:55)→ W 22 ( ρ Ga , ν ) and it can be expensive to calculate W 22 ( ρ Ga , ν ) for all possible a .
To solve this problem, we resort to another transport distance: (2, ν )
-transport metric.
Definition 1 (Craig (2016)) .
The (2, ν )-transport metric is given by W 2 ,ν ( µ i , µ j ) :=
(cid:0)(cid:82)
∥
T ∗ i ( x ) − T ∗ j ( x ) ∥
22 d ν ( x ) (cid:1) 1 / 2 , where T ∗ i is the optimal map from ν to µ i .
When ν has density with respect to Lebesgue measure, then W 2 ,ν is a valid metric (Craig, 2016, Prop.
Moreover, we can derive the closed-form formula of the map a (cid:55)→ W 22 ,ν
( ρ Ga , ν ) .
Proposition 1.
W 2 2 ,ν ( ρ G a , ν ) =
(cid:80)
[Equation]
=
[Equation]
W 2 2 ,ν ( µ
[Equation]
, ν ) −
[Equation]
2 (cid:80)
[Equation]
̸ =
[Equation]
a
[Equation]
a
[Equation]
W 2 2 ,ν ( µ i , µ j ) .
This equation implies that given distributions { µ i } , ν in Euclidean space, we can trivially solve the optimal a ∗ that minimizes W 22 ,ν ( ρ Ga , ν ) by a quadratic programming solver 2 .
The proof (Appendix C) relies on Brenier’s theorem.
Inspired by this, we also define a transport metric for datasets:
[Figure]
Definition 2.
The squared (2, Q )-dataset distance is given by W 22 ,Q ( P i , P j ) :
= (cid:82) (cid:0) ∥ x
i − x j ∥ 22 + W 22 ( α y i , α y j ) (cid:1)
d Q ( z ) , where [ x i ; y i ]
= T ∗ i ( z ) and T ∗ i is the OTDD map from Q to P i .
Denote P 2 ,Q ( X × P ( X )) as the set of all probability measures P that satisfy d OT ( P, Q ) < ∞ and the OTDD map from Q to P exists.
The following result shows that (2, Q )-dataset distance is a proper distance.
The proof is also left to the Appendix C.
Proposition 2.
W 2 ,Q is a valid metric on P 2 ,Q
( X × P ( X )) .
Unfortunately, in this case W 22 ,Q ( P i , P j ) does not have an analytic form like before because Brenier’s theorem may not hold for a general transport cost problem.
However, we still borrow this idea and define an approximated projection (cid:98) P a as the minimizer of function
[Equation] [Equation]
which is an analog of Proposition 1.
Unlike the Wasserstein distance, W 22 ,Q ( · , · ) is easier to compute because it does not involve optimization, so it is relatively cheap the locate the minimizer of
W 2 ( P a , Q ) .
Experimentally, we observe that W 22 ,Q ( P a , Q ) is predictive of model transferability across tasks.
Figure 1(a) illustrates this projection on toy 3D datasets, color-coded by class.
Experiments 4.1 Learning OTDD maps on synthetic datasets
Figure 1(b) illustrates the role of the optimal map in estimating the projection of a dataset into the generalized geodesic hull of three others.
Using maps T ∗ i estimated via barycentric projection (3) results in a better preservation of class structure, whereas using non-optimal maps T i based on random couplings (as the usual mixup does) destroys class structure.
4.2 Transfer learning
Next, we use our framework to generate new pretraining datasets for few-shot learning.
Given m labeled pretraining datasets { P i } , we consider a few-shot test dataset, in which only partial data is labelled, e.g. 5 samples per class.
Suppose the training resource and time are both limited such that the user can choose only one dataset to train the model, in the mean time, the user expects the model to have the best generalization ability.
To this end, we assume the training dataset is chosen from the generalized geodesic { P a } .
With a choice of the one-hot weight vector a , P a recovers the original dataset P i for some i .
Otherwise, P a will be the interpolation of datasets { P i } .
We first show that the generalization ability of training models has a strong correlation with the distance W 22 ,
Q ( P a , Q ) .
Then we compare our framework with several baseline methods.
Connection to generalization.
The closed-form expression of W 22 ,ν ( ρ Ga , ν ) (Prop. 1) provides the distance between a base distribution ν
and the distribution along generalized geodesic ρ Ga in
Euclidean space.
We study its analog (5) for labelled datasets Q and { P i } in Figure 2.
To investigate
[Table]
the generalization abilities of models trained on different datasets, we discretize the simplex ∆to obtain 36 interpolation parameters a , and train a 5-layer LeNet classifier on each P a .
Then we fine-tune all of these classifiers on the few-shot test dataset Q with only 20 samples per each class.
We control the same number of training iterations and fine-tuning iterations across all experiments.
We fix the same colorbar range for all heatmaps across datasets to highlight the different impact of choosing training dataset.
For some test datasets, the choice of training dataset can affect the finetuning accuracy greatly.
For example, when Q is EMNIST and the training dataset is FMNIST, the fine-tuning accuracy is only ∼ 60% , but this can be improved to > 70% by choosing an interpolated dataset closer to MNIST.
This is reasonable because MNIST shares more similarity with EMNSIT than FMNIST or USPS.
To some test datasets like FMNIST and KMNIST, this difference is not so obvious because all training datasets are all far away from the test dataset.
Comparison with baselines.
Next, we compare our method with several baseline methods on
NIST datasets.
In each set of experiment, we select one dataset as the test dataset, and the rest NIST datasets are the training datasets.
We assume the test dataset is 5-shot, and to do this, we randomly choose 5 samples per class to be the labeled data, and treat the remaining samples as unlabeled.
Our method firstly trains a model on (cid:98) P a , and fine-tune the model on 5-shot test dataset.
To obtain (cid:98) P a , we use barycentric projection or neural map to approximate the OTDD maps from test dataset to the training datasets.
Our results are shown in the first two rows in Table 1.
The first baseline method is to create a synthetic dataset as training dataset by Mixup among datasets.
We randomly sample data from each training datasets, and do the convex combination of them with weight ˆ a .
We use the same convex combination method in Sec. 3.2, thus this baseline is equivalent to our framework with suboptimal OTDD maps.
The other two baselines (the bottom block in Table 1) skip the transfer learning part, and directly train the model or solve 1-NN on the few-shot test dataset.
Overall, transfer learning can bring additional knowledge from other domains and improve the test accuracy by at most 21 % .
Among the methods in the first block, training on datasets generated by OTDD barycentric projection outperforms others except USPS dataset, where the difference is only about 2.6 % .
A Related work
Mixup and related In-Domain Interpolation Generating training data through convex combinations was popularized by mixup (Zhang et al., 2018): a simple data augmentation technique that interpolates features and labels between pairs of points.
This and other works based on it (Zhang et al., 2021; Chuang & Mroueh, 2021) use mixup to improve in-domain model robustness and generalization by increasing in-distribution diversity of the training data.
Although sharing some intuitive principles with mixup, our method interpolates entire datasets —rather than individual datapoints— with the goal of improving across-distribution diversity and therefore out-of-domain generalization.
Dataset synthesis in machine learning Generating data beyond what is provided as training dataset is a crucial component of machine learning in practice.
Basic transformations such as rotations, cropping, and pixel transformations can be found in most state-of-the-art computer vision models.
More recently, Generative Adversarial Nets (GAN) have been used to generate synthetic data in various contexts (Bowles et al., 2018; Yoon et al., 2019), a technique that has proven particularly successful in the medical imaging domain (Sandfort et al., 2019).
Since GANs are trained to replicate the dataset on which they are trained, these approaches are typically confined to generate in-distribution diversity, and typically operate on features only.
Discrete OT, Neural OT, Gradient Flows Barycentric projection (Ambrosio et al., 2008; Perrot et al., 2016) is a typical effective method to approximate optimal transport map with discrete regularized OT.
Other than this, neural net based optimal map in Euclidean space has made great progress (Makkuva et al., 2020; Fan et al., 2021; Rout et al., 2022) recently and reveal its power in image generation (Rout et al., 2022), style transfer (Korotin et al., 2022b), etc.
However, the study of the optimal map between two datasets is relatively scarse.
Some conditional Monge map solvers (Asadulaev et al., 2022; Bunne et al., 2022a) utilize the label information in a semi-supervised manner, where they assume the label to label correspondence between two distributions is known.
Our dataset mapping is distinct from them because we do not enforce the label to label mapping.
Based on the optimal coupling or map, geodesics and interpolation in general metric spaces have been studied extensively in the optimal transport and metric geometry literatures (McCann, 1997; Agueh & Carlier, 2011;
Ambrosio et al., 2008; Santambrogio, 2015; Villani, 2008; Craig, 2016), albeit mostly in a theoretical setting.
Gradient flows (Santambrogio, 2015), as an alternative approach for interpolation between distributions, have become increasingly popular in machine learning to model existing processes (Bunne et al., 2022b;
Mokrov et al., 2021; Fan et al., 2022) or solving optimization problems over datasets (Alvarez-Melis & Fusi, 2021), but they are computationally heavy.
B Discussion
Complexity The complexity of solving OTDD barycentric projection by Sinkhorn algorithm is
O ( N 2 ) (Dvurechensky et al., 2018), where N is the number of data in both datasets.
This can be expensive for large-scale dataset.
In practice, we solve the batched barycentric projection, i.e. take a batch from source and target datasets and solve the projection from source batch to target batch, and we normally fix batch size B as 10 4 .
This reduces the complexity from O ( N 2 ) to O ( BN ) .
The complexity of solving OTDD neural map is O ( BKH ) , where K is number of iterations, and
H is the size of the network.
We always choose K
= O ( N ) in the experiments.
The complexity of solving all the (2 , Q ) -dataset distances in (5) is O ( m 2 N ) since we need to solve the dataset distance between each pair of training datasets.
Putting these pieces together, the complexity of approximating the interpolation parameter ˆ a for the minimizer of (5) is O ( N ( B + m 2 )) .
Limitation The generation of synthetic dataset relies on solving OTDD maps from test dataset to each training dataset.
These OTDD maps are tailored to the considered test dataset and can not be reused for a new test dataset.
Another limitation is our framework is based on model training and fine-tuning pipeline.
This can be resource demanding for large-scale models, like GPT model.
C Proofs
Proof of Lemma 1.
By Santambrogio (2017, Sec. 4.4), the result holds when m = 2 .
Then Proposition 7.5 in Agueh & Carlier (2011) extends the result to the case of m > 2 .
Proof of Proposition 1.
Since linear combination preserves cyclically monotonicity, (cid:80) mi =1 a i T ∗ i ( x ) is the optimal map from ν to ρ Ga (McCann, 1995).
Then according to the definition of W 2 ,ν ( · , · ) , we can write
[Equation]
For scalars p, q 1 , . . . , q m , it holds that
[Equation]
Plugging this equality into (6) gives
[Equation]
Proof of Proposition 2.
Firstly, W
[Equation]
Q is symmetric and nonnegative by definition.
It is nondegenerate since W 2 ,Q ( P i , P j ) ≥ d OT ( P i , P j ) and d OT is a metric.
Finally, we show it satisfies the triangular inequality.
Indeed,
[Equation]
where the first inequality is the triangular inequality and the second inequality is the Minkowski inequality.
D Implementation details of OTDD map
OTDD barycentric projection We use the implementation https://github.com/microsoft/ otdd to solve OTDD coupling.
The rest part is straightforward.
OTDD neural map To solve the problem (4), we parameterize f, G, ℓ to be three neural networks.
In NIST dataset experiments, we parameterize f as ResNet 3 from WGAN-QC (Liu et al., 2019), and take feature map G to be UNet 4 (Ronneberger et al., 2015).
We generate the labels ¯ y with a pre-trained classifier ℓ ( · ) , and use a LeNet or VGG-5 with Spinal layers 5 (Kabir et al., 2022) to parameterize ℓ ( · ) .
In 2D Gaussian mixture experiments, we use Residual MLP to represent all of them.
We remove the discriminator’s condition on label to simplify the loss function as
[Equation]
In this formula, we assume both y and ¯ y are hard labels, but in practice, the output of ℓ ( · ) is a soft label.
Simply taking the argmax to get a hard label can break the computational graph, so we replace the label loss W 22 ( α y , α ¯ y ) by y ⊤ M ¯ y , where M ∈ R C Q × C P is the label-to-label matrix where
M ( i, j ) :
= W 22 ( α y i , α y j ) , and y is the one-hot label from dataset Q .
The matrix M is precomputed before the training, and is frozen during the training.
We pre-train the feature map G to be identity map before the main adversarial training.
We use the
Exponential Moving Average 6 of the trained feature maps as the final feature map.
Data processing For all the NIST datasets, we rescale the images to size 32 × 32 , and repeat their channel 3 times and obtain 3-channel images.
We use the default train-test split from torchvision .
Hyperparameters For the experimental results in Sec.
4.2, we use the OTDD neural map and train them with learning rate 10 − 3 and batch size 64.
We train a LeNet for 2000 iterations, and fine-tune for 100 epochs.
Regard the comparison with other baselines in Sec.
4.2, for transfer learning methods, we train a SpinalNet for 10 4 iterations, and fine-tune it for 2000 iterations on test dataset.
Training from scratch on the test dataset takes also 2000 iterations.
E Additional results E.1 OTDD neural map visualization
In Figure 4, we in addition provide qualitative results of OTDD map from EMNIST (letter) (Cohen et al., 2017) dataset to all other *NIST dataset and USPS dataset.
At this point, we can confirm three traits of OTDD map, which are mentioned at the end of Sec.
In the first labelling method, all 0 MNIST digits are assigned as class "0", and they are labelled as class "7" in the bottom labelling.
1) We don’t assume a known source label to target label correspondence.
So we can map between two irrelevent datasets such as EMNIST and FashinMNIST.
The map is invariant to the permutation of label assignment.
For example, we show two different labelling in Figure 3, and the final OTDD map will be the same.
3) It doesn’t enforce the label to label mapping but would follow the feature similarity.
From
Figure 4 in the appendix, we notice many crossclass mapping behaviors.
For example, when the target domain is USPS (Hull, 1994) dataset, the lower-case letter "l" is always mapped to digit 1, and the capital letter "L" is mapped to other digits such as 6 or 0 because the map follows the feature similarity.
We also show the OTDD neural map between 2D Gaussian mixture models with 16 components in
Figure 5.
This example is very special so that we have the closed-form solution of OTDD map.
The feature map is a identity map and the pushforward label is equal to the corresponding class that has the same conditional distribution p ( x | y ) as source label.
For example, the sample from top left corner cluster is still mapped to the top left corner cluster, and the label is changed from blue to orange.
This map achieves zero transport cost.
Since the transport cost is always non-negative, this map is the optimal OTDD map.
However, Asadulaev et al.
; Bunne et al.
(2022a) enforce mapping to preserve the labels, so with their methods, the blue cluster would still map to the blue cluster.
Thus their feature map is highly non-convex and more difficult to learn.
We refer to Figure 5 in Asadulaev et al.
for their performance on the same example.
Compared with them, our pushforward dataset aligns with the target dataset better.
E.2 McCann’s interpolation between datasets
Our OTDD map can be extended to generate McCann’s interpolation between datasets.
We propose an anolog of McCann’s interpolation (1) in the dataset space.
We define McCann’s interpolation between datasets P 0 and P 1 as
[Equation]
where T ∗ is the optimal OTDD map from P 0 to P 1 and t is the interpolation parameter.
The superscript M of P Mt means McCann.
We use the same convex combination method in Sec.
3.2 to obtain samples from P Mt .
Assume ( x 0 , y 0 ) ∼ P 0 , ( x 1 , y 1 ) =
T ∗ ( x 0 , y 0 ) and P 0 , P 1 contain 7,classes respectively, i.e. y 0 ∈ R 7 , y 1 ∈ R 3 .
Then the combination of features is x t = (1 − t ) x 0 + tx 1 , and the combination of labels is
[Equation]
Thus ( x t , y t ) is a sample from ((1 − t )
Id + t T ∗ ) ♯P 0 .
We visualize McCann’s interpolation between two Gaussian mixture distributions in Figure 6.
This method can map the labelled data from a dataset to another, and do the interpolation between them.
Thus we can use it to map abundant data from an external dataset, to a scarce dataset for data augmentation.
For example, in Figure 7, the target dataset only has 30 samples, but the source dataset has 60000 samples.
We learn the OTDD neural map between them and solve their interpolation.
We find that P M 1 creates new data out of the domain of the original target distribution, which Mixup (Zhang et al., 2018) can not achieve.
Thus, the data from P Mt for t close to 1.0 can enrich the target dataset, and be potentially used in data augmentation for classification tasks.