Domain Generalization via Optimal Transport with Metric Similarity Learning

Metric learning aims to learn a discriminative feature embedding where similar samples are closer while different samples are further apart [Deng et al., 2020]. Hadsell et al. [2006] proposed the siamese network together with contrastive loss to guide the instances stay close with each other in the feature space if they have the same labels and push them apart vice-versa. Schroff et al. [2015] proposed the triplet loss aiming to learn a feature space where a positive pair has higher similarity than the negative pair when comparing by the same anchor with a given margin. Oh Song et al. [2016] showed that neither the contrastive loss nor triplet loss could efficiently explore the full pair-wise relations between instances under the mini-batch training setting. They further propose the lifted structure loss to fully utilize pair-wise relations across batches. However, it only choose equal number of positive pairs as negative ones randomly, and many informative pairs are discarded [Wang et al., 2019], which restricts the ability of finding the informative pairs.  Yi et al. [2014] proposed the binomial deviance loss which could measure the hard pairs. One remarkable work by Wang et al. [2019] combines the advantages both from lifted structure loss and binomial loss to leverage the pair-similarity. They proposed to leverage not only pair-similarities (positive or negative pairs with each other) but also self-similarity which enables the learner to collect and weight informative pairs (positive or negative pairs) under an iterative (mining and weighting) manner. For a pair of instances, the self-similarity is gained from itself. Such a multi-similarity has been shown could measure the similarity and could cluster the samplers more efficiently and accurately. In the context of domain generalization,  Dou et al. [2019] proposed to guide the learner to leverage from the local similarity in the semantic feature space, in which the authors argued may contain essential domain-independent general knowledge for domain generalization and adopt the constrative loss and triplet loss to encourage the clustering for solving this issue. Leveraging from the across-domain class similarity information can encourage the learner to extract robust semantic features that regardless of domains, which is an useful auxiliary information for the learner. If the learner could not separate the samples (from different source domains) with domain-independent class-specific cohesion and separation on the domain invariant feature space, it would still suffer from ambiguous decision boundaries. This ambiguous decision boundaries might still be sensitive to the unseen target domain. Matsuura & Harada [2020] implement unsupervised clustering on source domains and showed better classification performance. Our work is orthogonal to previous works, proposing to enforce more distinguishable invariant features space via Wasserstein adversarial training and encouraging to leverage from label similarity information for better classification boundary.

Following Redko et al. [2017] and Li et al. [2017a], suppose we have $m$ known source domains distributions $\{\mathcal{D}_{i}\}_{i=1}^{m}$, and $i^{th}$ domain contains $N_{i}$ labeled instances in total, denoted by $\{(\mathbf{x}^{(i)}_{j},y^{(i)}_{j})\}^{N_{i}}_{j=1}$, where $\mathbf{x}^{(i)}_{j}\in\mathbb{R}^{n}$ is the $j^{th}$ instance feature from the $i^{th}$ domain and $y^{(i)}_{j}\in\{1,\dots,K\}$ are the corresponding labels. For a hypothesis class $\mathcal{H}$, the expected source and target risk of a hypothesis $h\in\mathcal{H}$ over domain distribution $\mathcal{D}_{i}$ is the probabilities that $h$ wrongly predicts on the entire distribution $\mathcal{D}_{i}$: $\epsilon_{i}(h)=\mathbb{E}_{(\mathbf{x},y)\sim\mathcal{D}_{i}}\ell(h(\mathbf{x},y))$, where $\ell(\cdot)$ is the loss function. The empirical loss is also defined by: $\hat{\epsilon}_{i}(h)=\frac{1}{N_{i}}\sum_{j=1}^{N_{i}}\ell(h(\mathbf{x}_{j},y_{j}))$.

In the setting of domain generalization, we only have the access to the seen source domains $\mathcal{D}_{i}$ but have no information about the target domain. The learner is expected to extract the underlying invariant feature space across the source domains and generalize to a new target domain.

We follow Redko et al. [2017] and define $c:\mathbb{R}^{n}\times\mathbb{R}^{n}\to\mathbb{R}^{+}$ as the cost function for transporting one unit of mass $\mathbf{x}$ to $\mathbf{x}^{\prime}$, then the primal form of the Wasserstein distance between $\mathcal{D}_{i}$ and $\mathcal{D}_{j}$ could be computed by,

where $\Pi(\mathcal{D}_{i},\mathcal{D}_{j})$ is the probability coupling on $\mathbb{R}^{n}\times\mathbb{R}^{n}$ with marginals $\mathcal{D}_{i}$ and $\mathcal{D}_{j}$ referring to all the possible coupling functions. Throughout this paper, we adopt Wasserstein-1 distance only ($p=1$).

Computing the primal form of Wasserstein distance (Eq. 1) is computational inefficiently. Assuming $|\mathcal{D}_{i}|=n,|\mathcal{D}_{j}|=m$, the time complexity for directly computing Eq. 1 is $\mathcal{O}(n^{3}+m^{3})$. On the contrary, leveraging the Kantorovich-Rubinstein duality [Wainwright, 2019] of Wasserstein distances could help to get a more efficient approximation. Assume $f$ a $1$-Lipschitz-continuous $w.r.t.$ the cost function: $\|f(x)-f(x^{\prime})\|\leq c(x,x^{\prime})$, we can prove that for any function $f$,

The equality arrives when $f$ reaches the maximum of the right side,

Optimal transport theory and Wasserstein distance were recently investigated in the context of machine learning [Arjovsky et al., 2017] especially in the domain adaptation area [Courty et al., 2016, Zhou et al., 2020]. The general idea of implementing the optimal transport technique for domain generalization across domains is illustrated in Fig. 2. To learn domain invariant features, OT technique is implemented to achieve domain alignments for extracting invariant features. After the OT transition, the invariant features can be generalized to unseen domain.

For a pair of instances $(\mathbf{x}_{i},y_{i})$ and $(\mathbf{x}_{j},y_{j})$, the notion of positive pairs usually refers to the condition where pair $i,j$ have same labels ($y_{i}=y_{j}$), while the negative pairs usually refers to the condition $y_{i}\neq y_{j}$. The central idea of metric learning is to encourage a pair of instances who have the same labels to be closer, and push negative pairs to be apart from each other [Wu et al., 2017].

Follow the framework of Wang et al. [2019], we show the general pair-weighting process of metric learning. Assuming the feature extractor $f$ parameterized by $\boldsymbol{\theta}_{f}$ projects the instance $\mathbf{x}\in\mathbb{R}^{n}$ to a $d$-dimensional normalized space: $f(\mathbf{x};\boldsymbol{\theta}_{f}):\;\mathbb{R}^{n}\to[0,1]^{d}$. Then, for two samples $\mathbf{x}_{i}$ and $\mathbf{x}_{j}$, the similarity between them could be defined as the inner product of the corresponding feature vector:

To leverage the across-domain class similarity information can encourage the learner to extract the classification boundary that regardless of domains, which is an useful auxiliary information for the learner. We further elaborate it in section 4.2.

As optimal transport could constrain labelled source samples of the same class to remain close during the transportation process [Courty et al., 2016]. We deploy optimal transport with Wasserstein distance [Redko et al., 2017, Shen et al., 2018] for aligning the marginal feature distribution over all the source domains.

A brief workflow of the optimal transport for a pair of sourcce domains is illustrated in Fig. 3(b). The critic function $D$ estimates the empirical Wasserstein Distance between the each source domain through a pair of instances from the empirical sets $\mathbf{x}^{(i)}\in\mathbf{X}^{(i)}$ and $\mathbf{x}^{(j)}\in\mathbf{X}^{(j)}$. In practice [Shen et al., 2018], the dual term Eq. 2 of Wasserstein distance could be computed by,

As in domain generalization setting, there usually exists more that two source domains, we can sum all the empirical Wasserstein distance between each pair of source domains,

Throughout this pair-wise optimal transport process, the learner could extract a domain-invariant feature space, we then propose to apply metric learning approaches to leverage the class label similarity for domain independent clustering feature extraction. We then introduce the metric learning for domain agnostic clustering in the next section.

As aforementioned, only aligning the marginal features via adversarial training is not sufficient for DG since there may exist a ambiguous decision boundary [Dou et al., 2019]. When predicting on the target domain, the learner may still suffer from this ambiguous decision boundary. To this end, we adopt the metric learning techniques [Wang et al., 2019] to help cluster the instances and promote a better prediction boundary for better generalization.

To solve this, except to the supervised source classification and alignment of the marginal distribution across domains with the Wasserstein adversarial training defined above, we then further encourage robust domain-independent local clustering via leverage from label information using the metric learning objective. The brief workflow is illustrated in Fig. 3(c). Specifically, we adopt the metric learning objective to require the images regardless of their domains to follow the two aspects: 1) images from the same class are semantically similar, thereby should be mapped nearby in the embedding space (semantic clustering), while 2) instances from different classes should be mapped apart from each other in embedding space. Since goal of domain generalization aims to learn to hypothesis could predict well on all the domains, the clustering should also be achieved under a domain-agnostic manner.

To this end, we mix the instances from all the source domains together and encourage the clustering for domain agnostic features via the metric learning techniques to achieve a domain-independent clustering decision boundary. For this, during each training iteration, for a batch $\{\mathbf{x}_{1}^{(i)},y_{1}^{(i)},\dots,\mathbf{x}^{(i)}_{b},y^{(i)}_{b}\}_{i=1}^{m}$ from $m$ source domains with batch size $b$, we mix all the instances from each domain and denoted by $\{(\mathbf{x}_{i}^{B},y_{i}^{B})\}_{i=1}^{m^{\prime}}$ with total size $m^{\prime}$. We first measure the relative similarity between the negative and positive pairs, which is introduced in the next sub-section.

In order to evaluate our proposed approach, we implement experiments on three common used datasest: VLCS [Torralba & Efros, 2011], PACS [Li et al., 2017a] and Office-home [Venkateswara et al., 2017] dataset. The VLCS dataset contains images from 4 different domains: PASCAL VOC2007 (V), LabelMe (L), Caltech (C), and SUN09 (S). Each domain includes five classes: bird, car, chair, dog and person. PACS dataset is a recent benchmark dataset for domain generalization. It consists of four domains: Photo (P), Art painting (A), Cartoon (C), Sketch (S), with objects from seven classes: dog, elephant, giraffe, guitar, house, horse, person. Office-Home is a more challenging dataset, which contains four different domains: Art (Ar), Clipart (Cl), Product (Pr) and Real World (Rw), with $65$ categories in each domain. Previous work showed that matter the adversarial model is trained under supervised [Long et al., 2017], semi-supervised [Zhou et al., 2020] or unsupervised [Long et al., 2018] way, the model will suffer from learning the diverse feature. To test our domain generalization model on this dataset could also help to affirm the effectiveness of our approach.

To show the effectiveness of our proposed approach, we compare our algorithm on the benchmark datasets with the following recent domain generalization methods.

• 1.

  Deep All: We follow the standard evaluation protocol of Domain Generalization to set up the pre-trained Alexnet or ResNet-18 fine-tuned on the aggregation of all source domains with only the classification loss.

• 2.

  TF [Li et al., 2017b]: A low-rank parameterized Convolution Neural Network model which aims to reduce the total number of model parameters for an end-to-end Domain Generalization training.

• 3.

  CIDDG [Li et al., 2018c]: Matches the conditional distribution by change the class prior.

• 4.

  MLDG [Li et al., 2018a]: The meta-learning approach for domain generalization. It runs the meta-optimization on simulated meta-train/ meta-test sets with domain shift

• 5.

  CCSA [Motiian et al., 2017]: The contrastive semantic alignment loss was adopted together with the source classification loss function for both the domain adaptation and domain generalization problem.

• 6.

  MMD-AAE [Li et al., 2018b]: The Adversarial Autoencoder model was adopted together with the Mean-Max Discrepancy to extract a domain invariant feature for generalization.

• 7.

  D-SAM [D’Innocente & Caputo, 2018]: It aggregates domain-specific modules and merges general and specific information together for generalization.

• 8.

  JiGen [Carlucci et al., 2019]: It achieves domain generalization by solving the Jigsaw puzzle via the unsupervised task.

• 9.

  MASF [Dou et al., 2019]: A meta-learning style method which based on MLDG and combined with Consitrastive Loss/ Triplet Loss to encourage domain-independent semantic feature space.

• 10.

  MMLD [Matsuura & Harada, 2020]: An approach that mixes all the source domains by assigning a pseudo domain label for extract domain-independent cluster feature space.

Following the general evaluation protocol of domain generalization ($e.g.$ Dou et al. [2019], Matsuura & Harada [2020]), on PACS and VLCS dataset. We first test our algorithm on by using AlexNet [Krizhevsky et al., 2012] backbones by removing the last layer as feature extractor. For preparing the dataset, we follow the train/val./test split and the data pre-processing protocol of Matsuura & Harada [2020]. As for the classifier, we initialize a three-layers MLP whose input has the same number of inputs as the feature extractor’s output and to have the same number of outputs as the number of object categories (2048-256-256-$K$), where $K$ is the number of classes. For the critic network, we also adopt a three-layers MLP (2048-1024-1024-1). For the metric learning objective, we use the output of the second layer of classifier network (with size 256) for computing the similarity.

For stable training, we set coefficient $\lambda_{d}=\frac{2}{1+\exp(-10p)}-1$ to regularize the adversarial loss, where $p$ is the training progress, to regularize the adversarial loss. This regularization scheme $\lambda_{d}$ has been widely used in adversarial training based domain adaptation and generalization setting ($e.g.$ [Long et al., 2017, Wen et al., 2019, Matsuura & Harada, 2020]) and have been proved could help to stabilize the training process. For the setting of $\lambda_{s}$, we follow the setting of [Dou et al., 2019] and set the value to $10^{-4}$. In our preliminary validation results, the performance is not sensitive with $\lambda_{d}\in[0,1]$. We also tried to range $\lambda_{s}$ from $10^{-3}$ to $10^{-6}$ via reverse validation and didn’t observe obvious differences.

Then, we examined our algorithm on the office-home benchmark, which is more challenging than the previous PACS and VLCS datasets. We follow the setting of [Carlucci et al., 2019], which is the most recent work who also evaluated on office-home dataset, to have a fair comparison. For this Office-home dataset, we also used reverse validation to set the learning rate as $2e-4$ for the whole model. For the remaining hyper-parameters, we keep the same with PACS and VLCS experiments. To avoid over-training, we also adopt the early stopping technique. All the experiments are programmed with PyTorch [Paszke et al., 2019].

We first reported the empirical results on PACS and VLCS dataset using AlexNet as feature extractor in Table 3 and Table 4, respectively. For each generalization task, we train the model on all the source domains and test on the target domain and report the average of top 5 accuracy values. The empirical results refers to the average accuracy about training on source domains while testing on the target domain.

From the empirical results, we can observe our method outperforms the baselines both on the PACS and VLCS dataset, indicating an improvement on benchmark performances. This showed the effectiveness of our method. Then, we report the empirical results on Office-home dataset in Table 5. As stated before, Office-home is a more larger and challenging dataset contains more diverse features from $65$ different classes. To evaluate the performance on this dataset requires large amount of computational resources. Due to the limits, we follow the evaluation protocol of Carlucci et al. [2019] to report the empirical results. From those results, we could observe that our algorithm outperforms the previous Domain Generalization method, this also confirm the effectiveness of our proposed method.

To further show the effectiveness of our algorithm especially on more deep models, follow Dou et al. [2019], we also report the results of our algorithm by using ResNet-18 backbone on PACS dataset in Table 6. The ResNet-18 backbone, the output feature dim will be $512$. From the results, we could observe that our method could outperform the baselines on most generalization tasks and on average $+1.6\%$ accuracy improvement.

Then, we implement ablation studies on each component of our algorithm. We report the empirical results of ablation studies in Table 7, where we test the ablation studies on both the AlexNet backbone and ResNet-18 backbone. We compare the ablations by, (1) Deep All: Train the model using feature extractor on source domain datasets with classification loss only, that is, neither optimal transport nor metric learning techniques is adopted. (2) No $\mathcal{L}_{D}$: Train the model with classification loss and metric learning loss but without adversarial training component; (3) $\mathcal{L}_{MS}$ w.o. $w^{+}$: omit the positive weighting scheme in $\mathcal{L}_{MS}$ (4)  $\mathcal{L}_{MS}$ w.o. $w^{-}$ : omit the positive weighting scheme in $\mathcal{L}_{MS}$. (5) No $\mathcal{L}_{MS}$: Train the model with classification loss and adversarial loss but without metric learning component; (6) WADG-All: Train the model with full objective Eq. 14.

From the results, we could observe that one we omit the adversarial training, the accuracy would drop off rapidly ($\sim 3.5\%$ with AlexNet backbone and $\sim 5.8\%$ with ResNet-18 backbone). The contribution of the metric learning loss is relatively small compared with adversarial loss. Comparing the ablations  $\mathcal{L}_{MS}$ w.o. $w^{+}$ and  $\mathcal{L}_{MS}$ w.o. $w^{-}$, we could observe almost similar accuracy. This indicates that the positive and negative weighting scheme of the metric learning objective may have equivalent contribution. . Once we omit the metric learning loss, the performance will drop $\sim 2.1\%$ and $\sim 2.5\%$ with AlexNet and ResNet-18 backbone, respectively.

Then, to better understand the contribution of each component of our algorithm, the T-SNE visualization of the ablation studies of each components on PACS and VLCS dataset are represented in Fig. 4 for the generalization task of target domain Photo. and Fig. 5 for the generalization task of target domain Caltech, respectively. Since our goal is to not only align the feature distribution but also encourage a cohesion and separable boundary, in order to show the alignment and clustering performance, we report the T-SNE features of all the source domains and target domain to show the feature alignment and clustering across domains.
For PASC dataset, as we can see, the T-SNE features by Deep All could neither project the instances from different domains to align with each other nor cluster the features into groups. The T-SNE features by No $\mathcal{L}_{D}$ showed the metric learning loss could to some extent to cluster the features, but without the adversarial training, the features could not be aligned well. The T-SNE features by No $\mathcal{L}_{MS}$ showed that the adversarial training could help to align the features from different domains but could not have a good clustering performance. The T-SNE features by WADG-All showed that the full objective could help to not only align the features from different domains but also could cluster the features from different domains into several cluster groups, which confirms the effective of our algorithm.
As for the VLCS dataset, we could observe similar performance on the T-SNE on the VLCS dataset while the features are somehow overlap with each other. This is due to the features in Caltech domain is somehow easy to learn and predict. As also analyzed in [Li et al., 2017a], a supervised model on Caltech domain could achieved $\sim 100\%$ accuracy, which also confirms that the features in Caltech domain is easy to learn indicating the features might be more likely overlapping with each other. As we can see from Fig.5(d), the WADG method could help to separate the features with each other, which again confirms the effectiveness of our proposed method.