SGD Finds then Tunes Features in Two-Layer Neural Networks with near-Optimal Sample Complexity: A Case Study in the XOR problem
Authors: Margalit Glasgow
Venue: ICLR 2024 spotlight (2024)
Keywords: optimization
stochastic gradient descent
two-layer neural network
sample complexity
Primary area: learning theory
TL;DR: We show that standard minibatch SGD can learn the quadratic XOR function on neural networks in $d polylog(d)$ samples

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 Boolean hypercube labeled by the quadratic ``XOR'' function $y = -x_ix_j$ , it is possible to train to a population error $o(1)$
 with $\Theta(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 
 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 \em signal-finding \em phase where the network is small and many of the neurons evolve independently to find features, and a \em signal-heavy \em 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.

Published as a conference paper at ICLR 2024
SGD FINDS THEN TUNES FEATURES IN TWO-LAYER NEURAL NETWORKS WITH NEAR- OPPIMAL SAMPLE COMPLEXITY: A CASE STUDY IN THE XOR PROBLEM
Margalit Glasgow
Department of Computer Science
Stanford University Stanford. CA 94305, USA
(mglasgow) @ettnnordddde.eeoo
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 = -T it is possible to train to a population error o(1) with dpolylog(d)) (d) samples. Our result considers simultaneously training both layers 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 (dd) 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: signal-finding phase where the network is small and many the neurons evolve independently 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 learn features, since those neurons will be amplified by the simultaneous growth of their second layer weights.
find features, and signal-heavy phase, where
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 functions via SGD on neural networks. A particularly well-understood regime in this regard is the neural tangent kernel (NTK))Jacot al.) (021a)) where the network only moves a small distance from its initialization. However, in many cases, the NTK provably requires 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 lot of attention, is learning target functions that are inherently low-dimensional. depending only on a constant number of dimensions of the data (Chen & Meka 2020; Chen et al. 2020; Nichani al. 2022 Barak et al. 2022; Refinetti et al., 2021| Abbe et al. 2021a 2022: 2023). Such functions, often called sparse or multi-inde functions, can be written as f(x) := ((,x), where U E * has orthogonal rows, and is function on TR* Many works have shown that learning such target functions via SGD on neural networks possible in much fewer samples than achievable by kernel methods (Chen et al.) 2020) Bai & Lee 2019: Damian et al.) 2022 Abbe et al. 202 2022 2023). The results these papers apply to a large class of ground
et al.) 2022; Bietti et al., 2022 Mousavi-
* "htps://websstafford.cm cdu/ mglasgowindex.hmm html.

---

Published as a conference paper at ICLR 2024
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 it reflect standard training practices, for example using layer-wise training, changing learning rates, 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.]
The goal of this paper is to show that for simple but commonly-studied problem, standard minibatch SGD on 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) == -I,T with a sample complexity of Ö(d) ! test-beds for function which exhibits some of the core challenges of analyzing SGD 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.T).
We study the XOR function because it one of the simplest
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 (dd) samples to learn this function (Wei et al., 2019; Ghorbani et al. 2021]] Abbe et al.) (023). Wei et al. (Wei et al.) 2019) showed that ((d) samples statistically suffice, either by finding the global optimum of two-layer network, or by training an infinite-width network, both of which are computationally intractable.
Similar guarantees of (d) are given by Bai et al. (Bai & Lee 2019) and Chen et al. (Chen et al. 2020); however, such approaches rely on drastically modilying the network architecture and training algorithm to achieve quadratic neural tangent kernel. Abbe et al. (Abbe et al. 2023) proves a sample complexity of (dd 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. 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 better sample complexities.
Ben Arous et 2022). in some cases yielding
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 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. (Mei et al., [2018b] 2019), or showing convergence to a lower-dimensional SDE as in Ben Arous et al. (Ben Arous 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-heary phase. In this phase, we show inductively that throughout training, the signal components stay significantly larger than their counter- parts. This inductive hypothesis allows us to approximate the output of the network on sample 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. The division into signal-finding and signal-heavy phases is similar to the two phases of learning in e.g. |Arous et al. (2021).
Our Phase 2 analysis leverages the simultancous 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
'We consider this near-optimal in the sense that for algorithms that are rotationally invariant ((d) samples are required. See Section G for details.
Ben Arous et al. (022) considers setting of nigh-dimensional SGD where constont number of summary statistics sufficient to track the key features the SGD dynamics and the loss, which can only applied to constant-wind 2-layer neural networks. Their coupling between high-dimensional SGD and low-dimension SDE holds for @11) time, which is not enough time to learn the XOR function, which requires ((log(d) time.
2

---

Published as a conference paper at ICLR 2024
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 -igavl-hevyy 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 rate which is on the same order the rate of the signal components (see Remark 4.3)
1.1 RELATED WORK
Learning Multi-Index Functions via Neural Networks Most related to our work is a body of
work aiming understand the sample complexity of learning multi-index functions via SGD on
neural networks [Betti et al. (2022); Refinetti et al. (202T); [Chen et al. ;2020): Abbe et al. (2021a)
2022
)023): Damian al. (2022): Barak et al. (2022); ; Daniely & Malach (2020); Mousavi-Hosseeee
ct al. )
; Nichani et al. (202); ; Ge et al. (2017): Mahankali et al. (2023); Ba ct al. (2022);
Dandi al. 2023). Such functions are typically studied in either the Gaussian data setting where
~ (., Ta), or in the Boolean hypercube setting, where ~ Uniform( (f {+1}*) In both cases, we
have f(x) := g(Ux), where U projects r onto a lower dimensional space of dimension ks and g is
an arbitrary function on k variables. In the Boolean setting, U projects onto subset of k coordinates of a, so in the case of the XOR function we study, k = and g is a quadratie function.
Chen and Meka (Chen & Meka 2020) showed when k is constant, and is a degree- polynomial for constant D, there exists a polynomial-tin algorithm which learns such multi-index functions on Gaussian covariates in ((d) samples. Such algorithms can also be emulated in the same sample com- plexity via SGD on neural networks designed to emulate arbitrary Statistical Query algorithms (Abbe) & 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 ((dD) samples (Hsu 2021). A line of work by Abbe et al. (Abbe et al. 202 la
2022 2023) has conjectured that the sample complexity required for SGD is (ddmax(L-1 L denotes the "leap complexity". a measure of hicrarchical 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 miead-field analysis that is possible to learn with (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-wiss training give similar results for subclasses of multi-indes functions (Damian et al. 2022 Mousavi-Hosseim et al., 2022 Barak et al.
where
2022); Mousavi-Hossenne
ct al. (2022) studies a setting where k = 1 and L while Damian et al.
2002
;: Barak al. (2022)
study settings where L 2 2, and use just single gradient step on on the first laver, which req requires
((d samples. Numerous other works (Tan & Vershynin 2019 Bietti al. 2022 Wu et al. 2023 Arous et al. 202T) have made progress in the setting of single- index functions (k- 1) when L >>> I.
In some cases, the result achieve tight guarantees that depend on a quantity called the "information exponent" of g, which is equivalent to the leap complexity when k- = 1, though these methods require training only a single neuron in R.d. The recent work Mahankali et al. (2023) considers training a single-index target function with A and degree 4 on a 2-layer neural network via vanilla gradient descent, and shows a sample complexity of O (d which improves over kernel methods.
The above discussion highlights gap in our understanding when k 2 2 and L >>2 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 good stepping stone for understanding the behaviour of SGD on neural networks for more general functions with A: 2 2, T. > 2. Note that other works (Bai & Lee 2019 Chen et al. 2020h have achieved near-optimal sample complexity of Ö(d) for the XOR problems; these works use 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 O(d-1), which is larger than the rate achieved by Abbe et al. (2022) whenever
3

---

Published as a conference paper at ICLR 2024
L= and D 2 3. We emphasize that work achieves the near-optimal sample complexity (dd) with a standard two-layer neural network, trained with standard minibatch SGD.
We note that many
works have explored both empirically (eg. (Woodworth et al. 2020 Chizat
ct al. 2019)) and thcoretically veg.LLi et al. 2020, Allen- Li 2020; Suzuki & Akiyama 2020]
Telgarsky 2022 Jacot et al. 2021bb the cample-complxxtty advantages of "rich" SGD training the "lazy" NTK regime.
Simultaneous Training of Layers. While 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 & Bach, 2020; Lyu & Li, 2019 Lyu et al. 2021; Maennel et al. 2018). Under variety assumptions, these works have related the solutions found via gradient descent margin-maximine
solutions. A much finer understanding of the implicit bias of simultaneous training is provided for a line of work on diagonal neural
networks (Pesme & Flammarion 2023 Even al. 2023
1.2 ORGANIZATION PAPER
we describe the data and training model. In Section we state our result. In Section 4 we overview the proof techniques. We conclude in Section All proofs are in the Appendix.
In Section 2
1.3 NOTATION
For vector v, use || denote the E norm, and ||u||a denote the l norm. We use || to denote the spectral norm of matrix ... All big-O notation is with respect d co, and use Ö to suppress log factors in big-O notation. w(1) denotes growing to infinity with d. We use (r-1(r) to denote the sphere of radius in dimensions, and 1-)) to denote the indicator variable of event.
MODEL AND SETTING
2.1 DATA.
We study the setting where the data comes from the Boolean hypercube x Uniform((-)) (11,1)) and the label is given by y(z) = XOR(I1,I2) := -2.
Note that with μn el eg, and μ2 == es e2, we can model the distribution
(1i + 1)
w.p. 1/4
(-M++6.)) +1,1
w.p. 1/4 (-p2+ w.p. 1/4
(x, v)
(w2 -1)
w.p. 1/4
where & Uniform(O* (02 x xf f-1, (-1,1)d-2) 13d=2) so that 1 {21. We will often write
x =
where is the projection of onto the space spanned by e₁ and 62. and is the projection of orthogonal to and e2, We denote this distribution by Pa. and throughout, it is implicitly assumed that all probabilities and expectations over x are for Pa-
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.
2.2 TRAINING.
Model. We train both layers of two-layer ReLU network with width p:
1
a (u
p
j=1
4

---

Published as a conference paper at ICLR 2024
where (a)) max(0, a) is the ReLU function. We will use the variable p p Σラ= (1c to denote the empirical distribution of the neurons and their second layer weights. Thus denote
fa(x) Ev.a-pa aww" x),
We will often abuse notation and write probabilities and expectations using p, and use aw to denote its associated second layer weight. We note that is not necessarily the case the second layer weight aw is function of w: do this for the convenience of not indexing each pair as (wj az).
Initialization. We initialize the network with " Uniform(S*-1(00 (8)) for scale parameter 6, such that ||æ 8. We initialize the second layer as aj = where Cj ~ Uniform(=1).
Minibatch SGD. We train using minibatch SGD on the logistic loss function
lp(x) -2 log I [Ttex(-
and define the population loss Lp Euwplp(x). We will use the shorthand (x(x) denote the derivative of l (x()) with respect to fz(x):
exp( (x)) == ex exp(((())((())
1 exp(-v(z)f,
We use Pt. to denote the empirical distribution of the p neurons (wit), ( at iteration At each step, we perform the minibatch SGD update
((++1) = w(t) 1/ (D
a (++1)==0 a Vi
Here Lp = En EnMeM, L (x(( denotes the empirical loss with respect to minibatch Mr of m random samples chosen i.i.d. from Pa at step i, and for loss function and a parameter u in the AL denotes the scaled partial derivative of the loss with respect to u, defined in network, V.L p Pau particular for neuron (w, aw)) follows:
L
Σ
pl (
1
Σ z(9EM,
and (x oo' (o)
z0 EM
a
pl px
Σ
m
Daw
m
z(1 E,
z(EMM
3 MAIN RESULT
The following theorem is our main result.
Theorem 3.1. There exists a constant C such that the following holds for any large enough. Let = 1/ log C (d). Suppose we train a 2-layer neural network with minibatch SGD as Section 2.2 with minibatch size of m 2 d/0, width 1/0 S p S. dC, step size dCC < S O, and initialization scale 0. Then for some I. S C'log(d//n, with probubility - d d11),, we have
Ex [ (x)) < (log(d))
By setting 7 and d/0, Theorem 3.1 population loss in e (dpolylog(d) samples and iterations on a polynomial network.
states that
learn the XOR function up
Since the ReLU function is non-differentiable zero, define a' (0) 0.
*For convenience, we scale this derivative up by a factor of p to correspond to the conventional mean-field scaling. If we didn't perform this scaling, we could achieve the same result by scaling the learning rate
5

---

Published as a conference paper at ICLR 2024
Table I: Summary of Notation used in Proof Overview and Proofs
Hun w aw 2 0 12222 ow 0
12 w Wopp 12 w
aw 2 0 awe 0
W1.2 = Waig + Wopp
Usig
w__ = w W1:2
7/p So(W))(())
Mmin
( 7/a
Ymmx = max μu ( : 7pp
9a = 16(11)
Omin
min ppp)) ((((
Omax = maxRe pu 14 (x)
PROOF OVERVIEW
Throughout the following section, and in our proofs, will use the following shorthand to refer to the components of a neurons w. We decompose w = W1.2 + w1, where w1.2 is the projection of w in the direction spanned and 62 (and equivalently by /41 e1 e2 and μ/2 e1 e2), and to is the component of w in the orthogonal subspace. We further decompose W1.2 = Wsig + Wopp as follows:
Hu w Guupr w On < O.
Ow 2 0;
Wopp [ Pan w Ow 2 0: G
Wsig
Ou < 0.
Intuitively, we want the neurons to grow in the Wsig direction, but not the Wopp direction; in a network achieving the maximum normalized margin, will have w: Wsig exactly, and Wapp = wt 0. We summarize this notation in Table 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 z) that the loss (x(x) is well approximated by a first order approximation of the loss at f = 0, namely
Lo(:; p) := 2 log(1 -))))))((()(((a)
As long as this approximation holds, the neurons of the network evolve (approximately) independently, since (x) = ola(z:p) -y(x) does not depend the full network p. We will show under this
Øfp(a)
approximation that for many neurons, [| grows exponentially fast. Thus we will run this first phase for (log (d)/n) iterations until for all four clusters μ E { = there exists a large set of neurons Sp, on which wiig!" > 0, and the "margin" from this set of neurons is large, i.e.
3p = [1 E (pwamo(w" p)] >> Epllwi + Woppl|?
(4.1)
In the Phase 2, assume that Eq. 4.1 holds, and we leverage the dominance of the signal show that (1) The signal components Waig grow faster that Wopp + w.. and thus Eq. continues to hold; and (2) SGD balances the signal components in the 4 cluster directions such that the margins % balance, and become sufficiently large to guarantee (1) loss.
We proceed describe the analysis in the two phases in more detail. Full proofs are in the Appendix.
4.1 PIIASE
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 Lo := (o p) population gradients, and bounds on the difference V Lo V Lp.
Lo population gradients. Since the primary objective of this phase to grow the neurons in the signal direction, sketch here the computation of the gradient V.n Lo in the subspace spanned by M1, u22 The remaining estimates of Lo are simpler, and their main objective is to show that V Lo and V... Lo are sufficiently small, such that ||w doesn't change much throughout Phase and Jawl stays approximately the same as ||w]] For convenience, the reader may assume that jaw. = || exactly, which would hold if we took 7) to in gradient flow.
For a data sample Rad we denote @pan({#111 + By leveraging the symmetry of the data distribution and the fact that v(2)
E, where Span ((=m, H22)), and T

---

Published as a conference paper at ICLR 2024
(--2), compute
V w1 Lo -wwll
-wwBs 2 (=U(2) (o' g w² 3'ww" s- w²²22=
= --aB E.
2 lw0 El) sigsign(( (ongn((">> 2)2
1
(4.2)
aw E. sign( w 2)2P&ll ス |=1
2
1
2 Ow E. (() sign u² 2) N... 1 (G lwo =|]
1
2 Jwwa
2 ane (.)(2)) sign(w" 2)2
w
Here the two approximations come from the fact that has boolean coordinates and not Gaussian, and from an approximation of the Gaussian distribution, which holds whenever [w taking the expectation over E { = the last line of Eq 4.2 can be shown to evaluate to
is small. By
||w
Wsig +
Wopp-
(4.3)
[
2n
|| 2
Observe that near initialization, this gradient is quite small, since ||= is approximately [|
Va for
random initialization. Nevertheless, this gradient suggests that Wsig will grow exponentially fast.
Bounding the difference VL Vpp To bound || V Ww - V Lo la. first recall that
Lo V VuLp E (f,(2) (x)x)) x)=
Defining A,, = (e,(2) - ey (z))oo' ( (w=z), we can show using routine arguments (see Lemma D.2 for the details) that:
V V Lo la |ewl||Bn, Am S law SaulVEE /E,
(4.4)
la Ex)
a
|nw Bpllla x
polylog(d)
While this deviation bound is useful for showing that wi doesn't 't move too much this bound far exceeds the scale of the gradient in the Wsig- which is on the scale awl va near initialization. Fortunately,
can show in Lemma D.3 that the deviation is much smaller on the first two coordinates, namely,
|112 |w]
V wi L.
V wa Lollz < law O0log (d) E Illa w122 [
(4.5)
Note that since near initialization || |||||| for all neurons, this guarantee is much stronger than Eq. [4.4 In fact, since throughout this phase we can show that auy and ||=||| change relatively little, staying at the scale 1 //polylogdd), the approximation
in Eq. 4.5 is smaller than the gradient in the Wsig direction (Fq. 4.3 whenever say ||wag 2 100E,, [I|a, W1.2 which occurs on a substantial fraction of the neurons.
Lemma D.3 is the most important lemma in our Phase I analysis. At a high level, it shows that the approximation error W'aak Worz can be coupled with the growth of the signal, Way This is because we use symmetrization trick with the pairs = and -V... Ti
to show that both the error and the signal gradient only grow from samples x = z + where IT w| 1 [w]
In more detail, to prove Eq. 4.5
we also need to leverage the fact that for any & E {μ1 #2) and E { H we have (E 2) E 2')1 < 1pl, [], much smaller than can
7

---

Published as a conference paper at ICLR 2024
expect (x) - [x)] to be. Thus Aç+= A--sl I 4pEpplllaw whenever 15 m] 2 |22 w/
(such that o' d'ww" (E 2)) = d'(w" 2))). Following the symmetrization trick Eq. 4.2 we have
(V winz p Vun Lo)
= | 21
aw
|| Ez Ae+22ll
(EgB,((5::) Ac=:)=11
- 2v 2/2E,lllaw a. |l| + (2E w| S w (=(w))) |
careful computation comparing w Gaussian distribution then shows that
|| ||
E2 E1, w| & |z" ( ~ (1221 w/ S wlll (Ex
Ep [|l
Putting Phase 1 Together The building blocks above, combined with standard concentration bounds on Vi.,. suffice to show that a substantial mass of neurons will evolve according to Eg 4.3 4.3 leading exponential growth in Wogg After (log(d))m) iterations, for these neurons, can achieve |[Wsigg > | + Woppl- Formally, we show the following for some Ç = 1/polylog(d):
Lemma 4.1 (Output of Phase Informal; See Lemma f.r formal version). With high probability, for 1 Ó(1) and some ( = 1/polylog(d), after some T (llo(dd)m) iterations of minibatch SGD, with m (dd samples each minibatch, the network pr satisfies:
1. Ep [ + Wappl 1 |
2. For each μs E 1 on al least 0.1 fraction of all the neurons, we have Wigh
and ||Wsiel 1| 10
We remark that the analysis to prove Lemma 4.1 is somewhat subtle, since the tight approximation in
Eq 4.2
breaks down when ||aag||| approaches [w The details are given in Appendix D
4.1.1
PIIASE 2
The conclusion of Lemma 4.1 is sufficient condition of the network to begin the second phase. In the second phase, we have that (for most x)
(x(x) R (2),
(4.6)
where we recall that z is the component of x in the space spanned by μi and μ22 We refer to this as the clean loss derivative, and our main tool will be analyzing the evolution SGD under this clean surrogate for the loss derivative. Namely, we define:
w Lp == an E. (z) (w" x)x and Vol Le == Eo =a,
(4.7)
Before proceeding, we introduce the following definitions, which will useful in Phase 2 (summa- rized in Table T):
exp( -/max)
Ymin
min pE [ b22
7a
Imin
min
1,,44)]=
1 exp( -mmax)
exp( -Ymin)
Ymax
max ee
7/
Imax
max pE (E(L,,)pp))
(1)())) 1 exp( (mimm
To ensure the approximation in Eq. h.6lsolds 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 can--sinal-direttinn components of the network. Formally, we consider the following condition.
Definition 4.2 (Signal-Heavy Inductive Hypothesis). For parameters Ç = o(1) and H 1 with < 10/)), we say network is (C, -iignal-llll if there exists some set of heavy neurons on which exp ()) [| |[olell S <llWaigh and
1 & wol < S-min-
8

---

Published as a conference paper at ICLR 2024
Here we have defined Fi == E[1(w S. wight > 0) (w" 1i) and Amin min (=tn, Pi- Further,
Ep||[/1?]] S a CH s 2H,
for all neurons, have |
||/l||
We show via straightforward argument in Lemma E.4 that if the conclusion of Lemma 4.1 (from Phase 1) holds for some Ç, then the network is (0(ç (5) H) J-signal-heavy for H (log log (d)).
Assuming that the network is (C. H) --signall heavy, using similar approach to Eq. 4.4 we can show (see Lemma E.5 for the precise statement) that for any neuron (w, awc)
V Wopp - V 2 E
V (Ea(s fo E E |I| S < 7max:
and similarly [V 1, - Vol L 2 |Re|/Smmax.
By working with the clean gradients, is possible to approximately track (or bound) the evolution of Wsig. L., and Wopp on neurons in S. the set of neurons for which W >> || + Woppl-- In and E,8 we show the following for any w (Iet p be the direction of Wsig):
Lemmas E.6
E.7
1. The signal component Wisg grows quickly. We have -w V
(1)))(()) where v2_ Also aw grows at a similar rate. This growth is due to the fact that points with 2 = -14 will almost never activate the ReLU, while points with z = 11 almost always will.
7 -/2
2. A linear combination of || 2|? and |!pppl! ??? decreases. The argument here more subtle, but the key idea is to argue that if |ww 2 @TTp=| frequently, then |||_|| will decrease. Meanwhile, if IwT gl I | frequently, then Wopp will decrease (and there is sizeable
event 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 =1e, (1²)
for the "weakest" cluster, that is, in the direction μ that maximizes F/- will
have ((+1) 2 (1 (max) )
min
On neurons outside of S, we show in Lemma E. that they grow at most as fast as the rate of the weakest clusters, meaning we 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 T(" (1)), ,, and at (up to) an equal linear rate outside of S. Thus the neurons outside of S might eventually attain non-negligible mass. However, because the layers are trained simultaneously, this leads to positive feedback between the growth of || Wsseell and Jaw.), l. leading to exponential growth, maintaining 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 slightly worse parameter), the behavior of the weakest margin improves, and the network (measured by the size of the largest margin 7mmax doesn't 't become too big.
Lemma 4.4 (Phase 2 Inductive Step: Informal; See Lemma E.3 for formal version). If a network Pe is (C. -signal heavy with heavy set S, then after one minibatch gradient step, with probability 1 d d1),
1. P1i is (C(1 + 10mg H)) signal heavy.
2. (1+1) min
2 (1 + 21/7(1 gmax)) Imin
3. Ymax
< + 2777(1 + (mi)) Ymars where Ymmx := m (1p1. 7/
Theorem 3.1 is proved by iterating this lemma O(log log())/)) times, yielding Vmin 7min w(1).
5 CONCLUSION
In this work, we showed that in Ô(d) samples, is possible to learn the XOR function on Boolean data 2-yeyer neural network. Our results shows that by a careful analysis that compares that

---

Published as a conference paper at ICLR 2024
dynamics the dyamines under the surrogate Lg loss, 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 discuss some of the limits and possible extensions of our techniques in Section A
show that SGD find the signal features,
REFERENCES
Emmanuel Abbe and Colin Sandon. On the universality of deep learning. Advances in Neural Information Processing Systems, 33:20061-20072, 2020.
Emmanuel Abbe, Enric Bdixaa Matthew S Brennan, Guy Bresler, and Dheeraj Nagaraj. The staircase property: How hierarchical structure can guide deep learning. Advances in Neural Information Processing Systems, 34:26989-27000 2021a.
Emmanuel Abbe, Pritish Kamath, Eran Malach, Colin Sandon, and Nathan Srebro. On the power of differentiable learning versus pac and sq learning. Advances in Neural Information Processing Systems, 34:24340-24351, 2021b.
Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-staircase property: necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory, pp. 4782-4887 PMLR, 2022.
Emmanuel Abbe, Enric Boix-Adsera, and Theodor Misiakiewicz. Sgd learning neural networks: lapp complexity and saddle-to-saddle dynamics. arXiv preprint arXiv:2302. 11055, 2023,
Zeyuan Allen-Zhu and Yuanzhi Li. Backward feature correction: How deep learning performs deep learning. arXiv preprint AXXv:2000.00411: 2020.
Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensional inference. The Journal of Machine Learning Research, 22(1):4788 4838, 2021.
Jimmy Ba, Murat Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional asymptotics of feature learning: How gradient step improves the representation. Advances in Neural Information Processing Systems, 35:37932-37946, 2022.
Yu Bai and Jason D Lee. Beyond linearization: On quadratic and higher-order approximation of wide neural networks. arXiv preprint arXiv:1910.016/ 2019.
Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. Advances in Neural Information Processing Systems, 15:22750-,,764, 2022.
Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. High-dimensional limit theorems for sgd: Effective dynamics and critical scaling. Advances in Neural Information Processing Systems, 35: 55349-25362. 2022.
Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks. arXiv preprint arXiv:2210.1566 2022.
Minshuo Chen, Yu Bai, Jason D Lee, Tuo Zhao, Huan Wang. Caiming Xiong, and Richard Socher. Towards understanding hierarchical learning: Benefits of neural representations. Advances in Neural Information Processing Systems, 33:22134-22145. 2020.
Sitan Chen and Raghu Meka. Learning polynomials in few relevant dimensions. In Conference on Learning Theory, pp. 1161-1227 PMLR, 2020.
Lenaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Conference on Learning Theory, pp. 1305-1338. PMLR, 2020.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in neural information processing systems, 32, 2019.
10

---

Published as a conference paper at ICLR 2024
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference on Learning Theory, pp. 5413-5452. PMLR, 2022.
Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. How two-layer neural networks leam, one (giant) step at time. In NeurlPS 2023 Workshop Mathematics of Modern Machine Learning, 2023.
Amit Daniely and Eran Malach. Learning parities with neural networks. Advances in Neural Information Processing Systems, 33:20356-20365 20365, 2020.
Mathieu Even, Scott Pesme, Suriya Gunasekar, and Nicolas Flammarion. (s) gd over diagonal linear networks: Implicit regularisation, large stepsizes and edge of stability. arXiv preprint arXiv:2302.08882, 2023.
Spencer Frei, Niladri S Chatterji, and Peter L Bartlett. Random feature amplification: Feature learning and generalization in neural networks. arXiv preprint arXiv:2202.07626 2022.
Rong Ge, Jason Lee, and Tengyu Ma. Learning one-hideen-layrr neural networks with landscape design. arXiv preprint 1xXi:/7111000000 2017.
Behrooz Ghorbani, Song Mei. Theodor Misiaiewicz, and Andrea Montanari. Linearized two-layer neural networks in high dimension. The Annals, 49(2):1029-1054: 2021.
Daniel Hsu. Dimension lower bounds for linear approaches to function approximation. Daneil Hsu's homepage, 2021.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: convergence and generalization in neural networks. In Proceedings of the 53rd Annual ACM SIGACT Symposium on Theory of Computing, pp. 6-6, 2021a.
Arthur Jacot, François Ged, Berfin Şimşek, Clément Hongler, and Franck Gabriel. Saddle-to-saddle dynamics in deep linear networks: Small initialization training, symmetry, and sparsity. arXiv preprint arXiv: 00t.//...15590 15933, 2021b.
Yuanzhi Li, Tengyu Ma, and Hongyang R Zhang. Learning over-parametrrized two-layer relu neural networks beyond ntk. arXiv preprint arXiv:2007.04596, 2020.
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint arXiv: 1iX::1906005890 2019.
Kaifeng Lyu, Zhiyuan Li, Runzhe Wang, and Sanjeev Arora. Gradient descent two-layer nets: Margin maximization and simplicity bias. Advances in Neural Information Processing Systems, 34:12978 12991, 2021.
Hartmut Maennel, Olivier Bousquet, and Sylvain Gelly. Gradient descent quantizes relu network features. arXiv preprint arXiv:1803.08367. 2018.
Arvind Mahankali, Jeff Z Haochen, Kefan Dong, Margalit Glasgow, and Tengyu Ma. Beyond ntk with vanilla gradient descent: A mean-field analysis of neural networks with polynomial width, samples, and time. arXiv preprint arXiv:2306. 16361, 2023.
Song Mei, Yu Bai, and Andrea Montanari. The landscape of empirical risk for nonconvex losses. The Annals of Statistiics, 4666A):2747-2774 2018a.
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two- layer neural networks. Proceedings of the National Academy of Sciences, L15(33)E7665-7777, 2018b.
Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. In Conference on Learning Theory, pp. PMLR, 2019.
Alireza Mousavi-Hosseine Hosseini, Sejun Park, Manuela Girotti, Ioannis Mitliagkas, and Murat A Erdogdu. Neural networks efficiently learn low-dimensional representations with sgd. arXiv preprint a209 14863, 2022.
11

---

Published as a conference paper at ICLR 2024
Eshaan Nichani, Yu Bai, and Jason D Lee. Identifying good directions to escape the ntk regime and efficiently learn low-degree plus sparse polynomials. arXiv preprint arXiv:2206.00 036688 2022.
Scott Pesme and Nicolas Flammarion. Saddle-to-saddle dynamics in diagonal linear networks. arXiv preprint X1v:22304000488, 2023.
Maria Refnettti, Sebastian Goldt, Florent Krzakala, and Lenka Zdeborová. Classifying high- dimensional gaussian mixtures: Where kernel methods fail and neural networks succeed. In International Conference on Machine Learning, pp. 5936-8947, PMLR, 2021.
Taiji Suzuki and Shunta Akiyama. Benefit of deep learning with non-convex noisy gradient descent: Provable excess risk bound and superiority kernel methods. In International Conference Learning Representations, 2020.
Yan Shuo Tan and Roman Vershynin. Online stochastic gradient descent with arbitrary initialization solves non-smooth, non-convex phase retrieval. arXiv preprint arXiv:1900.122 1980.12377 2019.
Matus Telgarsky. Feature selection with gradient descent on two-layer networks in low-rotation regimes. arXiv preprint arXiv:2208.02788 02789, 2022.
Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets vs their induced kernel. Advances in Neural Information Processing Systems, 32, 2019.
Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory, pp. 3635-3673, PMLR, 2020.
Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Sham M Kakade. Learning high-dimensiooal single-neuron relu networks with finite samples. arXiv preprint SXXi:2303300225,, 2023.
12


Reviews:
[{"review_id": "hWYhTgH7Yy", "reviewer": "Reviewer_kd4Q", "summary": "The authors investigate the problem of learning a XOR function on Boolean data using two-layer ReLU networks trained with minibatch stochastic gradient descent (SGD). By investigating the training dynamics of SGD by jointly training the two layers on the logistic loss they identify two phases: a first phase in which the output network is close to zero (\"signal-finding\"), and a second stage (\"signal-heavy\") )in which the signal components stay larger than their counterparts. They provably show that with dpolylog(d) samples the training algorithm reaches a o(1) population error.", "questions": "- I believe these works are relevant for the comparison with the literature in this context [1,2,3]. In [1] the authors show that even one step of GD in the $n=O(d)$ regime with a large learning rate beats kernel methods, but only a single index approximation of the target is learned. in [2] the authors expand the findings of [1] to the $n=O(d^l)$ scenario and investigate the learning of multi-index functions with a finite number of GD steps by extending the staircase condition to Gaussian data and a large batch setting. Finally, [3] proves that a smoothed version of SGD matches correlation statistical query lower bound: if the target function has information exponent $k$, smoothed-SGD attains $O(d^{\\frac{k}{2}})$ sample-complexity.\n\n-  Ben Arous et al. [2021] is present in the references but not in the text. What is the relationship between the sample complexity obtained in this work for XOR (multi-index) and for $k = 2$ (information-exponent) single index target functions? \n\n- It is just briefly mentioned how the findings differs from the results of Ben Arous et al. [2022], a better comparison must be done. Could the author comment on this point?\n\n-  Suggestion of relevant reference for the Saddle-to-Saddle dynamics [4].\n\n- The analysis of the joint training of the two layers is nice and welcome. Could the author comment on the consequences in terms of transfer learning? After the training procedure, has the first layer learned the relevant feature subspace such that other tasks dependent on the subspace U could be learned?  \n\n\n[1] Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High-dimensional asymptotics of feature learning: How one gradient step improves the representation [2022].\n\n[2] Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. How two-layer neural networks learn: one (giant) step at a time [2023]. \n\n[3] Alex Damian, Eshaan Nichani, Rong Ge, and Jason D. Lee. Smoothing the Landscape Boosts the Signal for SGD: Optimal Sample Complexity for Learning Single Index Models [2023].\n\n[4] Arthur Jacot, François Ged, Berfin Şimşek, Clément Hongler, and Franck Gabriel. Saddle-to-saddle dynamics in deep linear networks: Small initialization training, symmetry, and sparsity [2021].", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The authors analyze an important problem: learning a target function that depends only on two relevant directions with two-layer neural networks. The training algorithm considers joint updates of the first and second layer, this protocol is more closely connected to real-life settings than the layerwise training explored in many publications in the multi-index learning literature.", "weaknesses": "The manuscript's biggest weakness is the clarity of the exposition and its relationship with previous works. A more in-depth discussion is provided below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UA8lsqzOoE", "reviewer": "Reviewer_chEd", "summary": "This paper studies the problem of learning a degree-2 parity function (XOR) with a 2-layer neural network with ReLU activations under the logistic loss function. Particularly, the training is done with mini-batch SGD where both layers are trained simultaneously. It is further shown that the sample complexity is $\\tilde O(d)$ which is also shown to be optimal for rotationally-invariant algorithms. \n\nIn the proof, it is shown that SGD increases the weights incident to the parity bits for a fraction of neurons. Afterwards, these neurons get reinforced making the learning possible.", "questions": "- Q1. What modifications do you expect to be necessary so that similar analysis can be carried out for higher degree parities?\n- Q2. Using $\\ell_2$ loss is already discussed in the appendix but for the Gaussian setting. What are the challenges of using $\\ell_2$ (or other loss functions) for the Boolean parity analysis?\n\nRemark. It would be beneficial for readers if further discussion about the proof techniques (and their comparison with recent work) is included.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- Including a rather extensive literature review and also mentioning the limitations of their analyses and avenues for future work (e.g., discussion about batch size and Gaussian setting).\n- Providing the analysis in the case that both layers are trained jointly with reasonable learning rates. Also, there are no uncommon modifications to the algorithm such as projections, or changes in the learning rate, ...\n- Interestingly they discuss that joint training of the layers may have a regulating effect that makes the learning possible (see Remark 4.3).", "weaknesses": "Nothing in particular, the limitations of the work are clearly stated in the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ym0cgoIvvU", "reviewer": "Reviewer_xNLA", "summary": "This paper studies the classification problem of learning XOR functions on Boolean hypercube. The setting is standard that using 2-layer ReLU neural networks with vanilla online minibatch SGD on logistic loss. The authors show that neural nets can achieve $o(1)$ population loss with $d \\mathrm{polylog}(d)$ samples. This is near-optimal sample complexity in the sense that all rotationally invariant algorithms need at least $d/\\mathrm{polylog}(d)$ samples. The proof shows a 2-stage dynamic that in the first stage the neurons learn the features and grow large to reduce loss in the second stage. The analysis approximates the real dynamics with dynamics induced by certain surrogate loss.", "questions": "-\tWhat is the initialization scale in Theorem 3.1? On the top of page 5, the scale is $\\theta$. However, it becomes $\\theta/\\sqrt{p}$ in the statement of Theorem 3.1.\n-\tI was wondering why stepsize $\\eta$ is required to be not very small ($>1/\\mathrm{poly}(d)$) in Theorem 3.1. Are there any fundamental difficulties in the proof that that stepsize cannot be too small?\n-\tI wonder what the order of $o(1)$ test loss is in Theorem 3.1. Is it $1/\\mathrm{polylog}(d)$ or $1/\\mathrm{poly}(d)$? It seems not very clear in the main text. Or in general, what is the dependency on test loss $\\epsilon$ for sample complexity and running time？\n\nMinor:\n\n-\tAbove eq 4.1: ‘ie.’ -> ‘i.e.,’\n-\tIn Definition 4.2, after ‘Here we have defined’: $w_{sig}^T$ -> $w_{sig}^T x$", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "-\tUnderstanding the optimization for neural networks in feature learning regime is an important and interesting problem. This paper takes a step to go beyond single-index models to more complex XOR functions.\n-\tThe paper is overall well-written and easy to follow. A proof sketch is provided to illustrate some important proof ideas.\n-\tThe paper shows that vanilla online minibatch SGD on 2-layer neural networks can learn XOR function efficiently with $d \\mathrm{polylog}(d)$ samples. This seems to be new in literature.\n-\tThe proof techniques introduced in this paper also seem to be novel, especially the observation that the variance of gradient on the signal entries can be shown to be much smaller than the signal (eq (4.5)).", "weaknesses": "-\tThe results essentially still rely on the small initialization to learn features at the beginning, which is similar to many previous works.\n-\tAs discussed in Appendix A, it might be difficult to generalize the techniques to other settings (though it is understandable).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5llsH533gG", "reviewer": "Reviewer_Y9yq", "summary": "This paper proves that 2-layer neural networks trained with SGD can learn the XOR function on the boolean hypercube using $\\tilde{O}(d)$ samples. Unlike prior work, the analysis doesn't require layer-wise training or the use of random biases, which is thanks to carefully tracking the evolution of the weights in two distinct phases and exploiting the symmetries of the problem. The analysis reveals that in the first phase the weight vectors of a significant number of neurons align with the XOR features, thus feature learning occurs. In the second phase, the neurons that recovered the relevant features will grow and dominate the prediction of the network, thus the XOR function is learned.", "questions": "* Perhaps the authors could highlight the $\\varepsilon$-dependency (test error) of the rate in Theorem 3.1. This question is less relevant in some high-dimensional settings where the goal is only to learn to constant accuracy, but may still provide some insights given that the authors consider the specific problem of learning XOR.\n\n* It seems that the algorithm spends more time in the first phase $\\Theta(\\log(d) / \\eta)$ than the second phase $\\Theta(\\log\\log(d) / \\eta)$. This might be interesting to point out, as it is also in line with observations from Ben Arous et al., 2021, who show that SGD spends most of its time in the *weak recovery* phase.\n\n* As a minor comment, it might be helpful to mention the scale of $\\zeta$ inside Lemma 4.1 to make the statement of the lemma more self-contained.\n***\nBen Arous et al. \"Online stochastic gradient descent on non-convex losses from high-dimensional inference.\" JMLR 2021.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* This paper presents a novel analysis of feature learning in two-layer neural networks and distinguishes itself from prior work by considering a standard (unmodified) SGD training algorithm and achieving sample complexity that is almost optimal among the class of rotationally invariant algorithms. \n\n* Plenty of intuition is provided on the mechanism of the proofs, and the two-phase dynamics could potentially be used for analyzing feature learning in other problems as well.", "weaknesses": "* The considered problem can be a bit restrictive, and for example could be generalized to learning $k$-sparse parities. This is not a major issue however, as this paper is one of the first attempts at providing a feature learning analysis of *standard* SGD without additional modifications.\n\n* As the authors have also pointed out, the batch size chosen is relatively large compared to practical settings, and it is interesting to know whether a smaller batch size can also learn XOR with near-optimal sample complexity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)