Cross-Domain Feature Augmentation for Domain Generalization

1 Introduction

Deep learning methods typically assume that training and testing data are independent and identically distributed. However, this assumption is often violated in the real world, leading to a decrease in model performance when faced with a different distribution [Torralba and Efros, 2011]. The field of domain generalization aims to mitigate this issue by learning a model from one or more distinct yet related training domains, with the goal of generalizing effectively to domains that have not been previously encountered. Studies suggest that the poor generalization on unseen distributions can be attributed to the failure of learning the invariance across different domains during the training phase [Muandet et al., 2013; Li et al., 2018]. To tackle this, research has focused on representation learning and data augmentation as key to learning invariance.

(b) DSU

(c) XDomainMix

Figure 1: Samples of images reconstructed from features produced by DSU [Li et al., 2022] and the proposed XDomainMix. The elephant reconstructed from XDomainMix's features shows a more complex background. The horse reconstructed from XDomainMix's features displays different characteristics such as a white belly. The top floor of the house generated by XDomainMix's features shows solid walls, instead of glass walls in the original. In contrast, images reconstructed from DSU's features have limited diversity and appear largely similar to the original images.

Invariant representation learning aims to align representation across domains [Shi et al., 2022], and learn invariant causal predictors [Arjovsky et al., 2019]. They usually impose regularization, which may result in a hard optimization problem [Yao et al., 2022a]. In contrast, data augmentation techniques propose to generate additional samples for the learning of invariance, and avoid the complexities in the regularization approach [Mancini et al., 2020]. Data augmentation can be generally classified into two types: input space and feature space augmentation. The former often encounters limitations due to a lack of diversity in the augmented data [Li et al., 2021b], while the latter offers more versatility and has yielded promising outcomes [Zhou et al., 2021].

Despite the versatility of feature space augmentation, existing methods such as MixStyle [Zhou et al., 2021] and DSU [Li et al., 2022] do not consider feature semantics during the augmentation process. Instead, they alter feature statistics which often leads to a limited range of diversity. This lack of diversity in the generated features motivates us to decompose the features according to feature semantics. We build on prior research which suggests that the features learned for each class can be viewed as a combination of class-specific and class-generic components [Chu et al., 2020]. The classspecific component carries information unique to a class, while the class-generic component carries information that is shared across classes. We observe that, even within the same class, features of samples from different domains can be distinguished, indicating that these features may contain domain-specific information. As such, we broaden our understanding of features to include domain-specific and domaingeneric components.

We introduce a method called XDomainMix that changes domain-specific components of a feature while preserving class-specific components. With this, the model is able to learn features that are not tied to specific domains, allowing it to make predictions based on features that are invariant across domains. Figure 1 shows samples of original images and reconstructed images based on existing feature augmentation technique (DSU) and the proposed XDomainMix. We visualize the augmented features using a pre-trained autoencoder [Huang and Belongie, 2017]. From the reconstructed images, we see that DSU's augmented features remain largely the same as that of the original image feature. On the other hand, the images reconstructed from the features obtained using XDomainMix have richer variety while preserving the salient features of the class.

Results of experiments on benchmark datasets show the superiority of XDomainMix for domain generalization. We quantitatively measure the invariance of learned representation and prediction to show that the models trained with XDomainMix's features are more invariant across domains compared to state-of-the-art feature augmentation methods. Our measurement of the divergence between original features and augmented features shows that XDomainMix results in more diverse augmentation.

2 RelatedWork

To learn invariance, existing domain generalization approaches can be categorized into representation learning methods and data augmentation methods. Works on learning invariant representation employ regularizers to align representations or gradients [Sun and Saenko, 2016; Li et al., 2020; Kim et al., 2021; Mahajan et al., 2021; Shi et al., 2022; Rame et al., 2022; Yao et al., 2022b] across different domains, enforce the optimal classifier on top of the representation space to be the same across all domains [Arjovsky et al., 2019; Ahuja et al., 2021], or uses distributionally robust optimization [Sagawa et al., 2020]. However, the use of regularization terms during learning of invariant representation could make the learning process more complex and potentially limit the model's expressive power.

Another approach is to employ data augmentation to learn invariance. Existing work that operates in the input space includes network-learned transformation [Zhou et al., 2020;

Li et al., 2021a], adversarial data augmentation [Volpi et al., 2018; Shankar et al., 2018], mixup [Mancini et al., 2020; Yao et al., 2022a], and Fourier-based transformation [Xu et al., 2021]. Each of these techniques manipulates the input data in different ways to create variations that help the model learn invariant features. However, the range of transformations that can be applied in the input space is often limited.

On the other hand, feature augmentation can offer more flexibility and potential for learning more effective invariant representations. Prior work has generated diverse distributions in the feature space by changing feature statistics [Zhou et al., 2021; Jeon et al., 2021; Wang et al., 2022; Li et al., 2022; Fan et al., 2023], adding noise [Li et al., 2021b], or mixing up features from different domains [Mancini et al., 2020; Qiao and Peng, 2021]. For example, MixStyle [Zhou et al., 2021] synthesizes new domains by mixing the feature statistics of two features. DSU [Li et al., 2022] extends the idea by modeling feature statistics as a probability distribution and using new feature statistics drawn from the distribution to augment features. In addition to generating diverse distributions, RSC [Huang et al., 2020] adopts a different approach by discarding the most activated features instead of generating diverse data. This encourages the model to use less-activated features that might be associated with labels relevant to data outside the domain.

In contrast to existing methods, our work carefully considers feature semantics by leveraging class-label information and domain information to augment features. This increases intra-class variability and helps the model to learn a broader understanding of each class, thus improving its ability to handle new, unseen data.

3 ProposedMethod

We consider the problem where we have a set of source domains DS = {S1, · · · , SN}, N > 1 and a set of unseen domains DU . Each source domain Si = {(x(i) j , y(i) j )}ni j=1 has a joint distribution on the input x and the label y. Domains in DU have distinct joint distributions from those of the domains in DS. We assume that all domains in DS and DU share the same label space but the class distribution across domains need not be the same. The goal is to learn a mapping g : x → y using the source domains in DS such that the error is minimized when g is applied to samples in DU .

In deep learning, g is typically realized as a composition of two functions: a feature extractor f : x → Z that maps input x to Z in the latent feature space, followed by a classifier c : Z → y that maps Z to the output label y. Ideally, f should extract features that are domain invariant yet retain class-specific information. The features of a given input can be decomposed into two components: class-specific and class-generic. The class-specific component consists of feature semantics that are strongly correlated with class labels, making them more important in discriminating a target class from other classes.

Furthermore, features can also be decomposed into domain-specific and domain-generic components. This is because samples from different domains, even if they belong to the same class, possess unique feature characteristics to their respective domains. We extend these notions to decompose the features extracted by f into four distinct components: class-specific domain-specific (Zc,d), class-specific domaingeneric (Zc,¬d), class-generic domain-specific (Z¬c,d), and class-generic domain-generic (Z¬c,¬d) component.

We determine whether an extracted feature contains information that is specific to a class or a domain by its importance to the prediction of the class and domain respectively. In other words, if a feature is important to the prediction of a specific class and a specific domain, it is considered a class-specific domain-specific component. If a feature is crucial in class prediction but not domain prediction, it falls into the classspecific domain-generic category. Similarly, features that are important for domain but not class predictions are categorized as class-generic domain-specific, and those not significant to either are labeled as class-generic domain-generic.

Our proposed feature augmentation strategy, XDomainMix, performs augmentation in the feature space by modifying the domain-specific component of features in a way that preserves class-related information. To discourage the use of domain-specific features for class prediction and encourage the exploitation of less activated features, class-specific domain-specific component is discarded with some probability during training. Details of feature decomposition and augmentation are described in the following subsections. Figure 2 shows an overview of our proposed method.

3.1 Feature Decomposition

Suppose the feature extractor f extracts Z = f(x) ∈ RK . Let zk be the kth dimension of Z. We determine whether zk is class-specific or class-generic via the class importance score, which is computed as the product of feature value and the derivative of class classifier c's predicted logit vc of x's ground truth class [Selvaraju et al., 2017; Chu et al., 2020] as they show how much zk contributes to vc:

To determine if zk is domain-specific, we employ a domain classifier d that has an identical architecture as the class classifier c. d is trained to predict domain labels of features extracted by the feature extractor. Similar to Equation 1, domain importance score is computed using the derivative of d's predicted logit vd of x's ground truth domain:

Let τc and τd be predefined thresholds for filtering feature dimensions that are class-specific and domain-specific respectively. We obtain a class-specific mask Mc ∈ RK and a domain-specific mask Md ∈ RK on Z for {zk }Kk=1 where their respective kth entries are given as follows:

Complementary class-generic mask and domain-generic mask are obtained by 1-Mc and 1-Md where 1 is the tensor of values 1 and of the same size as Z. Class-specific domainspecific (Zc,d), class-specific domain-generic (Zc,¬d), classgeneric domain-specific (Z¬c,d), and class-generic domaingeneric (Z¬c,¬d) feature components are obtained by

where ⊙ is element-wise multiplication. Note that Zc,d + Zc,¬d + Z¬c,d + Z¬c,¬d = Z.

3.2 Cross Domain Feature Augmentation

To achieve domain invariance and reduce reliance on domainspecific information presented in training domains during prediction, we manipulate domain-specific feature components to enhance diversity from a domain perspective. Further, the augmentation should increase feature diversity while preserving class semantics using existing data. This is achieved by mixing the class-specific domain-specific feature component of a sample with the class-specific domainspecific feature component of a same-class sample from other domains. For class-generic domain-specific feature component, it is mixed with the class-generic domain-specific feature component of a different-class sample from other domains, introducing further diversity.

Specifically, for the feature Z extracted from input x, we randomly sample two inputs xi and xj whose domains are different from x. Further, xi has the same class label as x while xj is from a different class. Let Zi be the feature extracted from input xi and Zj be the feature extracted from xj . Then we have

where λ1 and λ2 are the mixup ratios independently sampled from a uniform distribution U(0, 1).

Anew feature ˜ Z with the same class label as Z is generated by replacing the domain-specific component in Z as follows:

To further encourage the model to focus on domaininvariant features and exploit the less activated feature during class prediction, we discard the class-specific domainspecific feature component with some probability pdiscard as follows:

where p is randomly sampled from a uniform distribution U(0, 1).

3.3 Training Procedure

Prior research has shown that empirical risk minimization (ERM) [Vapnik, 1999] is a competitive baseline [Gulrajani and Lopez-Paz, 2021; Wiles et al., 2022]. The objective function of ERM is given by

where ℓ is the loss function to measure the error between the predicted class ˆ y(i) j and the ground truth y(i) j . N is the number of training domains and ni is the number of training samples in domain i.

We train the model in two phases. During the warm-up phase, the feature extractor f and class classifier c are trained on the original dataset for class label prediction following Lerm. The domain classifier d is trained using Z, the features extracted by f, to predict domain labels. The objective function of d is given by

where ℓ is the loss function to measure the error between the predicted domain d(Z (i) j ) and the ground truth i.

When the warm-up phase is completed, we use Equation 4 to decompose the features obtained from f. Augmented features are then derived using Equation 7. Both feature extractor f and class classifier c are trained using the original and augmented features with the following objective function:

where ˜ Z (i) j is the augmented feature derived from Z (i) j . c(Z (i) j ) is the predicted class given Z (i) j , and c( ˜ Z (i) j ) is the predicted class given ˜ Z (i) j . y(i) j is the ground truth class.

We also train the domain classifier d using Ld. Note that the domain classifier d is not trained using this set of augmented features, as the augmented features do not have assigned domain labels since they need not follow the distribution of existing domains.

4 Performance Study

We implement our proposed solution using PyTorch 1.12.0 and perform a series of experiments on NVIDIA Tesla V100 GPU to evaluate the effectiveness of the proposed XDomainMix. The following benchmark datasets are used:

• Camelyon17 [Bandi et al., 2018] fromWilds [Koh et al., 2021]. This dataset contains 455,954 tumor and normal tissue slide images from 5 hospitals (domains). Distribution shift arises from variations in patient population, slide staining, and image acquisition.
• FMoW [Christie et al., 2018] fromWilds. This dataset contains 141,696 satellite images from 62 land use categories across 16 years from 5 regions (domains).
• PACS [Li et al., 2017]. This dataset contains 9,991 images of 7 objects in 4 visual styles (domains): art painting, cartoon, photo, and sketch.
• TerraIncognita [Beery et al., 2018]. The dataset contains 24,788 images from 10 categories of wild animals taken from 4 different locations (domains).
• DomainNet [Peng et al., 2019]. This dataset contains 586,575 images from 365 classes in 6 visual styles (domains): clipart, infograph, painting, quickdraw, real, and sketch.

The class importance thresholds τc and domain importance thresholds τd in Equation 3 are set as follows: τc is set to be the 50%-quantile of the class importance scores of {zk } in a feature Z so that 50% dimensions are considered classspecific, while the remaining 50% are class-generic. τd controls the strength of the augmentation as it determines the identification of domain-specific feature components. We employ a cyclic changing scheme for τd to let the model learn gradually from weak augmentation to strong augmentation and give the domain classifier more time to adapt to a more domain-invariant feature extractor. The value is initially set to be 90%-quantile of domain importance scores of {zk } in a feature Z. As the training proceeds, τd is decreased by 10% quantile for every n step until it reaches the 50%-quantile, where it also remains for n steps. The same cycle is repeated where τd is set to be 90%-quantile again. Input xi, xj used in augmentation (Equation 5) are samples from the same training batch. pdiscard is set to 0.2.

For Camelyon17 and FMoW datasets, we follow the setup in LISA [Yao et al., 2022a]. Non-pretrained DenseNet-121 is used for Canmelyon17 and pretrained DenseNet-121 is used for FMoW. We use the same partitioning in Wilds [Koh et al., 2021] to obtain the training, validation, and test domains. The batch size is set to 32, and the model is trained for 2 epochs for Camelyon17 and 5 epochs for FMoW. The learning rate and weight decay are set to 1e-4 and 0. The warm-up phase is set to 4000 steps. We tune the step n in {100, 500} for changing τd. The best model is selected based on its performance in the validation domain.

For PACS, TerraIncognita and DomainNet datasets, we follow the setup in DomainBed [Gulrajani and Lopez-Paz, 2021], and use a pre-trained ResNet-50. Each domain in the dataset is used as a test domain in turn, with the remaining domains serving as training domains. The batch size is set to 32 (24 for DomainNet), and the model is trained for 5000 steps (15000 steps for DomainNet). We tune the learning rate in {2e-5, 3e-5, 4e-5, 5e-5, 6e-5} and weight decay in (1e-6, 1e-2) using the DomainBed framework. The warm-up phase is set to 3000 steps and n is set to 100 steps. The best model is selected based on its performance on the validations split of the training domains.

4.1 Domain Generalization Performance

We compare our proposed XDomainMix with ERM [Vapnik, 1999] and the following state-of-the-art methods:

• GroupDRO [Sagawa et al., 2020] minimizes worst-case loss for distributionally robust optimization.
• RSC [Huang et al., 2020] discards features that have higher activation to activate the remaining features appear to be applicable to out-of-domain data.
• MixStyle [Zhou et al., 2021] synthesizes new domains by mixing feature statistics of two features.
• DSU [Li et al., 2022] synthesizes new domains by renormalizing feature statistics of features with the ones drawn from a probability distribution.
• LISA [Yao et al., 2022a] selectively mixes up samples to learn an invariant predictor.
• Fish [Shi et al., 2022] aligns gradients across domains by maximizing the gradient inner product.

Classification accuracy, which is the ratio of the number of correct predictions to the total number of samples is reported. Following the instruction of datasets, average accuracy on the test domain over 10 runs is reported for the Camelyon17 dataset; worst-group accuracy on the test domain over 3 runs is reported for the FMoW dataset. For PACS, TerraIncognita and DomainNet datasets, the averaged accuracy of the test domains over 3 runs is reported.

Table 1 shows the results. Our method consistently achieves the highest average accuracy across all the datasets, outperforming SOTA methods. This result suggests that XDomainMix is able to train models with good domain generalization ability.

4.2 Model Invariance

One advantage of XDomainMix is that is able to learn invariance across training domains. We quantify the invariance in terms of representation invariance and predictions invariance. Representation invariance refers to the disparity between representations of the same class across different domains. The distance between second-order statistics (covariances) [Sun and Saenko, 2016] can be used to measure representation invariance. Prediction invariance considers the variation in predictions across different domains. We employ risk variance [Yao et al., 2022a] which measures how similar the model performs across domains.

Let {Z (i) j |y(i) j = yc} be the set of representations of class label yc from domain i. We use C (i) yc to denote the covariance matrix of the representations. Given the set of class labels Y and the set of training domains DS, the measurement result is given by 1 |Y||DS | ∑yc∈Y ∑i,i′∈DS ||C (i) yc -C (i′ ) yc ||2F . ||·||2F denotes the squared matrix Frobenius norm. A smaller distance suggests that same-class representations across domains are more similar.

Let Ri be the loss in predicting the class labels of inputs from domain i. The risk variance is given by the variance among training domains, Var{R1 ,R2 , ...,R|DS| }. Lower risk variance suggests a more consistent model performance across domains.

• (b) Prediction invariance measured by variance of risk across domains. The results are reported in the unit of 1e-3.

Table 3: Divergence of the augmented feature and original feature measured by MMD in the unit of 1e-2.

Table 2 shows the results. We see that our method has the smallest covariance distance in the FMoW, PACS, TerraIncognita, and DomainNet dataset, and the second-smallest in the Camelyon17 dataset. The results indicate that the representations of the same class learned by our method have the least divergence across domains. Additionally, XDomainMix has the lowest risk variance, suggesting that it is able to maintain consistent performance in predictions across domains. Overall, the results demonstrate that our approach is able to learn invariance at both the representation level and the prediction level.

4.3 Diversity of Augmented Features

To show that XDomainMix can generate more diverse features, we measure the distance between the original and augmented features using maximum mean discrepancy (MMD). A higher MMD suggests that the distance between the original and augmented features is further. The same set of original features is used to ensure fairness and comparability of the measurement result. Average and standard deviation over three runs are reported. Table 3 shows the results. Features augmented by XDomainMix consistently have the highest MMD compared to MixStyle and DSU which are two SOTA feature augmentation methods. This suggests that the features augmented by XDomainMix exhibit the most deviation from the original features, leading to a more varied augmentation. Visualization of sample images reconstructed from augmented features are given in Supplementary.

4.4 Experiments on the Identified Features

In this set of experiments, we demonstrate that XDomainMix is able to identify features that are important for class and domain prediction. We evaluate the model performance for class or domain prediction after eliminating those features with the highest importance score computed in Equations 1 and 2. A decrease in accuracy suggests that the features that have been removed are important for the predictions.

For comparison, we implement two alternative selection strategies: a random method that arbitrarily selects features to remove, and a gradient norm approach, where features are chosen for removal based on the magnitude of the gradient in the importance score computation. Samples in the validation set of PACS dataset are used in this experiment.

Figure 3 shows the results. Our method shows the largest drop in both class prediction and domain prediction accuracies compared to random removal and gradient norm methods. This indicates that XDomainMix is able to identify features that are specific to the domain and class effectively.

To visualize the extracted features, we map the high dimensional feature vectors obtained by f to a lower dimensional space, This transformation is carried out using two linear layers, as described in [Zhang et al., 2021]. Figure 4(a) provides a visualization for the model that is trained on the PACS dataset, with Art as the unseen domain. We see that the features identified as class-specific are well separated by class. This is in contrast to the features that are generic across classes, which are not as clearly delineated.

Figure 4: Visualization of features from different classes/domains, indicated by the different colors.

Similarly, we also visualize the extracted domain-specific and domain-invariant features. As shown in Figure 4(b), the domain-specific features are noticeably better separated compared to the domain-invariant features.

4.5 Ablation Study

To understand the contribution of each component in the augmentation, we perform ablation studies on Camelyon17 and FMoW datasets. Table 4 shows the result. Compared to the baseline, mixing class-specific domain-specific feature components (Zc,d) only, or mixing class-generic domain-specific feature components (Z¬c,d) only in the augmentation can improve the performance. This suggests that by manipulating domain-specific feature components, models that are better at domain generalization can be learned. Mixing Z¬c,d leads to greater improvement, indicating that enriching diversity by content from other classes is more helpful than simply intraclass augmentation.

Augmenting both Zc,d and Z¬c,d does not consistently lead to performance improvement, possibly due to dataset-specific characteristics. Probabilistically discarding Zc,d seems to encourage the model to use less domain-specific information and exploit less activated features in prediction, which improves the domain generalization performance.

5 Conclusion and FutureWork

In this work, we have developed a feature augmentation method to address the domain generalization problem. Our approach aims to enhance data diversity within the feature space for learning models that are invariant across domains by mixing domain-specific components of features from different domains while retaining class-related information. We have also probabilistically discarded domain-specific features to discourage the model from using such features for their predictions, thereby achieving good domain generalization performance. Our experiments on multiple datasets demonstrate the effectiveness of the proposed method.

While our method presents a promising approach to solving the domain generalization problem, there are several limitations. Our method needs more than one training domain to perform cross-domain feature augmentation. Our method assumes that the datasets across different domains share the same label space and similar class distributions, and its performance may be affected if this is not the case. Our method is mainly empirically validated, and a theoretical analysis or guarantee of its performance is still lacking. Further research is needed to provide a deeper theoretical understanding of the proposed method and its performance bounds.

Acknowledgements

This research/project is supported by the National Research Foundation, Singapore under its AI Singapore Programme (AISG Award No: AISG-GC-2019-001-2B). We thank Dr Wenjie Feng for the helpful discussions.