.tocmtchapter mtchaptersubsection mtappendixnone

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

Yongqiang Chen1, Wei Huang2††footnotemark: , Kaiwen Zhou1††footnotemark:
1The Chinese University of Hong Kong 2RIKEN AIP
<EMAIL_ADDRESS><EMAIL_ADDRESS>
Yatao Bian3, Bo Han4, James Cheng1
3Tencent AI Lab 4Hong Kong Baptist University
<EMAIL_ADDRESS><EMAIL_ADDRESS>
Equal Contribution. Work done during Yongqiang’s internship at Tencent AI Lab.

###### 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 objectives111Code is
available at https://github.com/LFhase/FeAT..

##  1 Introduction

Understanding feature learning in neural networks is crucial to understanding
how they generalize to different data distributions [62, 70, 11, 67, 2, 12].
Deep networks trained with empirical risk minimization (ERM) learn highly
predictive features that generalize surprisingly well to in-distribution (ID)
data [75, 24]. However, ERM also tends to learn _spurious_ features or
shortcuts such as image backgrounds [7, 23, 19, 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
[63, 38, 33]. In addition, when optimizing various penalty terms [61, 40, 53,
41, 55, 1, 76, 69, 57, 85, 15] that aim to regularize ERM to capture the
invariant features (termed as OOD objectives), there also exists a curious
phenomenon that the performance of OOD objectives largely relies on the pre-
training with ERM before applying the OOD objectives [83, 16]. As shown in
Fig. 1, 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].

Figure 1: (a) An illustration of FeAT (top row) compared to ERM (bottom row).
Different colors in samples denote the respective dominant features. As the
original data is dominated by spurious features (blue), ERM tends to learn
more spurious features but limited invariant features (orange). Thus the OOD
training with IRMv1 can only leverage limited invariant features and achieve
limited performance. In contrast, iteratively, FeAT divides
$\mathcal{D_{\mathrm{tr}}}$ into augmentation $D^{a}$ and retention sets
$D^{r}$ that contain features not learned and already learned by the current
model at the round, respectively. In each round, FeAT augments the model with
new features contained in the growing augmentation sets while retaining the
already learned features contained in the retention sets, which will lead the
model to learn richer features for OOD training and obtain a better OOD
performance. Then FeAT augments the model with new features while retaining
the already learned features, which leads to richer features for OOD training
and better OOD performance. (b) OOD Performance vs. the number of ERM pre-
training epochs in ColoredMNIST-025. The performance of various OOD objectives
largely relies on the quality of ERM-learned features. When there exist
underlying useful features poorly learned by ERM, the OOD performance will be
limited. In contrast, FeAT learns richer features with $2$ rounds (or $300$
epochs) and improves the OOD performance.

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 Feature Augmented Training (FeAT), an
iterative strategy to enforce the model to learn richer features. As shown in
Fig. 1, 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 ${\cal D}^{r}$) or not (Augmentation set ${\cal D}^{a}$), by examining
whether the model yields correct (${\cal D}^{r}$) or incorrect (${\cal
D}^{a}$) predictions for samples from the subsets, respectively. Intuitively,
${\cal D}^{a}$ and ${\cal 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 ${\cal D}^{a}$
and _retains_ the already learned features by minimizing ERM losses on all
${\cal 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 [62, 70, 11, 22, 2,
12, 32]. Beyond the empirical probing [65, 26, 28, 21], 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 [78, 86, 12, 32]. 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 [63,
33]. 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, 76, 57]. However, the final OOD performance has a large dependence
on the number of ERM pre-training epochs [83, 16]. 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 [59, 5] 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; $\|\cdot\|_{2}$ to denote the Euclidean norm of a vector or the
spectral norm of a matrix, while $\|\cdot\|_{F}$ for the Frobenius norm of a
matrix. ${\bf I}_{d}$ refers to the identity matrix in $\mathbb{R}^{d\times
d}$. Full details are deferred to Appendix A.

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

######  Definition 3.1.

${\cal D}=\\{{\cal D}_{e}\\}_{e\in{{\mathcal{E}}_{\text{all}}}}$ is composed
of multiple subsets ${\cal D}_{e}$ from different environments
$e\in{{\mathcal{E}}_{\text{all}}}$, where each ${\cal
D}_{e}=\\{({\mathbf{x}}_{i}^{e},y_{i}^{e})\\}_{i=1}^{n_{e}}$ is composed of
i.i.d. samples $({\mathbf{x}}_{i}^{e},y_{i}^{e})\sim\mathbb{P}^{e}$. Each data
$({\mathbf{x}}^{e},y^{e})\in{\cal D}_{e}$ with
${\mathbf{x}}^{e}\in\mathbb{R}^{2d}$ and $y^{e}\in\\{-1,1\\}$ is generated as
follows:

  1. (a)

Sample $y^{e}\in\\{-1,1\\}$ uniformly;

  2. (b)

Given $y^{e}$, each input
${\mathbf{x}}^{e}=[{\mathbf{x}}^{e}_{1},{\mathbf{x}}^{e}_{2}]$ contains a
feature patch ${\mathbf{x}}_{1}$ and a noise patch ${\mathbf{x}}_{2}$, that
are sampled as:

$\displaystyle{\mathbf{x}}_{1}$
$\displaystyle=y\cdot\textrm{Rad}(\alpha)\cdot{\mathbf{v}}_{1}+y\cdot\textrm{Rad}(\beta)\cdot{\mathbf{v}}_{2}\quad{\mathbf{x}}_{2}=\bm{\xi}$

where $\textrm{Rad}(\delta)$ is a random variable taking value $-1$ with
probability $\delta$ and $+1$ with probability $1-\delta$,
${\mathbf{v}}_{1}=[1,0,\ldots 0]^{\top}$ and ${\mathbf{v}}_{2}=[0,1,0,\ldots
0]^{\top}$.

  3. (c)

A noise vector $\bm{\xi}$ is generated from the Gaussian distribution
$\mathcal{N}(\mathbf{0},\sigma_{p}^{2}\cdot(\mathbf{I}_{d}-{\mathbf{v}}_{1}{\mathbf{v}}_{1}^{\top}-{\mathbf{v}}_{2}{\mathbf{v}}_{2}^{\top}))$

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, ${\mathbf{v}}_{1}$ and ${\mathbf{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 ${\mathbf{v}}$. Each
environment is denoted as
${{\mathcal{E}}_{\alpha}}\\!=\\!\\{(\alpha,\beta_{e}):0<\beta_{e}<1\\}$, where
${\mathbf{v}}_{1}$ is the invariant feature as $\alpha$ is fixed while
${\mathbf{v}}_{2}$ is the spurious feature as $\beta_{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 ${\mathbf{x}}_{1}$,
${\mathbf{x}}_{2}$, respectively,222When the environment $e$ is not explicitly
considered, we will omit it for clarity. and the second layer parameters of
the network are fixed as $\frac{1}{m}$ and $-\frac{1}{m}$, respectively. Then
the network can be written as
$f({\mathbf{W}},{\mathbf{x}})\\!=\\!F_{+1}({\mathbf{W}}_{+1},{\mathbf{x}})-F_{-1}({\mathbf{W}}_{-1},{\mathbf{x}})$,
where $F_{+1}({\mathbf{W}}_{+1},{\mathbf{x}})$ and
$F_{-1}({\mathbf{W}}_{-1},{\mathbf{x}})$ are defined as follows:

$\displaystyle
F_{j}({\mathbf{W}}_{j},{\mathbf{x}})=\frac{1}{m}\sum_{r=1}^{m}\left[{\psi}({\mathbf{w}}_{j,r}^{\top}{\mathbf{x}}_{1})+{\psi}({\mathbf{w}}_{j,r}^{\top}{\mathbf{x}}_{2})\right],$
(1)

where ${\psi(x)}$ is the activation function. We assume that all network
weights are initialized as $\mathcal{N}(0,\sigma_{0}^{2})$.

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

$\displaystyle\vspace{-0.05in}{L}({\mathbf{W}})$
$\displaystyle=\sum_{e\in\mathcal{E}_{\mathrm{tr}}}\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}\ell(y^{e}_{i}\cdot
f({\mathbf{W}},{\mathbf{x}}^{e}_{i})),\vspace{-0.05in}$ (2)

where $\ell(z)\\!=\\!\log(1\\!+\\!\exp(-z))$ and $\\{{\cal
D}_{e}\\}_{e\in{{\mathcal{E}}_{\text{tr}}}}\\!=\\!\\{\\{{\mathbf{x}}_{i}^{e},y_{i}^{e}\\}_{i=1}^{n_{e}}\\}_{e\in{{\mathcal{E}}_{\text{tr}}}}$
is the trainset with $\sum_{e\in{{\mathcal{E}}_{\text{tr}}}}n_{e}\\!=\\!n$.

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

$\min_{w,\varphi}\
\sum_{e\in{{\mathcal{E}}_{\text{tr}}}}L_{e}(w\circ\varphi),\text{s.t.}\
w\in\operatorname*{arg\,min}_{\bar{w}:{\mathcal{Z}}\rightarrow{\mathcal{Y}}}L_{e}(\bar{w}\circ\varphi),\
\forall e\in{{\mathcal{E}}_{\text{tr}}}.$ (3)

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

$\min_{\varphi}\sum_{e\in{{\mathcal{E}}_{\text{tr}}}}L_{e}(\varphi)+\lambda|\nabla_{w|w=1}L_{e}(w\cdot\varphi)|^{2}.$
(4)

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

$\displaystyle L_{\text{IRMv1}}({\mathbf{W}})$
$\displaystyle=\sum_{e\in\mathcal{E}_{\mathrm{tr}}}\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}\ell\left(y_{i}^{e}\cdot
f({\mathbf{W}},{\mathbf{x}}_{i}^{e})\right)+\sum_{e\in\mathcal{E}_{\mathrm{tr}}}\frac{\lambda}{n_{e}^{2}}\left(\sum_{i=1}^{n_{e}}{\ell_{i}^{\prime
e}}\cdot y_{i}^{e}\cdot f({\mathbf{W}},{\mathbf{x}}_{i}^{e})\right)^{2},$ (5)

where ${\ell_{i}^{\prime e}}=\ell^{\prime}(y_{i}^{e}\cdot
f({\mathbf{W}},{\mathbf{x}}_{i}^{e}))=-\frac{\exp(-y_{i}^{e}\cdot
f({\mathbf{W}},{\mathbf{x}}_{i}^{e}))}{1+\exp(-y_{i}^{e}\cdot
f({\mathbf{W}},{\mathbf{x}}_{i}^{e}))}$. Due to the complexity of IRMv1, in
the analysis below, we introduce $C_{\text{IRMv1}}^{e}$ for the ease of
expressions. Specifically, we define $C_{\text{IRMv1}}^{e}$ as

$C_{\text{IRMv1}}^{e}\triangleq\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}{\ell^{\prime}\big{(}y^{e}_{i}\hat{y}^{e}_{i}\big{)}\cdot
y^{e}_{i}\hat{y}^{e}_{i}},$

where $\hat{y}^{e}_{i}\triangleq f({\mathbf{W}},{\mathbf{x}}^{e}_{i})$ is the
logit of sample ${\mathbf{x}}_{i}$ from environment $e$. The convergence of
$C_{\text{IRMv1}}$ indicates the convergence of IRMv1 penalty. The following
lemma will be useful in our analysis.

######  Lemma 3.2.

(Cao et al. [12]) Let ${\mathbf{w}}_{j,r}(t)$333We use
${\mathbf{w}}_{j,r}(t)$, ${\mathbf{w}}_{j,r}^{(t)}$ and
${\mathbf{w}}_{j,r}^{t}$ interchangeably. for $j\in\\{+1,-1\\}$ and
$r\in\\{1,2,\ldots,m\\}$ be the convolution filters of the CNN at $t$-th
iteration of gradient descent. Then there exists unique coefficients
$\gamma^{inv}_{j,r}(t),\gamma^{spu}_{j,r}(t)\geq 0$ and $\rho_{j,r,i}(t)$ such
that,

$\displaystyle{\mathbf{w}}_{j,r}(t)$
$\displaystyle={\mathbf{w}}_{j,r}(0)+j\cdot\gamma^{inv}_{j,r}(t)\cdot{\mathbf{v}}_{1}+j\cdot\gamma^{spu}_{j,r}(t)\cdot{\mathbf{v}}_{2}+\sum_{i=1}^{n}\rho_{j,r,i}(t)\cdot\|\bm{\xi}_{i}\|^{-2}_{2}\cdot\bm{\xi}_{i}.$
(6)

We refer Eq. (6) as the signal-noise decomposition of
${\mathbf{w}}_{j,r}{(t)}$ [12]. We add normalization factor
$\|\bm{\xi}_{i}\|_{2}^{-2}$ in the definition so that
$\rho_{j,r}^{(t)}\approx\langle{\mathbf{w}}_{j,r}^{(t)},\bm{\xi}_{i}\rangle$.
Note that $\|{\mathbf{v}}_{1}\|_{2}=\|{\mathbf{v}}_{2}\|_{2}=1$, the
corresponding normalization factors are thus neglected.
Furthermore,$\gamma^{inv}_{j,r}\approx\langle{\mathbf{w}}_{j,r},{\mathbf{v}}_{1}\rangle$
and
$\gamma^{spu}_{j,r}\approx\langle{\mathbf{w}}_{j,r},{\mathbf{v}}_{2}\rangle$
respectively denote the degrees of invariant and spurious feature learning.

##  4 Theoretical Understanding of Feature Learning in OOD Generalization

(a) $C_{\text{IRMv1}}$, w/ PT

(b) FL, w/ PT

(c) $C_{\text{IRMv1}}$, w/o PT

(d) FL, w/o PT

Figure 2: The convergences of $C_{\text{IRMv1}}$ and feature learning
coefficients (FL) with or without ERM pre-training (PT). The invariant and
spurious feature learning terms 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. The training environments are
$\mathcal{E}_{tr}=\\{(0.25,0.1),(0.25,0.2)\\}$. The black dashed line
indicates the end of pre-training. More details are given in Appendix D.1.

###  4.1 ERM Feature Learning

With the setup in Sec. 3, we first study the feature learning of the ERM
objective. We consider a two training environments setup
$\mathcal{E}_{tr}=\\{(\alpha,\beta_{1}),(\alpha,\beta_{2})\\}$ where the
signal of invariant feature is weaker than the average of spurious signals
(i.e., $\alpha>\frac{\beta_{1}+\beta_{2}}{2}$), which corresponds to Figure 2.
For a precise characterization of the training dynamic, we adopted a minimal
setup where $\psi(x)=x$ in Figure 2 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 $\rho>0$, let
$\underline{n}\triangleq\min_{e\in\mathcal{E}_{tr}}{n_{e}}$. Suppose that we
run $T$ iterations of GD for the ERM objective. With sufficiently large
$\underline{n}$ and $\psi(x)=x$, assuming that (i)
$\alpha,\beta_{1},\beta_{2}<\frac{1}{2}$, and (ii)
$\alpha>\frac{\beta_{1}+\beta_{2}}{2}$, with properly chosen $\sigma_{0}^{2}$
and $\sigma_{p}^{2}$, there exists a constant $\eta$, such that with
probability at least $1-2\rho$, 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\in\\{0,\ldots,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 ERM training 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 [63, 38, 33]. 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 [63, 38, 33]. 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]) [79, 71, 80,
46]. 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
${\mathbf{c}}(t)\triangleq\left[C^{1}_{\text{IRMv1}}({\mathbf{W}},t),C^{2}_{\text{IRMv1}}({\mathbf{W}},t),\cdots,C^{|{{\mathcal{E}}_{\text{tr}}}|}_{\text{IRMv1}}({\mathbf{W}},t)\right],$
and $\lambda_{0}=\lambda_{\min}(\mathbf{H}^{\infty})$, where
$\mathbf{H}^{\infty}_{e,e^{\prime}}\triangleq\frac{1}{2mn_{e}n_{e^{\prime}}}\sum_{i=1}^{n_{e}}\psi^{\prime}(\langle{\mathbf{w}}_{j,r}(0),\mathbf{x}^{e}_{1,i}\rangle)\mathbf{x}^{e\top}_{1,i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}\psi^{\prime}(\langle{\mathbf{w}}_{j,r}(0),{\mathbf{x}}_{1,i^{\prime}}^{e^{\prime}}\rangle){\mathbf{x}}_{1,i^{\prime}}^{e^{\prime}}$.
Suppose that dimension $d=\Omega(\log(m/\delta))$, network width
$m=\Omega(1/\delta)$, regularization factor $\lambda\geq
1/(\sigma_{0}{|\mathcal{E}_{tr}|}^{3/2})$, noise variance
$\sigma_{p}=O(d^{-2})$, weight initial scale
$\sigma_{0}=O(\frac{|\mathcal{E}_{tr}|^{7/2}\beta^{3}L}{d^{1/2}m^{2}\lambda^{2}_{0}\log(1/\epsilon)})$,
then with probability at least $1-\delta$, after training time
$T=\Omega\left(\frac{\log(1/\epsilon)}{\eta\lambda\lambda_{0}}\right)$, we
have $\|{\mathbf{c}}(T)\|_{2}\leq\epsilon,\
{\gamma_{j,r}^{inv}}(T)\\!=\\!{o(1)},\ {\gamma_{j,r}^{spu}}(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
$\left\lvert\langle{\mathbf{w}}_{j,r},{\mathbf{v}}_{1}\rangle\right\rvert=\Omega(1)$,
$\left\lvert\langle{\mathbf{w}}_{j,r},{\mathbf{v}}_{2}\rangle\right\rvert=\Omega(1)$,
and
$\left\lvert\langle{\mathbf{w}}_{j,r},\bm{\xi}\rangle\right\rvert=O(\sigma_{0}\sigma_{p}\sqrt{d})$.
Then, we have the following hold.

######  Proposition 4.3.

Given the same setting as Theorem 4.2, suppose that $\psi(x)=x$,
$\gamma_{j,r}^{inv}(t_{1})=\gamma_{j,r}^{inv}(t_{1}-1)$, and
$\gamma_{j,r}^{spu}(t_{1})=\gamma_{j,r}^{spu}(t_{1}-1)$ at the end of ERM pre-
train $t_{1}$, $\delta>0$, and $n>C\log(1/\delta)$, with $C$ being a positive
constant, then with a high probability at least $1-\delta$, we have
$\sum_{e}C^{e}_{\text{IRMv1}}(t_{1})=0,$
$\gamma_{j,r}^{inv}(t_{1}+1)>\gamma_{j,r}^{inv}(t_{1}),$ and
$\gamma_{j,r}^{spu}(t_{1}+1)<\gamma_{j,r}^{spu}(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 [83, 16].

###  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
$\psi(x)=x$, $\gamma_{j,r}^{inv}(t_{1})=o(1)$, and
$\gamma_{j,r}^{spu}(t_{1})=\Theta(1)$ at the end of ERM pre-training $t_{1}$.
Suppose that $\delta>0$, and $n>C\log(1/\delta)$, with $C$ being a positive
constant, then with a high probability at least $1-\delta$, we have
$\gamma_{j,r}^{inv}(t_{1}+1)<\gamma_{j,r}^{inv}(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 pre-training 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
[59, 5] and rich feature construction (or Bonsai) [83], and other approaches
[81, 58].

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 Feature Augmented Training (FeAT), that directly augment the feature
learning in an iterative manner.

Algorithm 1 FeAT: Feature Augmented Training

1: Input: Training data $\mathcal{D_{\mathrm{tr}}}$; the maximum augmentation
rounds $K$; predictor $f:=w\circ\varphi$; length of inner training epochs $t$;
termination threshold $p$;

2: Initialize groups
$G^{a}\leftarrow{\mathcal{D_{\mathrm{tr}}}},G^{r}\leftarrow\\{\\}$;

3: for $k\in[1,\ldots,K]$ do

4: Randomly initialize $w_{k}$;

5: for $j\in[1,\ldots,t]$ do

6: Obtain $\ell_{\text{FeAT}}$ with $G$ via Eq. 7;

7: Update $w_{k},\varphi$ with $\ell_{\text{FeAT}}$;

8: end for

9: // Early Stop if $f_{k}=w_{k}\circ\varphi$ 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 $\mathcal{D_{\mathrm{tr}}}$ into groups ${\cal D}_{k}^{a},{\cal
D}_{k}^{r}$ according to whether $f_{k}$ classifies the examples in
$\mathcal{D_{\mathrm{tr}}}$ correctly or not;

14: Update groups $G^{a}\leftarrow G^{a}\cup\\{{\cal
D}_{k}^{a}\\},G^{r}\leftarrow G^{r}\cup\\{{\cal D}_{k}^{r}\\}$;

15: end for

16: Synthesize the final classifier
$w\leftarrow\frac{1}{K}\sum_{i=1}^{K}w_{i}$;

17: return $f=w\circ\varphi$;

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 explores 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\circ\varphi$ that consists of a featurizer $\varphi$
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
$2k-1$ groups, where $G^{a}=\\{{\cal D}_{i}^{a}\\}_{i=0}^{k-1}$ is the grouped
sets for new feature augmentation, and $G^{r}=\\{{\cal
D}_{i}^{r}\\}_{i=1}^{k-1}$ is the grouped sets for already learned feature
retention (notice that ${\cal D}_{0}^{r}$ is the empty set), where ${\cal
D}_{i}^{a}$ and ${\cal D}_{i}^{r}$ 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

$\ell_{\text{FeAT}}=\max_{{\cal D}_{i}^{a}\in G^{a}}\ell_{{\cal
D}_{i}^{a}}(w_{k}\circ\varphi)+\lambda\cdot\sum_{{\cal D}_{i}^{r}\in
G^{r}}\ell_{{\cal D}_{i}^{r}}(w_{i}\circ\varphi),$ (7)

where $\ell_{{\cal D}_{i}}(w\circ\varphi)$ refers to the empirical risk of
$w\circ\varphi$ evaluated at the subset ${\cal D}_{i}$, and $\\{w_{i}|1\leq
i\leq 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 $2K-1$ subsets and a
larger memory cost in training the network, which may cause additional storage
burden when $\varphi$ 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 (${\cal D}_{k-1}^{a},{\cal
D}_{k-1}^{r},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

Table 1: OOD performance on ColoredMNIST datasets initialized with different
representations.

| ColoredMNIST-025 | ColoredMNIST-01
---|---|---
| ERM-nf | ERM | Bonsai | FeAT | ERM-nf | ERM | Bonsai | FeAT
ERM |  $17.14$ ($\pm$0.73) |  $12.40$ ($\pm$0.32) |  $11.21$ ($\pm$0.49) |  $\mathbf{17.27}$ ($\pm$2.55) |  $73.06$ ($\pm$0.71) |  $73.75$ ($\pm$0.49) |  $70.95$ ($\pm$0.93) |  $\mathbf{76.05}$ ($\pm$1.45)
IRMv1 |  $67.29$ ($\pm$0.99) |  $59.81$ ($\pm$4.46) |  $70.28$ ($\pm$0.72) |  $\mathbf{70.57}$ ($\pm$0.68) |  $76.89$ ($\pm$3.25) |  $73.84$ ($\pm$0.56) |  $76.71$ ($\pm$4.10) |  $\mathbf{82.33}$ ($\pm$1.77)
V-Rex |  $68.62$ ($\pm$0.73) |  $65.96$ ($\pm$1.29) |  $70.31$ ($\pm$0.66) |  $\mathbf{70.82}$ ($\pm$0.59) |  $83.52$ ($\pm$2.52) |  $81.20$ ($\pm$3.27) |  $82.61$ ($\pm$1.76) |  $\mathbf{84.70}$ ($\pm$0.69)
IRMX |  $67.00$ ($\pm$1.95) |  $64.05$ ($\pm$0.88) |  $70.46$ ($\pm$0.42) |  $\mathbf{70.78}$ ($\pm$0.61) |  $81.61$ ($\pm$1.98) |  $75.97$ ($\pm$0.88) |  $80.28$ ($\pm$1.62) |  $\mathbf{84.34}$ ($\pm$0.97)
IB-IRM |  $56.09$ ($\pm$2.04) |  $59.81$ ($\pm$4.46) |  $70.28$ ($\pm$0.72) |  $\mathbf{70.57}$ ($\pm$0.68) |  $75.81$ ($\pm$0.63) |  $73.84$ ($\pm$0.56) |  $76.71$ ($\pm$4.10) |  $\mathbf{82.33}$ ($\pm$1.77)
CLOvE |  $58.67$ ($\pm$7.69) |  $65.78$ ($\pm$0.00) |  $65.57$ ($\pm$3.02) |  $\mathbf{65.78}$ ($\pm$2.68) |  $75.66$ ($\pm$10.6) |  $74.73$ ($\pm$0.36) |  $72.73$ ($\pm$1.18) |  $\mathbf{75.12}$ ($\pm$1.08)
IGA |  $51.22$ ($\pm$3.67) |  $62.43$ ($\pm$3.06) |  $\mathbf{70.17}$ ($\pm$0.89) |  $67.11$ ($\pm$3.40) |  $74.20$ ($\pm$2.45) |  $73.74$ ($\pm$0.48) |  $74.72$ ($\pm$3.60) |  $\mathbf{83.46}$ ($\pm$2.17)
Fishr |  $69.38$ ($\pm$0.39) |  $67.74$ ($\pm$0.90) |  $68.75$ ($\pm$1.10) |  $\mathbf{70.56}$ ($\pm$0.97) |  $77.29$ ($\pm$1.61) |  $82.23$ ($\pm$1.35) |  $84.19$ ($\pm$0.66) |  $\mathbf{84.26}$ ($\pm$0.93)
Oracle |  $71.97$ ($\pm$0.34) |  $86.55$ ($\pm$0.27)

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
${{\mathcal{E}}_{\text{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
${{\mathcal{E}}_{\text{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-of-the-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$
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 [83, 16].
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.

Table 2: OOD generalization performances on Wilds benchmark.

Init. | Method |  Camelyon17 |  CivilComments |  FMoW |  iWildCam |  Amazon |  RxRx1
---|---|---|---|---|---|---|---
Avg. acc. (%) | Worst acc. (%) | Worst acc. (%) | Macro F1 | 10-th per. acc. (%) | Avg. acc. (%)
ERM |  DFR† |  $95.14$ ($\pm$1.96) |  $\mathbf{77.34}$ ($\pm$0.50) |  $41.96$ ($\pm$1.90) |  $23.15$ ($\pm$0.24) |  $48.00$ ($\pm$0.00) | -
ERM |  DFR-s† | - |  $82.24$ ($\pm$0.13) |  $56.17$ ($\pm$0.62) |  $52.44$ ($\pm$0.34) | - | -
[0.5pt/1pt] Bonsai |  DFR† |  $95.17$ ($\pm$0.18) |  $77.07$ ($\pm$0.85) |  $43.26$ ($\pm$0.82) |  $21.36$ ($\pm$0.41) |  $46.67$ ($\pm$0.00) | -
Bonsai |  DFR-s† | - |  $81.26$ ($\pm$1.86) |  $58.58$ ($\pm$1.17) |  $50.85$ ($\pm$0.18) | - | -
[0.5pt/1pt] FeAT |  DFR† |  $\mathbf{95.28}$ ($\pm$0.19) |  $\mathbf{77.34}$ ($\pm$0.59) |  $\mathbf{43.54}$ ($\pm$1.26) |  $\mathbf{23.54}$ ($\pm$0.52) |  $\mathbf{49.33}$ ($\pm$0.00) | -
FeAT |  DFR-s† | - |  $79.56$ ($\pm$0.38) |  $57.69$ ($\pm$0.78) |  $52.31$ ($\pm$0.38) | - | -
ERM | ERM |  $74.30$ ($\pm$5.96) |  $55.53$ ($\pm$1.78) |  $33.58$ ($\pm$1.02) |  $28.22$ ($\pm$0.78) |  $51.11$ ($\pm$0.63) |  $30.21$ ($\pm$0.09)
ERM | GroupDRO |  $76.09$ ($\pm$6.46) |  $69.50$ ($\pm$0.15) |  $33.03$ ($\pm$0.52) |  $28.51$ ($\pm$0.58) |  $52.00$ ($\pm$0.00) |  $29.99$ ($\pm$0.13)
ERM | IRMv1 |  $75.68$ ($\pm$7.41) |  $68.84$ ($\pm$0.95) |  $33.45$ ($\pm$1.07) |  $28.76$ ($\pm$0.45) |  $52.00$ ($\pm$0.00) |  $30.10$ ($\pm$0.05)
ERM | V-REx |  $71.60$ ($\pm$7.88) |  $69.03$ ($\pm$1.08) |  $33.06$ ($\pm$0.46) |  $28.82$ ($\pm$0.47) |  $52.44$ ($\pm$0.63) |  $29.88$ ($\pm$0.35)
ERM | IRMX |  $73.49$ ($\pm$9.33) |  $68.91$ ($\pm$1.19) |  $33.13$ ($\pm$0.86) |  $28.82$ ($\pm$0.47) |  $52.00$ ($\pm$0.00) |  $30.10$ ($\pm$0.05)
[0.5pt/1pt] Bonsai | ERM |  $73.98$ ($\pm$5.30) |  $63.34$ ($\pm$3.49) |  $31.91$ ($\pm$0.51) |  $28.27$ ($\pm$1.05) |  $48.58$ ($\pm$0.56) |  $24.22$ ($\pm$0.44)
Bonsai | GroupDRO |  $72.82$ ($\pm$5.37) |  $70.23$ ($\pm$1.33) |  $33.12$ ($\pm$1.20) |  $27.16$ ($\pm$1.18) |  $42.67$ ($\pm$1.09) |  $22.95$ ($\pm$0.46)
Bonsai | IRMv1 |  $73.59$ ($\pm$6.16) |  $68.39$ ($\pm$2.01) |  $32.51$ ($\pm$1.23) |  $27.60$ ($\pm$1.57) |  $47.11$ ($\pm$0.63) |  $23.35$ ($\pm$0.43)
Bonsai | V-REx |  $76.39$ ($\pm$5.32) |  $68.67$ ($\pm$1.29) |  $33.17$ ($\pm$1.26) |  $25.81$ ($\pm$0.42) |  $48.00$ ($\pm$0.00) |  $23.34$ ($\pm$0.42)
Bonsai | IRMX |  $64.77$ ($\pm$10.1) |  $69.56$ ($\pm$0.95) |  $32.63$ ($\pm$0.75) |  $27.62$ ($\pm$0.66) |  $46.67$ ($\pm$0.00) |  $23.34$ ($\pm$0.40)
[0.5pt/1pt] FeAT | ERM |  $77.80$ ($\pm$2.48) |  $68.11$ ($\pm$2.27) |  $33.13$ ($\pm$0.78) |  $28.47$ ($\pm$0.67) |  $\mathbf{52.89}$ ($\pm$0.63) |  $\mathbf{30.66}$ ($\pm$0.42)
FeAT | GroupDRO |  $\mathbf{80.41}$ ($\pm$3.30) |  $\mathbf{71.29}$ ($\pm$0.46) |  $33.55$ ($\pm$1.67) |  $28.38$ ($\pm$1.32) |  $52.58$ ($\pm$0.56) |  $29.99$ ($\pm$0.11)
FeAT | IRMv1 |  $77.97$ ($\pm$3.09) |  $70.33$ ($\pm$1.14) |  $\mathbf{34.04}$ ($\pm$0.70) |  $\mathbf{29.66}$ ($\pm$1.52) |  $\mathbf{52.89}$ ($\pm$0.63) |  $29.99$ ($\pm$0.19)
FeAT | V-REx |  $75.12$ ($\pm$6.55) |  $70.97$ ($\pm$1.06) |  $34.00$ ($\pm$0.71) |  $29.48$ ($\pm$1.94) |  $\mathbf{52.89}$ ($\pm$0.63) |  $30.57$ ($\pm$0.53)
FeAT | IRMX |  $76.91$ ($\pm$6.76) |  $71.18$ ($\pm$1.10) |  $33.99$ ($\pm$0.73) |  $29.04$ ($\pm$2.96) |  $\mathbf{52.89}$ ($\pm$0.63) |  $29.92$ ($\pm$0.16)
†DFR/DFR-s use an additional OOD dataset to evaluate invariant and spurious
feature learning, respectively.

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 [39, 69, 83, 16] 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 contrast, 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.

Table 3: Performances at different FeAT rounds. ColoredMNIST-025 | Round-1 | Round-2 | Round-3
---|---|---|---
Training Acc. |  85.08$\pm$ 0.14 |  71.87$\pm$ 0.96 |  84.93$\pm$ 1.26
Retention Acc. | - |  88.11$\pm$ 4.28 |  43.82$\pm$ 0.59
OOD Acc. |  11.08$\pm$ 0.30 |  70.64$\pm$ 0.62 |  10.07$\pm$ 0.26

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 representations 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 ColoredMNIST-025. 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.

(a) ERM 300 epochs

(b) FeAT 2 rounds

Figure 3: GradCAM visualization on ColoredMNIST-025, where the shortcuts are
now concentrated to a colored path at the up left. Three visualizations are
drawn for each sample: the original figure, the gray-colored gradcam, and the
gradcam. It can be found that ERM can not properly capture the desired
features while FeAT can stably capture the desired features.

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

## Acknowledgements

We thank the reviewers for their valuable comments. This work was supported by
the RIKEN Incentive Research Project 100847-202301062011, by CUHK direct grant
4055146. BH was supported by the NSFC Young Scientists Fund No. 62006202, NSFC
General Program No. 62376235, Guangdong Basic and Applied Basic Research
Foundation No. 2022A1515011652, HKBU Faculty Niche Research Areas No. RC-FNRA-
IG/22-23/SCI/04, and Tencent AI Lab Rhino-Bird Gift Fund.

## References

  * Ahuja et al. [2021] 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.
  * Allen-Zhu and Li [2020] Z. Allen-Zhu and Y. Li.  Towards understanding ensemble, knowledge distillation and self-distillation in deep learning.  _arXiv preprint arXiv:2012.09816_ , 2020.
  * Andreassen et al. [2021] 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.
  * Arjovsky et al. [2019] M. Arjovsky, L. Bottou, I. Gulrajani, and D. Lopez-Paz.  Invariant risk minimization.  _arXiv preprint arXiv:1907.02893_ , 2019.
  * Arpit et al. [2022] 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.
  * Bándi et al. [2019] 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.
  * Beery et al. [2018] S. Beery, G. V. Horn, and P. Perona.  Recognition in terra incognita.  In _Computer Vision European Conference, Part XVI_ , pages 472–489, 2018.
  * Beery et al. [2020] S. Beery, E. Cole, and A. Gjoka.  The iwildcam 2020 competition dataset.  _arXiv preprint arXiv:2004.10340_ , 2020.
  * Borkan et al. [2019] 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.
  * Brown et al. [2020] 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.
  * Brutzkus et al. [2018] 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.
  * Cao et al. [2022] 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.
  * Chen et al. [2023a] 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_ , 2023a.
  * Chen et al. [2022a] 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_ , 2022a.
  * Chen et al. [2022b] 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_ , 2022b.
  * Chen et al. [2022c] 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_ , 2022c.
  * Chen et al. [2023b] 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_ , 2023b.
  * Christie et al. [2018] 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.
  * DeGrave et al. [2021] 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.
  * Deng et al. [2023] Y. Deng, Y. Yang, B. Mirzasoleiman, and Q. Gu.  Robust learning with progressive data expansion against spurious correlation.  _arXiv preprint_ , arXiv:2306.04949, 2023.
  * Elhage et al. [2022] 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.
  * Frei et al. [2021] 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.
  * Geirhos et al. [2020] 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.
  * Goodfellow et al. [2016] I. Goodfellow, Y. Bengio, and A. Courville.  _Deep Learning_.  MIT Press, 2016.  http://www.deeplearningbook.org.
  * Gulrajani and Lopez-Paz [2021] I. Gulrajani and D. Lopez-Paz.  In search of lost domain generalization.  In _International Conference on Learning Representations_ , 2021.
  * Gupta et al. [2022] A. Gupta, N. Saunshi, D. Yu, K. Lyu, and S. Arora.  New definitions and evaluations for saliency methods: Staying intrinsic and sound, 2022.
  * He et al. [2016] 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.
  * Hermann and Lampinen [2020] 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.
  * Hu [2021] W. Hu.  _Understanding Deep Learning via Analyzing Dynamics of Gradient Descent_.  Princeton University, 2021.
  * Hu et al. [2020] 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.
  * Huang et al. [2017] 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.
  * Huang et al. [2023] 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.
  * Izmailov et al. [2022] 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.
  * Ji et al. [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.
  * Kalimeris et al. [2019] 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.
  * Kamath et al. [2021] 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.
  * Kingma and Ba [2015] D. P. Kingma and J. Ba.  Adam: A method for stochastic optimization.  In _International Conference on Learning Representations_ , 2015.
  * Kirichenko et al. [2022] 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.
  * Koh et al. [2021] 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.
  * Koyama and Yamaguchi [2020] M. Koyama and S. Yamaguchi.  Out-of-distribution generalization with maximal invariant predictor.  _arXiv preprint arXiv:2008.01883_ , 2020.
  * Krueger et al. [2021] 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.
  * Kumar et al. [2022a] 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_ , 2022a.
  * Kumar et al. [2022b] 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_ , 2022b.
  * Lee et al. [2023] Y. Lee, H. Yao, and C. Finn.  Diversify and disambiguate: Out-of-distribution robustness via disagreement.  In _International Conference on Learning Representations_ , 2023.
  * Lin et al. [2023a] 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_ , 2023a.
  * Lin et al. [2023b] 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_ , 2023b.
  * Miller et al. [2021] 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.
  * Min et al. [2023] R. Min, Z. Qin, L. Shen, and M. Cheng.  Towards stable backdoor purification through feature shift tuning.  _arXiv preprint arXiv:2310.01875_ , 2023.
  * Naganuma et al. [2022] 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.
  * Namkoong and Duchi [2016] 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.
  * Ni et al. [2019] 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.
  * Nicolicioiu et al. [2023] A. M. Nicolicioiu, A. L. Nicolicioiu, B. Alexe, and D. Teney.  Learning diverse features in vision transformers for improved generalization.  2023\.
  * Parascandolo et al. [2021] 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.
  * Paszke et al. [2019] 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.
  * Pezeshki et al. [2021] M. Pezeshki, S. Kaba, Y. Bengio, A. C. Courville, D. Precup, and G. Lajoie.  Gradient starvation: A learning proclivity in neural networks.  In _Advances in Neural Information Processing Systems_ , pages 1256–1272, 2021.
  * Radford et al. [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.
  * Rame et al. [2021] A. Rame, C. Dancette, and M. Cord.  Fishr: Invariant gradient variances for out-of-distribution generalization.  _arXiv preprint arXiv:2109.02934_ , 2021.
  * Ramé et al. [2022] 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.
  * Rame et al. [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.
  * Recht et al. [2019] 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.
  * Rojas-Carulla et al. [2018] 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.
  * Rosenblatt [1957] F. Rosenblatt.  The perceptron - a perceiving and recognizing automaton.  Technical Report 85-460-1, Cornell Aeronautical Laboratory, 1957.
  * Rosenfeld et al. [2022] 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.
  * Sagawa* et al. [2020] S. Sagawa*, P. W. Koh*, T. B. Hashimoto, and P. Liang.  Distributionally robust neural networks.  In _International Conference on Learning Representations_ , 2020.
  * Samek et al. [2019] 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.
  * Sanh et al. [2019] 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.
  * Shah et al. [2020] 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.
  * Shen et al. [2022] R. Shen, S. Bubeck, and S. Gunasekar.  Data augmentation as feature manipulation.  In _International Conference on Machine Learning_ , pages 19773–19808, 2022.
  * Shi et al. [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.
  * Shwartz-Ziv and Tishby [2017] R. Shwartz-Ziv and N. Tishby.  Opening the black box of deep neural networks via information.  _arXiv preprint arXiv:1703.00810_ , 2017.
  * Taori et al. [2020] 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.
  * Taylor et al. [2019] 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.
  * Teney et al. [2022a] 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, 2022a.
  * Teney et al. [2022b] 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_ , 2022b.
  * Vapnik [1991] V. Vapnik.  Principles of risk minimization for learning theory.  In _Advances in Neural Information Processing Systems_ , pages 831–838, 1991.
  * Wald et al. [2021] 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.
  * Wang et al. [2023] 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.
  * Wen and Li [2021] 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.
  * Wenzel et al. [2022] 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.
  * Wortsman et al. [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.
  * Ye et al. [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.
  * Zhang and Bottou [2022] J. Zhang and L. Bottou.  Learning useful representations for shifting tasks and distributions.  _arXiv preprint arXiv:2212.07346_ , 2022.
  * Zhang et al. [2022a] 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, 2022a.
  * Zhang et al. [2022b] 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_ , 2022b.
  * Zhou et al. [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.
  * Zou et al. [2021] 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

.tocmtappendix mtchapternone mtappendixsubsection

###### Contents

  1. 1 Introduction
  2. 2 Related Work
  3. 3 Preliminaries and Problem Definition
  4. 4 Theoretical Understanding of Feature Learning in OOD Generalization
    1. 4.1 ERM Feature Learning
    2. 4.2 IRM Feature Learning
    3. 4.3 Limitations of ERM Feature Learning
  5. 5 Feature Augmentated Training
    1. 5.1 Rich Features for OOD Generalization
    2. 5.2 The FeAT Algorithm
  6. 6 Empirical Study
  7. 7 Conclusions
  8. A Notations
  9. B Limitations and Future Directions
  10. C Related Work
  11. D Proofs for theoretical results
    1. D.1 Implementation details of the synthetic CNN experiments
    2. D.2 Proof for Theorem 4.1
      1. D.2.1 Proof of Lemma D.2
      2. D.2.2 Proof of Lemma D.3
      3. D.2.3 ERM Feature Learning with Non-Linear Activation Functions
    3. D.3 Proof for Theorem 4.2
    4. D.4 Proof for Proposition 4.3
    5. D.5 Proof for Corollary 4.4
  12. E More Details about iFeAT
  13. F More Details about the Experiments
    1. F.1 More details about ColoredMNIST experiments
    2. F.2 More details about Wilds experiments
      1. F.2.1 Dataset description.
      2. F.2.2 Training and evaluation details.
    3. F.3 Software and hardware
    4. F.4 Computational analysis
    5. F.5 Feature learning analysis

##  Appendix A Notations

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 ${\bf 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,\ldots,n\\}$.

Table 4: Notations for key concepts involved in this paper

Symbols | Definitions
---|---
${\mathcal{X}}=\mathbb{R}^{n}$ | the input space
${\mathcal{Y}}=\mathbb{R}$ | the label space
${\mathcal{Z}}=\mathbb{R}^{d}$ | the latent space
$m$ | the hidden dimension
$F_{j}(\cdot)$ | the $j$-th filter of the CNN model
${\mathbf{W}}_{j}$ | the weights of $j$-th filter of the CNN model, containing $m$ hidden units ${\mathbf{w}}_{j,r}$
$\psi(\cdot)$ | the activation function of the CNN model
$\varphi$ | the featurizer $\varphi:{\mathcal{X}}\rightarrow{\mathcal{Z}}$ learns a latent representation for each input example
$w$ | the classifier $w:{\mathcal{Z}}\rightarrow{\mathcal{Y}}$
$w_{j}$ | the classifier learned at $j$-th round
$f\in{\mathcal{F}}$ | the predictor $f=w\circ\varphi:{\mathcal{X}}\rightarrow{\mathcal{Y}}$ is composed of a featurizer and classifier
| when $w$ is linear, $f$ can be simply represented via dot product
$w\cdot\varphi$
${{\mathcal{E}}_{\text{all}}}$ | the set of indices for all environments
${{\mathcal{E}}_{\text{tr}}}$ | the subset of indices of training environments
$e$ | the index set of a specific environment
$\mathcal{E}_{\alpha}$ | the set of environments following the data model as Def. 3.1, where each is specified as $(\alpha,\beta_{e})$
${\cal D}^{e},{\cal D}_{e}$ | the dataset from environment $e$, containing $n_{e}$ samples $\\{{\mathbf{x}}_{i}^{e},y_{i}^{e}\\}$ considered as i.i.d. from $\mathbb{P}^{e}$
${\cal D}$ | the overall dataset containing $n$ samples from all environments, ${\cal D}=\\{{\cal D}^{e}\\}_{e\in{{\mathcal{E}}_{\text{all}}}}$
${\cal D}^{a}$ | the augmentation set, we use ${\cal D}^{a}_{i}$ to denote the augmentation set separated at $i$-th round
${\cal D}^{r}$ | the retention set, we use ${\cal D}^{r}_{i}$ to denote the retention set separated at $i$-th round
$G$ |  $G=\\{G^{r},G^{a}\\}$ with $2k-1$ groups at round $k$, where $G^{a}=\\{{\cal D}_{i}^{a}\\}_{i=0}^{k-1}$ is the grouped sets,
| for new feature augmentation and $G^{r}=\\{{\cal D}_{i}^{r}\\}_{i=1}^{k-1}$
is the grouped sets for already learned feature retention
$L_{e}$ | the empirical risk calculated based on ${\cal D}^{e}$, e.g., square loss or logistic loss
$\ell_{\text{FeAT}}$ | the FeAT objective, including $\ell_{{\cal D}_{i}^{a}}$ the empirical risk at ${\cal D}^{a}_{i}$ and $\ell_{{\cal D}_{i}^{r}}$ at ${\cal D}^{r}_{i}$
$L_{\text{IRMv1}}({\mathbf{W}})$ | the IRMv1 loss
$\ell^{\prime e}$ | the first order derivative of $L_{e}$ with respect to the $i$-th sample from environment $e$
$C_{\text{IRMv1}}^{e}$ |  $C_{\text{IRMv1}}^{e}\triangleq\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}{\ell^{\prime}\big{(}y^{e}_{i}\hat{y}^{e}_{i}\big{)}\cdot y^{e}_{i}\hat{y}^{e}_{i}}$, a useful quantity to analyze IRMv1 dynamics
$\gamma_{j,r}^{inv},\gamma_{j,r,1}$ | the invariant feature learning quantity in Eq. 6
$\gamma_{j,r}^{spu},\gamma_{j,r,2}$ | the spurious feature learning quantity in Eq. 6
$\rho_{j,r,i}(t)$ | the noise feature learning quantity in Eq. 6

##  Appendix 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 [28, 21]. 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 [15, 17, 14, 77, 34].

##  Appendix C Related Work

On Feature Learning and Generalization. Understanding feature learning in deep
networks is crucial to understanding their generalization [62, 70, 11, 22, 2,
12]. Earlier attempts are mostly about empirical probing [65, 26, 28, 21].
Hermann and Lampinen [28], Elhage et al. [21], 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 [78, 86,
12] 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 [63, 33]. 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, 80, 46].

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, 76, 57]. 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 [83, 16]. 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 [59, 5, 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.

##  Appendix 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}^{e}_{i}$ (which
is a function of ${\mathbf{W}}$) of sample $i$ in the environment $e$ can be
explicitly written as

$\hat{y}^{e}_{i}=f({\mathbf{W}},{\mathbf{x}}^{e}_{i})=F_{+1}({\mathbf{W}}_{+1},{\mathbf{x}}^{e}_{i})-F_{-1}({\mathbf{W}}_{-1},{\mathbf{x}}^{e}_{i})=\sum_{j\in\\{\pm
1\\}}{\frac{j}{m}\sum_{r=1}^{m}{\left[{\mathbf{w}}_{j,r}^{\top}({\mathbf{x}}^{e}_{i,1}+{\mathbf{x}}^{e}_{i,2})\right]}},$

where ${\mathbf{W}}\triangleq\\{{\mathbf{W}}_{+1},{\mathbf{W}}_{-1}\\}$ and
${\mathbf{W}}_{j}\triangleq\begin{bmatrix}{\mathbf{w}}^{\top}_{j,1}\\\
\vdots\\\ {\mathbf{w}}^{\top}_{j,m}\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\textrm{Rad}(1-\beta_{e})\cdot{\mathbf{v}}_{2},\
\ \quad{\mathbf{x}}_{i,2}=\bm{\xi},$

where half of the dataset uses $\textrm{Rad}(1-\beta_{1})$ and the other half
uses $\textrm{Rad}(1-\beta_{2})$. Here
$\bm{\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}^{e}_{i}$. Thus, for environment $e$, the penalty is

$\left(\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}{\nabla_{w|w=1}\ell\big{(}y^{e}_{i}(w\cdot\hat{y}^{e}_{i})\big{)}}\right)^{2}=\left(\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}{\ell^{\prime}\big{(}y^{e}_{i}\hat{y}^{e}_{i}\big{)}\cdot
y^{e}_{i}\hat{y}^{e}_{i}}\right)^{2}.$

Then, the IRMv1 objective is (we set $n_{1}=n_{2}=2500$ in the simulation)

$L_{\text{IRMv1}}({\mathbf{W}})=\sum_{e\in\mathcal{E}_{tr}}{\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}{\ell\big{(}y^{e}_{i}\hat{y}^{e}_{i}\big{)}}}+\lambda\sum_{e\in\mathcal{E}_{tr}}{\left(\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}{\ell^{\prime}\big{(}y^{e}_{i}\hat{y}^{e}_{i}\big{)}\cdot
y^{e}_{i}\hat{y}^{e}_{i}}\right)^{2}}.$

We used constant stepsize GD to minimize $L_{\text{IRMv1}}({\mathbf{W}})$, and
we chose $\lambda=10^{8}$ (heavy regularization setup).

Let
$C_{\text{IRMv1}}^{e}\triangleq\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}{\ell^{\prime}\big{(}y^{e}_{i}\hat{y}^{e}_{i}\big{)}\cdot
y^{e}_{i}\hat{y}^{e}_{i}}$. The gradient of $L_{\text{IRMv1}}({\mathbf{W}})$
with respect to each ${\mathbf{w}}_{j,r}$ can be explicitly written as

$\displaystyle\nabla_{{\mathbf{w}}_{j,r}}L_{\text{IRMv1}}({\mathbf{W}})$
$\displaystyle={}$
$\displaystyle\sum_{e\in\mathcal{E}_{tr}}{\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}{\ell^{\prime}\big{(}y^{e}_{i}\hat{y}^{e}_{i}\big{)}\cdot
y^{e}_{i}\cdot\frac{j}{m}({\mathbf{x}}^{e}_{i,1}+{\mathbf{x}}^{e}_{i,2})}}$
$\displaystyle+2\lambda\sum_{e\in\mathcal{E}_{tr}}{\frac{C_{\text{IRMv1}}^{e}}{n_{e}}\sum_{i=1}^{n_{e}}{\Big{(}\ell^{\prime\prime}\big{(}y^{e}_{i}\hat{y}^{e}_{i}\big{)}\cdot\hat{y}^{e}_{i}\cdot\frac{j}{m}({\mathbf{x}}^{e}_{i,1}+{\mathbf{x}}^{e}_{i,2})+\ell^{\prime}\big{(}y^{e}_{i}\hat{y}^{e}_{i}\big{)}\cdot
y^{e}_{i}\cdot\frac{j}{m}({\mathbf{x}}^{e}_{i,1}+{\mathbf{x}}^{e}_{i,2})\Big{)}}}$
$\displaystyle={}$
$\displaystyle\sum_{e\in\mathcal{E}_{tr}}{\frac{j}{n_{e}m}\sum_{i=1}^{n_{e}}{\ell^{\prime}\big{(}y^{e}_{i}\hat{y}^{e}_{i}\big{)}\cdot
y^{e}_{i}\cdot({\mathbf{x}}^{e}_{i,1}+{\mathbf{x}}^{e}_{i,2})}}$
$\displaystyle+2\lambda\sum_{e\in\mathcal{E}_{tr}}{\frac{jC_{\text{IRMv1}}^{e}}{n_{e}m}\sum_{i=1}^{n_{e}}{\ell^{\prime\prime}\big{(}y^{e}_{i}\hat{y}^{e}_{i}\big{)}\cdot\hat{y}^{e}_{i}\cdot({\mathbf{x}}^{e}_{i,1}+{\mathbf{x}}^{e}_{i,2})}}$
$\displaystyle+2\lambda\sum_{e\in\mathcal{E}_{tr}}{\frac{jC_{\text{IRMv1}}^{e}}{n_{e}m}\sum_{i=1}^{n_{e}}{\ell^{\prime}\big{(}y^{e}_{i}\hat{y}^{e}_{i}\big{)}\cdot
y^{e}_{i}\cdot({\mathbf{x}}^{e}_{i,1}+{\mathbf{x}}^{e}_{i,2})}}$
$\displaystyle={}$
$\displaystyle\sum_{e\in\mathcal{E}_{tr}}{\frac{j(1+2\lambda
C_{\text{IRMv1}}^{e})}{n_{e}m}\sum_{i=1}^{n_{e}}{\ell^{\prime}\big{(}y^{e}_{i}\hat{y}^{e}_{i}\big{)}\cdot
y^{e}_{i}\cdot({\mathbf{x}}^{e}_{i,1}+{\mathbf{x}}^{e}_{i,2})}}$
$\displaystyle{}+2\lambda\sum_{e\in\mathcal{E}_{tr}}{\frac{jC_{\text{IRMv1}}^{e}}{n_{e}m}\sum_{i=1}^{n_{e}}{\ell^{\prime\prime}\big{(}y^{e}_{i}\hat{y}^{e}_{i}\big{)}\cdot\hat{y}^{e}_{i}\cdot({\mathbf{x}}^{e}_{i,1}+{\mathbf{x}}^{e}_{i,2})}}.$

Observe that $C_{\text{IRMv1}}^{e}$ is in fact the scalar gradient
$C_{\text{IRMv1}}^{e}=\nabla_{w|w=1}L^{e}_{\text{ERM}}({\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_{\text{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^{t}_{j,r}\triangleq\langle{\mathbf{w}}_{j,r}^{t},j{\mathbf{v}}_{1}\rangle$
and
$\Gamma^{t}_{j,r}\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

$\displaystyle\alpha,\beta_{1},\beta_{2}<\frac{1-\epsilon_{C}-\delta(\frac{1}{4}+\frac{\epsilon_{C}}{2})}{2}$
$\displaystyle\text{($\alpha,\beta_{1},\beta_{2}$ are sufficiently smaller
than $\frac{1}{2}$)},$
$\displaystyle\alpha>\frac{\beta_{1}+\beta_{2}}{2}+\epsilon_{C}+\frac{\delta(1+\epsilon_{C})}{2}$
$\displaystyle\text{($\alpha$ is sufficiently larger than
$\frac{\beta_{1}+\beta_{2}}{2}$)},$

and choosing

$\displaystyle\sigma_{0}^{2}$
$\displaystyle=O\left(\underline{n}^{-2}\log^{-1}{(m/\rho)}\right),$
$\displaystyle\sigma_{p}^{2}$
$\displaystyle=O\left(\min\left\\{d^{-1/2}\log^{-1/2}{(nm/\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^{t}_{j,r}$ and
$\Gamma^{t}_{j,r}$ are converging and the increment of the spurious feature
$\Gamma^{t+1}_{j,r}-\Gamma^{t}_{j,r}$ is larger than that of the invariant
feature $\Lambda^{t+1}_{j,r}-\Lambda^{t}_{j,r}$ at any iteration
$t\in\\{0,\ldots,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_{\text{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

$\displaystyle{\mathbf{w}}^{t+1}_{j,r}$
$\displaystyle={\mathbf{w}}^{t}_{j,r}-\frac{j\eta}{m}\sum_{e\in\mathcal{E}_{tr}}{\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}{\ell^{\prime}\big{(}y^{e}_{i}\hat{y}^{e}_{i}\big{)}\cdot
y^{e}_{i}\cdot({\mathbf{x}}^{e}_{i,1}+{\mathbf{x}}^{e}_{i,2}})}$
$\displaystyle={\mathbf{w}}^{t}_{j,r}-\frac{j\eta}{m}\sum_{e\in\mathcal{E}_{tr}}{\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}{\ell^{\prime}\big{(}y^{e}_{i}\hat{y}^{e}_{i}\big{)}\cdot(\textrm{Rad}(\alpha)_{i}\cdot{\mathbf{v}}_{1}+\textrm{Rad}(\beta_{e})_{i}\cdot{\mathbf{v}}_{2}+y^{e}_{i}\bm{\xi}^{e}_{i})}}.$

Define the quantities of interest (the feature learning terms):
$\Lambda^{t}_{j,r}\triangleq\langle{\mathbf{w}}_{j,r}^{t},j{\mathbf{v}}_{1}\rangle,\Gamma^{t}_{j,r}\triangleq\langle{\mathbf{w}}_{j,r}^{t},j{\mathbf{v}}_{2}\rangle,\Xi^{t,e}_{j,r,i}\triangleq\langle{\mathbf{w}}_{j,r}^{t},j\bm{\xi}^{e}_{i}\rangle$.
From our data generating procedure (Definition 3.1), we know that the first
two coordinates of $\bm{\xi}^{e}_{i}$ are zero. Thus, we can write down the
update rule for each feature learning term as follows.

$\displaystyle\Lambda^{t+1}_{j,r}$
$\displaystyle=\Lambda^{t}_{j,r}-\frac{\eta}{m}\sum_{e\in\mathcal{E}_{tr}}{\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}{\ell^{\prime}\big{(}y^{e}_{i}\hat{y}^{e}_{i}\big{)}\cdot\textrm{Rad}(\alpha)_{i}}},$
$\displaystyle\Gamma^{t+1}_{j,r}$
$\displaystyle=\Gamma^{t}_{j,r}-\frac{\eta}{m}\sum_{e\in\mathcal{E}_{tr}}{\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}{\ell^{\prime}\big{(}y^{e}_{i}\hat{y}^{e}_{i}\big{)}\cdot\textrm{Rad}(\beta_{e})_{i}}},$
$\displaystyle\Xi^{t+1,e^{\prime}}_{j,r,i^{\prime}}$
$\displaystyle=\Xi^{t,e^{\prime}}_{j,r,i^{\prime}}-\frac{\eta}{m}\sum_{e\in\mathcal{E}_{tr}}{\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}{\ell^{\prime}\big{(}y^{e}_{i}\hat{y}^{e}_{i}\big{)}\cdot
y_{i}^{e}\cdot\langle\bm{\xi}^{e}_{i},\bm{\xi}^{e^{\prime}}_{i^{\prime}}\rangle}}.$

More explicitly, we can write

$\displaystyle\Lambda^{t+1}_{j,r}$
$\displaystyle=\Lambda^{t}_{j,r}+\frac{\eta}{m}\sum_{e\in\mathcal{E}_{tr}}{\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}{\frac{\textrm{Rad}(\alpha)_{i}}{1+\exp\\{y^{e}_{i}\hat{y}^{e}_{i}\\}}}},$
(8) $\displaystyle\Gamma^{t+1}_{j,r}$
$\displaystyle=\Gamma^{t}_{j,r}+\frac{\eta}{m}\sum_{e\in\mathcal{E}_{tr}}{\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}{\frac{\textrm{Rad}(\beta_{e})_{i}}{1+\exp\\{y^{e}_{i}\hat{y}^{e}_{i}\\}}}},$
(9) $\displaystyle\Xi^{t+1,e^{\prime}}_{j,r,i^{\prime}}$
$\displaystyle=\Xi^{t,e^{\prime}}_{j,r,i^{\prime}}+\frac{\eta}{m}\sum_{e\in\mathcal{E}_{tr}}{\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}{\frac{y_{i}^{e}\cdot\langle\bm{\xi}^{e}_{i},\bm{\xi}^{e^{\prime}}_{i^{\prime}}\rangle}{1+\exp\\{y^{e}_{i}\hat{y}^{e}_{i}\\}}}}.$
(10)

Notice that the updates (8), (9) for $\Lambda_{j,r},\Gamma_{j,r}$ are
independent of $j,r$. Denoting

$\displaystyle\Delta_{\Lambda}^{t}$
$\displaystyle\triangleq\frac{1}{m}\sum_{e\in\mathcal{E}_{tr}}{\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}{\frac{\textrm{Rad}(\alpha)_{i}}{1+\exp\\{y^{e}_{i}\hat{y}^{e}_{i}\\}}}},$
$\displaystyle\Delta_{\Gamma}^{t}$
$\displaystyle\triangleq\frac{1}{m}\sum_{e\in\mathcal{E}_{tr}}{\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}{\frac{\textrm{Rad}(\beta_{e})_{i}}{1+\exp\\{y^{e}_{i}\hat{y}^{e}_{i}\\}}}},$

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

$\displaystyle\Lambda^{t+1}_{j,r}$
$\displaystyle=\Lambda^{t}_{j,r}+\eta\cdot\Delta_{\Lambda}^{t}=\eta\cdot\sum_{k=0}^{t}{\Delta_{\Lambda}^{k}}+\Lambda^{0}_{j,r},$
(11) $\displaystyle\Gamma^{t+1}_{j,r}$
$\displaystyle=\Gamma^{t}_{j,r}+\eta\cdot\Delta_{\Gamma}^{t}=\eta\cdot\sum_{k=0}^{t}{\Delta_{\Gamma}^{k}}+\Gamma^{0}_{j,r}.$

Then, we write the logit $\hat{y}^{e}_{i}$ as

$\displaystyle\hat{y}^{e}_{i}={}$ $\displaystyle\sum_{j\in\\{\pm
1\\}}{\frac{j}{m}\sum_{r=1}^{m}{\left[\left\langle{\mathbf{w}}^{t}_{j,r},y_{i}^{e}\cdot\textrm{Rad}(\alpha)_{i}\cdot{\mathbf{v}}_{1}+y_{i}^{e}\cdot\textrm{Rad}(\beta_{e})_{i}\cdot{\mathbf{v}}_{2}+{\mathbf{x}}^{e}_{i,2}\right\rangle\right]}}$
$\displaystyle={}$ $\displaystyle\sum_{j\in\\{\pm
1\\}}{\frac{j}{m}\sum_{r=1}^{m}{\left[jy_{i}^{e}\cdot\textrm{Rad}(\alpha)_{i}\cdot\Lambda^{t}_{j,r}+jy_{i}^{e}\cdot\textrm{Rad}(\beta_{e})_{i}\cdot\Gamma^{t}_{j,r}+j\cdot\Xi^{t,e}_{j,r,i}\right]}}$
$\displaystyle={}$ $\displaystyle\sum_{j\in\\{\pm
1\\}}{\frac{1}{m}\sum_{r=1}^{m}{\left[y_{i}^{e}\cdot\textrm{Rad}(\alpha)_{i}\cdot\Lambda^{t}_{j,r}+y_{i}^{e}\cdot\textrm{Rad}(\beta_{e})_{i}\cdot\Gamma^{t}_{j,r}+\Xi^{t,e}_{j,r,i}\right]}}$
$\displaystyle={}$ $\displaystyle
y_{i}^{e}\cdot\textrm{Rad}(\alpha)_{i}\cdot\sum_{j\in\\{\pm
1\\}}{\sum_{r=1}^{m}{\frac{\Lambda^{t}_{j,r}}{m}}}+y_{i}^{e}\cdot\textrm{Rad}(\beta_{e})_{i}\cdot\sum_{j\in\\{\pm
1\\}}{\sum_{r=1}^{m}{\frac{\Gamma^{t}_{j,r}}{m}}}+\sum_{j\in\\{\pm
1\\}}{\sum_{r=1}^{m}{\frac{\Xi^{t,e}_{j,r,i}}{m}}}$ $\displaystyle={}$
$\displaystyle y_{i}^{e}\cdot\textrm{Rad}(\alpha)_{i}\cdot
2\eta\cdot\sum_{k=0}^{t-1}{\Delta_{\Lambda}^{k}}+y_{i}^{e}\cdot\textrm{Rad}(\beta_{e})_{i}\cdot
2\eta\cdot\sum_{k=0}^{t-1}{\Delta_{\Gamma}^{k}}$
$\displaystyle+y_{i}^{e}\cdot\textrm{Rad}(\alpha)_{i}\cdot\sum_{j\in\\{\pm
1\\}}{\sum_{r=1}^{m}{\frac{\Lambda^{0}_{j,r}}{m}}}+y_{i}^{e}\cdot\textrm{Rad}(\beta_{e})_{i}\cdot\sum_{j\in\\{\pm
1\\}}{\sum_{r=1}^{m}{\frac{\Gamma^{0}_{j,r}}{m}}}+\sum_{j\in\\{\pm
1\\}}{\sum_{r=1}^{m}{\frac{\Xi^{t,e}_{j,r,i}}{m}}}.$

Denoting $\mathbb{Q}^{e}_{i}\triangleq\textrm{Rad}(\alpha)_{i}\sum_{j\in\\{\pm
1\\}}{\sum_{r=1}^{m}{\frac{\Lambda^{0}_{j,r}}{m}}}+\textrm{Rad}(\beta_{e})_{i}\sum_{j\in\\{\pm
1\\}}{\sum_{r=1}^{m}{\frac{\Gamma^{0}_{j,r}}{m}}}+y^{e}_{i}\cdot\sum_{j\in\\{\pm
1\\}}{\sum_{r=1}^{m}{\frac{\Xi^{t,e}_{j,r,i}}{m}}}$, we have

$\hat{y}_{i}^{e}=y_{i}^{e}\cdot\left(\textrm{Rad}(\alpha)_{i}\cdot
2\eta\cdot\sum_{k=0}^{t-1}{\Delta_{\Lambda}^{k}}+\textrm{Rad}(\beta_{e})_{i}\cdot
2\eta\cdot\sum_{k=0}^{t-1}{\Delta_{\Gamma}^{k}}+\mathbb{Q}^{e}_{i}\right),$

We need the following concentration lemma to control the scale of
$\mathbb{Q}^{e}_{i}$, whose proof is given in Appendix D.2.1.

######  Lemma D.2.

Denote
$\underline{n}\triangleq\min_{e\in\mathcal{E}_{tr}}{n_{e}},n\triangleq\sum_{e\in\mathcal{E}_{tr}}{n_{e}}$.
For $\rho>0$, if

$\displaystyle\sigma_{0}^{2}$
$\displaystyle=O\left(\underline{n}^{-2}\log^{-1}{(m/\rho)}\right),$
$\displaystyle\sigma_{p}^{2}$
$\displaystyle=O\left(\min\left\\{d^{-1/2}\log^{-1/2}{(nm/\rho)},T^{-1}\eta^{-1}m\left(d+n\sqrt{d\log(n^{2}/\rho)}\right)^{-1}\right\\}\right),$

then with probability at least $1-\rho$, for any
$e\in\mathcal{E}_{tr},i\in[n_{e}]$, it holds that
$\left\lvert\mathbb{Q}_{i}^{e}\right\rvert=O(\underline{n}^{-1})$.

Then $\Delta_{\Lambda}^{t}$ and $\Delta_{\Gamma}^{t}$ can be explicitly
written as

$\displaystyle\ \ \ \ \Delta_{\Lambda}^{t}=$
$\displaystyle\sum_{e\in\mathcal{E}_{tr}}{\frac{1}{n_{e}m}\sum_{i=1}^{n_{e}}{\frac{\textrm{Rad}(\alpha)_{i}}{1+\exp{\left\\{\textrm{Rad}(\alpha)_{i}\cdot
2\eta\cdot\sum_{k=0}^{t-1}{\Delta_{\Lambda}^{k}}\right\\}}\cdot\exp{\left\\{\textrm{Rad}(\beta_{e})_{i}\cdot
2\eta\cdot\sum_{k=0}^{t-1}{\Delta_{\Gamma}^{k}}\right\\}}\cdot\exp{\left\\{\mathbb{Q}^{e}_{i}\right\\}}}}},$
$\displaystyle\ \ \ \ \Delta_{\Gamma}^{t}=$
$\displaystyle\sum_{e\in\mathcal{E}_{tr}}{\frac{1}{n_{e}m}\sum_{i=1}^{n_{e}}{\frac{\textrm{Rad}(\beta_{e})_{i}}{1+\exp{\left\\{\textrm{Rad}(\alpha)_{i}\cdot
2\eta\cdot\sum_{k=0}^{t-1}{\Delta_{\Lambda}^{k}}\right\\}}\cdot\exp{\left\\{\textrm{Rad}(\beta_{e})_{i}\cdot
2\eta\cdot\sum_{k=0}^{t-1}{\Delta_{\Gamma}^{k}}\right\\}}\cdot\exp{\left\\{\mathbb{Q}^{e}_{i}\right\\}}}}}.$

We are going to analyze the convergences of two sequences
$\\{\Delta_{\Gamma}^{t}+\Delta_{\Lambda}^{t}\\}$ and
$\\{\Delta_{\Gamma}^{t}-\Delta_{\Lambda}^{t}\\}$. Notice that

$\displaystyle\ \ \ \ \Delta_{\Gamma}^{t}+\Delta_{\Lambda}^{t}=$
$\displaystyle\sum_{e\in\mathcal{E}_{tr}}{\frac{1}{n_{e}m}\sum_{i=1}^{n_{e}}{\frac{\textrm{Rad}(\beta_{e})_{i}+\textrm{Rad}(\alpha)_{i}}{1+\exp{\left\\{\textrm{Rad}(\alpha)_{i}\cdot
2\eta\cdot\sum_{k=0}^{t-1}{\Delta_{\Lambda}^{k}}\right\\}}\cdot\exp{\left\\{\textrm{Rad}(\beta_{e})_{i}\cdot
2\eta\cdot\sum_{k=0}^{t-1}{\Delta_{\Gamma}^{k}}\right\\}}\cdot\exp{\left\\{\mathbb{Q}^{e}_{i}\right\\}}}}},$
$\displaystyle\ \ \ \ \Delta_{\Gamma}^{t}-\Delta_{\Lambda}^{t}=$
$\displaystyle\sum_{e\in\mathcal{E}_{tr}}{\frac{1}{n_{e}m}\sum_{i=1}^{n_{e}}{\frac{\textrm{Rad}(\beta_{e})_{i}-\textrm{Rad}(\alpha)_{i}}{1+\exp{\left\\{\textrm{Rad}(\alpha)_{i}\cdot
2\eta\cdot\sum_{k=0}^{t-1}{\Delta_{\Lambda}^{k}}\right\\}}\cdot\exp{\left\\{\textrm{Rad}(\beta_{e})_{i}\cdot
2\eta\cdot\sum_{k=0}^{t-1}{\Delta_{\Gamma}^{k}}\right\\}}\cdot\exp{\left\\{\mathbb{Q}^{e}_{i}\right\\}}}}}.$

We can further write these two terms as

$\displaystyle\Delta_{\Gamma}^{t}+\Delta_{\Lambda}^{t}={}$
$\displaystyle\sum_{e\in\mathcal{E}_{tr}}{\frac{2}{n_{e}m}\sum_{\begin{subarray}{c}i\in[n_{e}]\\\
{\text{Rad}}(\beta_{e})_{i}=+1\\\
{\text{Rad}}(\alpha)_{i}=+1\end{subarray}}{\frac{1}{1+\exp{\left\\{2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}+\Delta_{\Lambda}^{k})}\right\\}}\cdot\exp{\left\\{\mathbb{Q}^{e}_{i}\right\\}}}}}$
$\displaystyle-\sum_{e\in\mathcal{E}_{tr}}{\frac{2}{n_{e}m}\sum_{\begin{subarray}{c}i\in[n_{e}]\\\
{\text{Rad}}(\beta_{e})_{i}=-1\\\
{\text{Rad}}(\alpha)_{i}=-1\end{subarray}}{\frac{1}{1+\exp{\left\\{-2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}+\Delta_{\Lambda}^{k})}\right\\}}\cdot\exp{\left\\{\mathbb{Q}^{e}_{i}\right\\}}}}},$
$\displaystyle\Delta_{\Gamma}^{t}-\Delta_{\Lambda}^{t}={}$
$\displaystyle\sum_{e\in\mathcal{E}_{tr}}{\frac{2}{n_{e}m}\sum_{\begin{subarray}{c}i\in[n_{e}]\\\
{\text{Rad}}(\beta_{e})_{i}=+1\\\
{\text{Rad}}(\alpha)_{i}=-1\end{subarray}}{\frac{1}{1+\exp{\left\\{2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}-\Delta_{\Lambda}^{k})}\right\\}}\cdot\exp{\left\\{\mathbb{Q}^{e}_{i}\right\\}}}}}$
$\displaystyle-\sum_{e\in\mathcal{E}_{tr}}{\frac{2}{n_{e}m}\sum_{\begin{subarray}{c}i\in[n_{e}]\\\
{\text{Rad}}(\beta_{e})_{i}=-1\\\
{\text{Rad}}(\alpha)_{i}=+1\end{subarray}}{\frac{1}{1+\exp{\left\\{-2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}-\Delta_{\Lambda}^{k})}\right\\}}\cdot\exp{\left\\{\mathbb{Q}^{e}_{i}\right\\}}}}}.$

According to Lemma D.2, for all $e\in\mathcal{E}_{tr},i\in[n_{e}],\rho>0$,
letting $\delta\triangleq\exp\\{O(\underline{n}^{-1})\\}-1$, we have
$1+\delta\geq\exp{\\{\mathbb{Q}^{e}_{i}\\}}\geq(1+\delta)^{-1}$ with
probability at least $1-\rho$. Let
$C^{e}_{j\ell}\triangleq|\\{i\mid\textrm{Rad}(\alpha)_{i}=j,\textrm{Rad}(\beta_{e})_{i}=\ell,i\in\mathcal{E}_{e}\\}|$
for any $j\in\\{\pm 1\\},\ell\in\\{\pm 1\\},e\in\mathcal{E}_{tr}$, and then
define
$\overline{C}_{j\ell}\triangleq\sum_{e\in\mathcal{E}_{tr}}{\frac{C^{e}_{j\ell}}{n_{e}}}$.
We can upper bound and formulate $\Delta_{\Gamma}^{t}+\Delta_{\Lambda}^{t}$
and $\Delta_{\Gamma}^{t}-\Delta_{\Lambda}^{t}$ as

$\displaystyle\ \ \ \ \Delta_{\Gamma}^{t}+\Delta_{\Lambda}^{t}\leq{}$
$\displaystyle\frac{2}{m}\left(\frac{\overline{C}_{+1+1}}{1+\exp{\left\\{2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}+\Delta_{\Lambda}^{k})}\right\\}}\cdot(1+\delta)^{-1}}-\frac{\overline{C}_{-1-1}}{1+\exp{\left\\{-2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}+\Delta_{\Lambda}^{k})}\right\\}}\cdot(1+\delta)}\right)$
$\displaystyle={}\frac{2}{m}\cdot\frac{\overline{C}_{+1+1}(1+\delta)-\overline{C}_{-1-1}\cdot\exp{\left\\{2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}+\Delta_{\Lambda}^{k})}\right\\}}}{1+\delta+\exp{\left\\{2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}+\Delta_{\Lambda}^{k})}\right\\}}},$
(12) $\displaystyle\ \ \ \ \Delta_{\Gamma}^{t}-\Delta_{\Lambda}^{t}\leq{}$
$\displaystyle\frac{2}{m}\left(\frac{\overline{C}_{-1+1}}{1+\exp{\left\\{2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}-\Delta_{\Lambda}^{k})}\right\\}}\cdot(1+\delta)^{-1}}-\frac{\overline{C}_{+1-1}}{1+\exp{\left\\{-2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}-\Delta_{\Lambda}^{k})}\right\\}}\cdot(1+\delta)}\right)$
$\displaystyle={}\frac{2}{m}\cdot\frac{\overline{C}_{-1+1}(1+\delta)-\overline{C}_{+1-1}\cdot\exp{\left\\{2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}-\Delta_{\Lambda}^{k})}\right\\}}}{1+\delta+\exp{\left\\{2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}-\Delta_{\Lambda}^{k})}\right\\}}}.$
(13)

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

$\displaystyle\Delta_{\Gamma}^{t}+\Delta_{\Lambda}^{t}\geq{}$
$\displaystyle\frac{2}{m}\cdot\frac{\overline{C}_{+1+1}-\overline{C}_{-1-1}(1+\delta)\cdot\exp{\left\\{2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}+\Delta_{\Lambda}^{k})}\right\\}}}{1+\exp{\left\\{2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}+\Delta_{\Lambda}^{k})}\right\\}}\cdot(1+\delta)},$
(14) $\displaystyle\Delta_{\Gamma}^{t}-\Delta_{\Lambda}^{t}\geq{}$
$\displaystyle\frac{2}{m}\cdot\frac{\overline{C}_{-1+1}-\overline{C}_{+1-1}(1+\delta)\cdot\exp{\left\\{2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}-\Delta_{\Lambda}^{k})}\right\\}}}{1+\exp{\left\\{2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}-\Delta_{\Lambda}^{k})}\right\\}\cdot(1+\delta)}}.$
(15)

The upper and lower bounds (12), (13), (14) and (15) imply that the
convergences of $\\{\Delta_{\Gamma}^{t}+\Delta_{\Lambda}^{t}\\}$ and
$\\{\Delta_{\Gamma}^{t}-\Delta_{\Lambda}^{t}\\}$ are determined by recursive
equations of the form
$\mathcal{Q}^{t}=\frac{C_{1}-C_{2}\cdot\exp{\\{\eta\sum_{k=0}^{t-1}{\mathcal{Q}^{k}}\\}}}{1+C_{3}\cdot\exp{\\{\eta\sum_{k=0}^{t-1}{\mathcal{Q}^{k}}\\}}}$.
We first establish that with suitably chosen $\eta$, the sequences
$\\{\Delta_{\Gamma}^{t}+\Delta_{\Lambda}^{t}\\}$ and
$\\{\Delta_{\Gamma}^{t}-\Delta_{\Lambda}^{t}\\}$ are guaranteed to be
positive. Observed that for the $\mathcal{Q}^{t}$-type recursive equation, the
sign of $\mathcal{Q}^{0}$ is independent of $\eta$, and only determined by the
constants $C_{1},C_{2},C_{3}$. At iteration $0$, (14) and (15) give

$\displaystyle\Delta_{\Gamma}^{0}+\Delta_{\Lambda}^{0}\geq{}$
$\displaystyle\frac{2}{m}\cdot\frac{\overline{C}_{+1+1}-\overline{C}_{-1-1}(1+\delta)}{2+\delta},$
(16) $\displaystyle\Delta_{\Gamma}^{0}-\Delta_{\Lambda}^{0}\geq{}$
$\displaystyle\frac{2}{m}\cdot\frac{\overline{C}_{-1+1}-\overline{C}_{+1-1}(1+\delta)}{2+\delta}.$
(17)

To proceed, we need the following concentration lemma to control the
deviations of the constants $\overline{C}_{+1+1}$, $\overline{C}_{+1-1}$,
$\overline{C}_{-1+1}$ and $\overline{C}_{-1-1}$ from their expectations, whose
proof is given in Appendix D.2.2.

######  Lemma D.3.

For $\rho>0$, considering two environments and denoting
$\epsilon_{C}\triangleq\sqrt{\frac{2\log{(16/\rho)}}{\underline{n}}}$, with
probability at least $1-\rho$, we have

$\displaystyle\left\lvert\overline{C}_{+1+1}-(1-\alpha)(2-\beta_{1}-\beta_{2})\right\rvert$
$\displaystyle\leq\epsilon_{C},$ (18)
$\displaystyle\left\lvert\overline{C}_{+1-1}-(1-\alpha)(\beta_{1}+\beta_{2})\right\rvert$
$\displaystyle\leq\epsilon_{C},$
$\displaystyle\left\lvert\overline{C}_{-1+1}-\alpha(2-\beta_{1}-\beta_{2})\right\rvert$
$\displaystyle\leq\epsilon_{C},$
$\displaystyle\left\lvert\overline{C}_{-1-1}-\alpha(\beta_{1}+\beta_{2})\right\rvert$
$\displaystyle\leq\epsilon_{C}.$

Using Lemma D.3, with probability at least $1-\rho$, the constants
$\overline{C}_{+1+1}$, $\overline{C}_{+1-1}$, $\overline{C}_{-1+1}$ and
$\overline{C}_{-1-1}$ are close to their expectations.

Based on our assumptions that

$\displaystyle\alpha,\beta_{1},\beta_{2}<\frac{1-\epsilon_{C}-\delta(\frac{1}{4}+\frac{\epsilon_{C}}{2})}{2}$
$\displaystyle\text{($\alpha,\beta_{1},\beta_{2}$ are sufficiently smaller
than $\frac{1}{2}$)},$
$\displaystyle\alpha>\frac{\beta_{1}+\beta_{2}}{2}+\epsilon_{C}+\frac{\delta(1+\epsilon_{C})}{2}$
$\displaystyle\text{($\alpha$ is sufficiently larger than
$\frac{\beta_{1}+\beta_{2}}{2}$)},$

it can be verified that with probability at least $1-2\rho$,
$\Delta_{\Gamma}^{0}+\Delta_{\Lambda}^{0}>0,\Delta_{\Gamma}^{0}-\Delta_{\Lambda}^{0}>0.$

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

$\eta<\min{\left\\{\frac{1}{2(\Delta_{\Gamma}^{0}+\Delta_{\Lambda}^{0})}\log{\frac{\overline{C}_{+1+1}}{\overline{C}_{-1-1}(1+\delta)}},\frac{1}{2(\Delta_{\Gamma}^{0}-\Delta_{\Lambda}^{0})}\log{\frac{\overline{C}_{-1+1}}{\overline{C}_{+1-1}(1+\delta)}}\right\\}},$

it holds that
$\Delta_{\Gamma}^{1}+\Delta_{\Lambda}^{1}>0,\Delta_{\Gamma}^{1}-\Delta_{\Lambda}^{1}>0$.
By recursively applying this argument, we see the requirement for $\eta$ to
ensure that $\Delta_{\Gamma}^{t}+\Delta_{\Lambda}^{t}>0$ and
$\Delta_{\Gamma}^{t}-\Delta_{\Lambda}^{t}>0$ for any $t\in\\{0,\ldots,T\\}$ is

$\eta<\min{\left\\{\frac{1}{2\sum_{k=0}^{T-1}{(\Delta_{\Gamma}^{k}+\Delta_{\Lambda}^{k})}}\log{\frac{\overline{C}_{+1+1}}{\overline{C}_{-1-1}(1+\delta)}},\frac{1}{2\sum_{k=0}^{T-1}{(\Delta_{\Gamma}^{k}-\Delta_{\Lambda}^{k})}}\log{\frac{\overline{C}_{-1+1}}{\overline{C}_{+1-1}(1+\delta)}}\right\\}}.$
(19)

In other words, for the $\mathcal{Q}^{t}$-type recursive equation, as long as
$\mathcal{Q}^{0}\geq 0$, there always exists a sufficiently small $\eta$ to
guarantee that the whole sequence $\\{\mathcal{Q}^{t}\\}$ is positive. From
now on, we will focus on the case where the two sequences
$\\{\Delta_{\Gamma}^{t}+\Delta_{\Lambda}^{t}\\}$ and
$\\{\Delta_{\Gamma}^{t}-\Delta_{\Lambda}^{t}\\}$ decrease to an
$\epsilon_{\Delta}>0$ error, i.e.,
$\min_{t\in\\{0,\ldots,T\\}}{\\{\Delta_{\Gamma}^{t}+\Delta_{\Lambda}^{t},\Delta_{\Gamma}^{t}-\Delta_{\Lambda}^{t}\\}}=\epsilon_{\Delta}$.

Then, we show that the two sequences
$\\{\Delta_{\Gamma}^{t}+\Delta_{\Lambda}^{t}\\}$ and
$\\{\Delta_{\Gamma}^{t}-\Delta_{\Lambda}^{t}\\}$ decrease monotonically, which
thus leads to a more refined upper bound for $\eta$ at (19). Inspect the upper
bounds (12), (13) at iteration $t+1$, which can be written as

$\displaystyle\ \ \ \ \Delta_{\Gamma}^{t+1}+\Delta_{\Lambda}^{t+1}\leq{}$
$\displaystyle\frac{2}{m}\cdot\frac{\overline{C}_{+1+1}-\overline{C}_{-1-1}\cdot\exp{\left\\{2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}+\Delta_{\Lambda}^{k})}\right\\}}\cdot\exp{\left\\{2\eta\cdot(\Delta_{\Gamma}^{t}+\Delta_{\Lambda}^{t})\right\\}}(1+\delta)^{-1}}{1+\exp{\left\\{2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}+\Delta_{\Lambda}^{k})}\right\\}}\cdot\exp{\left\\{2\eta\cdot(\Delta_{\Gamma}^{t}+\Delta_{\Lambda}^{t})\right\\}}(1+\delta)^{-1}}\triangleq\spadesuit^{t+1},$
$\displaystyle\ \ \ \ \Delta_{\Gamma}^{t+1}-\Delta_{\Lambda}^{t+1}\leq{}$
$\displaystyle\frac{2}{m}\cdot\frac{\overline{C}_{-1+1}-\overline{C}_{+1-1}\cdot\exp{\left\\{2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}-\Delta_{\Lambda}^{k})}\right\\}}\cdot\exp{\left\\{2\eta\cdot(\Delta_{\Gamma}^{t}-\Delta_{\Lambda}^{t})\right\\}}(1+\delta)^{-1}}{1+\exp{\left\\{2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}-\Delta_{\Lambda}^{k})}\right\\}}\cdot\exp{\left\\{2\eta\cdot(\Delta_{\Gamma}^{t}-\Delta_{\Lambda}^{t})\right\\}}(1+\delta)^{-1}}\triangleq\clubsuit^{t+1}.$

Requiring that
$\eta>\max{\left\\{\frac{1}{\Delta_{\Gamma}^{t}+\Delta_{\Lambda}^{t}}\log{(1+\delta)},\frac{1}{\Delta_{\Gamma}^{t}-\Delta_{\Lambda}^{t}}\log{(1+\delta)}\right\\}},\forall
t\in\\{0,\ldots,T\\}\Rightarrow\eta>\epsilon_{\Delta}^{-1}\log{(1+\delta)}$,
we have

$\displaystyle\spadesuit^{t+1}$
$\displaystyle<\frac{2}{m}\cdot\frac{\overline{C}_{+1+1}-\overline{C}_{-1-1}\cdot\exp{\left\\{2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}+\Delta_{\Lambda}^{k})}\right\\}}\cdot\exp{\left\\{2\eta\cdot(\Delta_{\Gamma}^{t}+\Delta_{\Lambda}^{t})\right\\}}(1+\delta)^{-1}}{1+\exp{\left\\{2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}+\Delta_{\Lambda}^{k})}\right\\}}\cdot(1+\delta)}$
$\displaystyle<\Delta_{\Gamma}^{t}+\Delta_{\Lambda}^{t},$
$\displaystyle\clubsuit^{t+1}$
$\displaystyle<\frac{2}{m}\cdot\frac{\overline{C}_{-1+1}-\overline{C}_{+1-1}\cdot\exp{\left\\{2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}-\Delta_{\Lambda}^{k})}\right\\}}\cdot\exp{\left\\{2\eta\cdot(\Delta_{\Gamma}^{t}-\Delta_{\Lambda}^{t})\right\\}}(1+\delta)^{-1}}{1+\exp{\left\\{2\eta\cdot\sum_{k=0}^{t-1}{(\Delta_{\Gamma}^{k}-\Delta_{\Lambda}^{k})}\right\\}}\cdot(1+\delta)}$
$\displaystyle<\Delta_{\Gamma}^{t}-\Delta_{\Lambda}^{t},$

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
$\eta$ for the sequences $\\{\Delta_{\Gamma}^{t}+\Delta_{\Lambda}^{t}\\}$ and
$\\{\Delta_{\Gamma}^{t}-\Delta_{\Lambda}^{t}\\}$ to be positive and
monotonically decreasing:

$\displaystyle\epsilon_{\Delta}^{-1}\log{(1+\delta)}<\eta<\min\bigg{\\{}$
$\displaystyle\frac{m(2+\delta)}{4T(\overline{C}_{+1+1}(1+\delta)-\overline{C}_{-1-1})}\log{\frac{\overline{C}_{+1+1}}{\overline{C}_{-1-1}(1+\delta)}},$
(20)
$\displaystyle\frac{m(2+\delta)}{4T(\overline{C}_{-1+1}(1+\delta)-\overline{C}_{+1-1})}\log{\frac{\overline{C}_{-1+1}}{\overline{C}_{+1-1}(1+\delta)}}\bigg{\\}},$

which uses the upper bounds (12) and (13) at iteration $0$. The constants
$\overline{C}_{+1+1}$, $\overline{C}_{+1-1}$, $\overline{C}_{-1+1}$ and
$\overline{C}_{-1-1}$ can be substituted using the concentration bounds at
(18) to generate an upper bound for $\eta$ that only involves
$\alpha,\beta_{1},\beta_{2},m,\delta,T,\epsilon_{C}$. Here we omit the precise
upper bound for clarity. Note that the left hand side of (20) approaches $0$
if $\delta\rightarrow 0$, which means that there exists a constant choice of
$\eta$ in (20) if $\underline{n}$ is sufficiently large in Lemma D.2 and D.3.

To conclude, in view of (11), the convergences of the sequences
$\\{\Delta_{\Gamma}^{t}+\Delta_{\Lambda}^{t}\\}$ and
$\\{\Delta_{\Gamma}^{t}-\Delta_{\Lambda}^{t}\\}$ imply that
$\Lambda^{t}_{j,r}$ and $\Gamma^{t}_{j,r}$ are converging, and the positive
sequence $\\{\Delta_{\Gamma}^{t}-\Delta_{\Lambda}^{t}\\}$ indicates that the
increment of the spurious feature $\Gamma^{t+1}_{j,r}-\Gamma^{t}_{j,r}$ is
larger than that of the invariant feature
$\Lambda^{t+1}_{j,r}-\Lambda^{t}_{j,r}$ at any iteration
$t\in\\{0,\ldots,T-1\\}$. ∎

####  D.2.1 Proof of Lemma D.2

First, we recall some concentration inequalities for sub-Gaussian random
variables. Since
$\bm{\xi}^{e}_{i}\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}))$,
for $(i^{\prime},e^{\prime})\neq(i,e)$, using Bernstein’s inequality for sub-
exponential random variables, we have for sufficiently small $a\geq 0$,

$\begin{gathered}\textup{Pr}\left\\{\lvert\langle\bm{\xi}^{e}_{i},\bm{\xi}^{e^{\prime}}_{i^{\prime}}\rangle\rvert\geq
a\right\\}\leq 2\exp\left\\{-\frac{a^{2}}{4\sigma_{p}^{4}(d-2)}\right\\},\\\
\textup{Pr}\left\\{\big{\lvert}\lVert\bm{\xi}^{e}_{i}\rVert_{2}^{2}-\sigma_{p}^{2}(d-2)\big{\rvert}\geq
a\right\\}\leq
2\exp\left\\{-\frac{a^{2}}{512\sigma_{p}^{4}(d-2)}\right\\}.\end{gathered}$

Moreover, for $\xi_{r}\sim\mathcal{N}(0,\sigma_{0}^{2})$ (indicating the
initial weights ${\mathbf{w}}^{0}_{j,r}$), the standard Gaussian tail gives

$\textup{Pr}\left\\{\left\lvert\frac{1}{m}\sum_{r=1}^{m}{\xi_{r}}\right\rvert\geq
a\right\\}\leq 2\exp\left\\{-\frac{ma^{2}}{2\sigma_{0}^{2}}\right\\}.$

Denote
$n\triangleq\sum_{e\in\mathcal{E}_{tr}}{n_{e}},\underline{n}\triangleq\min_{e\in\mathcal{E}_{tr}}{n_{e}}$,
by properly choosing $a$ for each tail bound and applying a union bound, we
can conclude that for $\rho>0$, with probability at least $1-\rho$, it holds
that $\forall i,e,i^{\prime},e^{\prime},r$,

$\displaystyle\lvert\langle\bm{\xi}^{e}_{i},\bm{\xi}^{e^{\prime}}_{i^{\prime}}\rangle\rvert\leq
2\sigma_{p}^{2}\sqrt{(d-2)\log{\frac{8n^{2}}{\rho}}},$
$\displaystyle\lVert\bm{\xi}^{e}_{i}\rVert_{2}^{2}\leq\sigma_{p}^{2}(d-2)+16\sigma_{p}^{2}\sqrt{2(d-2)\log{\frac{8n}{\rho}}},$
$\displaystyle\left\lvert\frac{1}{m}\sum_{r=1}^{m}{\xi_{r}}\right\rvert\leq\sigma_{0}\sqrt{\frac{2}{m}\log{\frac{32m}{\rho}}},$
$\displaystyle\lvert\langle\bm{\xi}_{r},\bm{\xi}^{e^{\prime}}_{i^{\prime}}\rangle\rvert\leq
2\sigma_{p}\sigma_{0}\sqrt{(d-2)\log{\frac{16nm}{\rho}}}.$

We start with bound the growth of $\Xi^{t,e}_{j,r,i}$. By bounding the update
rule (10), with probability at least $1-\rho$, we have

$\displaystyle\left\lvert\Xi^{t+1,e^{\prime}}_{j,r,i^{\prime}}\right\rvert$
$\displaystyle\leq\left\lvert\Xi^{t,e^{\prime}}_{j,r,i^{\prime}}\right\rvert+\frac{\eta}{m}\sum_{e\in\mathcal{E}_{tr}}{\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}{\frac{1}{1+\exp\\{y^{e}_{i}\hat{y}^{e}_{i}\\}}\cdot\lvert\langle\bm{\xi}^{e}_{i},\bm{\xi}^{e^{\prime}}_{i^{\prime}}\rangle\rvert}}$
$\displaystyle\leq\left\lvert\Xi^{t,e^{\prime}}_{j,r,i^{\prime}}\right\rvert+\frac{\eta}{m}\sum_{e\in\mathcal{E}_{tr}}{\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}{\lvert\langle\bm{\xi}^{e}_{i},\bm{\xi}^{e^{\prime}}_{i^{\prime}}\rangle\rvert}}$
$\displaystyle=\left\lvert\Xi^{0,e^{\prime}}_{j,r,i^{\prime}}\right\rvert+(t+1)\cdot\frac{\eta}{m}\sum_{e\in\mathcal{E}_{tr}}{\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}{\lvert\langle\bm{\xi}^{e}_{i},\bm{\xi}^{e^{\prime}}_{i^{\prime}}\rangle\rvert}}$
$\displaystyle=\lvert\langle\bm{\xi}_{r},\bm{\xi}^{e^{\prime}}_{i^{\prime}}\rangle\rvert+(t+1)\cdot\left(\frac{\eta}{mn_{e^{\prime}}}\lVert\bm{\xi}^{e^{\prime}}_{i^{\prime}}\rVert^{2}_{2}+\sum_{(i,e)\neq(i^{\prime},e^{\prime})}{\frac{\eta}{mn_{e}}\lvert\langle\bm{\xi}^{e}_{i},\bm{\xi}^{e^{\prime}}_{i^{\prime}}\rangle\rvert}\right)$
$\displaystyle\leq 2\sigma_{p}\sigma_{0}\sqrt{(d-2)\log{\frac{16nm}{\rho}}}$
$\displaystyle\ \
+\frac{T\eta\sigma_{p}^{2}}{m\underline{n}}\left((d-2)+16\sqrt{2(d-2)\log{\frac{8n}{\rho}}}+2n\sqrt{(d-2)\log{\frac{8n^{2}}{\rho}}}\right).$

Then, we can bound $\left\lvert\mathbb{Q}^{e}_{i}\right\rvert$ as

$\displaystyle\left\lvert\mathbb{Q}^{e}_{i}\right\rvert$ $\displaystyle\leq
2\cdot\left\lvert\frac{1}{m}\sum_{r=1}^{m}{\xi_{r}}\right\rvert+2\cdot\left\lvert\frac{1}{m}\sum_{r=1}^{m}{\xi_{r}}\right\rvert+\frac{2}{m}\sum_{r=1}^{m}{\left\lvert\Xi^{t,e}_{j,r,i}\right\rvert}$
$\displaystyle\leq
4\sigma_{0}\sqrt{\frac{2}{m}\log{\frac{32m}{\rho}}}+4\sigma_{p}\sigma_{0}\sqrt{(d-2)\log{\frac{16nm}{\rho}}}$
$\displaystyle\ \
+\frac{2T\eta\sigma_{p}^{2}}{m\underline{n}}\left((d-2)+16\sqrt{2(d-2)\log{\frac{8n}{\rho}}}+2n\sqrt{(d-2)\log{\frac{8n^{2}}{\rho}}}\right).$

Thus, with sufficient small $\sigma_{0},\sigma_{p}$, i.e.,

$\displaystyle\sigma_{0}^{2}$
$\displaystyle=O\left(\underline{n}^{-2}\log^{-1}{(m/\rho)}\right),$
$\displaystyle\sigma_{p}^{2}$
$\displaystyle=O\left(\min\left\\{d^{-1/2}\log^{-1/2}{(nm/\rho)},T^{-1}\eta^{-1}m\left(d+n\sqrt{d\log(n^{2}/\rho)}\right)^{-1}\right\\}\right),$

we ensured that
$\left\lvert\mathbb{Q}^{e}_{i}\right\rvert=O(\underline{n}^{-1})$.

####  D.2.2 Proof of Lemma D.3

For $e\in\mathcal{E}_{tr}$, using Hoeffding’s inequality, it holds that

$\textup{Pr}\left\\{\left\lvert\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}{\mathbf{1}_{\\{\textrm{Rad}(\alpha)_{i}=+1,\textrm{Rad}(\beta_{e})_{i}=+1\\}}-(1-\alpha)(1-\beta_{e})}\right\rvert\geq
a\right\\}\leq 2\exp{\\{-2a^{2}n_{e}\\}}.$

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

$\textup{Pr}\left\\{\left\lvert\overline{C}_{+1+1}-(1-\alpha)(2-\beta_{1}-\beta_{2})\right\rvert\leq
a\right\\}\geq 1-4\exp{\left\\{-\frac{a^{2}\underline{n}}{2}\right\\}}.$

Thus, for $\rho>0$, with probability at least $1-\frac{\rho}{4}$, we can
conclude that

$\left\lvert\overline{C}_{+1+1}-(1-\alpha)(2-\beta_{1}-\beta_{2})\right\rvert\leq\sqrt{\frac{2\log{(16/\rho)}}{\underline{n}}}.$

Using the above arguments for other constants $\overline{C}_{+1-1}$,
$\overline{C}_{-1+1}$ and $\overline{C}_{-1-1}$, and applying a union bound,
we obtain the desired results.

####  D.2.3 ERM Feature Learning with Non-Linear Activation Functions

It was numerically observed that in the early stage of (stochastic) GD
training, the learning dynamics of neural networks can be mimicked by training
a simple linear model [35]. Hu et al. [30] rigorously proved this phenomenon
for training two-layer neural network with $\ell_{2}$ loss function in the
Neural Tangent Kernel (NTK) region. We briefly summarize their results here:
For a two-layer fully-connected neural network (with fixed second layer
$\\{v_{r}\\}$):

$f_{FC}({\mathbf{W}},{\mathbf{x}})\triangleq\frac{1}{\sqrt{m}}\sum_{r=1}^{m}{v_{r}\psi\left({\mathbf{w}}_{r}^{\top}{\mathbf{x}}/\sqrt{d}\right)},$
(21)

considering the $\ell_{2}$ training loss
$\ell_{2}(\hat{y},y)\triangleq\frac{1}{2}(\hat{y}-y)^{2}$ and the ERM
objective
$L_{\textup{ERM}}({\mathbf{W}})=\frac{1}{n}\sum_{i=1}^{n}{\ell_{2}\big{(}f_{FC}({\mathbf{W}},{\mathbf{x}}_{i}),y_{i}\big{)}}$,
when using GD: ${\mathbf{W}}^{t+1}={\mathbf{W}}^{t}-\eta\cdot\nabla
L_{\textup{ERM}}({\mathbf{W}}^{t})$ to minimize the ERM objective, the
following holds.

######  Theorem D.4 (Theorem 3.2 of [30]).

Let $\alpha_{nl}\in(0,\frac{1}{4})$ be a fixed constant, and $\psi(\cdot)$ be
a smooth (with bounded first and second derivatives) or piece-wise linear
activation function. Suppose that $n$ and $m$ satisfy
$n=\Omega(d^{1+\alpha_{nl}})$ and $m=\Omega(d^{1+\alpha_{nl}})$. Suppose that
$\eta\ll d$. Then there exists a universal constant $c>0$ such that with high
probability, for all $t=O(\frac{d}{\eta}\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:

$\frac{1}{n}\sum_{i=1}^{n}{\big{(}f^{t}_{FC}({\mathbf{x}}_{i})-f^{t}_{lin}({\mathbf{x}}_{i})\big{)}^{2}}=O(d^{-\Omega(\alpha_{nl})}).$

The linear model
$f_{lin}(\bm{\beta},{\mathbf{x}})=\bm{\beta}^{\top}\bm{R}({\mathbf{x}})$ is a
linear function of the transformed data
$\bm{R}({\mathbf{x}})=\frac{1}{\sqrt{d}}\begin{bmatrix}\zeta{\mathbf{x}}\\\
\nu\end{bmatrix}$, where $\zeta$ and $\nu$ are constants related to
$\psi^{\prime}$ 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}(\bm{\beta},{\mathbf{x}})$ similar to the analysis444Note that when
$\psi(x)=x$, our CNN model can be viewed as a linear model with re-
parameterized weight matrices. Thus, the discussion in Appendix D.2 can be
viewed as studying the feature learning terms for a linear model with logistic
loss function. 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^{-\Omega(\alpha_{nl})})$.

Recall that our CNN model is
$f({\mathbf{W}},{\mathbf{x}})\\!=\\!F_{+1}({\mathbf{W}}_{+1},{\mathbf{x}})-F_{-1}({\mathbf{W}}_{-1},{\mathbf{x}})$,
where $F_{+1}({\mathbf{W}}_{+1},{\mathbf{x}})$ and
$F_{-1}({\mathbf{W}}_{-1},{\mathbf{x}})$ are defined as follows:

$\displaystyle
F_{j}({\mathbf{W}}_{j},{\mathbf{x}})=\frac{1}{m}\sum_{r=1}^{m}\left[{\psi}({\mathbf{w}}_{j,r}^{\top}{\mathbf{x}}_{1})+{\psi}({\mathbf{w}}_{j,r}^{\top}{\mathbf{x}}_{2})\right],j\in\\{\pm
1\\}.$

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}=\pm\frac{1}{\sqrt{m}}\\}$ and transforming the dataset as
$\left\\{\sqrt{d}\begin{bmatrix}y\cdot\textrm{Rad}(\alpha)\cdot{\mathbf{v}}_{1}+y\cdot\textrm{Rad}(\beta)\cdot{\mathbf{v}}_{2}\\\
0\end{bmatrix},\sqrt{d}\bm{\xi}\right\\}$. Then by tuning the norms of
${\mathbf{v}}_{1},{\mathbf{v}}_{2}$ and $\bm{\xi}$, we 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 $\ell_{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

$\displaystyle{\mathbf{c}}(t)\triangleq\left[C^{1}_{\text{IRMv1}}({\mathbf{W}},t),C^{2}_{\text{IRMv1}}({\mathbf{W}},t),\cdots,C^{|{{\mathcal{E}}_{\text{tr}}}|}_{\text{IRMv1}}({\mathbf{W}},t)\right],$

and $\lambda_{0}=\lambda_{\min}(\mathbf{H}^{\infty})$, where we define

$\displaystyle\mathbf{H}^{\infty}_{e,e^{\prime}}\triangleq\frac{1}{2mn_{e}n_{e^{\prime}}}\sum_{i=1}^{n_{e}}\psi^{\prime}(\langle{\mathbf{w}}_{j,r}(0),\mathbf{x}^{e}_{1,i}\rangle)\mathbf{x}^{e\top}_{1,i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}\psi^{\prime}(\langle{\mathbf{w}}_{j,r}(0),{\mathbf{x}}_{1,i^{\prime}}^{e^{\prime}}\rangle){\mathbf{x}}_{1,i^{\prime}}^{e^{\prime}}.$

Suppose that activation function is smooth, $\psi^{\prime}(0)\leq\beta$,
$|\psi^{\prime}(x)-\psi^{\prime}(x^{\prime})|<\beta|x-x^{\prime}|$ and
Lipschitz $|\psi(0)|<L$, $|\psi(x)-\psi(x^{\prime})|<L|x-x^{\prime}|$. Assume
that dimension $d=\Omega(\log(m/\delta))$, network width $m=\Omega(1/\delta)$,
regularization factor $\lambda\geq
1/(\sigma_{0}\sqrt{|\mathcal{E}_{tr}|}^{3})$, noise variance
$\sigma_{p}=O(d^{-2})$, weight initial scale
$\sigma_{0}=O(\frac{|\mathcal{E}_{tr}|^{7/2}\beta^{3}L}{d^{1/2}m^{2}\lambda^{2}_{0}\log(1/\epsilon)})$,
then with probability at least $1-\delta$, after training time
$T=\Omega\left(\frac{\log(1/\epsilon)}{\eta\lambda\lambda_{0}}\right)$, we
have:

$\displaystyle\|{\mathbf{c}}(T)\|_{2}\leq\epsilon,\quad{\gamma^{inv}_{j,r}(T)=o(1),\quad\gamma^{spu}_{j,r}(T)=o(1)}.$

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

######  Lemma D.6 ([12]).

Suppose that $\delta>0$ and $d=\Omega(\log(4n/\delta))$. Then with probability
at least $1-\delta$,

$\displaystyle\sigma_{p}^{2}d/2\leq\|\bm{\xi}_{i}\|_{2}^{2}\leq
3\sigma_{p}^{2}d/2$

for all $i,i^{\prime}\in[n]$.

######  Lemma D.7 ([12]).

Suppose that $d\geq\Omega(\log(mn/\delta))$, $m=\Omega(\log(1/\delta))$. Then
with probability at least $1-\delta$,

$\displaystyle|\langle{\mathbf{w}}_{j,r}^{(0)},\mathbf{v}_{1}\rangle|\leq\sqrt{2\log(8m/\delta)}\cdot\sigma_{0}\|\mathbf{v}_{1}\|_{2},$
$\displaystyle|\langle{\mathbf{w}}_{j,r}^{(0)},\mathbf{v}_{2}\rangle|\leq\sqrt{2\log(8m/\delta)}\cdot\sigma_{0}\|\mathbf{v}_{2}\|_{2},$
$\displaystyle|\langle{\mathbf{w}}_{j,r}^{(0)},\bm{\xi}_{i}\rangle|\leq
2\sqrt{\log(8mn/\delta)}\cdot\sigma_{0}\sigma_{p}\sqrt{d}$

for all $r\in[m]$, $j\in\\{\pm 1\\}$ and $i\in[n]$.

######  Lemma D.8.

Suppose that $\delta>0$ and $d=\Omega(\log(4m/\delta))$. Then with probability
at least $1-\delta$, for all $r\in[m]$ and $j\in\\{-1,1\\}$, we have

$\displaystyle\sigma_{0}^{2}d/2\leq\|\mathbf{w}_{j,r}(0)\|_{2}^{2}\leq
3\sigma_{0}^{2}d/2.$

###### Proof of Lemma D.8.

By Bernstein’s inequality, with probability at least $1-\delta/(2m)$ we have

$\displaystyle\big{|}\|\mathbf{w}_{j,r}(0)\|_{2}^{2}-\sigma_{0}^{2}d\big{|}=O(\sigma_{0}^{2}\cdot\sqrt{d\log(4m/\delta)}).$

Therefore, as long as $d=\Omega(\log(4m/\delta))$, we have

$\displaystyle\sigma_{0}^{2}d/2\leq\|\mathbf{w}_{j,r}(0)\|_{2}^{2}\leq
3\sigma_{0}^{2}d/2.$

∎

###### Proof of Theorem D.5.

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

$\displaystyle\frac{d{\mathbf{w}}_{j,r}(t)}{dt}$
$\displaystyle=-\eta\cdot\nabla_{{\mathbf{w}}_{j,r}}L_{\mathrm{IRMv1}}({\mathbf{W}}(t))$
$\displaystyle=-\frac{\eta}{nm}\sum_{e\in\mathcal{E}_{\mathrm{tr}}}\sum_{i=1}^{n_{e}}\ell^{\prime}_{i}(t){\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(t),y_{i}^{e}{\mathbf{v}}_{i}^{e}\rangle)\cdot
j{\mathbf{v}}_{i}^{e}-\frac{\eta}{nm}\sum_{e\in\mathcal{E}_{\mathrm{tr}}}\sum_{i=1}^{n_{e}}\ell^{\prime}_{i}(t){\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(t),\bm{\xi}_{i}\rangle)\cdot
jy_{i}^{e}\bm{\xi}_{i}$
$\displaystyle\quad-\frac{2\eta\lambda}{nm}\sum_{e\in\mathcal{E}_{\mathrm{tr}}}C_{\text{IRMv1}}^{e}\sum_{i=1}^{n_{e}}\ell^{\prime\prime}_{i}\hat{y}_{i}^{e}{\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(t),y_{i}^{e}{\mathbf{v}}_{i}^{e}\rangle)jy_{i}^{e}{\textnormal{v}}^{e}_{i}-\frac{2\eta\lambda}{nm}\sum_{e\in\mathcal{E}_{\mathrm{tr}}}C_{\text{IRMv1}}^{e}\sum_{i=1}^{n_{e}}\ell^{\prime\prime}_{i}\hat{y}_{i}^{e}{\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(t),\bm{\xi}_{i}\rangle)j\bm{\xi}_{i}$
$\displaystyle\quad-\frac{2\eta\lambda}{nm}\sum_{e\in\mathcal{E}_{\mathrm{tr}}}C_{\text{IRMv1}}^{e}\sum_{i=1}^{n_{e}}\ell^{\prime}_{i}(t){\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(t),y_{i}^{e}{\mathbf{v}}_{i}^{e}\rangle)j{\mathbf{v}}_{i}^{e}-\frac{2\eta\lambda}{nm}\sum_{e\in\mathcal{E}_{\mathrm{tr}}}C_{\text{IRMv1}}^{e}\sum_{i=1}^{n_{e}}\ell^{\prime}_{i}(t){\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(t),\bm{\xi}_{i}\rangle)jy_{i}^{e}\bm{\xi}_{i}$
$\displaystyle=-\frac{\eta}{nm}\sum_{e\in\mathcal{E}_{\mathrm{tr}}}(1+2\lambda
C_{\text{IRMv1}}^{e}(t))\sum_{i=1}^{n_{e}}\ell^{\prime}_{i}(t){\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(t),y_{i}^{e}{\mathbf{v}}_{i}^{e}\rangle)\cdot
j{\mathbf{v}}_{i}^{e}$
$\displaystyle\quad-\frac{\eta}{nm}\sum_{e\in\mathcal{E}_{\mathrm{tr}}}(1+2\lambda
C_{\text{IRMv1}}^{e}(t))\sum_{i=1}^{n_{e}}\ell^{\prime}_{i}(t){\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(t),\bm{\xi}_{i}\rangle)\cdot
jy_{i}^{e}\bm{\xi}_{i}$
$\displaystyle\quad-\frac{2\eta\lambda}{nm}\sum_{e\in\mathcal{E}_{\mathrm{tr}}}C_{\text{IRMv1}}^{e}\sum_{i=1}^{n_{e}}\ell^{\prime\prime}_{i}\hat{y}_{i}^{e}{\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(t),y_{i}^{e}{\mathbf{v}}_{i}^{e}\rangle)jy_{i}^{e}{\textnormal{v}}^{e}_{i}-\frac{2\eta\lambda}{nm}\sum_{e\in\mathcal{E}_{\mathrm{tr}}}C_{\text{IRMv1}}^{e}\sum_{i=1}^{n_{e}}\ell^{\prime\prime}_{i}\hat{y}_{i}^{e}{\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(t),\bm{\xi}_{i}\rangle)j\bm{\xi}_{i},$

where $C_{\text{IRMv1}}^{e}=\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}\ell_{i}^{\prime
e}\hat{y}_{i}^{e}y_{i}^{e}$ and
${\mathbf{v}}_{i}^{e}=\textrm{Rad}(\alpha)_{i}\cdot{\mathbf{v}}_{1}+\textrm{Rad}(\beta_{e})_{i}\cdot{\mathbf{v}}_{2}$.
Note that $\ell^{\prime\prime}$ has the opposite sign to $\ell^{\prime}$.

Then we look at the dynamics of $C^{e}_{\text{IRMv1}}(t)$ according to the
gradient flow update rule:

$\displaystyle\frac{dC^{e}_{\text{IRMv1}}({\mathbf{W}},t)}{dt}$