Year: 2021

Abstract:
Despite the impressive success on many tasks, deep learning models are shown to rely on spurious features, which will catastrophically fail when generalized to outof-distribution (OOD) data. To alleviate this issue, Invariant Risk Minimization (IRM) is proposed to extract domain-invariant features for OOD generalization. Nevertheless, recent work shows that IRM is only effective for a certain type of distribution shift (e.g., correlation shift) while fails for other cases (e.g., diversity shift). Meanwhile, another thread of method, Adversarial Training (AT), has shown better domain transfer performance, suggesting that it is potential to be an effective candidate for extracting domain-invariant features. In this paper, we investigate this possibility by exploring the similarity between the IRM and AT objectives. Inspired by this connection, we propose Domain-wise Adversarial Training (DAT), an AT-inspired method for alleviating distribution shift by domainspecific perturbations. Extensive experiments show that our proposed DAT can effectively remove the domain-varying features and improve OOD generalization under both correlation shift and diversity shift.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORKS
text: 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.

heading: 3 ON THE RELATIONSHIP BETWEEN IRM AND AT

heading: 3.1 PRELIMINARY
text: 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)

heading: 3.2 RELATING ADVERSARIAL TRAINING TO IRM
text: Motivation As shown above, it seems that IRM and AT are two distinct learning paradigms, while in fact, we can show that IRM is closely related to a certain kind of adversarial training. To see this, we first notice that AT can be rephrased into a regularized ERM loss with a penalty on sample-wise robustness through linearization:
RAT(β · Φ) = E(x,y)∼D max ||δx||≤ε `(β>Φ(x+ δx), y)
= E(x,y)∼D [ `(β>Φ(x), y) + max
||δx||≤ε
( `(β>Φ(x+ δx), y)− `(β>Φ(x), y) ) ] ≈ R(β · Φ) + εE(x,y)∼D
∥∥∇x`(β>Φ(x), y)∥∥︸ ︷︷ ︸ PenaltyAT ,
(7)
which resembles the gradient penalty adopted in IRMv1. One main difference is that AT’s penalty is calculated w.r.t. sample-wise gradients, while IRM’s penalty w.r.t. the population loss. This difference motivates us to adopt a population-level perturbation δ instead, which, in the literature of adversarial learning, corresponds to Universal AT (Eq. 6) that could be rephrased as follows:
RUAT(β · Φ) = max ‖δ‖≤ε E(x,y)∼D`(β>Φ(x+ δ), y)
= E(x,y)∼D`(β>Φ(x), y) + max ‖δ‖≤ε
E(x,y)∼D ( `(β>Φ(x+ δ), y)− `(β>Φ(x), y) ) ≈ R(β · Φ) + ε
∥∥∇xE(x,y)∼D`(β>Φ(x), y)∥∥︸ ︷︷ ︸ PenaltyUAT .
(8)
DAT Inspired by the connection above between UAT with IRM, we adapt UAT to the OOD setting with multiple domains and propose Domain-wise Adversarial Training (DAT), which adopts a domain-wise perturbation δe for each training domain e ∈ E . Formally, we have
min β,Φ ∑ e∈E E(x,y)∼De` ( β>Φ(x+ δe)), y ) s.t. δe ∈ arg max
‖δ‖≤ε E(x,y)∼De`(β
>Φ(x+ δ), y),∀e ∈ E . (9)
The objective could be interpreted as the loss function of universal training in a domain-wise fashion, which lead to solving for the perturbation vector individually for each domain. We name this method as Domain-wise Adversarial Training (DAT), a generalization of UAT for learning invariance from multiple domains. In practice, we solve the problem above with alternating update of model parameters β,Φ and perturbations δe. Specifically, for each mini-batch Be sampled from domainDe, we update δe withBe using gradient ascent to find the best adversarial perturbations. Then the adversarial samples are used to train the model. A detailed description is shown in Algorithm 1.
Algorithm 1 Domain-wise Adversarial Training Input: Dataset of multiple environments De, e ∈ E , desired lp norm of the perturbation ε, pertur-
bation step size α Output: Model (Φ, β)
Randomly initiate θ, perturbation δe,∀e ∈ E for iterations = 1, 2, 3, . . . do
for each environment e do 1. Sample batch Be from environment e 2. Update the perturbation δe using one-step gradient ascent with step size α 3. Project the perturbation δe to the `p ball of radius ε 4. Generate adversarial examples xadv ← x+ δe,∀x ∈ Bi 5. Update Φ and β with gradient descent (Adam) on xadv
end for end for

heading: 3.3 CONNECTION BETWEEN IRM AND DAT
text: Here, we establish a formal connection between IRM and DAT. To begin with, we note that DAT can also be reformulated as a regularized ERM in the multi-domain scenario.
RDAT(β · Φ) = ∑ e∈E max ||δe||≤ε E(x,y)∼De`(β >Φ(x+ δe), y)
= ∑ e∈E [ E(x,y)∼De`(β >Φ(x), y) + max ||δe||≤ε E(x,y)∼De ( `(β>Φ(x+ δe), y)− `(β>Φ(x), y) ) ] ≈ ∑ e∈E [ Re(β · Φ) + ε ∥∥∇xE(x,y)∼De`(β>Φ(x), y)∥∥︸ ︷︷ ︸ PenaltyDAT ] .
(10)
Based on this reformulation, we can show that there exists an intrinsic relationship between DAT and IRM as in the following proposition:
Proposition 1 Consider each De as the corresponding distribution of a particular training domain e. For any β · Φ as a deep network with any activation function, the penalty term of IRMv1, PenaltyIRM (Eq. 4), could be expressed as the square of a reweighted version of the penalty term of the above approximate target, PenaltyDAT(Eq. 10), on each environment e with coefficients related to the distribution De, which could be stated as follows:
PenaltyIRM = ‖EDe [Lxx]‖ 2 , and PenaltyDAT = ‖EDeLx‖ (11)
where Lx = (1− σ(yβ>Φxx))yβ>Φx.
The proposition above shows that IRM can be regarded as an instance-reweighted version of DAT, which indicates that DAT can also be applied to improve out-of-distribution generalization performance. We explore this possibility in the next section.
Besides, assume that each domain only contains one sample, we can further obtain a stronger result, that is, the equivalence between IRM and (linearized) DAT as follows.
Remark 1 (Equivalence under Single-sample Environments) When the environments degenerate into a single data point, we have the following relationship: If ε is sufficiently small, then for β · Φ as a deep network with any activation function, the penalty term of IRMv1 (Eq. 4) on each sample and the square of the maximization term of Linearized version of Eq. 9 (LDAT, obtained by first-order approximation of DAT)
PenaltyLDAT = 〈 ∇x` ( βTΦ(x), y ) ,±δ̂x 〉 (12)
on each sample with perturbation δ̂x = ±εx only differ by a fixed multiple ε2. Which is formally stated as
Penalty2LDAT = ε 2 · PenaltyIRM. (13)
Proofs of Proposition 1 and Remark 1 can be found in Appendix C.

heading: 4 UNDERSTANDING DOMAIN-WISE ADVERSARIAL TRAINING
text: The above connection between DAT and IRM highlights that our DAT is potentially helpful for addressing OOD problems. In this section, we further explore how domain-wise perturbations could help alleviate distribution shift in real-world scenarios. In particular, we notice that the domain-wise perturbations could successfully remove the domain-varying background information, which usually corresponds to the spurious features for image classification tasks. And we empirically verify this property by applying DAT to a well-designed OOD task based on background shift.

heading: 4.1 COMPARING ERM TO DAT FOR BACKGROUND REMOVAL
text: ERM Learns Spurious Background Features Our understanding of DAT hinges on the insight that an image is composed of a foreground object and the corresponding background, and typically, the object is the invariant feature while the background is only spuriously correlated with the labels. However, models relying on the spurious background information will easily fail when encountering images from a different domain. This phenomenon is also empirically verified by Xiao et al. (2021), who find that models trained on an ImageNet-like dataset with ERM require image backgrounds for correctly classifying large portions of test sets. These findings point out the limitations of ERM and motivate us to find a solution that could effectively learn background-invariant classifier.
Removing Background Information with DAT Comparing to the failure of ERM above, we notice that Domain-wise Adversarial Training (DAT) can be applied to eliminate the domain-wise background information with its domain-specific perturbations.
Take the NICO images in Figure 2 for an example, where samples from the environment “on grass” have a common background dominated by the green grass with low frequency, while the foreground object (e.g., the cats) has complex and instance-specific shape and texture with much higher frequency. In fact, Moosavi-Dezfooli et al. (2017) show that a universal perturbation vector lies in a low dimensional subspace, which fits the background statistics and could be used to eliminate the low-frequency background factor. Therefore, when applying our DAT to these samples, the domainwise perturbation will capture the common domain-specific background. And consequently, the domain-wise AT will help remove the dependence on these spurious background features.

heading: 4.2 EMPIRICAL VERIFICATION WITH CONTROLLED EXPERIMENTS
text: To verify the above analysis, we construct a synthetic OOD task to evaluate a classifier’s dependence on background information. It is based on two datasets introduced by Xiao et al. (2021), Mixed-Same and Mixed-Rand, which are constructed from a subset of ImageNet images with the background of each image replaced by another background that is either from the same class (MixedSame) or a random class (Mixed-Rand). As they are perfect candidates for evaluating a classifier in terms of its background dependence, we construct a new OOD task by using Mixed-Same as the training domain and evaluating the learned classifier on Mixed-Rand as the test domain. If the clas-
sifier relies too much on background information, it will perform poorly in the test domain where objects and backgrounds are disentangled. In particular, experiment results show that ERM achieves a test accuracy of 71.9%, while DAT achieved 72.6% on the test domain with random background, which means that DAT has better generalization ability by effectively removing background information. Sample images from the dataset and corresponding attention heatmaps are shown in Figure 3, which demonstrate that ERM may lose its focus when background correlation is broken while DAT doesn’t. Details for the experiment are shown in Appendix A.1

heading: 5 EXPERIMENTS

heading: 5.1 EXPERIMENTAL SETUP
text: For the experiments, we follow the setting in Ye et al. (2021) and evaluate OOD generalization on both the two types of distribution shift, diversity shift and correlation shift. In particular, we select four representative tasks. For correlation shift, we use CMNIST (Arjovsky et al., 2019), a synthetic dataset on digit classification which has color as the spurious feature, and NICO (He et al., 2020), a real-world dataset on object classification with context as the spurious feature. As for diversity shift, we use PACS (Li et al., 2017) and Terra Incognita (Beery et al., 2018), which are both real-world datasets with four domains. To ensure fair evaluation, we conduct all our experiments following the evaluation pipeline of DomainBed (Gulrajani & Lopez-Paz, 2021). Specifically, we use the same strategy for dataset splitting and model selection as in Ye et al. (2021) for each of these tasks. For datasets except for NICO, one of the domains is used as the test domain. We train the models in each run treating one of the domains as the test domain with the rest of the domains as training domains. Then report the average accuracy of all domains. For NICO, the training, test, evaluation, and test domains are predefined. We train the models on training domains, evaluate them on the evaluation domain for model selection, and report their accuracy on the test domain. More details of the experimental settings could be found in Appendix A.2.
When training models using DAT, we first perform standard data augmentation Gulrajani & LopezPaz (2021), then proceed with the update on perturbation and model parameters as shown in 1, where the perturbed samples are clipped to the legal range after data augmentation.
As discussed in Section 4, DAT can reduce the influence of background even when no domain labels are given, which corresponds to the single-source domain generalization setting. To see how this could help in this setting, we conduct experiments which strengthen this claim. Experimental setting and results can be found in Appendix B.

heading: 5.2 EVALUATION ON BENCHMARK DATASETS
text: We compare our results with previous works, including vanilla ERM, invariance-based methods including IRM (Ahuja et al., 2020), robust optimization methods including GroupDRO (Sagawa et al., 2020), distribution matching methods including MMD (Li et al., 2018b) and CORAL (Sun & Saenko, 2016), a method based on domain classifier (DANN, (Ganin et al., 2016)) and various other algorithms. The results for the CMNIST dataset are adopted from Gulrajani & Lopez-Paz
(2021), which is the average of three domains, while the results of the other datasets are adopted from Ye et al. (2021). Apart from that, we implement and test four AT based algorithms including sample-wise adversarial training (AT (Goodfellow et al., 2014)), universal adversarial training (UAT (Shafahi et al., 2020) (which do not distinguish between different training domains), WRM (Sinha et al., 2018), and Adversarial Data Augmentation (ADA (Volpi et al., 2018)).
From Table 1, we observe that on a variety of OOD benchmark datasets, DAT has a performance consistently better than ERM. In particular, DAT achieves state-of-the-art performance under both diversity and correlation shift. In comparison, IRM only outperforms DAT on CMNIST and samplewise AT fails on three of the four datasets. Compared to other AT-based algorithms (e.g., samplewise AT, UAT, WRM, and ADA), DAT has superior performance by considering a domain-wise perturbation that removes domain-varying spurious features. The results demonstrate DAT’s effectiveness in dealing with domain discrepancy. RSC Huang et al. (2020) has the most superior performance under diversity shift, this may due to that RSC masks large gradients, which prevents the model from becoming over-confident by only capturing a few dom-
inant features that are not invariant. However, RSC does not use domain labels, which makes it fail to consider the change of correlation between features and label class when generalizing to a new domain and leads to poor performance under correlation shift. In comparison, DAT solves the same problem in a reversed way, by perturbing inputs to reduce dominant features, and results in a “weaker” augmentation compared to RSC under diversity shift. But DAT considers domain difference explicitly, thus having better performance than RSC under correlation shift. We also note that DAT performs slightly worse than IRM on CMNIST, which could be due to the data generation process of CMNIST. As a synthetic dataset, CMNIST uses colored digits instead of colored background as the spurious features. In contrast, in real-world datasets like NICO (Figure 2), the spurious features mainly lie in the background instead of the object. This suggests that our method is more suitable for alleviating real-world distribution shifts.

heading: 5.3 ANALYSIS
text: We conduct experiments to understand better what our algorithm learns and how the magnitude of hyperparameters affects its performance. Qualitative Analysis through Semantic Graphs We use GradCam (Selvaraju et al., 2017; Gildenblat & contributors, 2021) to visualize the attention heatmaps of models trained by ERM, sample-wise AT, and DAT on the NICO dataset. The results are shown in Figure 4, where we can see that DAT puts more attention on the object itself instead of the strongly correlated background, while ERM and sample-wise AT tend to use environmental features instead.
Perturbation Radius and Step Size We investigate the effect of perturbation radius ε and perturbation step size α on the NICO dataset. The results are shown in Table 2. From the results, we find that the perturbation radius ε has a huge impact on the OOD accuracy. When ε is too large (> 10−1), it begins to hurt the invariant feature and causes performance degradation (from 72.9% to 67.9%). The step size α has a smaller influence, and choosing a value about 1/100 to 1/10 times the size of ε would be appropriate.
Norm of Perturbations We also investigate the effect of the norm used in DAT. We perform experiments on the NICO dataset using `∞-norm bounded DAT. The results are shown in Table 3. Although there are slight differences, DAT bounded by both p = 2 and p = ∞ can achieve state-of-the-art performance (72.6% and 72.5%) on the NICO dataset.
6 DISCUSSIONS
Comparison with Sample-wise AT Previous works (Yi et al., 2021; Hendrycks et al., 2021) try to exploit sample-wise AT as a data augmentation strategy to get higher OOD performance. However, the performance only improves when the distribution shift is dominated by diversity shift, e.g., noise, and blurring. Otherwise, the performance might be degraded, as shown in Table 1. One possible explanation is that because sample-wise AT fails to capture the domain-level variations as DAT. As a result,
it could possibly add perturbations to the invariant features and lead to inferior performance, especially under correlation shift.
Comparison with Invariant Causal Prediction A thread of methods including ICP (Peters et al., 2016), IRM (Arjovsky et al., 2019), and IGA (Koyama & Yamaguchi, 2020) try to find the invariant data representations that could induce an invariant classifier. They have superior performance on synthetic datasets like CMNIST but fail to outperform ERM on real-world datasets (tasks dominated by both correlation shift and diversity shift). We believe their failures might be attributed to the lack of prior information in their invariant learning principles. While in our DAT, we effectively exploit the foreground-background difference in image classification tasks through domain-wise perturbations.

heading: 7 CONCLUSION
text: In this work, we carefully analyze the similarity between IRM and adversarial training in a domainwise fashion and establish a formal connection between OOD and adversarial robustness. Based on this connection, we propose a new adversarial training method for domain generalization, namely Domain-wise Adversarial Training (DAT). We show that it could effectively remove the spurious background features in image classification and obtain superior performance on benchmark datasets. Notably, our DAT could outperform ERM consistently on tasks dominated by both correlation shift and diversity shift, while previous methods typically fail in one of the two cases.

heading: A EXPERIMENT SETTINGS AND MORE DETAILED RESULTS
text: We train all the models based on the DomainBed benchmark presented in Gulrajani & Lopez-Paz (2021).
A.1 MIXED-SAME AND MIXED-RAND
To construct the dataset, we use the validation set of Mixed-Same (4185 images) provided by Xiao et al. (2021) as the training environment and use the validation set of Mixed-Rand (4185 images) as the testing (OOD) environment.
For training, we use pretrained ResNet-18 as backbone and train it for 5000 iterations. The hyperparameters are same to those proposed in Gulrajani & Lopez-Paz (2021), except that the learning rate searching interval is altered to [10−5.5, 10−4.5], while we set hyperparameters search space of DAT as ε ∈ [10−0.5, 100], α ∈ [10−1.5, 10−1].
A.2 TYPICAL OOD DATASETS
We conduct experiments on four typical OOD datasets, including CMNIST, NICO, Terra Incognita, and PACS. We split the NICO dataset using the same strategy in Ye et al. (2021) for a fair comparison.
For network architecture, models trained on CMNIST adopt the four-layer convolution network MNIST-CNN provided in DomainBed benchmark (Gulrajani & Lopez-Paz, 2021), while other datasets use ResNet-18 as the backbone. For Terra Incognita and PACS, pretrained ResNet-18 is used. While on NICO, we use unpretrained ResNet-18 as the pretraining dataset contains photos that are largely overlapped with ImageNet classes.
For model selection, models trained on PACS and Terra Incognita are selected by training-domain validation, which selects the model with highest training accuracy averaged across all training domains. For NICO, an extra OOD validation set is used. CMNIST, due to the large correlation shift, uses test-domain validation as its model selection criteria.
DAT, AT, and UAT models are trained for an appropriate number of iterations to ensure convergence, that is, 5000 for Terra Incognita, 8000 for NICO, 10000 for CMNIST and PACS.
When training NICO using DAT, we clamp the single sample adversarial loss to (0, 2) to obtain a better domain-wise perturbation as suggested in Shafahi et al. (2020).
We use the same hyperparameters as those proposed in Ye et al. (2021) whenever possible while setting the hyperparameters search space for DAT, sample-wise AT and UAT as in Table 6 and 5. We use DAT with perturbation bounded by `2-norm in our experiments. For sample-wise AT, we use a 10-step `2 PGD (Madry et al., 2018).
For WRM, we set the step number to 10 in the inner maximization, while searching for the learning rate in the inner maximization in the range [10−2, 10−1] and γ in the range [10−0.5, 100.5]. We train the models for 5000 iterations to ensure convergence.
For ADA, we set the step number to 10 in the inner maximization, while searching for the adversarial learning rate in the range [10−1, 100.5], γ in the range [0.5, 1.5], the number of steps in the min-phase
in the range [100.5, 102.5], and the number of whole adversarial phases in the range [100.5, 102]. We train the models for 5000 iterations to ensure convergence.
The domain split results for ERM, UAT, and DAT are shown in Table 7, 8, and 9, where the results of ERM on CMNIST are from Gulrajani & Lopez-Paz (2021) and all other results come from our experiments.

heading: B ON SINGLE-DOMAIN DOMAIN GENERALIZATION
text: It can be seen that our proposed DAT can also work with a single domain. In this case, DAT reduces to UAT as a special case. In this section, we present a deeper understanding on the effectiveness of DAT on single-domain generalization, as well as the comparison between DAT and ensembled UAT. Specifically, in Section B.1, we perform experiments under single-source domain generalization setting, which show that DAT is still effective under this scenario by removing the common background in the domain. In Section B.2, we compare DAT with the ensemble of UAT, the results suggest that DAT has extra benefits by training a single model on multiple domains.
B.1 SINGLE-SOURCE DOMAIN GENERALIZATION
To show that our method can also be applied to single-domain generalization scenarios, we conduct experiments with similar settings to Appendix A.2 on the four benchmark datasets, except that we train on one of the domains and report its test accuracy on all other domains. For the NICO dataset, as it comes with a validation domain, we train the models on one of the training domains, using the validation domain for model selection, then report its accuracy on the test domain.
We train all models using ERM/DAT for 5000 iterations to ensure convergence and set the hyperparameter searching space of DAT as shown in Table 10. All other settings are the same as those shown in Appendix A.2.
We list the domain-average results in Table 11.
We can see that our DAT enjoys superior performance on all four datasets compared to ERM, either under correlation shift or diversity shift. Although the difference is not as significant as in the multiple domain setting, it shows that our DAT works for both single-domain and multi-domain generalization scenarios. In particular, its advantages are more significant with multiple domains, where the domain-wise perturbation mechanism is more effective.
The domain split results for ERM and DAT under single-source domain generalization setting are shown in Table 12a to 12h.
B.2 COMPARISON WITH THE ENSEMBLE OF UAT
To show that models trained using DAT on multi-domains are not a trivial ensemble of models trained on single domains, we train voting classifiers on CMNIST and NICO using UAT. The models are trained separately on each domain and then perform an equal weight voting to get prediction results. The results are shown in Table 13.
From the results, we can see that invariance learned by DAT indeed performs much better than a trivial ensemble of independent UAT models. We believe that training DAT on multiple domains can help extract the domain-invariant features through sharing the same model across domains while removing domain-varying features with domain-specific perturbations. Spurious background information from multiple domains can help the classifier to identify the truly domain-invariant features. In comparison, single-domain UAT only has access to a single domain and can only learn invariance of a single domain, which leads to inferior performance.

heading: C PROOFS
text: We give a proof of Remark 1 first, then use a similar technique to prove Proposition 1.
Remark 1 (Equivalence under Single-sample Environments) When the environments degenerate into a single data point, we have the following relationship: If ε is sufficiently small, then for β · Φ as a deep network with any activation function, the penalty term of IRMv1 (Eq. 4) on each sample and the square of the maximization term of Linearized version of DAT (LDAT, obtained by first-order approximation of DAT)
PenaltyLDAT = 〈 ∇x` ( βTΦ(x), y ) ,±δ̂x 〉 (14)
on each sample with δ̂x = ±εx only differ by a fixed multiple ε2. Which is formally stated as
Penalty2LDAT = ε 2 · PenaltyIRM. (15)
Proof of Remark 1 Note that we presume Φ is piece-wise linear. We represent the output logit of the given deep network as β>Φ(x) = β>Φxx where Φx is the matrix related to sample x due to the fact of different activation patterns.
We first derive the form of PenaltyLDAT by first-order approximation of Eq. 9. `(β>Φ(x+ δe), y)− `(β>Φ(x), y) ≈ 〈 ∇x`(β>Φ(x), y), δe 〉 (16)
Letting δe = ±εx, we have that PenaltyLDAT = 〈 ∇x` ( βTΦ(x), y ) ,±δ̂x 〉 (17)
The penalty term of IRMv1 on each sample is as follows:∥∥∇w|w=1.0 − log(σ(w · (yβ>Φ(x))))∥∥2 = ∥∥∥∥−σ′(w · (yβ>Φ(x)))σ(w · (yβ>Φ(x))) yβ>Φ(x) ∥∥∥∥2
= ∥∥−(1− σ(yβ>Φ(x)))yβ>Φ(x)∥∥2
= (1− σ(yβ>Φ(x)))2 ∥∥β>Φxx∥∥2
(18)
For LDAT, note that the gradient w.r.t. x is equal to
∇x[− log(σ(yβ>Φ(x))] = −(1− σ(yβ>Φ(x)))yβ>Φx So the square of the penalty term of LDAT on each sample with perturbation ±εx is:
[〈∇x − log(σ(yβ>Φ(x))),±εx〉]2 = ε2(1− σ(β>Φ(x)))2 ∥∥β>Φxx∥∥2 (19)
which is identical to Eq. 4 with a difference of multiple ε2.
Proposition 1 Consider each De as the corresponding distribution of a particular training domain e. For any β · Φ as a deep network with any activation function, the penalty term of IRMv1, PenaltyIRM (Eq. 4), could be expressed as the square of a reweighted version of the penalty term of the above approximate target, PenaltyDAT(Eq. 10), on each environment e with coefficients related to the distribution De, which could be stated as follows:
PenaltyIRM = ‖EDe [Lxx]‖ 2 , and PenaltyDAT = ‖EDeLx‖ (20)
where Lx = (1− σ(yβ>Φxx))yβ>Φx.
Proof of Proposition 1 For IRMv1, from the proof of Remark 1 we know
∇w|w=1.0`(w · (β · Φ), y) = −(1− σ(yβ>Φxx))yβ>Φxx (21)
So that PenaltyIRM (suppose derivation and integration are commutable)∥∥∇w|w=1.0Re(w · (β · Φ))∥∥2 = ∥∥EDe(1− σ(yβ>Φxx))yβ>Φxx∥∥2 (22)
For DAT, on each environment
max ‖δe‖≤ε
EDe [ `(β>Φ(x+ δe), y)− `(β>Φ(x), y) ) ] ≈ max ‖δe‖≤ε EDe [〈 ∇x`(β>Φ(x), y), δe
〉] = max ‖δe‖≤ε 〈 EDe [ ∇x`(β>Φ(x), y) ] , δe 〉
=ε ∥∥EDe [∇x`(β>Φ(x), y)]∥∥
(23)
So PenaltyDAT is∥∥EDe [∇x`(β>Φ(x), y)]∥∥ = ∥∥EDe [−(1− σ(yβ>Φ(x)))yβ>Φx]∥∥ (24) So we know PenaltyIRM can be regarded as the square of a reweighted version of PenaltyDAT with coefficient on each sample inside the expectation x.

Review summaries:
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.

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. 

This paper focuses on domain generalization (DG). It proposes a method based on adversarial training (AT): the main idea is learning universal adversarial perturbations at domain level (domain adversarial perturbations, DAT), and rely on those within a standard AT routine. The relation between the proposed DAT and invariant risk minimization (IRM) is provided, and the method is empirically tested on standard domain generalization benchmarks.