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 tackles the problem of out-of-distribution (OOD) generalization of deep learning models with a novel adversarial training formulation (DAT) that introduces a shared adversarial perturbation per training domain.  The paper further establishes a connection between DAT, Invariant Risk Minimization (IRM), and standard adversarial training (AT).  The authors posit that DAT inherits the advantages of IRM (which is effective for correlation shift while it performs more poorly on diversity shift) and AT, which improves upon Empirical Risk Minimization (ERM) for diversity shift. Experiments on several OOD benchmark datasets show that DAT outperforms prior works and ERM on average.