SGD Finds then Tunes Features in Two-Layer Neural Networks with near-Optimal Sample Complexity: A Case Study in the XOR problem

We thank the reviewer for their positive assessment and insightful comments and questions!

• Convergence rate of test error. A minor revision to our calculation (which is now in the uploaded revision) yields that $\gamma_{min} = \Omega(\log\log(d)),$ and so the loss is $O(1/polylog(d))$.

  For $\epsilon$ smaller than $O(1/polylog(d))$, we expect the number of iterations should be roughly $\tilde{O}(1/\epsilon)$, matching the convergence rate of logistic regression. However, we do not prove this in our paper because it would require modifying Lemma E.3, the Phase 2 inductive lemma.

• Indeed, our findings are similar to Ben Arous et al. 2021, in that that there is a long search ("weak-recovery") phase in which the directions of the multi-index function are learned but the loss does not decrease, and then after that phase, the loss decays quickly. We will point this comparison in the introduction.

• We have revised to mention the scale of $\zeta$ (it is $1/\log^c(d)$).

We thank the reviewer for their positive assessment and insightful comments and questions!

Initialization scale. We apologize for the typo - the initialization scale should always be $\theta$; we have fixed this in the revision. We note that our assumption on a $1/polylog(d)$ initialization is much weaker than comparable results that leverage a small initialization (eg. Abbe et al. 2022 initializes the first layer weights at 0 exactly --- see "Choice of Hyperparameters" in section E.1). Similarly to Abbe et al 2022 (and other works), our small initialization keeps the output of the network close to 0 for the first phase of training, which indeed simplifies the analysis. However, we note that because our initialization is still within log(d) factor of 1, we still have to be careful .

It may be possible to remove the small initialization assumption by claiming that since the second layer weights are mean-0 and random, the network would still whp output a value close to 0 throughout phase 1, though we leave this question to future work.

Step size If the step size is smaller, we would need more iterations of SGD (and thus more samples) to learn the XOR function. Escaping the saddle takes $T = \Theta(\log(d))$ time, where $T = \eta t$ and $t$ is the number of iterations. This is because learning the signal mimics a power iteration where the signal components of the neurons grow like $w_{sig}^{(t)} \approx w_{sig}^{(0)} \exp(\eta t)$.

Test loss An exact bound on the loss can be found by looking at the details of the proof of Theorem 3.1 (at the end of Section F), which repeatedly applies Lemma E.3, the Phase 2 inductive lemma. As stated, the calculation currently given there (which was not optimized to give the best dependence on $d$) shows that we obtain $\gamma_{min} = \Omega(\log\log\log\log(d)),$ and so by the final calculation in this proof, we have that the loss is $O(\exp(-\gamma_{min}) = O(1/\log\log\log(d))$.

A minor revision to our calculation (which is now in the uploaded revision) yields that $\gamma_{min} = \Omega(\log\log(d)),$ and so the loss is $O(1/polylog(d))$.

For $\epsilon$ smaller than $O(1/polylog(d))$, we expect the number of iterations should be roughly $(1/\epsilon)$, matching the convergence rate of logistic regression. However, for simplicity, we do not prove this in our paper because it would require modifying Lemma E.3, the Phase 2 inductive lemma.

We thank the reviewer for their positive assessment and insightful comments and questions!

Modifications for higher degree parities. To represent a higher order parity via relu-activated neurons, you need more neurons, and the roles of the neurons are not ``exchangable'' (e.g. exhibit a certain symmetry with respect to the data distribution). In the 2-parity problem, this means we expect $w_{sig}$ to grow in each of these 4 directions at approximately the same rate. This phenomenon may not occur for higher order parities. For example, for a 3-parity $y = -x_1x_2x_3$, the rate of learning the direction $(1, 1, 1)$ may be differently from the rate of learning the $(1, 1, -1)$ direction.

Thus to carry out this analysis for higher order parities you would need to:

1. Determine the set of neurons necessary to represent the parity. [Note, there may be multiple options]
2. Use a mean-field like analysis to predict the rates at which these different directions will be learning, near initialization. This step can also confirm which set of neurons will represent the parity, if there were many options in Phase 1.
3. If the neurons grow at the same rate, use an analysis similar to our paper with more directions to learn.
4. If the neurons do not grow at the same rate, strengthen Phase 2 of our analysis so that it can hold for \Theta(\log(d)) iterations. Then modify our analysis to look like this:

• Phase 1: Learn the directions learned fastest.
• Phase 1.5 (A combination of our phase 1 and 2, repeat as needed): Learn the directions which are learned slower, while showing that the signal in the already-learned directions is not forgotten.
• Phase 2. Grow the signal components of the neurons, as in our Phase 2.

$\ell_2$ loss for Boolean data. We expect using $\ell_2$ loss could work, though it would require modifying our Phase 2 to yield a slightly stronger result. Namely, we would not just need to show that the margin in each of the 4 cluster directions is large. Rather, we would need to show that the components learned in each of the 4 directions are exactly equal. Presumably the same sort of Phase 2 "signal-heavy" inductive hypothesis should work. There may be some additional challenges due to the fact that the signal may decrease if it grows too large.

Discussion of Comparison of Proof Techniques Due to space limitations, we will add this comparison to the Discussion section in Appendix A. We will discuss:

• How our decomposition into two phases of analysis is similar to other works, such as Mahankali et al. 2023, Ben Arous et al. 2021.
• How we leverage the small initialization of the network, similarly (but not to the same extent as) other works such as Abbe et al. 2022.

We thank the reviewer for their suggestions, particularly to enhance our comparisons to related work.

We have added reference to [1, 2, 3, 4] in the the introduction:

• We discuss [1, 2] in the subsection "Learning Multi-Index Functions via Neural Networks" of the related work, where we discuss works that take 1 step of GD on the first layer. We note that [2] was not available online at the time of our submission.
• We add [3] and Ben Arous et al. 2021 in our discussion of learning single-index functions in the related work (middle of page 3). Thank you for catching that we missed citing Ben Arous et al. 2021 in the body.
• We have added [4] to the discussion at the end of the subsection "Learning Multi-Index Functions via Neural Networks" in the related work, where we discuss differences between kernel and rich regimes.

What is the relationship between the sample complexity obtained in this work for XOR (multi-index) and for (information-exponent) single index target functions? As explained in the related work, it is conjectured that the sample complexity of SGD on a 2-layer NN is $\tilde{\Theta}(d^{\max(L-1, 1)})$, where $L$ is the leap complexity, which is $2$ for XOR. In the case of single-index functions, the leap complexity equals the information exponent, and thus the sample complexity of learning XOR via SGD matches that of learning a single index function with information exponent $k = 2$. We have clarified this connection between leap complexity and information exponent in the related work.

Better comparison to Ben Arous et al. 2022. Ben Arous et al. 2022 considers a setting of high-dimensional SGD where there are a constant number of summary statistics which are sufficient to track the key features of the SGD dynamics and the loss. There are several differences between the setting they study and our, and the reasons why their tools cannot be applied to our setting.

• Their work can be applied to a 2-layer network with constant width, but not one with width growing with $d$, as we study, because this would not have a constant number of summary statistics.
• They study the XOR problem in a setting where the cluster separation $\lambda$ grows to infinity.
• Their coupling between high dimensional SGD and a low-dimension SDE only holds for $\Theta(1)$ time, which is not enough time to learn the XOR function --- because the initialization of the network is near a saddle, $\Theta(\log(d))$ time is necessary. We have clarified this comparison in a footnote where Ben Arous et al. 2022 is referenced.

Yes, transfer learning could occur after our training procedure since the directions of $U$ will be learned by the first layer weights. This would follow from a similar proof that to Damian et al. 2022. We do not include this because the reasoning would be the same, and is fairly straightforward.

The reviewer also says that a key weakness is the "clarity of exposition", but there are no details provided. If the reviewer has any suggestions that would be appreciated. Otherwise, would the reviewer consider increasing their score?