Benign Overfitting and Grokking in ReLU Networks for XOR Cluster Data

We thank the reviewer for their review and address their comments and questions below. We hope that the reviewer will consider raising their score in light of our response.

1: My main concern about this paper lies in its assumptions… Do the authors have any ideas for improvements in the high-dimensional setting?

Note that all existing benign overfitting results require input dimension much larger than the number of samples ($p>\mathrm{poly}(n)$), e.g. [Bartlett et al. (2019), Frei et al. (2022), Cao et al. (2022), Frei et al. (2023), Xu & Gu (2023), Kou et al. (2023)]. Compared with these papers, our non-linearly separable setting is significantly harder. One major technical challenge is to give concentration bounds for summation of dependent random variables (see Equations (A.17)-(A.22) and Lemma A.17 in Section A.6). There are few techniques that can analyze the concentration of such summation, which is necessary for our trajectory analysis. We note that such difficulty is due to the non-linearly separable nature of our data distribution, which was not present in previous papers on linearly separable data.

Furthermore, we are not sure what the reviewer means regarding potential ideas for “improvements in the high-dimensional setting.” We do not have a complete understanding of whether grokking or benign overfitting happen in lower-dimensional settings, and we believe a full characterization of whether these two phenomena occur as a function of the dimension, number of samples, step-size, initialization variance, and signal-to-noise ratio (determined by $\|\mu_i\|$) is an extremely challenging problem. Indeed, prior work by Frei et al. (2023) (Figure 2, page 15) showed that in the high-SNR, low-dimensional setting for two-layer ReLU nets trained on the XOR cluster distribution we consider, overfitting is not benign and grokking does not occur.

We believe it is significant that we have established a proof of both benign overfitting and grokking in ReLU networks for nonlinear data, and that our work helps elucidate how high-dimensionality and low-SNR can play a crucial role in both of these phenomena. We believe these contributions are of substantial interest to the ICLR community. We hope the reviewer will consider raising their score.

2: What behavior does the test error exhibit when the training time is between 1 and $n^{0.01}$?

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.

We thank the reviewer the detailed review and for appreciating the contributions of our paper. It is encouraging to see that the reviewer thinks our paper is “important” and “structured and clear” and offers “a new theoretical examination of benign overfitting and grokking.”

We are happy to address the reviewer’s questions below, and hope that the reviewer will consider raising the score in light of our response.

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?

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), Frei et al. (2022), Cao et al. (2022), Frei et al. (2023), Xu & Gu (2023), Kou et al. (2023)]. 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)]).

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) (page 15), where they showed empirically that an upper bound for the number of samples is required (see the discussion section there).

Could you elaborate on the mechanism of overfitting after the neural network learns the directions of u1 and u2?

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), Xu & Gu (2023)] 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.

Could you clarify the role and significance of Lemma 4.6 in the context of the paper's objectives and findings?

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

We thank the reviewer for the detailed review and for appreciating the contributions of our paper. It is encouraging to see that the reviewer thinks our paper is well-written and regards our paper as providing “notable results on benign overfitting” and “useful theoretical understanding on grokking.”

We are happy to address the reviewer’s questions below.

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), Frei et al. (2023), Xu & Gu (2023)].

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.

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), Wei et al. (2019), Barak et al. (2022), and Telgarsky (2023) 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), Frei et al. (2022), Cao et al. (2022), Frei et al. (2023), Xu & Gu (2023), Kou et al. (2023)]) 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)] 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.

---

Dear authors,

I appreciate authors' efforts on the extensive clarification.

Best,