SlowGuess/NeurIPS_2023_OCR huggingface dataset:
[1] 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.
[2] Z. Allen-Zhu and Y. Li. Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. arXiv preprint arXiv:2012.09816, 2020.
[3] 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.
[4] M. Arjovsky, L. Bottou, I. Gulrajani, and D. Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv:1907.02893, 2019.
[5] D. Arpit, H. Wang, Y. Zhou, and C. Xiong. Ensemble of averages: Improving model selection and boosting performance in domain generalization. In Advances in Neural Information Processing Systems, 2022.
[6] P. Bandy, 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. Cetin, E. Halici, H. Jackson, R. Chen, F. Both, J. Franke, H. Kusters-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.
[7] S. Beery, G. V. Horn, and P. Perona. Recognition in terra incognita. In Computer Vision European Conference, Part XVI, pages 472-489, 2018.
[8] S. Beery, E. Cole, and A. Gjoka. The iwildcam 2020 competition dataset. arXiv preprint arXiv:2004.10340, 2020.
[9] 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.
[10] 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.
[11] 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.
[12] 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.
[13] 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.
[85] 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.
[86] 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.
Appendix of FeAT
A Notations 18
B Limitations and Future Directions 19
C Related Work 19
D Proofs for theoretical results 20
D.1 Implementation details of the synthetic CNN experiments 20
D.2 Proof for Theorem 4.1 22
D.3 Proof for Theorem 4.2 30
D.4 Proof for Proposition 4.3 38
D.5 Proof for Corollary 4.4 40
E More Details about iFeAT 42
F More Details about the Experiments 42
F.1 More details about COLOREDMNIST experiments 42
F.2 More details about WILDS experiments 44
F.3 Software and hardware 47
F.4 Computational analysis 47
F.5 Feature learning analysis 47
We use bold-faced letters for vectors and matrices otherwise for scalar. We use  $\| \cdot \| _2$  to denote the Euclidean norm of a vector or the spectral norm of a matrix, while denoting  $\| \cdot \| _F$  as the Frobenius norm of a matrix. For a neural network, we denote  $\psi (x)$  as the activation function. Let  $\mathbf{I}_d$  be the identity matrix with a dimension of  $\mathbb{R}^{d\times d}$ . When comparing two sequences  $\{a_n\}$  and  $\{b_n\}$ , we employ standard asymptotic notations such as  $O(\cdot)$ ,  $o(\cdot)$ ,  $\Omega (\cdot)$ , and  $\Theta (\cdot)$  to describe their limiting behavior. Lastly, sequences of integers are denoted as  $[n] = \{1,2,\dots ,n\}$ .
Table 4: Notations for key concepts involved in this paper
Symbols	Definitions
X = Rn	the input space
Y = R	the label space
Z = Rd	the latent space
m	the hidden dimension
Fj(·)	the j-th filter of the CNN model
Wj	the weights of j-th filter of the CNN model, containing m hidden units wj,r
ψ(·)	the activation function of the CNN model
φ	the featurizer φ: X → Z learns a latent representation for each input example
w	the classifier w: Z → Y
wj	the classifier learned at j-th round
f ∈ F	the predictor f = w ∘ φ: X → Y is composed of a featurizer and classifier when w is linear, f can be simply represented via dot product w · φ
Eall	the set of indices for all environments
Etr	the subset of indices of training environments
e	the index set of a specific environment
Eα	the set of environments following the data model as Def. 3.1, where each is specified as (α,βe)
De, De	the dataset from environment e, containing ne samples {xi, yi} considered as i.i.d. from Pe
D	the overall dataset containing n samples from all environments, D = {De}e∈Eall
Da	the augmentation set, we use Da to denote the augmentation set separated at i-th round
Dr	the retention set, we use Dr to denote the retention set separated at i-th round
G	G = {Gr, Ga} with 2k-1 groups at round k, where Ga = {Di}k-1i=0 is the grouped sets, for new feature augmentation and Gr = {Dr}k-1i=1 is the grouped sets for already learned feature retention
Le	the empirical risk calculated based on De, e.g., square loss or logistic loss
lFeAT	the FeAT objective, including lDi the empirical risk at Di and lDi at Dr
LIRMv1(W)	the IRMv1 loss
l'e	the first order derivative of Le with respect to the i-th sample from environment e
CErMv1	CErMv1 ≈ 1/nc ∑ne i=1ncl(yei,yei) · yei,yei, a useful quantity to analyze IRMv1 dynamics
γinvj,r, γj,r,1	the invariant feature learning quantity in Eq. 6
γspu j,r, γj,r,2	the spurious feature learning quantity in Eq. 6
ρj,r,i(t)	the noise feature learning quantity in Eq. 6
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].
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 assembling 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. Nicoliciou 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  $\psi (x) = x$ , the logit  $\hat{y}_i^e$  (which is a function of  $\mathbf{W}$ ) of sample  $i$  in the environment  $e$  can be explicitly written as
$$
\hat {y} _ {i} ^ {e} = f (\mathbf {W}, \mathbf {x} _ {i} ^ {e}) = F _ {+ 1} (\mathbf {W} _ {+ 1}, \mathbf {x} _ {i} ^ {e}) - F _ {- 1} (\mathbf {W} _ {- 1}, \mathbf {x} _ {i} ^ {e}) = \sum_ {j \in \{\pm 1 \}} \frac {j}{m} \sum_ {r = 1} ^ {m} \left[ \mathbf {w} _ {j, r} ^ {\top} (\mathbf {x} _ {i, 1} ^ {e} + \mathbf {x} _ {i, 2} ^ {e}) \right],
$$
where  $\mathbf{W} \triangleq \{\mathbf{W}_{+1}, \mathbf{W}_{-1}\}$  and  $\mathbf{W}_j \triangleq \begin{bmatrix} \mathbf{w}_{j,1}^\top \\ \vdots \\ \mathbf{w}_{j,m}^\top \end{bmatrix}$  for  $j \in \{\pm 1\}$ . We initialized all the network weights as  $\mathcal{N}(0, \sigma_0^2)$  and we set  $\sigma_0 = 0.01$ .
The test dataset  $(\mathbf{x},y)$  is generated through
$$
\mathbf {x} _ {i, 1} = y _ {i} \cdot \mathbf {v} _ {1} + y _ {i} \cdot \operatorname {R a d} \left(1 - \beta_ {e}\right) \cdot \mathbf {v} _ {2}, \quad \mathbf {x} _ {i, 2} = \boldsymbol {\xi},
$$
where half of the dataset uses  $\mathrm{Rad}(1 - \beta_1)$  and the other half uses  $\mathrm{Rad}(1 - \beta_2)$ . Here  $\pmb{\xi} \sim \mathcal{N}(0, \sigma_p^2 \cdot (\mathbf{I}_d - \mathbf{v}_1\mathbf{v}_1^\top - \mathbf{v}_2\mathbf{v}_2^\top))$  and we chose  $\sigma_p = 0.01$ .
From the definition of IRMv1, we take derivative wrt. the scalar 1 of the logit  $1 \cdot \hat{y}_i^e$ . Thus, for environment  $e$ , the penalty is
$$
\left(\frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} \nabla_ {w | w = 1} \ell \big (y _ {i} ^ {e} (w \cdot \hat {y} _ {i} ^ {e}) \big)\right) ^ {2} = \left(\frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} \ell^ {\prime} \big (y _ {i} ^ {e} \hat {y} _ {i} ^ {e} \big) \cdot y _ {i} ^ {e} \hat {y} _ {i} ^ {e}\right) ^ {2}.
$$
Then, the IRMv1 objective is (we set  $n_1 = n_2 = 2500$  in the simulation)
$$
L _ {\mathrm {I R M v 1}} (\mathbf {W}) = \sum_ {e \in \mathcal {E} _ {t r}} \frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} \ell \left(y _ {i} ^ {e} \hat {y} _ {i} ^ {e}\right) + \lambda \sum_ {e \in \mathcal {E} _ {t r}} \left(\frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} \ell^ {\prime} \left(y _ {i} ^ {e} \hat {y} _ {i} ^ {e}\right) \cdot y _ {i} ^ {e} \hat {y} _ {i} ^ {e}\right) ^ {2}.
$$
We used constant stepsize GD to minimize  $L_{\mathrm{IRMv1}}(\mathbf{W})$ , and we chose  $\lambda = 10^8$  (heavy regularization setup).
Let  $C_{\mathrm{IRMv1}}^{e} \triangleq \frac{1}{n_e} \sum_{i=1}^{n_e} \ell'(y_i^e \hat{y}_i^e) \cdot y_i^e \hat{y}_i^e$ . The gradient of  $L_{\mathrm{IRMv1}}(\mathbf{W})$  with respect to each  $\mathbf{w}_{j,r}$  can be explicitly written as
$$
\begin{array}{l} \nabla_ {\mathbf {w} _ {j, r}} L _ {\mathrm {I R M v 1}} (\mathbf {W}) \\ = \sum_ {e \in \mathcal {E} _ {t r}} \frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} \ell^ {\prime} \left(y _ {i} ^ {e} \hat {y} _ {i} ^ {e}\right) \cdot y _ {i} ^ {e} \cdot \frac {j}{m} \left(\mathbf {x} _ {i, 1} ^ {e} + \mathbf {x} _ {i, 2} ^ {e}\right) \\ + 2 \lambda \sum_ {e \in \mathcal {E} _ {t r}} \frac {C _ {\mathrm {I R M v 1}} ^ {e}}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} \left(\ell^ {\prime \prime} (y _ {i} ^ {e} \hat {y} _ {i} ^ {e}) \cdot \hat {y} _ {i} ^ {e} \cdot \frac {j}{m} (\mathbf {x} _ {i, 1} ^ {e} + \mathbf {x} _ {i, 2} ^ {e}) + \ell^ {\prime} (y _ {i} ^ {e} \hat {y} _ {i} ^ {e}) \cdot y _ {i} ^ {e} \cdot \frac {j}{m} (\mathbf {x} _ {i, 1} ^ {e} + \mathbf {x} _ {i, 2} ^ {e})\right) \\ = \sum_ {e \in \mathcal {E} _ {t r}} \frac {j}{n _ {e} m} \sum_ {i = 1} ^ {n _ {e}} \ell^ {\prime} \left(y _ {i} ^ {e} \hat {y} _ {i} ^ {e}\right) \cdot y _ {i} ^ {e} \cdot \left(\mathbf {x} _ {i, 1} ^ {e} + \mathbf {x} _ {i, 2} ^ {e}\right) \\ + 2 \lambda \sum_ {e \in \mathcal {E} _ {t r}} \frac {j C _ {\mathrm {I R M v 1}} ^ {e}}{n _ {e} m} \sum_ {i = 1} ^ {n _ {e}} \ell^ {\prime \prime} \left(y _ {i} ^ {e} \hat {y} _ {i} ^ {e}\right) \cdot \hat {y} _ {i} ^ {e} \cdot \left(\mathbf {x} _ {i, 1} ^ {e} + \mathbf {x} _ {i, 2} ^ {e}\right) \\ + 2 \lambda \sum_ {e \in \mathcal {E} _ {t r}} \frac {j C _ {\mathrm {I R M v 1}} ^ {e}}{n _ {e} m} \sum_ {i = 1} ^ {n _ {e}} \ell^ {\prime} \left(y _ {i} ^ {e} \hat {y} _ {i} ^ {e}\right) \cdot y _ {i} ^ {e} \cdot \left(\mathbf {x} _ {i, 1} ^ {e} + \mathbf {x} _ {i, 2} ^ {e}\right) \\ = \sum_ {e \in \mathcal {E} _ {t r}} \frac {j (1 + 2 \lambda C _ {\mathrm {I R M v 1}} ^ {e})}{n _ {e} m} \sum_ {i = 1} ^ {n _ {e}} \ell^ {\prime} \left(y _ {i} ^ {e} \hat {y} _ {i} ^ {e}\right) \cdot y _ {i} ^ {e} \cdot \left(\mathbf {x} _ {i, 1} ^ {e} + \mathbf {x} _ {i, 2} ^ {e}\right) \\ + 2 \lambda \sum_ {e \in \mathcal {E} _ {t r}} \frac {j C _ {\mathrm {I R M v 1}} ^ {e}}{n _ {e} m} \sum_ {i = 1} ^ {n _ {e}} \ell^ {\prime \prime} \left(y _ {i} ^ {e} \hat {y} _ {i} ^ {e}\right) \cdot \hat {y} _ {i} ^ {e} \cdot (\mathbf {x} _ {i, 1} ^ {e} + \mathbf {x} _ {i, 2} ^ {e}). \\ \end{array}
$$
Observe that  $C_{\mathrm{IRMv1}}^e$  is in fact the scalar gradient  $C_{\mathrm{IRMv1}}^e = \nabla_{w|w=1} L_{\mathrm{ERM}}^e(\mathbf{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_{\mathrm{IRMv1}}^e$  in our simulations.
The invariant and spurious feature learning terms that we tracked are the mean of  $\langle \mathbf{w}_{j,r}, j\mathbf{v}_1 \rangle$  and  $\langle \mathbf{w}_{j,r}, j\mathbf{v}_2 \rangle$  for  $j \in \{\pm 1\}, r \in [m]$ , respectively.
D.2 Proof for Theorem 4.1
Theorem D.1 (Formal statement of Theorem 4.1). For  $\rho >0$ , denote  $\underline{n}\triangleq \min_{e\in \mathcal{E}_{tr}}n_e$ ,  $n\triangleq \sum_{e\in \mathcal{E}_{tr}}n_e$ ,  $\epsilon_C\triangleq \sqrt{\frac{2\log(16 / \rho)}{\underline{n}}}$  and  $\delta \triangleq \exp \{O(\underline{n}^{-1})\} -1$ . Define the feature learning terms  $\Lambda_{j,r}^{t}\triangleq \langle \mathbf{w}_{j,r}^{t},j\mathbf{v}_{1}\rangle$  and  $\Gamma_{j,r}^{t}\triangleq \langle \mathbf{w}_{j,r}^{t},j\mathbf{v}_{2}\rangle$  for  $j\in \{\pm 1\}$ ,  $r\in [m]$ . Suppose we run  $T$  iterations of  $GD$  for the ERM objective. With sufficiently large  $\underline{n}$  and  $\psi (x) = x$ , assuming that
$$
\alpha , \beta_ {1}, \beta_ {2} <   \frac {1 - \epsilon_ {C} - \delta \left(\frac {1}{4} + \frac {\epsilon_ {C}}{2}\right)}{2} (\alpha , \beta_ {1}, \beta_ {2} \text {a r e s u f f i c i e n t l y s m a l l e r t h a n} \frac {1}{2}),
$$
$$
\alpha > \frac {\beta_ {1} + \beta_ {2}}{2} + \epsilon_ {C} + \frac {\delta (1 + \epsilon_ {C})}{2} (\alpha i s s u f f i c i e n t l y l a r g e r t h a n \frac {\beta_ {1} + \beta_ {2}}{2}),
$$
and choosing
$$
\sigma_ {0} ^ {2} = O \left(\underline {{n}} ^ {- 2} \log^ {- 1} (m / \rho)\right),
$$
$$
\sigma_ {p} ^ {2} = O \left(\min \left\{d ^ {- 1 / 2} \log^ {- 1 / 2} (n m / \rho), T ^ {- 1} \eta^ {- 1} m \left(d + n \sqrt {d \log (n ^ {2} / \rho)}\right) ^ {- 1} \right\}\right),
$$
there exists a constant  $\eta$ , such that for any  $j \in \{\pm 1\}, r \in [m]$ , with probability at least  $1 - 2\rho$ ,  $\Lambda_{j,r}^{t}$  and  $\Gamma_{j,r}^{t}$  are converging and the increment of the spurious feature  $\Gamma_{j,r}^{t+1} - \Gamma_{j,r}^{t}$  is larger than that of the invariant feature  $\Lambda_{j,r}^{t+1} - \Lambda_{j,r}^{t}$  at any iteration  $t \in \{0, \dots, T-1\}$ .
Proof of Theorem D.1. We begin with checking the feature learning terms in the ERM stage using constant stepsize GD:  $\mathbf{W}^{t + 1} = \mathbf{W}^t -\eta \cdot \nabla_{\mathbf{W}}L_{\mathrm{IRMv1}}(\mathbf{W}^t)$ . Note that with  $\psi (x) = x$  the update rule for each  $\mathbf{w}_{j,r},\forall j\in \{+1, - 1\} ,r\in [m]$  can be written as
$$
\begin{array}{l} \mathbf {w} _ {j, r} ^ {t + 1} = \mathbf {w} _ {j, r} ^ {t} - \frac {j \eta}{m} \sum_ {e \in \mathcal {E} _ {t r}} \frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} \ell^ {\prime} (y _ {i} ^ {e} \hat {y} _ {i} ^ {e}) \cdot y _ {i} ^ {e} \cdot (\mathbf {x} _ {i, 1} ^ {e} + \mathbf {x} _ {i, 2} ^ {e}) \\ = \mathbf {w} _ {j, r} ^ {t} - \frac {j \eta}{m} \sum_ {e \in \mathcal {E} _ {t r}} \frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} \ell^ {\prime} \big (y _ {i} ^ {e} \hat {y} _ {i} ^ {e} \big) \cdot (\operatorname {R a d} (\alpha) _ {i} \cdot \mathbf {v} _ {1} + \operatorname {R a d} (\beta_ {e}) _ {i} \cdot \mathbf {v} _ {2} + y _ {i} ^ {e} \pmb {\xi} _ {i} ^ {e}). \\ \end{array}
$$
Define the quantities of interest (the feature learning terms):  $\Lambda_{j,r}^{t}\triangleq \langle \mathbf{w}_{j,r}^{t},j\mathbf{v}_{1}\rangle ,\Gamma_{j,r}^{t}\triangleq$ $\langle \mathbf{w}_{j,r}^{t},j\mathbf{v}_{2}\rangle ,\Xi_{j,r,i}^{t,e}\triangleq \langle \mathbf{w}_{j,r}^{t},j\pmb{\xi}_{i}^{e}\rangle .$  From our data generating procedure (Definition 3.1), we know that the first two coordinates of  $\pmb{\xi}_i^e$  are zero. Thus, we can write down the update rule for each feature learning term as follows.
$$
\Lambda_ {j, r} ^ {t + 1} = \Lambda_ {j, r} ^ {t} - \frac {\eta}{m} \sum_ {e \in \mathcal {E} _ {t r}} \frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} \ell^ {\prime} \left(y _ {i} ^ {e} \hat {y} _ {i} ^ {e}\right) \cdot \operatorname {R a d} (\alpha) _ {i},
$$
$$
\Gamma_ {j, r} ^ {t + 1} = \Gamma_ {j, r} ^ {t} - \frac {\eta}{m} \sum_ {e \in \mathcal {E} _ {t r}} \frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} \ell^ {\prime} \left(y _ {i} ^ {e} \hat {y} _ {i} ^ {e}\right) \cdot \operatorname {R a d} (\beta_ {e}) _ {i},
$$
$$
\Xi_ {j, r, i ^ {\prime}} ^ {t + 1, e ^ {\prime}} = \Xi_ {j, r, i ^ {\prime}} ^ {t, e ^ {\prime}} - \frac {\eta}{m} \sum_ {e \in \mathcal {E} _ {t r}} \frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} \ell^ {\prime} \left(y _ {i} ^ {e} \hat {y} _ {i} ^ {e}\right) \cdot y _ {i} ^ {e} \cdot \langle \pmb {\xi} _ {i} ^ {e}, \pmb {\xi} _ {i ^ {\prime}} ^ {e ^ {\prime}} \rangle .
$$
More explicitly, we can write
$$
\Lambda_ {j, r} ^ {t + 1} = \Lambda_ {j, r} ^ {t} + \frac {\eta}{m} \sum_ {e \in \mathcal {E} _ {t r}} \frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} \frac {\operatorname {R a d} (\alpha) _ {i}}{1 + \exp \left\{y _ {i} ^ {e} \hat {y} _ {i} ^ {e} \right\}}, \tag {8}
$$
$$
\Gamma_ {j, r} ^ {t + 1} = \Gamma_ {j, r} ^ {t} + \frac {\eta}{m} \sum_ {e \in \mathcal {E} _ {t r}} \frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} \frac {\operatorname {R a d} \left(\beta_ {e}\right) _ {i}}{1 + \exp \left\{y _ {i} ^ {e} \hat {y} _ {i} ^ {e} \right\}}, \tag {9}
$$