Improving Out-of-Distribution Generalization by Adversarial Training with Structured Priors
Abstract
Deep models often fail to generalize well in test domains when the data distribution differs from that in the training domain.
Among numerous approaches to address this Out-of-Distribution (OOD) generalization problem, there has been a growing surge of interest in exploiting Adversarial Training (AT) to improve OOD performance.
Recent works have revealed that the robust model obtained by conducting sample-wise AT also retains transferability to biased test domains.
In this paper, we empirically show that sample-wise AT has limited improvement on OOD performance.
Specifically, we find that AT can only maintain performance at smaller scales of perturbation while Universal AT (UAT) is more robust to larger-scale perturbations.
This provides us with clues that adversarial perturbations with universal (low dimensional) structures can enhance the robustness against large data distribution shifts that are common in OOD scenarios.
Inspired by this, we propose two AT variants with low-rank structures to train OOD-robust models.
Extensive experiments on DomainBed benchmark show that our proposed approaches outperform Empirical Risk Minimization (ERM) and sample-wise AT. Our code is available at https://github.com/NOVAglow646/NIPS22-MAT-and-LDAT-for-OOD .
Introduction
Existing deep learning methods have achieved good performance on visual classification tasks under the same distribution of training sets and test sets.
However, when the data distribution of the test set is different from that of the training set, the classification performance of the deep neural networks (DNNs) may decrease sharply [1].
This is mainly because DNNs may capture spurious features such as the background and style information to assist the fast fitting during the training process [2].
However, in real-world scenarios, test data may differ from training data in the background and style information, thus DNNs that rely on unstable spurious features to make predictions will fail.
Solving the above problem is known as the out-of-distribution (OOD) generalization.
Another scenario where DNNs may fail is that they are often vulnerable to adversarial examples [3].
Adversarial training (AT) is originally proposed as an effective way to defend against adversarial attacks [4].
Moreover, there is work showing that adversarial training helps to solve the OOD generalization problem because OOD data can be seen as stronger perturbations to some extent [5].
The reason why AT can defend against adversarial attacks meanwhile benefit OOD generalization is that it can make DNNs robust to the interference of spurious features, such as randomly injected noise (in adversarial examples) or the spurious correlation between labels and background information (in OOD generalization).
In other words, AT enables DNNs to make predictions using intrinsic features rather than spurious features.
A potential problem, however, is that existing AT methods ignore the specific design of perturbations when used for solving OOD generalization problems.
They usually simply conduct sample-wise AT [6], which only brings limited performance improvement to OOD generalization.
The essential reason for the failure of this type of approach is that the perturbations it uses cannot distinguish invariant and spurious features.
As a result, it improves the robustness at the expense of the decreasing standard accuracy [7].
Moreover, we empirically find that when adapting Universal AT (UAT [8]) to OOD problems, i.e., conducting AT with domain-wise perturbations, it shows stronger input-robustness when facing larger-scale perturbations compared to the sample-wise AT (see Section 3.2).
Since the sample injected with large-scale perturbations can be regarded as OOD samples [5], we draw inspiration from this phenomenon that AT with universal (low-dimensional) structures can be the key to solving OOD generalization.
Therefore, we propose to use structured low-rank perturbations related to domain information in AT, which can help the model to filter out background and style information, thus benefiting OOD generalization.
We make the following contributions in our work:
• We identify the limitations of sample-wise AT on OOD generalization through a series of experiments.
To alleviate this problem, we further propose two simple but effective AT variants with structured priors to improve OOD performances.
• We theoretically prove that our proposed structured AT approach can accelerate the convergence of reliance on spurious features to 0 when using finite-time-stopped gradient descent, thus enhancing the robustness of the model against spurious correlations.
• By conducting experiments on the DomainBed benchmark [9], we demonstrate that our methods outperform ERM and sample-wise AT on various OOD datasets.
Related Work
Solving OOD Generalization with AT.
According to [3], the performance of deep models is susceptible to small-scale perturbations injected in the input images, even if these perturbations are imperceptible to humans.
Adversarial training (AT) is an effective approach to improve the robustness to input perturbations [4, 10, 11].
However, many recent works have begun to focus on the connection between AT and OOD due to the fact that OOD data can be regarded as one kind of large-scale perturbation.
These works seek to exploit the robustness provided by AT to improve OOD generalization.
For instance, [6] applied sample-wise AT to OOD generalization.
They theoretically found that if a model is robust to input perturbation on training samples, it also generalizes well on OOD data.
[5] theoretically established a link between the objective of AT and the OOD robustness.
They revealed that the AT procedure can be regarded as a heuristic solution to the worst-case problem around the training domain distribution.
Nevertheless, the discussion of [6] and [5] is restricted to the framework of using Wasserstein distance to measure the distribution shift, which is less practical for the real-world OOD setting where domain shifts are diverse.
Additionally, they only studied the case of sample-wise AT and did not further investigate the effect of different forms of AT (not sample-wise) on OOD performance.
Other works such as [12] focus on the structure design of the perturbations.
They used multi-scale perturbations within one sample, but they did not exploit the universal information within one training domain.
In our work, we focus on real-world OOD scenarios where there are additional clues lying in the distribution shifts, i.e, the low-rank structures in the spurious features (such as background and style information) across one domain.
We further design a low-rank structure in the perturbations to specifically eliminate such low-rank spurious correlations.
OOD Evaluation Benchmark.
The DomainBed benchmark [9] provides a fair way of evaluating different state-of-the-art OOD methods, which has been widely accepted by the community.
By conducting rigorous experiments in a consistent setting, they revealed that many algorithms that claim to outperform previous methods cannot even outperform ERM.
Unlike previous works using AT to address OOD generalization, such as [6] and [5], we adopt the Domainbed benchmark for a fair comparison of our approach with existing state-of-the-art methods in this paper.
Weakness of Sample-wise AT for OOD Generalization 3.1 Preliminaries
Out-of-distribution (OOD) Generalization .
Assuming x ∈ X as the random data in the input space
X and y ∈ Y as the target random data in the label space Y , we have the predictor f = w ◦ ϕ
( x ) where ϕ : X → Z denotes the feature extractor and w : Z → Y denotes the classifier.
Now we give the formal definition of the OOD generalization problem.
We have a set of m training domains E = { E 1 , E 2 , ..., E m } , where each domain E e is characterized by a input dataset
E e := { ( x ei , y ei ) } n e i =
1 containing n e i.i.d input samples drawn from the distribution of P e , and a test domain E m +1 with data following the distribution of P te , where P te ̸
= P i , i = 1 , 2 , ..., m .
L : X → R + denotes the loss function.
The ultimate goal of OOD generalization is to find an optimal predictor f that minimizes the risk on the unseen test domain:
[Equation]
Adversarial Training (AT) 3 .
According to [4], AT can be expressed as the following optimization problem:
[Equation]
where δ ∈ S is the random injected perturbation with l p norm bounded by ϵ .
The inner maximization problem can be optimized by fast gradient sign method (FGSM [13]), a simple one-step scheme:
[Equation]
where sgn ( · ) is the sign function, or by projected gradient descent (PGD [4]), a more powerful multi-step variant: [Equation] ∇ x L ( f ( x ) , y ))) , (
[Equation]
) where (cid:81)
S is the projection operator onto the set S , γ is the step size and t denotes the iteration.
3.2 Weakness of AT for OOD Generalization
We now highlight some weaknesses of sample-wise AT for OOD generalization based on a series of empirical evidence.
We first conduct a toy experiment on the DomainBed benchmark [9] to evaluate the OOD performance of AT.
We run ERM and AT on four OOD datasets: PACS [14], OfficeHome [15], VLCS [16], and NICO [17] with a fixed set of hyperparameters (detailed experimental settings can be found in Appendix C.1).
The results are shown in Table 1.
We can see that the improvement of OOD performance by AT is limited with an average improvement of only 0.1%.
[Table] We further investigate the reason behind the limitations of performance improvements on OOD datasets of AT.
Although previous works have revealed that the robust features obtained by AT can improve OOD generalization ([6] [5] [18]), we find that sample-wise AT only tolerates small-scale perturbations.
Thus, we design an experiment on NICO dataset with multiple scales of perturbations.
The scale is calculated with the l 2 norm of the perturbation matrix (experiment details are shown in Appendix C.1).
As shown in Figure 1, AT suffers severe performance degradation when using large perturbations.
This provides clues to understanding the failure of AT in OOD scenarios.
The 3 For simplicity, we denote ‘AT’ distribution shifts in OOD data usually have much larger scales than the invisible perturbations commonly used in AT.
Hence, AT methods designed for small perturbations cannot handle these large-scale domain shifts that often appear in OOD data.
However, our experiment shows that this problem can be alleviated by adapting universal AT (UAT [8]) to the OOD setting, i.e., using a perturbation for each domain.
[Figure]
Figure 1 shows that UAT remains its generalization performance when the perturbation scale is large.
There are two empirical explanations for this: First, the background and style information usually have a low-rank structure, such as the grassland and snowfield that have recurring parts.
Second, similar spurious features often appear within one specific domain, such as PACS [14] and VLCS [16] datasets.
As stated in [8], the universal perturbation lies in a low dimensional space.
Hence using universal (domain-wise) perturbations will help to resist such low-rank shifts and improve the robustness of the model.
Inspired by this, we proposed two new AT variants with more sophisticated low-rank structures on different dimensions to improve OOD generalization in the next section.
The Proposed Structured AT Method
In order to construct low-rank structured perturbations, we start by analyzing the structure of samplewise perturbations.
Assume that each input data x has a shape of N × N × C .
N is the size of the input image and C is the number of channels.
For simplicity, we assume C = 3 .
We reparameterize the sample-wise perturbations as a series of 2-D matrices { D 11 , D 12 , ..., D 1 m } , { D 21 , D 22 , ..., D 2 m } , { D 31 , D 32 , ..., D 3 m } where D ce ∈ R n e × N 2 denotes the perturbations in the e -th domain for the input channel c , n e is the number of the samples in the domain E e , and m is the number of domains.
The i -th row of D ce represents the c -th channel of the i -th sample in the domain E e (see the first column in Figure 2 for illustration).
By such reparameterization, it is natural to find that there are two orientations to reduce the rank of the perturbations:
Along the dimension of the number of samples (along the red arrow in the upper left corner of Figure 2).
This corresponds to reducing the number of the perturbations used within one domain.
Along the dimension of the input scale (along the blue arrow in the upper left corner of Figure 2).
This corresponds to reducing the rank of the perturbation used for a specific input sample.
In the following parts, we propose two AT variants with structured priors that reduce the rank in these two directions.
4.1 MAT: Adversarial Training with Combinations of Multiple Perturbations
In this part, we propose domain-wise Multiple-perturbation Adversarial Training (MAT).
It aims to conduct rank minimization along the dimension of the number of samples.
Instead of using sample-wise perturbations, MAT constructs a combination of multiple perturbations and shares this mixed perturbation within a domain.
Specifically, we choose to train the linear combination of k perturbations for each domain E e to conduct AT.
Here k is a hyperparameter and k is far less than the number of samples in domain E e .
The optimization problem can be reformulated as:
[Equation]
Domain (cid:8) (cid:2869) …
Sample-wise AT (cid:1866) (cid:2869) (cid:1830) (cid:2869) (cid:1840)(cid:3400)(cid:1840) (rank= (cid:1840) )
Domain (cid:8) (cid:2870) (cid:1866) (cid:2870) (cid:1830) (cid:2870)
Domain (cid:8) (cid:2871) (cid:1866) (cid:2871) (cid:1830) (cid:2871) …
MAT (cid:1840)(cid:3400)(cid:1840) (rank= (cid:1840) )
(cid:1863) (cid:1863) (cid:1863) …
LDAT (cid:1840)(cid:3400)(cid:1840) (rank= (cid:1864) )
(cid:883) (cid:883) (cid:883) …
where
[Equation]
Here e denotes the subscript of a training domain and α ei is the weight that can be learned for each perturbation δ ei .
The detailed training procedure of MAT is in Algorithm 1.
We first initialize k perturbations δ ei and their correspondent coefficients α ei for each training domain e with Gaussian noise.
Then we transform δ ei and α ei to make sure (cid:80) ki
=1 α ei = 1 , α ei ≥ 0 , and || δ ei || 2 ≤ ϵ .
For the inner maximization, we conduct a one-step gradient ascent to optimize δ i and α i .
MAT works as a low-rank version of sample-wise AT.
In sample-wise AT, we maintain n e perturbations for each domain E e , where n e is the number of training samples in domain E e .
As for MAT, it reduces the number of perturbations available to samples from n e to k and obtains low-rank structures (see the third column in Figure 2 for illustration).
Therefore, it fulfills rank reduction along the sample-number dimension.
4.2 LDAT: Adversarial Training with Low-rank Decomposed Perturbations
Based on MAT, we further propose Adversarial Training with Low-rank Decomposed perturbations (LDAT).
Analogous to MAT, LDAT still shares one perturbation in a specific domain.
Moreover,
LDAT imposes a low-rank constraint on the perturbation itself, which corresponds to the dimension of the input scale.
Technically, we obtain the domain-wise low-rank perturbation matrix δ ∈ R N × N × C by multiplying two matrices: δ = AB .
Here A ∈ R N × l × C and B ∈ R l × N × C where l is a hyperparameter and l ≪ N .
Since rank ( AB ) ≤ rank ( A ) and rank ( AB ) ≤ rank ( B ) hold for arbitrary matrices
A , B , we have rank ( δ ) ≤ l .
Therefore LDAT reduces the rank of the perturbation from a large value N to a relatively small value l (see the last column in Figure 2 for illustration). 1: Randomly initiate θ , perturbation δ ei weight α ei such that (cid:80) ki =1 α ei = 1 , α ei ≥ 0 , || δ ei || 2 ≤ ϵ , ∀ i ∈ {
..., k } and ∀ e ∈ { m } . do 3: do 4: Randomly select batch B Compute the adversarial sample: x e ′ u
= x eu + (cid:80) kj =1 α
ej δ ej , ∀ u ∈ { 1 , ..., b } 6:
Update δ e by δ ei ← δ ei + γ 1 b (cid:80) bu =1 ∇ δ ei L ( f θ ( x e ′ u ) , y eu ) , ∀ i ∈ { 1 , ..., k } , ∀ u ∈ { 1 , ..., b } .
Update α ei by α ei ← α ei + η 1 b (cid:80) bu =1 ∇ α ei L ( f θ ( x e ′ u ) , y eu ) , ∀ i ∈ { 1 , ..., k } , ∀ u ∈ { 1 , ..., b } .
Project δ ei to the l 2 ball of radius ϵ .
Compute the adversarial sample: x e ′ u
= x eu + (cid:80) kj =1 α
ej δ ej , ∀ u ∈ { 1 , ..., b } 10:
Update model parameter: θ ← θ − r 1 b (cid:80) bu =1 ∇ θ L ( f θ ( x e ′ u ) , y eu ) , ∀ u ∈ { 1 , ..., b } .
11: end for 12: end for The formal definition of the LDAT objective is:
[Equation]
where
[Equation]
We provide the detailed training procedure of LDAT in Appendix D due to the space limitation of the main text. In comparison to MAT, LDAT reduces the number of perturbations available to the samples in a domain from k to 1.
In addition, it reduces the rank of the perturbation for a single channel of a sample from N to l .
4.3 Theoretical Analysis
In this part, we theoretically explain why the domain-wise perturbation proposed in MAT and LDAT can help to improve the robustness of the model against spurious correlations following [2] and [19].
In general, we prove that MAT and LDAT can prevent the model from relying more on spurious features to make predictions as the spurious correlations in the training data increase.
Consequently, the model trained with MAT or LDAT will generalize better on OOD data.
Notations.
Let x ∈ X denote the random data in the input space X and let y ∈ Y denote the target random data in label space Y .
For simplicity, let Y ∈ { 1 , − 1 } in this section.
Let D denote an underlying class of distributions over X × Y .
Let x inv and x sp denote the invariant features and the spurious features respectively.
Also for simplicity, assume that there exists an identity mapping Φ : X inv × X sp → X such that each D ∈ D is induced by a distribution over X inv × X sp (so x can be denoted as x = ( x inv , x sp ) ).
Let x sp take values in { + β, − β } for some β > 0 .
A Simple OOD Task.
Consider a simple OOD task where we have two training domains representing the grass and desert backgrounds respectively.
Both domains have two classes: the cow class and the camel class.
In the grass/desert domain, the cow/camel class predominates.
During test time, the correlation between the labels and the background flips.
We can abstract this cow-camel dataset into the following model: a training dataset S with four groups of data points drawn from the four quadrants of the feature space {− 1 , +1 } × {− β, + β } respectively (shown in Figure 3).
We set the invariant features x inv =
y and the spurious features x sp to be yβ with probability p ∈
[0 . 5 , 1) and − yβ with probability 1 − p .
Note that p measures the intensity of spurious correlations in a certain environment.
When p = 0 .
5 , there are no correlations between the labels and the spurious features.
(cid:28646)(cid:28675)(cid:28680)(cid:28677)(cid:28668)(cid:28674)(cid:28680)(cid:28678) (cid:28636)(cid:28673)(cid:28681)(cid:28660)(cid:28677)(cid:28668)(cid:28660)(cid:28673)(cid:28679)
[Figure] Consider a linear classifier h ( x ) =
w inv x inv + w sp x sp .
Following [19], let us consider MAT/LDAT trained with gradient descent algorithm stopped in finite time t .
In order to characterize the dependence of the model on spurious features during the training process, we investigate the convergence rate of
[Equation] [Equation]
inv | to 0 on the above dataset, which denotes the ratio between the output of the spurious component to that of the invariant component.
We prove that after adding the domain-wise perturbations in finite-time-stopped gradient descent, the lower bound of the convergence rate of this ratio does not increase monotonically with p .
Hence, the model will not learn a large prediction weight based on spurious features even if the spurious correlation is strong ( p is large).
In the following theorem, we denote the domain-wise perturbation in MAT/LDAT as δ .
Theorem 4.1 applies to both MAT and LDAT since they both use domain-wise perturbations.
See Appendix A for a formal statement and full proof of Theorem 4.1.
Theorem 4.1.
(informal)
Let H be the set of linear classifiers h ( x ) =
w inv ( t ) x inv + w sp ( t ) x sp .
Consider the above 2-D OOD dataset S .
Assume that the empirical distribution of x inv given x sp · y > 0 is identical to the empirical distribution of x inv given x sp · y < 0 .
δ is the optimal perturbation obtained by optimizing the object in Eq.
(7) or Eq.
Let w inv ( t ) x inv + w sp ( t ) x sp be initialized to the origin, and trained with MAT/LDAT to minimize the exponential loss on S .
Then, for any ( x, y ) ∈ S , we have:
[Equation]
where ϵ :
= δy 2 β is a real number close to 0, c 1 := 2(2 M (1+ δ ) − 1) ( β + δy )
[Equation]
M = max x ∈ S ˆ w · x denotes the maximum value of the margin of the max-margin classifier ˆ w on S .
Ω( · ) is the lower bound of a given function within a constant factor.
Therefore, the lower bound of the convergence rate does not increase monotonically with p under the condition that
[-2 + ( c 1 1 − c 1 ) 
-][Equation]
< 0 .
To sum up, since we can prevent this lower bound from growing monotonically with p , we accelerate the convergence rate of w [Equation] |
[Equation]
inv [Equation] t ) x inv | to 0 when there is stronger spurious correlation (larger p ).
Recall that the ratio
[Equation]
| reflects the degree of reliance on spurious features.
Therefore, faster convergence of this ratio to 0 (smaller lower bound) means that the model will end up relying less on the spurious correlations within a finite training time.
In other words, the OOD robustness can be enhanced by using domain-wise perturbations.
Remark.
Here, we demonstrate that MAT and LDAT show stronger OOD robustness compared to ERM.
We compare the result in Theorem 4.1 to that in Theorem 2 of [2].
The full statement of the Theorem 2 in [2] is in Appendix B. According to the Theorem 2 in [2], even if the max-margin classifier does not rely on x sp for any level of spurious correlation p ∈ [0 . 5 , 1) , ERM trained by gradient descent stopped in finite time still fails to avoid using spurious features.
Moreover, when conducting ERM with finite-time-stopped gradient descent, the lower bound of the convergence rate of w sp ( t ) β | w inv ( t ) x inv | to 0 is
[Equation]
(11) where c := 2(2 M − 1) β, and M follows the definition in Theorem 4.1.
This lower bound grows monotonically with p , thus ERM will have slower convergence for larger spurious correlations.
However, with domain-wise perturbations, we can modify the lower bound so that it does not increase monotonically with the spurious correlation p .
Thus, we can draw the conclusion that using a perturbation for each domain is helpful to reduce dependence on spurious features compared to ERM.
Experiments 5.1 Experimental Setup
We conduct experiments on the DomainBed benchmark [9], a testbed for OOD generalization that implements consistent experimental protocols across various approaches to ensure fair comparisons.
We evaluate on PACS [14], OfficeHome [15], VLCS [16], NICO [17], and Colored MNIST [1].
There are several changes in our experimentation setting comparing to DomainBed:
Backbone Network.
We use ResNet-18 [20] as our backbone for datasets excluding Colored MNIST instead of ResNet-50 used in [9] for efficiency.
Hyperparameter Search Space.
We use a smaller hyperparameter search space than [9].
We conduct a random search of 8 trials for PACS, OfficeHome, and VLCS while 6 trials for NICO and Colored MNIST in the hyperparameter search space, instead of 20 trials adopted in [9] for feasibility.
See Appendix C.2 for more details.
Model Selection Strategy.
Since hyperparameter choice has a significant impact on the OOD performance, it is critical to use appropriate model selection method.
For PACS, OfficeHome, and
VLCS datasets, we use training-domain validation proposed in [9] since it is more in line with the
OOD scenario.
For NICO, we adopt OOD validation following [21].
For Colored MNIST, we use test-domain validation [9] since it can enlarge the gaps in OOD performance among the algorithms while the gap induced by training-domain validation on Colored MNIST is marginal.
Hyperparameters for MAT and LDAT.
To retain low-rank structures in perturbations, we set the upper bound of the search space of the perturbation number k in MAT to be 20.
Similarly, the upper bound of the rank of the perturbation used in LDAT l is 20.
Specifically, the search space of k and l is { 5 , 10 , 15 , 20 } (except on CMNIST, where k ∈ [5 , 20] and l ∈ [10 , 20] ).
The complete setup of the hyperparameters for MAT and LDAT is provided in Appendix C.2.
Table 2 summarizes the results on the five OOD datasets.
The results of other approaches for PACS,
OfficeHome, NICO, and Colored MNIST datasets are adopted from [21].
The results on VLCS of other algorithms are missing (denoted as "-") because [21] does not experiment on this dataset.
Comparison with ERM and Sample-wise AT.
From Table 2, we observe that both MAT and
LDAT outperform ERM (on both our runs and the results in [21]) and AT on average.
In particular,
MAT achieves consistently better results than ERM on all five datasets.
Additionally, the average performance of AT is worse than ERM, which is consistent with our observations in Section 3.2.
Comparison with Existing State-of-the-Art Approaches.
Although the results from [21] use a different training protocol from ours: they use a larger search space and 20 random search for the hyperparameter combinations, the comparison between ERM ([21]) and ERM (our runs) indicates that their corresponding performances are close.
A similar comparison has been made in [22].
We find that MAT outperforms all previous algorithms and LDAT ranked fourth among all methods, merely after VREx [23] and IRM [1] (see avg 1 in Table 2). And even when excluding Colored MNIST (toy example), our methods still outperform ERM by 0 . 4 ∼ 1 .
4% , whereas other methods show no improvement over ERM (see avg 3 ).
From these results, we can see that the promotion of our proposed methods is higher than the other works, and our methods clearly outperform ERM.
We also extend our evaluation to compare with adversarial augmentation based method [5] in Appendix C.3.
Single-training domain generalization experiments are shown in Appendix C.4, which shows that our methods can maintain OOD performance without the reliance on multi-source training data.
Comparison between MAT and LDAT.
From Table 2 we can see that MAT outperforms LDAT on average.
Since LDAT reduces the number of the perturbations used in a domain from k to(shown in Figure 2), LDAT can be regarded as a low-rank version of MAT.
This indicates that the oversimplified perturbations may be less effective than the ones maintaining some flexibility. [Table] Origin
ERM AT MAT LDAT
Since both MAT and LDAT outperform most existing state-of-the-art methods and they both exploit low-rank structures, these two methods mutually corroborate the effectiveness of low-rank structure for OOD generalization. The respective advantages of the two methods are as follows: • MAT is a complex and high-rank version of LDAT, which has a stronger ability to describe more complex spurious background information.
As shown by the attention heatmap in Figure 4, MAT can better capture the object than LDAT when faced with a more complex background (the example in the second row of Figure 4).
• LDAT costs less memory than MAT during the training process, although there is no significant difference in training time between the two methods.
When the memory is limited, LDAT is preferred.
5.3 Empirical Understanding
Visualization.
To empirically show that MAT and LDAT can reduce the reliance on spurious features, we visualize the pixel attention heatmap of ERM, AT, MAT, and LDAT on NICO dataset using
GradCam [34].
It reflects the contribution of different components of the feature map to the prediction results.
We pick the model with the best performance for each method.
The results in Figureindicate that the model trained by MAT and LDAT focuses more on the object itself, while ERM and
AT adopt the background information that spuriously correlates to the class to make predictions.
Parameter Analysis.
The number of the perturbations k used in a domain in MAT and the rank of the perturbation l in LDAT are two key hyperparameters.
We conduct further experiments to analyze the impact on the performances of these two parameters.
We adopt a fixed set of parameters except for k and l and evaluate on PACS dataset.
The results in Figure 5 show that MAT and LDAT are able to keep their performances over ERM as long as k and l are far less than the number of the samples N .
Additionally, we can observe from the trend that when k and l are too small ( = 5 ), the performances degenerate.
This implies the oversimplified structures of the perturbations can be less effective for generalization.
When k and l take larger value (about 1000), the performance will drop (see Table 3).
For the selection of optimal parameters (range), we observe that a set of parameters that are good on one dataset also work well on other datasets (see Table 6 in the appendix), so in practice we adopt the strategy of searching for the optimal parameters roughly on one dataset and then applying them to other datasets.
Additional analysis on the impact of the learning rate for the perturbations is in Appendix C.2. [Figure] MAT) and the rank of the perturbation in a domain l (in LDAT). [Table]
Conclusion
In this work, we empirically revealed the limitations of sample-wise AT on OOD tasks.
Due to the lack of constraints on the perturbation and the utilization of domain features, sample-wise AT fails to generalize well when facing large-scale perturbations which is close to the real-world OOD scenarios.
We further proposed two AT variants with structured priors, named MAT and LDAT, which add low-rank perturbations to improve model’s robustness against the distribution shift of spurious correlations. We theoretically proved the domain-wise perturbations used in MAT and LDAT can benefit OOD generalization, and validated the effectiveness of the proposed methods on OOD tasks through a series of experiments on Domainbed benchmark. Acknowledgment
Yisen Wang is partially supported by the NSF China (No. 62006153), Project 2020BD006 supported by PKU-Baidu Fund, Open Research Projects of Zhejiang Lab (No. 2022RC0AB05), and Huawei Technologies Inc.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information on how to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or [N/A] .
You are strongly encouraged to include a justification to your answer , either by referencing the appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes] See Section .
• Did you include the license to the code and datasets?
[No] The code and the data are proprietary.
• Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below.
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] See Section 1.
(b) Did you describe the limitations of your work?
[Yes] See Section 5.2.
We mention that the oversimplified perturbation (LDAT) will be less effective than the ones with more flexibility (MAT).
(c) Did you discuss any potential negative societal impacts of your work?
[No] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] Full assumptions are in Appendix A. (b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix A. 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[No] It will release upon acceptance.
(b) Did you specify all the training details (e g., data splits, hyperparameters, how they were chosen)?
[Yes] See Section C.1 and C.2.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] See Table 2 for an example.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[No] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] We use DomainBed benchmark and GradCam and cite their creators.
(b) Did you mention the license of the assets?
[No] All assets we use are open source.
(c) Did you include any new assets either in the supplemental material or as a URL?
[No] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[No] All datasets we use are open source.
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[No] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[No] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[No] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[No]
A Proof of Theorem 4.
Theorem 4.1 (formal) Let H be the set of linear classifiers h ( x ) = w inv ( t ) x inv + w sp ( t ) x sp . Consider any task that satisfies all the constraints in Section 3.1. in [2].
Consider a dataset S drawn from D such that the empirical distribution of x inv given x sp · y > 0 (denoted as ( x inv , y ) ∼ D inv ) is identical to the empirical distribution of x inv given x sp · y < 0 .
δ is the optimal perturbation obtained by optimizing object (7) or (9).
δ can be seen as a random variable.
Let w inv ( t ) x inv + w sp ( t ) x sp be initialized to the origin, and trained with MAT/LDAT with an infinitesimal learning rate to minimize the exponential loss on S .
Then, for any ( x, y ) ∈ S , we have:
[Equation]
where ϵ :
= δy 2 β is a real number close to 0, c 1 :=
2(2 M (1+ δ ) − 1) ( β + δy )
[Equation]
M = max x ∈ S ˆ w · x denotes the maximum value of the margin of the max-margin classifier ˆ w on
S .
Ω( · ) is the lower bound of a given function within a constant factor.
Therefore, the lower bound of the convergence rate does not increase monotonically with p under the condition that
[Equation]
ϵc
[Equation]
+ 32 ϵ < 0 .
Proof.
For brevity, we use w c and w s to represent w inv and w sp respectively.
Also, we use x c to represent x inv .
We use x e to represent x sp since the spurious feature is correlated to the environment
E e .
Let S min and S maj denote the subset of datapoints in S where x e · y < 0 and x e · y >respectively.
Let D inv denote the distribution over ( x c , y ) induced by drawing ( x, y ) uniformly from
S min .
The corresponding marginal distribution of y is denoted as D y .
By the assumption of the theorem, this distribution would be the same if x was drawn uniformly from S maj .
Then, the loss function that is being minimized in this setting corresponds to:
[Equation]
The update on w s can be written as:
[Equation]
Proof of bounds on w c ( t ) x c .
Using the result of [19] and [2], we get
[Equation]
for a sufficiently large t and for all x ∈ S .
Proof of the upper bound on w s .
To calculate the lower bound of w s , we prove the upper bound as auxiliary first.
Note that ∆ w s decreases monotonically with w s .
Assume that β > | δ | (this is reasonable since the perturbation radius is usually smaller than the scale of the spurious feature).
Let ∆ w s = 0 , we get
[Equation]
Since ∆ w s decrease monotonically with w s , (which can be inferred from ∂ ∆ w s ∂w s = − e − ( x c + δ ) w c y [ − pe − ( δy + β ) w s ( δy + β ) 2 + (1 − p ) e ( β − δy )
w s ( β − δy ) 2 ] ≤ 0 ), when w s < w 0 , ∆ w s > 0 and when w s > w 0 , ∆ w s < 0 .
As a result, for any system that is initialized at 0, w s can never cross the point w 0 .
Thus, we get the upper bound of w s ( t ) :
[Equation]
Proof of the lower bound on w s .
We lower bound w s via the upper bound on w s as:
[Equation]
Next, using the upper bound on | w c ( t ) x c | , we get:
[Equation]
For brevity, we denote the term p ( δy + β ) e − ( δy + β ) w s as T and denote the term [ p ( β + δy )] 12 − δy 2 β
[(1 − p )( β − δy )]
12 + δy 2 β as L .
And in the following proof, we omit the expectation marker E also for simplicity.
It is clear that both T and L > 0 .
Here, T − L > 0 since w s < w 0 .
Note that ∆ w s = − ∂ L [ f ( x ) ,y ] ∂w s =
∂w s ∂t , rearranging this and integrating, we get:
[Equation]
Since for a sufficiently large t , 1 − 1 (1+ t
[Equation]
2 M
[Equation]
[Equation]
δ
[Equation]
1 > 12 , we have:
[Equation]
we can further lower bound the right hand side by applying the inequality x ≥
ln( x + 1) for positive x :
[Equation]
Thus,
[Equation]
Note that the denominator on the left side of the inequality is greater than 0 since T − L > 0 .
Rearrange this inequality:
[Equation]
Putting L =
[ p ( β + δy )] 12 − δy 2 β
[(1 − p )( β − δy )]
12 + δy 2 β back into the inequality,
[Equation]
Let c 1 := 2(2 M (1+ δ ) − 1) ( β + δy ), c 2 :=
[Equation]
M (1+
[Equation]
−
[Equation]
( δy
[Equation]
β
[Equation]
32 − ϵ
[Equation]
β
[Equation]
12+
[Equation]
ϵ
[Equation]
= δy 2 β .
Put the expectation mark back into this inequality, finally, we get the lower bound on w s :
[Equation]
To show that the lower bound on the dependency on spurious correlations induced by MAT and LDAT does not increase monotonically with p under some conditions, we take the derivative of the obtained lower bound g ( p ) in Theorem 4.1 with respect to p :
[Equation]
Since the denominator of ∂g ( p ) ∂p is positive, we pick out the numerator: c 2 +( 12 + ϵ ) p 12 − ϵ (1 − p ) 12 + ϵ − c 1 ( 12 − ϵ )
{+p − 12 − ϵ (1 − p ) 
+}
{++ ϵ + c 1 ( 12 + ϵ ) p 12 − ϵ (1 − p ) − 
+}
{++ ϵ + ( 12 + ϵ ) p 32 − ϵ (1 − p ) − 
+}
+ ϵ .
In order to study the positive and negative change of the numerator, we continue to derive it with respect to p and obtain p
[Equation]
Since we assume β > | δ | and c 1 > 0 , ∂ 2 g ( p ) ∂p> 0 and ∂g ( p ) ∂p increase with p monotonically.
The minimum of ∂g ( p ) ∂p is reached when p = 12 .
This minimum equals to 2 ϵc 1 + c 2 + 34 + 32 ϵ .
When 2 ϵc 1 + c 2 + 3 4 + 3 2 ϵ < 0 , the lower bound does not increase with p monotonically when p is within a certain range ( ∈ (0 . 5 , 1) ).
B Detailed Statement of Theorem 2 in [2]
We now introduce the Theorem 2 in [2].
Before introducing it, we first introduce the concept of the easy-to-learn tasks in [2], i.e. tasks with a set of constraints.
The motivation of restricting ourselves to the constrained set of tasks is that it prevents us from designing complex examples where ERM is forced to rely on spurious features due to a not-so-fundamental factor.
Each constraint forbids a specific failure mode of ERM in OOD scenarios.
The Theorem 2 in [2] shows that even under such favorable conditions for ERM, this classical method can also be perturbed by the spurious features.
Notations.
For convenience, we will give some notations here again.
Consider an input space
X and a label space Y ∈ {− 1 , 1 } .
Let D ∈ D denote a distribution over X × Y .
p D denotes the probability density function (PDF) of D .
Let H denote a class of classifiers h : X → R .
Consider a dataset S drawn from D .
Let L D ( h ) :
=
E ( x,y ) ∼D
[ h ( x ) · y < 0] the loss of h on D .
Let h ⋆ = arg
min h ∈ H max D∈ D L D ( h ) denote the optimal classifier in the worst case.
With an abuse of notation, we also denote the PDF of the distribution over X inv × X sp as p D ( · ) .
Let D train denote the distribution of the pooled training data.
Assume that there exists a mapping Φ : X inv × X sp → X such that each D ∈ D is induced by a distribution over X inv × X sp .
Definition B.1.
Easy-to-learn tasks .
Tasks that satisfy the following constraints are easy-to-learn.
Fully predictive invariant features. For all D ∈ D , L D ( h ⋆ ) = 0 . Identical invariant distribution. Across all D ∈ D , p D ( x inv ) is identical.
Conditional independence.
For all D ∈ D , x sp ⊥ x inv .
Two-valued spurious features.
We set x sp =
R and the support of x sp in D train is {− β, + β } .
Identity mapping.
Φ is the identity mapping i.e., x = ( x inv , x sp ) .
Theorem B.2.
(The Theorem 2 in [2]) Let H be the set of linear classifiers h ( x ) =
w inv · x inv + w sp x sp .
Then, for any easy-to-learn task, continuous-time gradient descent training of w inv ( t ) x inv + w sp ( t ) x sp to minimize the exponential loss, satisfies:
[Equation]
where M = max x ∈ S
ˆ wx where ˆ w is the max-margin classifier on S .
c := 2(2 M − 1) β.
[Table]
C Experiment Details and Supplementary Experimental Results C.1 Settings of the Toy Experiments
For the experiment in Table 1, Figure 1, Figure 5, Table 7, and Table 8, we use a fixed set of hyperparameters (see Table 4) instead of conducting a random search of 20 trials over the hyperparameter distribution (the setting in [9]) for efficiency.
We report the average across three independent runs.
For model selection method, training-domain validation [9] is used for PACS, OfficeHome, and
VLCS.
For NICO, an OOD validation set is adopted following [21].
C.2 Experiment Setting and Additional Results of Table
Overall setup.
We conduct a random search of 8 trials for PACS, OfficeHome, VLCS and 6 random trials for NICO and Colored MNIST in the hyperparameter search space, instead of 20 trials adopted in [9] for feasibility.
We then average the best results for each hyperparameter combination and dataset (according to each model selection criterion) across test domains (except for Colored MNIST where we test on one biased domain only).
Finally, we report the average of this number across three independent runs, and its corresponding standard error.
We run all datasets for 8000 epochs during the training process.
Hyperparameter Search Space.
We use a smaller hyperparameter search space than that in [9].
The search space for PACS, OfficeHome, VLCS, NICO and Colored MNIST is shown in Table 5.
To determine the search space of a hyperparameter for benchmark running, we first fix other parameters and conduct a grid search to determine the approximate range of the better performances.
Take the learning rate η of the MAT matrices as an example, we fix k = 20 and try different values of η on
PACS.
We find that the results of η = 0 .
01 (82.2 ± 0.4%) and η = 0 .
001 (82.3 ± 0.5%) are better than that of η = 0 .
1 (81.6 ± 0.2%), so we adopt the random search space of { 0 . 01 , 0 . 001 } .
The same is true for the other parameters.
In practical applications, as for the choice of the optimal parameter, we find that through experiments that for MAT and LDAT, the value of the rank k (MAT) and l (LDAT) with good test accuracy (outperforms ERM) on one data set also has good one on other datasets, as shown in the Table 6.
In
Table 6, k = 10 for MAT and l = 15 for LDAT outperform ERM on all three datasets, as marked in bold .
Thus, we could find an optimal set of parameters with the model selection methods and then apply them to other datasets.
Model Selection Stategy.
For PACS, OfficeHome and VLCS datasets, we use training-domain validation proposed in [9].
This model selection method first randomly collect 20% of each training domain to form a validation set.
Then, it chooses the hyperparameter maximizing the accuracy on the validation set.
For NICO, we adopt the OOD validation proposed in [21].
This method chooses the model maximizing the accuracy on a validation set that follows neither the distribution of the training domain or the distribution of the test domain.
For Colored MNIST, we use test-domain validation, i.e., using a validation set that follows the distribution of the test domain.
This is because it can enlarge the gaps in OOD performance among the algorithms while the gap induced by training-domain validation on Colored MNIST is marginal.
[Table]
All except Colored
[Table]
[Table]
Backbone Network.
We use ResNet-18 [20] pretrained on ImageNet [35] for PACS, OfficeHome and VLCS.
We use unpretrained ResNet-18 for NICO since it contains images largely overlapped with ImageNet classes.
As for Colored MNIST, We use a small CNN-architecture following [9].
Impact of Learning Rate.
We further investigate the impact on OOD performances of the learning rate for the perturbation weights in MAT and the learning rate for the decomposed factors in LDAT.
We use the experimental setting introduced in Appendix C.1.
The results of MAT and LDAT are shown in Table 7 and 8 respectively.
In Table 7 and 8 we observe that the learning rate for the
[Table]
[Table]
perturbations has a marginal effect on the OOD accuracy.
Both MAT and LDAT outperform ERM and AT on PACS when using different values of learning rate.
C.3 Comparing to Existing Data Augmentation Baseline
To better verify the improvement of our proposed method on the existing adversarial augmentation methods for OOD, we reproduce the algorithm in [5].
[5] proposed a minimax iterative training procedure to generate adversarial data that follows fictitious target distributions (GUT).
As discussed in Section 2, their work is restricted in the framework of using Wasserstein distance to measure the distribution shift, which is less practical for the real-world OOD setting where domain shifts are diverse.
Additionally, They focus only on sample-wise operations and ignore the use of common features within a domain.
The experimental results on NICO dataset is in Table 9.
The unique hyperparameters of GUT follow the Settings in [5] except we set the T max to be 5 instead of 15 for efficiency.
We can see that both our proposed method outperform GUT.
C.4 Comparing to Existing Data Augmentation Baseline under Single-training Domain Generalization Setting
We conduct experiments to further verify the effectiveness of MAT and LDAT under single-training domain generalization setting, i.e., using only one training domain and generalize on the others.
We compare our work with Neuron Coverage-Guided Domain Generalization (NCDG) [36].
The results are in Table 10.
Both MAT and LDAT outperform NCDG under the scenario of single-source domain generalization.
D Detailed Description of LDAT
In this section, we describe the detailed training procedure of LDAT (see Algorithm 2).
We conduct a single-step gradient ascent for the inner maximization for the perturbations LDAT.
We adopt l 2 norm for the perturbations.
[Table]
[Table]
Algorithm
Require:
Labeled training data of m domains E 1 , ..., E m , where E e := { ( x ei , y ei ) } n e i =1 , rank of the perturbations l , factor A , B learning rate ρ ,
FGSM step size γ , perturbation radius ϵ , number of training epochs T , learning rate for model parameters r , batch size b .
Ensure:
Updated model f θ with parameter θ .
1: Randomly initiate θ , perturbation δ e , factor A e , B e such that || δ e || 2 ≤ ϵ , ∀ e ∈ { 1 , ..., m } .
2: for iterations in 1 , 2 , ..., T do 3: for e in 1 , 2 , ..., m do 4:
Randomly select batch B e = { ( x eu , y eu ) } bu =1 from domain E e .
Compute the adversarial sample: x e ′ u
= x eu + A e B e , ∀ u ∈ { 1 , ..., b } 6:
Update A e by A e ← A e + ρ 1 b (cid:80) bu =1 ∇ A e L ( f θ ( x e ′ u ) , y eu ) , ∀ u ∈ { 1 , ..., b } .
Update B e by B e ← B e + ρ 1 b (cid:80) bu =1 ∇ B e L ( f θ ( x e ′ u ) , y eu ) , ∀ u ∈ { 1 , ..., b } .
Project δ ei to the l 2 ball of radius ϵ .
Compute the adversarial sample: x e ′ u
= x eu + A e B e , ∀ u ∈ { 1 , ..., b } 10:
Update model parameter: θ ← θ − r 1 b (cid:80) bu =1 ∇ θ L ( f θ ( x e ′ u ) , y e u ) , ∀ u ∈ { 1 , ..., b } .
11: end for 12: end for