Improving Out-of-Distribution Robustness via Selective Augmentation

1 INTRODUCTION . To deploy machine learning algorithms in real-world applications , we must pay attention to distribution shifts , i.e . when the test distribution is different from the training distribution , which substantially degrades model performance . In this paper , we refer this problem as out-of-distribution ( OOD ) generalization and specifically consider performance gaps caused by two kinds of distribution shifts : domain shifts and subpopulation shifts . In domain shifts , the test data is sampled from different domains than the training data , which requires the trained model to generalize well to test domains without seeing training data from those domains . Take health risk prediction as an example . We may want to train a model on patients from a few sampled hospitals and then deploy the model to a broader set of hospitals ( Koh et al. , 2021 ) . In subpopulation shifts , the proportions of subpopulations in the test distribution differ from the proportions in the training distribution . When subpopulation shift occur , models perform poorly when they falsely rely on spurious correlations , which may occur when some subpopulations are under-represented in the training set . For example , in financial risk prediction , a machine learning model trained on the entire population may associate the labels with demographic features ( e.g. , religion and race ) , making the model fail on the test set when such an association does not hold in reality . To improve model robustness under these two kinds of distribution shifts , methods for learning invariant representations have shown effectiveness in various applications . These methods learn features or prediction mechanisms that are invariant to different domains while still containing sufficient information for the targeted task ( Li et al. , 2018 ; Arjovsky et al. , 2019 ) . Concretely , some prior works learn invariant representations by aligning and regularizing the domain-specific representations ( Li et al. , 2018 ; Sun & Saenko , 2016 ) . Other works aim to find invariant representations by balancing the risk across domains using regularizers ( Arjovsky et al. , 2019 ; Krueger et al. , 2021 ; Rosenfeld et al. , 2021 ) , which further increases the dependency between the invariant representations and labels . However , designing regularizers that are widely suitable to datasets from diverse domains is especially challenging and insuitable regularizers may adversely limit the model ’ s expressive power , leading to inconsistent performance among various real-world datasets . For example , on the WILDS datasets ( Koh et al. , 2021 ) , invariant risk minimization ( IRM ) ( Arjovsky et al. , 2019 ) outperforms empirical risk minimization ( ERM ) on CivilComments , but fails to improve robustness on a variety of other datasets like Camelyon17 and RxRx1 . A similar phenomenon is also reflected in the performance of CORAL ( Sun & Saenko , 2016 ) . Instead of explicitly imposing regularization to learn invariant representations , we turn towards an implicit solution . Inspired by mixup ( Zhang et al. , 2018 ) , we aim to alleviate the effects of domainrelated spurious information through data interpolation , leading to a simple algorithm called LISA ( Learning Invariant Representations with Selective Augmentation ) . Concretely , LISA linearly interpolates the features for a pair of samples and applies the same interpolation strategy on the corresponding labels . Critically , the pairs are selectively chosen according to two sample selection strategies . In selection strategy I , LISA interpolates samples with the same label but from different domains , aiming to eliminate domain-related spurious correlations . In selection strategy II , LISA interpolates samples with the same domain but different labels , where the model should to ignore the domain information and generate different predicted values as the interpolation ratio changes . In this way , LISA encourages the model to learn domain-invariant predictors without explicitly constraining or regularizing the representation . The primary contributions of this paper are as follows : ( 1 ) We develop a method that tackles the problem of distribution shifts by canceling out the domain-related spurious correlations via data interpolation . ( 2 ) We conduct broad empirical experiments to evaluate the effectiveness of LISA on nine benchmark datasets from diverse domains . In these experiments , we make the following observations . First , we find that LISA consistently outperforms seven prior methods in addressing both domain shifts and subpopulation shifts . Second , we identify that the performance gains of LISA are indeed caused by canceling out domain-specific information and learning invariant representations , rather than simply involving more data via interpolation . Third , when the degree of distribution shift increases , LISA achieves more significant performance gains . ( 3 ) Finally , we provide theoretical analysis of the phenomena distilled from the empirical studies , where we provably demonstrate that LISA can achieve smaller worst-domain error compared with ERM and vanilla mixup . We also note that to the best of our knowledge , this is the first theoretical analysis of how mixup ( with or without the selection strategies ) affects mis-classification error . 2 PRELIMINARIES In this paper , we consider the setting where one predicts the label y ∈ Y based on the input feature x ∈ X . Given a parameter space Θ and a loss function ℓ , we are supposed to train a model fθ under the training distribution Ptr , where θ ∈ Θ . In empirical risk minimization ( ERM ) , assume the empirical distribution over training data is P̂tr , ERM optimizes the following objective : θ∗ : = argmin θ∈Θ E ( x , y ) ∼P̂ [ ℓ ( fθ ( x ) , y ) ] . ( 1 ) In a traditional machine learning setting , a test set , sampled from a test distribution Pts , is used to evaluate the generalization of the trained model θ∗ , where the test distribution is assumed to be the same as the training distribution , i.e. , P tr = P ts . In this paper , we are interested in the setting when distribution shift occurs , i.e. , P tr ̸= P ts . Specifically , follow Muandet et al . ( 2013 ) ; Albuquerque et al . ( 2019 ) ; Koh et al . ( 2021 ) , we regard the overall data distribution containing D = { 1 , . . . , D } domains and each domain d ∈ D is associated with a data distribution Pd over a set ( X , Y , d ) = { ( xi , yi , d ) } N d i=1 , where N d is the number of samples in domain d. Then , we formulate the training distribution as the mixture of D domains , i.e. , P tr = ∑ d∈D r tr d Pd , where { rtrd } denotes the mixture probabilities in training set . Here , the training domains are defined as Dtr = { d ∈ D|rtrd > 0 } . Similarly , the test distribution could be represented as P ts = ∑ d∈D r ts d Pd , where { rtsd } is the mixture probabilities in test set . The test domains are defined as Dts = { d ∈ D|rtsd > 0 } . In domain shifts , we investigate the problem that the test domains are disjoint from the training domains , i.e. , Dtr ∩ Dts = ∅ . In general , we assume the test domains share some common properties with the training domains . For example , in Camelyon17 ( Koh et al. , 2021 ) data , we train the model on some hospitals and test it in a new hospital . We evaluate the worst-domain or average performance of the classifier among all test domains . In subpopulation shifts , the test set have domains that have been seen in the training set , but with a different proportion of subpopulations , i.e. , Dts ⊆ Dtr but { rtsd } ≠ { rtrd } . Under this setting , follow Sagawa et al . ( 2020a ) , we specially consider group-based spurious correlations , where each group g ∈ G is defined to be associated with a domain d and a label y , i.e. , g = ( d , y ) . We assume the domain identification is spuriously correlated with the label . For example , we illustrate the Waterbirds dataset in Figure 1 , where the background d ( water or land ) is spuriously correlated with the label y ( waterbird or landbird ) . Based on the group definition , we evaluate the model via the worst test group error , i.e. , maxg E ( x , y ) ∼g [ ℓ0−1 ( fθ ( x ) , y ) ] , where ℓ0−1 represents the 0-1 loss . 3 LEARNING INVARIANT REPRESENTATIONS WITH SELECTIVE AUGMENTATION . This section presents LISA , a simple way to improve robustness to domain shifts or subpopulation shifts . The key idea behind LISA is to encourage the model to alleviate the effects of domain-related spurious correlations by selective data interpolation . Before detailing how to select interpolated samples , we first provide a general formulation for data interpolation . In LISA , we perform linear interpolation between training samples . Specifically , given samples ( xi , yi , di ) and ( xj , yj , dj ) drawn from domains di and dj , we apply mixup ( Zhang et al. , 2018 ) , a simple data interpolation strategy , separately on the input features and corresponding labels as : xmix = λxi + ( 1− λ ) xj , ymix = λyi + ( 1− λ ) yj , ( 2 ) where the interpolation ratio λ ∈ [ 0 , 1 ] is sampled from a Beta distribution Beta ( α , β ) and yi and yj are one-hot vectors for classification problem . Notice that the mixup approach in equation 2 can be replaced by CutMix ( Yun et al. , 2019 ) , which shows stronger empirical performance in visionbased applications . In text-based applications , we replace the feature interpolation in equation 2 with Manifold Mixup ( Verma et al. , 2019 ) , where the interpolation strategy is performed on the output representation of a pre-trained model , e.g. , the output of BERT ( Devlin et al. , 2019 ) . After obtaining the interpolated features and labels , we replace the original features and labels in ERM with the interpolated ones . Then , the optimization process in equation 1 is reformulated as : θ∗ : = argmin θ∈Θ E { ( xi , yi , di ) , ( xj , yj , dj ) } ∼P̂ [ ℓ ( fθ ( xmix ) , ymix ) ] . ( 3 ) Without additional selection strategies , vanilla mixup will regularize the model and reduce overfitting ( Zhang et al. , 2021b ) , allowing it to attain good in-distribution generalization . However , vanilla mixup may not be able to cancel out spurious correlations , causing the model to still fail at attaining good OOD generalization ( see empirical comparisons in Section 4.3 and theoretical discussion in Section 5 ) . In LISA , we instead adopt a new strategy where mixup is only applied across specific domains or groups , which leans towards learning invariant representations and thus better OOD performance . Specifically , the two kinds of sample selection strategies are presented as follows : • Selection Strategy I : Interpolating samples with the same label . In selection strategy I , LISA interpolates samples with the same label but different domains ( i.e. , di ̸= dj , yi = yj ) . This produces datapoints that have both domains partially present , effectively eliminating spurious correlations between domain and label in cases where the pair of domains correlate differently with the label . Additionally , if domain information does not fully reflect the spurious correlations in some datasets , we can also enlarge the interpolation scope to cover more potentially spurious correlations by only interpolating samples within the same class regardless domain information ( i.e. , yi = yj ) . As a result , LISA with selection strategy I should learn domain-invariant representations for each class and thus achieve better OOD robustness . • Selection Strategy II : Interpolating samples with the same domain . Supposing domain information is highly spuriously correlated with the label information , selection strategy II applies the interpolation strategy on samples with the same domain but different labels , i.e. , di = dj , yi ̸= yj . Intuitively , even within the same domain , the model is supposed to generate different predicted labels since the interpolation ratio λ is randomly sampled , corresponding to different labels ymix . This causes the model to make predictions that are less dependent on the domain , again improving OOD robustness . In LISA , we randomly perform selection strategies I or II during the training process with probability psel and 1−psel for each batch of data , where psel is treated as a hyperparameter in our experiments . The choice of psel depends on the number of domains and the relation between domain information and spurious correlations . Empirically , using selection strategy I only brings much more benefits when there are more domains , and/or the domain information can not fully reflect the spurious correlations . LISA with selection strategy II can benefit the performance when domain information is highly spuriously correlated with the label , where we find a balanced ratio ( i.e. , psel = 0.5 ) performs the best . The pseudocode of the training procedure of LISA is shown in Algorithm 1 . Algorithm 1 Training Procedure of LISA Require : Training data , Step size η 1 : while not converge do 2 : Sample λ ∼ Beta ( α , β ) 3 : Sample a set of samples B1 uniformly from the training data 4 : Randomly select a sample selection strategy I or II with the probability psel and 1− psel 5 : if use selection strategy I then 6 : For each sample ( xi , yi , di ) , find another one ( xj , yj , dj ) from the dataset with the same label ( yi = yj ) but different domains ( di ̸= dj ) , and construct set B2 . 7 : else if use selection strategy II then 8 : For each sample ( xi , yi , di ) , find another one ( xj , yj , dj ) from the same domain ( di = dj ) but different labels ( yi ̸= yj ) , constructing set B2 . 9 : Update θ with data λB1 + ( 1− λ ) B2 .

Summary: This paper propose a mixup-style data augmentation method under the data distribution shift context. In particular, data distributions are formulated as mixture of distributions (i.e., domains), and two distribution shift scenarios are considered: (1) domain shift, where the test domain and train domain are disjoint. (2) subpopulation shift, where test distribution has different mixture proportion than train distribution. It's  assumed that domain identification spuriously correlates with labels. To tackle this problem, this paper proposes two mixup strategies: (I) mixup two examples with same label but different domains; (II) mixup two examples with same domain but different labels. It's claimed that such mixup could cancel out the spurious correlations. Extensive experiments on a variety of datasets show its superiority compared to empirical risk minimization (ERM) and alternative data augmentation methods. The paper further provide theoretical analysis that under certain conditions, the proposed method has asymptotically smaller worst case classification errors than ERM and vanilla mixup.