Domain-wise Adversarial Training for Out-of-Distribution Generalization

1 INTRODUCTION . Modern deep learning techniques have achieved remarkable success on many tasks ( He et al. , 2016 ; Brown et al. , 2020 ) . Yet , under some scenarios , deep models will suffer a catastrophic performance degradation since they tend to seize on spurious correlations in the training data ( Beery et al. , 2018 ) . One of those representative scenarios is the Out-of-Distribution ( OOD ) generalization , where one expects the trained model to perform well at the test time even when the training and testing data come from different distributions1 . Another representative scenario under which deep models are unstable is the adversarial example . Researchers have found that deep models are quite brittle since one can inject imperceptible perturbations into the input and cause the model to make wrong predictions with extremely high confidence ( Szegedy et al. , 2014 ) . These two issues have some similarities to each other . They both arise because deep networks do not learn the essential causal associations ( or intrinsic features ) . Nevertheless , in their corresponding fields , different approaches have been proposed . For OOD generalization , a large class of methods called Invariant Causal Prediction ( ICP ) ( Peters et al. , 2016 ) are proposed . Among them , Invariant Risk Minimization ( IRM ) ( Arjovsky et al. , 2019 ) attracts significant attention , which intends to extract features that are invariant across different data distributions and expects the model to ignore information related to the environment . While for adversarial robustness against adversarial examples , Adversarial Training ( AT ) ( Madry et al. , 2018 ) is the most effective approach at the current stage ( Athalye et al. , 2018 ) . It trains a model on adversarial examples that are generated by injecting perturbations optimized for each image into natural examples . These two fields seem to be independent , and their connections are rarely studied . That is exactly what we are exploring in this paper . Although IRM and its variants have shown promise on certain tasks , e.g. , CMNIST ( Arjovsky et al. , 2019 ) , recent studies ( Gulrajani & Lopez-Paz , 2021 ) show that on a large-scale controlled experiment on OOD generalization , all these methods fail to exceed the simplest i.i.d . baseline , i.e. , Empirical Risk Minimization ( ERM ) . Through a dissection of the benchmark datasets , Ye et al . ( 2021 ) 1In current literature , the terminologies of “ domain ” , “ environment ” and “ distribution ” are often used interchangeably , so we do the same in the whole paper . notice that there are actually two types of distribution shifts : correlation shift ( same support , different correlation ) and diversity shift ( different support , same correlation ) . IRM variants can only perform well under ( some ) correlation shift while performing poorly under diversity shift . Thus , we need to seek better alternatives for OOD generalization , while AT seems to be a promising candidate from both theoretical and empirical aspects . Theoretically , by learning invariance w.r.t . local input perturbations , AT can be regarded as Distributionally Robust Optimization ( DRO ) ( Sinha et al. , 2018 ; Volpi et al. , 2018 ; Rahimian & Mehrotra , 2019 ; Duchi et al. , 2021 ) over ` p-bounded distributional shift . Thus , AT could reliably extract robust features , e.g. , the object shape , from the input . Empirically , several recent works show that AT has better domain transferability than ERM ( Salman et al. , 2020 ) . These findings naturally leads to the questions : Will AT perform better than IRM ? Will AT be helpful for OOD generalization ? In this paper , we take a further step to answering these intriguing questions . We first reveal the connections between IRM and AT , and find that IRM can be regarded as an instance-reweighted version of Domain-wise Adversarial Training ( DAT ) , a new version of adversarial training that we propose for generalizing Universal Adversarial Training ( Moosavi-Dezfooli et al. , 2017 ) to multiple domains . Inspired by this connection , we further explore how DAT performs for OOD data . We first notice that DAT is suitable for solving domain generalization problems as it can effectively remove the relatively static background information with domain-wise perturbations . We further verify this intuition on both synthetic tasks ( Xiao et al. , 2021 ) and real-world datasets , where we show clear advantages over ERM . At last , we conduct extensive experiments on benchmark datasets and show our DAT can outperform ERM consistently on tasks dominated by both correlation shift and diversity shift , and as a result , achieve state-of-the-art performance ( in average ) on these datasets . We summarize our contributions as follows : • We develop a new kind of adversarial training , Domain-wise Adversarial Training ( DAT ) , for domain generalization , and we establish the intrinsic similarity between IRM and domain-wise AT objectives . • We analyze how DAT will benefit learning invariant features and verify our hypothesis through both synthetic data and real-world datasets . • Extensive experiments on benchmark datasets show that DAT does not only perform better than ERM under correlation shift like IRM but also outperforms ERM under diversity shift like ( sample-wise ) AT . Therefore , our methods achieve state-of-the-art results by surpassing ERM at both kinds of distribution shifts . 2 RELATED WORKS . IRM and Its Variants Invariant Risk Minimization ( IRM ) ( Arjovsky et al . ( 2019 ) ) develops a paradigm to extract causal ( invariant ) features and find the optimal invariant classifier on top over several given training environments . The work of Kamath et al . ( 2021 ) reveals the gap between IRM and IRMv1 , show that even in a simple model that echos the idea of the original IRM objective , IRMv1 can fail catastrophically . Rosenfeld et al . ( 2021 ) prove that when the number of training environments is not large enough , IRM can face the risk of using environmental features . There also exists a predictor feasible for IRMv1 that is very similar to the ERM solution . AT and Its Variants Szegedy et al . ( 2014 ) report one can inject imperceptible perturbations to fool deep models . Athalye et al . ( 2018 ) reveal that among the proposed defenses , adversarial training was the only effective one . Adversarial Training ( AT ) ( Goodfellow et al. , 2014 ; Madry et al. , 2018 ) is the representative approach to train robust models . Recently , Kamath et al . ( 2021 ) show adversarially learned features can transfer better than standardly trained models , while various works ( Volpi et al. , 2018 ; Sinha et al. , 2018 ; Ford et al. , 2019 ; Qiao et al. , 2020 ; Yi et al. , 2021 ; Gokhale et al. , 2021 ) adopt sample-wise adversarial training or adversarial data augmentation to improve OOD robustness . However , most of the discussions are limited to distributional robustness w.r.t . Wasserstein distance , making it less practical for accounting real-world OOD scenarios as discussed in this work . Evaluating OOD Robustness Gulrajani & Lopez-Paz ( 2021 ) recently noticed that the evaluation criteria are crucial for fair comparison of OOD robustness , where under fair settings , no algorithm has performed consistently better than vanilla ERM . Ye et al . ( 2021 ) further identifies two kinds of distribution shifts in current benchmark datasets : diversity shift and correlation shift . Diversity shift refers to the shift of the distribution support of spurious feature z , while correlation shift refers to the change of conditional probability of label y given spurious feature z on the same support . We present an illustrative example in Figure 1 . In particular , they show that there seems to be a tradeoff between the two tasks as an algorithm that performs well at one task tends to perform poorly on the other . Instead , in this work , we show that our DAT can achieve superior performance on both correlation shift and diversity shift tasks . Universal Adversarial Training Universal adversarial perturbation proposed by Moosavi-Dezfooli et al . ( 2017 ) is a type of adversarial attack that adopts a universal perturbation for all images . Universal adversarial training ( UAT ) ( Shafahi et al. , 2020 ) is then proposed to defend against this attack by training against universally perturbed data . Instead , in our work , we firstly show that we can adapt UAT for solving domain generalization problems . 3 ON THE RELATIONSHIP BETWEEN IRM AND AT . 3.1 PRELIMINARY . Notation Let Φ : X ⊂ Rn → Rd denotes the representation of a θ-parameterized piecewise linear classifier , i.e. , Φ ( · ) = φL ( WLφL−1 ( . . . ) + bL−1 ) + bL , where φL is the activation function , and WL , bL denote the layer-wise weight matrix and bias vector , collectively denoted by θ. Additionally , let β be the linear classifier on top , and the output of the network is β · Φ ( x ) = β > Φ ( x ) . Let ` ( ŷ , y ) = − log σ ( yŷ ) be the sample logistic loss . We consider a two-class ( y = ±1 ) classification setting with output dimension d = 1 , and our discussion can be easily extended to the general cases . ERM The traditional Empirical Risk Minimization ( ERM ) algorithm optimizes over the loss on i.i.d . data , i.e. , min β , Φ R ( β · Φ ) , where R ( β · Φ ) = E ( x , y ) ∼D ` ( β > Φ ( x ) , y ) . ( 1 ) In OOD generalization problem , one faces a set of ( training ) environments E , where each environment e ∈ E corresponds to a unique data distribution De . When facing multiple environments , the ERM objective simply mixes the data together and takes the form ( ERM ) min β , Φ ∑ e Re ( β · Φ ) , where Re ( β · Φ ) = E ( x , y ) ∼De ` ( β > Φ ( x ) , y ) . ( 2 ) IRM and IRMv1 Instead of simply mixing the data together , IRM seeks to learn an invariant representation Φ such that it can be minimizing with the same classifier β . Formally we have ( IRM ) min β , Φ ∑ e∈E Re ( β · Φ ) s.t . β ∈ arg min β̄ Re ( β̄ · Φ ) , ∀e ∈ E . ( 3 ) Since this bi-level optimization problem is difficult to solve , the practical version IRMv1 as regularized ERM , where the gradient penalty is calculated w.r.t . a dummy variable w : ( IRMv1 ) min β , Φ ∑ e∈E [ Re ( β · Φ ) + λ · ||∇w|w=1.0Re ( w · ( β · Φ ) ) ||2︸ ︷︷ ︸ PenaltyIRM ] . ( 4 ) AT Adversarial Training ( Madry et al. , 2018 ) ) instead replaces ERM with a minimax objective , ( AT ) min β , Φ RAT ( β · Φ ) = min β , Φ E ( x , y ) ∼D max ||δx||p≤ε ` ( β > Φ ( x+ δx ) , y ) , ( 5 ) where one maximizes the inner loss by injecting sample-wise perturbations δx and solve the outer minimization w.r.t . parameters β , Φ on the perturbed sample ( x+ δx , y ) . Typically , the perturbation is constrained within an ` p ball with radius ε . In this way , AT can learn models that are robust to ` p perturbations . UAT Instead of injecting sample-wise perturbations , Moosavi-Dezfooli et al . ( 2017 ) notice that we can also adopt a universal perturbation δ for all samples , which results in the Unversarial Adversarial Training ( UAT ) objective as follows : ( UAT ) min β , Φ RUAT ( β · Φ ) = min β , Φ max ||δ||p≤ε E ( x , y ) ∼De ` ( β > Φ ( x+ δ ) , y ) , ( 6 )

Summary: The paper investigates whether adversarial training (AT) could be used for extracting domain-invariant features, and whether AT can benefit OOD generalization.  The paper first shows the relationship between AT and IRM (invariant risk minimization).  A new version of AT is proposed -- called "Domain-wise Adversarial Training (DAT)"; with IRM being a version of DAT.  Results suggest that DAT can outperform ERM under correlation-shift and diversity shift.