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

Reviewer comment (Theory):
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.

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

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

Reviewer comment (Theory):
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.

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

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

Reviewer comment (Novelty):
The results essentially still rely on the small initialization to learn features at the beginning, which is similar to many previous works.

Author response:
**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.

Reviewer comment (Theory):
I 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?

Author response:
**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)$.

Reviewer comment (Theory):
I 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?

Author response:
**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.

Reviewer comment (Theory):
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?

Author response:
**$\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.

Reviewer comment (Theory):
It would be beneficial for readers if further discussion about the proof techniques (and their comparison with recent work) is included.

Author response:
**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.

Reviewer comment (Writing):
The manuscript's biggest weakness is the clarity of the exposition and its relationship with previous works.

Author response:
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?

Reviewer comment (Novelty):
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.

Author response:
**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.

Reviewer comment (Theory):
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?

Author response:
**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*.

Reviewer comment (Novelty):
It is just briefly mentioned how the findings differ from the results of Ben Arous et al. [2022], a better comparison must be done. Could the author comment on this point?

Author response:
**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.*

Reviewer comment (Novelty):
Suggestion of relevant reference for the Saddle-to-Saddle dynamics [4].

Author response:
**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.

Reviewer comment (Evaluation):
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?

Author response:
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.