Title: Understanding and Improving Feature Learning for Out-of-Distribution Generalization
Venue: NeurIPS 2023
Accepted: yes
Keywords: Out-of-Distribution Generalization

Feature Learning

Invariant Risk Minimization
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

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 F eature A ugmented T raining (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 1 .

1 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 minimization (ERM) learn highly predictive features that generalize surprisingly well to in-distribution (ID) data [24, 75]. 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-of-distribution (OOD) data, and suffers from serious performance degeneration [39]. Therefore, it is 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 already learned sufficiently good features that are able to generalize to OOD data [33, 38, 63]. In 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 a curious phenomenon that the

Pre-training

OOD training

ϕ

Featurizer

ERM

ˆ

y

w

ϕ

Featurizer

FeAT

D

tr

D

r

D

a

ˆ

y

w

Invariant Features

Spurious Features

Learned  Features

Underlying Features

IRMv1

D

tr

IRMv1

D

tr

D

tr

0

50

100

150

200

250

FeAT

0

25

50

IRMv1

0

50

100

150

200

250

FeAT

VREx

0

50

100

150

200

250

FeAT

0

25

50

IRMX

0

50

100

150

200

250

FeAT

CLoVE

0

50

100

150

200

250

FeAT

0

25

50

IGA

0

50

100

150

200

250

FeAT

0.0

0.2

0.4

0.6

0.8

1.0

pretrain epochs

0.0

0.2

0.4

0.6

0.8

1.0

test acc

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-training epochs has a large influence on the final OOD performance. These seemingly contradicting phenomena raise a 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 a two-layer CNN network, when trained with ERM and a widely used OOD objective, IRMv1 [4], respectively. We use a 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 of 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 learn these features during OOD training and the OOD performance will be limited. Given a 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 a bottleneck to OOD generalization [60].

To remedy the issue, we propose F eature A ugmented T raining (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 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, 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 a and retains the already learned features by minimizing ERM losses on all D r . Along with the growth of the augmentation and retention sets, FeAT is able to learn richer features for OOD training and obtain a 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, 70]. 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 to 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 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-learn features and improve OOD generalization. Deng et al. [20] finds neural networks tend to learn spurious features under imbalanced groups. In contrast, we study the direct effects of ERM and OOD objectives to feature learning and provide a 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 a 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-retention on the 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 elusive. In addition to a new rich feature learning algorithm, our work provides theoretical explanations for the success of rich feature learning in OOD generalization.

3 Preliminaries and Problem Definition

Notations. We use old-faced letters for vectors and matrices otherwise for scalar; ∥ · ∥ 2 to denote the Euclidean norm of a vector or the spectral norm of a matrix, while ∥ · ∥ F for the Frobenius norm of a matrix. I d refers to the identity matrix in R d × d . Full details are deferred to Appendix A.

Our data model D = { x i , y i } n i =1 is adapted from [2, 12] and further characterizes each data point x i as invariant and spurious feature patches from the two-bit model [16, 36].

Definition 3.1. D = {D e } e ∈E all is composed of multiple subsets D e from different environments e ∈ E all, where each D e = { ( x e i , y e i ) } n e i =1 is composed of i.i.d. samples ( x e i , y e i ) ∼ P e . Each data ( x e , y e ) ∈ D e with x e ∈ R 2 d and y e ∈ {-1 , 1 } is generated as follows:

(a) Sample y e ∈ {-1 , 1 } uniformly;

(b) Given y e , each input x e = [ x e 1 , x e 2 ] contains a feature patch x 1 and a noise patch x 2 , that are sampled as:

where Rad ( δ ) is a random variable taking value -1 with probability δ and +1 with probability 1 -δ , v 1 = [1 , 0 , . . . 0] ⊤ and v 2 = [0 , 1 , 0 , . . . 0] ⊤ .

(c) A noise vector ξ is generated from the Gaussian distribution N ( 0 , σ 2 p · ( I d -v 1 v ⊤ 1 -v 2 v ⊤ 2 ))

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 are noises that are irrelevant to the label. In particular, v 1 and v 2 are feature vectors that simulate the invariant and spurious features, respectively. Although our data model focuses on two feature vectors, the discussion and results can be further generalized to multiple invariant and spurious features with fine-grained characteristics [68]. Following previous works [12], we 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 E α = { ( α, β e ) : 0 < β e < 1 } , where v 1 is the invariant feature as α is fixed while v 2 is the spurious feature as β e varies across e .

CNN model. We consider training a two-layer convolutional neural network with a hidden layer width of m . The filters are applied to x 1 , x 2 , respectively, 2 and the second layer parameters of the network are fixed as 1 m and -1 m , respectively. Then the network can be written as f ( W , x ) = F +1 ( W +1 , x ) -F -1 ( W -1 , x ) , where F +1 ( W +1 , x ) and F -1 ( W -1 , x ) are defined as follows:

where ψ ( x ) is the activation function. We assume that all network weights are initialized as N (0 , σ 2 0 ) .

ERMobjective. We train the CNN model by minimizing the empirical cross-entropy loss function:

where ℓ ( z )=log(1+exp( -z )) and {D e } e ∈E tr = {{ x e i , y e i } n e i =1 } e ∈E tr is the trainset with ∑ e ∈E tr n e = n .

OOD objective. The goal of OOD generalization is, given the data from training environments {D e } e ∈E tr , to find a predictor f : X → Y that generalizes well to all (unseen) environments, or minimizes max e ∈E all L e ( f ) , where L e is the empirical risk under environment e . The predictor f = w ◦ φ is usually composed of a featurizer φ : X → Z that learns to extract useful features, and a classifier w : Z → 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 an invariant representation φ , such that there exists a classifier acting on φ that is simultaneously optimal in E tr . Hence, IRM leads to a challenging bi-level optimization problem as

Due to the optimization difficulty of Eq. (3), Arjovsky et al. [4] relax Eq. (3) into IRMv1 as follows:

Given the convolutional neural network (Eq. 1) and logistic loss (Eq. 2), IRMv1 can be written as

where ℓ ′ e i = ℓ ′ ( y e i · f ( W , x e i )) = -exp( -y e i · f ( W , x e i )) 1+exp( -y e i · f ( W , x e i )) . Due to the complexity of IRMv1, in the analysis below, we introduce C e IRMv1 for the ease of expressions. Specifically, we define C e IRMv1 as

where ˆ y e i ≜ f ( W , x e i ) is the logit of sample x i from environment e . The convergence of C IRMv1 indicates the convergence of IRMv1 penalty. The following lemma will be useful in our analysis.

0

2000

4000

6000

Epoch

-0.15

-0.10

-0.05

0.00

0.05

CIRMv

1

C

1

IRMv

1

C

2

IRMv

1

(a)

C

IRMv1, w/ PT

0

2000

4000

6000

Epoch

0

2000

4000

6000

8000

10000

Feature Learning

0.0

0.2

0.4

0.6

0.8

1.0

Loss

inv. feat.

spu. feat.

train loss

test loss

(b) FL, w/ PT

0

2000

4000

6000

Epoch

0.0000

0.0005

0.0010

0.0015

CIRMv

1

C

1

IRMv

1

C

2

IRMv

1

(c)

C

IRMv1, w/o PT

0

2000

4000

6000

Epoch

-30

-20

-10

0

10

20

30

Feature Learning

0.650

0.675

0.700

0.725

0.750

0.775

Loss

inv. feat.

spu. feat.

train loss

test loss

(d) FL, w/o PT

Lemma 3.2. (Cao et al. [12]) Let w j,r ( t ) 3 for j ∈ { +1 , -1 } and r ∈ { 1 , 2 , . . . , m } be the convolution filters of the CNN at t -th iteration of gradient descent. Then there exists unique coefficients γ inv j,r ( t ) , γ spu j,r ( t ) ≥ 0 and ρ j,r,i ( t ) such that,

We refer Eq. (6) as the signal-noise decomposition of w j,r ( t ) [12]. We add normalization factor ∥ ξ i ∥ -2 2 in the definition so that ρ ( t ) j,r ≈ ⟨ w ( t ) j,r , ξ i ⟩ . Note that ∥ v 1 ∥ 2 = ∥ v 2 ∥ 2 = 1 , the corresponding normalization factors are thus neglected. Furthermore, γ inv j,r ≈ ⟨ w j,r , v 1 ⟩ and γ spu j,r ≈ ⟨ w j,r , v 2 ⟩ respectively denote the degrees of invariant and spurious feature learning.

4 Theoretical Understanding of Feature Learning in OOD Generalization

4.1 ERMFeature Learning

With the setup in Sec. 3, we first study the feature learning of the ERM objective. We consider a two training environments setup E tr = { ( α, β 1 ) , ( α, β 2 ) } where the signal of invariant feature is weaker than the average of spurious signals (i.e., α > β 1 + β 2 2 ), which corresponds to Figure 2. For a precise characterization of the training dynamic, we adopted a minimal setup where ψ ( 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 ρ > 0 , let n ≜ min e ∈E tr n e . Suppose that we run T iterations of GD for the ERM objective. With sufficiently large n and ψ ( x ) = x , assuming that (i) α, β 1 , β 2 < 1 2 , and (ii) α > β 1 + β 2 2 , with properly chosen σ 2 0 and σ 2 p , there exists a constant η , such that with probability at least 1 -2 ρ , 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 ∈ { 0 , . . . , 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.1 states that ERMtraining learns both invariant feature and spurious feature at the same time, and if the average of spurious signals is 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 a classifier based on a subset of unbiased data on top of the ERM-trained featurizer achieves impressive OOD generalization performance [33, 38, 63].

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 a 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-learned 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 (1) with data model (3.1), define c ( t ) ≜ [ C 1 IRMv1 ( W , t ) , C 2 IRMv1 ( W , t ) , · · · , C |E tr | IRMv1 ( W , t ) ] , and λ 0 = λ min ( H ∞ ) , where H ∞ e,e ′ ≜ 1 2 mn e n e ′ ∑ n e i =1 ψ ′ ( ⟨ w j,r (0) , x e 1 ,i ⟩ ) x e ⊤ 1 ,i ∑ n e ′ i ′ =1 ψ ′ ( ⟨ w j,r (0) , x e ′ 1 ,i ′ ⟩ ) x e ′ 1 ,i ′ . Suppose that dimension d = Ω(log( m/δ )) , network width m = Ω(1 /δ ) , regularization factor λ ≥ 1 / ( σ 0 |E tr | 3 / 2 ) , noise variance σ p = O ( d -2 ) , weight initial scale σ 0 = O ( |E tr | 7 / 2 β 3 L d 1 / 2 m 2 λ 2 0 log(1 /ϵ ) ) , then with probability at least 1 -δ , after training time T = Ω ( log(1 /ϵ ) ηλλ 0 ) , we have ∥ c ( T ) ∥ 2 ≤ ϵ, γ inv j,r ( T )= o (1) , γ spu j,r ( T )= o (1) .

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, Theorem 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 at the beginning of the training. Then, what would happen when given a properly pre-trained network?

After ERM pre-training, according to Theorem 4.1, we have |⟨ w j,r , v 1 ⟩| = Ω(1) , |⟨ w j,r , v 2 ⟩| = Ω(1) , and |⟨ w j,r , ξ ⟩| = O ( σ 0 σ p √ d ) . Then, we have the following hold.

Proposition 4.3. Given the same setting as Theorem 4.2, suppose that ψ ( x ) = x , γ inv j,r ( t 1 ) = γ inv j,r ( t 1 -1) , and γ spu j,r ( t 1 ) = γ spu j,r ( t 1 -1) at the end of ERM pre-train t 1 , δ > 0 , and n > C log(1 /δ ) , with C being a positive constant, then with a high probability at least 1 -δ , we have ∑ e C e IRMv1 ( t 1 ) = 0 , γ inv j,r ( t 1 +1) > γ inv j,r ( t 1 ) , and γ spu j,r ( t 1 +1) < γ spu j,r ( t 1 ) .

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-training [16, 83].

4.3 Limitations of ERM Feature Learning

Combining results from both Sec. 4.1 and Sec. 4.2, we know that the invariant features will be learned during ERM pre-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 can 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 ψ ( x ) = x , γ inv j,r ( t 1 ) = o (1) , and γ spu j,r ( t 1 ) = Θ(1) at the end of ERM pre-training t 1 . Suppose that δ > 0 , and n > C log(1 /δ ) , with C being a positive constant, then with a high probability at least 1 -δ , we have γ inv j,r ( t 1 +1) < γ inv j,r ( t 1 ) .

Corollary 4.4 shows that IRMv1 requires sufficiently well-learned features for OOD generalization. It 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 pretraining stage for OOD generalization. Otherwise, the OOD training is less likely to 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, we propose F eature A ugmented T raining (FeAT), that directly augment the feature learning in an iterative manner.

Intuitively, the potentially useful features presented in the training data are features that have non-trivial correlations with labels, or using the respective feature to predict the labels is able to achieve a non-trivial training performance . Moreover, the invariance principle assumes that the training data comes from different environments [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 differentiate the distinct sets of useful features entangled in the training data into different subsets, where ERM can effectively learn the dominant features presented in the corresponding subset as shown in Theorem 4.1.

The intuition naturally motivates an iterative rich feature learning algorithm, i.e., FeAT, that identifies the subsets containing distinct features and ex-

Algorithm 1 FeAT: F eature A ugmented T raining

1: Input: Training data D tr ; the maximum augmentation rounds K ; predictor f := w ◦ φ ; length of inner training epochs t ; termination threshold p ;

2: Initialize groups G a ←D tr , G r ←{} ;

3: for k ∈ [1 , . . . , K ] do

4: Randomly initialize w k ;

5: for j ∈ [1 , . . . , t ] do

6: Obtain ℓ FeAT with G via Eq. 7;

7: Update w k , φ with ℓ FeAT;

8: end for

9: // Early Stop if f k = w k ◦ φ fails to find new features.

10: if Training accuracy of f k is smaller than p then

11: Set K = k -1 and terminate the loop;

12: end if

13: Split D tr into groups D a k , D r k according to whether f k classifies the examples in D tr correctly or not;

15: end for

14: Update groups G a ← G a ∪{D a k } , G r ← G r ∪{D r k } ;

16: Synthesize the final classifier w ← 1 K ∑ K i =1 w i ;

17: return f = w ◦ φ ;

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 ◦ φ 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 f yields the correct prediction, denoted as G r k , and the subset of samples that contains the features that have not been learned, denoted as G a k .

At the k -th round, given the grouped subsets G = { G r , G a } with 2 k -1 groups, where G a = {D a i } k -1 i =0 is the grouped sets for new feature augmentation, and G r = {D r i } k -1 i =1 is the grouped sets for already learned feature retention (notice that D r 0 is the empty set), where D a i and D r i are the corresponding augmentation and retention set elicited at i -th round. FeAT performs distributionally robust optimization (DRO) [50, 83] on G a to explore new features that have not been learned in previous rounds. Meanwhile, FeAT also 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. Then, the FeAT objective at round k is

where ℓ D i ( w ◦ φ ) refers to the empirical risk of w ◦ φ evaluated at the subset D i , and { w i | 1 ≤ i ≤ 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 in 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 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 in Bonsai could also lead to suboptimal performance and meanwhile require more computational overhead.

Practical implementations. Algorithm 1 requires to store 2 K -1 subsets and a larger memory cost in training the network, which may cause additional storage burden when φ contains a massive amount of parameters [39]. Hence, we propose a lightweight variant of FeAT (denoted as iFeAT) which only retains the latest subsets and historical classifiers ( D a k -1 , D r k -1 , w k -1 at the k -th round). Throughout the whole experiment, we will use iFeAT and find that iFeAT already achieves state-of-the-art. More details are given in Appendix E.

As iFeAT stores only the latest augmentation and retention subsets, inspecting the training performance 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 to 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 if so).

Compared to ERM, the additional computational and memory overhead introduced in FeAT mainly lie in the FeAT training and partitioning. At each training step, FeAT needs ( k -1) additional forward and backward propagation, the same as Bonsai, while FeAT 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. More details are given in Appendix F.4.

6 Empirical Study

We conduct extensive experiments on COLOREDMNIST [16] and WILDS [39] to verify the effectiveness of FeAT in learning richer features than ERM and the state-of-the-art algorithm Bonsai [83].

Proof-of-concept study on COLOREDMNIST. We first conduct a proof-of-concept study using COLOREDMNIST [16] and examine the feature learning performance of FeAT under various conditions. We consider both the original COLOREDMNIST with E tr = { (0 . 25 , 0 . 1) , (0 . 25 , 0 . 2) } (denoted as COLOREDMNIST-025), where spurious features are better correlated with labels, and the modified COLOREDMNIST (denoted as COLOREDMNIST-01) with E tr = { (0 . 1 , 0 . 2) , (0 . 1 , 0 . 25) } , 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-art rich feature learning algorithm Bonsai [83]. Based on the features learned by ERM, Bonsai, and FeAT, we adopt various state-ofthe-art OOD objectives including IRMv1 [4], VREx [41], IRMX [16], IB-IRM [1], CLOvE [76], IGA [40] and Fishr [57] for OOD training, in order to evaluate the practical quality of the learned features. The feature representations are frozen once initialized for the OOD training as fine-tuning the featurizer can distort the pre-trained features [43]. We also compare FeAT with the common training approach that uses unfrozen ERM features, denoted as ERM-NF. For Bonsai, we trained 2

† DFR/DFR-s use an additional OOD dataset to evaluate invariant and spurious feature learning, respectively.

rounds following Zhang et al. [83], while for FeAT the automatic termination stopped at round 2 in COLOREDMNIST-025 and round 3 in COLOREDMNIST-01. For ERM, we pre-trained the model with the same number of overall epochs as FeAT in COLOREDMNIST-01, while early stopping 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.1.

The results are reported in Table 1. It can be found that ERM will learn insufficiently good features under both stronger spurious correlations and invariant correlations, confirming our discussion in Sec. 4.3. Besides, Bonsai learns richer features in COLOREDMNIST-025 and boosts OOD performance, but Bonsai sometimes leads to suboptimal performances in COLOREDMNIST-01, 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 RXRX1 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 a proper measure for the quality of learned invariant features [33]. When the original dataset does not provide a proper OOD validation set, e.g., CAMELYON17, we use an alternative implementation based on a 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 neural architectures become more complicated, the ERM learned features can have a lower quality 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 underperforms ERM in rich feature learning and OOD generalization in multiple datasets. In con-

pred:0 label:0

pred:1 label:1

pred:1 label:1

pred:0 label:0

pred:1 label:1

pred:0 label:0

pred:1 label:1

pred:0 label:0

pred:0 label:1

pred:1 label:0

pred:1 label:0

pred:1 label:0

pred:0 label:1

pred:1 label:0

pred:0 label:1

pred:1 label:0

(a) ERM 300 epochs

pred:0 label:0

pred:1 label:1

pred:1 label:1

pred:0 label:0

pred:1 label:1

pred:0 label:0

pred:1 label:1

pred:0 label:0

pred:1 label:1

pred:0 label:0

pred:0 label:0

pred:0 label:0

pred:0 label:1

pred:0 label:0

pred:1 label:1

pred:0 label:0

(b) FeAT 2 rounds

trast, FeAT consistently achieves the best invariant feature learning performance across various challenging realistic datasets. Meanwhile, compared to ERM and Bonsai, FeAT also reduces over-fitting to the spurious feature learning led by spurious correlations. As a result, FeAT achieves consistent improvements when the learned features are applied to various OOD objectives.

The termination check in FeAT. As elaborated in Sec. 5.2, a key difference 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 in various subsets of COLOREDMNIST-025 at different rounds in Table 3. By inspecting the retention accuracy, after FeAT learns sufficiently good features at Round 2 , it is not necessary to proceed with Round 3 as it will destroy the already learned features and lead to degenerated retention and OOD performance. More details and results are given in Appendix F.1.

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 ColoredMNIST025. As shown in Fig. 3, 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, FeAT effectively captures the meaningful features for all samples and generalizes to OOD examples well. More analysis including results on WILDS benchmark can be found in Appendix F.5.

7 Conclusions

In this paper, we conducted a 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 in 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.

References

K. Ahuja, E. Caballero, D. Zhang, J.-C. Gagnon-Audet, Y. Bengio, I. Mitliagkas, and I. Rish. Invariance principle meets information bottleneck for out-of-distribution generalization. In Advances in Neural Information Processing Systems , 2021.

Z. Allen-Zhu and Y. Li. Towards understanding ensemble, knowledge distillation and selfdistillation in deep learning. arXiv preprint arXiv:2012.09816 , 2020.

A. Andreassen, Y. Bahri, B. Neyshabur, and R. Roelofs. The evolution of out-of-distribution robustness throughout fine-tuning. arXiv preprint arXiv:2106.15831 , 2021.

M. Arjovsky, L. Bottou, I. Gulrajani, and D. Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv:1907.02893 , 2019.

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.

P. Bándi, O. Geessink, Q. Manson, M. V. Dijk, M. Balkenhol, M. Hermsen, B. E. Bejnordi, B. Lee, K. Paeng, A. Zhong, Q. Li, F. G. Zanjani, S. 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, E. Halici, H. Jackson, R. Chen, F. Both, J. Franke, H. Küsters-Vandevelde, 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.

S. Beery, G. V. Horn, and P. Perona. Recognition in terra incognita. In Computer Vision European Conference, Part XVI , pages 472-489, 2018.

S. Beery, E. Cole, and A. Gjoka. The iwildcam 2020 competition dataset. arXiv preprint arXiv:2004.10340 , 2020.

D. Borkan, L. Dixon, J. Sorensen, N. Thain, and L. Vasserman. Nuanced metrics for measuring unintended bias with real data for text classification. In Companion of The 2019 World Wide Web Conference , pages 491-500, 2019.

T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei. Language models are few-shot learners. In Advances in Neural Information Processing Systems , 2020.

A. Brutzkus, A. Globerson, E. Malach, and S. Shalev-Shwartz. SGD learns over-parameterized networks that provably generalize on linearly separable data. In International Conference on Learning Representations , 2018.

Y. Cao, Z. Chen, M. Belkin, and Q. Gu. Benign overfitting in two-layer convolutional neural networks. In Advances in Neural Information Processing Systems , 2022.

H. S. Chen, Y. Lee, A. Setlur, S. Levine, and C. Finn. Project and probe: Sample-efficient domain adaptation by interpolating orthogonal features. arXiv preprint arXiv:2302.05441 , 2023.

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.

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.

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.07766 , 2022.

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.

G. A. Christie, N. Fendley, J. Wilson, and R. Mukherjee. Functional map of the world. In IEEE Conference on Computer Vision and Pattern Recognition , pages 6172-6180, 2018.

A. J. DeGrave, J. D. Janizek, and S. Lee. AI for radiographic COVID-19 detection selects shortcuts over signal. Nature Machince Intelligence , 3(7):610-619, 2021.

Y. Deng, Y. Yang, B. Mirzasoleiman, and Q. Gu. Robust learning with progressive data expansion against spurious correlation. arXiv preprint , arXiv:2306.04949, 2023.

N. Elhage, T. Hume, C. Olsson, N. Schiefer, T. Henighan, S. Kravec, Z. Hatfield-Dodds, R. Lasenby, D. Drain, C. Chen, R. Grosse, S. McCandlish, J. Kaplan, D. Amodei, M. Wattenberg, and C. Olah. Toy models of superposition. arXiv preprint arXiv:2209.10652 , 2022.

S. Frei, Y . Cao, and Q. Gu. Provable generalization of sgd-trained neural networks of any width in the presence of adversarial label noise. In International Conference on Machine Learning , pages 3427-3438, 2021.

R. Geirhos, J. Jacobsen, C. Michaelis, R. S. Zemel, W. Brendel, M. Bethge, and F. A. Wichmann. Shortcut learning in deep neural networks. Nature Machine Intelligence , 2(11):665-673, 2020.

I. Goodfellow, Y. Bengio, and A. Courville. Deep Learning . MIT Press, 2016. http: //www.deeplearningbook.org .

I. Gulrajani and D. Lopez-Paz. In search of lost domain generalization. In International Conference on Learning Representations , 2021.

A. Gupta, N. Saunshi, D. Yu, K. Lyu, and S. Arora. New definitions and evaluations for saliency methods: Staying intrinsic and sound, 2022.

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.

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.

W. Hu. Understanding Deep Learning via Analyzing Dynamics of Gradient Descent . Princeton University, 2021.

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.

G. Huang, Z. Liu, L. van der Maaten, and K. Q. Weinberger. Densely connected convolutional networks. In IEEE Conference on Computer Vision and Pattern Recognition , pages 2261-2269, 2017.

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.

P. Izmailov, P. Kirichenko, N. Gruver, and A. G. Wilson. On feature learning in the presence of spurious correlations. In Advances in Neural Information Processing Systems , 2022.

Y. Ji, L. Zhang, J. Wu, B. 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-distribution (ood) dataset curator and benchmark for ai-aided drug discovery - a focus on affinity prediction problems with noise annotations. arXiv preprint arXiv:2201.09637 , 2022.

D. Kalimeris, G. Kaplun, P. Nakkiran, B. Edelman, T. Yang, B. Barak, and H. Zhang. SGD on neural networks learns functions of increasing complexity. Advances in neural information processing systems , 32, 2019.

P. 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.

D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations , 2015.

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.

P. W. Koh, S. Sagawa, H. Marklund, S. M. Xie, M. Zhang, A. Balsubramani, W. Hu, M. Yasunaga, R. L. Phillips, I. Gao, T. Lee, E. David, I. Stavness, W. Guo, B. Earnshaw, I. Haque, S. M. Beery, J. Leskovec, A. Kundaje, E. Pierson, S. Levine, C. Finn, and P. Liang. WILDS: A benchmark of in-the-wild distribution shifts. In International Conference on Machine Learning, , pages 5637-5664, 2021.

M. Koyama and S. Yamaguchi. Out-of-distribution generalization with maximal invariant predictor. arXiv preprint arXiv:2008.01883 , 2020.

D. Krueger, E. Caballero, J. Jacobsen, A. Zhang, J. Binas, D. Zhang, R. L. Priol, and A. C. Courville. Out-of-distribution generalization via risk extrapolation (rex). In International Conference on Machine Learning , pages 5815-5826, 2021.

A. Kumar, A. Raghunathan, R. M. Jones, T. Ma, and P. Liang. Fine-tuning can distort pretrained features and underperform out-of-distribution. In International Conference on Learning Representations , 2022.

A. Kumar, A. Raghunathan, R. M. Jones, T. Ma, and P. Liang. Fine-tuning can distort pretrained features and underperform out-of-distribution. In International Conference on Learning Representations , 2022.

Y. Lee, H. Yao, and C. Finn. Diversify and disambiguate: Out-of-distribution robustness via disagreement. In International Conference on Learning Representations , 2023.

Y. Lin, L. Tan, Y. Hao, H. Wong, H. Dong, W. Zhang, Y. Yang, and T. Zhang. Spurious feature diversification improves out-of-distribution generalization. arXiv preprint arXiv:2309.17230 , 2023.

Y. Lin, L. Tan, H. Lin, Z. Zheng, R. Pi, J. Zhang, S. Diao, H. Wang, H. Zhao, Y. Yao, and T. Zhang. Speciality vs generality: An empirical study on catastrophic forgetting in fine-tuning foundation models. arXiv preprint arXiv:2309.06256 , 2023.

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 out-of-distribution and in-distribution generalization. In International Conference on Machine Learning , pages 7721-7735, 2021.

R. Min, Z. Qin, L. Shen, and M. Cheng. Towards stable backdoor purification through feature shift tuning. arXiv preprint arXiv:2310.01875 , 2023.

H. Naganuma, K. Ahuja, I. Mitliagkas, S. Takagi, T. Motokawa, R. Yokota, K. Ishikawa, and I. Sato. Empirical study on optimizer selection for out-of-distribution generalization. arXiv preprint arXiv:2211.08583 , 2022.

H. Namkoong and J. C. Duchi. Stochastic gradient methods for distributionally robust optimization with f-divergences. In Advances in Neural Information Processing Systems , pages 2208-2216, 2016.

J. Ni, J. Li, and J. McAuley. Justifying recommendations using distantly-labeled reviews and fine-grained aspects. In Conference on Empirical Methods in Natural Language Processing and International Joint Conference on Natural Language Processing , pages 188-197, 2019.

A. M. Nicolicioiu, A. L. Nicolicioiu, B. Alexe, and D. Teney. Learning diverse features in vision transformers for improved generalization. 2023.

G. Parascandolo, A. Neitz, A. Orvieto, L. Gresele, and B. Schölkopf. Learning explanations that are hard to vary. In International Conference on Learning Representations , 2021.

A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. DeVito, 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.

M. Pezeshki, S. Kaba, Y. Bengio, A. C. Courville, D. Precup, and G. Lajoie. Gradient starvation: Alearning proclivity in neural networks. In Advances in Neural Information Processing Systems , pages 1256-1272, 2021.

A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. 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.

A. Rame, C. Dancette, and M. Cord. Fishr: Invariant gradient variances for out-of-distribution generalization. arXiv preprint arXiv:2109.02934 , 2021.

A. Ramé, K. Ahuja, J. Zhang, M. Cord, L. Bottou, and D. Lopez-Paz. Model ratatouille: Recycling diverse models for out-of-distribution generalization. arXiv preprint arXiv:2212.10445 , 2022.

A. Rame, M. Kirchmeyer, T. Rahier, A. Rakotomamonjy, patrick gallinari, and M. Cord. Diverse weight averaging for out-of-distribution generalization. In Advances in Neural Information Processing Systems , 2022.

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.

M. Rojas-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.

F. Rosenblatt. The perceptron - a perceiving and recognizing automaton. Technical Report 85-460-1, Cornell Aeronautical Laboratory, 1957.

E. Rosenfeld, P. Ravikumar, and A. Risteski. Domain-adjusted regression or: Erm may already learn features sufficient for out-of-distribution generalization. arXiv preprint arXiv:2202.06856 , 2022.

S. Sagawa*, P. W. Koh*, T. B. Hashimoto, and P. Liang. Distributionally robust neural networks. In International Conference on Learning Representations , 2020.

W. Samek, G. Montavon, A. Vedaldi, L. K. Hansen, and K.-R. Muller. Explainable AI: Interpreting, Explaining and Visualizing Deep Learning . Springer Publishing Company, Incorporated, 1st edition, 2019. ISBN 3030289532.

V. Sanh, L. Debut, J. Chaumond, and T. Wolf. Distilbert, a distilled version of BERT: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01108 , 2019.

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.

R. Shen, S. Bubeck, and S. Gunasekar. Data augmentation as feature manipulation. In International Conference on Machine Learning , pages 19773-19808, 2022.

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.

R. Shwartz-Ziv and N. Tishby. Opening the black box of deep neural networks via information. arXiv preprint arXiv:1703.00810 , 2017.

R. Taori, A. Dave, V. Shankar, N. Carlini, B. Recht, and L. Schmidt. Measuring robustness to natural distribution shifts in image classification. In Advances in Neural Information Processing Systems , 2020.

J. Taylor, B. Earnshaw, B. Mabey, M. Victors, and J. Yosinski. Rxrx1: An image set for cellular morphological variation across many experimental batches. In International Conference on Learning Representations , 2019.

D. Teney, E. Abbasnejad, S. Lucey, and A. van den Hengel. Evading the simplicity bias: Training a diverse set of models discovers solutions with superior OOD generalization. In IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 16740-16751, 2022.

D. Teney, Y. Lin, S. J. Oh, and E. Abbasnejad. ID and OOD performance are sometimes inversely correlated on real-world datasets. arXiv preprint arXiv:2209.00613 , 2022.

V. Vapnik. Principles of risk minimization for learning theory. In Advances in Neural Information Processing Systems , pages 831-838, 1991.

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.

Z. Wang, Y. Chen, Y. Duan, W. Li, B. Han, J. Cheng, and H. Tong. Towards out-of-distribution generalizable predictions of chemical kinetics properties. arXiv preprint , arXiv:2310.03152, 2023.

Z. Wen and Y. Li. Toward understanding the feature learning process of self-supervised contrastive learning. In International Conference on Machine Learning , pages 11112-11122, 2021.

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 out-of-distribution generalization in transfer learning. In Advances in Neural Information Processing Systems , 2022.

M. Wortsman, G. Ilharco, J. W. Kim, M. Li, S. Kornblith, R. Roelofs, R. 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.

H. Ye, J. Zou, and L. Zhang. Freeze then train: Towards provable representation learning under spurious correlations and feature noise. arXiv preprint arXiv:2210.11075 , 2022.

J. Zhang and L. Bottou. Learning useful representations for shifting tasks and distributions. arXiv preprint arXiv:2212.07346 , 2022.

J. Zhang, D. Lopez-Paz, and L. Bottou. Rich feature construction for the optimizationgeneralization dilemma. In International Conference on Machine Learning , pages 26397-26411, 2022.

Y. Zhang, M. Gong, T. Liu, G. Niu, X. Tian, B. Han, B. Schölkopf, and K. Zhang. Adversarial robustness through the lens of causality. In International Conference on Learning Representations , 2022.

X. Zhou, Y. Lin, R. Pi, W. Zhang, R. Xu, P. Cui, and T. Zhang. Model agnostic sample reweighting for out-of-distribution learning. In International Conference on Machine Learning , pages 27203-27221, 2022.

D. Zou, Y. Cao, Y. Li, and Q. Gu. Understanding the generalization of adam in learning neural networks with proper regularization. arXiv preprint arXiv:2108.11371 , 2021.

Contents

Appendix of FeAT

A Notations

We use bold-faced letters for vectors and matrices otherwise for scalar. We use ∥ · ∥ 2 to denote the Euclidean norm of a vector or the spectral norm of a matrix, while denoting ∥ · ∥ F as the Frobenius norm of a matrix. For a neural network, we denote ψ ( x ) as the activation function. Let I d be the identity matrix with a dimension of R d × d . When comparing two sequences { a n } and { b n } , we employ standard asymptotic notations such as O ( · ) , o ( · ) , Ω( · ) , and Θ( · ) to describe their limiting behavior. Lastly, sequences of integers are denoted as [ n ] = { 1 , 2 , . . . , n } .

B Limitations and Future Directions

As a 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 a feature and the capacity of features that a 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, 70]. 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 a network can be influenced by several other factors, such as the difficulty of learning a feature and the capacity of features that a 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 a 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-learn 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 a 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 to 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-tuning 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 et 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 a solution for mitigating the trade-offs in optimization. Theorem 4.1 implies that, in cases where invariant features are more informative than spurious features, the higher ID performance indicates a 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 a better fit to spurious features, thus bringing a lower OOD performance [74]. Similarly, when fine-tuning a pre-trained model, if the model does not learn the features sufficiently well, ID-OOD performance will be in a positive correlation. However, when spurious correlations are present as easy-to-learn features, ERM can lead to a better fit for spurious features and distort the previously learned invariant features [42, 46, 80].

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 of 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. In contrast, FeAT relieves the requirements by performing direct augmentation-retention on the grouped subsets, and thus obtains better performance. More crucially, 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 on 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.

Besides, there exists a 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 to 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 to 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 ψ ( x ) = x , the logit ˆ y e i (which is a function of W ) of sample i in the environment e can be explicitly written as

where W ≜ { W +1 , W -1 } and W j ≜    w ⊤ j, 1 . . . w ⊤ j,m    for j ∈ {± 1 } . We initialized all the network weights as N (0 , σ 2 0 ) and we set σ 0 = 0 . 01 .

The test dataset ( x , y ) is generated through

where half of the dataset uses Rad (1 -β 1 ) and the other half uses Rad (1 -β 2 ) . Here ξ ∼ N (0 , σ 2 p · ( I d -v 1 v ⊤ 1 -v 2 v ⊤ 2 )) and we chose σ p = 0 . 01 .

From the definition of IRMv1, we take derivative wrt. the scalar 1 of the logit 1 · ˆ y e i . Thus, for environment e , the penalty is

Then, the IRMv1 objective is (we set n 1 = n 2 = 2500 in the simulation)

We used constant stepsize GD to minimize L IRMv1 ( W ) , and we chose λ = 10 8 (heavy regularization setup).

Let C e IRMv1 ≜ 1 n e ∑ n e i =1 ℓ ′ ( y e i ˆ y e i ) · y e i ˆ y e i . The gradient of L IRMv1 ( W ) with respect to each w j,r can be explicitly written as

Observe that C e IRMv1 is in fact the scalar gradient C e IRMv1 = ∇ w | w =1 L e ERM ( W ) that we want to force zero, whose effect can be understood as a dynamic re-weighting of the ERM gradient. Due to its importance in the analysis and interpretation of IRMv1, we tracked C e IRMv1 in our simulations.

The invariant and spurious feature learning terms that we tracked are the mean of ⟨ w j,r , j v 1 ⟩ and ⟨ w j,r , j v 2 ⟩ for j ∈ {± 1 } , r ∈ [ m ] , respectively.

D.2 Proof for Theorem 4.1

Theorem D.1 (Formal statement of Theorem 4.1) . For ρ > 0 , denote n ≜ min e ∈E tr n e , n ≜ ∑ e ∈E tr n e , ϵ C ≜ √ 2 log (16 /ρ ) n and δ ≜ exp { O ( n -1 ) } -1 . Define the feature learning terms Λ t j,r ≜ ⟨ w t j,r , j v 1 ⟩ and Γ t j,r ≜ ⟨ w t j,r , j v 2 ⟩ for j ∈ {± 1 } , r ∈ [ m ] . Suppose we run T iterations of GD for the ERM objective. With sufficiently large n and ψ ( x ) = x , assuming that

and choosing

there exists a constant η , such that for any j ∈ {± 1 } , r ∈ [ m ] , with probability at least 1 -2 ρ , Λ t j,r and Γ t j,r are converging and the increment of the spurious feature Γ t +1 j,r -Γ t j,r is larger than that of the invariant feature Λ t +1 j,r -Λ t j,r at any iteration t ∈ { 0 , . . . , T -1 } .

Proof of Theorem D.1. We begin with checking the feature learning terms in the ERM stage using constant stepsize GD: W t +1 = W t -η · ∇ W L IRMv1 ( W t ) . Note that with ψ ( x ) = x the update rule for each w j,r , ∀ j ∈ { +1 , -1 } , r ∈ [ m ] can be written as

Define the quantities of interest (the feature learning terms): Λ t j,r ≜ ⟨ w t j,r , j v 1 ⟩ , Γ t j,r ≜ ⟨ w t j,r , j v 2 ⟩ , Ξ t,e j,r,i ≜ ⟨ w t j,r , j ξ e i ⟩ . From our data generating procedure (Definition 3.1), we know that the first two coordinates of ξ e i are zero. Thus, we can write down the update rule for each feature learning term as follows.

More explicitly, we can write

Notice that the updates (8), (9) for Λ j,r , Γ j,r are independent of j, r . Denoting

we can conclude that for any j ∈ { +1 , -1 } , r ∈ [ m ] ,

Then, we write the logit ˆ y e i as

Denoting Q e i ≜ Rad ( α ) i ∑ j ∈{± 1 } ∑ m r =1 Λ 0 j,r m + Rad ( β e ) i ∑ j ∈{± 1 } ∑ m r =1 Γ 0 j,r m + y e i · ∑ j ∈{± 1 } ∑ m r =1 Ξ t,e j,r,i m , we have

We need the following concentration lemma to control the scale of Q e i , whose proof is given in Appendix D.2.1.

Lemma D.2. Denote n ≜ min e ∈E tr n e , n ≜ ∑ e ∈E tr n e . For ρ > 0 , if

then with probability at least 1 -ρ , for any e ∈ E tr , i ∈ [ n e ] , it holds that | Q e i | = O ( n -1 ) .

Then ∆ t Λ and ∆ t Γ can be explicitly written as

We are going to analyze the convergences of two sequences { ∆ t Γ +∆ t Λ } and { ∆ t Γ -∆ t Λ } . Notice that

We can further write these two terms as

According to Lemma D.2, for all e ∈ E tr , i ∈ [ n e ] , ρ > 0 , letting δ ≜ exp { O ( n -1 ) } -1 , we have 1 + δ ≥ exp { Q e i } ≥ (1 + δ ) -1 with probability at least 1 -ρ . Let C e jℓ ≜ |{ i | Rad ( α ) i = j, Rad ( β e ) i = ℓ, i ∈ E e }| for any j ∈ {± 1 } , ℓ ∈ {± 1 } , e ∈ E tr , and then define C jℓ ≜ ∑ e ∈E tr C e jℓ n e .

We can upper bound and formulate ∆ t Γ +∆ t Λ and ∆ t Γ -∆ t Λ as

Based on similar arguments, we can also establish lower bounds for these two terms,

The upper and lower bounds (12), (13), (14) and (15) imply that the convergences of { ∆ t Γ +∆ t Λ } and { ∆ t Γ -∆ t Λ } are determined by recursive equations of the form Q t = C 1 -C 2 · exp { η ∑ t -1 k =0 Q k } 1+ C 3 · exp { η ∑ t -1 k =0 Q k } . We first establish that with suitably chosen η , the sequences { ∆ t Γ +∆ t Λ } and { ∆ t Γ -∆ t Λ } are guaranteed to be positive. Observed that for the Q t -type recursive equation, the sign of Q 0 is independent of η , and only determined by the constants C 1 , C 2 , C 3 . At iteration 0 , (14) and (15) give

To proceed, we need the following concentration lemma to control the deviations of the constants C +1+1 , C +1 -1 , C -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 ϵ C ≜ √ 2 log (16 /ρ ) n , with probability at least 1 -ρ , we have

Using Lemma D.3, with probability at least 1 -ρ , the constants C +1+1 , C +1 -1 , C -1+1 and C -1 -1 are close to their expectations.

Based on our assumptions that

it can be verified that with probability at least 1 -2 ρ , ∆ 0 Γ +∆ 0 Λ > 0 , ∆ 0 Γ -∆ 0 Λ > 0 .

Then, at iteration 1 , from (14) and (15), we see that as long as we require

it holds that ∆ 1 Γ + ∆ 1 Λ > 0 , ∆ 1 Γ -∆ 1 Λ > 0 . By recursively applying this argument, we see the requirement for η to ensure that ∆ t Γ +∆ t Λ > 0 and ∆ t Γ -∆ t Λ > 0 for any t ∈ { 0 , . . . , T } is

In other words, for the Q t -type recursive equation, as long as Q 0 ≥ 0 , there always exists a sufficiently small η to guarantee that the whole sequence {Q t } is positive. From now on, we will focus on the case where the two sequences { ∆ t Γ +∆ t Λ } and { ∆ t Γ -∆ t Λ } decrease to an ϵ ∆ > 0 error, i.e., min t ∈{ 0 ,...,T } { ∆ t Γ +∆ t Λ , ∆ t Γ -∆ t Λ } = ϵ ∆ .

Then, we show that the two sequences { ∆ t Γ +∆ t Λ } and { ∆ t Γ -∆ t Λ } decrease monotonically, which thus leads to a more refined upper bound for η at (19). Inspect the upper bounds (12), (13) at iteration t +1 , which can be written as

∆

t

+1

Γ

+∆

t

+1

Λ

Requiring that η > max { 1 ∆ t Γ +∆ t Λ log (1 + δ ) , 1 ∆ t Γ -∆ t Λ log (1 + δ ) } , ∀ t ∈ { 0 , . . . , T } ⇒ η > ϵ -1 ∆ log (1 + δ ) , we have

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 η for the sequences { ∆ t Γ +∆ t Λ } and { ∆ t Γ -∆ t Λ } to be positive and monotonically decreasing:

which uses the upper bounds (12) and (13) at iteration 0 . The constants C +1+1 , C +1 -1 , C -1+1 and C -1 -1 can be substituted using the concentration bounds at (18) to generate an upper bound for η that only involves α, β 1 , β 2 , m, δ, T, ϵ C . Here we omit the precise upper bound for clarity. Note that

the left hand side of (20) approaches 0 if δ → 0 , which means that there exists a constant choice of η in (20) if n is sufficiently large in Lemma D.2 and D.3.

To conclude, in view of (11), the convergences of the sequences { ∆ t Γ +∆ t Λ } and { ∆ t Γ -∆ t Λ } imply that Λ t j,r and Γ t j,r are converging, and the positive sequence { ∆ t Γ -∆ t Λ } indicates that the increment of the spurious feature Γ t +1 j,r -Γ t j,r is larger than that of the invariant feature Λ t +1 j,r -Λ t j,r at any iteration t ∈ { 0 , . . . , T -1 } .

D.2.1 Proof of Lemma D.2

First, we recall some concentration inequalities for sub-Gaussian random variables. Since ξ e i ∼ N (0 , σ 2 p · ( I d -v 1 v ⊤ 1 -v 2 v ⊤ 2 )) , for ( i ′ , e ′ ) = ( i, e ) , using Bernstein's inequality for sub-exponential random variables, we have for sufficiently small a ≥ 0 ,

̸

Moreover, for ξ r ∼ N (0 , σ 2 0 ) (indicating the initial weights w 0 j,r ), the standard Gaussian tail gives

Denote n ≜ ∑ e ∈E tr n e , n ≜ min e ∈E tr n e , by properly choosing a for each tail bound and applying a union bound, we can conclude that for ρ > 0 , with probability at least 1 -ρ , it holds that ∀ i, e, i ′ , e ′ , r ,

We start with bound the growth of Ξ t,e j,r,i . By bounding the update rule (10), with probability at least 1 -ρ , we have

̸

Then, we can bound | Q e i | as

Thus, with sufficient small σ 0 , σ p , i.e.,

we ensured that | Q e i | = O ( n -1 ) .

D.2.2 Proof of Lemma D.3

For e ∈ E tr , using Hoeffding's inequality, it holds that

Considering two environments, using a union bound, we can conclude that

Thus, for ρ > 0 , with probability at least 1 -ρ 4 , we can conclude that

Using the above arguments for other constants C +1 -1 , C -1+1 and C -1 -1 , and applying a union bound, we obtain the desired results.

D.2.3 ERMFeature 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 ℓ 2 loss function in the Neural Tangent Kernel (NTK) region. We briefly summarize their results here: For a two-layer fullyconnected neural network (with fixed second layer { v r } ):

considering the ℓ 2 training loss ℓ 2 (ˆ y, y ) ≜ 1 2 (ˆ y -y ) 2 and the ERM objective L ERM ( W ) = 1 n ∑ n i =1 ℓ 2 ( f FC ( W , x i ) , y i ) , when using GD: W t +1 = W t -η · ∇ L ERM ( W t ) to minimize the ERM objective, the following holds.

Theorem D.4 (Theorem 3.2 of [30]) . Let α nl ∈ (0 , 1 4 ) be a fixed constant, and ψ ( · ) be a smooth (with bounded first and second derivatives) or piece-wise linear activation function. Suppose that n and m satisfy n = Ω( d 1+ α nl ) and m = Ω( d 1+ α nl ) . Suppose that η ≪ d . Then there exists a universal constant c > 0 such that with high probability, for all t = O ( d η log d ) simultaneously, the learned neural network f t FC and the linear model f t lin (defined below) at iteration t are close on average on the training data:

The linear model f lin ( β , x ) = β ⊤ R ( x ) is a linear function of the transformed data R ( x ) = 1 √ d [ ζ x ν ] , where ζ and ν are constants related to ψ ′ 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 f lin ( β , x ) similar to the analysis 4 in Appendix D.2, we obtain similar results as in Theorem D.1 in the early stage of GD training, but with an error of O ( d -Ω( α nl ) ) .

Recall that our CNN model is f ( W , x )= F +1 ( W +1 , x ) -F -1 ( W -1 , x ) , where F +1 ( W +1 , x ) and F -1 ( W -1 , x ) are defined as follows:

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 { v r = ± 1 √ m } and transforming the dataset as

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 a slightly different training dynamic. To fix such gap, we can leverage Proposition 6.4.1 in [29] for the early stage behavior of 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 ℓ 2 loss function, which follows a similar line of proofs and has a simpler dynamic.

D.3 Proof for Theorem 4.2

Theorem D.5 (Restatement of Theorem 4.2) . Consider training a CNN model with the same data as in Theorem 4.1, define

and λ 0 = λ min ( H ∞ ) , where we define

Suppose that activation function is smooth, ψ ′ (0) ≤ β , | ψ ′ ( x ) -ψ ′ ( x ′ ) | < β | x -x ′ | and Lipschitz | ψ (0) | < L , | ψ ( x ) -ψ ( x ′ ) | < L | x -x ′ | . Assume that dimension d = Ω(log( m/δ )) , network width m = Ω(1 /δ ) , regularization factor λ ≥ 1 / ( σ 0 √ |E tr | 3 ) , noise variance σ p = O ( d -2 ) , weight initial scale σ 0 = O ( |E tr | 7 / 2 β 3 L d 1 / 2 m 2 λ 2 0 log(1 /ϵ ) ) , then with probability at least 1 -δ , after training time T = Ω ( log(1 /ϵ ) ηλλ 0 ) , we have:

Before proving Theorem D.5, we first provide some useful lemmas as follows:

Lemma D.6 ([12]) . Suppose that δ > 0 and d = Ω(log(4 n/δ )) . Then with probability at least 1 -δ ,

for all i, i ′ ∈ [ n ] .

Lemma D.7 ([12]) . Suppose that d ≥ Ω(log( mn/δ )) , m = Ω(log(1 /δ )) . Then with probability at least 1 -δ ,

for all r ∈ [ m ] , j ∈ {± 1 } and i ∈ [ n ] .

Lemma D.8. Suppose that δ > 0 and d = Ω(log(4 m/δ )) . Then with probability at least 1 -δ , for all r ∈ [ m ] and j ∈ {-1 , 1 } , we have

Proof of Lemma D.8. By Bernstein's inequality, with probability at least 1 -δ/ (2 m ) we have

Therefore, as long as d = Ω(log(4 m/δ )) , we have

Proof of Theorem D.5. The proof is by induction method. First we show the gradient flow of weights by IRMv1 objective function (5):

where C e IRMv1 = 1 n e ∑ n e i =1 ℓ ′ e i ˆ y e i y e i and v e i = Rad ( α ) i · v 1 + Rad ( β e ) i · v 2 . Note that ℓ ′′ has the opposite sign to ℓ ′ .

Then we look at the dynamics of C e IRMv1 ( t ) according to the gradient flow update rule:

where we define H e,e ′ ( t ) = ∑ j ∑ m r =1 〈 ∂C e IRMv1 ( W ,t ) ∂ w j,r ( t ) , ∂C e ′ IRMv1 ( W ,t ) ∂ w j,r ( t ) 〉 and g e ( t ) = ∑ j = ± 1 ∑ m r =1 〈 ∂C e IRMv1 ( W ,t ) ∂ w j,r ( t ) , ∂L ( W ,t ) ∂ w j,r ( t ) 〉 . Thus H ( t ) is an |E tr | × |E tr | matrix. We can write the dynamics of c ( t ) = [ C 1 IRMv1 ( W , t ) , C 2 IRMv1 ( W , t ) , · · · , C |E tr | IRMv1 ( W , t ) ] in a compact way:

Our next step is to show H ( t ) is stable during training. To proceed with the analysis, we write down the expression for ∂C e IRMv1 ( W ,t ) ∂ w j,r ( t ) ∈ R d :

When we consider non-linear activation function ψ ( x ) , the entry of matrix H ( t ) can be computed as follows:

where we calculate each item as follows:

where we have used R ≜ ∥ w j,r ( t ) ∥ 2 . Besides, inequality (a) results from applying the smoothness property of the activation function and Cauchy-Schwarz inequality; inequality (b) is by smoothness property of the activation function and Cauchy-Schwarz inequality. Besides, we have used | ℓ e i | ≤ 1 for all i ∈ n e and e ∈ E all; inequality (c) is by the fact that ∥ v e i ∥ 2 ≤ 2 for all i ∈ n e and e ∈ E all and Lemma D.8.

Similarly, we calculate the upper bound for I 2 as follows:

Next, we give the upper bound of I 3 :

where we have used γ which is defined as follows:

where inequality (a) is by the Lipschitz property of non-linear activation function and we have used the bound for ℓ ′ i ( t ) + 1 2 :

and we provide the bound of ℓ ′′ i ( t ) -1 4 :

Similarly, we give the upper bound of I 4 :

Together, we obtain the upper bound for ∣ ∣ ∣ H 1 e,e ′ ( t ) -H 1 , ∞ e,e ′ ∣ ∣ ∣ :

Then we calculate the upper bound for the residual terms:

where inequality (a) is by the smoothness property of the activation function and Cauchy-Schwarz inequality, and inequality (b) is by triangle inequality and the fact that ∥ v e i ∥ 2 ≤ 2 for all i ∈ n e and e ∈ E all, and | ℓ ′′ i | ≤ 1 for all i ∈ [ n ] and Lemma D.8. Similarly, we further provide the upper bound of residual terms:

Similarly, we further have that:

Keep going on, we provide the computation results further:

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

Similarly, the next H term can be calculated as follows:

Finally, we have the upper for the last term:

To summarize, we have that,

where we have used σ p = O ( d -2 ) , R = o (1) , and σ 0 = O ( √ R/d ) . Furthermore, we show that the perturbation term in Equation (22) is bounded during training. In particular, we show the complete expression:

Similar to the computation process for matrix H , we adopt a divide and conquer manner:

The techniques used are the same when deriving upper bound for matrix H . Next, we have

The last second term can be calculated as follows:

Finally, we show the upper bound of last term:

In a summary, we have the following inequality:

where we have used σ p = O ( d -2 ) and R = o (1) . With all the bounds at hand, we are ready to have the dynamics for ∥ c ( t ) ∥ 2 2

which requires that ∥ H ( t ) -H ∞ ∥ 2 ≤ λ 0 . This leads to the following inequality:

which leads to the conclusion for R as follows:

Besides, we have the inequality that

Combined with Equation (24), we obtain the condition for λ as follows:

By inequality (23), taking the convergence time T = Ω ( log( σ 0 /ϵ ) ηλλ 0 ) we have that:

∥

c

(

T

)

∥

2

ϵ.

According to the gradient descent for IRMV1 objective function, the evolution of coefficients can be expressed as:

Then we have,

Similarly, we have,

At the time step T , the feature learning satisfies that:

To make sure that γ inv j,r ( T ) = o (1) and γ spu j,r ( T ) = o (1) , we need the following condition:

combined with inequality (24) and inequality (26), we have:

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 ψ ( x ) = x , γ j,r, 1 ( t 1 ) = γ j,r, 1 ( t 1 -1) , and γ j,r, 2 ( t 1 ) = γ j,r, 2 ( t 1 -1) at the end of ERM pre-train t 1 and E tr = { (0 . 25 , 0 . 1) , (0 . 25 , 0 . 2) } . Suppose that δ > 0 , and n > C log(1 /δ ) , with C being a positive constant, then with a high probability at least 1 -δ , we have

Proof of Proposition D.9. According to the gradient descent for IRMV1 objective function, the evolution of coefficients can be expressed as:

where

ℓ

To simplify the notation, we further define

and

. Similarly, we define and

In the limit of n →∞ , we have:

and,

Because Rad( α ) i and Rad( β ) i are random variables, applying Hoeffding's inequality, we have with probability at least 1 -δ ,

Similarly, we can apply the concentration bound to other quantities and obtain the same bound.

By the assumption that γ j,r, 1 ( t 1 ) = γ j,r, 1 ( t 1 -1) and γ j,r, 2 ( t 1 ) = γ j,r, 2 ( t 1 -1) , we have that ∑ e A e 1 ( t 1 ) = ∑ e B e 1 ( t 1 ) = 0 :

Solving the above equations, we have,

where we denote γ ∞ 1 ( t 1 ) ≜ lim n →∞ γ 1 ( t 1 ) and γ ∞ 2 ( t 1 ) ≜ lim n →∞ γ 1 ( t 2 ) , G m = ((1 -A ) + √ ( A -1) 2 +4 A ) / (2 A ) and G b = ((1 -B ) + √ ( B -1) 2 +4 B ) / (2 B ) , with A = α ( β 1 + β 2 ) / ((1 -α )(2 -β 1 -β 2 )) and B = α (2 -β 1 -β 2 ) / ((1 -α ) ∗ ( β 1 + β 2 )) .

By the convexity of function f ( x ) = e x , with a constant C , we have:

Then we know that,

Therefore, we have that:

Then the evolution of coefficients reduces to

Taking the solution of γ j,r, 1 ( t 1 ) , γ j,r, 2 ( t 1 ) and value of α, β 1 , β 2 , we arrive at the conclusion that with a high a probability at least 1 -δ and n > C 1 log(1 /δ ) with C 1 being a positive constant, we have:

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 ) = x , γ j,r, 1 ( t 1 ) = o (1) , and γ j,r, 2 ( t 1 ) = Θ(1) at the end of ERM pre-train t 1 and E tr = { (0 . 25 , 0 . 1) , (0 . 25 , 0 . 2) } . Suppose that δ > 0 , and n > C log(1 /δ ) , with C being a positive constant, then with a high probability at least 1 -δ , we have

Proof of Corollary D.10. Recall that the feature learning update rule:

Taking the value of γ j,r, 1 ( t 1 ) , γ j,r, 2 ( t 1 ) and, we can conclude that:

On the other hand,

Finally, taking the value of environment of ( α, β 1 , β 2 ) = (0 . 25 , 0 . 1 , 0 . 2) , we conclude that with a high a probability at least 1 -δ and n > C 1 log(1 /δ ) with C 1 being a positive constant, we have:

E More Details about iFeAT

As mentioned in Sec. 5.2 that, when the featurizer is implemented as a deep net that have a massive amount of parameters, backpropagating through Algorithm 1 can allocate too much memory for propagating with 2 K -1 batches of data. It is common for many realistic benchmarks such as Camelyon17 and FMoW in wilds benchmark [39] that adopts a DenseNet [31] with 121 layers as 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 a high computational and memory overhead as the round grows. For example, in RxRx1, 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 2.

Algorithm 2 FeAT: F eature A ugmented T raining

F 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.

F.1 More details about COLOREDMNIST experiments

Datasets. In the controlled experiments with COLOREDMNIST, we follow the evaluation settings as previous works [4, 16, 83]. In addition to the original COLOREDMNIST with E tr = { (0 . 25 , 0 . 1) , (0 . 25 , 0 . 2) } (denoted as COLOREDMNIST-025) where spurious features are better correlated with labels, we also incorporate the modified one (denoted as COLOREDMNIST-01) with E tr = { (0 . 1 , 0 . 2) , (0 . 1 , 0 . 25) } where invariant features are better correlated with labels, since both cases can happen at real world.

Architecture and optimization. To 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, we use the Adam [37] optimizer with a learning rate of 1 e -3 and a

weight decay of 1 e -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 1 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 a penalty weight of { 1 e 1 , 1 e 2 , 1 e 3 , 1 e 4 , 1 e 5 } . We adopt the implementations from Zhang et al. [83] for various OOD objectives, including IRMv1 [4],VREx [41],IB-IRM [1],CLOvE [76],IGA [40] and Fishr [57] Besides, we also incorporate the state-of-the-art OOD objective proposed by Chen et al. [16] that is able to resolve both COLOREDMNIST-025 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], 5 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-025. While in COLOREDMNIST-01, round 1 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 COLOREDMNIST-025, 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-025 as we empirically find more epochs will incur severe performance degeneration in COLOREDMNIST-025. While in COLOREDMNIST-01, we train the 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

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 in Table. 1. In contrast, FeAT significantly and consistently learns richer features given both COLOREDMNIST-025 and COLOREDMNIST-01 than ERM, which shows the superiority of FeAT.

The termination check in FeAT. Akey difference between FeAT and previous rich feature learning algorithms is that FeAT is 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 in various subsets of COLOREDMNIST025 and COLOREDMNIST-01 at different rounds. We use a termination accuracy of 130% , which trades off the exploration (i.e., training accuracy as 80% ) and the retention (i.e., retention accuracy as 50% ) properly. As shown in Table 6, in COLOREDMNIST-025 (COLOREDMNIST-01), after FeAT learns sufficiently good features at Round 2 ( 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% .

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 the feature learning performance given data from realistic scenarios, we select 6 challenging datasets from WILDS [39] benchmark. The datasets contain various realistic distribution shifts, ranging from domain distribution shifts, subpopulation shifts and the their mixed. A summary of 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].

Amazon. We follow the WILDS splits and data processing pipeline for the Amazon dataset [51]. It 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 d are defined according to the reviewer/customer who wrote the product reviews. The evaluation metric used for the task is 10 th percentile of per-user accuracies in the OOD test sets, and the backbone model is a DistilBert [66], following the WILDS protocol [39].

Camelyon17. We follow the WILDS splits and data processing pipeline for the Camelyon17 dataset [6]. It provides 450 , 000 lymph-node scans from 5 hospitals. The task in Camelyon17 is to take the input of 96 × 96 medical images to predict whether there exists a tumor tissue in the image. The domains d refers to the 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 a DenseNet-121 [31] as the backbone for the featurizer.

CivilComments. We follow the WILDS splits and data processing pipeline for the CivilComments dataset [9]. It 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, LGBTQ, 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]. It 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 is split according to 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 FMoWis the worst region accuracy and the backbone model for the featurizer is a 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 F1 as the evaluation metric and a 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 RxRx1 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 as the evaluation metric and a ResNet-50 [27] as the backbone for the featurizer.

F.2.2 Training and evaluation details.

We follow previous works to implement and evaluate different methods used in our experiments [39]. The information of the referred paper and code is listed as in Table. 9.

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, 31, 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 .

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].

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 { 1 e -2 , 1 e -1 , 1 , 1 e 1 , 1 e 2 } , 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 0 , 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], we 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 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 a 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 a relatively large distance from the desired featurizer in the task. We did not train the model for much too long epochs as Izmailov et al. [33] find that it 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 { 1 e -2 , 1 e -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.

For ERM, we train the model simply by the overall number of steps, except for RxRx1 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 2 K -1 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 ).

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 in FeAT mainly lie in the FeAT training and partitioning. At each training step, FeAT needs ( k -1) additional forward and backward propagation, the same as Bonsai, while FeAT 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.

We calculated the computational overhead: 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, 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-025, as shown in Fig. 4 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, FeAT effectively captures the meaningful features for all samples and generalizes to OOD examples well.

We also visualize the saliency maps of ERM, Bonsai, and FeAT on all real-world datasets used in our work with https://github.com/pytorch/captum . The visualizations are shown as in Fig. 5 to Fig. 11, for which the labels and the predictions of different algorithms are given in 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.

pred:0 label:0

pred:1 label:1

pred:1 label:1

pred:0 label:0

pred:1 label:1

pred:0 label:0

pred:1 label:1

pred:0 label:0

pred:0 label:1

pred:1 label:0

pred:0 label:0

pred:1 label:0

pred:0 label:1

pred:1 label:0

pred:1 label:1

pred:1 label:0

pred:0 label:0

pred:1 label:1

pred:1 label:1

pred:0 label:0

pred:1 label:1

pred:0 label:0

pred:1 label:1

pred:0 label:0

pred:0 label:1

pred:1 label:0

pred:1 label:0

pred:1 label:0

pred:0 label:1

pred:1 label:0

pred:0 label:1

pred:1 label:0

pred:0 label:0

pred:1 label:1

pred:1 label:1

pred:0 label:0

pred:1 label:1

pred:0 label:0

pred:1 label:1

pred:0 label:0

pred:0 label:1

pred:1 label:0

pred:0 label:0

pred:1 label:0

pred:0 label:1

pred:1 label:0

pred:0 label:1

pred:1 label:0

pred:0 label:0

pred:1 label:1

pred:1 label:1

pred:0 label:0

pred:1 label:1

pred:0 label:0

pred:1 label:1

pred:0 label:0

pred:1 label:1

pred:0 label:0

pred:0 label:0

pred:0 label:0

pred:0 label:1

pred:0 label:0

pred:1 label:1

pred:0 label:0

Legend:

Method

Word Importance

ERM [CLS] " adolf hilt ##er was actually left wing . the nazi party was very socialist . " if you knew anything about the history of nazi germany - - persecution of jews , intellectuals , communists , social democrats , etc - - you ' d realize how ridiculous your statements are . [SEP]

ERM [CLS] no one can force anyone to not say what they want , free speech lives get over it , , what is not toll ##erated is a call to harm people , that ##s all , and  fl ying some lunatic nazi  fl ag as stupid and disgusting as it is , , is not doing harm . . [SEP]

ERM [CLS] the catholic church is o ffi cially part of # radical ##bol ##sh ##ev ##ik ##ter ##ror ##ism communist organizers as a catholic i ' m disgusted by the indo ##ct ##rina ##tion of what is becoming a cult the pope and this mor ##on mc ##el ##roy ignore mass genocide while prop ##ping up the cult of mohammed bent on the genocide of christians [SEP]

ERM [CLS] does the term " cognitive di ##sson ##ance " apply to bailey ? [SEP]

Bonsai [CLS] " adolf hilt ##er was actually left wing . the nazi party was very socialist . " if you knew anything about the history of nazi germany - - persecution of jews , intellectuals , communists , social democrats , etc - - you ' d realize how ridiculous your statements are . [SEP]

Bonsai [CLS] no one can force anyone to not say what they want , free speech lives get over it , , what is not toll ##erated is a call to harm people , that ##s all , and  fl ying some lunatic nazi  fl ag as stupid and disgusting as it is , , is not doing harm . . [SEP]

Bonsai [CLS] the catholic church is o ffi cially part of # radical ##bol ##sh ##ev ##ik ##ter ##ror ##ism communist organizers as a catholic i ' m disgusted by the indo ##ct ##rina ##tion of what is becoming a cult the pope and this mor ##on mc ##el ##roy ignore mass genocide while prop ##ping up the cult of mohammed bent on the genocide of christians [SEP]

Bonsai [CLS] does the term " cognitive di ##sson ##ance " apply to bailey ? [SEP]

FeAT [CLS] " adolf hilt ##er was actually left wing . the nazi party was very socialist . " if you knew anything about the history of nazi germany - - persecution of jews , intellectuals , communists , social democrats , etc - - you ' d realize how ridiculous your statements are . [SEP]

FeAT [CLS] no one can force anyone to not say what they want , free speech lives get over it , , what is not toll ##erated is a call to harm people , that ##s all , and  fl ying some lunatic nazi  fl ag as stupid and disgusting as it is , , is not doing harm . . [SEP]

FeAT [CLS] the catholic church is o ffi cially part of # radical ##bol ##sh ##ev ##ik ##ter ##ror ##ism communist organizers as a catholic i ' m disgusted by the indo ##ct ##rina ##tion of what is becoming a cult the pope and this mor ##on mc ##el ##roy ignore mass genocide while prop ##ping up the cult of mohammed bent on the genocide of christians [SEP]

FeAT [CLS] does the term " cognitive di ##sson ##ance " apply to bailey ? [SEP]

(a) CIVILCOMMENTS

Negative

Neutral

Positive

Method

Word Importance

(a) AMAZON-ERM

Bonsai [CLS] american dream ##z try ##s to be american idol with a political agenda on the bush ad ##mins ##tra ##tion which in itself goes all over the place & never provides entertainment . hugh grant is a actor who pretty much hit his peak well over a decade ago & his take as simon from idol is rather just there . mandy moore is ok nothing to say more than she may be a bit player as a actress at best . this  fi lm takes on so much to be a satire , but it ends up with more questions then any answers . too long & over ##bl ##own . william def ##oe was trying to look like dick cheney but ended up looking more like henry fond ##a . dennis qu ##aid was channel ##ing george bush with mixed results . just a ok  fi lm . [SEP]

Bonsai [CLS] this  fi ts my iphone perfectly , and i already have an inc ##ip ##io silicon sleeve that i keep it in permanently . it ' s a little s ##nu ##g with the sleeve , but easy enough to put in and take out . there are 2 belt loops , that i don ' t use , and one belt clip that i do use . the belt clip seems sturdy enough . when i bought this , it was $ 2 . 98 with no s & h . this certainly isn ' t worth a whole lot more than that , but it ' s a great value for the price . i ' d actually be willing to pay up to $ 7 . 00 for it . it ' s not too bulky , at least not more than a holster should be . the magnet keeps it closed . it stays on my belt just  fi ne , and it ' s sturdy enough to keep the iphone safe for a short drop , especially with the silicon sleeve . it ' s not terribly fashionable , but you can ' t beat the price . [SEP]

[CLS] the mind of till ##ie cole must be quite an interesting place . in her third book of the scarred souls series cole once again explores the deepest parts of human de ##pr ##avi ##ty . it ' ll shock and outrage you , it ' ll make you feel so much more for her characters . ra ##va ##ge is told mainly from the points of views of 194 and z ##oya with a couple chapters from lu ##ka . it is not a stand alone book , in my opinion this series should be read in order . i ' ve been a fan of this series from the beginning . the darkness of human nature till ##ie cole continually portrays never cease ##s to engage me . her stories come from a place of raw honesty and empathy . her writing mature ##s with each book . the story of 194 and z ##oya and their roles in this tangled web of a world cole has created was as consuming as the others .

Bonsai however , for me , something was o ff  in their relationship , in the romance . i missed the connection i felt with the previous books main characters . while cole dug deep into their past and pulled out their  fl aws and vu ##ln ##era ##bilities so well as individuals , an aspect of their chemistry and need for each other just didn ' t quite click . this series consistently moves forward though adding more depth and layers to the ultimate story and goals of all of these characters . and in that aspect , the story really excelled and has me just as intrigued in how it all works out as i ever was and i very much look forward to the next installment of the scarred souls series . ra ##va ##ge left much open to be explored and i ' m con fi dent in cole ' s story telling to take me even deeper into her interesting ##ly de ##pr ##ave ##d world . [SEP]

Bonsai [CLS] just keeps getting better and better ! if you like alpha men , strong heroine ##s and laughter with a dash of seriousness , don ' t miss this series ! ! [SEP]

(b) AMAZON-Bonsai

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.

(a) AMAZON-FeAT

(a) CAMELYON17-ERM

(c) CAMELYON17-FeAT

(d) CAMELYON17-ERM

(f) CAMELYON17-FeAT

(g) CAMELYON17-ERM

(i) CAMELYON17-FeAT

168

(j) CAMELYON17-ERM

(l) CAMELYON17-FeAT

(a) FMOW-ERM

(b) FMOW-Bonsai

(c) FMOW-FeAT

(d) FMOW-ERM

(e) FMOW-Bonsai

(f) FMOW-FeAT

(g) FMOW-ERM

(h) FMOW-Bonsai

(i) FMOW-FeAT

(j) FMOW-ERM

(k) FMOW-Bonsai

(l) FMOW-FeAT

4

S

(a) IWILDCAM-ERM

(b) IWILDCAM-Bonsai

(c) IWILDCAM-FeAT

(d) IWILDCAM-ERM

(e) IWILDCAM-Bonsai

(f) IWILDCAM-FeAT

(g) IWILDCAM-ERM

(i) IWILDCAM-FeAT

(j) IWILDCAM-ERM

(k) IWILDCAM-Bonsai

(l) IWILDCAM-FeAT

-088:08:40AMM

(a) RXRX1-ERM

(c) RXRX1-FeAT

(d) RXRX1-ERM

(f) RXRX1-FeAT

(g) RXRX1-ERM

(i) RXRX1-FeAT

(j) RXRX1-ERM

(l) RXRX1-FeAT

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\n- 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\nQUESTIONS:\n- 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\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-18"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n**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\n\nWEAKNESSES:\n### 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\n\nQUESTIONS:\n**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/)?\n\nLIMITATIONS:\nThe limitations are adequately addressed in my opinion, except for the issues raised in the weakness section.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. Interesting theoretocal findings. \n\n2. Promosing empirical results.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}