Paper: Benign Overfitting and Grokking in ReLU Networks for XOR Cluster Data (ICLR-2024)

Reviewer comment (Experiments):
Do the authors have any ideas for improvements in the high-dimensional setting?

Author response:
Thank you for your question. It is important to note that previous studies on benign overfitting, including the one you referred to, only focused on linearly separable data. Our data distribution, however, is not linearly separable, which leads to qualitatively different behavior and a more complex trajectory analysis. Specifically, a new technical challenge in our analysis is providing concentration bounds for sum of dependent random variables that are not weakly dependent (i.e. $\text{Cor}(X_t, X_s) \rightarrow 0$ as |t-s| goes to infinity), as detailed in Equations (A.17)-(A.22) and Lemma A.17 in Section A.6. There are few probabilistic tools that can handle such a summation, and in order to handle this concentration we need the stronger condition $p>n^5$. The proof of this concentration inequality is the only place that requires $p > n^5$, and $p>n^2 \log(n)$ suffices for the remaining  proofs.

Our paper focuses on theoretically capturing the qualitative behavior of (non-linear) benign overfitting and grokking (which we believe is a significant contribution, as Reviewers 6CDP and mTpn both noted in their reviews), instead of optimizing all the technical assumptions. We leave the relaxation of this assumption for future work.

Reviewer comment (Theory):
What behavior does the test error exhibit when the training time is between 1 and $n^{0.01}$?

Author response:
Regarding the test error behavior between time $1$ and time $n^{0.01}$: for technical reasons, we cannot prove a generalization bound for the behavior between times 1 and $n^{0.01}$. That said, we are able to prove that the expectation of the margin keeps increasing with step $t$, which intuitively should help generalization. Empirically, Figure 1 shows that the test accuracy increases gradually from near-random to near-perfect.

Reviewer comment (Theory):
The assumption made in A1 seems to contradict common understanding. Generally, increasing the number of samples, even with limited noisy labels, tends to enhance the generalization capability of neural networks. However, in Assumption A1, having a larger number of training samples seems to adversely affect the model, as indicated by its presence on the right-hand side of the inequality. This aspect might require further clarification or justification within the context of the study.

Author response:
We appreciate the question regarding the potential negative impact of increasing the number of samples on the model's performance, and agree that it is a bit counterintuitive. We wish to mention that an assumption on the upper bound of $n$ was needed in all previous benign overfitting results (since they all require high-dimensional data), e.g. [[Bartlett et al. (2019)](https://www.pnas.org/doi/10.1073/pnas.1907378117), [Frei et al. (2022)](https://arxiv.org/abs/2202.05928), [Cao et al. (2022)](https://arxiv.org/abs/2202.06526), [Frei et al. (2023)](https://jmlr.org/beta/papers/v24/22-1132.html), [Xu & Gu (2023)](https://proceedings.mlr.press/v206/xu23k.html), [Kou et al. (2023)](https://arxiv.org/abs/2303.04145)]. Also note that in practical settings it is possible that increasing the sample size could hurt generalization (see e.g. “deep double descent” [[Nakkiran et al. (2019)](https://arxiv.org/abs/1912.02292)]). 

Empirically, such an upper bound for sample size is indeed necessary. We have added an experiment for large sample size $n$ (see Figure 5 in Appendix A.8), which empirically shows that increasing $n$ could harm generalization, so it is not just a limitation of the analysis. 

Additionally, we also refer the reviewer to Figure 2 in [Frei et al. (2023)](https://jmlr.org/beta/papers/v24/22-1132.html) (page 15), where they showed empirically that an upper bound for the number of samples is required (see the discussion section there).

Reviewer comment (Theory):
The role of Lemma 4.6 in the paper is unclear. A clearer explanation of how it relates to other parts of the paper and how it contributes to the overall arguments and conclusions is needed for better understanding.

Author response:
The role of this technical lemma is solely to explain the approximation in Equation (4.3). We will make this point clear in the updated version of the paper.

To provide some further explanation, recall that (4.3) is the key equation which shows that the network’s decision boundary at step 1 is approximately linear and thus it fails to generalize. The key step in (4.3) (with the “a.s.”) follows from Lemma 4.6. In the formal proof of the non-generalization result, we use concentration properties that have a similar form to Lemma 4.6 (see Lemma A.12 in Section A.5.2, which is used to prove Theorem 4.8).

Reviewer comment (Novelty):
The paper seems to lack a comparative discussion with some relevant works, specifically references [1,2].

Author response:
Thank you for bringing up the related papers. We note that these papers appeared on arXiv *after* the ICLR submission deadline. That said, we are happy to make a comparison with these papers:

[1] uses a simplified convolutional neural network on the XOR problem, and proves benign overfitting with a similar generalization upper bound as ours. Compared with [1], a key difference in our paper is that our setting also exhibits the delayed generalization/grokking phenomenon, which was not observed in [1] or any other benign overfitting papers before. Our results give a full characterization of this grokking phenomenon, offering the first rigorous proof of grokking in a neural network setting.

[2] uses a two-layer neural network on the (p, 2) parity task and gives a generalization bound for that problem. It’s worth noting that their XOR data distribution is different from ours: their noise terms are Boolean $\pm 1$, while our noise terms are Gaussian. As they mentioned in the Conclusion section, their results only hold for Boolean data and cannot be extended to Gaussian noise XOR data. [2] belongs to another line of work studying (p, 2) parity [[Ji & Telgarsky (2019)](https://arxiv.org/abs/1909.12292); [Wei et al. (2019)](https://arxiv.org/abs/1810.05369); [Barak et al. (2022)](https://arxiv.org/abs/2207.08799); [Telgarsky (2023)](https://openreview.net/forum?id=swEskiem99); [Suzuki et al. (2023)](https://openreview.net/forum?id=tj86aGVNb3&referrer=%5BAuthor%20Console%5D(%2Fgroup%3Fid%3DNeurIPS.cc%2F2023%2FConference%2FAuthors%23your-submissions)] and is not directly comparable with the Gaussian noise XOR problem we study.

Reviewer comment (Theory):
Could you provide more insights or justification regarding Assumption A1? Specifically, could you clarify why an increase in the number of training samples seems to negatively influence the model, contrary to the common understanding that more samples generally improve a model's generalization capability?

Author response:
We appreciate the question regarding the potential negative impact of increasing the number of samples on the model's performance, and agree that it is a bit counterintuitive. We wish to mention that an assumption on the upper bound of $n$ was needed in all previous benign overfitting results (since they all require high-dimensional data), e.g. [[Bartlett et al. (2019)](https://www.pnas.org/doi/10.1073/pnas.1907378117), [Frei et al. (2022)](https://arxiv.org/abs/2202.05928), [Cao et al. (2022)](https://arxiv.org/abs/2202.06526), [Frei et al. (2023)](https://jmlr.org/beta/papers/v24/22-1132.html), [Xu & Gu (2023)](https://proceedings.mlr.press/v206/xu23k.html), [Kou et al. (2023)](https://arxiv.org/abs/2303.04145)]. Also note that in practical settings it is possible that increasing the sample size could hurt generalization (see e.g. “deep double descent” [[Nakkiran et al. (2019)](https://arxiv.org/abs/1912.02292)]). 

Empirically, such an upper bound for sample size is indeed necessary. We have added an experiment for large sample size $n$ (see Figure 5 in Appendix A.8), which empirically shows that increasing $n$ could harm generalization, so it is not just a limitation of the analysis. 

Additionally, we also refer the reviewer to Figure 2 in [Frei et al. (2023)](https://jmlr.org/beta/papers/v24/22-1132.html) (page 15), where they showed empirically that an upper bound for the number of samples is required (see the discussion section there).

Reviewer comment (Theory):
Could you elaborate on the mechanism of overfitting after the neural network learns the directions of $u_1$ and $u_2$.

Author response:
A high-level intuition of this mechanism is: the training data points can be easily memorized by the neural network due to the near-orthogonality property ($\| x_i \|^2 \gg |\langle x_i, x_j \rangle|, \forall j \neq i$) and the high-dimensionality ($p\gg n$). The network has the capacity to learn the correct directions $\mu_1$ and $\mu_2$ and to memorize all the training data points.

In our paper, we prove that once the model memorizes a data point, it will never forget it. Specifically, for $a_j y_i > 0$, if at step $s$, the neuron $w_j$ positively correlates with the data point $x_i$, i.e. $\langle w_j^{(s)}, x_i\rangle > 0$, then it will stay that way afterward since $\langle w_j^{(t+1)} – w_j^{(t)}, x_i \rangle > 0, \forall t \ge s$. We refer the reviewer to (F1)-(F2) in Corollary A.5, which are the key properties used to prove the overfitting result. We note that previous benign overfitting results [[Cao et al. (2022)](https://arxiv.org/abs/2202.06526), [Xu & Gu (2023)](https://proceedings.mlr.press/v206/xu23k.html)] also utilized analogous properties to prove the overfitting result.

We think a more interesting question is how features can continue to be learned after the network has already achieved a perfect fit to data - this question seems to be the crux of why grokking has received so much attention. What we show is that even though the network has fitted the data, the small signal in the data continues to be picked up by gradient descent.  In high dimensions, the noise components are close to completely orthogonal, so over time the small increments in the signal learned by each step of gradient descent grow more than the fixed level of noise coming from the nearly-orthogonal noise variables.

Reviewer comment (Theory):
Could you clarify the role and significance of Lemma 4.6 in the context of the paper's objectives and findings?

Author response:
The role of this technical lemma is solely to explain the approximation in Equation (4.3). We will make this point clear in the updated version of the paper.

To provide some further explanation, recall that (4.3) is the key equation which shows that the network’s decision boundary at step 1 is approximately linear and thus it fails to generalize. The key step in (4.3) (with the “a.s.”) follows from Lemma 4.6. In the formal proof of the non-generalization result, we use concentration properties that have a similar form to Lemma 4.6 (see Lemma A.12 in Section A.5.2, which is used to prove Theorem 4.8).

Reviewer comment (Experiments):
Could you make a comparison of your results and techniques with [1] and [2]?

Author response:
Thank you for bringing up the related papers. We note that these papers appeared on arXiv *after* the ICLR submission deadline. That said, we are happy to make a comparison with these papers:

[1] uses a simplified convolutional neural network on the XOR problem, and proves benign overfitting with a similar generalization upper bound as ours. Compared with [1], a key difference in our paper is that our setting also exhibits the delayed generalization/grokking phenomenon, which was not observed in [1] or any other benign overfitting papers before. Our results give a full characterization of this grokking phenomenon, offering the first rigorous proof of grokking in a neural network setting.

[2] uses a two-layer neural network on the (p, 2) parity task and gives a generalization bound for that problem. It’s worth noting that their XOR data distribution is different from ours: their noise terms are Boolean $\pm 1$, while our noise terms are Gaussian. As they mentioned in the Conclusion section, their results only hold for Boolean data and cannot be extended to Gaussian noise XOR data. [2] belongs to another line of work studying (p, 2) parity [[Ji & Telgarsky (2019)](https://arxiv.org/abs/1909.12292); [Wei et al. (2019)](https://arxiv.org/abs/1810.05369); [Barak et al. (2022)](https://arxiv.org/abs/2207.08799); [Telgarsky (2023)](https://openreview.net/forum?id=swEskiem99); [Suzuki et al. (2023)](https://openreview.net/forum?id=tj86aGVNb3&referrer=%5BAuthor%20Console%5D(%2Fgroup%3Fid%3DNeurIPS.cc%2F2023%2FConference%2FAuthors%23your-submissions)] and is not directly comparable with the Gaussian noise XOR problem we study.

Reviewer comment (Experiments):
Have you considered variations in the network architecture, such as not fixing the second layer, and if so, how do these variations influence the results?

Author response:
We appreciate the reviewer’s question. We have added an experiment that uses gradient descent to train both layers with the same initialization as our existing experiments. The results (Figure 7 in Appendix A.8) show that the learning dynamics is similar to that with the second layer fixed and that the benign overfitting and grokking phenomena are still observed.

Theoretically, our developed framework cannot be directly extended to gradient descent on both layers, as some concentration results for random variables with complicated dependency structure need to be obtained. The techniques used in [2] also cannot be directly extended to our setting. As [2] mentioned in their Conclusion Section, for Boolean data, it is easier to compute $\nabla_a L$ and $\nabla_W L$, and the dynamics of the gradients do not depend on interactions within neurons. But for Gaussian noise XOR, the dynamics can be quite different and the neurons corresponding to the signal may go in the wrong directions. It is definitely worth exploring the learning dynamics of neural networks trained by gradient descent on both layers for future research.

Reviewer comment (Experiments):
In the model settings used in the paper, there doesn’t appear to be an upper bound on the network width. If an extremely wide network setting were used, would the findings align with the "lazy training" regime? Could you discuss how the results might be influenced by varying the width of the network to such extremes?

Author response:
Yes, our result holds for any width $m \geq C n^{0.01}$ and does not require an upper bound on $m$. The same results hold even if $m$ is extremely large.

We emphasize that the setting we consider is different from the “lazy training” / “neural tangent kernel (NTK)” regime, even if the width goes to infinity. This is because we consider a smaller initialization scale than that in the NTK regime. One way to see this difference is that our weight matrix W goes far from its initialization after one step of gradient descent (since step size $\alpha$ is much larger than the initialization scale), while in the NTK regime the weight matrix will always stay close to its initialization. Therefore, even with infinite width, the setting we consider is not in the lazy training/NTK regime.

Reviewer comment (Theory):
In my understanding, it is crucial to take a small initialization scale compared to the step size $\alpha$ to obtain the perfect overfitting at the first gradient step. I think this is acceptable as theory, but it should be better to justify such a small initialization. I also want to know what happens if the initialization scale is much larger than used in Figure 3.

Author response:
**Justification of the small initialization:**

The small initialization is indeed crucial as it makes the analysis of the training dynamics easier. Such an assumption has been made in a number of related theory papers, such as [[Frei et al. (2022)](https://arxiv.org/abs/2202.05928), [Frei et al. (2023)](https://jmlr.org/beta/papers/v24/22-1132.html), [Xu & Gu (2023)](https://proceedings.mlr.press/v206/xu23k.html)].

We share the reviewer’s desire for understanding what happens under a larger initialization scale and have done some experiments in this regime. Specifically, in our paper, Figure 1 (right) gives the train and test accuracy under large initialization (relative to step size). The large initialization version of Figure 3 is given in Figure 4 in Appendix A.7. We observe similar benign overfitting and grokking behaviors for large initialization, though the overfitting and generalization require more steps than the case of small initialization.

We have also run additional experiments with a fixed learning rate and varying initialization scales; see Figure 6 in Appendix A.8. We again observe similar qualitative phenomena.

Reviewer comment (Novelty):
Compared to [Ji & Telgarsky (2019)](https://arxiv.org/abs/1909.12292); [Wei et al. (2019)](https://arxiv.org/abs/1810.05369); [Barak et al. (2022)](https://arxiv.org/abs/2207.08799); [Telgarsky (2023)](https://openreview.net/forum?id=swEskiem99); [Suzuki et al. (2023)](https://openreview.net/forum?id=tj86aGVNb3&referrer=%5BAuthor%20Console%5D(%2Fgroup%3Fid%3DNeurIPS.cc%2F2023%2FConference%2FAuthors%23your-submissions), where $x$ is a $d$-dimensional (essentially) rotationally invariant input and $y=\rm{sgn}(x_1x_2)$, this paper considers large signal $\|\mu\|$ as an input.

Author response:
**Large signal-to-noise (SNR) ratio:**

We agree that our SNR is larger than that in the papers that the reviewer mentioned. However, those papers consider a different problem and are not directly comparable to our work. In particular, the XOR problem they studied is the discrete (p, 2) parity task, while we study the Gaussian noise XOR problem. Instead of a large SNR, those papers all require a large sample size $n$:  [Ji & Telgarsky (2019)](https://arxiv.org/abs/1909.12292), [Wei et al. (2019)](https://arxiv.org/abs/1810.05369), [Barak et al. (2022)](https://arxiv.org/abs/2207.08799), and [Telgarsky (2023)](https://openreview.net/forum?id=swEskiem99) all need $n = \Omega(p^2)$, and [Suzuki et al. (2023)](https://openreview.net/forum?id=tj86aGVNb3&referrer=%5BAuthor%20Console%5D(%2Fgroup%3Fid%3DNeurIPS.cc%2F2023%2FConference%2FAuthors%23your-submissions) need $n = \Omega(p)$. On the other hand, we study the high-dimensional regime $n \ll p$, which is in line with all existing benign overfitting results (e.g. [[Bartlett et al. (2019)](https://www.pnas.org/doi/10.1073/pnas.1907378117), [Frei et al. (2022)](https://arxiv.org/abs/2202.05928), [Cao et al. (2022)](https://arxiv.org/abs/2202.06526), [Frei et al. (2023)](https://jmlr.org/beta/papers/v24/22-1132.html), [Xu & Gu (2023)](https://proceedings.mlr.press/v206/xu23k.html), [Kou et al. (2023)](https://arxiv.org/abs/2303.04145)]) and intuitively explains why a relatively large SNR is needed.

We also note that there is a concurrent work studying (p, 2) parity [[Glasgow (2023)](https://arxiv.org/abs/2309.15111)] which discusses the difference between the (p, 2) parity problem and the Gaussian noise XOR problem in their Conclusion and Discussion section and why their results cannot be extended to Gaussian noise XOR.

Reviewer comment (Experiments):
Is it possible to see an additional experiment when the initialization scale is not so small?

Author response:
**Justification of the small initialization:**

The small initialization is indeed crucial as it makes the analysis of the training dynamics easier. Such an assumption has been made in a number of related theory papers, such as [[Frei et al. (2022)](https://arxiv.org/abs/2202.05928), [Frei et al. (2023)](https://jmlr.org/beta/papers/v24/22-1132.html), [Xu & Gu (2023)](https://proceedings.mlr.press/v206/xu23k.html)].

We share the reviewer’s desire for understanding what happens under a larger initialization scale and have done some experiments in this regime. Specifically, in our paper, Figure 1 (right) gives the train and test accuracy under large initialization (relative to step size). The large initialization version of Figure 3 is given in Figure 4 in Appendix A.7. We observe similar benign overfitting and grokking behaviors for large initialization, though the overfitting and generalization require more steps than the case of small initialization.

We have also run additional experiments with a fixed learning rate and varying initialization scales; see Figure 6 in Appendix A.8. We again observe similar qualitative phenomena.