Understanding and Improving Feature Learning for Out-of-Distribution Generalization

Summary

The paper studies OOD generalization in the presence of spurious correlation. First, it provides a theoretical analysis showing that during ERM training, both spurious and invariant features are learned but at different rates, which, in turn, influences the performance of the following optimization with OOD objectives, such as IRM. Then, it leverages the theoretical analysis to propose a new feature pre-training method to facilitate learning a more diverse set of features. The empirical evaluation shows that the proposed method improves standard ERM pre-training and other feature pre-training baselines in various settings.

Strengths

The paper provides a good combination of theoretical and empirical studies. Both parts are well-motivated, sound and offer a valuable contribution to the field: - The theoretical results deepen our understanding of how ERM learns different features and how this learning schedule interacts with further OOD training and generalization. - The proposed method, motivated by the theoretical analysis, is an efficient and effective way of improving OOD performance, as shown on several datasets.

Weaknesses

- Discussions provided after each formal statement helps understand their implications. However, the overall clarity and presentation of the paper could be improved to be more friendly for a more general audience. For example: - 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? - It is not entirely clear what “OOD objectives” mean before L122. - 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. - 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.

Questions

- According to the provided theory, pre-training features with ERM for longer seems always beneficial. However, as shown in Fig.1 left and mentioned in experiments (e.g., L306-307), overfitting can occur, and OOD performance decreases with further pre-training. Does your theory also explain this phenomenon? - 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.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

- As mentioned in the weaknesses, the presented proofs are limited to linear models. Additional work is needed to extend these results to the case of non-linear CNN.

Summary

The paper consists of two parts. The first part provides a theoretical analysis of the training dynamics for a simple model and data distribution under ERM and IRM. Specifically, the authors explore the questions of feature learning, when one of the features changes its correlation to the target between environments (spurious feature) and the other does not (core feature). The second part proposes a training method for extracting reach features, and shows promising results across an array of benchmarks.

Strengths

**S1.** The paper provides an extensive theoretical analysis of gradient descent dynamics under ERM and IRMv1 for a specific model and data distribution. **S2.** The paper proposes a novel training method (FAT) which shows strong performance in terms of feature learning

Weaknesses

### W1: Theory The 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. Unfortunately, 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 $f(W, x) = F_{+1}(W_{+1}, x) - F_{-1}(W_{-1}, x)$, where $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. We can then rewrite the model as a linear model $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. In other words, the "CNN" model used by the authors is simply a reparameterized linear model. Moreover, the reparameterization does not really change the gradient dynamics in a non-trivial way, as $\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$. In other words, all of the weights $w_{j, r}$ are updated with $\pm$ the gradient of the linear model divided by $m$. So **understanding the training dynamics in the "CNN" model is equivalent to understanding the training dynamics of a linear model**. To summarize, the authors study the training dynamics of linear model with a logistic loss. The authors do not clarify this connection in the paper, which in my opinion is misleading. This omission also unnecessarily complicates the presentation. I also believe there are several important implications of the fact that the authors analyze the training dynamics of a linear model: - 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. - 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. - 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. Please correct me if I am wrong in the reasoning above! ### W2: Theory $\leftrightarrow$ Methodology The connection between the theory and the proposed method (FAT) is not clear to me. It 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. However, as I mentioned above, the relevance of the theory to feature learning in neural networks is questionable. In particular, the theory does not connect as far as I can tell to any of the details of the method. The method could very easily be presented without the theory. ### W3: Presentation Because 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. The presentation of the method is not very clear. In particular, many of the design decisions are not explained. For example, why do we need to reinitialize the weights of the classifier in each "round"? The datasets $\mathcal{D_i^a}$ are never defined in the text, but used in line 262. Overall, it is quite hard to follow the description of the algorithm, and the intuition behind it. Moreover, the iFAT method which is actually used in all of the experiments is not even described in the paper. I 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. ### W4: Performance Overall, 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. **References** [1] [_Last Layer Re-Training is Sufficient for Robustness to Spurious Correlations_](https://openreview.net/forum?id=Zb6c8A-Fghk); P. Kirichenko, P. Izmailov, A. G. Wilson; ICLR 2023

Questions

**Q1.** In Eq. 5, why is the second term divided by $n_e^2$ and not $n_e$? **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? **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/)?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

1 poor

Contribution

2 fair

Limitations

The limitations are adequately addressed in my opinion, except for the issues raised in the weakness section.

Summary

This work aims to understand and compare feature learning in ERM and certain OOD generalization objectives. Additionally, it proposes an approach to enhance feature learning for improved OOD generalization. First, 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. To 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. The 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.

Strengths

1. This work aims to improve feature learning for better OOD generalization, which has been identified as a limitation of existing OOD objectives in recent research [1]. While the authors acknowledge that the idea to do this is not new [2], the proposed approach FAT is novel and demonstrates effectiveness in enhancing OOD performance. FAT incorporates a termination step to halt further training once the model has acquired sufficiently rich features, and the authors also propose an efficient version, iFAT, to make the approach practical for implementation. 2. 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. 3. In their comparisons, the authors extensively evaluate ERM and Bonsai [2]—the primary contenders for feature learning—as well as various OOD objectives applied on top of these methods, across multiple datasets (as shown in Table 2). The results demonstrate that FAT enhances OOD performance.

Weaknesses

1. The claim that ERM does not learn sufficiently good feature representations, while FAT improves them, requires stronger support and evidence. 2. Some sections of the writing require clarifications to improve understanding and readability.

Questions

1. Improvement in feature learning compared to ERM: - 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). - 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. 2. Clarity: - 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. - The description of the method in the introduction (lines 55-58) is unclear and difficult to comprehend without referencing Algorithm 1. - Concerning the algorithm, the rationale behind Line 16 (returning the average of all linear classifiers) is unclear and could benefit from further explanation. 3. Minor: - Please address any typos, e.g., in the Fig. 1 caption, and inconsistencies in notation, e.g., in Equation (5), Lemma 3.2, etc. - 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. References: [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. [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.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

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

Summary

This paper explores the relationshp between the ERM training and OOD generalization in feature learning. The authors analyze the corresponding learned features by ERM and OOD objectives. To answer the question, they conduct the investigation of feature learning in a two-layer CNN network training with ERM and IRMv1. They adopt the data models proposed by [2, 11], and include features with different correlation degrees to the labels to simulate invariant and spurious features like [26]. The theoretocal results extend [51] from data augmentation to ERM learning process. They find that ERM fails since it learns the spurious features more quickly than invariant features, when spurious correlations are stronger than invariant correlations. However, invariant feature learning also happens with RRM so long as the invariant feature has a non-trivial correlation strength with the labels. Moreovoer, they find that IRMv1 requires sufficiently well-learned features for OOD generalization. Compared with the former workshop version, this conference submission adds the so-called Feature Augmented Training (FAT) to learn features for OOD generalization. The proposed method iteratively augments the model to learn new features while retaining the already learned features. In each round, the retention and augmentation operations are performed distributionally robust optimization on different subsets of the training data that capture distinct features. The experimental results verify the promising OOD generalization performance of the proposed method.

Strengths

1. Interesting theoretocal findings. 2. Promosing empirical results.

Weaknesses

1. The symbol system seems hard to follow. Some superscripts and subscripts are redendunt: L_S in Eq.(2), \ell‘ in Eq.(5), o_d in Thm 4.2. Please double-check these symbols and make them clear. 2. Some symbols are not well-defined, such as the \cdot in the predictor f, little o, big O, \Omega, \Theta. 3. The activation function and the variance share the same symbol. 4. 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. 5. It is better to discuss the computational complexity of the proposed method and compare with the other competitors.

Questions

1. It is better to claim Lma 3.2 in the supplementary material. 2. The details of Fig.2 are imoportant. It should be presented in the main body.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

Yes.