SGD Finds then Tunes Features in Two-Layer Neural Networks with
  near-Optimal Sample Complexity: A Case Study in the XOR problem
In this work, we consider the optimization process of minibatch stochastic
gradient descent (SGD) on a 2-layer neural network with data separated by a
quadratic ground truth function. We prove that with data drawn from the
$d$-dimensional Boolean hypercube labeled by the quadratic ``XOR'' function $y
= -x_ix_j$, it is possible to train to a population error $o(1)$ with $d
\:\text{polylog}(d)$ samples. Our result considers simultaneously training both
layers of the two-layer-neural network with ReLU activations via standard
minibatch SGD on the logistic loss. To our knowledge, this work is the first to
give a sample complexity of $\tilde{O}(d)$ for efficiently learning the XOR
function on isotropic data on a standard neural network with standard training.
Our main technique is showing that the network evolves in two phases: a
$\textit{signal-finding}$ phase where the network is small and many of the
neurons evolve independently to find features, and a $\textit{signal-heavy}$
phase, where SGD maintains and balances the features. We leverage the
simultaneous training of the layers to show that it is sufficient for only a
small fraction of the neurons to learn features, since those neurons will be
amplified by the simultaneous growth of their second layer weights.
Margalit Glasgow
SGD Finds then Tunes Features in Two-Layer Neural Networks with Near-Optimal Sample Complexity: A Case Study in the XOR problem

###### Abstract

In this work, we consider the optimization process of minibatch stochastic gradient descent (SGD) on a 2-layer neural network with data separated by a quadratic ground truth function. We prove that with data drawn from the \(d\)-dimensional Boolean hypercube labeled by the quadratic "XOR" function \(y=-x_{i}x_{j}\), it is possible to train to a population error \(o(1)\) with \(d\text{polylog}(d)\) samples. Our result considers simultaneously training both layers of the two-layer-neural network with ReLU activations via standard minibatch SGD on the logistic loss. To our knowledge, this work is the first to give a sample complexity of \(\bar{O}(d)\) for efficiently learning the XOR function on isotropic data on a standard neural network with standard training. Our main technique is showing that the network evolves in two phases: a _signal-finding_ phase where the network is small and many of the neurons evolve independently to find features, and a _signal-heavy_ phase, where SGD maintains and balances the features. We leverage the simultaneous training of the layers to show that it is sufficient for only a small fraction of the neurons to learn features, since those neurons will be amplified by the simultaneous growth of their second layer weights.

## 1 Introduction

Stochastic gradient descent (SGD) is the primary method of training neural networks in modern machine learning. Despite the empirical success of SGD, there are still many questions about why SGD is often able to efficiently find good local minima in the non-convex optimization landscape characteristic of training neural networks.

A growing body of work aims to theoretically understand the optimization dynamics and sample complexity of learning natural classes of functions via SGD on neural networks. A particularly well-understood regime in this regard is the neural tangent kernel (NTK)[Jacot et al., 2021], where the network only moves a small distance from its initialization. However, in many cases, the NTK provably requires a poor sample complexity to generalize [Abbe et al., 2022].

More recent work aims to prove convergence guarantees for SGD on neural networks with tight sample complexity guarantees. A natural test-bed for this, which has garnered a lot of attention, is learning target functions that are inherently low-dimensional, depending only on a constant number of dimensions of the data [Chen and Meka, 2020, Chen et al., 2020, Nichani et al., 2022, Barak et al., 2022, Bietti et al., 2022, Mousavi-Hosseini et al., 2022, Refinetti et al., 2021, Abbe et al., 2021, 2022, 2023]. Such functions, often called _sparse_ or _multi-index_ functions, can be written as \(f(x):=g(Ux)\), where \(U\in\mathbb{R}^{k\times d}\) has orthogonal rows, and \(g\) is a function on \(\mathbb{R}^{k}\). Many works have shown that learning such target functions via SGD on neural networks is possible in much fewer samples than achievable by kernel methods [Chen et al., 2020, Bai and Lee, 2019, Damian et al., 2022, Abbe et al., 2021, 2022, 2023]. The results in these papers apply to a large class of ground truth functions, and have greatly enhanced our understanding of the sample complexity necessary for learning via SGD on neural networks.
The limitation of the aforementioned works is that they typically modify the SGD algorithm in ways that don't reflect standard training practices, for example using layer-wise training, changing learning rates, or clipping. While providing strong guarantees on certain subclasses of multi-index functions, such modifications may limit the ability of SGD to learn broader classes of multi-index functions with good sample complexity. We discuss this more in the context of related work in Section 1.1.

The goal of this paper is to show that for a simple but commonly-studied problem, standard minibatch SGD on a two-layer neural network can learn the ground truth function in near-optimal sample complexity. In particular, we prove in Theorem 3.1 that a polynomial-width ReLU network trained via online minibatch SGD on the logistic loss will classify the boolean XOR function \(f(x):=-x_{i}x_{j}\) with a sample complexity of \(\tilde{O}(d)\).1 We study the XOR function because it one of the simplest test-beds for a function which exhibits some of the core challenges of analyzing SGD on neural networks: a random initialization is near a saddle point, and the sample complexity attainable by kernel methods is suboptimal (see further discussion in Section 1.1).

Footnote 1: We consider this near-optimal in the sense that for algorithms that are rotationally invariant, \(\tilde{\Theta}(d)\) samples are required. See Section F for details.

Despite its simplicity, the prior theoretical understanding of learning the XOR function via SGD on standard networks is lacking. It is well-known that the NTK requires \(\Theta(d^{2})\) samples to learn this function (Wei et al., 2019; Ghorbani et al., 2021; Abbe et al., 2023). Wei et al. (2019) showed that \(\tilde{O}(d)\) samples statistically suffice, either by finding the global optimum of a two-layer network, or by training an infinite-width network, both of which are computationally intractable. Similar guarantees of \(\tilde{O}(d)\) are given by Bai and Lee (2019) and Chen et al. (2020); however, such approaches rely on drastically modifying the network architecture and training algorithm to achieve a quadratic neural tangent kernel. Abbe et al. (2023) proves a sample complexity of \(\tilde{O}(d)\) for the XOR problem, but uses an algorithm which assumes knowledge of the coordinate system under which the data is structured, and is thus not rotationally invariant. It is also worth noting that several works have studied the XOR problem with non-isotropic data, where the cluster separation grows to infinity (Frei et al., 2022; Ben Arous et al., 2022), in some cases yielding better sample complexities.

The main approach in our work is showing that while running SGD, the network naturally evolves in two phases. In the first phase, which we call the _signal-finding_ phase, the network is small, and thus we can show that a sufficient fraction of the neurons evolve independently, similarly to how they would evolve if the output of the network was zero. Phase 1 is challenging because it requires moving away from the saddle near where the network is initialized, which requires super-constant time (here we use "time" to mean the number of iterations times step size). This rules out using the mean field model approach as in Mei et al. (2018, 2019), or showing convergence to a lower-dimensional SDE as in Ben Arous et al. (2022), which both break down after constant time when directly applied to our setting.

After the signal components in the network have become large enough to dominate the remaining components, the network evolves in what we call the _signal-heavy_ phase. In this phase, we show inductively that throughout training, the signal components stay significantly larger than their counterparts. This inductive hypothesis allows us to approximate the output of the network on a sample \(x\) by its _clean_ approximation, given by a network where all the non-signal components have been removed. Under this approximation, the dynamics of the network are easier to compute, and we can show that the signal components will grow and rebalance until all four of the clusters in the XOR problem have sufficiently small loss.

Our Phase 2 analysis leverages the simultaneous training of both layers to show that the dominance of the signal components will be maintained throughout training. In particular, we show once individual neurons become signal heavy, their second layer weights become large, and thus a positive feedback cycle between the first and second layer weights of that neuron causes it to grow faster than non-signal-heavy neurons. This allows us to maintain the signal-heavy inductive hypothesis. If we only trained the first layer, and all second layer weights had equal absolute value, then unless we have strong control over the balance of the clusters, it would be possible for the non-signal components to grow at a rate which is on the same order as the rate of the signal components (see Remark 4.3).
### Related Work

Learning Multi-Index Functions via Neural NetworksMost related to our work is a body of work aiming to understand the sample complexity of learning multi-index functions via SGD on neural networks (Bietti et al., 2022; Refinetti et al., 2021; Chen et al., 2020; Abbe et al., 2021; Abbe et al., 2022; Damien et al., 2022; Barak et al., 2022; Daniely and Malach, 2020; Mousavi-Hosseini et al., 2022; Nichani et al., 2022; Ge et al., 2017; Mahankali et al., 2023). Such functions are typically studied in either the Gaussian data setting where \(x\sim\mathcal{N}(0,I_{d})\), or in the Boolean hypercube setting, where \(x\sim\text{Uniform}(\{\pm 1\}^{d})\). In both cases, we have \(f(x):=g(Ux)\), where \(U\) projects \(x\) onto a lower dimensional space of dimension \(k\), and \(g\) is an arbitrary function on \(k\) variables. In the Boolean setting, \(U\) projects onto a subset of \(k\) coordinates of \(x\), so in the case of the XOR function we study, \(k=2\) and \(g\) is a quadratic function.

Chen and Meka (Chen and Meka, 2020) showed when \(k\) is constant, and \(g\) is a degree-\(D\) polynomial for constant \(D\), there exists a polynomial-time algorithm which learns such multi-index functions on Gaussian covariates in \(\tilde{O}(d)\) samples. Such algorithms can also be emulated in the same sample complexity via SGD on neural networks designed to emulate arbitrary Statistical Query algorithms (Abbe and Sandon, 2020; Abbe et al., 2021), though these networks bear little similarity to standard neural networks used in practice.

The sample complexity of learning multi-index functions via SGD on standard neural networks is an open and active area of research. It is known that the neural tangent kernel (and more generally, kernel methods) require \(\Omega(d^{D})\) samples (Hsu, 2021). A line of work by Abbe et al. (Abbe et al., 2021, 2022, 2023) has conjectured that the sample complexity required for SGD is \(\tilde{\Theta}(d^{\max(L-1,1)})\), where \(L\) denotes the "leap complexity", a measure of hierarchical structure upper bounded by \(D\), and which equals \(2\) for the XOR function. If true, this conjecture would place the sample complexity of SGD on standard neural networks squarely between that of kernel methods and arbitrary polynomial-time algorithms. When \(L=1\), Abbe et al. (2022) showed via a mean-field analysis that is possible to learn with \(\Theta(d)\) samples via layer-wise training, where the first layer is trained until it learns the subspace \(U\), and then the second layer is trained as a linear model. For \(L>1\), Abbe et al. (2023) provided a layer-wise SGD algorithm achieving the conjectured complexity, but which assumes knowledge of the coordinate system under which the data is structured. This means the algorithm is not-rotationally invariant, barring the network from learning more general multi-index functions. Other works have also used layer-wise training to give similar results for subclasses of multi-index functions (Damian et al., 2022; Mousavi-Hosseini et al., 2022; Barak et al., 2022); Mousavi-Hosseini et al. (2022) studies a setting where \(k=1\) and \(L=1\), while Damian et al. (2022); Barak et al. (2022) study settings where \(L\geq 2\), and use just a single gradient step on on the first layer, which requires \(\Omega(d^{L})\) samples. Numerous other works (Tan and Vershynin, 2019; Bietti et al., 2022; Wu et al., 2023) have made progress in the setting of single-index functions (\(k=1\)) when \(L>1\), in some cases achieving tight guarantees that depend on a quantity called the "information exponent" of \(g\), though these methods require training only a single neuron in \(\mathbb{R}^{d}\). The recent work Mahankali et al. (2023) considers training a single-index target function with \(k=2\) and degree \(4\) on a \(2\)-layer neural network via vanilla gradient descent, and shows a sample complexity of \(O(d^{3+\epsilon})\), which improves over kernel methods.

The above discussion highlights a gap in our understanding when \(k\geq 2\) and \(L\geq 2\). Indeed, such a setting is challenging because it requires learning multiple neurons, and escaping one (or more) saddles (Abbe et al., 2023). For this reason, we believe the XOR function (with \(k,L=2\)) is a good stepping stone for understanding the behaviour of SGD on neural networks for more general functions with \(k\geq 2,L\geq 2\). We note that several other works (Bai and Lee, 2019; Chen et al., 2020) have achieved a near-optimal sample complexity of \(\tilde{O}(d)\) for the XOR problems; these works use a non-standard architecture and training algorithm which puts SGD into a quadratic NTK regime. While such a regime can often attain sample complexities beating the standard (linear) NTK, in general this method yields complexities of \(\tilde{O}(d^{D-1})\), which is larger than the rate achieved by Abbe et al. (2022) whenever \(L=1\) and \(D\geq 3\). We emphasize that our work achieves the near-optimal sample complexity \(\tilde{O}(d)\) with a standard two-layer neural network, trained with standard minibatch SGD.

We note that many more works have explored both empirically (eg. (Woodworth et al., 2020; Chizat et al., 2019)) and theoretically (eg.(Li et al., 2020; Allen-Zhu and Li, 2020; Suzuki and Akiyama, 2020; Telgarsky, 2022)) the sample-complexity advantages of "rich" SGD training over the "lazy" NTK regime.

Simultaneous Training of LayersWhile many of the works mentioned above use layer-wise training algorithms, the standard empirical practice is to train all layers simultaneously. Several theoretical works
explore this setting, uncovering implicit biases of ReLU (or other homogeneous) networks trained simultaneously (Wei et al., 2019; Chizat and Bach, 2020; Lyu and Li, 2019; Lyu et al., 2021; Maennel et al., 2018). Under a variety of assumptions, these works have related the solutions found via gradient descent to margin-maximizing solutions. A much finer understanding of the implicit bias of simultaneous training is provided for a line of work on diagonal neural networks (Pesme and Flammarion, 2023; Even et al., 2023).

### Organization of Paper

In Section 2, we formally describe the data and training model we study. In Section 3 we state our result. In Section 4, we give an overview of proof technique. In Section 5, we discuss the limitations of our work, takeaways, and open questions. All proofs are given in the Appendix.

### Notation

For a vector \(v\), we use \(\|v\|\) to denote the \(\ell_{2}\) norm, and \(\|v\|_{1}\) to denote the \(\ell_{1}\) norm. We use \(\|M\|_{2}\) to denote the spectral norm of a matrix \(M\). All big-O notation is with respect to \(d\to\infty\), and we use \(\tilde{O}\) to suppress log factors in big-O notation. \(\omega(1)\) denotes growing to infinity with \(d\). We use \(\mathbb{S}^{d-1}(r)\) to denote the sphere of radius \(r\) in \(d\) dimensions, and \(\mathbf{1}(\cdot)\) to denote the indicator variable of an event.

## 2 Model and Setting

### Data.

We study the setting where the data comes from the Boolean hypercube \(x\sim\text{Uniform}(\{-1,1\}^{d})\), and the label \(y\) is given by \(y(x)=\text{XOR}(x_{1},x_{2}):=-x_{1}x_{2}\).

Note that with \(\mu_{1}:=e_{1}-e_{2}\), and \(\mu_{2}:=e_{1}+e_{2}\), we can model the distribution as

\[(x,y)=\begin{cases}(\mu_{1}+\xi,1)&w.p.\;1/4\qquad(-\mu_{1}+\xi,1)\qquad w.p. \;1/4\\ (\mu_{2}+\xi,-1)&w.p.\;1/4\qquad(-\mu_{2}+\xi,-1)\quad w.p.\;1/4\end{cases},\]

where \(\xi\sim\text{Uniform}(0^{2}\times\{-1,1\}^{d-2})\) so that \(\xi\perp\{\mu_{1},\mu_{2}\}\). We will often write

\[x=z+\xi,\]

where \(z\) is the projection of \(x\) onto the space spanned by \(e_{1}\) and \(e_{2}\), and \(\xi\) is the projection of \(x\) orthogonal to \(e_{1}\) and \(e_{2}\). We denote this distribution by \(P_{d}\), and throughout, it is implicitly assumed that all probabilities and expectations over \(x\) are for \(x\sim P_{d}\).

**Remark 2.1**.: _While for simplicity, we state our results for the setting where the data comes from an axis-aligned Boolean hypercube, and where ground truth depends on the first two dimensions, the minibatch SGD algorithm and the initialization of the network will be rotationally invariant. Thus all our results hold for a Boolean hypercube with any basis._

### Training.

Model.We train both layers of a two-layer ReLU network with width \(p\):

\[\frac{1}{p}\sum_{j=1}^{p}a_{j}\sigma(w_{j}^{T}x),\]

where \(\sigma(\alpha)=\max(0,\alpha)\) is the ReLU function. We will use the variable \(\rho:=\frac{1}{p}\sum_{j=1}^{p}\mathbf{1}_{(w_{j},a_{j})}\) to denote the empirical distribution of the neurons and their second layer weights. Thus we denote

\[f_{\rho}(x):=\mathbb{E}_{w,a\sim\rho}a\cdot\sigma(w^{T}x),\]

We will often abuse notation and write probabilities and expectations using \(w\sim\rho\), and use \(a_{w}\) to denote its associated second layer weight. We note that it is not necessarily the case the second layer weight \(a_{w}\) is a _function_ of \(w\); we do this for the convenience of not indexing each pair as \((w_{j},a_{j})\).
Initialization.We initialize the network with \(w_{j}\sim\text{Uniform}(\mathbb{S}^{d-1}(\theta))\) for a scale parameter \(\theta\), such that \(\|w_{j}\|=\theta\). We initialize the second layer as \(a_{j}=\epsilon_{j}\|w_{j}\|\), where \(\epsilon_{j}\sim\text{Uniform}(\pm 1)\).

Minibatch SGD.We train using minibatch SGD on the logistic loss function

\[\ell_{\rho}(x):=-2\log\left(\frac{1}{1+\exp(-y(x)f_{\rho}(x))}\right),\]

and define the population loss \(L_{\rho}:=\mathbb{E}_{x\sim P}\ell_{\rho}(x)\). We will use the shorthand \(\ell_{\rho}^{\prime}(x)\) to denote the derivative of \(\ell_{\rho}(x)\) with respect to \(f_{\rho}(x)\):

\[\ell_{\rho}^{\prime}(x):=-\frac{2y(x)\exp(-y(x)f_{\rho}(x))}{1+\exp(-y(x)f_{ \rho}(x))}.\]

We use \(\rho_{t}\) to denote the empirical distribution of the \(p\) neurons \((w^{(t)},a_{w}^{(t)})\) at iteration \(t\). At each step, we perform the minibatch SGD update

\[w^{(t+1)}:=w^{(t)}-\eta\nabla\hat{L}_{\rho}(w^{(t)})\qquad a_{w}^{(t+1)}:=a_{w }^{(t)}-\eta\nabla\hat{L}_{\rho}(a_{w}^{(t)}).\]

Here \(\hat{L}_{\rho}=\frac{1}{m}\sum_{x^{(i)}\in M_{t}}\ell_{\rho}(x^{(i)})\) denotes the empirical loss with respect to a minibatch \(M_{t}\) of \(m\) random samples chosen i.i.d. from \(P_{d}\) at step \(t\), and for a loss function \(L\) and a parameter \(u\) in the network, \(\nabla_{u}L:=p\frac{\partial L}{\partial u}\) denotes the scaled partial derivative of the loss with respect to \(u\), defined in particular for a neuron \((w,a_{w})\), as follows: 23

Footnote 2: Since the ReLU function is non-differentiable at zero, we define \(\sigma^{\prime}(0)=0\).

Footnote 3: For convenience, we scale this derivative up by a factor of \(p\) to correspond to the conventional scaling in the mean-field model. Of course if we didn’t perform this scaling, we would achieve the same result by scaling the learning rate \(\eta\).

\[\nabla_{w}\hat{L}_{\rho} =\frac{1}{m}\sum_{x^{(i)}\in M_{t}}\frac{\partial}{\partial w}p \ell_{\rho}(x^{(i)})=\frac{1}{m}\sum_{x^{(i)}\in M_{t}}a_{w}\ell_{\rho_{t}}^{ \prime}(x^{(i)})\sigma^{\prime}(w^{T}x^{(i)})x^{(i)};\] \[\nabla_{a_{w}}\hat{L}_{\rho} =\frac{1}{m}\sum_{x^{(i)}\in M_{t}}\frac{\partial}{\partial a_{w }}p\ell_{\rho}(x^{(i)})=\frac{1}{m}\sum_{x^{(i)}\in M_{t}}\ell_{\rho_{t}}^{ \prime}(x^{(i)})\sigma(x_{i}^{T}w).\]

## 3 Main Result

The following theorem is our main result.

**Theorem 3.1**.: _There exists a constant \(C>0\) such that the following holds for any \(d\) large enough. Let \(\theta:=1/\log^{C}(d)\). Suppose we train a 2-layer neural network with minibatch SGD as in Section 2.2 with a minibatch size of \(m\geq d/\theta\), width \(1/\theta\leq p\leq d^{C}\), step size \(d^{-C}\leq\eta\leq\theta\), and initialization scale \(\theta/\sqrt{p}\). Then for some \(t\leq C\log(d)/\eta\), with probability \(1-d^{-\omega(1)}\), we have_

\[\mathbb{E}_{x\sim P_{d}}[\ell_{\rho_{t}}(x)]\leq o(1).\]

By setting \(\eta=\theta\) and \(m=d/\theta\), Theorem 3.1 states that we can learn the XOR function up to \(\epsilon\) population loss in \(\Theta\left(d\text{polylog}(d)\right)\) samples and iterations on a polynomial-width network.

## 4 Proof Overview

Throughout the following section, and in our proofs, we will use the following shorthand to refer to the components of a neurons \(w\). We decompose \(w=w_{1:2}+w_{\perp}\), where \(w_{1:2}\) is the projection of \(w\) in the direction spanned \(e_{1}\) and \(e_{2}\) (and equivalently by \(\mu_{1}=e_{1}-e_{2}\) and \(\mu_{2}=e_{1}+e_{2}\)), and \(w_{\perp}\) is the component of \(w\) in the orthogonal subspace. We further decompose \(w_{1:2}=w_{\text{sig}}+w_{\text{opp}}\) as follows:

\[w_{\text{sig}}=\begin{cases}\frac{1}{2}\mu_{1}\mu_{1}^{T}w&a_{w}\geq 0;\\ \frac{1}{2}\mu_{2}\mu_{2}^{T}w&a_{w}<0.\end{cases}\qquad w_{\text{opp}}= \begin{cases}\frac{1}{2}\mu_{2}\mu_{2}^{T}w&a_{w}\geq 0;\\ \frac{1}{2}\mu_{1}\mu_{1}^{T}w&a_{w}<0.\end{cases}\]
Intuitively, we want the neurons to grow in the \(w_{\rm sig}\) direction, but not the \(w_{\rm opp}\) direction; in a network achieving the maximum normalized margin, we will have \(w=w_{\rm sig}\) exactly, and \(w_{\rm opp}=w_{\perp}=0\). We summarize this notation in Table 1, along with future shorthand we will introduce in this section.

The main idea of our proof is to break up the analysis of SGD into two main phases. In the first phase, the network is small, and thus we have (for most \(x\)) that the loss \(\ell_{\rho}(x)\) is well approximated by a first order approximation of the loss at \(f_{\rho}=0\), namely

\[\ell_{0}(x;\rho):=-2\log(1/2)-y(x)f_{\rho}(x).\]

As long as this approximation holds, the neurons of the network evolve (approximately) independently, since \(\ell_{0}^{\prime}(x):=\frac{\partial\ell_{0}(x;\rho)}{\partial f_{\rho}(x)}=-y(x)\) does not depend on the full network \(\rho\). We will show under this approximation that for many neurons, \(\|w_{\rm sig}\|\) grows exponentially fast. Thus we will run this first phase for \(\Theta(\log(d)/\eta)\) iterations until for all four clusters \(\mu\in\{\pm\mu_{1},\pm\mu_{2}\}\), there exists a large set of neurons \(S_{\mu}\) on which \(w_{\rm sig}^{T}\mu>0\), and the "margin" from this set of neurons is large, ie.

\[\tilde{\gamma}_{\mu}:=\mathbb{E}_{\rho}[\mathbf{1}(w\in S_{\mu})a_{w}\sigma(w ^{T}\mu)]\gg\mathbb{E}_{\rho}\|w_{\perp}+w_{\rm opp}\|^{2}. \tag{4.1}\]

In the Phase 2, we assume that Eq. 4.1 holds, and we leverage the dominance of the signal to show that (1) The signal components \(w_{\rm sig}\) grow faster that \(w_{\rm opp}+w_{\perp}\), and thus Eq. 4.1 continues to hold; and (2) SGD balances the signal components in the 4 cluster directions such that the margins \(\tilde{\gamma}_{\mu}\) balance, and become sufficiently large to guarantee \(o(1)\) loss.

We proceed to describe the analysis in the two phases in more detail. Full proofs are in the Appendix.

### Phase 1

In Phase 1, we approximate the evolution of the network at each gradient step by the gradient step that would occur for a network with output \(0\). The main building blocks of our analysis are estimates of the \(L_{0}:=\mathbb{E}_{x}\ell_{0}(x;\rho)\) population gradients, and bounds on the difference \(\nabla L_{0}-\nabla L_{\rho}\).

\(L_{0}\) population gradients.Since the primary objective of this phase is to grow the neurons in the signal direction, we sketch here the computation of the gradient \(\nabla_{w_{1,2}}L_{0}\) in the subspace spanned by \(\mu_{1},\mu_{2}\). The remaining estimates of \(\nabla L_{0}\) are simpler, and their main objective is to show that \(\nabla_{w_{\perp}}L_{0}\) and \(\nabla_{a_{w}}L_{0}\) are sufficiently small, such that \(\|w_{\perp}\|\) doesn't change much throughout Phase 1, and \(|a_{w}|\) stays approximately the same as \(\|w\|\). For convenience, the reader may assume that \(|a_{w}|=\|w\|\) exactly.4,

Footnote 4: When \(\eta\to 0\) as in gradient flow, this equivalence holds exactly for ReLU networks, as long as the initialization satisfies \(|a_{w}|=\|w\|\).

For a data sample \(x\sim\text{Rad}^{d}\), we denote \(x=z+\xi\), where \(z\in\text{Span}(\{\pm\mu_{1},\pm\mu_{2}\})\), and \(\xi\perp\text{Span}(\{\pm\mu_{1},\pm\mu_{2}\})\).

\begin{table}
\begin{tabular}{c|c|c|c} \hline \hline \(w_{\rm sig}=\begin{cases}\frac{1}{2}\mu_{1}\mu_{1}^{T}w&a_{w}\geq 0\\ \frac{1}{2}\mu_{2}\mu_{2}^{T}w&a_{w}<0\end{cases}\) & \(w_{\rm opp}=\begin{cases}\frac{1}{2}\mu_{2}\mu_{2}^{T}w&a_{w}\geq 0\\ \frac{1}{2}\mu_{1}\mu_{1}^{T}w&a_{w}<0\end{cases}\) & \(\begin{cases}w_{1:2}=w_{\rm sig}+w_{\rm opp}\\ w_{\perp}=w-w_{1:2}\end{cases}\) \\ \hline \(\gamma_{\mu}=f_{\rho}(\mu)y(\mu)\) & \(\gamma_{\rm min}=\min_{\mu\in\{\pm\mu_{1},\pm\mu_{2}\}}\gamma_{\mu}\) & \(\gamma_{\rm max}=\max_{\mu\in\{\pm\mu_{1},\pm\mu_{2}\}}\gamma_{\mu}\) \\ \hline \(g_{\mu}=|\ell_{\rho}^{\prime}(\mu)|\) & \(g_{\rm min}=\min_{\mu\in\{\pm\mu_{1},\pm\mu_{2}\}}|\ell_{\rho}^{\prime}(\mu)|\) & \(g_{\rm max}=\max_{\mu\in\{\pm\mu_{1},\pm\mu_{2}\}}|\ell_{\rho}^{\prime}(\mu)|\) \\ \hline \hline \end{tabular}
\end{table}
Table 1: Summary of Notation used in Proof Overview and Proofs
By leveraging the symmetry of the data distribution and the fact that \(y(z)=y(-z)\), we can compute

\[\begin{split}\nabla_{w_{1:2}}L_{0}&=-a_{w}\mathbb{E}_{ x=z+\xi}y(x)\sigma^{\prime}(w^{T}x)z\\ &=-a_{w}\mathbb{E}_{\xi}\frac{1}{2}\mathbb{E}_{z}y(z)\left(\sigma ^{\prime}(w^{T}\xi+w^{T}z)-\sigma^{\prime}(w^{T}\xi-w^{T}z)\right)z\\ &=-a_{w}\mathbb{E}_{\xi}\frac{1}{2}\mathbb{E}_{z}y(z)\mathbf{1}(|w ^{T}z|\geq|w^{T}\xi|)\operatorname{sign}(w^{T}z)z\\ &=-\frac{1}{2}a_{w}\mathbb{E}_{z}y(z)\operatorname{sign}(w^{T}z)z \mathbb{P}_{\xi}[|w^{T}z|\geq|w^{T}\xi|]\\ &\approx-\frac{1}{2}a_{w}\mathbb{E}_{z}y(z)\operatorname{sign}(w ^{T}z)z\mathbb{P}_{G\sim\mathcal{N}(0,\|w\|^{2})}[G\leq|w^{T}z|]\\ &\approx-\frac{1}{2}a_{w}\mathbb{E}_{z}y(z)\operatorname{sign}(w ^{T}z)z\sqrt{\frac{2}{\pi}}\frac{|w^{T}z|}{\|w\|}.\end{split} \tag{4.2}\]

Here the two approximations come from the fact that \(\xi\) has boolean coordinates and not Gaussian, and from an approximation of the Gaussian distribution, which holds whenever \(\frac{|w^{T}z|}{\|w\|}\) is small. By taking the expectation over \(z\in\{\pm\mu_{1},\pm\mu_{2}\}\), the last line of Eq 4.2 can be shown to evaluate to

\[-\frac{|a_{w}|}{\|w\|\sqrt{2\pi}}w_{\text{sig}}+\frac{|a_{w}|}{\|w\|\sqrt{2\pi} }w_{\text{opp}}. \tag{4.3}\]

Observe that near initialization, this gradient is quite small, since \(\frac{\|w_{\text{sig}}\|}{\|w\|}\) is approximately \(\frac{1}{\sqrt{d}}\) for a random initialization. Nevertheless, this gradient suggests that \(w_{\text{sig}}\) will grow exponentially fast.

Bounding the difference \(\nabla L_{0}-\nabla L_{\rho}\).To bound \(\|\nabla_{w}L_{\rho}-\nabla_{w}L_{0}\|_{2}\), first recall that

\[\nabla_{w}L_{0}-\nabla_{w}L_{\rho}=\mathbb{E}_{x}a_{w}(\ell_{\rho}^{\prime}(x )-\ell_{0}^{\prime}(x))\sigma^{\prime}(w^{T}x)x.\]

Defining \(\Delta_{x}:=(\ell_{\rho}^{\prime}(x)-\ell_{0}^{\prime}(x))\sigma^{\prime}(w^{ T}x)\), we can show using routine arguments (see Lemma C.2 for the details) that:

\[\|\nabla_{w}L_{\rho}-\nabla_{w}L_{0}\|_{2}=|a_{w}|\|\mathbb{E}_{ x}\Delta_{x}x\| \leq|a_{w}|\sqrt{\mathbb{E}_{x}\Delta_{x}^{2}} \tag{4.4}\] \[\approx|a_{w}|\sqrt{\mathbb{E}_{x}f_{\rho}(x)^{2}}\] \[\lessapprox|a_{w}|\mathbb{E}_{\rho}[\|a_{w}w\|]\approx\frac{|a_{ w}|}{\operatorname{polylog}(d)}.\]

While this deviation bound is useful for showing that \(w_{\perp}\) doesn't move too much, this bound far exceeds the scale of the gradient in the \(w_{\text{sig}}\), which is on the scale \(\frac{|a_{w}|}{\sqrt{d}}\) near initialization. Fortunately, we can show in Lemma C.3 that the deviation is much smaller on the first two coordinates, namely,

\[\|\nabla_{w_{1:2}}L_{\rho}-\nabla_{w_{1:2}}L_{0}\|_{2}\leq|a_{w}|O(\log^{2}(d) )\left(\mathbb{E}_{\rho}[\|a_{w}w_{1:2}\|]+\mathbb{E}_{\rho}[\|a_{w}w\|] \frac{\|w_{1:2}\|}{\|w\|}\right) \tag{4.5}\]

Note that since near initialization \(\|w_{1:2}\|\ll\|w\|\) for all neurons, this guarantee is much stronger than Eq. 4.4. In fact, since throughout this phase we can show that \(a_{w}\) and \(\|w\|\) change relatively little, staying at the scale \(1/\text{polylog}(d)\), the approximation error in Eq. 4.5 is smaller than the gradient in the \(w_{\text{sig}}\) direction (Eq. 4.3) whenever say \(\|w_{\text{sig}}\|\geq 100\mathbb{E}_{\rho}[\|a_{w}w_{1:2}\|]\), which occurs on a substantial fraction of the neurons.

Lemma C.3 is the most important lemma in our Phase 1 analysis. At a high level, it shows that the approximation error \(\|\nabla_{w_{1:2}}L_{\rho}-\nabla_{w_{1:2}}L_{0}\|_{2}\) can be coupled with the growth of the signal, \(-(\nabla_{w}L_{0})^{T}\frac{w_{\text{sig}}}{\|w_{\text{sig}}\|}\). This is because we use a symmetrization trick with the pairs \(z+\xi\) and \(-z+\xi\) to show that both the error and the signal gradient only grow from samples \(x=z+\xi\) where \(|z^{T}w|\geq|\xi^{T}w|\).

In more detail, to prove Eq. 4.5, we also need to leverage the fact that for any \(\xi\in\{\mu_{1},\mu_{2}\}^{\perp}\) and \(z\in\{\pm\mu_{1},\pm\mu_{2}\}\), we have \(|\ell_{\rho}^{\prime}(\xi+z)-\ell_{\rho}^{\prime}(\xi-z^{\prime})|\leq 4p\mathbb{E}_{ \rho}[\|a_{w}w_{1:2}\|]\), much smaller than we can expect
\(|\ell^{\prime}_{\rho}(x)-\ell^{\prime}_{0}(x)|\) to be. Thus \(|\Delta_{\xi+z}-\Delta_{\xi-z}|\leq 4p\mathbb{E}_{\rho}[\|a_{w}w_{1:2}\|]\) whenever \(|\xi^{T}w|\geq|z^{T}w|\) (such that \(\sigma^{\prime}(w^{T}(\xi+z))=\sigma^{\prime}(w^{T}(\xi-z))\)). Following the symmetrization trick in Eq. 4.2, we have

\[\left\|\frac{1}{a_{w}}\left(\nabla_{w_{1:2}}L_{\rho}-\nabla_{w_{1: 2}}L_{0}\right)\right\| =\|\mathbb{E}_{x}\Delta_{x}z\|\] \[=\|\mathbb{E}_{\xi}\mathbb{E}_{z}\Delta_{\xi+z}z\|\] \[=\frac{1}{2}\|\mathbb{E}_{\xi}\mathbb{E}_{z}(\Delta_{\xi+z}- \Delta_{\xi-z})z\|\] \[\leq 2\sqrt{2}\mathbb{E}_{\rho}[\|a_{w}w_{1:2}\|]+\sqrt{2} \mathbb{E}_{\xi}\mathbb{E}_{z}\mathbf{1}(|\xi^{T}w|\leq|z^{T}w|)|\Delta_{x}|.\]

A careful computation comparing \(w^{T}\xi\) to a Gaussian distribution then shows that

\[\mathbb{E}_{z}\mathbf{1}(|\xi^{T}w|\leq|z^{T}w|)|\Delta_{x}|\approx\left( \mathbb{P}_{x}[|\xi^{T}w|\leq|z^{T}w|]\right)(\mathbb{E}_{x}|\Delta_{x}|) \lessapprox\frac{\|w_{1:2}\|}{\|w\|}\mathbb{E}_{\rho}[\|a_{w}w\|].\]

Putting Phase 1 TogetherThe building blocks above, combined with standard concentration bounds on \(\nabla\hat{L}_{\rho}\), suffice to show that a substantial mass of neurons will evolve according to Eq 4.3, leading to exponential growth in \(w_{\text{sig}}\). After \(\Theta(\log(d)/\eta)\) iterations, for these neurons, we can achieve \(\|w_{\text{sig}}\|\gg\|w_{\perp}+w_{\text{opp}}\|\). Formally, we show the following for some \(\zeta\leq 1/\text{polylog}(d)\):

**Lemma 4.1** (Output of Phase 1: Informal; See Lemma C.1 for formal version).: _With high probability, for \(\eta\leq\tilde{O}(1)\), after some \(T=\Theta(\log(d)/\eta)\) iterations of minibatch SGD, with \(m=\tilde{\Theta}(d)\) samples in each minibatch, the network \(\rho_{T}\) satisfies:_

1. \(\mathbb{E}_{\rho_{T}}[\|w_{\perp}+w_{opp}\|^{2}]\leq\theta\)_._
2. _For each_ \(\mu\in\{\pm\mu_{1},\pm\mu_{2}\}\)_, on at least a_ \(0.1\) _fraction of all the neurons, we have_ \(w_{\text{sig}}^{T}\mu>0\)_, and_ \(\|w_{\text{sig}}\|^{2}\geq\zeta^{-1}\theta\)_._

We remark that the analysis to prove Lemma 4.1 is somewhat subtle, since the tight approximation in Eq 4.2 breaks down when \(\|w_{\text{sig}}\|\) approaches \(\|w_{\perp}\|\). The details are given in Appendix C.

#### 4.1.1 Phase 2

The conclusion of Lemma 4.1 is a sufficient condition of the network to begin the second phase. In the second phase, we have that (for most \(x\))

\[\ell^{\prime}_{\rho}(x)\approx\ell^{\prime}_{\rho}(z), \tag{4.6}\]

where we recall that \(z\) is the component of \(x\) in the space spanned by \(\mu_{1}\) and \(\mu_{2}\). We refer to this as the _clean_ loss derivative, and our main tool will be analyzing the evolution of SGD under this clean surrogate for the loss derivative. Namely, we define:

\[\nabla^{\text{cl}}_{w}L_{\rho}:=a_{w}\mathbb{E}_{x}\ell^{\prime}_{\rho}(z) \sigma^{\prime}(w^{T}x)x\quad\text{and}\quad\nabla^{\text{cl}}_{a_{w}}L_{\rho }:=\mathbb{E}_{x}\ell^{\prime}_{\rho}(z)\sigma(w^{T}x). \tag{4.7}\]

Before proceeding, we introduce the following definitions, which will be useful in Phase 2 (summarized in Table 1):

\[\gamma_{\min} :=\min_{\mu\in\{\pm\mu_{1},\pm\mu_{2}\}}\gamma_{\mu} g_{\min}:=\min_{\mu\in\{\pm\mu_{1},\pm\mu_{2}\}}|\ell^{\prime}_{ \rho}(\mu)|=\frac{\exp(-\gamma_{\max})}{1+\exp(-\gamma_{\max})}\] \[\gamma_{\max} :=\max_{\mu\in\{\pm\mu_{1},\pm\mu_{2}\}}\gamma_{\mu} g_{\max}:=\max_{\mu\in\{\pm\mu_{1},\pm\mu_{2}\}}|\ell^{\prime}_{\rho}( \mu)|=\frac{\exp(-\gamma_{\min})}{1+\exp(-\gamma_{\min})}\]

To ensure the approximation in Eq. 4.6 holds throughout the entire the second phase, we will maintain a certain inductive hypothesis, which ensures the the scale of the signal-direction components of the network continue to dominate the scale of the non-signal-direction components of the network. Formally, we consider the following condition.
**Definition 4.2** (Signal-Heavy Inductive Hypothesis).: _For parameters \(\zeta=o(1)\) and \(H>1\) with \(\zeta\leq\exp(-10H)\), we say a network is \((\zeta,H)\)-signal-heavy if there exists some set of heavy neurons \(S\) on which \(\exp(6H)\|w_{\perp}\|+\|w_{opp}\|\leq\zeta\|w_{sig}\|\), and_

\[\mathbb{E}_{\rho}\mathbf{1}(w\notin S)\|w\|^{2}\leq\zeta\tilde{\gamma}_{min}.\]

_Here we have defined \(\tilde{\gamma}_{\mu}:=\mathbb{E}[\mathbf{1}(w\in S,w_{sig}^{T}>0)a_{w}\sigma(w ^{T}\mu)]\) and \(\tilde{\gamma}_{min}:=\min_{\mu\in\{\pm\mu_{1},\pm\mu_{2}\}}\tilde{\gamma}_{\mu}\)._

_Further,_

\[\mathbb{E}_{\rho}[\|w\|^{2}]\leq\mathbb{E}_{\rho}[|a_{w}|^{2}]+\zeta H\leq 2H,\]

_and for all neurons, we have \(|a_{w}|\leq\|w\|\)._

We show via a straightforward argument in Lemma D.4 that if the conclusion of Lemma 4.1 (from Phase 1) holds for some \(\zeta\), then the network is \((\Theta(\zeta^{1/3}),H)\)-signal-heavy, for \(H=\Theta(\log\log(d))\).

Assuming that the network is \((\zeta,H)\)-signal-heavy, using a similar approach to Eq. 4.4, we can show (see Lemma D.5 for the precise statement) that for any neuron \((w,a_{w})\),

\[\frac{1}{|a_{w}|}\|\nabla_{w}L_{\rho}-\nabla_{w}^{\mathrm{cl}}L_{\rho}\|_{2} \lessapprox\sqrt{\mathbb{E}_{x}(f_{\rho}(x)-f_{\rho}(z))^{2}}\lessapprox \mathbb{E}_{\rho}[\|a_{w}w_{\perp}\|]\leq\zeta\gamma_{\max},\]

and similarly \(\|\nabla_{a_{w}}L_{\rho}-\nabla_{a_{w}}^{\mathrm{cl}}L_{\rho}\|_{2}\lessapprox \|w\|\zeta\gamma_{\max}\).

By working with the clean gradients, it is possible to approximately track (or bound) the evolution of \(w_{\text{sig}}\), \(w_{\perp}\), and \(w_{\text{opp}}\) on neurons in \(S\), the set of neurons for which \(\|w_{\text{sig}}\|\gg\|w_{\perp}+w_{\text{opp}}\|\). In Lemmas D.6, D.7, and D.8 we show the following for any \(w\in S\) (let \(\mu\) be the direction of \(w_{\text{sig}}\)):

1. **The signal component \(w_{\text{sig}}\) grows quickly.** We have \(-w_{\text{sig}}^{T}\nabla_{w}^{\mathrm{cl}}L_{\rho}\approx|a_{w}\ell_{\rho}^{ \prime}(\mu)|\tau\), where \(\tau:=\frac{\sqrt{2}}{4}\). Also \(a_{w}\) grows at a similar rate. This growth is due to the fact that points with \(z=-\mu\) will almost never activate the ReLU, while points with \(z=\mu\) almost always will.
2. **A linear combination of \(\|w_{\perp}\|^{2}\) and \(\|w_{\text{opp}}\|^{2}\) decreases.** The argument here is more subtle, but the key idea is to argue that if \(|w_{\perp}^{T}\xi|\geq|w_{\text{opp}}^{T}z|\) frequently, then \(\|w_{\perp}\|^{2}\) will decrease. Meanwhile, if \(|w_{\perp}^{T}\xi|\leq|w_{\text{opp}}^{T}z|\) frequently, then \(w_{\text{opp}}\) will decrease (and there is a sizeable event on which they both decrease).

Since most of the mass of the network is in \(S\), this shows that the signal will grow at the exponential rate \(\tau|\ell_{\rho}^{\prime}(\mu)|\) -- or for the "weakest" cluster, that is, in the direction \(\mu\) that maximizes \(\tilde{\gamma}_{\mu}\), we will have \(\tilde{\gamma}_{\min}^{(t+1)}\gtrapprox(1+2\eta\tau g_{\max})\,\tilde{\gamma}_ {\min}^{(t)}\).

On neurons outside of \(S\), we show in Lemma D.11 that they grow _at most_ as fast as the rate of the weakest clusters, meaning we can essentially ignore these neurons.

**Remark 4.3**.: _If we did not train the second layer weights (and for instance they all had norm \(1\)), then our tools would not suffice to maintain the signal-heavy hypothesis in Definition 4.2. Indeed, the neurons in \(S\) would grow at a linear rate of \(\tau|\ell_{\rho}^{\prime}(\mu)|\), and at (up to) an equal linear rate outside of \(S\). Thus the neurons outside of \(S\) might eventually attain a non-negligible mass. However, because the layers are trained simultaneously, this leads to positive feedback between the growth of \(\|w_{\text{sig}}\|\) and \(|a_{w}|\), leading to exponential growth, which maintains the mass ratios between the neurons in and out of \(S\)._

Combining the ideas above, we prove the following lemma, which shows that after one SGD step, the network stays signal-heavy (with a slightly worse parameter), the behavior of the weakest margin improves, and the network (measured by the size of the largest margin \(\gamma_{\max}\)) doesn't become too big.

**Lemma 4.4** (Phase 2 Inductive Step: Informal; See Lemma D.3 for formal version).: _If a network \(\rho_{t}\) is \((\zeta,H)\)-signal heavy with heavy set \(S\), then after one minibatch gradient step, with probability \(1-d^{-\omega(1)}\),_

1. \(\rho_{t+1}\) _is_ \((\zeta(1+10\eta\zeta H),H)\)_-signal heavy._
2. \(\tilde{\gamma}_{min}^{(t+1)}\geq(1+2\eta\tau(1-o(1))g_{\max})\,\tilde{\gamma}_ {min}^{(t)}\)__
3. \(\tilde{\gamma}_{max}^{(t+1)}\leq(1+2\eta\tau(1+o(1))g_{\min})\,\tilde{\gamma}_ {max}^{(t)}\)_, where_ \(\tilde{\gamma}_{max}^{(t)}:=\max_{\mu\in\{\pm\mu_{1},\pm\mu_{2}\}}\tilde{\gamma} _{\mu}^{(t)}\)_._

Theorem 3.1 is proved by iterating this lemma for \(\Theta(\log\log(d)/\eta)\) steps, yielding \(\gamma_{\min}\approx\tilde{\gamma}_{\min}=\omega(1)\).
Conclusion and Discussion

In this work, we showed that in \(\tilde{O}(d)\) samples, it is possible to learn the XOR function on Boolean data on a 2-layer neural network. Our results shows that by a careful analysis that compares that dynamics to the dyamincs under the surrogate \(L_{0}\) loss, we can show that SGD find the signal features, and escape the region of the saddle where it was initialized. Then, after learning the feature direction, we show that SGD will enlarge and balance the signal components to learn well-classify points from all 4 clusters.

We now discuss some of the limits and possible extensions of our techniques.

Minibatch SGD vs SGD vs GD.In this work, we study minibatch SGD, with a batch size of \(m\geq d\text{polylog}(d)\). This affords us enough samples at each iteration to have strong enough convergence to the population loss. Extending our results to SGD with a batch size of 1 is an interesting open question, and it is possible that this could be achieved using the drift-martingale techniques in Tan and Vershynin (2019); Arous et al. (2021); Abbe et al. (2023). Such methods allow larger fluctuations from the population loss at each step, but show that the fluctuations concentrate over time, even when SGD is run for \(T=\omega(1/\eta)\) steps, enough time to escape a saddle.

We remark that in this problem, using minibatch SGD with fresh samples can achieve stronger sample complexities than that required to show uniform convergence of the empirical gradient to the population gradient (as in Ge et al. (2017); Mei et al. (2018)), which in our setting, is \(\Omega(d^{2})\) samples. This means proving the convergence of GD on the empirical loss would require tools beyond uniform convergence.

Boolean Data vs Gaussian Data.One limitation of this work is that our results only hold for boolean data, and not gaussian data \(x\sim\mathcal{N}(0,I_{d})\). As a matter of convenience, it is easier to compute the population gradients \(\nabla_{w}L_{0}\) and \(\nabla_{w}^{\text{cl}}L_{\rho}\) with Boolean data, and the gradient does not depend on interactions between \(w_{\text{sig}}\) and \(w_{\text{opp}}\). With some willingness to perform various Gaussian integrals, we believe the analysis in Phase 1 could be extended to the Gaussian setting. This would require changing Lemma C.17 to reflect the population gradients, and modifying the definition of "strong" neurons (Def. C.13) to be a more restrictive set that only includes neurons where \(\|w_{\text{opp}}\|\ll\|w_{\text{sig}}\|\), such that \(w_{\text{sig}}\) grows at the maximum possible rate. We do not know of any way to directly extend Phase 2 to the Gaussian case. This is because if the cluster margins \(\gamma_{u}\) become very imbalanced, it is possible \(w_{\text{sig}}\) could grow in the wrong direction.

Classification vs Regression.In our classification setting, it suffices to show that the margin on each cluster grows large. We accomplish this in our Phase 2 analysis by showing that there is a large mass of neurons primarily in the \(\mu\)-direction for each \(\mu\in\{\pm\mu_{1},\pm\mu_{2}\}\). Adapting this strategy may be possible for XOR regression on Boolean data, but on Gaussian data, representing the ground truth function would require more specialization among the neurons. To see this, consider the following simpler example: to represent the single-index function \(f^{*}(x)=(e_{1}^{T}x)^{2}\) on Gaussian data on a ReLU network without biases, the neurons cannot all be oriented in the \(\pm e_{1}\) direction, otherwise the output would be \(a\sigma(x_{1})+b\sigma(-x_{1})\) for scalars \(a,b\). Studying the power of SGD to perform this specialization is an exciting open direction. We believe that our Phase 1 analysis may be a useful first step in this regard to show that the network can become signal-heavy. More powerful techniques would need to be developed to show specialization once the network contains sufficient signal.
