Understanding and Improving Feature Learning for Out-of-Distribution Generalization
Authors: Yongqiang Chen
Wei Huang
Kaiwen Zhou
Yatao Bian
Bo Han
James Cheng
Venue: NeurIPS 2023 poster (2023)
Keywords: Out-of-Distribution Generalization
Feature Learning
Invariant Risk Minimization
TL;DR: We show ERM learns both invariant and spurious features and propose a new algorithm to learn richer features than ERM for facilitating OOD generalization.

Abstract: A common explanation for the failure of out-of-distribution (OOD) generalization is that the model trained with empirical risk minimization (ERM) learns spurious features instead of invariant features. However, several recent studies challenged this explanation and found that deep networks may have already learned sufficiently good features for OOD generalization. Despite the contradictions at first glance, we theoretically show that ERM essentially learns both spurious and invariant features, while ERM tends to learn spurious features faster if the spurious correlation is stronger. Moreover, when fed the ERM learned features to the OOD objectives, the invariant feature learning quality significantly affects the final OOD performance, as OOD objectives rarely learn new features. Therefore, ERM feature learning can be a bottleneck to OOD generalization. To alleviate the reliance, we propose Feature Augmented Training (FeAT), to enforce the model to learn richer features ready for OOD generalization. FeAT iteratively augments the model to learn new features while retaining the already learned features. In each round, the retention and augmentation operations are performed on different subsets of the training data that capture distinct features. Extensive experiments show that FeAT effectively learns richer features thus boosting the performance of various OOD objectives.

Understanding and Improving Feature Learning for Out-of-Distribution Generalization
Yongqiang Chen', Wei Huang?: Kaiwen Zhou'
'hhe Chinese University of Hong Kong PRIKEN AIP
Kyqchen, kwzhou, cheng) (cseennn/ese.cuke.e..e ecu.. weihuang.vrerikenn.jpp
Yatao Bian", Bo Han James Cheng!
Tencent Al Lab "Hong Kong Baptist University
yataobbaangmaiil.cmm bhanm1@comp [haam11@comphkkuueduu...
Abstract
A common explanation for the failure of out-of- of-distribution (OOD) generalization is that the model trained with empirical risk minimization (ERM) learns spurious features instead of invariant features. However, several recent studies challenged this explanation and found that deep networks may have already learned sufficiently good features for OOD generalization. Despite the contradictions at first glance, we theoretically show that ERM essentially learns both spurious and invariant features, while ERM tends learn spurious features faster if the spurious correlation is stronger. Moreover, when fed the ERM learned features to the OOD objectives, the invariant feature learning quality significantly affects the final OOD performance, as OOD objectives rarely learn new features. Therefore, ERM feature learning can be bottleneck to OOD generalization. To alleviate the reliance, propose Feature Augmented Training (FeAT), to enforce the model to learn richer features ready for OOD generalization. FeAT iteratively augments the model learn new features while retaining the already learned features. In each round, the retention and augmentation operations are performed different subsets of the training data that capture distinct features. Extensive experiments show that FeAT effectively learns richer features thus boosting the performance of various OOD objectives!.
Introduction
Understanding feature learning in neural networks is crucial to understanding how they generalize to different data distributions [2, 11, 12, 62, 67, 70]. Deep networks trained with empirical risk mini- mization (ERM) learn highly predictive features that generalize surprisingly well to in in-distribution (ID) data [24, 75J. However, ERM also tends to learn spurious features or shortcuts such as image backgrounds [7. 19, 23, 84] whose correlations with labels do not hold in the out-o--distribution (OOD) data, and suffers from serious performance degeneration [39]. Therefore, it widely believed that the reason for the OOD failures of deep networks is that ERM fails to learn the desired features that have invariant correlations with labels across different distributions [7].
However, several recent works find that ERM-trained models have alreudy learned sufficiently good features that are able to generalize to OOD data [33, 38, 63]. addition, when optimizing various penalty terms [1, 15, 40, 41, 53, 55, 57, 61, 69, 76, 85] that aim to regularize ERM to capture the invariant features (termed as OOD objectives). there also exists curious phenomenon that the
"qqual Contribution. Work done during Yongqiang's internship Tencent Al Lab.
1 Cooe available at /attps://githan wwt.//oasse/ee..eee
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

InMvt
VSF
PeAT
Featurizer vp
IRMv1
BRAOX
C1
Spurions Features Leursd Fentures
Invariant Features Linderlying Features
IGA
RAY
Featurizer Yp
w
ERM
IRMv1
Pre-training
OOD training
pretrain epochs
(a) An illustration of FeAT (top row) compared to ERM (bottom row). Different colors in samples denote the respective dominant features. As the original data dominated by spurious features (blue), ERM tends to learn more spurious features but limited invariant features (orange). Thus the OOD training with IRMv1 can only leverage limited invariant features and achieve limited performance. In contrast, iteratively, FeAT divides Der into augmentation D and retention sets D" that contain features not learned and already learned by the current model at the round, respectively. In each round, FeAT augments the model with new features contained in the growing augmentation sets while retaining the already learned features contained in the retention sets, which will lead the model to learn richer features for OOD training and obtain a better OOD performance. Then FeAT augments the model with new features while retaining the already learned features, which leads the The performance of various OOD
Figure 1:
to richer features for OOD training and better OOD performance. (b) OOD Performance number of ERM pre-training epochs in COLOREDMNIST objectives largely relies on the quality of ERM-learned features. When there exist underlying useful features poorly learned by ERM, the OOD performance will be limited. In contrast, FeAT learns richer features with 2 rounds (or 300 epochs) and improves the OOD performance.
performance of OOD objectives largely relies on the pre-training with ERM before applying the OOD objectives [16, 83). As shown in Fig. 1(b), the number of ERM pre-trane pre-training epochs has large influence on the final OOD performance. These seemingly contradicting phenomena raise challenging research question:
What features are learned by ERM and OOD objectives, respectively. and how do the learned features generalize to in-distribution and out-of-distribution data?
To answer the question, we conduct a theoretical investigation of feature learning in two-layer CNN network, when trained with ERM and widely used OOD objective, IRMvl [4], respectively. We variation of the data models proposed in [2. 12], and include features with different correlation degrees to the labels to simulate invariant and spurious features [36].
First, we find that ERM essentially learns both spurious features and invariant features (Theorem 4.1). The degrees of spurious and invariant feature learning are mostly controlled by their correlation strengths with labels. Moreover, merely training with IRMv1 cannot learn new features (Theorem 4.2). Therefore, the quality ERM feature learning affects the final OOD performance significantly. Hence, as the number of ERM pre-training epochs increases, the model learns invariant features better and thus the final OOD performance will increase (Fig. 1). However, when ERM does not capture all useful features for OOD generalization, i.e., there exist some useful features that are poorly learned by ERM, the model can hardly Iearn these features during OOD training and the OOD performance will be limited. Given limited number of pre-training steps, it could often happen due to low invariant correlation strength, the feature learning biases of ERM [67], or the model architectures [28]. Consequently, ERM feature learning can be bottleneck to OOD generalization [60].
To remedy the issue, we propose Feature Augmented Training (FeAT), an iterative strategy to enforce the model to learn richer features. As shown in Fig. 1(a), in each round, FeAT separates the train set into two subsets according to whether the underlying features each set are already learned (Retention set D") or not (Augmentation set Do), by examining whether the model yields correct (D") or incorrect (D4) predictions for samples from the subsets, respectively. Intuitively, Da and Dr will contain distinct features that separated in different rounds. Then, FeAT performs distributionally

---

robust optimization (DRO) [50, 83] on all subsets, which augments the model to learn new features by minimizing the maximal ERM losses on all D and retains the already learned features by minimizing ERM losses on all 'D". Along with the growth of the augmentation and retention sets, FeAT is able to learn richer features for OOD training and obtain better OOD performance. FeAT terminates when the model cannot learn any new predictive features (Algorithm 1).
We conduct extensive experiments on both COLOREDMNIST [4, 16] and 6 datasets from the challenging benchmark, WILDS [39], and show that FeAT effectively learns richer features and thus consistently improves the OOD performance when applied to various OOD objectives (Sec. 6).
2 Related Work
We discuss the most related work to ours and leave more details in Appendix C.
On Feature Learning and Generalization. Understanding feature learning in deep networks is crucial to understanding their generalization [2. 11, 12, 22, 32, 62. 70J. Beyond the empirical probing [21, 26, 28, 65], Allen-Zhu and Li [2] proposed a new theoretical framework for analyzing the feature learning process of deep networks, which has been widely adopted study various deep learning phenomena [12, 32, 78, 86]. However, how the learned features from ID data can generalize to OOD data remains elusive. The only exceptions are [68] and [42]. Kumar et [42] find fine-tuning can distort the pre-trained features while fine-tuning can be considered as a special case in our framework. Shen et al. [68] focus on how data augmentation helps promote good but hard-to-learn features and improve OOD generalization. Deng et al. [20] finds neural networks tend to learn spurious features under imbalanced groups. contrast, we study the direct effects of ERM and OOD objectives to feature learning and provide theoretical explanation for the curious phenomenon [33, 63). To the best of our knowledge, we are the first to analyze the feature learning of ERM and OOD objectives and their interactions in the general OOD generalization setting.
Rich Feature Learning. Recently many OOD objectives have been proposed to regularize ERM such that the model can focus on learning invariant features [4. 41, 55, 57, 76]. However, the final OOD performance has large dependence on the number of ERM pre-training epochs [16, 83]. To remedy the issue, Zhang et al. [83] proposed Bonsai to construct rich feature representations as network initialization for OOD training. Although both Bonsai and FeAT perform DRO on grouped subsets, Bonsai rely on multiple initializations of the whole network to capture diverse features from the subsets, and complicated ensembling of the features, which requires more training epochs for convergence. In contrast, FeAT relieves the requirements via direct augmentation--eeentte grouped subsets, and thus obtains better performance. More crucially, although rich feature learning algorithms such as Bonsai and weight averaging [5, 59] have gained some successes, explanations about the reliance of OOD performance on ERM pre-training and why rich feature learning mitigates the issue remain clusive. In addition to new rich feature learning algorithm, work provides theoretical explanations for the success of rich feature learning in OOD generalization.
on the
3 Preliminaries and Problem Definition
Notations. We use old-faced letters for vectors and matrices otherwise for scalar; || ll2 to denote the Euclidean norm of vector or the spectral norm of a matrix, while || 1| for the Frobenius norm of matrix. I2 refers to the identity matrix in 1d Full details are deferred to Appendix A.
Our data model D (x,,8) yi is adapted from [2, 12] and further characterizes each data point x; as invariant and spurious feature patches from the two-bit bit model [16, 36].
Definition 3.1. D {De is composed of multiple subsets De from different environments e E Eall where each D. = {x}, )e is composed of i.i.d. samples (xf, ~ Fr. Each data (x", 3/²) E D. with x E 2d and y E (-1, 1) is generated as follows:
(a) Sample y" E f-1, 1) uniformly;
(b) Given y", each input x" = [x] xx] contains a feature patch x₁ and a noise patch X2, that are sampled as:
X1 = y Rad(a:) vı y- Rad( V2 x2 = where Rad(ā) is a random variable taking value -1 with probability 8 and +1 with probability - 8, vi T and v2 [0, 1, 0,... OOT.
3

---

(c) A noise vector & is generated from the Gaussian distribution N (L. vivi V2V22
Definition 3.1 is inspired by the structure of image data in image classification with CNN [2], where the inputs consist of different patches, some of the patches consist of features that are related to the class label of the image, and the others noises that are irrelevant to the label. In partieular, V1 and V2 are feature vectors that simulate the invariant and spurious features, respectively. Although our data model focuses two feature vectors, the discussion and results be further generalized to multiple invariant and spurious features with fine-grained characteristics [68]. Following previous works [12], assume that the noise patch is generated from the Gaussian distribution such that the noise vector is orthogonal to the signal vector v. Each environment is denoted as & = {(a, 3e) 0 < 3e < 1), where v1 is the invariant feature as is fixed while V2 is the spurious feature as S. varies across
CNN model. We consider training a two-layer convolutional neural network with a hidden layer width of m. The filters are applied to X1, X2, respectively. and the second layer parameters of the network are fixed as 1 and 1, respectively. Then the network can be written as f(W, x) F+1 (W x1, - (-1(W- (L1(W_1,,88 x1, where F41 (W+1++++44 +1, and F-1(W-WW)) -1. are defined as follows:
W J. x) = 772 Σ
[W(w], - * [[
where (x) is the activation function. We assume that all network weights are initialized as // o.).
ERM objective. We train the CNN model by minimizing the empirical cross-entropy loss function:
Re
L(W)= = Σ
DEE The Elyt f(W,x())
1m1
where ((s)=log(1+enn)) and (Do)oee, = (f I is the trainset with Eees, Yh
OOD objective. The goal of OOD generalization is, given the data from training environments (DeteeE,, to find predictor f : X' + y/ that generalizes well to all (unseen) environments, or minimizes max.eEa L.()), where Le is the empirical risk under environment e. The predictor = w & is usually composed of a featurizer
X + 2 that learns to extract useful features, and a classifier w + y that makes predictions from the extracted features.
Since we are interested in cases where the OOD objective succeeds in learning the invariant features. In the discussion below, without loss of generality, we study one of the most widely discussed OOD objective, IRMv1 objective, from IRM framework [4], and the data model where IRMv1 succeeds. Specifically, the IRM framework approaches OOD generalization by finding invariant representation yp such that there exists a classifier acting on yp that is simultancously optimal in &g. Hence, IRM leads to a challenging bi-level optimization problem as
min Σ
Le (), s.t. E arg min Le (w sp), Ve E E w:2-y
3EE
Due to the optimization difficulty of Eq. (3), Arjovsky et al. [4] relax Eq, (3) into IRMvl as follows:
min S
Le( +
CEE
Given the convolutional neural network (Bq. 1) and logistic loss (Eq, 2), IRMvl can be written as
LIRMVI (W)
Σ
e (uf f(W. xf)) +
t of f(W,xt)
(5)
eEEL
am
REEu
where 6:c = E'( f(W, x())
1+exp(- J(W x) Due the complexity of IRMv1, in the analysis below, we introduce CiRMAI for the ease of expressions. Specifically, we define CRRMA as
exp(- /(W-
1
CiRMVI
The
where of A /(W, xf) is the logit of sample x; from environment e. The convergence of CiRMvI indicates the convergence of IRMv1 penalty. The following lemma will be useful in our analysis.
When the environment is not explicitly considered, will omit it for clarity.

---

33002
tran
1005
00030
0.73
1900
6000
6003
3030
2003
200
Opeck
2002 4700
Spech
00.3
(a) CIRM-L, w/ WT
(b) FL., w/ PT
(c) CIIMAL w/o PT
(d) FL., w/o PT
Figure 2: The convergences of CIRMvL and feature learning coefficients (FL) with or with- out ERM pre-training (PT). The invariant and spurious feature learning terms are the mean of (w and (W (VV)) for j E {+1).r im], respectively. The training environments are Err = (00.55 0.1), (0.25, 0.2)). The black dashed line indicates the end of pre-training, More details are given in Appendix D.1.
Lemma 3.2. (Cao al. [12]) Let (t) for j [+1,-1} and E (1,2....,m) m] be the convolution filters of the CNN at t-th iteration of gradient descent. Then there exists unique coefficients ino (1), Tj. (t) 2 0 and Ptra(t) such that,
3 (t) v1 +J 20 (t) V2 + Σ Po. (t).(tt)-- -t) E.- W w,r (0) + + Yir
(6)
We refer Eq. (6) as the signal-noise decomposition of (t) [12]. We add normalization factor ||&:1122 in the definition so that Pj.- Rs (w Et) Note that [|v = ||ve|12 = 1, the corresponding spu
normalization factors are thus neglected. Furthermore, inz s (Wj.r. v1) and 7jrr s (Wirr V2) respectively denote the degrees of invariant and spurious feature learning.
4
Theoretical Understanding of Feature Learning in OOD Generalization
4.1 ERM Feature Learning
With the setup in Sec. 3, we first study the feature learning of the ERM objective. We consider a two training environments setup &,r = {(a 31), (a, 82)} where the signal of invariant feature is weaker than the average of spurious signals (i.e., > f. ), which corresponds to Figure 2. For precise characterization of the training dynamic, we adopted a minimal setup where v'(x)) = x in Figure 2(a) and the following theorem, which already captures the key phenomenon in ERM feature learning. We study ERM feature learning with non-linear activations in Appendix D.2.3.
Theorem 4.1. (Informal) For p O, let n & minege, nee Suppose that we run T iterations of GD for the ERM objective. With sufficiently large n and v(xx) = x, assuming that (i) a, $1, < 2/2 and (ii) a > 81, with properly chosen oe and oi. there exists a constant such that with probability least 2p, both invariant and spurious features are converging and the increment of the spurious feature is larger than that of the invariant feature at any iteration t E {O...., T - 1} (the detailed
quantitative result of this gap can be found at (15) in Appendix D.2).
As the formal statement of Theorem 4.1 is too complicated and lengthy, we leave it and its proof in Appendix D.2, while giving an informal but more intuitive version here. Theorem 4. states that FRM training learns both invariant feature and spurious feature at the same time, and if the average of spurious signals stronger, the coefficient of spurious feature learning will dominate that of invariant feature learning in the whole training process, corresponding to Figure 2(b). We establish the proof based on inspecting a novel recursive equation, which might be of independent interest. Note that Theorem 4.1 can be directly generalized to handle any number of environments.
Speaking of implications, Theorem 4.1 provides answers to the seemingly contradicting phenomena that ERM fails in OOD generalization [7, 19] but still learns the invariant features [33, 38, 63]. In fact, ERM fails since it learns the spurious features more quickly, when spurious correlations are stronger than invariant correlations. Nevertheless, invariant feature learning also happens, even when the spurious correlations are strong, so long as the invariant feature has a non-trivial correlation strength with the labels. Therefore, simply re-training classifier based on subset of unbiased data on top of the ERM-trained featurizer achieves impressive OOD generalization performance [33, 38, 63].
We use wj.r (t) and w interchangeably.
5

---

Theorem 4.1 also provides an explanation for the ID-OOD performance correlations when fine-tuning or training neural networks (especially large pre-trained models like CLIP [56], GPT [10]) [46, 71, 79, 80]. We provide detailed discussion in Appendix C.
4.2 IRM Feature Learning
Although Theorem 4.1 states that ERM learns both invariant and spurious features, the following questions remain unanswered: (1) whether IRMv1 learns new features or simply amplifies the already learned ERM features, and (2) how the quality of the ERM-leamed features affects the feature learning when IRMv1 is incorporated. We first study IRMv1 training from scratch (w/o pre-training).
Theorem 4.2. Consider training a CNN model (I) with data model (3.1), define c(t) & ChRMU (,t), PRRM (laniWWWr WWt).... Clen (W,t). and Xo Amin (H.), where H & Tiy mm EEs ((W,,, (0), (i,)),i,)))))) (1))) F1=1 ((w,,r (O) Xx, Suppose that dimension d ((log(m/65). network width variance op 2d weight initial scale do Ol 24/22m² 1ec_77/2g0 log(1/e) ). then with probability at least
1
((1/6), regularization factor X 2 ( noise
1 1-6, 6, after training time T = 2 (1)) we have ||c(T)]l2 S 2f (T) =o(), 75r (T)=
The proof is given in Appendix D.3. We highlight that Theorem 4.2 allows any number of training environments, which indicates a fundamental limitation of pure IRMv1 training. Intuitively. Theo- rem 4.2 implies that, when a heavy regularization of IRMv1 is applied, the model will not learn any features, corresponding to Figure 2(d). Instead, IRMv1 suppresses any feature learning, even the beginning of the training. Then, what would happen when given a properly pre-trained network?
After ERM pre-training, according to Theorem 4, we have [W,,r v1)1 = 2(1), 1(W v2)l ((1), and (W E) O(oopyV)). Then, we have the following hold.
Proposition 4.3. Given the same setting as Theorem 4.2, suppose that vb(x) = x, 7jr ine (11) = ine (41 -1), and (41) = Tj.r 7,,5((1 ((1-1) at the end of ERM pre-train L1, ō > 0. and > Cllo([168), with C being positive constant, then with a high probability at least 1 8, we have Ee C.iMim (f1) = 0, Pjor
inv (t1 + 1) > 75pr and spu
72r (t1 + 1) < epu (t1).
The proof is given in Appendix D.4, which takes converged feature learning terms from Theorem 4.1 as the inputs. Proposition 4.3 demonstrates that with sufficient ERM pre-training, IRMv1 can enhance the learning of invariant features while suppressing the learning of spurious features, which is verified in Figure 2(b) and 2(a). Thus, when given the initialization with better learned invariant features, i.e. longer ERM pre-training epochs, IRMv1 improves the invariant feature better. Proposition 4.3 explains why the final OOD performance highly depends on the ERM pre- pre-training [16, 83].
4.3 Limitations of ERM Feature Learning
Combining results from both Sec. 4. and Sec, 4.2, we know that the invariant features will be learned during ERM pre-traine vre-training and discovered during OOD training. However, given poorly learned invariant features, can IRMv1 still improve it? In practice, there often exist some invariant features that are not properly learned by ERM. For example. in our data model Def. 3.1 when the invariant correlation is much weaker than the spurious correlation, given a limited number of training steps, the spurious feature learning dominate the invariant feature learning. Besides, when considering other factors such as the simplicity bias of ERM [67] or the inductive biases of the network architecture (28), it is more likely that there exist invariant features that are not properly learned [60]. Then we have:
Corollary 4.4. Consider training the CNN with the data generated from Def. 3.1, suppose that (xx) = xr y (fz) = o(1). and spu
fjr (f1) = ((1) at the end of ERM pre-training t1- Suppose that 8 > 0, and n > C' Clog(1 18), with C being a positive constant, then with a high probability at least Sn 1 5, we have 15 (41 + 1) < 3/. (l1).
Corollary 4.4 shows that IRMv1 requires sufficiently well-learned features for OOD generalization. is also consistent with the experimental results in Fig. 2(b), 2(c), and Fig. 1, where all the OOD objectives only achieve a performance comparable to random guesses.

---

5 Feature Augmentated Training
5.1 Rich Features for OOD Generalization
The results in Sec. 4 imply the necessity of learning all potentially useful features during the pre- training stage for OOD generalization. Otherwise, the OOD training is less likely enhance the poorly learned features. It also explains the success of learning diverse and rich features by weight averaging [5, 59] and rich feature construction (or Bonsai) [83], and other approaches [58, 81].
Despite the empirical success, however, the learning of rich features in both Bonsai and weight averaging is unstable and expensive. On the one hand, they may discard previously learned useful features or fail to explore all the desired features as it is hard to evaluate the quality of the intermediate learned features. On the other hand, they also need multiple initializations and training of the whole networks with different random seeds to encourage the diversity of feature learning, which brings more instability and computational overhead, especially when applied to large and deep networks.
5.2 The FeAT Algorithm
To overcome the limitations of previous rich feature learning algorithms, propose Feature Augmented Training (FeAT), that directly augment the feature learning in an iterative manner.
Intuitivelyy the potentially useful fea- tures presented in the training data are features that have non-trivial cor- relations with labels, or using the re- spective feature to predict the labels is able to achieve non-trivial training
Algorithm 1 FeAT: Feature Augmented Training
Input: Training data Du:; the maximum augmentation
rounds K; predictor f == w o pp; length of inner training
epochs 6; termination threshold p;
2: Initialize groups Go <- Dtr, G" + 0:
performance. Moreover, the invari- ance principle assumes that the train- ing data comes from different envi- ronments [4], which implies that each set of features can only dominate the correlations with labels in a subset of data. Therefore, it is possible to dif- ferentiate the distinct sets of useful features entangled in the training data into different subsets, where ERM can effectively learn the dominant features
3: for A [1..... K] do
4: Randomly initialize wk:
5: for j [I.....t] do
6: Obtain FEAT with G via Eq. 7; Update wa sp with FCAT:: 8: end for
9:
// Early Stop if fa og fails to find
new features.
10:
Training accuracy Ju smaller than p then Set K = k 1 and terminate the loop:
11:
12:
end if
presented in the corresponding subset as shown in Theorem 4.1.
13:
Split Du into groups Di. DD' according to whether fx classifies the examples in Di, correctly or not;
14: Update groups G + G" (DR)), G" + G (PD):
The intuition naturally motivates an it- erative rich feature learning algorithm,
15: end for
16: Synthesize the final classifier < EK EE=10 10
i.e., FeAT, that identifies the subsets containing distinct features and ex- plores to learn new features in multiple rounds. The details of FeAT are given in Algorithm 1, where we are given a randomly initialized or pre-trained model f w ip that consists of a featurizer and a classifier w. In round k. FeAT first identifies the subset that contains the already learned features by collecting the samples where yields the correct prediction, denoted as Gy and the subset of samples that contains the features that have not been learned, denoted as GV.
17: return f w o pp;
At the k-th round, given the grouped subsets G (G", Ga) with 2k: groups, where Go (Dr) 1 is the grouped sets for new feature augmentation, and G" [DD the grouped sets for already learned feature retention (notice that D is the empty set), where D. and DJ are the corresponding augmentation and retention set elicited at i-th round. FeAT performs distributionally robust optimization (DRO) [50, 83] on G to explore new features that have not been learned in previous rounds. Meanwhile, FeAT also needs retain the already learned features by minimizing the empirical risk at G", for which we store and use the historical classifiers wi with the current featurizer to evaluate the feature retention degree. Then, the FeAT objective at round k is
fTeAT = max lop (we o (p) + A EG
Σ
EDy (w: o sp),
(7)
DrEG
7

---

Table 1: OOD performance on COLOREDMNIST datasets initialized with different representations.
COLOREDMNIST-OR
COLORRDMMNSTT-D
ERM-NF
ERM
BONSAI
REAT
ERM-NF
ERM
BONSAI
FEAT
ERM
17.14
12:40 (=0.22)
11.21 (1:000)
17. 27 (+2.5)
73.06 (#0771)
73.75 (=0,99
70.95 (:052)
76.05 (#1445)
IRMVI
(1:0.9)
(2:446)
(1:0.72
70.57 (1:058
76.89 (2:225)
73.84 (1:.56)
76.71 (:24.00)
(22)
V-REX
G8.62 (2:0.3)
65.96 (=1129)
70.31 (1:9:6)
70.82 (3:0.9)
83.52 (
81:20 (4:)27)
82.61 (#:1.76)
84.70 (::009)
IRXX
67.00 (1:1-99
64.05 (1:0:6)
70.46 (1:0-2)
70.78 (1:051)
81.61
75.97 (=0-66)
80.2 00.28// (::147)
84.34 (1:07)
18-RMM
56.09 (#:2.00)
39.81 (::4.0)
70.28 (2:0.72)
70.57 (1:06))
75.81 (#:)6:)
73. (2:0.5)
76.7 71 (:24.00 20)
82. (271.77)
CLOVE
58.67 (#:7.99
65.78 (=0.00)
65.57 (4:2.02)
05.78 (4:2.8)
75.66 (#:10.6)
74.73 (4:0.26)
72. T3 (#:1.25)
75.12 (#:000)
(+:7.7)
62.43 (4:200)
(4:0.0)
(4:7.0)
74.20 (=2-45)
73. (=0,48)
(#:3.00)
(#:2.77
FISER
69.38 (#:0299
67.74 (=4.90)
75 (1:1.10
70.58 (=:0.77
77.29 (#:15))
82.29 (=1.25
84.19 (:0.66)
84.26 (2:.0)
ORACLE
71.97
(#:9.27
where &D (w o (p refers to the empirical risk of o yp evaluated at the subset D;, and {w; S <is
k-1) } are the historical classifiers trained in round i. The final classifier is the average of all historical classifiers as they already capitalize all the learned features each round.
Relations with previous rich feature learning algorithms. Compared with previous rich feature learning algorithms, FeAT directly trades off the exploration of new features and the retention of the already learned features. Although Bonsai also adopts DRO to explore new features, the isolation of new feature exploration and already learned feature synthesis makes the feature learning in Bonsai more unstable. In other words, Bonsai can not evaluate the intermediate feature learning results due to eaeeeeeeeplliiii
the indirect feature exploration and synthesis. Consequently, Bonsai can not control when to stop the new feature exploration, and thus may fail to explore all of the desired features or discard important features. Besides, the multiple re-initializations and re-training of the whole network Bonsai could also lead to suboptimal performance and meanwhile require more computational overhead.
Practical implementations. Algorithm 1 requires store 2 2K - subsets and larger memory cost in training the network, which may cause additional storage burden when contains massive amount of parameters [39]. Hence, we propose lightweight variant of FeAT (denoted as iFeAT) which only retains the latest subsets and historical classifiers (D) Dj 1, WR-1 at the k-th round). Throughout the whole experiment, we will use iFeAT and find that iFeAT already achieves state-of-the-arr, More details are given in Appendix E.
As iFeAT stores only the latest augmentation and retention subsets, inspecting the training perfor- mance for termination check (line 10 of Algorithm 1) may not be suitable. However, one can still inspect the performance in either an OOD validation set check the quality of the intermediate feature representations, or the retention set to check whether learning new features leads to a severe contradiction of the already learned features (FeAT should terminate so).
Compared to ERM, the additional computational and memory overhcad introduced in FeAT mainly lie in the FeAT training and partitioning. At each training step. FeAT needs (k- (k-1) additional forward and backward propagation, the same as Bonsai, while FeAT only needs min(1, & 1) additional propagation. Besides, Bonsai additionally require another round of training with (K 1) additional propagation, given K total rounds. More details are given in Appendix F.4.
6 Empirical Study
We conduct extensive experiments on COLOREDMNIST [16] and WILDS [39] to verify the effec- tiveness of FeAT in learning richer features than ERM and the state-of-the-an art algorithm Bonsai [83].
Proof-of-concept study on COLOREDMNIST. We first conduct a proof-of-concep: study us- ing COI.OREDMNIST [16] and examine the feature learning performance of FeAT under various conditions. We consider both the original COLOREDMNIST with & = {00.25, 0.1), (0.25, 0.2)} where spurious features are better correlated with labels, and the
(denoted as COLOREDMNIST-O223 modified COLOREDMNIST (denoted as COLOREDMNIST-ON where invariant features are better correlated with labels. We compare the OOD performance of the features learned by FEAT, with that of ERM and the state-of-the-an art rich feature learning algorithm Bonsai [83]. Based on the features learned by ERM, Bonsai, and FEAT, we adopt various of the-art OOD objectives including IRMv1 1 [4], VREx [41], IRMX [16]. IB-IRM [1], CLOVE [76]. IGA [40] and Fishr [57] for OOD training. order to evaluate the practical quality of the learned features. The feature representations are frozen once initialized for the OOD training as fine-tuning
with & = (0.11, 0.2), (0.1, 0.25)
the featurizer can distort the pre-trained features [43]. We also compare FeAT with the common
training approach that uses unfrozen ERM features, denoted ERM-NF. For Bonsai, we trained 2
8

---

Table 2: OOD generalization performances on WILDS benchmark.
CAMELYON17
COMM.COMMNNTS
FMoW
(WILDCAM
AMAZON
RxRx1
INIT.
METHOD
Avg. (8)
Worst acc. (%)
Worst acc. (8)
Macro F1
10:th per acc. (R)
Avg. acc. (%)
ERM
DER!
(+1.6)
77.34 (=0050)
41.96 (=2.90)
23.15 (#0.4)
48 (#000)
ERM
DFR s!
82.24 (#0.3)
56. (=0.2)
52.44 (+0.44)
Bonsai
DFR
(1.17
17.07 (2088)
43.26 (2))2)
21.30 (#940)
46.67 (=200)
Bonsai
DFR-s
81.26 (1.7
58.58 (1z 12)
50.85 (1:0000
FeAT
DFR
95.28 (10019)
77.34 (+00.0
43. 54.54/10 (++1.20)
23. (+:0.52
49. 49.33 (#:000)
FeAT
DFR.s
79.56 (100 28)
57.09 (20.78)
52.31 (#> 28)
ERM
FRM
74.30 (1:5 96)
55.53 ( 28)
33. 58 (12)
28. 22 (#:0.00)
51.11 (=:)0))
30.21 (1:000)
ERM
GreuDDRO
76.09 (=:5.5)
69.50 (50015)
(=0052)
(#00 58)
(=0.00)
20.99 (+)..1)
FRM
IRMvl
75.68 (1:7.1)
68.84 (12 95)
33.45 (1t: m77)
28. 76 (#:945)
52.00 02.00 (=:0000)
30.10 (1:005)
ERM
V-REx
71.00 (17.889
69.08 (.07.0
33, 0G (1.0.06)
(1.0.47)
52. (1/067)
(.10.35)
HRM
IRMX
73.19 (49937)
68.91 (.19))
33. 13 (#00.6)
28. (=:0,77)
522.0 02)00) (2:0.00)
30.10 (1:002)
Bonsai
HRM
73.98 (4:.30)
63.34 (::)
31.91 (=0.51)
26. 27 (#00
48. (=0.56)
24.22 (+:0.4)
Bonsai
GreuDRO
72.82 (15)
70.23 (t: 24)
(=1
27. 67.16 (#: (8)
(+1 0o)
(=10.46
Bonsai
IRMvl
73.59 (16.16)
68.39 (1.2.0))
32. 51 (22)
27. 60 (1577)))
47 11 (1:00)
23. 35 (110.4))
Bonsai
V-REx
76.39 (4:5.32)
68.67 (#:7.29)
33.17 (#:7.20)
25.81 (#:0.42)
18. 00 (-:0,00)
23.34 (210.42)
Bonsai
IRMX
64.77 (110.1)
69.56 (1.0.95)
32.63 (120255)
27. 62 (1.0.6)
46.67 (1:000)
23. 23.34 (110.49)
FeAT
ERM
77.80 (+:2499)
68.11 (+22
33. (=0.78)
28.47 (#0667)
$2. 92.89/M (+:03)
30.66 (+:0-2)
FeAT
GroupDRO
80.41 (2:4.00)
71.29 (::0.46)
33.55 (7)
28.38 (1:1.9)
52.58 (=2.55) 56)
20.99 (:b9.1)
FcAT
IRMvI
77.97 (1.7.9)
70.33 (1.4.14
34.04 (110.70)
29.66 (1.1.52)
52. 89 (1.66.)
29. 99.99/ (110.19)
FeAT
V-REx
75.12 (3)
70.97 (1:7.00)
34.00 (=0071)
29.48 (1:1 14)
(::066))
30.57 (:b9.50)
FcAT
IRMX
76.91 (1:676)
71.18 (1.7.10)
33, (1007))
29.01 (1.2%)
52. 89 (1:06.))
29. 20.92/ (100.6)
DFR/DFR-s use an additional OOD datasct to evaluate invariant and spurious feature learning. respoctively.
rounds following Zhang et al. [83], while for FeAT the automatic termination stopped at round 2 in COLOREDMNIST-022 and round 3 in COLOREDMNIST-OL For ERM, we pre-trained the model with the same number of overall epochs as FEAT in COLOREDMMISTARE at the number of epochs of 1 round in COLOREDMNIST -025 to prevent over-fitting. All methods adopted the same backbone and the same training protocol following previous works [16, 83). More details are given in Appendix F.I.
while early stopping
The results are reported in Table 1. can be found that ERM will learn insufficiently good features under both stronger spurious correlations and invariant correlations, confirming our discussion and boosts OOD
in Sec. 4.3. Besides, Bonsai learns richer features in COLOREDMNIST performance, but Bonsai sometimes leads to suboptimal performances in COLOREDMNIST-OL. which could be caused by the unstable feature learning in Bonsai. In contrast, FeAT consistently improves the OOD performance of all OOD objectives for all the COLOREDMNIST datasets, demonstrating the advances of direct feature learning control in FeAT than Bonsai and ERM.
Experiments on real-world benchmarks. We also compare FeAT with ERM and Bonsai in 6 real-world OOD generalization datasets curated by Koh et al. [39] that contain complicated features and distribution shifts. The learned features are evaluated with several representative state-of-the-art OOD objectives in WILDS, including GroupDro [64], IRMv1 [4], VREx [41] as well as IRMX [16]. By default, we train ERM, Bonsai and FeAT the same number of steps, and kept the rounds of Bonsai and FeAT the same (though Bonsai still requires one more round for feature synthesis). The only exception is in RxRxI where both Bonsai and FeAT required more steps than ERM to converge. We use the same evaluation protocol following the practice in the literature [16, 39, 69, 83] to ensure a fair comparison. More details are given in Appendix F.2.
In addition to OOD objectives, we evaluate the learned features with Deep Feature Reweighting (DFR) [38]. DFR uses an additional OOD validation set where the spurious correlation does not hold, to perform logistic regression based on the learned features. Intuitively, DFR can serve as proper measure for the quality of learned invariant features [33]. When the original dataset does not provide a proper OOD validation set, c.g.. CAMELYON 17, we use an alternative implementation based on random split of the training and test data to perform the invariant feature quality measure [63]. Similarly, we also report DFR-s by regression with the environment labels (when available) to evaluate the spurious feature learning of different methods. More details are given in Appendix F.2.2.
The results are presented in Table 2. Similarly, when the tasks grow more challenging and neu- ral architectures become more complicated, the ERM learned features can have lower qual- ity as discussed Sec. 4.3. For example, ERM can not sufficiently learn all useful features in FMoW, while ERM can learn more spurious correlations in CivilComments. Moreover, it can also be observed the instability of Bonsai in learning richer features that Bonsai even under- performs ERM in rich feature learning and OOD generalization in multiple datasets. In con-

---

4
G
4
PAIT
9
3
8
3
3
3
X
4
8
3
(a) ERM 300 epochs
(b) FeAT rounds
Figure 3: GradCAM visualization on COLOREDMNISST022, where the shortcuts are now concen- trated to colored path at the up left. Three visualizations are drawn for each sample: the original figure, the gray-colored gradcam, and the gradcam. It can be found that ERM can not properly capture the desired features while FeAT can stably capture the desired features.
trast, FeAT consistently achieves the best invariant feature learning performance across various challenging realistic datasets. Meanwhile, compared to ERM and Bonsai, FeAT also reduces result, FeAT
over-fitting to the spurious feature learning led by spurious correlations. achieves consistent improvements when the learned features are applied to various OOD objectives.
The termination check in FeAT. As elaborated in Sec. 5.2, key differ- ence between FeAT and previous rich feature learning algorithms such as Bonsai is that FeAT is able to access the intermediate feature representa- tions and thus can perform the automatic termination check and learn the desired features stably. To verify, we list the FeAT performances various subsets of COLORRDDMISTT-OR in Table 3. By inspecting the retention accuracy, after FeAT learns sufficiently good features at Round 2, is not necessary to proceed with Round 3 as it will destroy already learned features and lead to degenerated retention and OOD performance. More details and results are given in Appendix F.1.
Table 3: Performances at different FeAT rounds.
COLOREDMNIST-023
ROUND-1
ROUND-2
ROUND-3
TRAINING Acc.
85.08 0.14
71.874 0.96
84,93= 1.26
RETENTION ACC.
88.114: 4.28
43.82:1: 0.59
OOD ACC.
11.08.1 0.30
70.641 0.62
10.07:L 0.26
different rounds
Computational analysis. We also analyze the computational and memory overhead of different methods, for which the details are given in Appendix F.4. Compared to ERM and Bonsai, iFeAT achieves the best performance without introducing too much additional overhead.
Feature learning analysis. We visualize the feature learning of ERM and FeAT on ColoredMNIST- 025. As shown in Fig. 3, ERM can learn both invariant and spurious features predict the label, aligned with our theory. However, ERM focuses more on spurious features and even forgets certain features with longer training epochs, which could be due multiple reasons such the simplicity biases of ERM. Hence predictions based on ERM learned features fail to generalize to OOD examples. In contrast, FeAT effectively captures the meaningful features for all samples and generalizes 001) examples well. More analysis including results on WILDS benchmark can be found in Appendix F.5.
Conclusions
In this paper, we conducted theoretical investigation of the invariant and spurious feature learning of ERM and OOD objectives. We found that ERM learns both invariant and spurious features when OOD objectives rarely learn new features. Thus, the features learned the ERM pre-training can greatly influence the final OOD performance. Having learned the limitations of ERM pre-training, we proposed FeAT to learn all potentially useful features. Our extensive experimental results verify that FeAT significantly boosts the OOD performance when used for OOD training.
10

---

Acknowledgements
We thank the reviewers for their valuable comments. This work was supported by the RIKEN Incentive
Research Project 100847-202301062201 by CUHK direct grant 4055146 BH was supported by the NSFC Young Scientists Fund No. 62006202, NSFC General Program No. 62376235, Guangdong Basic and Applied Basic Research Foundation No. 2022A151501 1652, HKBU Faculty Niche Research Areas No. I.--NNRA-IG/2223/SCCCII and Tencent Al Lab Rhino-Binn Gift Fund.
References
[1] K. Ahuja, E. Caballero, D. Zhang, J.-C. Gagnon-Audet, Y. Bengio, Mitliagkas, and I. Rish. Invariance principle meets information bottleneck for out-of-distribution generalization. In Advances in Neural Information Processing Systems, 2021.
[2] Z. Allen-Zhu and Y. Li. Towards understanding ensemble, knowledge distillation and self- distillation in deep learning, arXiv preprint arXiv:2012.09810 2020.
[3] A. Andreassen, Y. Bahri, B. Neyshabur, and R. Roelofs. The evolution of out-of-distribution robustness throughout fine-tuning, arXiv preprint vrXiv:2106.18831 2021.
[4] M. Arjovsky, L. Botton, I. Gulrajani, and D. Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv:190..02893 2019.
[5] D. Arpit, H. Wang, Y. Zhou, and C. Xiong. Ensemble of averages: Improving model selection and boosting performance in domain generalization. In Advances in Neural Information Processing Systems, 2022.
[6] P. Búndi, O. Geessink, Q. Manson, M. V. Dijk, M. Balkenhol, M. Hermsen, B. E. Bejnordi, Lee, K. Paeng, A. Zhong. Q. Li. G. Zanjani, Zinger, K. Fukuta, D. Komura, V. Ovtcharov, S. Cheng, S. Zeng. J. Thagaard, A. B. Dahl, H. Lin, H. Chen, L. Jacobsson, M. Hedlund, M. Çetin, Halici, H. Jackson, R. Chen, Both, Franke, H. Küisters-Vandvellee W. Vreuls, P. Bult, B. van Ginneken, J. van der Laak, and G. Litjens. From detection of individual metastases to classification of lymph node status at the patient level: The CAMELYON17 challenge. IEEE Trans. Medical Imaging. 38(2):550 560, 2019.
[7] S. Beery, G. Horn, and P. Perona. Recognition in terra incognita. Computer Vision European Conference, Part XVI, pages 472 489, 2018
[8] S. Beery, E. Cole, and A. Gjoka. The iwildcam 2020 competition dataset. arXiv preprint arXiv:2004.1 10340, 2020.
[9] D. Borkan, Dixon, J. Sorensen, Thain, and L. Vasserman. Nuanced metries for measuring unintended bias with real data for text classification. In Companion of The 2019 World Wide Web Conference, pages 491-500, 2019.
[10] T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, Shyam, G. Sastry, A. Askell, Agarwal, A. Herbert- Voss, G. Krueger, Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, Clark, C. Berner, S. McCandlish, A. Radford, Sutskever, and D. Amodei. Lan- guage models are few-shot learners. In Advances in Neural Information Processing Systems, 2020.
[11] A. Brutzkus, A. Globerson, E. Malach, and S. Shalev-Shwartz. SGD learns perrparraeteriized networks that provably generalize on linearly separable data. In International Conference on Learning Representations, 2018.
[12] Y. Cao, 7. Chen, M. Belkin, and Q. Gu. Benign overfitting in two-layer convolutional neural networks. In Advances in Neural Information Processing Systems, 2022.
[13] H. S. Chen, Lee, A. Setlur, Levine, and C. Finn. Project and probe: Sample-efficient domain adaptation by interpolating orthogonal features. arXiv preprint 2XXv:2302.004411 2023,

---

[14] Y. Chen, H. Yang, Y. Zhang, K. Ma, T. Liu, B. Han, and J. Cheng. Understanding and improving graph injection attack by promoting unnoticeability, In International Conference on Learning Representations, 2022.
[15] Y. Chen, Y. Zhang, Y. Bian, H. Yang, K. Ma, B. Xie. T. Liu, B. Han, and J. Cheng. Learning causally invariant representations for out-of-distribution generalization on graphs. In Advances in Neural Information Processing Systems, 2022.
[16] Y. Chen, K. Zhou, Y. Bian, B. Xie, K. Ma, Y. Zhang, H. Yang, B. Han, and J. Cheng. Pareto invariant risk minimization. arXiv preprint arXiv:2206.074 2022.
[17] Y. Chen, Y. Bian, K. Zhou, B. Xie, B. Han, and J. Cheng. Does invariant graph learning via environment augmentation learn invariance? In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[18) G. A. Christie, N. Fendley, J. Wilson, and R. Mukherjee. Functional map of the world. IEEE Conference Computer Vision and Pattern Recognition, pages 6172-6180, 2018.
[19] A. J. DeGrave, J. D. Janizek, and S. Lee. Al for radiographic COVID-19 detection selects shortcuts over signal. Nature Machince Intelligence, 3(7):610-619 2021.
[20] Y. Deng, Y. Yang, B. Mirzasoleiman, and Q Gu. Robust learning with progressive data expansion against spurious correlation. arXiv preprint, arXiv:2306.04949, 2023.
[21] N. Elhage, T. Hume, C. Olsson, N. Schiefer, T. Henighan, S. Kravec, Z. Hatfield-Dodds. R. Lasenby, Drain, C. Chen, R. Grosse, S. McCandlish, Kaplan, Amodei, M. Wattenberg, and C. Olah. Toy models of superposition. arXiv preprint arXiv:2209. 10652, 2022.
[22] S. Frei, Y. Cao, and Q. Gu. Provable generalization of sgd-trained neural networks of any width the presence of adversarial label noise. In International Conference on Muchine Learning, pages 3427-3438, 2021.
[23] R. Geirhos, J. Jacobsen, C. Michaelis, R. S. Zemel, W. Brendel, M. Bethge, and F. A. Wichmann. Shortcut learning deep neural networks. Nature Machine Intelligence, 2(1 2(16):6556633 2020.
[24] Goodfellow, Y. Bengio, and A. Courville. Deep Learning. / //www. www.deepplerningbookk org.
MIT Press, 2016. http:
[25] I. Gulrajani and D. Lopez-Paz. In search of lost domain generalization. In International Conference on Leurning Representations, 2021.
[26] A. Gupta, N. Saunshi, D. Yu, K. Lyu, and Arora. New definitions and evaluations for saliency methods: Staying intrinsic and sound, 2022.
[27] K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In IEEE Conference on Computer Vision and Pattern Recognition, pages 770-778, 2016.
[28] K. Hermann and A. Lampinen. What shapes feature representations? exploring datasets, architectures, and training. In Advances in Neural Information Processing Systems, pages 9995-10006, 2020.
[29] W. Hu. Understanding Deep Learning via Analyzing Dynamics of Gradient Descent. Princeton University, 2021.
[30] W. Hu, L. Xiao, B. Adlam, and J. Pennington. The surprising simplicity of the early-time learning dynamics of neural networks. Advances in Neural Information Processing Systems, 33: 17116-17128, 2020.
[31] G. Huang, 7. Liu, I. van der Maaten, and K. Q. Weinberger. Densely connected convolutional networks. In IEEE Conference on Computer Vision and Pattern Recognition, pages 2261-2269, 2017.
[32] W. Huang, Y. Cao, H. Wang, X. Cao, and T. Suzuki. Graph neural networks provably benefit from structural information: A feature learning perspective. arXiv preprint arXiv:2306. 13926, 2023.
12

---

[33] P. Izmailov, P. Kirichenko, N. Gruver, and A. G. Wilson. On feature learning in the presence of spurious correlations. Advances in Neural Information Processing Systems, 2022.
[34] Y. Ji, L. Zhang, Wu, Wu, L.-K. Huang, T. Xu, Y. Rong. L. Li, J. Ren, D. Xue, H. Lai, S. Xu, J. Feng, W. Liu, P. Luo, S. Zhou, J. Huang, P. Zhao, and Y. Bian. Drugood: Out-of-distribttion (ood) dataset curator and benchmark for ai-aided drug discovery focus on affinity prediction problems with noise annotations. arXiv preprint arXiv:2201. 00061.9977, 2022.
[35] D. Kalimeris, G. Kaplun, P. Nakkiran, B. Edelman, T. Yang, B. Barak, and H. Zhang. SGD neural networks learns functions of increasing complexity. Advances in neural information processing systems, 32, 2019.
[36] Kamath, A. Tangella, D. Sutherland, and N. Srebro. Does invariant risk minimization capture invariance? In International Conference on Artificial Intelligence and Statistics, pages 4069-4077, 2021.
[37] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In International Conference Leurning Representations, 2015.
[38] P. Kirichenko, P. Izmailov, and A. G. Wilson. Last layer re-training is sufficient for robustness to spurious correlations. arXiv preprint arXiv:2204 02930,997777 2022.
[39] P. W. Koh, S. Sagawa, H. Marklund, S. M. Xie, M. Zhang, A. Balsubramani, W. Hu, M. Ya- sunaga, R. L. Phillips, I. Gao, T. Lee, E. David, Stavness, W. Guo, B. Earnshaw, Haque, S. M. Beery, J. Leskovec, A. Kundaje, E. Pierson, S. Levine, C. Finn, and P. Liang. WILDS: A henchmark of in-the-wild distribution shifts. International Conference on Machine Learning. pages 5637-5664, 2021.
[40] M. Koyama and S. Yamaguchi. predictor. arXiv preprint FTXX::2000.01883, 2020.
generalization with maximal invariant
[41] D. Krueger, E. Caballero, J. Jacobsen, A. Zhang, Binas, D. Zhang, L. Priol, and A. C. Courville. dut-of-distrribution generalization via risk extrapolation (rex). In International Conference on Machine Learning, pages 5815-5826, 2021.
[42] A. Kumar, A. Raghunathan, R. M. Jones, T. Ma, and P. Liang. Fine-tuning can distort pre- trained features and underperform out-of-distribution. In International Conference on Learning Representations, 2022.
[43] A. Kumar, A. Raghunathan, R. M. Jones, T. Ma, and P. Liang. Fine-tuning can distort pre- trained features and underperform out-of-distribution. In International Conference on Learning Representations, 2022.
[44] Y. Lee, H. Yao, and C. Finn. Diversify and disambiguate: Out-of-distrriuttion robustness via disagreement. In International Conference on Learning Representattioss 2023.
[45] Y. Lin, L. Tan, Y. Hao, H. Wong, H. Dong, Zhang, Y. Yang, and T. Zhang. Spurious feature diversification improves out-of -uttof-distribution generalization. arXiv preprint arXiv:2309.12330. 2023.
[46] Y. Lin, L. Tan, H. Lin, Z. Zheng, R. Pi, Zhang,
S. Diao, H. Wang. H. Zhao, Y. Yao, and T. Zhang, Speciality vs generality: An empirical study on catastrophic forgetting in fine-tuming foundation models. arXiv preprint arXiv:2309.0625:
2023.
[47] J. P. Miller, R. Taori, A. Raghunathan, S. Sagawa, P. W. Koh, V. Shankar, P. Liang, Y. Carmon, and L. Schmidt, Accuracy on the line: on the strong correlation between -ut-of-distribution and in-distribution generalization. In International Conference on Machine Learning, pages 7721-7735, 2021.
[48] R. Min, Z. Qin, L. Shen, and M. Cheng. Towards stable backdoor purification through feature shift tuning. arXiv preprint arXiv:2310.01875 2023.
[49] H. Naganuma, K. Ahuja, I. Mitliagkas, S. Takagi. T. Motokawa, R. Yokota, K. Ishikawa, and Sato. Empirical study on optimizer selection for dut-o--distribution generalization. arXiv preprint arXiv:2211.088 2022.
13

---

[50] H. Namkoong and J. C. Duchi. Stochastic gradient methods for distributionally robust opti- mization with f-divergences. In Advances in Neural Information Processing Systems, pages 2216,, 2016.
[51] Ni, J. Li, and J. McAuley. Justifying recommendations using distantly-labeled reviews and fine- fine-grained aspects. In Conference on Empirical Methods in Natural Language Processing and International Joint Conference on Natural Language Processing, pages 188-197, 2019.
[52] A. M. Nicolicioiu, A. L.. Nicolicioiu, B. Alexe, and D. Teney. Learning diverse features in vision transformers for improved generalization. 2023.
[53] G. Parascandolo, A. Neitz, A. Orvieto, I. Gresele, and B. Schölkopf. Learning explanations that are hard to vary. In International Conference on Learning Representations, 2021.
[54] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, 7. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. Vitioo M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems, pages 8024-8035, 2019.
[55] M. Pezeshki, S. Kaba, Y. Bengio, A. C. Courville, D. Precup, and G. Lajoie. Gradient starvation: A learning proclivity in neural networks. In Advances in Neural Information Processing Systems, pages 1256 1272, 2021.
[56] A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, Mishkin, J. Clark, G. Krueger, and I. Sutskever. Learning transferable visual models from natural language supervision. In Proceedings of the 38th International Conference on Machine Learning, Proceedings of Machine Learning Research, pages 8748-8763, 2021.
[57] A. Rame, C. Dancette, and M. Cord. Fishr: Invariant gradient variances for out-of-distributio generalization. arXiv preprini (.XXx:2109.0299: 2021.
[58] A. Ramé, K. Ahuja, J. Zhang, M. Cord, L. Bottou, and D. Lopez-Paz. Model ratatouille: Recy- cling diverse models for out-of-distribution generalization. arXiv preprint vrXiv:2212.104453 2022.
[59] A. Rame, M. Kirchmeyer, T. Rahier, A. Rakotomamoniy, patrick gallinari, and M. Cord. Diverse weight averaging for -ittr--distiitttion generalization. In Advances Neural Information Processing Systems, 2022.
[60] B. Recht, R. Roelofs, L. Schmidt, and V. Shankar. Do ImageNet classifiers generalize to ImageNet? In International Conference on Machine Learning, pages 5389-5400, 2019.
[61] M. Kojas-Carulla. B. Schölkopf, R. Turner, and J. Peters. Invariant models for causal transfer learning. Journal of Machine Learning Research, 19(36):1-34, 2018.
[62] Rosenblatt. The perceptron a perceiving and recognizing automaton. Technical Report 85-460-1, Cornell Acronautical Laboratory, 1957.
[63] E. Rosenfeld, Ravikumar, and A. Risteski. Domain Domain-adjusted regression or: Erm may already learn features sufficient for oft-of- generalization. arXiv preprint arXiv:2202.0688 2022.
[64] S. Sagawa* P. W. Koh*, T. B. Hashimoto, and P. Liang. Distributionally robust neural networks. In International Conference on Learning Representations, 2020.
[65] W. Samek, G. Montavon, A. Vedaldi, L. K. Hansen, and K.-R. Muller. Explainable AI: Interpret- ing, Explaining and Visualizing Deep Learning. Springer Publishing Company, Incorporated, 1 st edition, 2019. ISBN 3030289532.
[66] V. Sanh, L. Debut, J. Chaumond, and Wolf. Distilbert, distilled version of BERT: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01108, 2019.
14

---

[67] H. Shah, K. Tamuly, A. Raghunathan, P. Jain, and P. Netrapalli. The pitfalls of simplicity bias in neural networks. In Advances in Neural Information Processing Systems, pages 9573-9585, 2020.
[68] R. Shen, S. Bubeck, and S. Gunasekar. Data augmentation as feature manipulation. In International Conference on Machine Learning, pages 19773-19808, 2022.
[69] Y. Shi, J. Seely, P. Torr, S. N, A. Hannun, N. Usunier, and G. Synnaeve. Gradient matching for domain generalization. In International Conference on Learning Representations, 2022.
[70] R. Shwartz-Ziv and N. Tishby, Opening the black box of deep neural networks via information. arXiv preprint arXiv: 1ee.110000000 00810, 2017.
[71] R. Taori, A. Dave, V. Shankar, N. Carlini, B. Recht, and I.. Schmidt. Measuring robustness to natural distribution shifts in image classification. In Advances in Neural Information Processing Systems, 2020.
[72] J. Taylor, B. Eamshaw, B. Mabey, M. Victors, and J. Yosinski. RxIxl: An image set for cellular morphological variation across many experimental batches. In International Conference on Learning Representations, 2019.
[73] D. Teney, E. Abbasnejad, S. Lucey, and A. van den Hengel. Evading the simplicity bias: Training diverse set of models discovers solutions with superior OOD generalization. In IEEE/CVF Conference on Computer Vision and Puttern Recognition, pages 16740-16751, 2022.
[74] D. Teney, Y. Lin, S. Oh, and E. Abbasnejad. ID and OOD performance are sometimes inversely correlated on real-world datasets. arXiv preprint arXiv:2209.00613. 2022.
[75] V. Vapnik. Principles of risk minimization for learning theory. In Advances in Neural Informa- tion Processing Systems, pages 831-838, 1991.
[76] Y. Wald, A. Feder, D. Greenfeld, and U. Shalit. On calibration and out-of-domain generalization. In Advances in Neural Information Processing Systems, pages 2215-2227, 2021.
[77] Z. Wang, Y. Chen, Y. Duan, W. Li, B. Han, J. Cheng, and Tong. Towards duttof-distributiin generalizable predictions of chemical kinetics properties. arXiv preprint, 4cXiv:2310.03155 2023.
[78] Z. Wen and Y. Li. Toward understanding the feature learning process of self-supervised contrastive learning. In International Conference on Machine Leurning, pages 111 11112-11122 2021.
[79] F. Wenzel, A. Dittadi, P. V. Gehler, C.-J. Simon-Gabriel, M. Horn, D. Zietlow, D. Kernert, C. Russell, T. Brox, B. Schiele, B. Schölkopf, and F. Locatello. Assaying -ut-of-dist generalization in transfer learning. In Advances in Neural Information Processing Systems, 2022.
[80] M. Wortsman, Ilharco, W. Kim, M. Li, S. Kornblith, Roelofs, G. Lopes, H. Hajishirzi, A. Farhadi, H. Namkoong, and L. Schmidt. Robust fine-tuning of zero-shot models. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 7949-7961, 2022.
[81] H. Ye, J. Zou, and Zhang. Freeze then train: Towards provable representation learning under spurious correlations and feature noise. arXiv preprint arXiv:2210.. 1005. 2022.
[82] J. Zhang and L.. Bottou. Leaminng useful representations for shifting tasks and distributions. arXiv preprint arXiv:22 2022.
[83] J. Zhang, D. Lopez-Paz, and L. Bottou. Rich feature construction for the optimization- generalization dilemma. In International Conference on Machine Learning, pages 26397-26411. 2022.
[84] Y. Zhang, M. Gong, T. Liu, G. Niu, X. Tian, B. Han, B. Schölkopf, and K. Zhang Adversarial ro- bustness through the lens of causality. In International Conference on Learning Representations, 2022.
15

---

[85] X. Zhou, Y. Lin, R. Pi, W. Zhang. R. Xu, P. Cui, and T. Zhang. Model agnostic sample reweighting for out-of-distributio learning. International Conference on Machine Learning. pages 27203-27221, 2022.
[86] D. Zou, Y. Cao, Y. Li, and Q. Gu. Understanding the generalization of adam in learning neural networks with proper regularization. arXiv preprint 1rxi::2108.1137,,, 2021.
16

---

Appendix of FeAT
Contents
A Notations
18
Limitations and Future Directions
19
C Related Work
19
Proofs for theoretical results
20
D.I Implementation details the synthetic CNN experiments
20
D.2 Proof for Theorem 4.1
22
D.3 Proof for Theorem 4.2
................................
30
D.4 Proof for Proposition 4.3
38
D.5 Proof for Corollary 4.4
40
More Details about iFeAT
42
F More Details about the Experiments
42
F.1
More details about COLOREDMNIST experiments
42
F.2
More details about WILDS experiments
44
F3 Software and hardware
47
F4 Computational analysis
47
F.5 Feature learning analysis
47
17

---

A Notations
We use bold-faced letters for vectors and matrices otherwise for scalar. We use || |l2 to denote the Euclidean norm of a vector or the spectral norm of a matrix, while denoting -Ile as the Frobenius norm of matrix. For a neural network, we denote (x))) as the activation function. Let La be the identity matrix with dimension of Rdxd When comparing two sequences {az} and {ba), we employ standard asymptotic notations such as O()), o(.), 92(.), and (-)) to describe their limiting behavior. Lastly, sequences of integers are denoted as [r]] (1..
Table 4: Notations for key concepts involved in this paper
Symbols
Definitions
X
R"
the input space
y
R
the label space
2 R
the latent space
m
hidden dimension
F-
the j-th filter of the CNN model
W
the weights of j-th filter of the CNN model, containing m hidden units Wj.r
v--)
the activation function of the CNN model
the featurizer : X' + Z learns latent representation for each input example
the classifier
y
Wy
classifier learned j-th round
C
the predictor - ur o yp X J composed of featurizer and classifier when to linear, f can simply represented via dot product 10 -0 set indices all environments
Exll
E
the subset of indices of training environments
the index set of a specific environment
En
the set of environments following the data model as Def. 3.1, where each is specitied as (o.
D
De
the dataset from environment containing N. samples {xf of considered as i..dd from pe overall dataset containing samples from all environments. D (2" Jeeea
D
D
the augmentation set, we use Df to denote the augmentation set separated i-th round the retention set, we use Dj to denote the retention set separated at 4-th round
D
G
G (G', (G*,G') with 2 1 groups round & where G (De) Fizo the grouped sets, new feature augmentation and G" (D) 1--1 is the grouped for already learned feature retention
Le
the empirical risk calculated based on 2", e.g., square loss or logistic loss
the FeAT objective, including fpy the empirical risk at D. and Epy at DJ
(W)
the IRMv1 loss
the first order derivative of L. with respect the i-th sample from enviromment CjeMei à w. E e (M/ Vf a useful quantity to analyze IRMv1 dynamics
Chese
invariant feature learning quantity in Eq.
750r (7baa
spw
the spurious feature learning quantity in Eq. 6
(f)ra()
the noise feature leamning quantity in Eq. 6
18

---

B Limitations and Future Directions
As pioneering work that studies feature learning of ERM and OOD objectives and their interactions in OOD generalization, our theoretical settings are limited to studying the influence of spurious and invariant correlation strengths on spurious and invariant feature learning, based on a two-layer CNN network. In fact, the feature learning of a network can be influenced by several other factors, such as the difficulty of learning feature and the capacity of features that model can learn [21. 28]. Future works can be built by extending our framework to consider the influence of a broad of factors on feature learning in OOD generalization.
Moreover, as there could exist cases where certain features should not be learned, it is also promising to explore how to prevent the feature learning of undesirable features during the early stages of OOD generalization and to further relieve the optimization dilemma in OOD generalization [16], to improve the robustness against backdoor attacks [48], and its further implications to OOD generalization [45]. Besides, it is also interesting to investigate feature learning for complicated data such as graphs [32], especially under various graph distribution shifts [14, 15, 17, 34, 77].
C Related Work
On Feature Learning and Generalization. Understanding feature learning in deep networks is
crucial to understanding their generalization [2, 11, 12, 22, 62, 70J. Earlier attempts are mostly about empirical probing [21. 26, 28, 65]. Elhage et al. [21]. Hermann and Lampinen [28], Shah et al. [67] find that the feature learning of network can be influenced by several other factors, such as the difficulty of learning a feature and the capacity of features that model can learn. Although our data model focuses on the correlation perspective, different correlation strengths in fact can simulate the difficulty or the simplicity of learning feature.
Beyond the empirical probing, Allen-Zhu and Li [2] proposed a new theoretical framework that characterizes the feature learning process of deep networks, which has been widely adopted to analyze behaviors of deep networks [12, 78, 86] However, how the learned features from ID data can be generalized to OOD data remains elusive. The only exceptions are [68] and [42]. Kumar et al. [42] find fine-tuning can distort the pre-trained features while fine-tuning can be considered as a special case in our framework. Shen et al. [68] focus on how data augmentation helps promote good but hard-to-leamy features and improve OOD generalization. Deng et al. [20] studies feature learning when the group-related features are more predictive for inferring group labels. In contrast, we study the direct effects of ERM and OOD objectives to feature learning and provide theoretical explanation to the phenomenon that ERM may have already learned good features [33, 63]. To the best of our knowledge, we are the first to analyze the feature learning of ERM and OOD objectives and their interactions in the general OOD generalization setting.
On the correlation between ID and OOD performances. The debate about feature learning and generalization under distribution shifts also extends the ID and OOD performance correlations along with training or fine-tuning neural nets across a variety of OOD generalization tasks. Andreassen et al. [3], Miller et al. [47], Wenzel et al. [79] found that there often exists a linear dependency between ID and OOD performance under a wide range of models and distribution shifts. While Kumar et al. [42], Wortsman et al. [80] found that fine-tume -uniinning pre- pre-trained models often leads to an increased in-distribution but decreased OOD performance. Teney et al. [74] observed cases where ID and OOD performance are inversely correlated. Chen et al. [16), Naganuma al. [49] studied the ID and OOD performance trade-offs from the optimization perspective.
Our work provides theoretical explanations for different correlation behaviors of ID and OOD performance, as well as provides solution for mitigating the trade-off's in optimization. Theorem 4.1 implies that, in cases where invariant features are more informative than spurious features, the higher ID performance indicates better fit to invariant features, thus promising a higher OOD performance, aligned with observations in [3, 47, 79]. While in cases where invariant features are less informative than spurious features, the higher ID performance implies better fit to spurious features, thus bringing lower OOD performance [74]. Similarly, when fine-tuning a pre-trained model, the model does not learn the features sufficiently well, ID-OOD performance will be in positive correlation. However, when spurious correlations are present as easy-to-learn features, ERM can lead to better fit for spurious features and distort the previously learned invariant features [42, 46, 80].
19

---

Rich Feature Learning. Recently many OOD objectives have been proposed to regularize ERM such that the model can focus on learning invariant features [4, 41, 55, 57, 76]. However, due to the intrinsic conflicts of ERM and OOD objectives, it often requires exhaustive hyperparameter tuning of ERM pre-training epochs and regularization weights [16, 83). Especially, the final OOD performance has a large dependence on the number of pre-training epochs. To remedy the issue, Zhang et al. [83] proposed Bonsai to construct rich feature representations with plentiful potentially useful features such as network initialization. Although both Bonsai and FeAT perform DRO on grouped subsets, Bonsai rely on multiple initializations
the whole network to capture diverse features from the subsets, and complicated ensembling of the features, which requires much more training epochs for the convergence. contrast, FeAT relieves the requirements by performing direct augmentation- rentettiiiiinnnnniiion although Bonsai and other rich feature learning algorithms such as weight averaging [5, 59, 82] have gained impressive successes in mitigating the dilemma, explanations about the reliance ERM pre-training and why rich feature learning mitigates the dilemma remain elusive. Our work provides novel theoretical explanations for the success of rich feature learning algorithms for OOD generalization. Complementary to the empirical observations made by existing works, our work provides the first theoretical explanation for the feature learning of ERM and OOD objectives for OOD generalization.
the grouped subsets, and thus obtains better performance. More crucially,
Besides, there exists rich literature on learning diverse representations for better generalization. Similar to weight average [59], Teney et al. [73] propose to train diverse models to resolve simplicity bias. Lee et al. [44] propose learn diverse solutions for the underspecified learning problem. Nicolicioiu et al. [52] propose to regularize attention heads in transformers to learn diverse features. Chen et al. [13] propose learn diverse classifiers for sample efficient domain adaption.
D Proofs for theoretical results
D.1 Implementation details of the synthetic CNN experiments
For linear activation function vixx) = the logit if (which is a function of W) of sample i in the environment e can be explicitly written as
m
0 = f(W.x, F+
(W. +1, x;) (
F-I(W x1, (1,
(x + xx
m
)E
W.J. 5,1
where W A [W+1. +1, W-1 and W,
for j E {+1). We initialized all the network
j.m_
weights NO of) and we set do 0.01.
The test dataset (x, y) generated through
X1,2 = E, where half of the dataset uses Rad(1 31) and the other half uses Rad(1 - B2). Here & ~ N(0, oz. (Ia Vivi v2vz)) and we chose op 0.01.
x = y/ Vi yı Rad B2)
From the definition of IRMv1, we take derivative wrt. the scalar of the logit gg. Thus, for environment the penalty is
-
Vulwer 0.)) = 1 Σ (( vi
Then, the IRMv1 objective is (we set n2 = 2500 in the simulation)
LIRMvI (W)
Σ
(v + Σ
(VEDE) vf
eEEEr
We used constant stepsize GD minimize LIRMAL (W), and we chose À = 10 (heavy regularization setup).
20

---

Let IRMVI - 1. E (v. be explicitly written as
The gradient of LIRMVI (W) with respect to each wjr can
wwir LIRMIL (W)
2 [x(x1 x9,2) m
E' (VE yf
Te
cEEir
+ 21
C IRMv1
ne
(v(ot) of - (xf (x
x(2) + (vEtE)
of _ _x + xx22)) m
The
eEEtr
Vi (xf,1 + x
cEEtr
jRiMM
+ 2X
e" (v(0E) 01 (x(, + x(22)
mem
eEEtr
CRimM
+ 2A
Σ (M of (x], + x(22)
mem
j(L + 2A0RRMML
1RMv1
e' (vEDE) of (xf, + x(
nmm
eEEt
jRiMv
+ 22
e" (v.O)) of (x + x
Observe that C'RMM is in fact the scalar gradient CiRMvl V I (W) that we want to force zero, whose effect be understood dynamic weightitging of the ERM gradient. Due to its importance in the analysis and interpretation of IRMv1, 1, we tracked CiRvvi in our simulations.
The invariant and spurious feature learning terms that we tracked are the mean of (Wj.r; and (V2) for j E {+1), r [m]. respectively.
21

---

D.2 Proof for Theorem 4.1
Theorem D.1 (Formal statement of Theorem 4.1). For p > denote n & min,ee., ne, n - 2og 2log (16/p) and ó - exp(O(n '}}
ExEEY Dey cg - V A/rr & (w), jvi) and
Define the feature learning terms j.r (v)) for E {=1], r E [m]] Suppose we run iterations of GD for the ERM objective. With sufficiently large n and (5(x) = x, assuming that
n
1 - Ec - 8(1 + 4
1
32
(a, B1, B2 are sufficiently smaller than 2
81 + B2
8(1 (11+cc)
B1 + B2
+
(ar is sufficiently larger than
2
2
2
and choosing
og 0 (n log (m/p)),
of O min d d/2 log 1/2 (- 'mmm dd ny mlllgg(22p)) "1))
there exists a constant 17, such that for any j E {=1}, E [m], with probability least 2p, A and 1 are converging and the increment of the spurious feature pt+1 1,, is larger than that of the invariant feature Arr -A,rr at any iteration E "0....,
Proof of Theorem D.I. We begin with checking the feature learning terms in the ERM stage using constant stepsize GD: W W 17 LIRvv (W*). Note that with (xx) æ the update rule for each w j.r, Vj E {+1, -1,, E [m]] can be written as
jn Σ
w
++1 3."
w.
E' (vt0)) of (xC, + xx,))
jir
The
e
i=1
in
w/
(v(ot) (Rad(o); V1 + Rad(3()s V2 + yf v(..).
fr
eEEr
Define the quantities of interest (the feature learning terms): AG.- - (w) (w),..vvi), r3. - f.r. - (E)) From our data generating procedure (Definition 3.1), know that the first two coordinates of E are zero. Thus, write down the update rule for each feature learning term as follows.
ne
A
Σ
(siti)) Rad(a),,
The i=1
eEEer
nhe (vA Rad(3.)s E' (vf0))
r++1
r
fir
m
Tee
eEEL
SH+Le' Jm
7
vf (ET, E5).
m
The
e
More explicitly, we can write
ne
Rad(a);
A
(8)
m
The
exp
aEE.
r++1
rdt
+
71
nhe
Rad( (B.);
(9)
fr
jir
m
The
+ exp }/
REEL
of (E: Eg')
SH+Le
= Et,e'
17
jo,e"
+
m
Th
1 (xf
(10)
1=1
22

---

Notice that the updates (8), (9) for As,, j.r are independent of j.r. Denoting
Rad(a);
A
De
exp Df)
eEEir
Rad (ad(Be);
4
1 (
eEEtr
we can conclude that for any j E [+1, -1,, [m].
1 = 13, + 1 D1 = 1) Σ
Λ + AI.
k=U
(II)
r#+1= jr + A = Σ
A + r
k=0
Then, we write the logit 0f as
m
0f Σ m
(w),r Vi Rad(a); V1 + of Rad(Be); V2 + x
}E(#1)
Σ
[jv] Rad(a); A3rr + jul Rad(Se); 12 + j.nd
jE(11)
m
Σ m jE(11)
[vf Rad(a); A + of Rad(Be)s 12 + +
of (ad(
m
+ of Rad(3e),
Σ
ir
ΣΣΣ
m
(E
jE(
je(
of Rad(a), 2m
A of Rad(3e) 2m
Σ
A
k=0
k=
L jr
oi Rad(a))
Σ
+ vf Rad(Be))
Σ
m
m
/E|=
Rad(a); EjE(H1) Em A + Rad(3e); Σje(ti) Em
D
Denoting Of A
+
ghe find
Eje(an)
Em
we have
(Rad(a)s 2m Σ A Rad( 2m
t-1 Σ + 03)
of =
k=0
We need the following concentration lemma to control the scale of Q whose proof given in Appendix D.2.1.
Lemma D.2. Denote n - min.ee., ne, 2 Eese,, For > 0, if
of = 0 (n 2log (m/P)),
o3 = o min d 1/2 1/2 log - 1/2 (nm/p), 11 Im (d+ n dlog(nz/))) -1 }))
then with probability least 1 - p. for any Exr, E [ne], it holds that [ O(n 11))
23

---

Then 4 and Af can be explicitly written as
A
Rad(a);
1 exp Rad(a)) 2n Ei ARj
exp RRd((B)) 2n S exp {e})
A
Rad(3e);
Σ
m
1 exp (Rad(a); 2n Et aR} exp (Rad(S)); 2n Ex 24} exp {Q
eEEi
i=1
We are going to analyze the convergences of two sequences (4/. + A and (4} A/). Notice that
A +
A
De
Rad(Be); + Rad(a);
Σ
mmm
exp Rad(a)) - 2n El= A exp Rad( 2n Ek= A/} exp (Q
1=1
A
A
The
Rad(3.); RRd(a))
m.m
i=1 1 + exp
Rad(a); 2n Et AR exp Rad((3); 2n Et-h A/} exp 109}
We can further write these two terms as
A +
A
Σ
2
Σ
eEE.
Te
VE n.]
+ exp
2n Et=b (AF + + exp (0.)
Rad( (ad(9.); (.. Rad(a);==1
Σ
eEEr
m&m
ve Rad(a),
1 1 exp
-2n Σt-l (af + A exp
(
A -5-4% A
2
Σ
mmm m
1 exp 2m Et-l (AF - 2/)} exp (0))
eEEer
iE [
Rad (ad(3.).
Radi =-
Σ
Σ
eEE, mm
rE(na]
+ exppp
-2n E6=1 (A A exp (
(a)
=++1
According to Lemma D.2, for all e E Er, E [nel,p > 0, letting 5 & cxp(OOn - 1, we have + 8 2 exp [Q{} 2 (1 + 8)-1 with probability at least - p. Let Cj. 2 ((i Rad(a)) = (3.) L, Eg) || for any {+1}, {=1}ee Eir. and then define C, Eeeer Cit
24

---

We can upper bound and formulate 4f + A and Ai - Ai as
4 + + I
2
++1+1
C-1-1
m
1 exp 2n Ek (A + } (1 (1+1)-1 + exp exp
-21 Ek (4 + 4R) (1 (1+8)
C+1+1(1 (1+ - C-1--- exp
exp { 2m. Ex= (Af + +
(12)
m
1 5 + exp 2n Ex (4F + }
A
A v
2
C-1+1
C+1-1
1 1xp exp 2m Ek=u (AF Ah) (1 (1+ 1+ exxp
-2n Σ (AF
AR)} (1 (1+
C-141(1 +6) C+1-1 exp 27. Σx= (4) - 40)
(13)
m
1 5 + cxp 2n cxp
E ( 2
Based on similar arguments, we can also establish lower bounds for these two terms,
- -1--1 6) exp 2n Ei (4 + A
2
4 A
(14)
m
exp exp 2n7 Σx= A
+ A (1+ 8)
A - A 2 C-141 - C1 + +0) exp 27 Σ ( A
(15)
m
1 exp 2m Σx=O (4F A (1 +1+0) 8)
The upper and lower bounds (12), (13), (14) and (15) imply that the convergences of f + 4 and
{Af A } determined by recursive equations of the form Q' CxD (n Σ1=1 2²) We
1+C (n 2*)
first establish that with suitably chosen 7), the sequences ( + A and { A are guaranteed to be positive. Observed that for the Q'-type recursive equation, the sign of Q is independent of and only determined by the constants C1, Cz. C3. At iteration 0, (14) and (15) give
2 C+1+1 - 11
A
A
(16)
m
2+0
2 C-141 - -C+1-1(1 +
4
(17)
m
2+5
To proceed, we need the following concentration lemma to control the deviations of the constants C+1+1, C+1-1, C1-1+1 and C--1-1 from their expectations, whose proof is given in Appendix D.2.2.
Lemma D.3. For 0, considering two environments and denoting EC A
2 (16/o)
with
probability least p. we have
C+1+1 -11 B 32) < EC
C+1-1 - a) 82) S EC
(18)
101-1-101 B1 -3) S cC2
7---1- 1011--((1 a(3 + -B) S 6C-
Using Lemma D.3, with probability at least 1 - p, the constants T+1+1. U+1-1, T- and T close to their expectations.
Based on our assumptions that
1 - - 8(
(a, B1, B2 are sufficiently smaller than 1 7
a, 31,33
7
8 + 32
8(1+ec) 2
(a sufficiently larger than 81 + B2
a
+ cc
2
2
25

---

it can be verified that with probability at least 1 - 2p. 4P + + > 0. - 4 > 0.
Then, at iteration from (14) and (15), we see that long require
C+1+1
C_1+1 C+1-1(1+8)
77 < min
log
log
2 + A
C-1-1(1+ + 5) 2 A
it holds that A + A > 0, A
A > O. By recursively applying this argument, we see the
requirement for 17 to ensure that A
A 0 and A -A for any E {(.....T' ......TJ] is
C+1+1
C-1+1 +
7) < min
log
log
2 2 k= (4F
A C-1-1(1 +0) 2
2 k== (4F AN
(1
(19)
In other words, for the Q'-type recursive equation, as long as Q 2 there always exists a sufficiently small to guarantee that the whole sequence {Q*} is positive. From now on, we will focus on the case where the two sequences { + A and f A decrease to an ca > error, i.e. (o...)
r) {4i + A 4 A Es.
Then, we show that the two sequences {Ai- + A and ( - A decrease monotonically, which thus leads to more refined upper bound for 1) (19). Inspect the upper bounds (12), (13) at iteration ++ 1, which can be written as
A+t + A S
C+1+1 - C-1-1- eepppeeeppp { 2m Et (A + At}] exp (2n (A) + (1))(1++ + +6)-1
2
H+
1 1+ 1xexp (2n- Et-o (A$ + } exp (2m (4/ (1)) +
A 1+1 1 S
2 C-141 - C+1-1 exp 2n Et=o (AF A exp (2n (4) - (1)(11 + +6)-1
a+
1 1 exp 2m Ei=o (A } exp (2 (4) A1) 5)
1 log (1 + 5), -.-25 log (1 + +0). vi (O....,T) = |o+ti
Requiring that 7) > max
< log (1 + 5), we have
- - --1-1-expp
2n Et=o (A + } exp (2n (4) + AW](++0)
H1_ 2
m
1 exp 27 Ekzo ( + } (1+
< A + A/
C-141 -11-1 11-1-cx0 exp 2n Eteo (4/ A exp (2n (4) - A11(+
2
1+]
{2n Ei (A }/) (1+5)
1 + exp exe
< -1
where the last inequalities use the lower bounds (14) and (15).
Based on the above discussion and (19), we can now clarify the requirements of 1) for the sequences { + } and { A to be positive and monotonically decreasing:
(2 + 5)
C+141
log (1 + 8) < min
log
4T
+11(1+++111 + 8) - C-1-1)
C
C-1-1(1+0)
(20)
m(2 + 6)
log C-141
4T -T(C-1+100 + 8) C+1-1)
C+1-1(1 + 5)
which uses the upper bounds (12) and (13) at iteration 0. The constants C+1+1. C+1-1, 1-141 and
C-1-1 can be substituted using the concentration bounds (18) to generate an upper bound for that only involves a, 31, 32, m, 6, T, Ec. Here omit the precise upper bound for clarity. Note that
26

---

the left hand side of (20) approaches 0 if - 0, which means that there exists constant choice of in (20) if sufficiently large in Lemma D.2 and D.3.
To conclude, view of (1),, the convergences the sequences {A} + A and (4}- - A imply that fr, and r converging, and the positive sequence { A indicates that the increment
of the spurious feature r++1 Itt is larger than that of the invariant feature A A at any 3u
j.r J.r iteration E {O....,T- 1..
D.2.1 Proof of Lemma D.2
First, recall some concentration inequalities for sub- Gabssia random variables. Since & (O, of vIa-vivj -v2v )), for (i', (i',')) o (i, e), using Bernstein's 's inequality for sub-exponential random variables, we have for sufficiently small a 2 0,
2
Pr PK(IKES- E 2 aa S 22exp
40 - 2)
a²
Pr (111e511 od 2) 2 a} S 2exp
5120 2)
Moreover, for & ~ NO og) (indicating the initial weights w the standard Gaussian tail gives
1
m
ma² 20
Er
exp
m
Denote n Eeee, Re, n - min.ge,, Ne. by properly choosing a for each tail bound and applying union bound, we can conclude that for p > 0, with probability at least 1 it holds that Vi, i', e',
(E) 45'1
< 20
(d 2) log
8n2
|e s on 2)) 160 p /2(d 2) log
8n
p
2 log
32m
KSr, (. s 20 (d- 2og
16nm
Σ &r
S 00g
p
P
We start with bound the growth of a
By bounding the update rule (10), with probability least
1 - P. we have
1, Juni
ete'
The
1
Jiree
m
Th
1 exple
eEEr
nhe
Thee Janee
m
The
nePer
i=1
0.ed ja.a
+ (t 11
Σ
(Er <
m
eEEr
i=1
1K5r (++1)-
|155 112 +
Σ
" mne
165 43')1
mne
16nm
< 22,0ou (d-2) 2og
Tno?
8n + 2m (d - 2) log
8n?
+ ma
(d 2) + 1 2(d - 2) log
p
27

---

Then, we can bound
[
as
1
1
2
[
2
+
2
&r
m
r=1
2
32m
16nm
log
+ +0,000
(d 2) log
m
no
8n
8n²
(d - 2) + 16 2(d - 2) log
+ 2 (d - 2) log
p
P
Thus, with sufficient small Co. pp. i.e.
og 0 (n
-2² 2log"* -1 (m/p)),
1/2 log 1/2 (nm/p); T -- (a (d+ ny (ndllg(n(/)) "}}
of O min d d/2
we ensured that [Q O(n
D.2.2 Proof of Lemma D.3
For e E Eir, using Hoeffding's inequality, it holds that
Pr 1
- - 8e) 2 20
S 2xp 1-2a²
==1
Considering two environments, using a union bound, we can conclude that
a 12
Pr (1 (2 B1 32)l a} 2 - 4exp
Thus, for p O, with probability at least - 4. conclude that
log (16/p)
C+1+1 -1- a2 B1 B2) S
n
Using the above arguments for other constants C+1-1, C-141 and C-1-1- and applying union bound, we obtain the desired results.
D.2.3 ERM Feature Learning with Non-Linear Activation Functions
It was numerically observed that in the early stage of (stochastic) GD training, the learning dynamics of neural networks can be mimicked by training a simple linear model [35]. Hu et al. [30] rigorously proved this phenomenon for training two-layer neural network with l loss function in the Neural Tangent Kernel (NTK) region. We briefly summarize their results here: For a two-layer fully- connected neural network (with fixed second layer {v²
(W x)
Σ v.ab (w,.x// NV.x/Va). r=1
(21)
Vm
[ y)² and the ERM objective LERM (W) when using GD: W = W ERM objective, the following holds.
considering the la training loss ()
& (ffre x,), (.).v)) vi). (W.,,)
VLER((W' to minimize the
Theorem D,4 (Theorem 3.2 of [30]). Let ant E (0. 1) be fixed constant, and vb(-) be a smooth
(with bounded first and second derivatives) or piece-wise linear activation funcctioo Suppose that ((d)+ani) and
2(dl+omd). Suppose that 1) < d. Then there exists universal constant c 0 such that with high probability, for all t = O d log d) simultaneously, the learned neural network fic and the linear model finm (defined below) at iteration are close on average on the training data:
and satisfy
(fbo(xi) Simm (xx)) = O(d-R(anw)))
28

---

The linear model (1, = ,3 R(x) is a linear function of the transformed data R(x)
va where [ and are constants related to vl and the dataset distribution (see (5) in [30] for
formal definitions).
We show that we can relate our data model to the dataset setup in [30], and thus by analyzing the feature learning terms for the linear model fism (B, x) similar to the analysis* in Appendix D.2. obtain similar results as in Theorem D.1 in the early stage GD training, but with an error of O(d-1(ow))
Recall that our CNN model is f(W, =)==++1 +1, F-1 (W -1,11,,, x), where F-1 F-1 -1, x) and F-1 T--(W-I1 x) are defined follows:
(W x)
[b(w]]
[w]],x))) x1) + v/(w [W]] ,J E (+13.
We can cast this CNN model into an instance of the two-layer fully connected neural network defined at (21) by specifying the values of {u, +//3 and transforming the dataset as
{ Va y Rad(a) vi - (ad(3).v2)
vV J. Then by tuning the norms of V1, V2 and &.
0
obtain a dataset that satisfies the input assumptions in [30]. Note that this cast drops the shared variable of our CNN model and thus might lead to slightly different training dynamic. To fix such gap, we can leverage Proposition 6.4.1 in [29] for the early stage behavior training a CNN model.
Based on the above ideas, to formalize the convergence results of the feature learning terms in the non-linear case, it remains to re-derive the analysis in Appendix D.2 based on & loss function, which follows similar line of proofs and has simpler dynamic.
*Note that when (xx) = z, our CNN model can be viewed as linear model with re- parammttrrized weight matrices. Thus, the discussion in Appendix D.2 can be viewed as studying the feature learning terms for linear model with logistic loss function.
29

---

D.3 Proof for Theorem 4.2
Theorem D.5 (Restatement of Theorem 4.2). Consider training CNN model with the same data as in Theorem 4.1, define
c(t) & [Giuun (W, t), CiRMa (W,t)
CRMM (W, 1t]
and A Amin (Ho), where we define
The H (0,,,
Rhe & ((w,,, xxxx))x
H e,e'
Σ
Suppose that activation function is smooth, v'(0)) 3. (2) (x')) < B| x'1 and Lipschitz |x(( < L. 1x(xx) - (xxr) < L|r x'). Assume that dimension d = (mlog(mm/)) network
width m (1(6))) regularization factor X Z 1/(oov/E noise variance σp = O 2), weight 18,, 1/2.11/
initial scale op = O( 2m (olgg1//)) 1og(1/0)), then with probability at least 1 6, after training time
log(1/e)) we have: 1/2da
T= 22
||((T)l2 SC. c,
inv Vir (T) = o(1)
(T) = o(1)
Before proving Theorem D.5, we first provide some useful lemmas as follows: Lemma D.6 ([12)). Suppose that and 22(oog(41) 4/))). Then with probability least 1
d/ < | 30
for all i, [ni]. Lemma D.7 ([12)). Suppose that d 2 Q(og(mmn00)) m ((log((/8)) Then with probability at least 1 8,
(w j.r vi)l 5 / (olog8m//// 5o ||2,
(0)
(( Sor v2) S 20mgg(8m/ (8m/0) To V2
(0)
Kw s.r Ei) 2 2 loglo8mmm// anopVd vd
(0)
for all [m], {=1} and i [n]]
Lemma D.8. Suppose that 8 > 0 and d = ((log(4m/8)). Then with probability at least 1 8, for all E [m] and j [1,, we have
o [w),) 30/d/2.
Proof of Lemma D.8. By Bernstein's 's inequality, with probability at least - 8/(2m) we have
[|],,,-(0]] agd = O / /dlog(4m/0))
Therefore, as long as d = 2(log(4m/8)), we have
d//22 < (0))3 < 3//d/2.
30

---

Proof of Theorem D.5. The proof is by induction method. First we show the gradient flow of weights by IRMv1 objective function (5):
dw 5it(() 11 Ww. LIRMVI (W
dt
7
Σ
(()))((W, (t), jv n
ett (ww,, -J9lE:
nm
TEm
BEE
BEEL
2nd
2nd
The
CRMVI
Σ CO (t) A.V.N
CRMMII
Σ
Cf (.),
nm
num
eEEu
EEBu
2nd
2nd
Σ Σ COW
CIRMOL
GRRMAI Σ (l0) (t), ())
7mm eEEt
TLYR
eEEu
Σ
(1 (1+22
IRMvl (1)
(OW (1, (1) jvf
nm eEEr
"
Σ
(1+ 2AC
IRMVI (t)
Σ (tw,. juLE:
nm eEEt.
2nd
Di ((w),))) Mi 2nd
MRMvl
Σ
C
IRMvI
o ((w,, (1), E.)
eE
1=1
BEEt
where CRRMvI opposite sign f"
E
e V.v. and v = Rad(a); + Rad(
V2. Note that e" has the
The
Then we look at the dynamics of CRRMa (t) according the gradient flow update rule:
dCiRMvi(W
m
iRiMMi (W,t)
dw;,r
(t)
dt
0CiRMvI (W.1)
aCiemvI (W,))
2AC IRMV1
(W,t) t) ΣΣ
(t)
(t)
r=1
0CRMM
IRMvI (W,t) () OL(W, 1)
Ow
21
Σ
CiÍ
CRMVI (W. f) H. He) ge (t);
where we define (e) = Σ
80 DC (W.t) (w...()
ace'
(W.)) (W.I) Ow...
and (t)
Σj E m2 the dynamics of c(t) CiRMv (W, t), CRRMA CiOMMWWWt, (W,t).... CRMMI W,t) in compact way:
(0mm. aL(W,))) Thus H(t) is an Eul x EEll matrix. We can write ows.r(()
dc(t) 2A H(()((() + g(()
(22)
dt
Our next step is to show H(t) is stable during training. To proceed with the analysis, we write down the expression for 8C (W.)) gd:
(w)
ORI (W(I)
1 me
(())
jvl nam
((w,
juft.
(t)
nmm
"0 (f) iol +
(ww -E:
m
mmm
31

---

When we consider non-linear activation function (xx), the entry of matrix (t) can be computed as follows:
dC IRMV1
.)
aC IRMV1
(W.t)
H.
jw (t)
(t)
Re
Ret
GO) T
(t)
(
Σ
& (t)) (1)00) Vs
m
1
MEONEE Σ
SROIN T
s (0 (OivE]
m
m
1
Tdd VEOMEE Σ
VEONE +
nam
m
1
(1
Σ
Σ
10 (1)
nem
m
r=1
P=1
P=1
H e,e' The matrix H is composed of eight elements. In addition, we define
H (1) - (1 + (1)
(0, + H c.e (1) + H (1) + H (t).
m
nee
nd
H e,0'
ΣΣ
(w
2 (( v
mem
m
nee
(((ws,,(0), vg
1 ((w,,, (0), v5))ss'
2
1=1
Then show that:
H
(t)
e,e'
(IEOO T
Σ
V0400
(O)=
2
VODOOO Σ
(1))
Ml
(O
(OR
Σ
(OME
2
VO4OV
Σ
(O)E
2
Σ VOGOVE
2
POCM
Σ
Σ
2OO
1
me.
ne
(O)COM
2
(0(0))
8'(0)
2
me.De
nc
2
Σ
- (000-000))()
Σ COD
- ((0)) Vi
The
+ To
(())) T Σ (0)
(c,(1)) vg"
+
Σ 2'()) (((()) + 2/) vi Σ
3 (0)
Mo
x=1
1 + 1g + 13 13 +1
32

---

where we calculate each item as follows:
R2
1 De
(W) - (10) Σ
400W
Bl[w;: (1) - w..- (, Σ
40000
2,32
|w (1) w., (0)2(
I| (t)aall)
Rel
(c) 32 R(R + good
m
where we have used - property of the activation function and Gauchy-Schwand inequality; inequality (b) by smoothness property of the activation function and Cauchy-Schwarz inequality. Besides, we have used [ S for all E ne and E Eat; inequality (c) is by the fact that ||v S for all E ne and E &all and Lemma D.8.
(.)((t))lz Besides, inequality (a) results from applying the smoothness
Similarly, calculate the upper bound for 12 as follows:
The W
(O T Thtt
1 Re!
(0)) v - ('0))
-'
32,3 R(R R + 3 oodd
m
Next, we give the upper bound of Iz:
13
2
1 W (O) T Σ (0)) (4((1++ 2
vz¹
mn, ner
2
Thet
(e((t) 1 2 v.
B [0)a
T
B w (0))lellv
mn, Te!
64,32 Food)
m
where we have used y which is defined as follows:
m
10 (1)
((w)( + (()a))
m
(a)
2LR,
where inequality (a) is by the Lipschitz property of non-linear activation function and we have used the bound for (t( + :
1
exp(-1f S(W, x;,t)) ())
1
(t)
2
+ exp( -- f(W, x,, t))
2
1
2 1 1+ cxpluf : f(W, x; t))
1
< max
2 1 + exx(2LLR
2
+ exp( 2LR
1
1
max
= O(LR).
2
2 7 2LR
2
2 - 2LR
33

---

and we provide the bound of (t) ::
(t) 1
exp(->f f(W,((WWWW x,, t)) cxp( V8 f(W, x,,t))
expluf S(W, x; 1)) + exp(- f(W, xi, 1))
2+ 2exp (2LR) =((L)))))
Similarly, give the upper bound of Ia:
2
(0 (())) + 1
4
vi 8
(0)
De!
2
F==1
2
1
B Bll (0)1l2) (01l.
2
vv?
B
mn, ne
648 LR)((odd)
m
Together, we obtain the upper bound for H (t) H
64 godd
1. ond)
H H (t)
H 1.00 e,e'
+
Then we calculate the upper bound for the residual terms:
m
THe VO
2
(t)
s0 (t)ool vg
mmm
m
j
2
Ne
Th (t)) Σ
('M) Oile v
me
2
M
The
|| (ww.0
mn
12
where inequality (a) is by the smoothness property of the activation function and Cauchy-Schwarz inequality, and inequality (b) by triangle inequality and the fact that ||v for all i n, and E Eau- and 1 < 1 for all i E [n] and Lemma D.8. Similarly, we further provide the upper bound of residual terms:
TE (E)
H (t)
COO Σ
=
m
me
r=1
2
ne
me
2
Re
me
64,32 LR3
34

---

Similarly, we further have that:
H (0)e =
PEONT T
Σ
20 (Ojing
m
me
Me
2,
Iw (MllMMM
12ll2llllll2 (Oliv
64,3 LR'
Keep going on, we provide the computation results further:
H (1)
ΣΣ
Σ s4 Es
mmm
me
j rm
(a)
2.3²
mnn Thee
(b) 2 R² of
m
where inequality (a) is by smoothness property of non-linear activation function and Cauchy inequality, inequality (b) is by Lemma D.6 and Lemma D.8. Next, we calculate the
H (t)
3 (OOEE
nem
nemm
r=1
2,3²
I (1ilallE (1la (22 (oess)
332dL*1
m
Similarly, the next H term be calculated follows:
H (0 = Σ nem
Σ
& wean E
me
2,3²
ww www.l
I (t)lel)))
me
4.3 d'Ln
m
Finally, we have the upper for the last term:
m
H (. (t)
Σ
SEROEES (1)
nam
m
2.2
(Wlllee (12 (0E5)
4/2 o
35

---

To summarize, we have that,
32,8 R(R + 3 oo0d 1283 L
1283 LR
( (t)
H
+
28 RPodd
88
of
m
m
m
BLL
0
m
where we have used o, Od 2). R o(1). and ou O R/d/. Furthermore, we show that the perturbation term in Equation (22) is bounded during training. In particular, we show the complete expression:
iRMMI (W,1) OL(W,)) () (w
(t)
Σ ==
(t,,r(t)
Σ
Σ ((t)) (t). ivf
nmm
nmm
if (1) vi)) ish
((0) (1,(())))))) iv.
mm
n.
n
((t)e' (t), (t), juLE
(())) iu
The
n
n
Ei)) jute Σ ((s,
(t).
Mem
I + 12 I3 Ta-
Similar to the computation process for matrix H, we adopt divide and conquer manner:
2,32
1 S Me
| (0lball
32 [ + 0))2
m
The techniques used are the same when deriving upper bound for matrix H. Next, we have
2,32 me
122 <
643 RR
m
The last second term can be calculated as follows:
2,32
The
(tllblll
me
23 R2o3d
m
Finally, we show the upper bound of last term:
2
14
Iw (1)lal11))
mn, The
4
36

---

In summary, we have the following inequality:
- 32,3 R2
643 LR
28 R8 R
Ige (t)(t)
+
m
m
m
m
00 1BR.R
where we have used 0p the dynamics for |
2) and (1). With all the bounds at hand, we are ready to have
dlc(t) 22cc - (()H((())))()) c())((() S -And)le((112.
(23)
dl
which requires that [H([)) - S Ao. This leads to the following inequality:
H(( H 11 Ht) H Ir S (t)
Err LR
< Ao.
which leads the conclusion for R follows:
Anm
R
(24)
E 3/2
Besides, we have the inequality that
E L2
ggl S
AAollc(0)
(25)
m Combined with Equation (24), we obtain the condition for À as follows:
x z 1/(o0v ( ) (lug(en/s)) mddo
(26)
By inequality (23), taking the convergence time T = Ω (
we have that:
||c(T)lla c.
According to the gradient descent for IRMVI objective function, the evolution of coefficients can be expressed as:
The Σ
inv (t + 1) = 73p (t) Gor
inn
7
(1+220Hm ( (t)) The 1=1
(())((())aa(a))
eEEs
y7t
Σ (U((t) of Rad(()))
20
IRMVI
30r 21 (1 1) 7g.r (1)
Σ (1+ 2AC (1))
Rh Σ (()(((
eE
yru
20 IRMVI
Σ
(((0)) Rad(
m
The
cEE.
1=1
Then we have,
1 ns
h (t 1)1 S h (t)l
+
Σ(1 2122 2ACiRMM (t) (()(()))))))
m
The
ecEir
20
(1(()) - Rad(
IRMvl
The
eEE
inn (t)l + " TEB cctt)l22
< ho
37

---

Similarly, we have,
(t) + C T Eer ABREL
8pu Nirr (1 + 1)1 S (t),
|c(()]l2-
At the time step T. the feature learning satisfies that:
Forr (T) c C""Veer ASR LT
JV Etr AB
||c(0)]l2;
2pM (2) S C
||(()]l2.
m
m
To make sure that (rr = o(1) and 75 (T) o(1) need the following condition:
" Eer ABR
||c(0)l12 < d
(27)
m
combined with inequality (24) and inequality (26), we have:
1/2 00 S d X log log(1/))
D.4 Proof for Proposition 4.3
Proposition D.9 (Restatement of Proposition 4.3). Consider training the CNN model with the same data as Theorem 4.1, suppose that (xx) = 2, 11... Tr.. - 1), and (),r2((1)) = 21, 1) the end of ERM pre-train t and Etr {00.25, 0.1), (0.25, 0.2)). Suppose that 0, and n >> C (og(())), with being a positive constant, then with a high probability at least
1 6, have
CiRMa (t1) =
Viaalt +1) Vi.r, (t1).
Brr(th + +1 (t1r)
Proof of Proposition D.9. According to the gradient descent for IRMV1 objective function, the evolution of coefficients can be expressed as:
(t. + (t) 7
(1+
2ACTRMVI (t))
Σ Radd(())))
m
BEE
20
Co yf Rad(())
m
eE
n
(1 E 2XC IRMVI (t)) Σ (d)))
nEEu
20
Co vf vadd(33))
IRMv1
The
where (" f(W, x{))
exp( /fW
To simplify the notation, we further define
the Af (ad
Re
and
A5 The
38

---

Similarly, we define
1 11
Bi Re
e (3e)s
and
1
B2 Te 1=1
Rad
In the limit of n - oo, we have:
lim A}(t1) =
n-000
+ -1
lim Af((1)
+ - -0]((1)))) a)
a-+0c
+ +1/ (1
- -
lim (t) = 1/(1 [1000----00110000--------1------
n-000
- - -1
lim BT(4)= 11(11 -) 11(1+ -1
1/(1 + ) 32) - J032.
and,
lim Ab(t1) = - 3 + - Tt 73 (1 -2)
- +)))))a( 81)(1 - 12) 0 aL 1 72),
lim A3(t2) = - (1 4 - a) 10)(++>) + - (1 -1 - 72)
R-+00
+
2 - 1010-)))1--)) (2) th - 11 + 72),
lim Bb(tt) = - (1 21 (1))) 31)(1 +2) - - -, + 72) n-000
+ - ( - 3 +72),
lim BB(1g) = / (11 - 2 - a) - -))())))((1111111111))))))))1111 +22) - ~ 7/ اد 1 0))22(711 + 72)
n-000
+ - (1 a -
- (-) + ー /1 + - 2 + 72)
Because (ad(a); and R3d(3); random variables, applying Hoeffding's 's inequality, have with probability at least 1- 6,
A1(()) lim A((1) (lo((1/8) =--00
Similarly, we can apply the concentration bound to other quantities and obtain the same bound.
By the assumption that (t1) (t 1) and (t,)((t)) 7am2(t1 1). we have that L. (f((1) = E. (t1 0:
lim (4)(4))+ + (((4)) = -1((1 + (1 - - 2 B)) - 1/(1 + - a) + +2) + +1/(1e - 1/(1 - - )a + +3) =0
lim (Bi(t)) + Bi((t) 1/0++ 3)) 1/(1 21 12 (1 a)(( + 32) n-+00
+ (1) 32) + 1/(1 + 7 ((1 + +82) =0
Solving the above equations, we have,
1
7 (tz) (o
16 (t1) = log Gm/ /Gb)
2
2
39

---

where we denote (t,) 1im (1()) and 72 (t1) - lim,- (12), Gm = ((1
A) + 11
and Gb = ((1 B) + y - 1)2 41) with A
aB3 - 32) ([[[-a)(22 a) 31 (2)) and B = a(2 - 31 (22-3))((11100)
- a) (31 + (2)).
By the convexity of function f(x) c², with constant C, we have:
Hog(11 1/6)
mm 1 < |eT -²
c 111+
1/(1 - 1log(1 (1/6)
h22 22 e²2 c²
Then we know that,
CRMMv1
Aovy = 71 + 72
t=1
C 1RMv1
Ciltu? = 71 A? + +2
Therefore, we have that:
GRRMM
C CRRMvl =0
Then the evolution of coefficients reduces to
(t + 1) 7 (t) - 7
Σ
(1+ 2AC IRM (t))
20 IRMv1 At())
eEStr
(t + (t)n2(t)
7
Σ (1+ 2AC CRMvI (t)
n
20 iRMMV (t)
m
m
eEET
Taking the solution of Yi.r. (f1), (t.)) and value of a, 31, 32, we arrive at the conclusion that with high probability least 1 - ō and Ci (1g(1/0) with C₁ being positive constant, we have:
(1, + +1 > 71.,1 (11) 26, + 1) < (t,))
D.5 Proof for Corollary 4.4
Corollary D.10 (Restatement of Corollary 4.4). Consider training the CNN model with the data
generated from Def. 3.1, suppose that (x)) = rr [t1r o(1), and ().) = O(1) at the end of ERM pre-train t and &,, = (0.25, 0.1), (0.25, 0.2) Suppose that 0, and C Clog(L/8)),
with C being a positive constant, then with a high probability at least 1 - 8, we have
Ma (t1 + 1) < (t1).
Proof of Corollary D.10. Recall that the feature learning update rule:
(t + 1) (t)ra())
7
'(
2ACTR (t)) ne
Σ (())Rad(a); i=1
m
eEE..
yr
Σ ORiMM
20
1
R2 C vl Rad(a)s.
m
ne
cEE.
nh
15 + +1) ():n2(1)
7
Σ
(1
2AC MRMMI (t))
> (ad
m
The
nEEu
Re
Σ
20
Co 1f Rad
Th
i=1
40

---

Taking the value of Tr., (5.r2((f)) and, we can conclude that:
lim (1((1) -
n-000
= =/(1++77aa-111---))++1+++7777711 1)(1))
= (00--l(((+=c(](0)))+++++++1"
= (2a- [2a-1)[1((1))
lim
n-000
= (2a - +/]]0+++++((()++)++)+++++
lim (1(+)) = -1100-+-000000000---0-+1+++11111
= =-][0=+00000))1(11++=>"
lim
= -1/t))=--++))] + 7) a1)
10000
1/(1 +[1++
= +//1++
On the other hand,
lim = httpa/aaaaaaaaa/aa..r..///./a/aa
n-+20
= |,
|| +171
lim
11)
lim ||,_,//|_|||||,.0,"000000-,,,-,, -133)))))
10000
lim |1000-/0||||00000000000000111111 ((
Finally, taking the value of environment of (a. 31, 3,) = (0.25, 0.1, 0.2), we conclude that with high a probability at least 1- 5 and n C2 log (/5) with C₁ being a positive constant, we have:
(1 + 1) < (11).
41

---

E More Details about iFeAT
As mentioned in Sec. 5.2 that, when the featurizer implemented deep net that have massive amount of parameters, backpropagating through Algorithm can allocate too much memory for propagating with 2K batches of data. It common for many realistic benchmarks such as Camelyon17 7 and FMoW in wilds benchmark [39] that adopts a DenseNet [31] with 121 layers the featurizer. To relieve the exceeding computational and memory overhead, we propose a lightweight version of FeAT, denoted as FeAT. Instead of storing all of historical subsets and classifiers, iFeAT iteratively use the augmentation and retention sets and historical classifier from only the last round. In contrast, previous rich feature learning algorithm [59, 83] incurs high computational and memory overhead as the round grows. For example, in RxRxI, we have to reduce the batch size of Bonsai to allow the proceeding of rounds 3.
We elaborate the detailed algorithmic description of iFeAT in Algorithm
Algorithm 2 FeAT: Feature Augmented Training
Input: Training data Drr: the maximum augmentation rounds K;; predictor f f= w o yp; length of inner training epochs termination threshold p;
2: Initialize groups G +- Dir, G
0::
3: for & E (1.... K] do
4: Randomly initialize wi;
5: for j el do
6: Obtain FEAT with G via Eq. 7:
7: Update wk, to with FEAT: end for
8:
9:
// Early Stop if fü w o P fails to find new features.
10:
if Training accuracy of f& is smaller than then
11:
Set K k- 1 and terminate the loop;
12:
end if
if > then
14:
// Hence it doesnot need to maintain all historical classifiers. Update wa + (w wx);
16:
end if'
Split Dir into groups Di, Dj according to fa;
18:
// Hence it doesnot need to maintain all historical subsets.
Update groups G {Dg). G" + [D]]:
20: end for
return ʃ w ip:
More Details about the Experiments
In this section, we provide more details and the implementation, evaluation and hyperparameter setups in complementary to the experiments in Sec. 6.
F1 More details about COLOREDMNIST experiments
Datasets. In the controlled experiments with COLOREDMNIST, we follow the evaluation set- tings as previous works [4, 16, 83]. In addition to the original COLOREDMNIST with & = (00.25, 0.1), (0.25, 0.2)) (denoted COLOREDMNIST--ON correlated with labels, we also incorporate the modified one (denoted as COLOREDMNIST-O1) with & = 100.1, 0.2). (0.1,0.25) where invariant features are better correlated with labels, since both cases can happen at real world.
where spurious features are better
Architecture and optimization.
ensure a fair comparison, we use 4-Layer MLP with a hidden
dimension of 256 as the backbone model for all methods, where we take the first 3 layers as the featurizer and the last layer as the classifier, following the common practice [25, 39]. For the optimization of the models, use the Adam [37] optimizer with learning rate of Le and
42

---

weight decay of le 3. We report the mean and standard deviation of the performances of different methods with each configuration of hyperparameters 10 times with the random seeds from to 10.
Implementation of ERM-NF and OOD objectives. For the common pre-training protocol with ERM, our implementation follows the previous works [83]. Specifically, we first train the model with {0, 50, 100, 150, 200, 250} epochs and then apply the OOD regularization of various objectives with penalty weight of (1e1, le2, 1e3, le4, 1e5). We adopt the implementations from Zhang et al. [83] for various OOD objectives, including IRMv1 [4], VREx [41],IB-IR] [1],CLOvE [76].IGA [40] and Fishr [57] Besides, we also incorporate the state-of-the-an OOD objective proposed by Chen et al. [16] that is able to resolve both COLOREDMNIST-022 and COLOREDMNIST-01.
Evaluation of feature learning methods. For the sake of fairness in comparison. by default, we train all feature learning methods by the same number of epochs and rounds (if applicable). For the implementation Bonsai, we strictly follow the recommended setups provided by Zhang et al. [83], where we train the model with Bonsai by 2 rounds with 50 epochs for round 1, 500 epochs for round 2, and 500 epochs for the synthesize round in COLOREDMNIST-OP COLOREDMNIST-0 round contains 150 epochs, round 2 contains 400 epochs and the synthesize round contains 500 epochs. For the implementation of FeAT, we train the model with 2 rounds of FeAT in WOLOREDMNIST-ODE
While in
and 3 rounds of FeAT in COLOREDMNIST-01, where each round contains 150 epochs. While for the retain penalty, we find using a fixed number of 0.01 already achieved sufficiently good performance. ERM only contains 1 round, for which we train the model with 150 epochs in COLOREDMNIST performance degeneration in COLORED NNISN-NSSTT model with ERM by 500 epochs to match up the overall training epochs of FeAT and Bonsai. We provide a detailed distribution of the number of epochs in each round in Table 5. It can be found
as we empirically find more epochs will incur severe While in COI.OREDMNISTT001 we train the
Table 5: Number of epochs in each round of various feature learning algorithms.
CMNIST-025
ROUND-1 ROUND-2
ROUND-3
SYN. ROUND
ERM
150
BONSAI
50
150
500
FEAT
150
150
CMNIST-01
ROUND-1
ROUND-2 ROUND-3
SYN. ROUND
ERM
500
BONSAL
150
400
500
FEAT
150
150
150
that, although Bonsai costs 2 - 3 times of training epochs more than ERM and FeAT, Bonsai does not necessarily find better feature representations for OOD training, as demonstrated Table, In contrast, FeAT significantly and consistently learns richer features given both COLOREDMNIST-005 and COLOREDMNIST-O than ERM, which shows the superiority of FeAT.
The termination check in FeAT. A key difference between FeAT and previous rich feature learning algorithms that FeAT able to perform the automatic termination check and learn the desired features stably. As elaborated in Sec. 5.2, FeAT can terminate automatically by inspecting the retention accuracy. To verify, we list the FeAT performances various subsets of COLLRED 025 and FOI.OREDMNIST-0F at different rounds. We use termination accuracy of 130%. which trades off the exploration (i.e., training accuracy as 80%) and the retention (i.e., retention accuracy after
as 50%) properly. As shown in Table 6, in COLOREDMNIST-OP FeAT learns sufficiently good features at Round (3), respectively, it is not necessary to proceed with Round 3 (4) as it will destroy the already learned features and lead to degenerated retention performance (i.e., the sum of training and retention accuracies is worse than 130%.
(COLOREDMNISST-1)
43

---

Table 6: Performances in various sets at different FeAT rounds.
COLOREDMNIST-O2
ROUND-1
ROUND-2
ROUND-3
TRAINING ACC.
85.084 0.14
71.87 0.96
84.934 1.26
RETENTION ACC.
88.11 4.28
43.824 0.59
OOD ACC.
11.08 0.30
70.64 0.62
10.07:: 0.26
COLOREUMNIST-01
ROUND-1
ROUND-2
ROUND-3
ROUND-4
TRAINING ACC.
88.63. 0.15
74.25 1.23
86.07 0.36
77.29 0.24
RETENTION AcC.
85,91 1.78
48.05 1.39
29.09: 1.15
73.50.
17.32. 2.69
12.48.1:
F.2
More details about WILDS experiments
In this section, we provide more details about the WILDS datasets used in the experiments as well as the evaluation setups.
F.2.1 Dataset description.
To evaluate feature learning performance given data from realistic scenarios, we select challeng- ing datasets from WILDS [39] benchmark. The datasets contain various realistic distribution shifts, ranging from domain distribution shifts, subpopulation shifts and the their mixed. summary the basic information and statistics of the selected WILDS datasets can be found in Table. 7, Table. 8, respectively. In the following, we will give a brief introduction to each of the datasets. More details can be found in the WILDS paper [39].
Table 7: A summary of datasets information from WILDS.
Dataset
Data (2))
Class information
Domains
Metrie
Architecture
AMAZON
Product reviews
Star ratings (5 classes)
7,676 reviewers
10-eth percentile acc.
DistiBBERT
CAMELYON17
slides
classes)
hospitals
Avg.
DenseNei-
COMMENTS
Online comments
Tooxiciyy (2 classes)
8 demographic groups
Wr. group acc.
DistilBEER
FMOW
Satellite images
Land (62 classes)
regions
Wr. group
DenseNet-121
(WILICAM
Photos
Animal species (186 classes)
324 locations
Macro
ResNer-50
RXRX1
Cell images
Genetic treatments (1,139 classes)
experimental butches
Avg.
ResNet 50
Table 8: A summary of datasets statistics from WILDS.
# Examples
Domains
Dataset
train
val
train
val
AMAZON
1,000.124
100,050
100,050
5,008
1,334
1,334
CAMELYON17
302.436
34,904
85,054
CIMILCOMENNTS
269,038
45,180
133,782
FMoW
76,863
19,915
22.108
11
3
IWILDCAM
129,809
14.961
42,791
243
32
48
RXRx1
40,612
9.854
34.432
33
Amazon. We follow the WILDS splits and data processing pipeline for the Amazon dataset [51]. provides 1.4 million comments collected from 7, 676 Amazon customers. The task is to predict the score (1-5 stars) for each review. The domains are defined according to the reviewer/customer who wrote the product reviews. The evaluation metric used for the task is 10th percentile of per-user accuracies in the OOD test sets, and the backbone model is DistilBert [66], following the WILDS protocol [39].
Camelyon17. We follow the WILDS splits and data processing pipeline for the Camelyon 1 7 dataset [6]. provides 450, 000 lymph-node scans from 5 hospitals. The task in Camelyon17 is to take the input of 96 x 96 medical images to predict whether there exists tumor tissue in the image. The domains d refers to index of the hospital where the image was taken. The training data are sampled from the first 3 hospitals where the OOD validation and test data are sampled from the 4-th and 5-th hospital, respectively. We will use the average accuracy as the evaluation metric and DenseNet- 121 [31] the backbone for the featurizer.
44

---

CivilComments. We follow the WILDS splits and data processing pipeline for the CivilComments dataset [9]. provides 450, 000 comments collected from online articles. The task is to classify whether an online comment text is toxic or non-toxic. The domains d are defined according to the demographic features, including male, female, I.BBTQ, Christian, Muslim, other religions, Black, and White. CivilComments is used to study the subpopulation shifts, here we will use the worst group/domain accuracy as the evaluation metric. As for the backbone of the featurizer, we will use a DistillBert [66] following WILDS [39].
FMoW. We follow the WILDS splits and data processing pipeline for the FMoW dataset [18]. provides satellite images from 16 years and 5 regions. The task in FMoW is to classify the images into 62 classes of building or land use categories. The domain split according the year that the satellite image was collected, as well as the regions in the image which could be Africa, America, Asia, Europe or Oceania. Distribution shifts could happen across different years and regions. The training data contains data collected before 2013, while the validation data contains images collected within 2013 to 2015, and the test data contains images collected after 2015. The evaluation metric for FMoW is the worst region accuracy and the backbone model for the featurizer is DenseNet-121 [31].
iWildCam. We follow the WILDS splits and data processing pipeline for the iWildCam dataset [8]. It is consist of 203, 029 heat or motion-activated photos of animal specifies from 323 different camera traps across different countries around the world. The task of iWildCam is to classify the corresponding animal specifies in the photos. The domains is split according to the locations of the camera traps which could introduce the distribution shifts. We will use the Macro Fl as the evaluation metric and ResNet-50 [27] as the backbone for the featurizer.
RxRx1. We follow the WILDS splits and data processing pipeline for the RxRx1 dataset [72]. The input is an image of cells taken by fluorescent microscopy. The cells can be genetically perturbed by siRNA and the task of RxRxI is to predict the class of the corresponding siRNA that have treated the cells. There exists 1, 139 genetic treatments and the domain shifts are introduced by the experimental batches. We will use the average accuracy of the OOD experimental batches the evaluation metric and 50 [27] as the backbone for the featurizer.
F.2.2 Training and evaluation details.
We follow previous works implement and evaluate different methods used in our experiments [39]. The information of the referred paper and code is listed as in Table. 9.
Table 9: The information of the referred paper and code.
Paper
Commit
Code
WILDS [39]
v2.0.0
attps / /s stanford. edu/
Fish [69]
https // htt.b con aon//ageeen/fiah
Bonsai [83]
httpe: // thub con/ hun//juu Tjulianyu/RRe
DFR [33, 38]
kttpa: //tth//g.ttttttt//o con/ www.///iiiioorrrree...eoao http://w/w/p///sooeeeee feature, feature_Learminng
The general hyperparemter setting inherit from the referred codes and papers, and are as listed in Table 10. We use the same backbone models to implement the featurizer [27, 66). By default, we repeat the experiments by 3 runs with the random seeds of 0, 1, 2. While for Camelyon17, we follow the official guide to repeat 10 times with the random seeds from 0 to 9.
Table 10: General hyperparameter settings for the experiments on WILDS.
Datasel
AMAZON
CAMELYON17
CIVICOMMENTS
FMoW
IWILDCAM
RXRXI
Num. soads
Learning rale
26-6
le.4
le-5
le-4
le-4
le-3
Weight decay
0.01
0
le-5
Scheduler
n/a
Cosine Warmup
Batch size
64
32
16
32
16 ReaNet50 Adam 10
72
Aschitecture
DistilBert
DeaseNet111 SGD
DisulBert Adam
DenseNet121 Adam
ResNer50
Optimizer
Adum
Adam
Domains minibatch
10
Group by
Countries 200
Hospitals
Demographics tuiicity
regions
Trap lncutions
Experimental batches
Training epochs
10
12
90
OOD objective implementations. We choose 4 representative OOD objectives to evaluate the quality of learned features, including GroupDRO [64], IRMv1 [4], VREx [41] and IRMX [16].
45

---

We implement the OOD objectives based on the code provided by Shi et al. [69]. For each OOD objective, by default, we follow the WILDS practice to sweep the penalty weights from the range of (le- 2, 1,-1,11 1, lel, 1e2]. and perform the model and hyperparameter selection via the performance in the provided OOD validation set of each dataset. Due to the overwhelming computational overhead required by large datasets and resource constraints, we tune the penalty weight in iWildCam according to the performance with seed O, which we empirically find yields similar results as full seed tunning. Besides in Amazon, we adopt the penalty weights tuned from CivilComments since the two datasets share a relatively high similarity, which we empirically find yields similar results as full seed tunning, too. On the other hand, it raises more challenges for feature learning algorithms in iWildCam and Amazon.
Deep Feature Reweighting (DFR) implementations. For the implementation of DFR [33, 38), use the code provided in Izmailov et al. [33]. By default, DFR considers the OOD validation as an unbiased dataset and adopts the OOD validation set to learn a new classifier based on the frozen features from the pre-trained featurizer. We follow the same implementation and evaluation protocol when evaluating feature learning quality in FMoW and CivilComments. However, since Camelyon17 1 does not have the desired OOD validation set, we follow the "cheating" protocol as in Rosenfeld et al. [63] to perform the logistic regression based the train and test sets. Note that when "cheating". the model is not able to access the whole test sets, Instead, the logistic regression is conducted on random split of the concatenated train and test data. Moreover, for Amazon and iWildCam, we find the original implementation fails to converge possibly due to the complexity of the task, and the relatively poor feature learning quality. Hence we implement a new logistic regression based on PyTorch [54] optimized with SGD, and perform DFR using "cheating" " protocol based on the OOD validation set and test set. Besides, we find neither the two aforementioned implementations or dataset choices can lead to DFR convergence in RxRx, which we will leave for future investigations.
Feature learning algorithm implementations. We implement all the feature learning methods based on the Fish code framework. For the fairness of comparison, we set all the methods to train the same number of steps or rounds (if applicable) in WILDS datasets. The only exception is in RxRx1, where both Bonsai and FeAT require more steps to converge, since the initialized featurizer has relatively large distance from the desired fcaturizer in the task. We did not train the model for much too long epochs as Izmailov al. [33] find that only requires 2 5 epochs for deep nets to learn high-quality invariant features. The final model is selected based on the OOD validation accuracy during the training. Besides, we tune the retain penalty in FeAT by searching over (le - 2, le 1, 0.5, 1. 2, 10), and finalize the retain penalty according to the OOD validation performance. We list the detailed training steps and rounds setups, as well as the used retain penalty in FeAT in Table 11.
Table 11: Hyperparameter setups of feature learning algorithms for the experiments on WILDS.
Datasel
AMAZON
CAMELYON17
CIVI.COMMENTS
l'MW
IWILDCAM
RxRx1
Overall steps
31,000
10,000
50.445
9,600
48,000
20,000
Approx. epochs
4
10
Num. of rounds
3
2
2
2
10
Steps per round
10,334
5,000
16,815
4,800
10
10
FeAT Retain penalty
2.0
1c-2
lc-2
1.0
0.5
For ERM, we train the model simply by the overall number of steps, except for RxRxI where we train the model by 15, 000 steps following previous setups [69]. Bonsai and FeAT directly adopt the setting listed in the Table 11. Besides, Bonsai will adopt one additional round for synthesizing the pre-trained models from different rounds. Although Zhang et al. [83] requires Bonsai to train the two rounds for synthesizing the learned features, we empirically find additional training steps in synthesizing will incur overfitting and worse performance. Moreover, as Bonsai requires propagating
2K - batches of the data that may exceed the memory limits, we use a smaller batch size when training Bonsai in iWildCam (8) and RxRx1 (56).
46

---

F.3 Software and hardware
We implement our methods with PyTorch [54]. For the software and hardware configurations, we ensure the consistent environments for each datasets, We run all the experiments on Linux servers with NVIDIA V100 graphics cards with CUDA 10.2.
F,4
Computational analysis
Compared to ERM,, the additional computational and memory overhead introduced FeAT mainly lie in the FeAT training and partitioning. At each training step. FeAT needs (k: (x-1) additional forward and backward propagation, the same as Bonsai, while FeAT only needs min(1, 1) additional
propagation. Besides, Bonsai additionally requires another round of training with (K 1) additional propagation, given K total rounds.
We calculated the computational overhead: The results aligned with our discussion. Bonsai requires
Table 12: Training and memory overhead of different algorithms.
Camelyon Training time
CivilComments Training time
Memory (%)
Memory (%)
ERM
56.21 29 mins
22.56:10.00
24.22:10.33 hrs
36.46:10.00
Bonsai
214.55 mins
51.75=+0.01
58.47++0.9 1 hrs
64.43++0.31
FeAT
101.14 12.79 mins
51.92=00.04
18.19-1.115 hrs
56.21=00.4
much more time for the additional synthetic round and much more memory when there 3 or more rounds. In contrast, FeAT achieves the best performance without introducing too much additional computational overhead.
F.5 Feature learning analysis
We first visualize the feature learning of ERM and FeAT on ColoredMNIST-005, as shown in Fig. be found that ERM can learn both invariant and spurious features predict the label, aligned with our theory.
However, ERM focuses
spurious features and even forgets certain features with longer training epochs, which could be due to multiple reasons such as the simplicity biases of ERM. Hence predictions based ERM learned features fail to generalize OOD examples. In contrast, FeAT effectively captures the meaningful features for all samples and generalizes to OOD examples well.
Table 13: Labels and predictions for the visualized samples.
Label
ERM
Bonsai
FeAT
Label
FRM
Bonsai
FeAT
1
0
113
68
0
113
0
113
0
0
113
CAMELYON17
IWILDCAM
0
36
36
36
36
0
0
36
36
36
36
40
40
40
40
1138
812
812
812
40
40
1138
1133
1125
1133
FMOW
RXRXI
40
2
29
29
35
43
1119
143
40
40
40
40
35
35
1054
35
toxic
toxic
toxic
toxic
toxic
toxic
toxic
toxic
COMMENTS
AMAZON
toxic
nontoxic
nontoxic
nontoxic
nontoxic
5
We also visualize the saliency maps of ERM, Bonsai, and FeAT on all real-world datasets used in our work with https: / thub. com/pytors ch/ captum. The visualizations are shown as in Fig. 5 to Fig. 11, for which the labels and the predictions of different algorithms are given Table. 13. It can be found that, across various tasks and data modalities, FeAT effectively learns more meaningful and diverse features than ERM and Bonsai, which serve as strong evidence for the consistent superiority of FeAT in OOD generalization.
47

---

4
4
keds
lite
PINLL Banc
PI23 HERE
PRE
aned Cleed
4
4
4
4
POOL
(a) ERM 150 epochs
(b) ERM 300 epochs
6
4
4
PISL REd
PRO 2
7
Hest bee
FRIO Intco
Ond EI
4
FRISL
3
3
8
D
4
4
4
4
3
(c) ERM epochs
(d) FeAT rounds
Figure 4: GradCAM visualization on COLOREDMNIST-025,, where the shortcuts are now concen- trated to colored path the up left. Three visualizations are drawn for each sample: the original figure, the gray-colored gradcam, and the gradcam. It can be found that ERM can not properly capture the desired features or even forget certain features with longer training epochs. FeAT can stably capture the desired features.
48

---

Lagend: Negative Nectmil
Poshbw
Method
Word Importance
(CLS) adalf hit thtar was actually left wing. the neot partty was very socialist "If you knaw anything about the history of nazl germary
persecutton jews, intelactuals communistis aocial democrats, ene you d realize how fidiosious your st statements ans. [SEP)
ERM
want, speech get ... rot (CLS) can anyone Slaid and disgastting
people
flying linaric
dring
(CLS) the cachollc shurch ts officially part of rodical adbol ddsh addv Adlk aeter udror Adlom communios organi lzens as catholic i! disgusted by the indo seat aarira Aition of whatis becoming cutt the popa and this mor #donn mo kitel Airryy inore mass genacide while prop ddping the
ERM
out of mohammed bent on the genocide of christ lins [SEP)
ERM
CCS) does the form cognitiro di kilsoon aslance apply to bally
[CLS]
wes actually left gin..
pairty was very socialist. you anything history of germery
Bonsal
persecution
intellctuels, communists
democrats, eco you'
how Bilouubur your statemenss are [SEP)
(CLS) no one can forca anyone to what they want, fres speech Evea get It, when not toll #dterttd is call to harm people, that ftts all and flying some lunazie raal flao es Buaidd and discusting os It is.. not daing harn.. ... [SEP]
[CLS) the cathalic church officially part of tibll #ew tiik #ftor Aelisnn communisc organizers as cathalic disguabed by aittion
Bonsal
the
bacoming the pape this
frew tirryy graae mass genocide while prop difpingg
mohammed bent the genoolde chriatians [SEP]
Bonaei
[CL
cognitive
"apply [SEP]
[CLS) adolf hilt kiler actually left
socialist.
anything
history of
PeAT
persacution jesai, telllchuusss communists. aocial democrats. end
realize how fidios dours your stafemenns are. [SEP)
(CLS) no one can force anyone to not say what they want, free soeech lives get over it, what not toll ascrrated a call to harm people that Ads all, and flying some lunazis real flag as sualdd and disgusting os It is not daing herm .. [SEP)
POAT
(CLS) the cathelic church
satool
organ cathalic
FeAT
this mehammed bent genoelde christians [SEP]
#ation
becoming
while ddping
FeAT
(CLS] does the torm cognitto di #esson #efance apply baley
DEP
(a) COVILCOMMENTS
Figure 5: Saliency map of feature learning on WILDS CIVIL COVILCOMMNNTS benchmark. The green- colored tokens are the learned features that contributed most to the target class, while the red-colored tokens contributed to the other classes. It can be found that FeAT is able to learn more meaningful and diverse features than ERM and Bonsai.
49

---

Method
Word Importance
[CLS] ameriiaan dream XV2 iids to be american lool with a political agends the bush ad aittninn iitira aittion which in itself goes all over actor pretty much peak well decade
take slmon than she may bit player actress bast this film takes so mush to be saxire but up with more susstions then any anowers. loo long & over ddbl Iolown william def lifoe was trying to look like dick chaney but anded up looking mora like honry fond nita., connis qu #talld was channal ftting george bush with mixed results. Just 6k nim [[SEP]
the place never provides entertainmant.
from tdoi rather just there. manely
nothing
ERM
(CLS) fros my ishone per serfectly, and already have inc dilipp Utio ailison sleeve i keep permanently it a little e udnu R5g the slerve, but easy enough to put and take out there are beht loope, , that
use, and one bett dlp sears sturey anough. when bought this, It was $ 2. 98 with no a & h. this cartainty lsn "I worth whole lot more than that but it s great value for the price I' dd actually be willing to pay up to $ 5..0 for it. It 's not too bullyy at least not more than holster should be ! the magnat keeps it closed It stays on my bot just hnn), and It' is sturoy enough to keepp the lohone sate for a short drop , espesially with the silioo deeve il is not lerriby fashionable, but you can 't beat the price. [SEP)
use the
ERM
(CL8) the mind of till dht cole must be quill an intereating place in her third book of the acamred souls series cole once again explores the decpest parts of human de titor inoov nty...t it' shock and cutrage you, It Ill make you teol so much more for har characters... ra ihivva irige is told mainly from the points of views of 1044 and iillyss with couple chaplere from lu liikas it is not a stand alone book, in my opinion this veries should be read in order . 'e been a fan of this series from the beginiig. the darknees of humnnn restrre Gil dilie cole continually portrays never caase IPs to angage me. her storles come from a plase of raw honesty and empathy. her writing mature tt's with each book. the story of 194 and Alyyya and their roles this tangled web of world cole has created was consuming as the others. however, for something was in their relationsnip. the romance missed the connection folt with the previcus books main characters.. while day deep into their pret and pulled out their flaws cherissrry and need for each other just utimato story and goals of of these characters... and In that aspect, story really excelled and has mo Just as Intrigued in how It all works out as ever was and il l'very much look forward to the nest ineealmment the scarred soule veries (rr ilivv liige left much open to be explorad and m confident in cala 's story tailing to take me owen deeper into her Intorasting fitly de léppr thtavv into worls. [SSP]
ERM
ittii ittoiitiie
individusis aspect of their
outte dllck this series conslstently moves forward though adaing more death and layers the
[CLS] just keops getting bettor and better you like alpha man, , strong hercine IP' and laughter with eash of sericusness, , son miss
ERM
this series! [SEP)
(a) AMAZON-ERM
(CLS) amarican dream #92 try tifs to be american lool with a political aganda on the bush ad nitinin nttrr itttion which in itsolf goos all over the place & never provides entertainment. hugh grant is actor who pretty much hit his peak well over decade ago & his take as simon from idoi rather just there. manily moore is eki nothing to mora than she may be a bit player as actross at bast this hilm takes on so much to be sotine, but ands up with more questionn then answers. too long & over dabi istown willisan def thhnn was trying to look
Bonsal
like dick chaney but ended looking like henry fond antaa, cennis qu #tlldd channel Attin george bush with mixed results, Just
ak film [[SEP]
[CLS] this tits my ionone perfectly and already have an Inc #tl #ttl silicon steeve that I knep it in parmanently It 's little a #dnu #iig and take out there batt loops, that I use. and bett cito that seems sturay enough. when bought It was
with the cleeve, but easy enough
the belt h- this certainty Isn worth whole lot than but for not bully.. at not more than helster should be sturay enough to keop the lshono safe short drop. , espectally with
Bonssi
great value the price. i' 'd actually be willing to pay up the magnat keeps closed, It stays my boit just tinal, the silicon steave. It" not terribly fashlonaole, but you can boat the price [SEP]
(CLS) mind deepest parts human de Wolr atlyy it shook outrage you, make you feel mush more for her characters. ra #nge tole mainty from the points of views of 194 and ihtoya with couple chapters from trka. It is a stand alone book, my opinion this series should be read in order.. .I 've been a fan of this series from the beginning. - the darieeees of human nature till ddlle cole continually portrays never cease léss to engage me. . her storles come from a place of row honesty and empathy. her writing mature hits with each book, the story of 194 and 2 itdlya and their roles in this tangled web of a world cole has created was as consuming as the others. however, for mo. something was off in thair relationentp. In the romance... I missed the connection I falt with the provious booics main charactors. while cole dug deep into thair past and pulled out thair flaws and #din aitorra #tolltiies so wall as Individuais, an aspect of their cherietry and need for each other just didn' I cuite dlick, this series consistently movee forward though adbing utimate story goals of these characters. and that aspect, the story really excelled and works
musi quite interesting place her third
scarred veries cole once again explores the
Bonsai
depth and layers the just Intrigued in how all very much look forward the next Instalment the seared souls saries, ihiva trige left much open be explored and I' m confident cole 's story tellin to take me even deeper into Interesting Ally udpr Stave ueo worle [SEP)
(CLS) just keeps getting better and better Ii like alpha , strong heroine NUs and laughter with eash of sericusness. miess this serias) [[SFF]
Bonsal
(b) AMAzON-Bonsa
Figure 6: Saliency map of feature learning on WILDS AMAZON benchmark. The green-colored tokens are the learned features that contributed most to the target class, while the red-colored tokens contributed to the other classes. It can be found that FeAT is able to learn more meaningful and diverse features than ERM and Bonsai.
50

---

Method
Importance
[CE, S] amarican dream 192 try Als to be american idei with political agenda on the bush ad ittmmin tetra tettion which itseff goes all over the place & never provides entertainment. hugh grant a actor who pretty much hit his peak well over decado ago & his take as simon from lool is rather just there. mandy moore ok nothing to more than she may be a bit player
actross at best this tim takes so much salire, but it ends with more questions then any anaserss foo long over pdbil Killon (willia def affoee was trying to look like dick cheney andad lookdng lkk henryy fond nito.. connts #ttaid channal ditingg george bush mixed rosults. just hil.. [SEP]
FeAT
[CL 6] this fits iphone per facilyy and already have inc tWip Attio silicon sieeve that keep permarently it ittle artnu difg with the sleeve, but enough put in take there are bell looge, that don 'I
one belt clip that do the bell cip that but it great should drop, especially
sturdy enough. whenl grice. actually willing to pay up magnet closed. stays on my belt just fine, and sturdy enough keep (phone safe silooo slease a bemblly fashionsble, but you can beat the prioe. [SEP]
with
cortainly lan worth whole
FeAT
not bulky, , at lsast not more
[CL s] the mind ttie cala be
again exploras interesting place. In thire book the scorrad series characters.. fllve Aflige outrage you
#dpr WVad kuty. It and #tay cnuple chapters from dtka It is a stand alone book, my opinion this series should be read orber i!" ve been fan series from the beginning. the darerwes of human nature ditie cole continuslly
deepest parts
you so
told mainly from the points of views
portrays nover coase htts to engage mo. her starles como from placo of raw honosty and ampathy. her writing mature itts with each book. the story of 194 and z foys and their roles in this langled web of a workd cole has created was as consuming as the others. however, for me. something was off in their relationship.. in the rorancce.. i missed the connection i felt with the previous books main characters. while cole oug into their past and pulled out their flavs and vu titn attora #ubiltiie so well incividuats inn aspect of their chemlstry and need for each just cich? quite click, this series consistendy these characters
FeAT
ultimate and gools all ffigge left open explored confidint in cole story telling to take deopar into her interesting #tly enpr staaee world.. [SaP]
forward though adding depth
aspect, the story really exselled and has instaliment scerred souls seriie.
intrigued
works
much look
[CLS] just
better and better!
heroine and laughter with dash soricusness, don 't mise this
FeAT
(SPP)
(a) AMAZON-FeAl
Figure 7: Saliency map of feature learning on WILDS AMAZON benchmark (part 2). The green- colored tokens are the learned features that contributed most to the target class, while the red-colored tokens contributed the other classes. It can be found that FeAT is able learn more meaningful and diverse features than ERM and Bonsai.
51

---

(a) CAMELYON17-ERG
(b) CAMELYON17 1---ON1171BBBa
(c) JAMELYON17-FeA.
(d) CAMELYON 17ERRM
(e) CAMELYON 1 77 7-Bonsai
(1) CAMELYON 77-FeAT
(g) CAMELYON 17-ERM
(b) CAMELYON 7-Bonsai
(i) CAMELYON I 77FeAT
(1) CAMBLYONI7-FFAT Figure 8: Saliency map of feature learning on CAMELYON 17 benchmark. The blue dots are the salient features. deeper blue color denotes more salient features. can be found that FeAT is able to learn more meaningful and diverse features than ERM and Bonsai.
(i) CAMELYON 1- 7-ERM
(k) CAMELYON B-Bonsai
52

---

(a) FMOW-ERS -EOW
(b) IMoW-Bonsa
(c) FMOW-FeAT
(d) FMOW-ERM
(e) FMoW-Booe -BBnsaii
(1) FMOW-FeAT
(g) FMOW-ERM
(h) FMoW-Bonsse
ii) FMOW-FeAT
()) FMOW-FeAR Figure 9: Saliency map of feature learning on FMoW benchmark. The blue dots are the salient features. A deeper blue color denotes more salient features. can be found that FeAT is able to lear more meaningful and diverse features than ERM and Bonsai.
(i) -MOW-ERM
(k) FMoW-Bonsai
53

---

(a) EWILDCAM-ERM
(b) IWILDCAM-BOULH
(c) IWILDCAM-FeAT
(d) WWILDCAM-ERM
(e) WWLLDCAM-Bonsa
(I) IWILDCAM-FeAT
(g) IWILDCAM-ERM
(h) WILLDCAM-BOnsai
(i) IWILDCAM-FEAT
(i) IWILDCAM-ERN
(k) IWILDCAM-Bon -onsai
(l) IWILDCAM-FAAT
Figure 10: Saliency map of feature learning IWILDCAM benchmark. The blue dots are the salient features. A deeper blue color denotes more salient features. can be found that FeAT is able to learn more meaningful and diverse features than ERM and Bonsai.
54

---

(a) R --EMM
(b) RxRxI-Bonse
(c) RXRXI-FEAT
(d) RXRXI-ERM
(e) RxRxI-Boonii -Boonaii
(I) RXRXI-FeAT
(g) RXRXI-ERM
(h)
ii) RXRxI-FeAT
(1) RXRXI-FFAT Figure 11: Saliency map of feature learning RxRxI benchmark. The blue dots are the salient features. A deeper blue color denotes more salient features. can be found that FeAT is able to learn more meaningful and diverse features than ERM and Bonsai.
( RXRXI-ERM
(k) XERx1-Bonsse
55


Reviews:
[{"review_id": "zio6jJbd6Q", "reviewer": "Reviewer_KJZ8", "summary": "This paper explores the relationshp between the ERM training and OOD generalization in feature learning. The authors analyze the corresponding learned features by ERM and OOD objectives. To answer the question, they conduct the investigation of feature learning in a two-layer CNN network training with ERM and IRMv1. They adopt the data models proposed by [2, 11], and include features with different correlation degrees to the labels to simulate invariant and spurious features like [26]. The theoretocal results extend  [51] from data augmentation to ERM learning process. They find that ERM fails since it learns the spurious features more quickly than invariant features, when spurious correlations are stronger than invariant correlations. However,  invariant feature learning also happens with RRM so long as the invariant feature has a non-trivial correlation strength with the labels. Moreovoer, they find that IRMv1 requires sufficiently well-learned features for OOD generalization. Compared with the former workshop version, this conference submission adds the so-called  Feature Augmented Training (FAT) to learn features for OOD generalization. The proposed method iteratively augments the model to learn new features while retaining the already learned features. In each round, the retention and augmentation operations are performed  distributionally robust optimization on different subsets of the training data that capture distinct features. The experimental results verify the promising OOD generalization performance of the proposed method. ", "questions": "1. It is better to claim Lma 3.2 in the supplementary material. \n\n2. The details of Fig.2 are imoportant. It should be presented in the main body.", "limitations": "Yes.", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. Interesting theoretocal findings. \n\n2. Promosing empirical results.\n\n", "weaknesses": "1. The symbol system seems hard to follow. Some superscripts and subscripts are redendunt: L_S in Eq.(2), \\ell‘ in Eq.(5), o_d in Thm 4.2. Please double-check these symbols and make them clear.\n\n2. Some symbols are not well-defined, such as the \\cdot in the predictor f, little o, big O, \\Omega, \\Theta. \n\n3. The activation function and the variance share the same symbol. \n\n4. Thm 4.1 (informal) does not provide the quantitative results to describe how much the incerment of spurious feature is larger than that of the invariant feature at any iteration. Althrough this is a informal statement, the quantitative results is important. \n\n5. It is better to discuss the computational complexity of the proposed method and compare with the other competitors.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time and your positive feedback! Please see our detailed responses to your comments and suggestions below.\n\n> W1 Some superscripts and subscripts are redendunt: L_S in Eq.(2), \\ell‘ in Eq.(5), o_d in Thm 4.2. Please double-check these symbols and make them clear. \n\nThanks for your suggestions. We have revised our draft accordingly:\n- We changed $L_S$ to $L$ in Eq.(2); \n- We changed $\\ell_i’$ to $\\ell\\_i’^e$ in Eq. (5) referring to the first order derivative of $L_e$ with respect to the $i$-th sample at $e$-th environment.  \n- By using $o_d$, we mean that the value is small with respect to the dimension $d$. This means that the requirements of the value don't grow with $d$, at least asymptotically. \n\nWe also double checked the other notations and simplified the superscripts and subscripts when unnecessary. Besides, we provided a table of key notations at the Appendix for readers easier to follow our notations.\n\n> W2 Some symbols are not well-defined, such as the \\cdot in the predictor f, little o, big O, \\Omega, \\Theta. \n\nWe have complemented our full introduction of notations in the Appendix A and provide pointers in the main text. Below is the specific definition for the notations mentioned in the review:\n- For $\\cdot$ in the predictor $f$, we refer to the product with a vector or a scalar. The use of $\\cdot$ in $f$ (Eq. (4)) is because of the relaxation of IRMv1 that relaxes $w$ to be a scalar.\n- The other notations are used to compare two sequences $\\{ a_n \\}$ and $\\{b_n \\}$. Specifically, we employ standard asymptotic notations such as $O(\\cdot)$, $o(\\cdot)$, $\\Omega(\\cdot)$, and $\\Theta(\\cdot)$ to describe their limiting behavior. Specifically, we write $a_n =O(b_n) $ if there exists a positive real number $C_1$ and a positive integer $N$ such that $|a_n| \\le C_1 |b_n| $ for all $n \\ge N$. Similarly, we write $a_n = \\Omega (b_n)$ if there exists $C_2 > 0$ and $N >0$ such that $|a_n| > C_2 |b_n |$ for all $n \\ge N$. We say $a_n = \\Theta(b_n)$ if $a_n = O(b_n)$ and $a_n = \\Omega(b_n)$. Besides, if $\\lim_{n \\rightarrow \\infty} |a_n/b_n| =0 $, we express this as $a_n = o(b_n)$. We use $\\widetilde{O}(\\cdot)$, $\\widetilde{\\Omega}(\\cdot)$, and $\\widetilde{\\Theta}(\\cdot)$ to hide logarithmic factors in these notations respectively. Moreover, we denote $a_n = \\textrm{poly} (b_n)$ if $a_n = O((b_n)^p)$ for some positive constant $p$ and $a_n = \\textrm{polylog}(b_n)$ if $a_n = \\textrm{poly}( \\log(b_n))$. \n\n> W3 The activation function and the variance share the same symbol. \n\nWe have revised our draft and changed the activation function symbol to $\\psi$.\n\n> W4 Quantitative results in Thm 4.1 (informal).\n\nThe quantitative results of how much the increment of the spurious feature is larger than that of the invariant feature can be found at Eq. (15) in Appendix C.2. This difference is primarily determined by the empirical distributions of $\\textup{Rad}(\\alpha)$ and $\\textup{Rad}(\\beta_e)$, reflected in the quantities $\\overline{C}_{j\\ell} = \\sum_e {\\frac{1}{n_e} |\\lbrace i\\mid \\textup{Rad}(\\alpha)_i = j, \\textup{Rad}(\\beta_e)_i = \\ell, i\\in \\mathcal{E}_e\\rbrace| }$  for $j\\in\\lbrace \\pm 1\\rbrace,\\ell\\in \\lbrace\\pm 1\\rbrace$. This (positive) difference decreases monotonically from Eq. (17) towards 0. We have added more quantitative results in Thm 4.1.  \n\n> W5 Computatial complexity.\n\nWe have revised the draft to include a discussion of computational complexity in Sec. 5.2: \n\nCompared to ERM, the additional computational and memory overhead introduced in FAT mainly lie in the FAT training and partitioning. At each training step, FAT needs $(k-1)$ additional forward and backward propagation, the same as Bonsai, while FAT only needs $\\min(1,k-1)$ additional propagation. Besides, Bonsai additionally require another round of training with $(K-1)$ additional propagation, given $K$ total rounds.\n\nWe calculated the computational overhead:\n|        | Camelyon17            |                | CivilComments      |                |\n|--------|-----------------------|----------------|--------------------|----------------|\n|        | Training time         | Memory (%)     | Training time      | Memory (%)     |\n| ERM    | 56.21$\\pm$8.29 mins   | 22.56$\\pm$0.00 | 24.22$\\pm$0.33 hrs | 36.46$\\pm$0.00 |\n| Bonsai | 214.55$\\pm$1.13 mins  | 51.75$\\pm$0.01 | 58.47$\\pm$0.91 hrs | 64.43$\\pm$0.31 |\n| FAT    | 101.14$\\pm$12.79 mins | 51.92$\\pm$0.04 | 28.19$\\pm$1.15 hrs | 56.21$\\pm$0.48 |\n\nThe results aligned with our discussion. Bonsai requires much more time for the additional synthetic round and much more memory when there are 3 or more rounds. In contrast, FAT achieves the best  OOD performance without introducing additional much too computational overhead.\n\n\n> Q1 Claim Lma 3.2 to supplementary part.\n\nLma 3.2 introduces key concepts of feature learning in neural networks by decomposing weights into two parts. While we agree moving Lma 3.2 will improve clarity, it could make readers who are not familiar with feature learning literature hard to follow our results (e.g., Reviewer Jcjd). \n\n> Q2 The details of Fig.2 are important.\n\nWe included the details in our updated draft: the invariant and spurious feature learning terms plotted in Fig. 2 are the mean of $\\langle \\mathbf{w}\\_{j,r}, j\\mathbf{v}\\_1 \\rangle$ and $\\langle \\mathbf{w}\\_{j,r}, j\\mathbf{v}\\_2 \\rangle$ for $j\\in \\lbrace\\pm 1\\rbrace, r\\in[m]$, respectively.\n\nWe’d appreciate it if you could take the above responses into consideration when making the final evaluation of our work. Please let us know if there are any outstanding questions.\n\n"}, {"review_id": "8eZsHqe4mt", "reviewer": "Reviewer_NB5i", "summary": "This work aims to understand and compare feature learning in ERM and certain OOD generalization objectives. Additionally, it proposes an approach to enhance feature learning for improved OOD generalization.\n\nFirst, the authors examine data consisting of invariant and spurious features. They theoretically show that when training a two-layer CNN with ERM, it learns both the features. However, when the spurious correlation in stronger, the spurious features are learned faster. Additionally, they explore the IRMv1 objective and illustrate that fine-tuning the ERM-trained model using the IRMv1 objective does not result in learning of new features.\n\nTo improve feature learning compared to ERM, the authors propose a technique called feature augmented training (FAT). At each training step, they partition the data into two sets based on the correctness of predictions. The set with accurate predictions is added to an augmentation set, while the set with incorrect predictions is added to a retention set. Both sets expand with each training round. The training objective involves applying Distributionally Robust Optimization (DRO) on the augmentation set to learn new features, combined with ERM on the retention set to retain the learned features.\n\nThe authors demonstrate that using OOD objectives on models trained with FAT leads to improvement in OOD performance on two variants of the colored MNIST dataset and six datasets from the WILDS benchmark.", "questions": "1. Improvement in feature learning compared to ERM:\n\n- In Section 6, it is mentioned that FAT, Bonsai [2] and ERM are generally trained for the same number of overall epochs. However, given that ERM is computationally less expensive than the other two methods, it would be valuable to investigate whether training ERM for a longer duration can further enhance the learned feature representations, or if the representations stop improving after a certain number of training epochs. This analysis would provide insight into the extent to which FAT and Bonsai truly improve feature learning, especially considering that they do not consistently lead to significant improvements in OOD performance in many cases (as observed in Table 2).\n\n- To better understand and compare feature learning in ERM and FAT, it would be beneficial to include additional analyses, such as visualizing the saliency maps using GradCAM.\n\n2. Clarity:\n\n- In the caption for Fig. 1, the description of FAT is unclear and it does not adequately explain the illustration of FAT in Fig. 1(a). Although the experimental results in Figure 1(b) provide informative insights, it would be more comprehensive to include details about the dataset and the number of training epochs used for FAT.\n\n- The description of the method in the introduction (lines 55-58) is unclear and difficult to comprehend without referencing Algorithm 1.\n\n- Concerning the algorithm, the rationale behind Line 16 (returning the average of all linear classifiers) is unclear and could benefit from further explanation.\n\n3. Minor:\n\n- Please address any typos, e.g., in the Fig. 1 caption, and inconsistencies in notation, e.g., in Equation (5), Lemma 3.2, etc.\n\n- It may be worth considering renaming the method from FAT to FeAT, as it better reflects the focus of the approach on improved feature learning.\n\nReferences:\n\n[1] P. Kirichenko, P. Izmailov, and A. G. Wilson. Last layer re-training is sufficient for robustness to spurious correlations. arXiv preprint arXiv:2204.02937, 2022.\n\n[2] J. Zhang, D. Lopez-Paz, and L. Bottou. Rich feature construction for the optimization-generalization dilemma. In International Conference on Machine Learning, pages 26397–26411, 529 2022.", "limitations": "\n[-] Reproducibility: The authors have not provided the code for implementing their approach. While the Appendix contains most of the necessary implementation details, sharing the code would greatly facilitate the verification of their findings.\n\n[+] The authors acknowledge the limitation that feature learning is influenced by various factors, and this work specifically aims to investigate feature learning in the context of spurious and invariant features.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. This work aims to improve feature learning for better OOD generalization, which has been identified as a limitation of existing OOD objectives in recent research [1]. While the authors acknowledge that the idea to do this is not new [2], the proposed approach FAT is novel and demonstrates effectiveness in enhancing OOD performance. FAT incorporates a termination step to halt further training once the model has acquired sufficiently rich features, and the authors also propose an efficient version, iFAT, to make the approach practical for implementation. \n\n2. The theoretical results presented in this work provide intriguing insights into feature learning with ERM and the observation [1] that OOD generalization objectives do not improve feature learning.\n\n3. In their comparisons, the authors extensively evaluate ERM and Bonsai [2]—the primary contenders for feature learning—as well as various OOD objectives applied on top of these methods, across multiple datasets (as shown in Table 2). The results demonstrate that FAT enhances OOD performance.", "weaknesses": "1. The claim that ERM does not learn sufficiently good feature representations, while FAT improves them, requires stronger support and evidence.\n\n2. Some sections of the writing require clarifications to improve understanding and readability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your support and constructive comments! Please see our detailed responses to your comments and suggestions below where we use reference numbers in our manuscript due to the character limit:\n\n> Q1.1 The effects of longer ERM training epoch.\n\nThank you for the insightful question. **In fact, the representations of ERM will stop improving after a certain number of ERM training epochs, shown in Fig. 10 of [25], or even decrease when being fed to OOD training for some datasets (Fig.1b)**, which highlights a drawback of ERM feature learning. Note that for deep models, the feature learning will quickly saturate within one epoch [25], within the training epochs we use (in Table 10). Therefore, the consistent improvements by FAT are non-trivial.\n\n> Q1.2 GradCAM\n\nWe followed your suggestion and plotted the saliency maps to better understand the feature learning of different algorithms. The figures are attached to the rebuttal pdf and to our latest draft as well.\n\nWe first visualize the feature learning of ERM and FAT on ColoredMNIST-025. It can be found that \n- ERM can learn both invariant and spurious features to predict the label, aligned with our theory. \n- However, ERM focuses more on spurious features and even **forgets** certain features with longer training epochs, which could be due to multiple reasons such as the simplicity biases of ERM[a]. Hence predictions based on ERM learned features fail to generalize to OOD examples. \n- In contrast, FAT effectively captures the meaningful features for all samples and generalizes to OOD examples well.\n\nWe also visualize the saliency maps of ERM, Bonsai and FAT on all real-world datasets used in our work. It can be found that, across various tasks and data modalities, FAT effectively learns more meaningful and diverse features than ERM and Bonsai, which serve as strong evidence for the consistent superiority of FAT in OOD generalization.\n\n> Q2.1 Caption of Fig.1,\n\nWe included more details when introducing FAT:\n\nFAT iteratively checks and divides $D_{tr}$ into augmentation $D^a$ and retention sets $D^r$ that contain features not learned and already learned by the current model at the round, respectively. Then FAT augments the model with new features while retaining the already learned features, which leads to richer features for OOD training and better OOD performance.\n\nWe also included the dataset and epochs/rounds used in FAT when introducing the experiments of Fig.1.\n\n> Q2.2 Introduction\n\nWe revised the introduction with a more intuitive explanation:\n\nIn each round, FAT separates the train set into two subsets according to whether the underlying features in each set are already learned (Retention set $D^r$) or not (Augmentation set $D^a$), by examining whether the model yields correct ($D^r$) or incorrect ($D^a$) predictions for samples from the subsets, respectively. Intuitively, $D^a$ and $D^r$ will contain distinct features that are separated in different rounds. Then, FAT performs distributionally robust optimization (DRO) on all subsets, which *augments* the model to learn new features by minimizing ERM losses on all $D^a$ and *retains* the already learned features by retaining ERM losses on all $D^r$.\n\n> Q2.3 Explanation of the algorithm return\n\nWe revised the paper to add the explanation. The rationale is that the average of historical classifiers could be a good initial point as they already capitalize all the learned features in each round.\n\n> Q3.1 Typos\n\nWe have double checked and corrected all the typos and inconsistencies.\n\n> Q3.2 A new name\n\nThanks for this wonderful suggestion! We switched the name to FeAT in the revised draft! Nevertheless, to avoid confusion to other reviewers, we will use FAT for consistency during rebuttal.\n\n> L1 Reproducibilty\n\nAs claimed “ We will provide a link to an anonymous repository of our code during the discussion phase.” during submission, we have provided the code in an anonymized link to AC, following the rebuttal guideline that we are not allowed to disclosure any external link in the rebuttal.\n\n**References**\n\n[a] The pitfalls of simplicity bias in neural networks, NeurIPS'20.\n\n\nPlease let us know if you have any further questions. We’d be grateful if you could take the above responses into consideration when making the final evaluation of our work. \n"}, {"review_id": "rjCsmnTRKG", "reviewer": "Reviewer_Jcjd", "summary": "The paper consists of two parts. The first part provides a theoretical analysis of the training dynamics for a simple model and data distribution under ERM and IRM. Specifically, the authors explore the questions of feature learning, when one of the features changes its correlation to the target between environments (spurious feature) and the other does not (core feature). The second part proposes a training method for extracting reach features, and shows promising results across an array of benchmarks.", "questions": "**Q1.** In Eq. 5, why is the second term divided by $n_e^2$ and not $n_e$?\n\n**Q2.** You mention in line 279 that FAT comes with additional memory requirements, especially if the feature extractor has many parameters. Why is that? What exactly do you need to store?\n\n**Q3.** How exactly do you run DFR on CivilComments and Camelyon17? The results, even for ERM+DFR, appear to be surprisingly good, and in particular better than what's currently reported in the [WILDS leaderboard](https://wilds.stanford.edu/leaderboard/)?", "limitations": "The limitations are adequately addressed in my opinion, except for the issues raised in the weakness section.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "**S1.** The paper provides an extensive theoretical analysis of gradient descent dynamics under ERM and IRMv1 for a specific model and data distribution.\n\n**S2.** The paper proposes a novel training method (FAT) which shows strong performance in terms of feature learning", "weaknesses": "### W1: Theory\n\nThe paper places a lot of emphasis on the theoretical analysis, which occupies pages 3-6. The authors use the theoretical results to provide intuition about feature learning in neural networks, and training dynamics of ERM and IRM. \n\nUnfortunately, the model used is extremely simple. While the authors argue that it is a convolutional neural network, it is in fact **a linear model**. Indeed, the authors define the model as\n$f(W, x) = F_{+1}(W_{+1}, x) - F_{-1}(W_{-1}, x)$, where \n$F_j(W_j, x) = \\frac 1 m \\sum_{r=1}^m [w_{j, r}^T x_1 + w_{j, r}^T x_2]$, where I omitted the activation $\\sigma$, which the authors set to be the identity mapping. \n\nWe can then rewrite the model as a linear model\n$f(W, x) = \\frac 1 m \\sum_{r=1}^m [w_{+1, r} - w_{-1, r}]^T (x_1 + x_2) = \\tilde w^T \\tilde x$, where $\\tilde w = \\frac 1 m \\sum_{r=1}^m [w_{+1, r} - w_{-1, r}] $ is the effective weight vector and $\\tilde x = x_1 + x_2$ is the effective feature vector.\nIn other words, the \"CNN\" model used by the authors is simply a reparameterized linear model.\nMoreover, the reparameterization does not really change the gradient dynamics in a non-trivial way, as \n$\\frac {\\partial L}{\\partial w_{j, r}} = \\frac {\\partial L}{\\partial \\tilde w} \\frac {\\partial \\tilde w} {\\partial w_{j, r}} = \\frac {j}{m} \\nabla_{\\tilde w} L$. \nIn other words, all of the weights $w_{j, r}$ are updated with $\\pm$ the gradient of the linear model divided by $m$.\nSo **understanding the training dynamics in the \"CNN\" model is equivalent to understanding the training dynamics of a linear model**.\n\nTo summarize, the authors study the training dynamics of linear model with a logistic loss.\nThe authors do not clarify this connection in the paper, which in my opinion is misleading. This omission also unnecessarily complicates the presentation.\n \nI also believe there are several important implications of the fact that the authors analyze the training dynamics of a linear model:\n- It is not clear what is even meant by feature learning. Specifically, the authors refer to how much weight the model assigns to each of the input features. However, this is more similar to last layer / classifier training in the context of [1] that the authors reference, rather than feature learning.\n- Consequently, it's unclear what conclusions can even be made from the experiments about _feature learning_ in _neural networks_. The connection here seems far-fetched.\n- The setup for the IRM is quite strange. Specifically, the feature that the model outputs in this context is the logit (a number) predicted by the model, and the classifier is just fixed to be $1$. As a result, the authors get a logistic regression model with some additional gradient penalty. I am not an expert on IRM, but it is not clear how relevant this model is to _feature learning_ in neural networks with IRM.\n\nPlease correct me if I am wrong in the reasoning above!\n\n### W2: Theory $\\leftrightarrow$ Methodology\n\nThe connection between the theory and the proposed method (FAT) is not clear to me.\nIt seems like the main conclusion from the theory comes down to the idea that we need to learn diverse features, which ERM by itself might not do.\nHowever, as I mentioned above, the relevance of the theory to feature learning in neural networks is questionable.\n\nIn particular, the theory does not connect as far as I can tell to any of the details of the method.\nThe method could very easily be presented without the theory.\n\n### W3: Presentation\n\nBecause so much emphasis is placed on the theoretical results (which are in my opinion of limited relevance), the authors have to describe the method and the experiments in a limited space. \nThe presentation of the method is not very clear.\nIn particular, many of the design decisions are not explained.\nFor example, why do we need to reinitialize the weights of the classifier in each \"round\"? \nThe datasets $\\mathcal{D_i^a}$ are never defined in the text, but used in line 262. \nOverall, it is quite hard to follow the description of the algorithm, and the intuition behind it.\n\nMoreover, the iFAT method which is actually used in all of the experiments is not even described in the paper.\n\nI would recommend to deemphasize the theory, and use most of the space in the paper to clearly describe the method, present detailed experimental results and ablations on the various design choices.\n\n### W4: Performance\n\nOverall, FAT seems to consistently provide good results in Table 2. However, it's worth noting that the improvements over ERM appear to be fairly small (<1%) except for Camelyon17.\n\n**References**\n\n[1] [_Last Layer Re-Training is Sufficient for Robustness to Spurious Correlations_](https://openreview.net/forum?id=Zb6c8A-Fghk);\nP. Kirichenko, P. Izmailov, A. G. Wilson;\nICLR 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time in reviewing our work. Please see our detailed responses to your comments and suggestions below where we use references in our draft due to the token limit.\n\n> W1.1 The linear activation in the CNN model\n\nWe respectfully disagree with the point:\n- First, we need to clarify that, as the first work analyzing the ERM and IRMv1 feature learning under distribution shifts, the results with IRM regularization (which is **non-linear**) we obtained are not obvious under **non-convexity** [a] despite the linear activation. In fact, The assumption of linear activation is widely used and standard in the literature of theory analysis in OOD generalization [4,32,40,61,64]. Linear CNN is also widely studied by the community [b,c,d]. \n\n- Besides, **we also show that our key theoretical results could be generalized to the non-linear setting**. Please find the details in [our response to Reviewer qcQH](https://openreview.net/forum?id=eozEoAtjG8&noteId=KmRvdQIZ2C) due to the character limits.\n\n- We adopt the linearity activation function primarily for the sake of simplicity and clarity when studying IRMv1. It's worth noting that the IRMv1 involves high-order derivatives, which can be a significant challenge to learning theory. \n\n- Going beyond the linearity of activations, the experiments in Fig. 1b show that, our theoretical results align with the empirical discoveries in more complex settings.\n\n> W1.2 The meaning of feature learning\n\nWe need to clarify that feature learning in neural networks refers to how the weights of a neural network **evolve to extract different features (Eq.6)**, especially when trained from scratch [2,11,51,58]. **The objectives are not limited to ERM**.\n\nThe last layer training method is applied to the **fixed trained features** and is limited to ERM, which serves as an indirect measure of feature learning for deep and complex networks where explicit quantities in Eq. 6 are not accessible.\n\nAnother key difference is that our analysis studies training a neural network based on **dataset containing spurious correlations (aligned with the setting in OOD generalization literature)**, while last layer training requires an **unbiased dataset** (i.e., without spurious correlations) to (indirectly) examine the invariant feature learning.\n\n> W1.3 The setup for IRM\n\nThe use of scalar in IRMv1 is because of the complicated formulation of the original IRM framework which involves a bilevel optimization (more details can be found in [13]). The gradient penalty in IRMv1 is **a practical variant to regularize** the feature learning in deep networks to focus on invariant features, and has gained lots of success[4]. \n\nTo avoid misunderstanding, we removed the sentence  “defined classifier w as the scalar 1” in the draft. We understand that Reviewer Jcjd may not be familiar with the literature of OOD generalization and IRM, nevertheless, we are happy to provide more details for any specific points unclear to Reviewer Jcjd.\n\n> W2 Theory and method.\n\nAs clarified, our theory precisely analyzes the feature learning of a CNN when optimized via ERM and IRMv1, and characterizes the functionalities of ERM and OOD objectives. Specifically, the second part of our theory implies that IRMv1 solely can not learn features but requires high quality feature representations for OOD generalization, hence the pre-training stage needs to learn high quality features. As pre-training with ERM may not learn all useful features for OOD generalization, our algorithm is thus motivated to strengthen the pre-training stage for better OOD generalization.\n\n**Without our theory, it’s unclear which stage needs to be improved and what objectives need to be used for improving OOD generalization.**\n\n> W3 Presentation\n\nAs shown in Eq. 7, we need to initialize (instead of “reinitialize”) a new classifier in each round to learn new features from the augmentation set while keeping the historical classifiers to retain the features already learned in each previous round.\n\nWe complemented the introduction of $D^a_i, D^r_i$ in line 262: $D_i^a$ and $D_i^r$ are the corresponding augmentation and retention set elicited at $i$-th round.\n\nWe also gave specific details in the revised paper, about the differences between iFAT and FAT, that iFAT stores only $D_{k-1}^a, D_{k-1}^r, w_{k-1}$ at the $k$-th round.\n\n> W4 Performance\n\nWe need to clarify that a large improvement in these real-world datasets is extremely more challenging than other ML tasks, as one may see in the Wilds leaderboard. The consistent improvements of FAT can serve as strong evidence for its superiority. \n\n> Q1 Eq 5\n\n$L_e$ is the averaged ERM risk among $n_e$ samples in environment $e$. The IRMv1 penalty (Eq.4) takes the square of the gradients wrt. $L_e$ and thus $1/n_e$ in $L_e$ is $1/n_e^2$ in Eq. 5.\n\n> Q2 Memory cost\n\nFAT needs to store all $D^a_i$, $D^r_i$ and $w_i$ from previous rounds. The storage of these items does not cost much memory, while training using FAT objective (Eq.7) with all previous subsets can lead to OOM for a large network, as Eq. 7 increases the batch size by a factor of $k$ for the $k$-th round.\n\n> Q3 DFR results\n\nNote that the results are aligned with DFR [28] and its seminal works [25,46]. We strictly follow their protocol with details provided in line 325 and Appendix E.2.2, that we use an unbiased subset to train the last layer with frozen ERM learned features. The results demonstrate ERM already learns invariant features, aligned with our Thm 4.1.\n\n**References**\n\n[a] Deep linear networks with arbitrary loss: All local minima are global, ICML’18.\n\n[b] Implicit bias of gradient descent on linear convolutional networks, NeurIPS’18.\n\n[c] Representation costs of linear neural networks: Analysis and design, NeurIPS’21.\n\n[d] Inductive bias of multi-channel linear convolutional networks with bounded weight norm, COLT’22.\n"}, {"review_id": "DTfK23ao5K", "reviewer": "Reviewer_qcQH", "summary": "The paper studies OOD generalization in the presence of spurious correlation. First, it provides a theoretical analysis showing that during ERM training, both spurious and invariant features are learned but at different rates, which, in turn, influences the performance of the following optimization with OOD objectives, such as IRM. Then, it leverages the theoretical analysis to propose a new feature pre-training method to facilitate learning a more diverse set of features. The empirical evaluation shows that the proposed method improves standard ERM pre-training and other feature pre-training baselines in various settings.\n", "questions": "- According to the provided theory, pre-training features with ERM for longer seems always beneficial. However, as shown in Fig.1 left and mentioned in experiments (e.g., L306-307), overfitting can occur, and OOD performance decreases with further pre-training. Does your theory also explain this phenomenon? \n- It was unclear why the memory cost was too significant to switch to iFAT. Couldn’t the subsets be stored as masks, which would not increase the memory costs considerably, given that K equals 2 or 3 in practice? The additional parameters are the linear layers w_i, which also should not contribute much to the memory cost.\n", "limitations": "- As mentioned in the weaknesses, the presented proofs are limited to linear models. Additional work is needed to extend these results to the case of non-linear CNN. \n", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper provides a good combination of theoretical and empirical studies. Both parts are well-motivated, sound and offer a valuable contribution to the field:\n- The theoretical results deepen our understanding of how ERM learns different features and how this learning schedule interacts with further OOD training and generalization.\n- The proposed method, motivated by the theoretical analysis, is an efficient and effective way of improving OOD performance, as shown on several datasets.", "weaknesses": "- Discussions provided after each formal statement helps understand their implications. However, the overall clarity and presentation of the paper could be improved to be more friendly for a more general audience. For example:\n    - Fig. 2 – more details can be presented in the caption, e.g., it is unclear what the “Feature Learning” axis means. Are those coefficients introduced in Lemma 3.2?\n    - It is not entirely clear what “OOD objectives” mean before L122.\n    - The description of the FAT method could be more elaborate. For example, the description of the retention mechanism via saving previous linear layers could be explained more explicitly.\n- The theoretical setting and proofs use the linear activation function, effectively rendering the model linear. While the authors mention in the main text that this framework can be extended to ReLU activation functions, it would be good to provide more details in the proof of where and how the authors think the referred works (L119) could be applied to extend their results to non-linear networks.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time in reviewing our paper and your positive feedback! We hope our response below would make you more confident in supporting our work. \n\n> W1.1  Axis in Fig.2.\n\nWe have revised the caption of Fig.2 to include the details.  As mentioned in Appendix C.1, the invariant and spurious feature learning terms plotted in Fig. 2 are the means of $\\langle \\mathbf{w}\\_{j,r}, j\\mathbf{v}\\_1 \\rangle$ and $\\langle \\mathbf{w}\\_{j,r}, j\\mathbf{v}\\_2 \\rangle$ for $j\\in \\lbrace\\pm 1\\rbrace, r\\in[m]$, respectively. \n\n> W1.2  OOD objective.\n\nWe have revised our draft to make it clearer: we refer OOD objective as the additional penalty terms developed to regularize ERM to capture the invariant features, as introduced in line 29. \n\n> W1.3 Retention mechanism.\n\nWe have revised our work to explicitly describe the retention mechanism: FAT needs to retain the already learned features by minimizing the empirical risk at $G^r$, for which we store and use the historical classifiers $w_i$ with the current featurizer to evaluate the feature retention degree.\n\n> W2 Extending results to non-linear.\n\nHere we showcase how our two key theorems can be extended to non-linear settings:\n  - For the ERM feature learning  (Thm 4.1), we have updated our draft to include the analysis for non-linear activation functions. The result is that for smooth or piecewise linear activation functions (ReLU, Leaky ReLU, Softplus, etc.), suppose we run $T = O(d\\log{d})$ iterations of GD for the ERM objective (in the early stage of training), the feature learning terms (e.g., $\\sigma(\\langle \\mathbf{w}\\_{j,r}, j\\mathbf{v}\\_1 \\rangle)$) will be proportional to the empirical distributions of the signal strengths (e.g., $\\textup{Rad}(\\alpha)$) up to an error of $O(d^{-\\Omega(\\zeta)})$ with $\\zeta\\in (0, \\frac{1}{4})$ being a fixed constant. The proof ideas are summarized as follows: We adapt the results in [A] (Thm 3.2), which shows that in the early stage of GD training, for a network that is sufficiently wide (NTK region), its output can be well approximated by a linear model, and the approximation error scales as $O(d^{-\\Omega(\\zeta)})$.  Based on this approximation, we extended our analysis in the linear case to allow non-linear activation functions. We observe a weaker dominance of the spurious feature in the non-linear case both theoretically and empirically.  \n\n  - For IRM feature learning (Thm 4.2), we have updated our draft to include the analysis for non-linear activation functions. To account for non-linear activation, we introduce an additional assumption, which requires that non-linear activation is Lipschitz and Smooth: activation function is smooth, $\\psi'(0) \\le \\beta$, $|\\psi'(x)-\\psi'(x')|<\\beta|x - x'|$ and Lipschitz $|\\psi(0)| < L $, $|\\psi(x)-\\psi(x')|< L |x - x'|$. Based on the additional assumption, we can extend our analysis in the linear case to allow non-linear activation functions.\n\n\n\n> Q1 ERM pre-training for longer.\n\nWe need to clarify that **our theory does not imply that longer ERM pre-training would benefit the feature learning** while suggesting the ERM learns both invariant and spurious feature learning till convergence. In other words, ERM feature learning will **saturate** after a certain number of steps, which aligns with the empirical evidence in Fig. 10 of [25]. \n\nThe performance decrease of ERM feature learning in Fig. 1b implies one of the ERM feature learning drawbacks, which may be because of the “simplicity bias” that ERM tends to learn simple functions [50]. Our additional analysis in the rebuttal pdf also suggests that ERM can **forget** certain useful features. \n\nThere are more other factors that could influence feature learning, such as network architectures and optimization algorithms. As the first work that theoretically characterizes feature learning of ERM and OOD objectives, incorporating all of those factors could unnecessarily improve the difficulty of the analysis. Nevertheless, as discussed in Appendix A, we believe it is a promising future extension of our theory. \n\n> Q2 memory cost\n\nWe need to clarify that FAT needs to store all $D^a_i$, $D^r_i$ and $w_i$ yielded in previous rounds. The storage of these items does not cause much memory issue, while **training using FAT objective (Eq.7) with all previous subsets can lead to OOM for a large network**, as Eq. 7 increases the batch size by a factor of $k$ for the $k$-th round. Note that typically, the batches are also sampled from each environment as shown in Table 9. Let the number of sampled domains be $d$, and the batch size from each domain be $b$, then each round will additionally introduce $2d\\times b$ samples in each minibatch, which is typically more than 200 and leads to OOM issue for a V100 GPU with around 32G GPU memory.\n\nWe also show the running time and memory cost of different methods in the \"global\" response.\n\n**References** \n\n[25] On feature learning in the presence of spurious correlations, NeurIPS'22.\n\n[50] The pitfalls of simplicity bias in neural networks, NeurIPS'20.\n\n[A] The surprising simplicity of the early-time learning dynamics of neural networks, NeurIPS’20.\n\nWe are happy to answer any outstanding questions, and we’d appreciate it if you could jointly consider our responses when making the final evaluation of our work.\n"}]

Author rebuttal:
Dear reviewers,


We thank the reviewers for their many helpful comments and suggestions. 

Most reviewers agree that our theoretical findings are interesting, important and useful (KJZ8, NB5i, qcQH). The insights we obtained deepen the understandings of feature learning under distribution shifts (qcQH) and serve as a solid motivation for learning rich features for better OOD generalization (NB5i, qcQH). All reviewers agree that the proposed solution FAT is novel and the empirical improvements are strong and promising.

Regarding the reviewers’ concerns, we believe they can be addressed and we have revised our draft according to the reviewers’ valuable suggestions. In the following, we give our responses to the reviewers’ concerns and suggestions.

1. Regarding the linear setting (Jcjd, qcQH), we’d like to clarify that, because this is the **first** work to theoretically analyze the feature learning of ERM and OOD objectives and their interactions in OOD generalization: 
  - We chose a **minimal setting** where we can observe all the necessary phenomena and which enables us to study complicated OOD objectives (e.g., IRMv1) that involve high-order derivatives. **The results with IRM regularization we obtained are not obvious under non-convexity, despite the linear activation**.
  - Nevertheless, **we also show that our key theoretical results could be generalized to the non-linear setting**. Please find the details in [our response to Reviewer qcQH](https://openreview.net/forum?id=eozEoAtjG8&noteId=KmRvdQIZ2C) due to the character limits.
  - Going beyond the linearity of activations, the experiments in Fig. 1b show that, our theoretical results align with the empirical discoveries in more complex settings.

2. Regarding more analysis of feature learning (NB5i, qcQH), we plotted the saliency maps to better understand the feature learning of different algorithms. The figures are attached to the rebuttal pdf and to our latest draft as well.
  - We first visualize the feature learning of ERM and FAT on ColoredMNIST-025. It can be found that ERM can learn both invariant and spurious features to predict the label, aligned with our theory. 
  - However, ERM focuses more on spurious features and even forgets certain features with longer training epochs, which could be due to multiple reasons such as the simplicity biases of ERM. Hence predictions based on ERM learned features fail to generalize to OOD examples. In contrast, FAT effectively captures the meaningful features for all samples and generalizes to OOD examples well.
  - We also visualize the saliency maps of ERM, Bonsai and FAT on all real-world datasets used in our work. It can be found that, across various tasks and data modalities, FAT effectively learns more meaningful and diverse features than ERM and Bonsai, which serve as strong evidence for the consistent superiority of FAT in OOD generalization.


3. Regarding the computational complexity (KJZ8) and memory cost (Jcjd, qcQH), compared to ERM, the additional computational and memory overhead introduced in FAT mainly lie in the FAT training and partitioning. 

- At each training step, FAT needs $(k-1)$ additional forward and backward propagation, the same as Bonsai, while FAT only needs $\min(1,k-1)$ additional propagation. Besides, Bonsai additionally requires another round of training with $(K-1)$ additional propagation, given $K$ total rounds.
- Although FAT needs to store all $D^a_i$, $D^r_i$ and $w_i$ yielded in previous rounds, the storage of these items does not cause much a memory issue, while **training using the FAT objective (Eq.7) with all previous subsets can lead to OOM for a large network**, as Eq. 7 increases the batch size by a factor of $k$ for the $k$-th round. Note that typically, the batches are also sampled from each environment as shown in Table 9. Let the number of sampled domains be $d$, and the batch size from each domain be $b$, then each round will additionally introduce $2d\times b$ samples in each minibatch, which is typically more than 200 and results in OOM for a V100 GPU with around 32G GPU memory.

- We calculated the computational and memory overhead in the table below. The results aligned with our discussion. Bonsai requires much more time for the additional synthetic round and much more memory when there are 3 or more rounds. In contrast, **FAT achieves the best  OOD performance without introducing too much additional computational overhead**.

|        | Camelyon17            |                | CivilComments      |                |
|--------|-----------------------|----------------|--------------------|----------------|
|        | Training time         | Memory (%)     | Training time      | Memory (%)     |
| ERM    | 56.21$\pm$8.29 mins   | 22.56$\pm$0.00 | 24.22$\pm$0.33 hrs | 36.46$\pm$0.00 |
| Bonsai | 214.55$\pm$1.13 mins  | 51.75$\pm$0.01 | 58.47$\pm$0.91 hrs | 64.43$\pm$0.31 |
| FAT    | 101.14$\pm$12.79 mins | 51.92$\pm$0.04 | 28.19$\pm$1.15 hrs | 56.21$\pm$0.48 |


Besides, we have revised our paper to correct the typos and inconsistencies, provide more intuitive explanations of several key concepts and designs, and add a table of notations for easy reference. 

In addition, we have provided a link to our code for reproducing the results in our paper to the AC.

Please let us know if you have any further questions. Thanks. 


Meta-review: The work advances the understandings of feature learning under distribution shifts and it provides a solid foundation for learning rich features for better OOD generalization with strong empirical studies. Most reviewers are positive about the novelty of the work both theoretically and empirically. 

In the final version, we suggest authors incorporate all the suggestions raised by the reviewers, especially to improve the readability of the paper raised by several reviewers.
Decision: Accept (poster)