# Feature Contamination: Neural Networks Learn Uncorrelated Features
and Fail to Generalize

Tianren Zhang  Chujie Zhao  Guanyu Chen  Yizhou Jiang  Feng Chen

###### Abstract

Learning representations that generalize under distribution shifts is critical
for building robust machine learning models. However, despite significant
efforts in recent years, algorithmic advances in this direction have been
limited. In this work, we seek to understand the fundamental difficulty of
out-of-distribution generalization with deep neural networks. We first
empirically show that perhaps surprisingly, even allowing a neural network to
_explicitly_ fit the representations obtained from a teacher network that
_can_ generalize out-of-distribution is insufficient for the generalization of
the student network. Then, by a theoretical study of two-layer ReLU networks
optimized by stochastic gradient descent (SGD) under a structured feature
model, we identify a fundamental yet unexplored feature learning proclivity of
neural networks, _feature contamination_ : neural networks can learn
_uncorrelated_ features together with predictive features, resulting in
generalization failure under distribution shifts. Notably, this mechanism
essentially differs from the prevailing narrative in the literature that
attributes the generalization failure to spurious correlations. Overall, our
results offer new insights into the non-linear feature learning dynamics of
neural networks and highlight the necessity of considering inductive biases in
out-of-distribution generalization.111Code is available at
https://github.com/trzhang0116/feature-contamination.

Out-of-Distribution Generalization, Distribution Shift, Spurious Correlations,
Neural Networks, Representation Learning

##  1 Introduction

The capability of generalizing under distribution shifts is crucial for
machine learning systems to be deployed in the wild (Amodei et al., 2016;
Beery et al., 2018; Koh et al., 2021). In the last decade, it has proved that
the conventional principle of empirical risk minimization (ERM), when combined
with deep neural networks, can lead to remarkable in-distribution (ID)
generalization performance given sufficient training data. Nevertheless, this
powerful paradigm can often fail in _out-of-distribution (OOD) generalization_
, where distribution shifts occur due to data variations that are not well-
covered in training (Torralba & Efros, 2011; Beery et al., 2018; Geirhos et
al., 2018; DeGrave et al., 2021).

In response, recent years have witnessed a surge of developing algorithms that
promote OOD generalization. However, the effectiveness of many proposed
algorithms has been called into question by recent work (Gulrajani & Lopez-
Paz, 2021; Koh et al., 2021), in which no tested algorithm exhibits a
significant advantage over ERM under fair comparisons. On the other hand, it
turns out that the most effective means of improving OOD generalization to
date is pre-training on a more diverse dataset (Taori et al., 2020; Wiles et
al., 2022), with notable examples including CLIP (Radford et al., 2021) and
GPT (Brown et al., 2020). Yet, using additional pre-training data also blurs
the notion of “OOD” itself since it essentially expands the training
distribution. Moreover, it has been observed that when the test distribution
differs from the pre-training distribution, pre-trained models can also suffer
from performance degradation (Bommasani et al., 2022; Liu et al., 2023; Li &
Flanigan, 2023).

The limited algorithmic success underlines the necessity of identifying and
understanding the fundamental factors behind OOD generalization. In
particular, a prevailing narrative in the literature attributes the OOD
generalization failure to _spurious correlations_ (Arjovsky et al., 2019;
Nagarajan et al., 2021; Schölkopf et al., 2021). This explanation is inspired
by the observation that the representations learned by ERM can absorb features
that have _correlational_ yet non-causal relationships with the output (Beery
et al., 2018; Geirhos et al., 2020), and it has motivated a main line of
algorithmic endeavor of designing better representation learning objectives in
recent years (Arjovsky et al., 2019; Krueger et al., 2021; Mitrovic et al.,
2021; Chen et al., 2022; Shi et al., 2022). However, despite being intuitive,
it remains elusive how much this failure mode actually contributes to the OOD
generalization failure in practice—as we will elaborate in the following
sections, there exists a major OOD generalization gap in many tasks that
_cannot_ be straightforwardly explained by spurious correlations, implying
that there must exist some more dominant factors.

On the theoretical side, a series of work has been devoted to analyzing the
failure modes of OOD generalization. However, existing analysis has two major
limitations: (i) conceptually, most studies only consider the failure mode due
to spurious correlations; (ii) technically, most studies either only consider
_linear_ models such as linear classification over prescribed features or
neural tangent kernels (Arjovsky et al., 2019; Sagawa et al., 2020b; Nagarajan
et al., 2021; Xu et al., 2021; Ahuja et al., 2021b, a; Pezeshki et al., 2021;
Chen et al., 2022; Wang et al., 2022; Rosenfeld et al., 2022; Abbe et al.,
2023; Chen et al., 2023), or only consider arbitrary _unstructured_ models
without taking into the account the role of optimization (Rosenfeld et al.,
2021; Kamath et al., 2021; Ye et al., 2021)—this makes them unable to capture
the inductive biases of today’s most widely used model class, i.e., _neural
networks_. As a result, it has been observed that many OOD generalization
algorithms that enjoy provable guarantees in their theoretical models do not
excel in practice (Gulrajani & Lopez-Paz, 2021).

Overall, the above results imply that current explanations and theoretical
models on OOD generalization may _not_ faithfully reflect real-world
distribution shifts. Motivated by the gap between theory and practice, we
argue that taking into account the inductive biases of neural networks is not
only important but also _necessary_ for understanding OOD generalization in
the era of deep learning.

###  1.1 Our Results and Implications

In this work, we set out to understand the fundamental difficulty of OOD
generalization with deep neural networks:

Empirically, inspired by the ongoing trend of designing specific
representation learning objectives for OOD generalization (Arjovsky et al.,
2019; Gulrajani & Lopez-Paz, 2021), we investigate what will happen in an
“ideal” setting where good representations are _explicitly given_ during
training. Concretely, we show on a range of distribution shift benchmarks that
perhaps surprisingly, even if we allow a neural network to explicitly fit the
representations obtained from a teacher network that _can_ generalize out-of-
distribution, the performance of the student network can still significantly
deteriorate under distribution shifts. Our results thus imply that only
considering the effect of the representation learning objective is
_insufficient_ for understanding OOD generalization in practice without
considering the inductive biases in optimization. Moreover, we show that the
above generalization failure _cannot_ be simply explained by spurious
correlations or other existing explanations in the literature on OOD
generalization.

Theoretically, we prove that in certain structured binary classification tasks
where the data is generated from generalizable _core features_ and other
_background features_ (formal definitions in Section 3), a randomly
initialized two-layer ReLU neural network trained by SGD can achieve ID
generalization given sufficient iterations, yet fails to generalize OOD. In
particular, we show that the above failure mode differs fundamentally from
prior work as it holds even when:

  * •

Background features are _uncorrelated_ with the label (this excludes the
failure mode due to spurious correlations).

  * •

Ground-truth labels can be perfectly predicted by core features (this excludes
the failure mode due to lacking informative features for prediction).

  * •

Core features and background features are distributed in orthogonal subspaces
(this excludes the failure mode due to non-linearly entangled features in the
input that may be hard to disentangle for the neural network in training).

Figure 1: OOD performance ($y$-axes) v.s. ID performance ($x$-axes) for three
model families including (i) linear probes on pre-trained representations
(purple stars), (ii) linear probes on distilled representations (orange
squares), and (iii) standard models trained on ID data (blue circles). The
$y$-axis of the sixth panel stands for the average accuracy on ImageNet-based
OOD test sets, averaged from the first five panels. Please refer to Section E
for more details on each model family.

Instead, we identify that the above failure stems from a fundamental yet
unexplored feature learning proclivity, which we name _feature contamination_
, of neural networks. In brief, feature contamination indicates that during
the learning of core features, SGD-trained neural networks also learn
background features simultaneously, even when background features are
_uncorrelated_ with the label and in the presence of weight decay. The reason
for this phenomenon is that the neurons in the network tend to have
_asymmetric activation_ for different classes, resulting in non-zero expected
gradient projections onto _both_ the core feature subspace and the background
feature subspace. This eventually leads to additional risks under distribution
shifts due to the coupling of core and background features in the neurons’
pre-activation. Moreover, we formally show that ReLU networks and linear
networks are _provably different_ in our setting with the latter exhibiting no
such behavior, suggesting a separation between linear and non-linear models.
Finally, we present empirical evidence on deep neural networks that connects
feature contamination to the empirical OOD generalization failure observed in
our experiments.

At a high level, we expect that feature contamination as a _novel inductive
bias of SGD-trained neural networks_ may also be used in more general
contexts. For example, it may serve as a new perspective for understanding the
_feature learning_ process of (deep) neural networks, complementing other
known inductive biases of neural networks such as the simplicity bias (Arpit
et al., 2017; Shah et al., 2020).

##  2 Good Representations Are Hard to Learn Even when Explicitly Given in
Training

We begin our analysis by an empirical study inspired by recent algorithmic
explorations in OOD generalization. Existing work has made various attempts to
learn OOD-generalizable models by designing auxiliary _representation
learning_ objectives beyond minimizing the prediction risk (see the baseline
models in Section E.4 for some examples). Typically, those objectives reflect
the premise of the properties of “good” representations. For example, a major
line of work focuses on learning invariant representations across multiple
training domains (Arjovsky et al., 2019; Chen et al., 2022; Shi et al., 2022),
with the aim of removing domain-specific spurious correlations. Given the
limited success of existing algorithms, here we would like to investigate the
fundamental limitations of such representation learning methods. However, a
main confounder in our study is that it is often unclear whether optimizing
certain objectives is indeed effective for shaping the representation to
satisfy the ideal properties—for example, it has been shown that optimizing
some invariant representation learning objectives may lead to representations
that are not truly invariant (Kamath et al., 2021; Rosenfeld et al., 2021).

To ablate the potential sub-optimality in the representation learning
objective, we focus on an “ideal” scenario where the model has _explicit
access to good representations_ in training. Concretely, we leverage large-
scale pre-trained models such as CLIP (Radford et al., 2021), which has shown
remarkable robustness against distribution shifts, to extract good
representations for each input: given a pre-trained CLIP encoder as a _teacher
encoder_ , we randomly initialize another _student encoder_ with the _same_
architecture. We then train the student encoder by minimizing the Euclidean
distance between its output representations and the representations extracted
by the teacher encoder, a process known as representation distillation (Hinton
et al., 2014; Tian et al., 2020). Finally, we evaluate the ID and the OOD
performance of both models. In total, our experiments span six different pre-
trained models and eight extensively benchmarked distribution shift datasets,
including five ImageNet-based natural distribution shift datasets (Taori et
al., 2020), two in-the-wild distribution shift datasets from WILDS (Koh et
al., 2021), and a domain generalization dataset DomainNet (Peng et al., 2019).
Please see Section E for more experimental details.

Evaluation protocol. We evaluate the ID and the OOD performance of pre-trained
and distilled encoders by training linear probes on top of their output
representations on the ID training set and then evaluate those linear probes
on both ID and OOD test sets. Note that under our protocol, the linear probes
still face OOD generalization tasks on the OOD test set, albeit with
representations instead of raw images as inputs. To compare the OOD
generalization ability of different models, we follow the evaluation protocol
of _effective robustness_ (Taori et al., 2020), which quantifies a model’s
distribution shift robustness as its OOD performance advantage over a baseline
representing the OOD performance of standard models trained on ID data.
Following Taori et al. (2020), we illustrate the effective robustness of our
models using scatter plots, with $x$-axes representing ID performance and
$y$-axes representing OOD performance.

Results. As shown in Figure 1, linear probes on distilled representations
exhibit consistent effective robustness gains over standard models. This is
not very surprising given that the distilled models have additional
supervision provided by the representations obtained from the teacher models
in training, while standard models do not. However, the upshot is that _even
with explicit access to good representations, the OOD generalization
performance of distilled models still lags far behind their pre-trained
counterparts_. For example, distilled models only close about half of the
average effective robustness gap between standard models and pre-trained
models in ImageNet-based datasets, with even worse performance on iWildCam and
DomainNet. Note that this is not due to the failure in distillation, as
distilled models do achieve similar _ID performance_ to that of the pre-
trained models. Our results thus suggest that the limited algorithmic success
in OOD generalization cannot be simply explained by not having a “good enough”
representation learning objective.

What is the cause of the above failure? First, one may argue that spurious
correlations can still play a role here, as the representations extracted by
pre-trained models may also contain spurious correlations to the label even if
they achieve generally good OOD performance. While we do acknowledge this
possibility, we emphasize that spurious correlations _cannot_ explain the
large _OOD performance gap_ between the distilled and the pre-trained models
since we would expect them to be similarly impacted by the spurious
correlations in their representations. Another plausible explanation is data
leakage, i.e., CLIP may have “seen” many OOD examples in its pre-training
stage and thus can extract richer predictive features for OOD examples (Zhang
& Bottou, 2023). However, this possibility is nullified by a recent study
(Mayilvahanan et al., 2024), which shows that CLIP’s distribution shift
robustness persists even when OOD examples are pruned from its pre-training
dataset.

In a nutshell, we argue that _existing explanations are insufficient to
account for the above OOD generalization gap_. This suggests that taking into
the account the inductive biases of SGD-trained neural networks are necessary
for understanding the OOD generalization failure in practice. In the following
sections, we will formally identify feature contamination as a novel OOD
generalization failure mode and further connect it to the results in this
section.

##  3 A Theoretical Model of OOD Generalization

Notation. We use $[d]$ to denote the set $\\{1,\ldots,d\\}$ for positive
integers $d$. For a set $\mathcal{S}$, we denote its cardinality by
$|\mathcal{S}|$. For a vector ${\mathbf{u}}$, we denote its $\ell^{2}$-norm by
$\lVert{\mathbf{u}}\rVert_{2}$. We denote the inner product of two vectors
${\mathbf{u}}$ and ${\mathbf{v}}$ by
$\langle{\mathbf{u}},{\mathbf{v}}\rangle$. We use the standard big-O notation:
$O(\cdot)$, $\Omega(\cdot)$, $\Theta(\cdot)$, $o(\cdot)$, as well as their
soft-O variants such as $\widetilde{\Theta}(\cdot)$ to hide logarithmic
factors. For some parameter $d$, we use $\mathsf{poly}(d)$ to denote
$\Theta(d^{C})$ with some unspecified large constant $C$. We use
$\mathbf{1}_{E}$ to denote the indicator function for an event $E$.

###  3.1 OOD Generalization Problem Setup

Task and data. We consider a binary classification task with an input space
$\mathcal{X}\subseteq\mathbb{R}^{d}$, a label space $\mathcal{Y}=\\{-1,1\\}$,
a model class $\mathcal{H}:\mathcal{X}\to\mathbb{R}$, and a loss function
$\ell:\mathcal{Y}\times\mathcal{Y}\to\mathbb{R}$. For every distribution
$\mathcal{D}$ over $\mathcal{X}\times\mathcal{Y}$ and model $h\in\mathcal{H}$,
the expected risk of $h$ on $\mathcal{D}$ is given by
$\mathcal{R}_{\mathcal{D}}(h)\vcentcolon=\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})\sim\mathcal{D}}\ell(h({\mathbf{x}}),{\mathbf{y}})$.
In an OOD generalization problem, there exist a set of distributions
$\mathbb{D}$ that consists of all possible distributions to which we would
like our model to generalize. In training, we have access to a training
distribution set $\mathbb{D}_{\mathrm{train}}\subsetneq\mathbb{D}$, where
$\mathbb{D}_{\mathrm{train}}$ may contain one or multiple training
distributions. Following prior work (Arjovsky et al., 2019; Sagawa et al.,
2020a; Nagarajan et al., 2021; Rosenfeld et al., 2021), we aim to select a
model $h\in\mathcal{H}$ to minimize the _OOD risk_ , defined as the worst-case
expected risk on $\mathbb{D}$:

$\mathcal{R}_{\mathrm{OOD}}(h)\vcentcolon=\max_{\mathcal{D}\in\mathbb{D}}\mathcal{R}_{\mathcal{D}}(h).$
(1)

It is clear that without further assumptions on $\mathbb{D}_{\mathrm{train}}$
and $\mathbb{D}$, OOD generalization is impossible since no model can
generalize to an arbitrary distribution. Fortunately, real-world distribution
shifts are often _structured_ with some structural similarities shared by
different distributions. We can thus hope that such structures can be captured
by certain algorithms to train models that can generalize OOD.

To formalize this, in this work we assume that both ID and OOD data are
generated by a dictionary
$\bm{M}=(\bm{m}_{1},\ldots,\bm{m}_{d_{0}})\in\mathbb{R}^{d\times d_{0}}$
consisting of $d_{0}$ features with each feature
$\bm{m}_{i}\in\mathbb{R}^{d}$. Throughout the paper, we work with the case
where $d_{0}$ is sufficiently large and
$d\in[\Omega(d_{0}^{2.01}),\mathsf{poly}(d_{0})]$. For simplicity, we assume
that every feature satisfies $\lVert\bm{m}_{i}\rVert_{2}=1$ and different
features are orthogonal: $\forall i\neq
j\in[d_{0}],\langle{\bm{m}_{i}},{\bm{m}_{j}}\rangle=0$.222Another advantage of
assuming orthogonal features is that this prevents the network from learning
background features due to their correlations with core features. We note that
our results can be extended to more general settings without orthogonality.

Among all features in $\bm{M}$, we assume that there are $d_{\mathrm{core}}$
features consistently correlating with the label in all distributions in
$\mathbb{D}$. We denote the index set of those features by
$\mathcal{S}_{\mathrm{core}}\subsetneq[d_{0}]$ and refer to them as core
features since they are consistently predictive of the label in all
distributions. We refer to the remaining features as background features and
denote their index set by
$\mathcal{S}_{\mathrm{bg}}=[d_{0}]\setminus\mathcal{S}_{\mathrm{core}}$ with
$d_{\mathrm{bg}}\vcentcolon=|\mathcal{S}_{\mathrm{bg}}|=d_{0}-d_{\mathrm{core}}$.
We assume that $d_{\mathrm{core}}=\Theta(d_{0})$ and
$d_{\mathrm{bg}}=\Theta(\frac{d_{0}}{\log d_{0}})$, so that the number of both
core features and background features is non-negligible. With the above
definitions, we introduce our ID and OOD data generation model in Definition
3.1.

######  Definition 3.1 (ID and OOD data generation).

Under the feature model stated above, consider a training distribution (ID
data distribution)
$\mathcal{D}_{\mathrm{train}}\in\mathbb{D}_{\mathrm{train}}$ and a test
distribution (OOD data distribution)
$\mathcal{D}_{\mathrm{test}}\in\mathbb{D}\setminus\mathbb{D}_{\mathrm{train}}$.333Note
that $\mathcal{D}_{\mathrm{train}}$ and $\mathcal{D}_{\mathrm{test}}$ can also
be (weighted) _mixtures_ of multiple distributions in
$\mathbb{D}_{\mathrm{train}}$ and
$\mathbb{D}\setminus\mathbb{D}_{\mathrm{train}}$, respectively. Each example
$({\mathbf{x}},{\mathbf{y}})\sim\mathcal{D}\in\\{\mathcal{D}_{\mathrm{train}},\mathcal{D}_{\mathrm{test}}\\}$
is generated as follows:

  1. 1.

Sample a label ${\mathbf{y}}$ from the uniform distribution over
$\mathcal{Y}$.

  2. 2.

Sample a weight vector
${\mathbf{z}}=({\mathbf{z}}_{1},\ldots,{\mathbf{z}}_{d_{0}})\in\mathbb{R}^{d_{0}}$
where different coordinates of ${\mathbf{z}}$ are independent random variables
generated as follows:

     * •

ID data ($\mathcal{D}=\mathcal{D}_{\mathrm{train}}$): for every $j\in[d_{0}]$,
sample ${\mathbf{z}}_{j}$ from some distribution $\mathcal{D}_{j}$ over
$[0,1]$ such that its moments satisfy
$\mu_{jp}\vcentcolon=\mathbb{E}_{\mathcal{D}_{j}}{{\mathbf{z}}_{j}^{p}}=\Theta(1)$
for $p\in[3]$ and its variance satisfies $\sigma^{2}_{j}=\Theta(1)$.

     * •

OOD data ($\mathcal{D}=\mathcal{D}_{\mathrm{test}}$): for every $j\in[d_{0}]$,
if $j\in\mathcal{S}_{\mathrm{core}}$, sample ${\mathbf{z}}_{j}$ from
$\mathcal{D}_{j}$ over $[0,1]$; if $j\in\mathcal{S}_{\mathrm{bg}}$, sample
${\mathbf{z}}_{j}$ from some distribution $\mathcal{D}^{\prime}_{j}$ over
$[-1,0]$ such that
$\mathbb{E}_{\mathcal{D}_{j}^{\prime}}{{\mathbf{z}}_{j}}=-\Theta(1)$.

  3. 3.

Generate
${\mathbf{x}}=\sum_{j\in\mathcal{S}_{\mathrm{core}}}{\mathbf{y}}{\mathbf{z}}_{j}\bm{m}_{j}+\sum_{j\in\mathcal{S}_{\mathrm{bg}}}{\mathbf{z}}_{j}\bm{m}_{j}$.

Remarks on data generation. Our data model formalizes a structured OOD
generalization setup reflecting several facets of real-world OOD
generalization problems:

  * •

The explicit separation of core and background features captures structural
assumptions that make OOD generalization tractable: under the distribution
shifts on background features, there still exists a set of core features that
enable robust classification. Hence, a model that discards background features
and retains core features can generalize OOD. This rules out the ill-posed
case where the ID data is not informative enough to learn a generalizable
model (Tripuraneni et al., 2020; Xu et al., 2021; Kumar et al., 2022), and is
also the key intuition of many OOD generalization algorithms aiming to learn
invariant representations (Gulrajani & Lopez-Paz, 2021).

  * •

The weights of background features are assumed to be independent of the label,
rendering background features and labels _uncorrelated_. This differs from
prior OOD generalization analysis (Arjovsky et al., 2019; Sagawa et al.,
2020b; Nagarajan et al., 2021; Rosenfeld et al., 2021) where background
features are assumed to be _spuriously correlated_ with the label and hence
useful for prediction. We intentionally make this assumption to “ablate” the
effect of spurious correlations in feature learning.444On the other hand, our
results can be extended to the settings where some background features have
spurious correlations.

###  3.2 Model and Training

Model. We consider a model class $\mathcal{H}$ representing width-$m$ two-
layer neural networks with ReLU activation. Formally, given hidden-layer
weights
${\mathbf{W}}=({\mathbf{w}}_{1},\ldots,{\mathbf{w}}_{m})\in\mathbb{R}^{d\times
m}$ and output-layer weights
${\mathbf{a}}=(a_{1},\ldots,a_{m})^{\top}\in\mathbb{R}^{m}$, the output of a
model $h\in\mathcal{H}$ given an input ${\mathbf{x}}\in\mathcal{X}$ is defined
as

$h({\mathbf{x}})=\sum\nolimits_{k\in[m]}a_{k}\cdot\mathsf{ReLU}{(\langle{{\mathbf{w}}_{k}},{{\mathbf{x}}}\rangle)},$
(2)

where $\mathsf{ReLU}(u)=\max\\{u,0\\},u\in\mathbb{R}$. Similar to practical
design choices, we consider an _overparameterized_ setting where
$m\in[\Theta(d_{0}),\Theta(d)]$. We initialize each weight vector
${\mathbf{w}}_{k},k\in[m]$ by sampling
${\mathbf{w}}_{i}^{(0)}\sim\mathcal{N}(\mathbf{0},\sigma^{2}_{0}\bm{I}_{d})$
with $\sigma_{0}^{2}=\frac{1}{d}$. We randomly initialize output-layer weights
${\mathbf{a}}$ by sampling
$a_{k}\sim\mathsf{Uniform}\\{-\frac{1}{m},\frac{1}{m}\\}$ independently for
each $k\in[m]$. To simplify our analysis, we keep output-layer weights
${\mathbf{a}}$ fixed during training, which is a common assumption in
analyzing two-layer neural networks (Allen-Zhu & Li, 2021; Karp et al., 2021;
Allen-Zhu & Li, 2023b).

Figure 2: A diagram of feature contamination in our binary classification
setting. Left: for models with non-linear activation functions such as ReLU,
activation asymmetry leads to non-zero gradient projections onto background
features. Right: for linear models, background features are cancelled out in
the gradients, exhibiting no feature contamination.

Training. We train the network using SGD to minimize a standard hinge loss
$\ell(y,y^{\prime})=\max\\{1-yy^{\prime},0\\}$ with step size $\eta>0$ for $T$
iterations. We also include a weight decay with strength
$\lambda=O(\frac{d_{0}}{m^{1.01}})$ for regularization. At each iteration
$t\in\\{0,\ldots,T\\}$, we i.i.d. sample a batch of examples
$\\{({\mathbf{x}}_{i}^{(t)},{\mathbf{y}}_{i}^{(t)})\\}_{i\in[N]}\sim\mathcal{D}_{\mathrm{train}}^{N}$
with batch size $N=\mathsf{poly}(d)$ and consider the following empirical
loss:

$\widehat{\mathcal{L}}(h^{(t)})=\frac{1}{N}\sum_{i\in[N]}\ell\left(h^{(t)}({\mathbf{x}}_{i}^{(t)}),{\mathbf{y}}_{i}^{(t)}\right)+\frac{\lambda}{2}\sum_{k\in[m]}\big{\lVert}{\mathbf{w}}_{k}^{(t)}\big{\rVert}_{2}^{2},$
(3)

where we use $h^{(t)}$ to denote the model at iteration $t$, with weights
${\mathbf{W}}^{(t)}=({\mathbf{w}}_{1}^{(t)},\ldots,{\mathbf{w}}_{m}^{(t)})$.
The SGD update for each weight vector ${\mathbf{w}}_{k},k\in[m]$ is then given
by

${\mathbf{w}}_{k}^{(t+1)}={\mathbf{w}}_{k}^{(t)}-\eta\nabla_{{\mathbf{w}}_{k}^{(t)}}\widehat{\mathcal{L}}(h^{(t)}).$
(4)

##  4 Main Theoretical Results

In this section, we present our main theoretical results, provide mathematical
reasoning of why feature contamination happens, and discuss its impact on
generalization. We also include numerical results and show that our findings
can be extended to more general data and neural network models.

Technical challenges. As we have discussed in Section 1, most existing
theoretical work on OOD generalization _separates generalization and
optimization_ and directly studies the _global minimizers_ of their training
objecives without considering optimization dynamics. By contrast, without a
unique global minimizer, our setup requires an explicit analysis on the SGD
optimization trajectory, which is known to be challenging due to its _non-
convex_ and _non-linear_ nature. Prior work has studied fine-tuning pre-
trained models for OOD generalization in the context of two-layer _linear
networks_ (Kumar et al., 2022; Lee et al., 2023). Analyzing non-linear
networks further requires a careful treatment on the activation property of
the neurons, which results in SGD dynamics that essentially deviate from
linear networks.

Figure 3: Numerical results. (a) _ID and OOD risks:_ During training, ID loss
quickly approaches zero, while OOD loss stays high. (b) _Activation
asymmetry:_ the difference of average neuron activation rates for different
classes largely increases during training. (c) _Feature contamination:_ the
average correlations between neuron weights and both core features and
_uncorrelated background features_ increase in training. (d) Feature
contamination also occurs in more general settings with different activation
functions. Please refer to Section F.1 for more details and results.

Our approach. At a high level, our analysis is based on the construction of
two neuron subsets $\mathcal{N}_{y}^{(t)}$ (see Definition A.4) for
$y\in\mathcal{Y}=\\{-1,1\\}$ at iteration $t\in\\{0,\ldots,T\\}$ so that each
subset has cardinality $\Theta(m)$ and its neurons are randomly initialized to
have large enough expected correlations with the examples from the class $y$
(i.e., “winning the lottery tickets” (Frankle & Carbin, 2019; Allen-Zhu & Li,
2021)). We then apply the Berry-Esseen theorem to bound the class-conditional
activation probabilities of ReLU for the neurons in those subsets. By a
careful treatment of the activation probabilities as the neurons evolve during
training, we can bound the expected gradients for each neuron in
$\mathcal{N}_{y}^{(t)}$ at every step $t$, hence iteratively tracking its
weight updates throughout training. This treatment allows us to characterize
the output of the network up to constant factors while avoiding the nuisance
of analyzing the activation probability of _every_ neuron in the network,
which turns out to be very challenging. For ease of presentation, in the
sequel we separate our main results into four parts and introduce them
progressively, with an illustration of our key ideas in Figure 2. Complete
proofs of all theoretical results are deferred to Appendix I.

_1\. Neuron activation is asymmetric._ Our key insight is that during
training, every neuron in $\mathcal{N}_{y}^{(t)}$ has the incentive to be
positively correlated with the examples from at most one class
${\mathbf{y}}_{\mathrm{pos}}=y$ (whether $y=1$ or $y=-1$ depends on the random
initialization of the neuron); we refer to those examples as _positive
examples_
$({\mathbf{x}}_{\mathrm{pos}},{\mathbf{y}}_{\mathrm{pos}})\sim\mathcal{D}_{\mathrm{train}}|{\mathbf{y}}={\mathbf{y}}_{\mathrm{pos}}$
for that neuron. Correspondingly, we refer to examples from the other class
${\mathbf{y}}_{\mathrm{neg}}=-y$ as _negative examples_
$({\mathbf{x}}_{\mathrm{neg}},{\mathbf{y}}_{\mathrm{neg}})\sim\mathcal{D}_{\mathrm{train}}|{\mathbf{y}}={\mathbf{y}}_{\mathrm{neg}}$
for the neuron. Due to randomness at initialization, we can show that
$|\mathcal{N}_{y}^{(0)}|=\Theta(m)$ for both $y\in\\{-1,1\\}$ and, after
sufficient SGD iterations, all neurons in $\mathcal{N}_{y}^{(t)}$ will
accumulate (in expectation) positive correlations with examples from $y$ and
negative correlations with examples from $-y$, resulting in class-wise
asymmetry in their activation as shown by Theorem 4.1.

######  Theorem 4.1 (Activation asymmetry).

For every $\eta\leq\frac{1}{\mathsf{poly}(d_{0})}$ and every
$y\in\mathcal{Y}$, there exists
$T_{0}=\widetilde{\Theta}(\frac{m}{\eta\sqrt{d}})$ such that w.h.p., for every
$t\geq T_{0}$, there exist $\Theta(m)$ neurons in which the weight
${\mathbf{w}}_{k}^{(t)}$ for each neuron satisfies:

$\displaystyle\mathbf{Pr}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0]}=1-O\left(d_{0}^{-\frac{1}{2}}\right),$ (5)
$\displaystyle\mathbf{Pr}_{{\mathbf{x}}|{\mathbf{y}}=-y\sim\mathcal{D}_{\mathrm{train}}}{[\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0]}=o(1).$

Figure 4: _Class-averaged_ activation rate histograms of a randomly
initialized CLIP-RN50 (left) and a distilled CLIP-RN50 (right). After
training, more classes have smaller average activation rates close to zero and
only a small number of classes have large average activation rates.

_2\. Activation asymmetry leads to feature contamination._ Note that for every
$k\in[m]$, the weight vector of the $k$-th neuron (we will also refer to it as
the _learned feature_ of the neuron) after $t$ iterations can be written as

${\mathbf{w}}_{k}^{(t)}=\sum_{j\in\mathcal{S}_{\mathrm{core}}}\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{j}}\rangle\bm{m}_{j}+\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{j}}\rangle\bm{m}_{j}+\mathrm{res},$
(6)

where the residual term satisfies
$\langle{\mathrm{res}},{\bm{m}_{j}}\rangle=0$ for every $j\in[d_{0}]$ and thus
can be neglected. Intuitively, Eq. (6) indicates that the learned feature can
be decomposed into its projections onto different feature vectors. Meanwhile,
as we will prove in Lemma A.3, at iteration $t$, the gradient projection onto
background features for every neuron $k\in\mathcal{N}_{y}^{(t)}$ satisfies:
for every $j\in\mathcal{S}_{\mathrm{bg}}$,

$\displaystyle\langle{-\nabla_{{\mathbf{w}}_{k}^{(t)}}\widehat{\mathcal{L}}(h^{(t)})},{\bm{m}_{j}}\rangle\propto$
(7)
$\displaystyle\qquad\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})}{(\mathbf{1}_{{\mathbf{y}}={\mathbf{y}}_{\mathrm{pos}}}-\mathbf{1}_{{\mathbf{y}}={\mathbf{y}}_{\mathrm{neg}}})\mathbf{1}_{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0}{\mathbf{z}}_{j}}.$

By Theorem 4.1, we then have that for at least $\Theta(m)$ neurons in
$\mathcal{N}_{y}^{(t)}$,
$\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}={\mathbf{y}}_{\mathrm{pos}}}{\mathbf{1}_{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0}}$ would be much larger than
$\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}={\mathbf{y}}_{\mathrm{neg}}}{\mathbf{1}_{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0}}$, resulting in a positive gradient projection onto _every_ background
feature $\bm{m}_{j}$ _regardless of its correlation with the label_. We refer
to this feature learning proclivity of neural networks as feature
contamination. Formally, Theorem 4.2 shows that this will result in the
neurons’ learned features accumulating both correlated core features and
_uncorrelated_ background features.

######  Theorem 4.2 (Learned features).

For every $\eta\leq\frac{1}{\mathsf{poly}(d_{0})}$ and every
$y\in\mathcal{Y}$, there exists $T_{1}=\Theta(\frac{m}{\eta d_{0}})$ such that
w.h.p., after $T_{1}$ iterations, there exist $\Theta(m)$ neurons in which the
weight ${\mathbf{w}}_{k}^{(T_{1})}$ for each neuron satisfies the following:

$\displaystyle\sum\nolimits_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j1}\langle{{\mathbf{w}}_{k}^{(T_{1})}},{\bm{m}_{j}}\rangle=y\cdot\Theta(1),$
(8)
$\displaystyle\sum\nolimits_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j1}\langle{{\mathbf{w}}_{k}^{(T_{1})}},{\bm{m}_{j}}\rangle=\widetilde{\Theta}(1).$

Figure 5: Average _neuron selectivity_ of random and distilled CLIP-RN50
(left) and CLIP-ViT-B/16 (right) models. Distilled models have larger
selectivity compared with random models and exhibit a selectivity drop in OOD
data. Please refer to Section F.3 for more details.  Figure 6: Empirical
evidence for feature contamination in _deep_ neural networks. (a) Grad-CAM
results for CLIP-RN50 and a distilled CLIP-RN50 on DomainNet. (b) Grad-CAM
results for CLIP-RN50 and a distilled CLIP-RN50 on controlled distribution
shifts. For the distilled model, while the weights of core objects are
dominant for ID images, they are _reduced_ under distribution shifts. (c)
Example images of the raw CIFAR-10 dataset and our modified version with
background color features that are _uncorrelated_ with the label, and the
visualization of deep features learned by a ResNet on both datasets.

_3\. Feature contamination induces large OOD risk._ Intuitively, this result
is a direct consequence of the _coupling_ of core features and background
features in the neurons’ pre-activation as shown by Theorem 4.2. With this
coupling, negative shifts of background features can _reduce the activation of
the neuron_ , resulting in OOD risk. In extreme cases, when the pre-activation
of a neuron is reduced to negative, the contribution of the core features
extracted by the neuron will also diminish. Formally, Theorem 4.3 quantifies
such impact of feature contamination on ID and OOD risks.

######  Theorem 4.3 (ID and OOD risks).

For every $\eta\leq\frac{1}{\mathsf{poly}(d_{0})}$, there exists
$T_{2}=\widetilde{\Theta}(\frac{m}{\eta d_{0}})$ such that w.h.p., after
$T_{2}$ iterations, the trained model $h^{(T_{2})}$ satisfies the following:

$\mathcal{R}_{\mathcal{D}_{\mathrm{train}}}(h^{(T_{2})})\leq o(1),\
\mathcal{R}_{\mathrm{OOD}}(h^{(T_{2})})=\widetilde{\Theta}(1).$ (9)

_4\. Linear networks are provably free from feature contamination._ Finally,
to further understand the role of non-linearity, we prove that if we “remove”
the non-linearity in the network by replacing each ReLU with the identity
function, then feature contamination will no longer occur.

######  Theorem 4.4 (Linear networks).

If we replace the ReLU functions in the network with identity functions and
keep other conditions the same as in Theorem 4.2, then with high probability,
we have
$|\langle{{\mathbf{w}}_{k}^{(T_{1})}},{\bm{m}_{j}}\rangle|\leq\widetilde{O}(\frac{1}{\sqrt{d}})$
for every $k\in[m]$ and every $j\in\mathcal{S}_{\mathrm{bg}}$.

The main intuition of Theorem 4.4 is that without non-linearity, the
activation magnitude for the examples from different classes would be no
longer asymmetric: for two-layer linear networks, the gradient projection onto
background features is akin to Eq. (7) but without the activation derivative
$\mathbf{1}_{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq 0}$. We
then have
$\langle{-\nabla_{{\mathbf{w}}_{k}^{(t)}}\widehat{\mathcal{L}}(h^{(t)})},{\bm{m}_{j}}\rangle\approx
0$ for every $j\in\mathcal{S}_{\mathrm{bg}}$ since positive gradients and
negative gradients on ${\mathbf{z}}_{j}$ will now cancel out. As a result, the
background features will not be accumulated during SGD.

Numerical results. As empirical evidence that corroborates our theory, we
provide numerical results in Figure 3 that demonstrate the existence of
activation asymmetry and feature contamination under our setting.

Extensions to more general settings. To show that feature contamination also
occurs _beyond_ our theoretical setting, we also conduct numerical experiments
with several relaxations from our setting, including (i) _non-linear_
relationships between core features and the label, (ii) more activation
functions other than ReLU, and (iii) optimizers with adaptive step sizes other
than vanilla SGD. As shown by Figure 3(d), feature contamination consistently
occurs in more general settings. In Section F.1, we also provide numerical
results in regression (representation distillation) tasks to show that feature
contamination also occurs beyond classification.

##  5 Feature Contamination in Practice

In this section, we present empirical evidence that connects our theoretical
results to practical OOD generalization failure and discuss possible
solutions.

Activation asymmetry in deep networks. In our theoretical model, feature
contamination stems from the asymmetric activation of the neurons. To examine
whether deep networks trained on real-world data also exhibit this behavior,
we compute the _class-averaged_ activation rate histograms of the ResNet and
ViT models trained in our experiments in Section 2. As shown in Figure 4 and
Section F.2, both models exhibit activation asymmetry after training, with low
average activation rates that are close to zero for most classes and high
average activation rates for only a small number of classes. Moreover, we also
adopt a more quantitative metric termed _neuron selectivity_ that measures the
difference of neuron activation magnitudes of different classes. As shown in
Figure 5, distilled models have considerably larger average neuron selectivity
than random models, which corroborates our theory. Please refer to Section F.3
for the implementation details of computing neuron selectivity.

Prediction heatmap visualization. In Figure 6(a), we visualize the prediction
heatmaps of CLIP-RN50 and a distilled CLIP-RN50 in the DomainNet dataset using
Grad-CAM (Chattopadhay et al., 2018). An intriguing phenomenon revealed by the
heatmaps is that for the distilled model, while the weights of core objects
are dominant for ID images, they are _reduced_ under distribution shifts,
resulting in OOD generalization failure. Quantitatively, a similar observation
is shown in Figure 5, where the selectivity of neurons drops in OOD data. Here
we argue that _feature contamination explains this phenomenon_ : the Grad-CAM
score of each feature map is calculated by differentiating the classification
score with respect to the feature map, thus being proportional to the
corresponding neurons’ activation. Hence, when the core features and the
background features are coupled in the neurons’ pre-activation (Theorem 4.2),
the shift of background features can reduce the activation, which in turn
reduces the Grad-CAM score of the feature map as visually observed. The
reduction in neuron selectivity in OOD data can also be explained in a similar
way.

Prediction heatmaps under controlled distribution shifts. In Figure 6(b), we
consider a synthetic OOD generalization task based on image style transfer
where we can manually control the “degree” of distribution shifts by
controlling the amount of style change. This setting closely matches our data
model in Definition 3.1 by keeping core features intact while _only_ changing
background features in the ID images. As shown in the figure, the prediction
heatmaps exhibit visually similar patterns as the heatmaps of natural OOD
images. This implies that our data model indeed captures some key
characteristics of real-world distribution shifts and backs the explanation of
feature contamination.

Visualizing contaminated deep features. To qualitatively show the impact of
feature contamination on the learned features of _deep_ neural networks, we
visualize the features learned by a ResNet on a modified CIFAR-10 dataset in
Figure 6(c), where the original images are padded with background red pixels
_uncorrelated_ with labels. Compared to the original CIFAR-10 dataset, the
learned features on the modified CIFAR-10 dataset exhibits evident color
differences, indicating that feature contamination also occurs in _deep
features_. See Section F.4 for more details and results.

Discussion on possible solutions. In our theoretical model, although gradient
descent accumulates both core and background features in the weight space due
to feature contamination, there still exists a _subspace_ (corresponding to
the span of core features) where background features are not accumulated.
Hence, constraining the SGD updates to this subspace would possibly lead to
ideal generalization. This is consistent with prior results (Idnani et al.,
2023) showing that projecting the network’s intermediate representations onto
certain subspaces may improve OOD generalization. However, how to effectively
find the correct subspace for projection without the explicit access to the
core and background feature subspaces remains an open problem.

##  6 Conclusion

In this section, we discuss potential implications of our results and list
important future directions.

Takeaway 1: OOD generalization algorithms need to consider inductive biases.
Many existing studies on OOD generalization motivate and analyze their
algorithms using linear or unstructured models that do not capture the
inductive biases of neural networks. Our results imply that OOD generalization
may not be feasible without considering such inductive biases, calling for
explicitly incorporating them into principled algorithm design.

Takeaway 2: Non-linearity in neural networks elicits new OOD generalization
challenges. As we formally show in Section 4, feature contamination
essentially stems from the gradient descent optimization process of non-linear
neural networks even with _uncorrelated_ background features, thus being
orthogonal to the prevailing narrative of spurious correlations. This provides
a new perspective on OOD generalization and may inspire new algorithmic
design.

Takeaway 3: Learned features may behave very differently from prescribed ones.
Many existing studies on OOD generalization explicitly or implicitly assume
that we can directly work on a set of _well-separated_ core/spurious features.
While this assumption helps build intuitions, our results highlight that it
can also be misleading since the features _learned_ by neural networks may
manifest in a _non-linearly coupled_ manner, thus often diverging from the
intuitions for prescribed, well-separated features.

###  6.1 Limitations, a Conjecture, and Future Work

While this work takes a step towards fully understanding OOD generalization in
practice, our results still leave much room for improvement such as extensions
to more general data distributions, multi-class classification, and more
complicated network architectures. Meanwhile, while our current setup focuses
on training from scratch, we envision that the viewpoint of feature
contamination may also be helpful in analyzing the effect of _pre-training_ on
OOD generalization. In particular, we have the following conjecture:

######  Conjecture.

Pre-training on a sufficiently diverse dataset does not remove uncorrelated
features, but _linearizes_ those features in the model’s representations,
hence mitigating feature contamination and improving OOD generalization.

We provide preliminary empirical evidence that supports the above conjecture
in Section G, as well as more discussion on related empirical observations in
recent work (Gandelsman et al., 2024; Mayilvahanan et al., 2024). Yet, we
believe that rigorously proving this conjecture requires a more fine-grained
treatment for the pre-training data distribution and the SGD dynamics and thus
leave it as future work.

## Acknowledgements

This work was supported in part by the National Key Research and Development
Program of China under STI 2030-Major Projects 2021ZD0200300, and in part by
the National Natural Science Foundation of China under Grant 62176133, and in
part by the Tsinghua-Guoqiang research program under Grant 2019GQG0006.

## Impact Statement

This paper presents work that aims to advance our understanding of the feature
learning process of neural networks and its impact on generalization under
distribution shifts, which may benefit building machine learning models that
are more generalizable, robust, and trustworthy.

## References

  * Abbe et al. (2023) Abbe, E., Bengio, S., Lotfi, A., and Rizk, K.  Generalization on the unseen, logic reasoning and degree curriculum.  In _International Conference on Machine Learning_ , 2023.
  * Ahuja et al. (2021a) Ahuja, K., Caballero, E., Zhang, D., Bengio, Y., Mitliagkas, I., and Rish, I.  Invariance principle meets information bottleneck for out-of-distribution generalization.  In _Advances in Neural Information Processing Systems_ , pp. 3438–3450, 2021a.
  * Ahuja et al. (2021b) Ahuja, K., Wang, J., Dhurandhar, A., Shanmugam, K., and Varshney, K. R.  Empirical or invariant risk minimization? A sample complexity perspective.  In _ICLR_ , 2021b.
  * Allen-Zhu & Li (2021) Allen-Zhu, Z. and Li, Y.  Feature purification: How adversarial training performs robust deep learning.  _arXiv preprint arXiv:2005.10190_ , 2021.
  * Allen-Zhu & Li (2023a) Allen-Zhu, Z. and Li, Y.  Physics of language models: Part 3.1, knowledge storage and extraction.  _arXiv preprint arXiv:2309.14316_ , 2023a.
  * Allen-Zhu & Li (2023b) Allen-Zhu, Z. and Li, Y.  Towards understanding ensemble, knowledge distillation and self-distillation in deep learning.  In _International Conference on Learning Representations_ , 2023b.
  * Amodei et al. (2016) Amodei, D., Olah, C., Steinhardt, J., Christiano, P., Schulman, J., and Mané, D.  Concrete problems in AI safety.  _arXiv preprint arXiv:1606.06565_ , 2016.
  * Arjovsky et al. (2019) Arjovsky, M., Bottou, L., Gulrajani, I., and Lopez-Paz, D.  Invariant risk minimization.  _arXiv preprint arXiv:1907.02893_ , 2019.
  * Arpit et al. (2017) Arpit, D., Jastrzebski, S., Ballas, N., Krueger, D., Bengio, E., Kanwal, M. S., Maharaj, T., Fischer, A., Courville, A., Bengio, Y., and Lacoste-Julien, S.  A closer look at memorization in deep networks.  In _International Conference on Machine Learning_ , pp. 233–242, 2017.
  * Barbu et al. (2019) Barbu, A., Mayo, D., Alverio, J., Luo, W., Wang, C., Gutfreund, D., Tenenbaum, J., and Katz, B.  ObjectNet: A large-scale bias-controlled dataset for pushing the limits of object recognition models.  In _Advances in Neural Information Processing Systems_ , pp. 9448–9458, 2019.
  * Beery et al. (2018) Beery, S., Van Horn, G., and Perona, P.  Recognition in terra incognita.  In _ECCV_ , volume 11220, pp. 472–489, 2018.
  * Bommasani et al. (2022) Bommasani, R., Hudson, D. A., Adeli, E., Altman, R., Arora, S., von Arx, S., Bernstein, M. S., Bohg, J., Bosselut, A., Brunskill, E., Brynjolfsson, E., Buch, S., Card, D., Castellon, R., Chatterji, N., Chen, A., Creel, K., Davis, J. Q., Demszky, D., Donahue, C., Doumbouya, M., Durmus, E., Ermon, S., Etchemendy, J., Ethayarajh, K., Fei-Fei, L., Finn, C., Gale, T., Gillespie, L., Goel, K., Goodman, N., Grossman, S., Guha, N., Hashimoto, T., Henderson, P., Hewitt, J., Ho, D. E., Hong, J., Hsu, K., Huang, J., Icard, T., Jain, S., Jurafsky, D., Kalluri, P., Karamcheti, S., Keeling, G., Khani, F., Khattab, O., Koh, P. W., Krass, M., Krishna, R., Kuditipudi, R., Kumar, A., Ladhak, F., Lee, M., Lee, T., Leskovec, J., Levent, I., Li, X. L., Li, X., Ma, T., Malik, A., Manning, C. D., Mirchandani, S., Mitchell, E., Munyikwa, Z., Nair, S., Narayan, A., Narayanan, D., Newman, B., Nie, A., Niebles, J. C., Nilforoshan, H., Nyarko, J., Ogut, G., Orr, L., Papadimitriou, I., Park, J. S., Piech, C., Portelance, E., Potts, C., Raghunathan, A., Reich, R., Ren, H., Rong, F., Roohani, Y., Ruiz, C., Ryan, J., Ré, C., Sadigh, D., Sagawa, S., Santhanam, K., Shih, A., Srinivasan, K., Tamkin, A., Taori, R., Thomas, A. W., Tramèr, F., Wang, R. E., Wang, W., Wu, B., Wu, J., Wu, Y., Xie, S. M., Yasunaga, M., You, J., Zaharia, M., Zhang, M., Zhang, T., Zhang, X., Zhang, Y., Zheng, L., Zhou, K., and Liang, P.  On the opportunities and risks of foundation models.  _arXiv preprint arXiv:2108.07258_ , 2022.
  * Brown et al. (2020) Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert-Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D. M., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D.  Language models are few-shot learners.  In _Advances in Neural Information Processing Systems_ , 2020\.
  * Chattopadhay et al. (2018) Chattopadhay, A., Sarkar, A., Howlader, P., and Balasubramanian, V. N.  Grad-cam++: Generalized gradient-based visual explanations for deep convolutional networks.  In _2018 IEEE winter conference on applications of computer vision (WACV)_ , pp. 839–847. IEEE, 2018.
  * Chen et al. (2022) Chen, Y., Rosenfeld, E., Sellke, M., Ma, T., and Risteski, A.  Iterative feature matching: Toward provable domain generalization with logarithmic environments.  In _Advances in Neural Information Processing Systems_ , 2022\.
  * Chen et al. (2023) Chen, Y., Huang, W., Zhou, K., Bian, Y., Han, B., and Cheng, J.  Understanding and improving feature learning for out-of-distribution generalization.  In _Advances in Neural Information Processing Systems_ , 2023\.
  * DeGrave et al. (2021) DeGrave, A. J., Janizek, J. D., and Lee, S.-I.  AI for radiographic COVID-19 detection selects shortcuts over signal.  _Nature Machine Intelligence_ , 3(7):610–619, 2021.  ISSN 2522-5839.
  * Deng et al. (2009) Deng, J., Dong, W., Socher, R., Li, L.-J., Li, K., and Fei-Fei, L.  Imagenet: A large-scale hierarchical image database.  In _CVPR_ , pp. 248–255, 2009.
  * Elhage et al. (2022) Elhage, N., Hume, T., Olsson, C., Schiefer, N., Henighan, T., Kravec, S., Hatfield-Dodds, Z., Lasenby, R., Drain, D., and Chen, C.  Toy models of superposition.  _arXiv preprint arXiv:2209.10652_ , 2022.
  * Frankle & Carbin (2019) Frankle, J. and Carbin, M.  The lottery ticket hypothesis: Finding sparse, trainable neural networks.  In _ICLR_ , 2019.
  * Gandelsman et al. (2024) Gandelsman, Y., Efros, A. A., and Steinhardt, J.  Interpreting CLIP’s image representation via text-based decomposition.  In _International Conference on Learning Representations_ , 2024\.
  * Ganin et al. (2016) Ganin, Y., Ustinova, E., Ajakan, H., Germain, P., Larochelle, H., Laviolette, F., Marchand, M., and Lempitsky, V.  Domain-adversarial training of neural networks.  _Journal of Machine Learning Research_ , 17(59):1–35, 2016.
  * Geirhos et al. (2018) Geirhos, R., Temme, C. R. M., Rauber, J., Schütt, H. H., Bethge, M., and Wichmann, F. A.  Generalisation in humans and deep neural networks.  In _Advances in Neural Information Processing Systems_ , pp. 7549–7561, 2018.
  * Geirhos et al. (2020) Geirhos, R., Jacobsen, J.-H., Michaelis, C., Zemel, R., Brendel, W., Bethge, M., and Wichmann, F. A.  Shortcut learning in deep neural networks.  _Nature Machine Intelligence_ , 2(11):665–673, 2020.  ISSN 2522-5839.
  * Gould et al. (2023) Gould, R., Ong, E., Ogden, G., and Conmy, A.  Successor heads: Recurring, interpretable attention heads in the wild.  _arXiv preprint arXiv:2312.09230_ , 2023.
  * Gulrajani & Lopez-Paz (2021) Gulrajani, I. and Lopez-Paz, D.  In search of lost domain generalization.  In _ICLR_ , 2021.
  * Gurnee & Tegmark (2024) Gurnee, W. and Tegmark, M.  Language models represent space and time.  In _International Conference on Learning Representations_ , 2024\.
  * HaoChen et al. (2022) HaoChen, J. Z., Wei, C., Kumar, A., and Ma, T.  Beyond separability: Analyzing the linear transferability of contrastive representations to related subpopulations.  _arXiv preprint arXiv:2204.02683_ , 2022.
  * Heinzerling & Inui (2024) Heinzerling, B. and Inui, K.  Monotonic representation of numeric properties in language models.  _arXiv preprint arXiv:2403.10381_ , 2024.
  * Hendrycks & Gimpel (2016) Hendrycks, D. and Gimpel, K.  Gaussian error linear units (GELUs).  _arXiv preprint arXiv:1606.08415_ , 2016.
  * Hendrycks et al. (2021a) Hendrycks, D., Basart, S., Mu, N., Kadavath, S., Wang, F., Dorundo, E., Desai, R., Zhu, T., Parajuli, S., Guo, M., Song, D., Steinhardt, J., and Gilmer, J.  The many faces of robustness: A critical analysis of out-of-distribution generalization.  In _ICCV_ , 2021a.
  * Hendrycks et al. (2021b) Hendrycks, D., Zhao, K., Basart, S., Steinhardt, J., and Song, D.  Natural adversarial examples.  In _CVPR_ , 2021b.
  * Hinton et al. (2014) Hinton, G., Vinyals, O., and Dean, J.  Distilling the knowledge in a neural network.  In _Advances in Neural Information Processing Systems Deep Learning Workshop_ , 2014.
  * Huang et al. (2020) Huang, Z., Wang, H., Xing, E. P., and Huang, D.  Self-challenging improves cross-domain generalization.  In _European Conference on Computer Vision_ , pp. 124–140, 2020.
  * Idnani et al. (2023) Idnani, D., Madan, V., Goyal, N., Schwab, D. J., and Vedantam, S. R.  Don’t forget the nullspace! Nullspace occupancy as a mechanism for out of distribution failure.  2023\.
  * Kamath et al. (2021) Kamath, P., Tangella, A., Sutherland, D. J., and Srebro, N.  Does invariant risk minimization capture invariance?  In _AISTATS_ , 2021.
  * Karp et al. (2021) Karp, S., Winston, E., Li, Y., and Singh, A.  Local signal adaptivity: Provable feature learning in neural networks beyond kernels.  In _Advances in Neural Information Processing Systems_ , pp. 24883–24897, 2021.
  * Kim et al. (2021) Kim, D., Yoo, Y., Park, S., Kim, J., and Lee, J.  SelfReg: Self-supervised contrastive regularization for domain generalization.  In _2021 IEEE/CVF International Conference on Computer Vision (ICCV)_ , pp. 9599–9608, 2021.
  * Kingma & Ba (2015) Kingma, D. P. and Ba, J.  Adam: A method for stochastic optimization.  In _International Conference on Learning Representations_ , 2015\.
  * Koh et al. (2021) Koh, P. W., Sagawa, S., Marklund, H., Xie, S. M., Zhang, M., Balsubramani, A., Hu, W., Yasunaga, M., Phillips, R. L., Gao, I., Lee, T., David, E., Stavness, I., Guo, W., Earnshaw, B. A., Haque, I. S., Beery, S., Leskovec, J., Kundaje, A., Pierson, E., Levine, S., Finn, C., and Liang, P.  Wilds: A benchmark of in-the-wild distribution shifts.  In _ICML_ , 2021.
  * Krueger et al. (2021) Krueger, D., Caballero, E., Jacobsen, J.-H., Zhang, A., Binas, J., Zhang, D., Priol, R. L., and Courville, A.  Out-of-distribution generalization via risk extrapolation (rex).  In _ICML_ , 2021.
  * Kumar et al. (2022) Kumar, A., Raghunathan, A., Jones, R., Ma, T., and Liang, P.  Fine-tuning can distort pretrained features and underperform out-of-distribution.  In _International Conference on Learning Representations_ , 2022\.
  * Lee et al. (2023) Lee, Y., Chen, A. S., Tajwar, F., Kumar, A., Yao, H., Liang, P., and Finn, C.  Surgical fine-tuning improves adaptation to distribution shifts.  In _International Conference on Learning Representations_ , 2023\.
  * Li & Flanigan (2023) Li, C. and Flanigan, J.  Task contamination: Language models may not be few-shot anymore.  _arXiv preprint arXiv:2312.16337_ , 2023.
  * Li et al. (2017) Li, D., Yang, Y., Song, Y.-Z., and Hospedales, T. M.  Deeper, broader and artier domain generalization.  In _2017 IEEE International Conference on Computer Vision (ICCV)_ , pp. 5543–5551, 2017.
  * Li et al. (2018) Li, D., Yang, Y., Song, Y.-Z., and Hospedales, T. M.  Learning to generalize: Meta-learning for domain generalization.  In _AAAI_ , 2018.
  * Liu et al. (2023) Liu, H., Ning, R., Teng, Z., Liu, J., Zhou, Q., and Zhang, Y.  Evaluating the logical reasoning ability of chatgpt and gpt-4.  _arXiv preprint arXiv:2304.03439_ , 2023.
  * Loshchilov & Hutter (2017) Loshchilov, I. and Hutter, F.  SGDR: Stochastic gradient descent with warm restarts.  In _International Conference on Learning Representations_ , 2017\.
  * Loshchilov & Hutter (2019) Loshchilov, I. and Hutter, F.  Decoupled weight decay regularization.  In _International Conference on Learning Representations_ , 2019\.
  * Marks & Tegmark (2023) Marks, S. and Tegmark, M.  The geometry of truth: Emergent linear structure in large language model representations of true/false datasets.  _arXiv preprint arXiv:2310.06824_ , 2023.
  * Mayilvahanan et al. (2024) Mayilvahanan, P., Wiedemer, T., Rusak, E., Bethge, M., and Brendel, W.  Does CLIP’s generalization performance mainly stem from high train-test similarity?  In _International Conference on Learning Representations_ , 2024\.
  * Miller et al. (2021) Miller, J., Taori, R., Raghunathan, A., Sagawa, S., Koh, P. W., Shankar, V., Liang, P., Carmon, Y., and Schmidt, L.  Accuracy on the line: On the strong correlation between out-of-distribution and in-distribution generalization.  In _ICML_ , 2021.
  * Mitrovic et al. (2021) Mitrovic, J., McWilliams, B., Walker, J. C., Buesing, L. H., and Blundell, C.  Representation learning via invariant causal mechanisms.  In _ICLR_ , 2021.
  * Morcos et al. (2018) Morcos, A. S., Barrett, D. G. T., Rabinowitz, N. C., and Botvinick, M.  On the importance of single directions for generalization.  In _International Conference on Learning Representations_ , 2018\.  URL http://arxiv.org/abs/1803.06959.
  * Nagarajan et al. (2021) Nagarajan, V., Andreassen, A., and Neyshabur, B.  Understanding the failure modes of out-of-distribution generalization.  In _ICLR_ , 2021.
  * Nam et al. (2021) Nam, H., Lee, H., Park, J., Yoon, W., and Yoo, D.  Reducing domain gap by reducing style bias.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 8690–8699, 2021.
  * Oliveira (2010) Oliveira, R. I.  Concentration of the adjacency matrix and of the laplacian in random graphs with independent edges.  _arXiv preprint arXiv:0911.0600_ , 2010.
  * Park et al. (2023) Park, K., Choe, Y. J., and Veitch, V.  The linear representation hypothesis and the geometry of large language models.  _arXiv preprint arXiv:2311.03658_ , 2023.
  * Peng et al. (2019) Peng, X., Bai, Q., Xia, X., Huang, Z., Saenko, K., and Wang, B.  Moment matching for multi-source domain adaptation.  In _ICCV_ , pp. 1406–1415, 2019.
  * Pezeshki et al. (2021) Pezeshki, M., Kaba, S.-O., Bengio, Y., Courville, A., Precup, D., and Lajoie, G.  Gradient starvation: A learning proclivity in neural networks.  In _Advances in Neural Information Processing Systems_ , pp. 1256–1272, 2021.
  * Radford et al. (2021) Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., Krueger, G., and Sutskever, I.  Learning transferable visual models from natural language supervision.  In _International Conference on Machine Learning_ , pp. 8748–8763, 2021.
  * Recht et al. (2019) Recht, B., Roelofs, R., Schmidt, L., and Shankar, V.  Do imagenet classifiers generalize to imagenet?  In _ICML_ , 2019.
  * Rosenfeld et al. (2021) Rosenfeld, E., Ravikumar, P., and Risteski, A.  The risks of invariant risk minimization.  In _ICLR_ , 2021.
  * Rosenfeld et al. (2022) Rosenfeld, E., Ravikumar, P., and Risteski, A.  Domain-adjusted regression or: ERM may already learn features sufficient for out-of-distribution generalization.  _arXiv preprint arXiv:2202.06856_ , 2022.
  * Sagawa et al. (2020a) Sagawa, S., Koh, P. W., Hashimoto, T. B., and Liang, P.  Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization.  In _ICLR_ , 2020a.
  * Sagawa et al. (2020b) Sagawa, S., Raghunathan, A., Koh, P. W., and Liang, P.  An investigation of why overparameterization exacerbates spurious correlations.  In _ICML_ , 2020b.
  * Salman et al. (2019) Salman, H., Li, J., Razenshteyn, I., Zhang, P., Zhang, H., Bubeck, S., and Yang, G.  Provably robust deep learning via adversarially trained smoothed classifiers.  In _Advances in Neural Information Processing Systems_ , volume 32, 2019.
  * Schölkopf et al. (2021) Schölkopf, B., Locatello, F., Bauer, S., Ke, N. R., Kalchbrenner, N., Goyal, A., and Bengio, Y.  Toward causal representation learning.  _Proceedings of the IEEE_ , 109(5):612–634, 2021\.  ISSN 1558-2256.
  * Shah et al. (2020) Shah, H., Tamuly, K., and Raghunathan, A.  The pitfalls of simplicity bias in neural networks.  In _Advances in Neural Information Processing Systems_ , 2020\.
  * Shen et al. (2022) Shen, K., Jones, R., Kumar, A., Xie, S. M., HaoChen, J. Z., Ma, T., and Liang, P.  Connect, not collapse: Explaining contrastive learning for unsupervised domain adaptation.  In _International Conference on Machine Learning_ , volume 19847-19878, 2022.
  * Shi et al. (2022) Shi, Y., Seely, J., Torr, P. H. S., Siddharth, N., Hannun, A., Usunier, N., and Synnaeve, G.  Gradient matching for domain generalization.  In _International Conference on Learning Representations_ , 2022\.
  * Singh et al. (2022) Singh, M., Gustafson, L., Adcock, A., De Freitas Reis, V., Gedik, B., Kosaraju, R. P., Mahajan, D., Girshick, R., Dollar, P., and Van Der Maaten, L.  Revisiting weakly supervised pre-training of visual perception models.  In _IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pp. 794–804, 2022.
  * Sun & Saenko (2016) Sun, B. and Saenko, K.  Deep coral: Correlation alignment for deep domain adaptation.  In _European conference on computer vision_ , pp. 443–450, 2016\.
  * Tan et al. (2020) Tan, S., Peng, X., and Saenko, K.  Class-imbalanced domain adaptation: An empirical odyssey.  _arXiv preprint arXiv:1910.10320_ , 2020.
  * Taori et al. (2020) Taori, R., Dave, A., Shankar, V., Carlini, N., Recht, B., and Schmidt, L.  Measuring robustness to natural distribution shifts in image classification.  In _Advances in Neural Information Processing Systems_ , pp. 18583–18599, 2020.
  * Tian et al. (2020) Tian, Y., Krishnan, D., and Isola, P.  Contrastive representation distillation.  In _International Conference on Learning Representations_ , 2020\.
  * Torralba & Efros (2011) Torralba, A. and Efros, A. A.  Unbiased look at dataset bias.  In _CVPR_ , pp. 1521–1528, 2011.
  * Tripuraneni et al. (2020) Tripuraneni, N., Jordan, M. I., and Jin, C.  On the theory of transfer learning: The importance of task diversity.  In _Advances in Neural Information Processing Systems_ , pp. 7852–7862, 2020.
  * Tropp (2012) Tropp, J. A.  User-friendly tail bounds for sums of random matrices.  _Foundations of computational mathematics_ , 12:389–434, 2012.
  * Vapnik (1999) Vapnik, V.  _The nature of statistical learning theory_.  1999\.
  * Wang et al. (2019) Wang, H., Ge, S., Xing, E. P., and Lipton, Z. C.  Learning robust global representations by penalizing local predictive power.  In _Advances in Neural Information Processing Systems_ , pp. 10506–10518, 2019.
  * Wang et al. (2022) Wang, H., Si, H., Li, B., and Zhao, H.  Provable domain generalization via invariant-feature subspace recovery.  In _International Conference on Machine Learning_ , pp. 23018–23033, 2022.
  * Wiles et al. (2022) Wiles, O., Gowal, S., Stimberg, F., Rebuffi, S.-A., Ktena, I., Dvijotham, K., and Cemgil, T.  A fine-grained analysis on distribution shift.  In _ICLR_ , 2022.
  * Wortsman et al. (2022) Wortsman, M., Ilharco, G., Kim, J. W., Li, M., Kornblith, S., Roelofs, R., Lopes, R. G., Hajishirzi, H., Farhadi, A., Namkoong, H., and Schmidt, L.  Robust fine-tuning of zero-shot models.  In _2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pp. 7949–7961, 2022.
  * Xu et al. (2021) Xu, K., Li, J., Zhang, M., Du, S. S., Kawarabayashi, K.-i., and Jegelka, S.  How neural networks extrapolate: From feedforward to graph neural networks.  In _ICLR_ , 2021.
  * Ye et al. (2021) Ye, H., Xie, C., Cai, T., Li, R., Li, Z., and Wang, L.  Towards a theoretical framework of out-of-distribution generalization.  In _Advances in Neural Information Processing Systems_ , 2021\.
  * Zhang et al. (2018) Zhang, H., Cisse, M., Dauphin, Y. N., and Lopez-Paz, D.  Mixup: Beyond empirical risk minimization.  In _ICLR_ , 2018.
  * Zhang & Bottou (2023) Zhang, J. and Bottou, L.  Learning useful representations for shifting tasks and distributions.  In _International Conference on Machine Learning_ , 2023.
  * Zhang et al. (2021) Zhang, M., Marklund, H., Dhawan, N., Gupta, A., Levine, S., and Finn, C.  Adaptive risk minimization: Learning to adapt to domain shift.  In _Advances in Neural Information Processing Systems_ , pp. 23664–23678, 2021.

Appendix

The appendix is divided into two parts for readability. In Appendix I, we
provide complete proofs of our theoretical results. In Appendix II, we present
experimental details and additional empirical results.

Appendix I: Proofs of Theoretical Results

In this part of the appendix, we provide complete proofs of our theorems in
the main text. A quick overview of the structure of this part is as follows:

  * •

In Section A, we introduce the preliminaries and some lemmas that characterize
the neuron properties at random initialization.

  * •

In Section B, we provide the proofs of our main theorems on activation
asymmetry (Theorem 4.1), feature contamination (Theorem 4.2), and ID/OOD risks
(Theorem 4.3).

  * •

In Section C, we provide the proof of Theorem 4.4 on linear neural networks.

  * •

In Section D, we provide the basic probability theory lemmas used in our
proofs for completeness.

##  Appendix A Preliminaries

#### Notation.

Throughout the appendix, we overload $\mathcal{D}_{\mathrm{train}}$ and
$\mathcal{D}_{\mathrm{test}}$ to allow them to denote (joint) training and
test distributions on both $\mathcal{X}\times\mathcal{Y}$ and
$\mathcal{Z}\times\mathcal{Y}$. We also use $\mathcal{D}_{\mathrm{train}}$ and
$\mathcal{D}_{\mathrm{test}}$ to denote the corresponding marginal
distributions on $\mathcal{X}$, $\mathcal{Y}$ and $\mathcal{Z}$. For
presentation brevity, unless otherwise stated, we use the shorthand
$\mathbb{E}_{(\cdot)}{}$ and $\mathbf{Pr}_{(\cdot)}{}$ to denote
$\mathbb{E}_{(\cdot)\sim\mathcal{D}_{\mathrm{train}}}{}$ and
$\mathbf{Pr}_{(\cdot)\sim\mathcal{D}_{\mathrm{train}}}{}$, respectively, and
use the shorthand $h$ to denote $h^{(t)}$ when it is clear from the context.
As in Definition 3.1, we denote the moments of each ${\mathbf{z}}_{j}$ on the
training distribution by
$\mu_{jp}\vcentcolon=\mathbb{E}_{{\mathbf{z}}\sim\mathcal{D}_{\mathrm{train}}}{{\mathbf{z}}_{j}^{p}}$
for every $j\in[d_{0}]$ and $p\in[3]$, and use the shorthand $\mu_{j}$ to
denote $\mu_{j1}$ when it is clear from the context.

###  A.1 Weight Decomposition and Gradient Calculations

We begin by recalling that each weight vector
${\mathbf{w}}_{k}\in\mathbb{R}^{d},k\in[m]$ (i.e., the learned feature of the
$k$-th neuron) in the network can be decomposed into the sum of its
projections to different feature vectors:

${\mathbf{w}}_{k}^{(t)}=\sum_{j\in\mathcal{S}_{\mathrm{core}}}\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{j}}\rangle\bm{m}_{j}+\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{j}}\rangle\bm{m}_{j}+\sum_{j\in[d]\setminus[d_{0}]}\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{j}}\rangle\bm{m}_{j},$
(10)

where $(\bm{m}_{d_{0}+1},\ldots,\bm{m}_{d})$ are an orthogonal complement of
$\bm{M}$. Since all possible inputs are generated to be in
$\mathrm{span}\\{\bm{m}_{1},\ldots,\bm{m}_{d_{0}}\\}$ as in Definition 3.1,
the last term in the RHS of Eq. (10) (i.e., the residual term in Eq. (6) in
the main text) can be neglected due to the orthogonality of different feature
vector $\bm{m}_{j}$s. Therefore, throughout the following analysis, we will
overload the notation ${\mathbf{w}}_{k}^{(t)}$ and let

${\mathbf{w}}_{k}^{(t)}=\sum_{j\in\mathcal{S}_{\mathrm{core}}}\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{j}}\rangle\bm{m}_{j}+\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{j}}\rangle\bm{m}_{j}.$
(11)

A direct consequence of Eq. (10) is that we can analyze the feature learning
process of the network by tracking the correlations between each weight vector
${\mathbf{w}}_{k}^{(t)}$ and different feature vectors $\bm{m}_{j},j\in[d]$ as
the training proceeds. To this end, we need to first analyze the gradient of
each neuron at every iteration.

Gradient of each neuron. Recall that at each iteration $t=0,\ldots,T-1$, the
SGD update for each weight vector ${\mathbf{w}}_{k},k\in[m]$ is given by

${\mathbf{w}}_{k}^{(t+1)}\leftarrow(1-\eta\lambda){\mathbf{w}}_{k}^{(t)}-\eta\nabla_{{\mathbf{w}}_{k}^{(t)}}\frac{1}{N}\sum_{i\in[N]}\ell\left(h^{(t)}({\mathbf{x}}_{i}^{(t)}),{\mathbf{y}}_{i}^{(t)}\right),$

where

$h^{(t)}({\mathbf{x}}_{i}^{(t)})=\sum_{k\in[m]}a_{k}\cdot\mathsf{ReLU}{\big{(}\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}_{i}^{(t)}}\rangle\big{)}}$

and $\ell(y,y^{\prime})=\max\\{1-yy^{\prime},0\\}$. We can then calculate the
gradient of each neuron ${\mathbf{w}}_{k}^{(t)}$ with regard to a certain
example $({\mathbf{x}},{\mathbf{y}})$:

######  Lemma A.1 (Gradient).

For every example $(x,y)\in\mathcal{X}\times\mathcal{Y}$, every $k\in[m]$, and
every iteration $t$, the following holds:

$\nabla_{{\mathbf{w}}_{k}^{(t)}}\ell\left(h(x),y\right)=-a_{k}y\mathbf{1}_{h(x)\leq
1}\mathbf{1}_{\langle{{\mathbf{w}}_{k}^{(t)}},{x}\rangle\geq 0}x.$ (12)

###### Proof.

The proof follows from simple calculation. ∎

We then introduce a lemma that bounds the empirical growth of the correlation
between each neuron ${\mathbf{w}}_{k}^{(t)}$ and each feature vector
$\bm{m}_{j}$ after an SGD update using population gradients.

######  Lemma A.2 (Gap between empirical and population gradients).

For every $k\in[m]$, every $j\in[d]$, and every iteration $t$, if the batch
size $N=\mathsf{poly}(d)$ for some sufficiently large polynomial, then the
following holds with probability at least $1-e^{-\Omega(d)}$:

$\bigg{|}\Big{\langle}{\nabla_{{\mathbf{w}}_{k}^{(t)}}\frac{1}{N}\sum_{i\in[N]}\ell\left(h^{(t)}({\mathbf{x}}_{i}^{(t)}),{\mathbf{y}}_{i}^{(t)}\right)},{\bm{m}_{j}}\Big{\rangle}-\Big{\langle}{\nabla_{{\mathbf{w}}_{k}^{(t)}}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})\sim\mathcal{D}_{\mathrm{train}}}{\ell\left(h({\mathbf{x}}),{\mathbf{y}}\right)}},{\bm{m}_{j}}\Big{\rangle}\bigg{|}\leq\frac{1}{\mathsf{poly}(d)}.$
(13)

###### Proof.

Recall that $\lVert\bm{m}_{j}\rVert_{2}=1$. Applying Cauchy-Schwarz inequality
gives

$\displaystyle\bigg{|}\Big{\langle}{\nabla_{{\mathbf{w}}_{k}^{(t)}}\frac{1}{N}\sum_{i\in[N]}\ell\left(h^{(t)}({\mathbf{x}}_{i}^{(t)}),{\mathbf{y}}_{i}^{(t)}\right)},{\bm{m}_{j}}\Big{\rangle}-\Big{\langle}{\nabla_{{\mathbf{w}}_{k}^{(t)}}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})\sim\mathcal{D}_{\mathrm{train}}}{\ell\left(h({\mathbf{x}}),{\mathbf{y}}\right)}},{\bm{m}_{j}}\Big{\rangle}\bigg{|}$
$\displaystyle\leq\Big{\lVert}\underbrace{\nabla_{{\mathbf{w}}_{k}^{(t)}}\frac{1}{N}\sum_{i\in[N]}\ell\left(h^{(t)}({\mathbf{x}}_{i}^{(t)}),{\mathbf{y}}_{i}^{(t)}\right)-\nabla_{{\mathbf{w}}_{k}^{(t)}}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})\sim\mathcal{D}_{\mathrm{train}}}{\ell\left(h({\mathbf{x}}),{\mathbf{y}}\right)}}_{{\mathbf{S}}^{(t)}}\Big{\rVert}_{2}.$

We define

${\mathbf{Z}}_{i}^{(t)}\vcentcolon=\frac{1}{N}\nabla_{{\mathbf{w}}_{k}^{(t)}}\ell\left(h^{(t)}({\mathbf{x}}_{i}^{(t)}),{\mathbf{y}}_{i}^{(t)}\right)-\frac{1}{N}\nabla_{{\mathbf{w}}_{k}^{(t)}}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})\sim\mathcal{D}_{\mathrm{train}}}{\ell\left(h({\mathbf{x}}),{\mathbf{y}}\right)},\
\forall i\in[N].$

It is easy to see that
${\mathbf{S}}^{(t)}=\sum_{i\in[N]}{\mathbf{Z}}_{i}^{(t)}$,
$\mathbb{E}{{\mathbf{Z}}_{i}^{(t)}}=0$ for every $i\in[N]$, and $\forall i\neq
j\in[N]$, ${\mathbf{Z}}_{i}^{(t)}$ and ${\mathbf{Z}}_{j}^{(t)}$ are
independent. By Lemma A.1, we have

${\mathbf{Z}}_{i}^{(t)}=\frac{1}{N}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})\sim\mathcal{D}_{\mathrm{train}}}{a_{k}{\mathbf{y}}\mathbf{1}_{h({\mathbf{x}})\leq
1}\mathbf{1}_{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0}\cdot{\mathbf{x}}}-\frac{1}{N}a_{k}{\mathbf{y}}_{i}^{(t)}\mathbf{1}_{h^{(t)}({\mathbf{x}}_{i}^{(t)})\leq
1}\mathbf{1}_{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}_{i}^{(t)}}\rangle\geq
0}\cdot{\mathbf{x}}_{i}^{(t)}.$

Recall that $a_{k}\in\\{-\frac{1}{m},\frac{1}{m}\\}$ and ${\mathbf{x}}$ is
generated by
${\mathbf{x}}=\sum_{j\in\mathcal{S}_{\mathrm{core}}}{\mathbf{y}}{\mathbf{z}}_{j}\bm{m}_{j}+\sum_{j\in\mathcal{S}_{\mathrm{bg}}}{\mathbf{z}}_{j}\bm{m}_{j}$
according to Definition 3.1. We then have
$\lVert{\mathbf{Z}}_{i}^{(t)}\rVert_{2}\leq\frac{2\sqrt{d_{0}}}{mN}$, which
also indicates that
$\mathbb{E}{\langle{{\mathbf{Z}}_{i}^{(t)}},{{\mathbf{Z}}_{i}^{(t)}}\rangle}\leq\frac{4d_{0}}{m^{2}N^{2}}$.
This gives

$\mathbb{E}{\langle{{\mathbf{S}}^{(t)}},{{\mathbf{S}}^{(t)}}\rangle}=\sum_{i\in[N]}\mathbb{E}{\langle{{\mathbf{Z}}_{i}^{(t)}},{{\mathbf{Z}}_{i}^{(t)}}\rangle}\leq\frac{4d_{0}}{m^{2}N}.$

Applying matrix Bernstein’s inequality (Lemma D.2), we have

$\mathbf{Pr}{\left[\lVert{\mathbf{S}}^{(t)}\rVert_{2}\geq\delta\right]}\leq(d+1)\exp\left(-\frac{3m^{2}N^{2}\delta^{2}}{24d_{0}+4\sqrt{d_{0}}\delta
mN}\right)$

hold with every $\delta=\frac{1}{\mathsf{poly}(d)}$. Therefore, we have that
for $N=\mathsf{poly}(d)$ with some sufficiently large polynomial, the
following holds with probability at least $1-e^{-\Omega(d)}$:

$\lVert{\mathbf{S}}^{(t)}\rVert_{2}\leq\frac{1}{\mathsf{poly}(d)}.$

This gives the desired result. ∎

Lemma A.2 directly leads to the following corollary:

######  Lemma A.3.

For every $k\in[m]$, every $j\in[d_{0}]$, and every iteration $t$, the
following holds:

$\displaystyle\Big{\langle}{\nabla_{{\mathbf{w}}_{k}^{(t)}}\frac{1}{N}\sum_{i\in[N]}\ell\left(h^{(t)}({\mathbf{x}}_{i}^{(t)}),{\mathbf{y}}_{i}^{(t)}\right)},{\bm{m}_{j}}\Big{\rangle}$
(14)
$\displaystyle\qquad=\big{\langle}{\nabla_{{\mathbf{w}}_{k}^{(t)}}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})\sim\mathcal{D}_{\mathrm{train}}}{\ell\left(h({\mathbf{x}}),{\mathbf{y}}\right)}},{\bm{m}_{j}}\big{\rangle}\pm\frac{1}{\mathsf{poly}(d)}$
$\displaystyle\qquad=-a_{k}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})\sim\mathcal{D}_{\mathrm{train}}}{{\mathbf{y}}\mathbf{1}_{h({\mathbf{x}})\leq
1}\mathbf{1}_{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0}\cdot{\mathbf{z}}_{j}}\pm\frac{1}{\mathsf{poly}(d)},\quad j\in[d_{0}].$

###### Proof.

The proof directly follows from combining Lemma A.1 and the generation process
of ${\mathbf{x}}$ in Definition 3.1. ∎

Lemma A.3 allows us to directly work with population gradients instead of
empirical gradients when analyzing the trajectory of SGD iterations in the
subsequent sections.

###  A.2 Neuron Characterization

In this section, we define two subsets of neurons that will be used throughout
our proofs.

######  Definition A.4 (Neuron characterization).

For each label $y\in\mathcal{Y}=\\{-1,1\\}$ and every iteration $t$, we define
the set $\mathcal{N}_{y}^{(t)}\subseteq[m]$ as:

$\displaystyle\mathcal{N}_{y}^{(t)}\vcentcolon=\Bigg{\\{}k\in[m]:$
$\displaystyle\sum_{j\in\mathcal{S}_{\mathrm{core}}}y\mu_{j}\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{j}}\rangle+\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{j}}\rangle\geq\Theta\left(\sqrt{\frac{d_{0}}{d}}\right),$
(15) $\displaystyle\operatorname{sign}{(a_{k})}=y\Bigg{\\}}.$

#### Intuition.

For each label $y\in\mathcal{Y}$ and iteration $t$, Definition A.4
characterizes a subset of neurons $\mathcal{N}_{y}^{(t)}$ in which

  * •

each neuron has (in expectation) enough positive correlations with the
examples from class $y$ (recall that
${\mathbf{x}}=\sum_{j\in\mathcal{S}_{\mathrm{core}}}{\mathbf{y}}{\mathbf{z}}_{j}\bm{m}_{j}+\sum_{j\in\mathcal{S}_{\mathrm{bg}}}{\mathbf{z}}_{j}\bm{m}_{j}$);

  * •

each neuron positively contributes to the classification of examples from
class $y$ (i.e., $\operatorname{sign}{(a_{k})}=y$).

In our main proof, we will show in an iterative fashion that each neuron in
$\mathcal{N}_{y}^{(t)}$ will accumulate either positive (if random
initialization gives $a_{k}=\frac{1}{m}$) or negative (if random
initialization gives $a_{k}=-\frac{1}{m}$) correlations with features in
$\mathcal{S}_{\mathrm{core}}$ (_core feature learning_), while also
accumulating positive correlations with features in
$\mathcal{S}_{\mathrm{bg}}$ (_feature contamination_).

For each neuron, we formally define the notion of _positive examples_ and
_negative examples_ which are informally mentioned in Section 4:

######  Definition A.5 (Positive examples and negative examples).

Let $(x,y)\in\mathcal{X}\times\mathcal{Y}$ be an example. For every $k\in[m]$,
we say that $(x,y)$ is a positive example of neuron $k$ if
$\operatorname{sign}(a_{k})=y$, and say that $(x,y)$ is a negative example of
neuron $k$ if $\operatorname{sign}(a_{k})=-y$.

###  A.3 Properties at Initialization

In this section, we introduce some useful properties of the neurons at
initialization $t=0$, which serve as a basis for our inductive proofs in the
subsequent sections.

######  Lemma A.6.

For every $j\in[d_{0}]$, every $\mathcal{S}\subseteq[d_{0}]$, and every
$\\{y_{j}\\}_{j\in\mathcal{S}}\in\\{-1,1\\}^{|\mathcal{S}|}$, the following
holds for every $\delta>0$ over random initialization:

$\displaystyle\mathbf{Pr}_{{\mathbf{W}}^{(0)}}{\Big{[}\sum_{j\in\mathcal{S}}y_{j}\mu_{j}\langle{{\mathbf{w}}_{k}^{(0)}},{\bm{m}_{j}}\rangle\geq\frac{\delta}{\sqrt{d}}\Big{]}}$
$\displaystyle\geq\frac{1}{\sqrt{2\pi}}\frac{\delta\sqrt{\sum_{j\in\mathcal{S}}\mu_{j}^{2}}}{\delta^{2}+\sum_{j\in\mathcal{S}}\mu_{j}^{2}}\exp{\left(-\frac{\delta^{2}}{2\sum_{j\in\mathcal{S}}\mu_{j}^{2}}\right)},$
(16)
$\displaystyle\mathbf{Pr}_{{\mathbf{W}}^{(0)}}{\Big{[}\sum_{j\in\mathcal{S}}y_{j}\mu_{j}\langle{{\mathbf{w}}_{k}^{(0)}},{\bm{m}_{j}}\rangle\geq\frac{\delta}{\sqrt{d}}\Big{]}}$
$\displaystyle\leq\frac{1}{\sqrt{2\pi}}\frac{\sqrt{\sum_{j\in\mathcal{S}}\mu_{j}^{2}}}{\delta}\exp{\left(-\frac{\delta^{2}}{2\sum_{j\in\mathcal{S}}\mu_{j}^{2}}\right)}.$

###### Proof.

Recall that different neurons are independently initialized by
${\mathbf{w}}_{k}^{(0)}\sim\mathcal{N}(\mathbf{0},\sigma_{0}^{2}\bm{I}_{d}),\forall
k\in[m]$ with $\sigma_{0}^{2}=\frac{1}{d}$. Using the fact that
$\lVert\bm{m}_{j}\rVert_{2}=1,\forall j\in[d_{0}]$ and $y_{j}^{2}=1,\forall
j\in\mathcal{S}$, we have

$\sum_{j\in\mathcal{S}}y_{j}\mu_{j}\langle{{\mathbf{w}}_{k}^{(0)}},{\bm{m}_{j}}\rangle\sim\mathcal{N}\Big{(}0,\
\frac{1}{d}\sum_{j\in\mathcal{S}}\mu_{j}^{2}\Big{)}$

Applying standard bounds for the Gaussian distribution function (Lemma D.3)
gives that for every $\delta>0$,

$\frac{1}{\sqrt{2\pi}}\frac{\delta}{\delta^{2}+1}\exp{\left(-\frac{\delta^{2}}{2}\right)}\leq\mathbf{Pr}_{{\mathbf{W}}^{(0)}}{\left[\frac{\sqrt{d}\sum_{j\in\mathcal{S}}y_{j}\mu_{j}\langle{{\mathbf{w}}_{k}^{(0)}},{\bm{m}_{j}}\rangle}{\sqrt{\sum_{j\in\mathcal{S}}\mu_{j}^{2}}}\geq\delta\right]}\leq\frac{1}{\sqrt{2\pi}}\frac{1}{\delta}\exp{\left(-\frac{\delta^{2}}{2}\right)}.$

A simple transformation completes the proof. ∎

######  Lemma A.7 (Neuron properties at initialization).

For each label $y\in\mathcal{Y}$, the following holds with probability at
least $1-e^{-\Omega(m)}$ over random initialization:

$\big{|}\mathcal{N}_{y}^{(0)}\big{|}=\Theta(m).$ (17)

###### Proof.

For each neuron $k\in[m]$, define events $E_{k1}$ and $E_{k2}$ to be

$\begin{gathered}E_{k1}\vcentcolon=\Bigg{\\{}\sum_{j\in\mathcal{S}_{\mathrm{core}}}y\mu_{j}\langle{{\mathbf{w}}_{k}^{(0)}},{\bm{m}_{j}}\rangle+\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\langle{{\mathbf{w}}_{k}^{(0)}},{\bm{m}_{j}}\rangle\geq{\Theta\left(\sqrt{\frac{d_{0}}{d}}\right)}\Bigg{\\}},\\\
E_{k2}\vcentcolon=\Big{\\{}\operatorname{sign}{(a_{k})}=y\Big{\\}}.\end{gathered}$

By $a_{k}\sim\mathsf{Uniform}\\{-\frac{1}{m},\frac{1}{m}\\}$, we immediately
have $\mathbf{Pr}{[E_{k2}]}=\frac{1}{2}$ for every $k\in[m]$. For $E_{k1}$, by
applying Lemma A.6 with $\delta=\Theta(\sqrt{d_{0}})$ we obtain

$\displaystyle\mathbf{Pr}_{{\mathbf{W}}^{(0)}}{[E_{k1}]}$
$\displaystyle\geq\frac{1}{\sqrt{2\pi}}\frac{\Theta\left(\sqrt{d_{0}\sum_{j\in[d_{0}]}\mu_{j}^{2}}\right)}{\Theta(d_{0})+\sum_{j\in[d_{0}]}\mu_{j}^{2}}\exp{\left(-\Theta\left(\frac{d_{0}}{\sum_{j\in[d_{0}]}\mu_{j}^{2}}\right)\right)},$
$\displaystyle\mathbf{Pr}_{{\mathbf{W}}^{(0)}}{[E_{k1}]}$
$\displaystyle\leq\frac{1}{\sqrt{2\pi}}\Theta\left(\sqrt{\frac{\sum_{j\in[d_{0}]}\mu_{j}^{2}}{d_{0}}}\right)\exp{\left(-\Theta\left(\frac{d_{0}}{\sum_{j\in[d_{0}]}\mu_{j}^{2}}\right)\right)}.$

Together with $\mu_{j}^{2}=\Theta(1)$ for every $j\in[d_{0}]$, we have that
$\mathbf{Pr}_{{\mathbf{W}}^{(0)}}{[E_{k1}]}=\Theta(1)$ for every $k\in[m]$.
Since events $E_{k1}$ and $E_{k2}$ are independent, we have that for each
neuron $k\in[m]$, the probability of it belonging to $\mathcal{N}_{y}^{(0)}$
is given by $\mathbf{Pr}{(k\in\mathcal{N}_{y}^{(0)})}=\mathbf{Pr}{(E_{k1}\cap
E_{k2})}=\Theta(1)$.

Since different neurons are independently initialized,
$|\mathcal{N}_{y}^{(0)}|$ follows a binomial distribution with trial number
$m$ and some success probability $\Theta(1)$. Applying the standard tail bound
for binomial variables (Lemma D.4) then gives
$|\mathcal{N}_{y}^{(0)}|\geq\Theta(m)$ with probability at least
$1-e^{-\Omega(m)}$. Together with the trivial upper bound that
$|\mathcal{N}_{y}^{(0)}|\leq m$, we have that
$|\mathcal{N}_{y}^{(0)}|=\Theta(m)$ with probability at least
$1-e^{-\Omega(m)}$. ∎

######  Lemma A.8 (Neuron properties at initialization, continued).

With probability at least $1-O(\frac{1}{m})$ over random initialization, for
every $y\in\mathcal{Y}$, the following holds:

$\max_{k\in[m]}\big{|}\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{\langle{{\mathbf{w}}_{k}^{(0)}},{{\mathbf{x}}}\rangle}\big{|}\leq{O}\left(\sqrt{\frac{d_{0}\log
m}{d}}\right).$

###### Proof.

Recall that different neurons are independently initialized by
${\mathbf{w}}_{k}^{(0)}\sim\mathcal{N}(\mathbf{0},\sigma_{0}^{2}\bm{I}_{d}),\forall
k\in[m]$ with $\sigma_{0}^{2}=\frac{1}{d}$. By $\lVert\bm{m}_{j}\rVert_{2}=1$,
we have

$\displaystyle\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{\langle{{\mathbf{w}}_{k}^{(0)}},{{\mathbf{x}}}\rangle}$
$\displaystyle=\sum_{j\in\mathcal{S}_{\mathrm{core}}}y\mu_{j}\langle{{\mathbf{w}}_{k}^{(0)}},{\bm{m}_{j}}\rangle+\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\langle{{\mathbf{w}}_{k}^{(0)}},{\bm{m}_{j}}\rangle$
$\displaystyle\sim\mathcal{N}\Big{(}0,\
\frac{1}{d}\sum_{j\in[d_{0}]}\mu_{j}^{2}\Big{)}.$

Applying Lemma D.5 over the i.i.d. random variables
$\langle{{\mathbf{w}}_{1}^{(0)}},{{\mathbf{x}}}\rangle,\ldots,\langle{{\mathbf{w}}_{m}^{(0)}},{{\mathbf{x}}}\rangle$
gives

$\mathbf{Pr}_{{\mathbf{W}}^{(0)}}{\left[\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{\langle{{\mathbf{w}}_{k}^{(0)}},{{\mathbf{x}}}\rangle}\geq
2\sqrt{\frac{\sum_{j\in[d_{0}]}\mu_{j}^{2}}{d}\log
m}\,\,\right]}\leq\frac{1}{m}.$

Finally, using $\sum_{j\in[d_{0}]}\mu_{j}^{2}=\Theta(d_{0})$ and
$m\in[\Theta(d_{0}),\Theta(d)]$ completes the proof. ∎

######  Lemma A.9 (Output magnitude at initialization).

For every $x\in\mathcal{X}$, the following holds with probability at least
$1-e^{-\Omega(d_{0})}$ over random initialization:

$\big{|}h^{(0)}(x)\big{|}=O\left(\frac{1}{\sqrt{d_{0}}}\right).$ (18)

###### Proof.

By
${\mathbf{w}}_{k}^{(0)}\sim\mathcal{N}(\mathbf{0},\sigma_{0}^{2}\bm{I}_{d})$
with $\sigma_{0}^{2}=\frac{1}{d}$ and $\lVert\bm{m}_{j}\rVert_{2}=1$, we have

$\sum_{k\in[m]}\frac{1}{m}\sum_{j\in[d_{0}]}\langle{{\mathbf{w}}_{k}^{(0)}},{\bm{m}_{j}}\rangle\sim\mathcal{N}\Big{(}0,\frac{d_{0}}{md}\Big{)}.$

Applying standard bounds for the Gaussian distribution function (Lemma D.3)
gives

$\frac{1}{\sqrt{2\pi}}\frac{\delta}{\delta^{2}+1}\exp{\left(-\frac{\delta^{2}}{2}\right)}\leq\mathbf{Pr}_{{\mathbf{W}}^{(0)}}{\left[\sum_{k\in[m]}\frac{1}{m}\sum_{j\in[d_{0}]}\langle{{\mathbf{w}}_{k}^{(0)}},{\bm{m}_{j}}\rangle\geq\delta\sqrt{\frac{d_{0}}{md}}\right]}\leq\frac{1}{\sqrt{2\pi}}\frac{1}{\delta}\exp{\left(-\frac{\delta^{2}}{2}\right)}$

for every $\delta>0$. Substituting $\delta$ by $\Theta(\sqrt{d_{0}})$ and
using the symmetry of Gaussian then yield

$\mathbf{Pr}_{{\mathbf{W}}^{(0)}}{\bigg{[}\,\bigg{|}\sum_{k\in[m]}\frac{1}{m}\sum_{j\in[d_{0}]}\langle{{\mathbf{w}}_{k}^{(0)}},{\bm{m}_{j}}\rangle\bigg{|}\geq\frac{\Theta(d_{0})}{\sqrt{md}}\bigg{]}}\leq\exp(-\Omega(d_{0})).$

We then have

$\displaystyle\big{|}h^{(0)}(x)\big{|}$
$\displaystyle=\Big{|}\sum_{k\in[m]}a_{k}\cdot\mathsf{ReLU}{\big{(}\langle{{\mathbf{w}}_{k}^{(0)}},{{\mathbf{x}}}\rangle\big{)}}\Big{|}$
$\displaystyle\leq\bigg{|}\sum_{k\in[m]}\frac{1}{m}\langle{{\mathbf{w}}_{k}^{(0)}},{{\mathbf{x}}}\rangle\bigg{|}$
$\displaystyle\leq\bigg{|}\sum_{k\in[m]}\frac{1}{m}\sum_{j\in[d_{0}]}\langle{{\mathbf{w}}_{k}^{(0)}},{\bm{m}_{j}}\rangle\bigg{|}$
$\displaystyle\leq\frac{\Theta(d_{0})}{\sqrt{md}}=O\left(\frac{1}{\sqrt{d_{0}}}\right).$

holds with probability at least $1-e^{-\Omega(d_{0})}$, where in the last
equality we use the fact that $m=\Omega(d_{0})$ and $d=\Omega(d_{0}^{2.5})$. ∎

In what follows, we will always assume that the high-probability events at
initialization in Lemma A.7, Lemma A.8, and Lemma A.9 hold—by a union bound
argument and the fact that $m=\Omega(d_{0})$, the probability that they all
hold is at least $1-O(\frac{1}{m})-e^{-\Omega(d_{0})}$.

##  Appendix B Activation Asymmetry, Feature Contamination, and OOD Failure:
Proofs of Theorem 4.1, Theorem 4.2, and Theorem 4.3

Before we delve into the main proofs, we first introduce some technical lemmas
that characterize the gradient updates starting from random initialization. We
begin by introducing an important lemma that characterizes the activation
probability of the ReLU function using the Berry-Esseen theorem:

######  Lemma B.1 (Activation probability).

Assume that the training (ID) data is generated according to Definition 3.1
and
$\left|\frac{\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{i}}\rangle}{\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{j}}\rangle}\right|=\Theta(1)$
for every $k\in[m]$ and for every $i,j\in[d_{0}]$. Then, for every label
$y\in\mathcal{Y}$, every $k\in[m]$, and every iteration $t$, the following
holds:

$\mathbf{Pr}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0]}=\Phi\left(\frac{\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y}{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle}}{\Theta(1)\sqrt{\sum_{j\in[d_{0}]}\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{j}}\rangle^{2}}}\right)\pm
O\left(\frac{1}{\sqrt{d_{0}}}\right),$ (19)

where $\Phi$ denotes the cumulative distribution function of
$\mathcal{N}(0,1)$.

###### Proof.

Recall Definition 3.1 that given label $y\in\mathcal{Y}$, ${\mathbf{x}}$ is
generated by

${\mathbf{x}}=\sum_{j\in\mathcal{S}_{\mathrm{core}}}y{\mathbf{z}}_{j}\bm{m}_{j}+\sum_{j\in\mathcal{S}_{\mathrm{bg}}}{\mathbf{z}}_{j}\bm{m}_{j}.$

Therefore,

$\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle=\sum_{j\in\mathcal{S}_{\mathrm{core}}}y{\mathbf{z}}_{j}\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{j}}\rangle+\sum_{j\in\mathcal{S}_{\mathrm{bg}}}{\mathbf{z}}_{j}\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{j}}\rangle.$

For every $j\in[d_{0}]$, define the random variable

${\mathbf{r}}_{j}\vcentcolon=y_{j}({\mathbf{z}}_{j}-\mu_{j})\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{j}}\rangle,$

where $y_{j}\vcentcolon=\left\\{\begin{aligned} &y,\
j\in\mathcal{S}_{\mathrm{core}}\\\ &1,\
j\in\mathcal{S}_{\mathrm{bg}}\end{aligned}\right.$. Recall that
$\mu_{j}\vcentcolon=\mathbb{E}_{{\mathbf{z}}\sim\mathcal{D}_{\mathrm{train}}}{{\mathbf{z}}_{j}}$.
It is then easy to derive that $\mathbb{E}{{\mathbf{r}}_{j}}=0$ and
$\mathbb{E}{{\mathbf{r}}_{j}^{2}}=\Theta(1)\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{j}}\rangle^{2}$
(recall that $\mathbb{E}{({\mathbf{z}}_{j}-\mu_{j})^{2}}=\Theta(1)$). We now
upper bound $\mathbb{E}{|{\mathbf{r}}_{j}^{3}|}$: first recall that by
Definition 3.1 we have $\mathbb{E}{{\mathbf{z}}_{j}^{3}}=\Theta(1)$. For every
$p\geq 1$, denote the $\ell_{p}$ norm of the random variable
${\mathbf{z}}_{j}$ by
$\lVert{\mathbf{z}}_{j}\rVert_{p}\vcentcolon=(\mathbb{E}{|{\mathbf{z}}_{j}|^{p}})^{\frac{1}{p}}$.
Applying Minkowsky’s inequality gives

$\displaystyle\lVert{\mathbf{z}}_{j}-\mu_{j}\rVert_{p}$
$\displaystyle\leq\lVert{\mathbf{z}}_{j}\rVert_{p}+\lVert\mu_{j}\rVert_{p}$
$\displaystyle\stackrel{{\scriptstyle(a)}}{{=}}\lVert{\mathbf{z}}_{j}\rVert_{p}+\lVert{\mathbf{z}}_{j}\rVert_{1}$
$\displaystyle\stackrel{{\scriptstyle(b)}}{{\leq}}2\lVert{\mathbf{z}}_{j}\rVert_{p},$

where $(a)$ is due to the fact that
$\lVert\mu_{j}\rVert_{p}=|\mu_{j}|=\lVert{\mathbf{z}}_{j}\rVert_{1}$ and $(b)$
is due to the power norm inequality indicating that $\lVert\cdot\rVert_{p}$ is
non-decreasing with regard to $p$. Letting $p=3$ and cubing the above
inequality gives

$\mathbb{E}{|{\mathbf{z}}_{j}-\mu_{j}|^{3}}\leq
8\mathbb{E}{|{\mathbf{z}}_{j}^{3}|}=8\mathbb{E}{{{\mathbf{z}}_{j}^{3}}}=\Theta(1),$

from which we obtain
$\mathbb{E}{|{\mathbf{r}}_{j}^{3}|}=\Theta(1)\cdot|\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{j}}\rangle|^{3}$.

We then define the normalized sum of ${\mathbf{r}}_{j}$:

${\mathbf{s}}_{d_{0}}\vcentcolon=\frac{\sum_{j\in[d_{0}]}{\mathbf{r}}_{j}}{\sqrt{\sum_{j\in[d_{0}]}\mathbb{E}{{\mathbf{r}}_{j}^{2}}}}.$

Since ${\mathbf{r}}_{i}$ and ${\mathbf{r}}_{j}$ are independent and zero-mean
for every $i\neq j\in[d_{0}]$, we can apply the Berry-Esseen theorem (Lemma
D.6) to the normalized sum ${\mathbf{s}}_{d_{0}}$ and obtain

$\displaystyle\sup_{\delta\in\mathbb{R}}\left|\mathbf{Pr}_{{\mathbf{x}}|{\mathbf{y}}=y}{[{\mathbf{s}}_{d_{0}}<\delta]}-\Phi(\delta)\right|$
$\displaystyle\leq
C_{0}\Bigg{(}\sum_{j\in[d_{0}]}\mathbb{E}{{\mathbf{r}}_{j}^{2}}\Bigg{)}^{-\frac{3}{2}}\sum_{j\in[d_{0}]}\mathbb{E}{|{\mathbf{r}}_{j}^{3}|}$
$\displaystyle=C_{0}\Bigg{(}\sum_{j\in[d_{0}]}\Theta(1)\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{j}}\rangle^{2}\Bigg{)}^{-\frac{3}{2}}\sum_{j\in[d_{0}]}\Theta(1)\left|\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{j}}\rangle\right|^{3}$
$\displaystyle\stackrel{{\scriptstyle(c)}}{{=}}O\left(\frac{1}{\sqrt{d_{0}}}\right),$

where $(c)$ is due to the assumption that
$\left|\frac{\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{i}}\rangle}{\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{j}}\rangle}\right|=\Theta(1)$.
Note that
$\sum_{j\in[d_{0}]}{\mathbf{r}}_{j}=\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle-\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y}{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle}$.
We then have for every $\delta\in\mathbb{R}$,

$\mathbf{Pr}_{{\mathbf{x}}|{\mathbf{y}}=y}{\left[\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y}{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle}+\delta\sqrt{\sum_{j\in[d_{0}]}\mathbb{E}{{\mathbf{r}}_{j}^{2}}}\right]}=1-\Phi(\delta)\pm
O\left(\frac{1}{\sqrt{d_{0}}}\right).$

Finally, letting
$\delta=-\dfrac{\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y}{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle}}{\sqrt{\sum_{j\in[d_{0}]}\mathbb{E}{{\mathbf{r}}_{j}^{2}}}}$
and using the symmetry of unit Gaussian $1-\Phi(\delta)=\Phi(-\delta)$ give
the desired result. ∎

We then define two terms that will be frequently used when analyzing
gradients.

######  Definition B.2.

For each label $y\in\mathcal{Y}$, every $k\in[m]$, every feature vector
$\bm{m}_{j},j\in[d_{0}]$, every iteration $t$, and every subset
$\mathcal{S}\subseteq[d_{0}]$, define

$\displaystyle g_{k,y,j}^{(t)}$
$\displaystyle\vcentcolon=\frac{1}{m}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})\sim\mathcal{D}_{\mathrm{train}}}{\mathbf{1}_{h^{(t)}({\mathbf{x}})\leq
1}\mathbf{1}_{{\mathbf{y}}=y}\mathbf{1}_{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0}\mu_{j}{\mathbf{z}}_{j}},$ (20) $\displaystyle g_{k,y,{\mathcal{S}}}^{(t)}$
$\displaystyle\vcentcolon=\sum_{j\in\mathcal{S}}g_{k,y,j}^{(t)}.$

Given the above notation, we now introduce two lemmas that separately bound
the gradient projection onto the core features and the gradient projection
onto the background features for neurons in $\mathcal{N}_{y}^{(t)}$, which
will be helpful for us to track the trajectory of SGD starting from network
initialization.

######  Lemma B.3 (Gradient projection onto core features, neurons in
$\mathcal{N}_{y}^{(t)}$).

For every iteration $t\leq O(\frac{m}{\eta d_{0}})$, the following holds for
every $y\in\mathcal{Y}$ and every neuron $k\in\mathcal{N}_{y}^{(t)}$:

$-\Big{\langle}{\nabla_{{\mathbf{w}}_{k}^{(t)}}\frac{1}{N}\sum_{i\in[N]}\ell\left(h^{(t)}({\mathbf{x}}_{i}^{(t)}),{\mathbf{y}}_{i}^{(t)}\right)},{\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}\bm{m}_{j}}\Big{\rangle}={y}\left(g_{k,y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}+g_{k,-y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}\right),$
(21)

where

$g_{k,y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}=\Theta\left(\frac{d_{0}}{m}\right).$
(22)

###### Proof.

Recall Definition 3.1 that given a label ${\mathbf{y}}\in\mathcal{Y}$,
${\mathbf{x}}$ is generated by

${\mathbf{x}}=\sum_{j\in\mathcal{S}_{\mathrm{core}}}{\mathbf{y}}{\mathbf{z}}_{j}\bm{m}_{j}+\sum_{j\in\mathcal{S}_{\mathrm{bg}}}{\mathbf{z}}_{j}\bm{m}_{j}.$

Then, applying Lemma A.3 to the LHS of Eq. (21) and using
$\operatorname{sign}(a_{k})={y}$ for every $k\in\mathcal{N}_{y}^{(t)}$, we
obtain

$\displaystyle-\Big{\langle}{\nabla_{{\mathbf{w}}_{k}^{(t)}}\frac{1}{N}\sum_{i\in[N]}\ell\left(h^{(t)}({\mathbf{x}}_{i}^{(t)}),{\mathbf{y}}_{i}^{(t)}\right)},{\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}\bm{m}_{j}}\Big{\rangle}$
(23)
$\displaystyle\qquad=-\big{\langle}{\nabla_{{\mathbf{w}}_{k}^{(t)}}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})\sim\mathcal{D}_{\mathrm{train}}}{\ell\left(h^{(t)}({\mathbf{x}}),{\mathbf{y}}\right)}},{\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}\bm{m}_{j}}\big{\rangle}\pm\frac{O(d_{0})}{\mathsf{poly}(d)}$
$\displaystyle\qquad=a_{k}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})\sim\mathcal{D}_{\mathrm{train}}}{{\mathbf{y}}\mathbf{1}_{h^{(t)}({\mathbf{x}})\leq
1}\mathbf{1}_{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0}\sum_{j\in\mathcal{S}_{\mathrm{core}}}{\mathbf{y}}\mu_{j}{\mathbf{z}}_{j}}\pm\frac{O(d_{0})}{\mathsf{poly}(d)}$
$\displaystyle\qquad=a_{k}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})\sim\mathcal{D}_{\mathrm{train}}}{\mathbf{1}_{h^{(t)}({\mathbf{x}})\leq
1}\mathbf{1}_{{\mathbf{y}}=y}\mathbf{1}_{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0}\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}{\mathbf{z}}_{j}}$
$\displaystyle\qquad\qquad+a_{k}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})\sim\mathcal{D}_{\mathrm{train}}}{\mathbf{1}_{h^{(t)}({\mathbf{x}})\leq
1}\mathbf{1}_{{\mathbf{y}}=-y}\mathbf{1}_{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0}\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}{\mathbf{z}}_{j}}\pm\frac{O(d_{0})}{\mathsf{poly}(d)}$
$\displaystyle\qquad={y}\left(g_{k,y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}+g_{k,-y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}\right)\pm\frac{O(d_{0})}{\mathsf{poly}(d)}.$

For $g_{k,y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}$, by the law of total
expectation we have

$\displaystyle g_{k,y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}$
$\displaystyle=\frac{1}{m}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})\sim\mathcal{D}_{\mathrm{train}}}{\mathbf{1}_{h^{(t)}({\mathbf{x}})\leq
1}\mathbf{1}_{{\mathbf{y}}=y}\mathbf{1}_{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0}\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}{\mathbf{z}}_{j}}$
$\displaystyle=\frac{1}{2m}\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y}{\bigg{[}\mathbf{1}_{h^{(t)}({\mathbf{x}})\leq
1}\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}{\mathbf{z}}_{j}\,\Big{|}\,\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0\bigg{]}}\mathbf{Pr}_{{\mathbf{x}}|{\mathbf{y}}=y}{[\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0]}$
$\displaystyle=\frac{1}{2m}\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y}{\bigg{[}\mathbf{1}_{h^{(t)}({\mathbf{x}})\leq
1}\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}{\mathbf{z}}_{j}\bigg{]}}$
$\displaystyle\qquad-\frac{1}{2m}\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y}{\bigg{[}\mathbf{1}_{h^{(t)}({\mathbf{x}})\leq
1}\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}{\mathbf{z}}_{j}\,\Big{|}\,\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle<0\bigg{]}}\mathbf{Pr}_{{\mathbf{x}}|{\mathbf{y}}=y}{[\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle<0]}.$

Applying Lemma B.1 gives

$\mathbf{Pr}_{{\mathbf{x}}|{\mathbf{y}}=y}{[\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle<0]}=\Phi\left(-\frac{\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y}{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle}}{\Theta(1)\sqrt{\sum_{j\in[d_{0}]}\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{j}}\rangle^{2}}}\right)\pm
O\left(\frac{1}{\sqrt{d_{0}}}\right).$

Recall that for ${\mathbf{x}}\sim\mathcal{D}_{\mathrm{train}}|{\mathbf{y}}=y$,

$\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle=\sum_{j\in\mathcal{S}_{\mathrm{core}}}y{\mathbf{z}}_{j}\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{j}}\rangle+\sum_{j\in\mathcal{S}_{\mathrm{bg}}}{\mathbf{z}}_{j}\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{j}}\rangle.$

By Definition A.4, we have for every $k\in\mathcal{N}_{y}^{(t)}$,
$\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y}{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle}\geq\Theta\Big{(}\sqrt{\dfrac{d_{0}}{d}}\Big{)}>0$,
which indicates that
$\Phi\left(-\dfrac{\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y}{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle}}{\Theta(1)\sqrt{\sum_{j\in[d_{0}]}\langle{{\mathbf{w}}_{k}^{(t)}},{\bm{m}_{j}}\rangle^{2}}}\right)<\dfrac{1}{2}$.
Together with $h^{(t)}({\mathbf{x}})\leq 1$, this gives

$\displaystyle g_{k,y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}$
$\displaystyle\geq\frac{1}{2m}\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y}{\bigg{[}\mathbf{1}_{h^{(t)}({\mathbf{x}})\leq
1}\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}{\mathbf{z}}_{j}\bigg{]}}$ (24)
$\displaystyle\qquad-\frac{1}{2m}\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y}{\bigg{[}\mathbf{1}_{h^{(t)}({\mathbf{x}})\leq
1}\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}{\mathbf{z}}_{j}\,\Big{|}\,\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle<0\bigg{]}}\cdot\left(\frac{1}{2}\pm
O\left(\frac{1}{\sqrt{d_{0}}}\right)\right)$
$\displaystyle\geq\sum_{j\in\mathcal{S}_{\mathrm{core}}}\frac{\mu_{j}^{2}}{2m}-\sum_{j\in\mathcal{S}_{\mathrm{core}}}\frac{\mu_{j}^{2}}{4m}\pm\sum_{j\in\mathcal{S}_{\mathrm{core}}}\frac{\mu_{j}^{2}}{\Theta(m\sqrt{d_{0}})}$
$\displaystyle=\Theta\left(\frac{d_{0}}{m}\right).$

Meanwhile, we also have the upper bound

$\displaystyle g_{k,y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}$
$\displaystyle=\frac{1}{m}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})}{\mathbf{1}_{h^{(t)}({\mathbf{x}})\leq
1}\mathbf{1}_{{\mathbf{y}}=y}\mathbf{1}_{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0}\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}{\mathbf{z}}_{j}}$ (25)
$\displaystyle\leq\frac{1}{2m}\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y}\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}{\mathbf{z}}_{j}$
$\displaystyle=\Theta\left(\frac{d_{0}}{m}\right).$

Combining Eqs. (24) and (25) gives

$\displaystyle
g_{k,y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}=\Theta\left(\frac{d_{0}}{m}\right).$

Finally, plugging the above equation and $m=O(d)$ into Eq. (23) completes the
proof. ∎

######  Lemma B.4 (Gradient projection onto background features, neurons in
$\mathcal{N}_{y}^{(t)}$).

For every iteration $t\leq O(\frac{m}{\eta d_{0}})$, the following holds for
every $y\in\mathcal{Y}$ and every neuron $k\in\mathcal{N}_{y}^{(t)}$:

$-\Big{\langle}{\nabla_{{\mathbf{w}}_{k}^{(t)}}\frac{1}{N}\sum_{i\in[N]}\ell\left(h^{(t)}({\mathbf{x}}_{i}^{(t)}),{\mathbf{y}}_{i}^{(t)}\right)},{\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\bm{m}_{j}}\Big{\rangle}=g_{k,y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}-g_{k,-y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)},$
(26)

where

$g_{k,y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}=\widetilde{\Theta}\left(\frac{d_{0}}{m}\right).$
(27)

###### Proof.

Similar to the proof of Lemma B.3, we apply Lemma A.3 to the LHS of Eq. (26)
and using $\operatorname{sign}(a_{k})={y}$ for every
$k\in\mathcal{N}_{y}^{(t)}$, which gives

$\displaystyle-\Big{\langle}{\nabla_{{\mathbf{w}}_{k}^{(t)}}\frac{1}{N}\sum_{i\in[N]}\ell\left(h^{(t)}({\mathbf{x}}_{i}^{(t)}),{\mathbf{y}}_{i}^{(t)}\right)},{\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\bm{m}_{j}}\Big{\rangle}$
(28)
$\displaystyle\qquad=-\big{\langle}{\nabla_{{\mathbf{w}}_{k}^{(t)}}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})\sim\mathcal{D}_{\mathrm{train}}}{\ell\left(h^{(t)}({\mathbf{x}}),{\mathbf{y}}\right)}},{\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\bm{m}_{j}}\big{\rangle}\pm\frac{O(d_{0})}{\mathsf{poly}(d)}$
$\displaystyle\qquad=a_{k}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})\sim\mathcal{D}_{\mathrm{train}}}{{\mathbf{y}}\mathbf{1}_{h^{(t)}({\mathbf{x}})\leq
1}\mathbf{1}_{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0}\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}{\mathbf{z}}_{j}}\pm\frac{O(d_{0})}{\mathsf{poly}(d)}$
$\displaystyle\qquad=\frac{1}{m}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})\sim\mathcal{D}_{\mathrm{train}}}{\mathbf{1}_{h^{(t)}({\mathbf{x}})\leq
1}\mathbf{1}_{{\mathbf{y}}=y}\mathbf{1}_{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0}\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}{\mathbf{z}}_{j}}$
$\displaystyle\qquad\qquad-\frac{1}{m}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})\sim\mathcal{D}_{\mathrm{train}}}{\mathbf{1}_{h^{(t)}({\mathbf{x}})\leq
1}\mathbf{1}_{{\mathbf{y}}=-y}\mathbf{1}_{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0}\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}{\mathbf{z}}_{j}}\pm\frac{O(d_{0})}{\mathsf{poly}(d)}$
$\displaystyle\qquad=g_{k,y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}-g_{k,-y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}\pm\frac{O(d_{0})}{\mathsf{poly}(d)}.$

Also, by a nearly identical argument to Lemma B.3 and
$d_{\mathrm{bg}}=\Theta\left(\dfrac{d_{0}}{\log d_{0}}\right)$, we have

$g_{k,y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}=\widetilde{\Theta}\left(\frac{d_{0}}{m}\right).$
(29)

This completes the proof. ∎

Next, we also introduce a lemma that bound the gradient projection onto core
features for all neurons:

######  Lemma B.5 (Gradient projection onto core features, all neurons).

For every iteration $t\leq O(\frac{m}{\eta d_{0}})$, the following holds for
every $y\in\mathcal{Y}$ and every neuron $k\in[m]$ with
$\operatorname{sign}{(a_{k})}=y$:

$-\Big{\langle}{\nabla_{{\mathbf{w}}_{k}^{(t)}}\frac{1}{N}\sum_{i\in[N]}\ell\left(h^{(t)}({\mathbf{x}}_{i}^{(t)}),{\mathbf{y}}_{i}^{(t)}\right)},{\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}\bm{m}_{j}}\Big{\rangle}=y\cdot
O\left(\frac{d_{0}}{m}\right).$ (30)

###### Proof.

By an identical proof to Lemma B.3, we have

$-\Big{\langle}{\nabla_{{\mathbf{w}}_{k}^{(t)}}\frac{1}{N}\sum_{i\in[N]}\ell\left(h^{(t)}({\mathbf{x}}_{i}^{(t)}),{\mathbf{y}}_{i}^{(t)}\right)},{\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}\bm{m}_{j}}\Big{\rangle}={y}\left(g_{k,y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}+g_{k,-y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}\right)\pm\frac{\Theta(d_{0})}{\mathsf{poly}(d)}.$

By Eq. (25), we have the upper bound
$g_{k,y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}\leq\Theta\left(\frac{d_{0}}{m}\right)$.
By a similar argument, we also have
$g_{k,-y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}\leq\Theta\left(\frac{d_{0}}{m}\right)$.
Plugging those upper bounds and $m=O(d)$ into the above equation completes the
proof. ∎

We then introduce a lemma that bounds the expected correlation between each
neuron in $\mathcal{N}_{y}^{(t)}$ and its positive examples.

######  Lemma B.6 (Correlation with positive examples, neurons in
$\mathcal{N}_{y}^{(t)}$).

For every iteration $t\leq O(\frac{m}{\eta d_{0}})$, every $y\in\mathcal{Y}$,
and every $k\in\mathcal{N}_{y}^{(t)}$, the following holds:

$\displaystyle\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t+1)}},{{\mathbf{x}}}\rangle]\geq(1-\lambda\eta)\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle]+\Theta\left(\frac{\eta
d_{0}}{m}\right).$ (31)

###### Proof.

Recall Definition 3.1 that given the label $y\in\mathcal{Y}$, ${\mathbf{x}}$
is generated by

${\mathbf{x}}=\sum_{j\in\mathcal{S}_{\mathrm{core}}}y{\mathbf{z}}_{j}\bm{m}_{j}+\sum_{j\in\mathcal{S}_{\mathrm{bg}}}{\mathbf{z}}_{j}\bm{m}_{j}.$

We can thus obtain

$\displaystyle\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t+1)}},{{\mathbf{x}}}\rangle]-\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle]$
$\displaystyle\qquad=\underbrace{y\Big{(}\big{\langle}{{\mathbf{w}}_{k}^{(t+1)}},{\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}\bm{m}_{j}}\big{\rangle}-\big{\langle}{{\mathbf{w}}_{k}^{(t)}},{\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}\bm{m}_{j}}\big{\rangle}\Big{)}}_{\Delta_{\mathrm{core}}^{(t)}}$
$\displaystyle\qquad\qquad+\Big{(}\underbrace{\big{\langle}{{\mathbf{w}}_{k}^{(t+1)}},{\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\bm{m}_{j}}\big{\rangle}-\big{\langle}{{\mathbf{w}}_{k}^{(t)}},{\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\bm{m}_{j}}\big{\rangle}}_{\Delta_{\mathrm{bg}}^{(t)}}\Big{)}.$

For $\Delta_{\mathrm{core}}^{(t)}$, by the SGD iteration (4) we have

$\Delta_{\mathrm{core}}^{(t)}=-\eta
y\Big{\langle}{\nabla_{{\mathbf{w}}_{k}^{(t)}}\frac{1}{N}\sum_{i\in[N]}\ell\left(h^{(t)}({\mathbf{x}}_{i}^{(t)}),{\mathbf{y}}_{i}^{(t)}\right)},{\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}\bm{m}_{j}}\Big{\rangle}-\lambda\eta
y\big{\langle}{{\mathbf{w}}_{k}^{(t)}},{\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}\bm{m}_{j}}\big{\rangle}.$

Applying Lemma B.3 gives

$\Delta_{\mathrm{core}}^{(t)}=\eta\left(g_{k,y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}+g_{k,-y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}\right)-\lambda\eta
y\big{\langle}{{\mathbf{w}}_{k}^{(t)}},{\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}\bm{m}_{j}}\big{\rangle},$

which results in the iterative expression

$y\big{\langle}{{\mathbf{w}}_{k}^{(t+1)}},{\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}\bm{m}_{j}}\big{\rangle}=y(1-\lambda\eta)\big{\langle}{{\mathbf{w}}_{k}^{(t)}},{\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}\bm{m}_{j}}\big{\rangle}+\eta\left(g_{k,y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}+g_{k,-y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}\right).$
(32)

For $\Delta_{\mathrm{bg}}^{(t)}$, by the SGD iteration (4) we have

$\Delta_{\mathrm{bg}}^{(t)}=-\eta\Big{\langle}{\nabla_{{\mathbf{w}}_{k}^{(t)}}\frac{1}{N}\sum_{i\in[N]}\ell\left(h^{(t)}({\mathbf{x}}_{i}^{(t)}),{\mathbf{y}}_{i}^{(t)}\right)},{\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\bm{m}_{j}}\Big{\rangle}-\lambda\eta\big{\langle}{{\mathbf{w}}_{k}^{(t)}},{\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\bm{m}_{j}}\big{\rangle}.$

Applying Lemma B.4 gives

$\Delta_{\mathrm{bg}}^{(t)}=\eta\left(g_{k,y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}-g_{k,-y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}\right)-\lambda\eta\big{\langle}{{\mathbf{w}}_{k}^{(t)}},{\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\bm{m}_{j}}\big{\rangle},$

which results in the iterative expression

$\big{\langle}{{\mathbf{w}}_{k}^{(t+1)}},{\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\bm{m}_{j}}\big{\rangle}=(1-\lambda\eta)\big{\langle}{{\mathbf{w}}_{k}^{(t)}},{\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\bm{m}_{j}}\big{\rangle}+\eta\left(g_{k,y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}-g_{k,-y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}\right).$
(33)

Combining Eqs. (36) and (37) gives

$\displaystyle\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t+1)}},{{\mathbf{x}}}\rangle]$
$\displaystyle=y\big{\langle}{{\mathbf{w}}_{k}^{(t+1)}},{\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}\bm{m}_{j}}\big{\rangle}+\big{\langle}{{\mathbf{w}}_{k}^{(t+1)}},{\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\bm{m}_{j}}\big{\rangle}$
$\displaystyle=y(1-\lambda\eta)\big{\langle}{{\mathbf{w}}_{k}^{(t)}},{\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}\bm{m}_{j}}\big{\rangle}+\eta\left(g_{k,y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}+g_{k,-y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}\right)$
$\displaystyle\qquad+(1-\lambda\eta)\big{\langle}{{\mathbf{w}}_{k}^{(t)}},{\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\bm{m}_{j}}\big{\rangle}+\eta\left(g_{k,y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}-g_{k,-y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}\right)$
$\displaystyle=y(1-\lambda\eta)\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle]}$
$\displaystyle\qquad+\eta\left(g_{k,y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}+g_{k,-y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}+g_{k,y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}-g_{k,-y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}\right)$
$\displaystyle\stackrel{{\scriptstyle(a)}}{{\geq}}(1-\lambda\eta)\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle]}+\Theta\left(\frac{\eta
d_{0}}{m}\right),$

where $(a)$ is due to
$g_{k,y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}=\Theta\left(\dfrac{d_{0}}{m}\right)$
(Lemma B.3),
$g_{k,y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}=O\left(\dfrac{d_{0}}{m}\right)$
(Lemma B.4), and

$\displaystyle
g_{k,-y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}-g_{k,-y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}$
$\displaystyle=\frac{1}{m}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})}{\mathbf{1}_{h^{(t)}({\mathbf{x}})\leq
1}\mathbf{1}_{{\mathbf{y}}=-y}\mathbf{1}_{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0}\bigg{(}\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}{\mathbf{z}}_{j}-\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}{\mathbf{z}}_{j}\bigg{)}}$
$\displaystyle\stackrel{{\scriptstyle(b)}}{{\geq}}\frac{1}{m}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})}{\mathbf{1}_{h^{(t)}({\mathbf{x}})\leq
1}\mathbf{1}_{{\mathbf{y}}=-y}\mathbf{1}_{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0}\left(\Theta(d_{0})-\Theta\left(\frac{d_{0}}{\log d_{0}}\right)\right)}$
$\displaystyle\geq 0,$

where $(b)$ is due to the fact that $d_{\mathrm{core}}=\Theta(d_{0})$ and
$d_{\mathrm{bg}}=\Theta\left(\dfrac{d_{0}}{\log d_{0}}\right)$. ∎

We also introduce a general upper bound on the expected correlation between
every neuron in the network and its positive examples.

######  Lemma B.7 (Correlation with positive examples, all neurons).

For every iteration $t\leq O(\frac{m}{\eta d_{0}})$, the following holds for
every $y\in\mathcal{Y}$ and every $k\in\mathcal{N}_{y}^{(t)}$ with
$\operatorname{sign}(a_{k})=y$:

$\displaystyle\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t+1)}},{{\mathbf{x}}}\rangle]\leq(1-\lambda\eta)\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle]+O\left(\frac{\eta
d_{0}}{m}\right).$ (34)

###### Proof.

By an identical proof to Lemma B.6, we have

$\displaystyle\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t+1)}},{{\mathbf{x}}}\rangle]$
$\displaystyle=y(1-\lambda\eta)\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle]}$
$\displaystyle\qquad+\eta\left(g_{k,y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}+g_{k,-y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}+g_{k,y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}-g_{k,-y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}\right)$

By Eq. (25), we have the upper bound
$g_{k,y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}\leq\Theta\left(\frac{d_{0}}{m}\right)$.
By a similar argument, we also have the upper bounds
$g_{k,-y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}\leq\Theta\left(\frac{d_{0}}{m}\right)$
and $g_{k,y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}\leq
O\left(\frac{d_{0}}{m}\right)$. Plugging those upper bounds and the trivial
lower bound $g_{k,-y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}\geq 0$ into the above
equation completes the proof. ∎

The above two lemmas directly lead to the following result saying that if a
neuron is initialized to have large enough correlation to its positive
examples (i.e., belonging to $\mathcal{N}_{y}^{(0)}$), then this large enough
correlation will be retained during training.

######  Lemma B.8 (Neuron properties during training).

For every label $y\in\mathcal{Y}$, every iteration
$t\leq\widetilde{O}(\frac{m}{\eta d_{0}})$, and every step size
$\eta\leq\frac{1}{\mathsf{poly}(d_{0})}$, we have
$\mathcal{N}_{y}^{(t+1)}\supseteq\mathcal{N}_{y}^{(t)}$.

###### Proof.

By Lemma A.8, we have at initialization

$\max_{k\in[m]}\big{|}\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{\langle{{\mathbf{w}}_{k}^{(0)}},{{\mathbf{x}}}\rangle}\big{|}\leq\widetilde{O}\left(\sqrt{\frac{d_{0}}{d}}\right),\
\forall y\in\mathcal{Y}.$

By Lemma B.7 and our choice of $T=\Theta(\frac{m}{\eta d_{0}})$, we have

$\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle]\leq
O\left(\frac{\eta
d_{0}T}{m}\right)+\max_{k\in[m]}\big{|}\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{\langle{{\mathbf{w}}_{k}^{(0)}},{{\mathbf{x}}}\rangle}\big{|}=O(1).$

By Lemma B.6, we have

$\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t+1)}},{{\mathbf{x}}}\rangle]\geq(1-\lambda\eta)\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle]+\Theta\left(\frac{\eta
d_{0}}{m}\right).$

Recall that $\lambda=O(\frac{d_{0}}{m^{1.01}})$. Therefore, as long as
$\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle]=\widetilde{O}(1)$,

$\displaystyle\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t+1)}},{{\mathbf{x}}}\rangle]-\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle]$
$\displaystyle\geq\Theta\left(\frac{\eta
d_{0}}{m}\right)-\lambda\eta\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle]$
$\displaystyle=\Theta\left(\frac{\eta d_{0}}{m}\right)>0.$

Finally, recall that
$\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle]=\sum_{j\in\mathcal{S}_{\mathrm{core}}}y\mu_{j}\langle{{\mathbf{w}}_{k}^{(0)}},{\bm{m}_{j}}\rangle+\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\langle{{\mathbf{w}}_{k}^{(0)}},{\bm{m}_{j}}\rangle$.
By Definition A.4, we immediately have
$\mathcal{N}_{y}^{(t+1)}\supseteq\mathcal{N}_{y}^{(t)}$. ∎

Finally, we introduce a lemma that bounds the expected correlation between
every neuron in the network and its negative examples.

######  Lemma B.9 (Correlation with negative examples, all neurons).

For every iteration $t$, every $y\in\mathcal{Y}$, and every $k\in[m]$ such
that $\operatorname{sign}{(a_{k})}=y$, the following holds:

$\displaystyle\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=-y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t+1)}},{{\mathbf{x}}}\rangle]$
$\displaystyle=(1-\lambda\eta)\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=-y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle]$
(35) $\displaystyle\qquad-\Theta\left(\frac{\eta
d_{0}}{m}\right)\mathbf{Pr}_{{\mathbf{x}}|{\mathbf{y}}=-y}{[\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0]}.$

###### Proof.

Similar to the proof of Lemma B.6, we have

$\displaystyle\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=-y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t+1)}},{{\mathbf{x}}}\rangle]-\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=-y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle]$
$\displaystyle\qquad=\underbrace{-y\Big{(}\big{\langle}{{\mathbf{w}}_{k}^{(t+1)}},{\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}\bm{m}_{j}}\big{\rangle}-\big{\langle}{{\mathbf{w}}_{k}^{(t)}},{\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}\bm{m}_{j}}\big{\rangle}\Big{)}}_{\Delta_{\mathrm{core}}^{(t)}}$
$\displaystyle\qquad\qquad+\Big{(}\underbrace{\big{\langle}{{\mathbf{w}}_{k}^{(t+1)}},{\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\bm{m}_{j}}\big{\rangle}-\big{\langle}{{\mathbf{w}}_{k}^{(t)}},{\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\bm{m}_{j}}\big{\rangle}}_{\Delta_{\mathrm{bg}}^{(t)}}\Big{)}.$

For $\Delta_{\mathrm{core}}^{(t)}$, by the SGD iteration (4) we have

$\Delta_{\mathrm{core}}^{(t)}=\eta
y\Big{\langle}{\nabla_{{\mathbf{w}}_{k}^{(t)}}\frac{1}{N}\sum_{i\in[N]}\ell\left(h^{(t)}({\mathbf{x}}_{i}^{(t)}),{\mathbf{y}}_{i}^{(t)}\right)},{\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}\bm{m}_{j}}\Big{\rangle}+\lambda\eta
y\big{\langle}{{\mathbf{w}}_{k}^{(t)}},{\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}\bm{m}_{j}}\big{\rangle}.$

Applying Lemma B.3 gives

$\Delta_{\mathrm{core}}^{(t)}=-\eta\left(g_{k,y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}+g_{k,-y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}\right)+\lambda\eta
y\big{\langle}{{\mathbf{w}}_{k}^{(t)}},{\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}\bm{m}_{j}}\big{\rangle},$

which results in the iterative expression

$-y\big{\langle}{{\mathbf{w}}_{k}^{(t+1)}},{\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}\bm{m}_{j}}\big{\rangle}=-y(1-\lambda\eta)\big{\langle}{{\mathbf{w}}_{k}^{(t)}},{\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}\bm{m}_{j}}\big{\rangle}-\eta\left(g_{k,y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}+g_{k,-y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}\right).$
(36)

For $\Delta_{\mathrm{bg}}^{(t)}$, by the SGD iteration (4) we have

$\Delta_{\mathrm{bg}}^{(t)}=-\eta\Big{\langle}{\nabla_{{\mathbf{w}}_{k}^{(t)}}\frac{1}{N}\sum_{i\in[N]}\ell\left(h^{(t)}({\mathbf{x}}_{i}^{(t)}),{\mathbf{y}}_{i}^{(t)}\right)},{\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\bm{m}_{j}}\Big{\rangle}-\lambda\eta\big{\langle}{{\mathbf{w}}_{k}^{(t)}},{\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\bm{m}_{j}}\big{\rangle}.$

Applying Lemma B.4 gives

$\Delta_{\mathrm{bg}}^{(t)}=\eta\left(g_{k,y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}-g_{k,-y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}\right)-\lambda\eta\big{\langle}{{\mathbf{w}}_{k}^{(t)}},{\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\bm{m}_{j}}\big{\rangle},$

which results in the iterative expression

$\big{\langle}{{\mathbf{w}}_{k}^{(t+1)}},{\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\bm{m}_{j}}\big{\rangle}=(1-\lambda\eta)\big{\langle}{{\mathbf{w}}_{k}^{(t)}},{\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\bm{m}_{j}}\big{\rangle}+\eta\left(g_{k,y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}-g_{k,-y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}\right).$
(37)

Combining Eqs. (36) and (37) gives

$\displaystyle\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=-y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t+1)}},{{\mathbf{x}}}\rangle]$
$\displaystyle=-y\big{\langle}{{\mathbf{w}}_{k}^{(t+1)}},{\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}\bm{m}_{j}}\big{\rangle}+\big{\langle}{{\mathbf{w}}_{k}^{(t+1)}},{\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\bm{m}_{j}}\big{\rangle}$
(38)
$\displaystyle=-y(1-\lambda\eta)\big{\langle}{{\mathbf{w}}_{k}^{(t)}},{\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}\bm{m}_{j}}\big{\rangle}-\eta\left(g_{k,y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}+g_{k,-y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}\right)$
$\displaystyle\qquad+(1-\lambda\eta)\big{\langle}{{\mathbf{w}}_{k}^{(t)}},{\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}\bm{m}_{j}}\big{\rangle}+\eta\left(g_{k,y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}-g_{k,-y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}\right)$
$\displaystyle=(1-\lambda\eta)\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=-y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle]$
$\displaystyle\qquad+\eta\left(g_{k,y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}-g_{k,y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}\right)-\eta\left(g_{k,-y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}+g_{k,-y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}\right).$

For
$g_{k,y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}-g_{k,y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}$,
we have

$\displaystyle
g_{k,y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}-g_{k,y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}$
$\displaystyle=\frac{1}{m}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})}{\mathbf{1}_{h^{(t)}({\mathbf{x}})\leq
1}\mathbf{1}_{{\mathbf{y}}=y}\mathbf{1}_{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0}\bigg{(}\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}{\mathbf{z}}_{j}-\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}{\mathbf{z}}_{j}\bigg{)}}$
(39)
$\displaystyle=\frac{1}{m}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})}{\mathbf{1}_{h^{(t)}({\mathbf{x}})\leq
1}\mathbf{1}_{{\mathbf{y}}=y}\mathbf{1}_{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0}\left(\Theta\left(\frac{d_{0}}{\log d_{0}}\right)-\Theta(d_{0})\right)}$
$\displaystyle\leq 0.$

For
$g_{k,-y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}+g_{k,-y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}$,
we have

$\displaystyle g_{k,-y,{\mathcal{S}_{\mathrm{core}}}}^{(t)}$
$\displaystyle+g_{k,-y,{\mathcal{S}_{\mathrm{bg}}}}^{(t)}=\frac{1}{m}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})\sim\mathcal{D}_{\mathrm{train}}}{\mathbf{1}_{h^{(t)}({\mathbf{x}})\leq
1}\mathbf{1}_{{\mathbf{y}}=-y}\mathbf{1}_{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0}\bigg{(}\sum_{j\in\mathcal{S}_{\mathrm{bg}}}\mu_{j}{\mathbf{z}}_{j}+\sum_{j\in\mathcal{S}_{\mathrm{core}}}\mu_{j}{\mathbf{z}}_{j}\bigg{)}}$
(40)
$\displaystyle=\frac{1}{m}\mathbb{E}_{({\mathbf{x}},{\mathbf{y}})\sim\mathcal{D}_{\mathrm{train}}}{\mathbf{1}_{h^{(t)}({\mathbf{x}})\leq
1}\mathbf{1}_{{\mathbf{y}}=-y}\mathbf{1}_{\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0}\sum_{j\in[d_{0}]}\mu_{j}{\mathbf{z}}_{j}}$
$\displaystyle=\frac{1}{2m}\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=-y}{\bigg{[}\mathbf{1}_{h^{(t)}({\mathbf{x}})\leq
1}\sum_{j\in[d_{0}]}\mu_{j}{\mathbf{z}}_{j}\Big{|}\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0\bigg{]}}\mathbf{Pr}_{{\mathbf{x}}|{\mathbf{y}}=-y}{[\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0]}$
$\displaystyle=\Theta\left(\frac{d_{0}}{m}\right)\mathbf{Pr}_{{\mathbf{x}}|{\mathbf{y}}=-y}{[\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0]}$

Finally, plugging Eqs. (39) and (40) into Eq. (38) gives the desired result. ∎

We are now ready to introduce the proofs of our main theoretical results.

###  B.1 Proof of Theorem 4.1

For ease of presentation, we first restate the theorem and then introduce its
proof.

######  Theorem B.1 (Activation asymmetry).

For every $\eta\leq\frac{1}{\mathsf{poly}(d_{0})}$ and every
$y\in\mathcal{Y}$, there exists
$T_{0}=\widetilde{\Theta}(\frac{m}{\eta\sqrt{d}})$ such that with high
probability, for every $t\geq T_{0}$, there exist $\Theta(m)$ neurons in which
the weight ${\mathbf{w}}_{k}^{(t)}$ for each neuron satisfies:

$\displaystyle\mathbf{Pr}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0]}=1-O\left(d_{0}^{-\frac{1}{2}}\right),$ (41)
$\displaystyle\mathbf{Pr}_{{\mathbf{x}}|{\mathbf{y}}=-y\sim\mathcal{D}_{\mathrm{train}}}{[\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle\geq
0]}=o(1).$

###### Proof.

For every $y\in\mathcal{Y}$, consider the neuron set $\mathcal{N}_{y}^{(t)}$
defined in Definition A.4. By Lemma A.7 and Lemma B.8, we have
$|\mathcal{N}_{y}^{(t)}|=\Theta(m)$ for every iteration
$t\leq\Theta(\frac{m}{\eta d_{0}})$. We then prove that after at most $T_{0}$
iterations, for every neuron $k\in\mathcal{N}_{y}^{(T_{0})}$ we have
$\mathbf{Pr}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[\langle{{\mathbf{w}}_{k}^{(T_{0})}},{{\mathbf{x}}}\rangle\geq
0]}=1-O\left(\dfrac{1}{\sqrt{d_{0}}}\right)$ and
$\mathbf{Pr}_{{\mathbf{x}}|{\mathbf{y}}=-y\sim\mathcal{D}_{\mathrm{train}}}{[\langle{{\mathbf{w}}_{k}^{(T_{0})}},{{\mathbf{x}}}\rangle\geq
0]}=o(1)$.

Part 1: proving
$\mathbf{Pr}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[\langle{{\mathbf{w}}_{k}^{(T_{0})}},{{\mathbf{x}}}\rangle\geq
0]}=1-O\left(\dfrac{1}{\sqrt{d_{0}}}\right)$.

Let $T_{0}=\Theta(\frac{m\sqrt{\log
md_{0}}}{\eta\sqrt{d}})=\widetilde{\Theta}(\frac{m}{\eta\sqrt{d}})$. By Lemma
B.6 and Lemma B.7 we have

$\displaystyle\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t+1)}},{{\mathbf{x}}}\rangle]$
$\displaystyle=(1-\lambda\eta)\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle]+\Theta\left(\frac{\eta
d_{0}}{m}\right)$
$\displaystyle\leq\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle]+\Theta\left(\frac{\eta
d_{0}}{m}\right),$

which gives
$\mathbb{E}_{{\mathbf{x}}|{\mathbf{y}}=y\sim\mathcal{D}_{\mathrm{train}}}{[}\langle{{\mathbf{w}}_{k}^{(t)}},{{\mathbf{x}}}\rangle]\leq\widetilde{\Theta}(\frac{d_{0}}{\sqrt{d}})=O(1)$.
Recall that $\lambda=o(\frac{d_{0}}{m})$, we then have