We are grateful to the reviewer for their time and thoughtful feedback. We sincerely appreciate the effort put into evaluating our work and offering constructive suggestions. Below, we respond to each of the points raised.

Theoretical Claims

We've rewritten the proofs of the two main theorems (2.1 and 3.1) more simply, this time with a proof sketch, which we can summarize as follows. The idea is to first show that, given a matrix $\tilde{\mathbf{X}} \in \mathbb{R}^{N \times n}$ and a vector $\mathbf{b}^* \in \mathbb{R}^{n}$ of sparsity $s$ ($\|\mathbf{b}^*\|_0 \le s$), we can write, for any vector $\mathbf{b} \in \mathbb{R}^{n}$, $\|\mathbf{b} - \mathbf{b}^*\|_2 \le C_1 \| \tilde{\mathbf{X}} \left(\mathbf{b} - \mathbf{b}^*\right)\|_2 + C_2 | \|\mathbf{b}\|_1 - \|\mathbf{b}^*\|_1|$ where $C_1$ and $C_2$ are constants that depend only on $s$ and $\tilde{\mathbf{X}}$ (notably on the restricted isometry constant of $\tilde{\mathbf{X}}$). With this decomposition in mind, we then show that :

• there is a memorization phase, after which the term $\| \tilde{\mathbf{X}} \left(\mathbf{b}^{(t)} - \mathbf{b}^*\right)\|_2$ vanish (or becomes proportional to the noise if there is any), since we have $\tilde{\mathbf{X}} \mathbf{b}^{(t)} \approx \mathbf{y}^* = \tilde{\mathbf{X}} \mathbf{b}^* + \xi$.
• and a second phase, during which $\|\mathbf{b}^{(t)}\|_1$ converges to $\|\mathbf{b}^*\|_1$: we show that this phase takes a time inversely proportional to $\alpha$ (the learning step) and $\beta_1$ (the regularization strength).

We arrive at the final result by combining these two steps. The same type of decomposition applies to matrix factorization, with a slight difficulty because not only singular values but also singular vectors are taken into account.

Weaknesses: Rewriting

The rewriting of the proofs allows us to considerably reduce the number of pages in the appendix and to move many results from the appendix to the main text.

Validation plot

In our case, for an iterate $\mathbf{b}$, the training error is $\|\tilde{\mathbf{X}} \mathbf{b} - \mathbf{y}^*\|_2^2$, and the test (or validation) error is $\|\mathbf{b} - \mathbf{b}^*\|_2^2$. Normally, the generalization error should be $\mathcal{E}(\mathbf{b}) = \mathbb{E}\left(\mathbf{x}^\top \mathbf{b} - \mathbf{y}^*(\mathbf{x})\right)^2>$ (expectation with respect to $\mathbf{x}$ and $\xi$), where $\mathbf{y}^*(\mathbf{x}) = \mathbf{x}^\top \mathbf{b}^* + \xi$. Assuming $\mathbb{E}\xi = 0$, and using $\Sigma = \mathbb{E}[\mathbf{x} \mathbf{x}^\top]$, we get $\mathcal{E}(\mathbf{b}) = \left(\mathbf{b} - \mathbf{b}^* \right)^\top \Sigma \left(\mathbf{b} - \mathbf{b}^* \right) + \mathbb{E}\xi^2$.

As we illustrate in the appendix, random matrices, notably Gaussian and Bernoulli with independent entries, allow for the lowest restricted isometry constant and, thus, better recovery of sparse vectors. Under this iid assumption, we have $\Sigma = \sigma^2\mathbf{I}_n$ for a certain $\sigma > 0$, which implies $\mathcal{E}(\mathbf{b}) = \sigma^2 \|\mathbf{b} - \mathbf{b}^*\|_2^2 + \mathbb{E}\xi^2$. So $\|\mathbf{b} - \mathbf{b}^*\|_2^2$ captures well the notion of test (or validation) error commonly used in the context of grokking, while $\mathbb{E}\xi^2$ capture the irreducible part of that error (we exclude it in the figure for simplicity).

Could the authors simplify Figure 1 and perhaps Figure 3 to relate these terms to the terms plotted in the figure?

Thanks for the suggestion, we will simplify the figures as best we can.

We thank the reviewer for their thoughtful feedback and the time spent reviewing our work. Our responses to the comments are provided below.

Weaknesses

1. We have rewritten both sections to emphasize the fundamental difference between the two main theorems (2.1 about sparse recovery and 3.1 about matrix factorization) with proof sketches. We first show (using standard results from compressed sensing literature) that, given a matrix $\tilde{\mathbf{X}} \in \mathbb{R}^{N \times n}$ and a vector $\mathbf{b}^* \in \mathbb{R}^{n}$ of sparsity $s$ ($\|\mathbf{b}^*\|_0 \le s$), we can write, for any vector $\mathbf{b} \in \mathbb{R}^{n}$, $\|\mathbf{b} - \mathbf{b}^*\|_2 \le C_1 \| \tilde{\mathbf{X}} \left(\mathbf{b} - \mathbf{b}^*\right)\|_2 + C_2 | \|\mathbf{b}\|_1 - \|\mathbf{b}^*\|_1|$ where $C_1$ and $C_2$ are constants that depend only on $s$ and $\tilde{\mathbf{X}}$. After a short amortization phase, the term $\| \tilde{\mathbf{X}} \left(\mathbf{b}^{(t)} - \mathbf{b}^*\right)\|_2$ vanish (or becomes proportional to the noise if there is any), since we have $\tilde{\mathbf{X}} \mathbf{b}^{(t)} \approx \mathbf{y}^* = \tilde{\mathbf{X}} \mathbf{b}^* + \xi$. Then comes a second phase, during which $\|\mathbf{b}^{(t)}\|_1$ converges to $\|\mathbf{b}^*\|_1$: our contribution is to show that this phase takes a time inversely proportional to $\alpha$ (the learning step) and $\beta_1$ (the regularization strength).

The same type of decomposition applies to matrix factorization, with a slight difficulty because not only singular values but also singular vectors are taken into account. In fact, although the two results (sparse recovery and matrix factorization) are similar, the distinction lies at several levels:

• in sparse recovery problems, we take into account only the iterate $\mathbf{b}^{(t)}$ (especially its $\ell_1$ norm), whereas, in matrix factorization, we take into account not only the singular value matrix $\Sigma^{(t)}$ (especially its $\ell_1$ norm) of the iterate $\mathbf{A}^{(t)} = \mathbf{U}^{(t)} \Sigma^{(t)} \mathbf{V}^{(t)\top}$, but also its singular vectors $\mathbf{U}^{(t)}$ (left) and $\mathbf{V}^{(t)}$ (right), which makes proofs in the second case slightly more difficult.
• if we take a matrix factorization problem and just optimize it with $\ell_1$, there's no grokking unless the matrix is extremely sparse so that the notion of sparsity prevails over the notion of rank, which shows the point of studying $\ell_*$ (nuclear norm regularization) separately.
• models (linear or not) trained with $\ell_2$, $\ell_1$ and $\ell_*$ have very different properties: the aim of our work is to show that in all these cases, if we want to have a model that generalizes, and which has the desired property (weights of low $\ell_2$ norm, sparse, or of low rank), then it's better to use a property-specific, weak regularization, and to train the model long enough.

This extends Luy et al.'s work, which focuses solely on $\ell_2$. Indeed, there will be no grokking by using only $\ell_2$ in the problems we study theoretically in the paper. This also shows that the mechanism behind grokking goes beyond the simple $\ell_2$ norm (we show that with other regularization, we can change the delay between memorization and generalization).

2. Our results cannot be obtained directly from Luy et al.'s work. They show that large-scale initialization and non-zero weight decay lead to grokking provided the model is homogenous. However, as we show, using only $\ell_2$ (with large-scale initialization or not) in the problems we study theoretically in the paper, there will be no grokking. With large-scale initialization and $\ell_2$ only, there is an abrupt transition in the generalization error during training, driven by changes in the $\ell_2$-norm of the model parameters. This transition, however, does not result in convergence to an optimal solution. We call this phenomenon "grokking without understanding," and we attribute it to the fact that the assumptions underlying the theoretical predictions of Luy et al. are violated in our setting.

Questions: Difference between the core idea of the paper and that of Lyu et al. (2023)

Luy et al.'s focus on $\ell_2$ and large initialization. But in our case, there will be no grokking by using only $\ell_2$ in the problems we study theoretically. So, we extend their work to other regularizers, and there is no need for large-scale initialization in our setting. In addition, we show how grokking can be necessary in specific contexts. Suppose we want a model with low-rank weights and opt to achieve that through $\ell_*$ regularization. In that case, our work suggests using the lowest possible regularization coefficient (the one the computational resources allow) and training the model as long as possible, far beyond the overfitting point (grokking). We also show that grokking can be extremely amplified/reduced by selecting the data appropriately.

We appreciate the reviewer’s time and insightful comments. Please find our responses to the points below.

Claims And Evidence

• We worked in a setting where the notion of sparsity (resp. low rank) is well defined, $\mathbb{R}^n$ (resp. $\mathbb{R}^{n_1 \times n_2}$). This is the number of non-zero elements of a vector (resp. the number of non-zero singular values of a matrix). If we decide to work, say, in $\mathbb{Z}/p\mathbb{Z}$ for some prime number $p$ (like in the first problem grokking was first reported on, modular arithmetic), then these definitions are not trivially valid, and the theory we developed no longer applies (indeed, the theory of sparse recovery and matrix factorization over finite fields is very complex and uses completely different tools from those used in $\mathbb{C}$). This is why we can not answer the question, “Is the modular arithmetic problem to recover the removed half of the binary relationships sparse?” using our current theory and leave it as a future direction (see section F of the Appendix). Although for tasks such as the sparse parity task (on which grokking is observed [1, 2]) one can construct sparse MLP that fits the data [1], we are not aware of any work that has done this for modular arithmetic. [3] constructs a two-layer MLP with square activations that does modular addition, but the model weights constructed are dense.

• The most crucial point of our work is to show that not only does the grokking mechanism go beyond the $\ell_2$ norm, but also why it is necessary to do so. Indeed, we show that the time to generalize is inversely proportional to the regularization strength used and that the generalization error after convergence decreases with it. Moreover, we show that grokking can also be highly accelerated by simply choosing appropriate data samples. To the best of our knowledge, no previous work on grokking has done all these.

  Although our theory relies on simple frameworks that allow us to extract the laws justifying our results, we also show this empirically for different types of models and datasets, including the basic setup on which grokking was first observed, the algorithmic dataset. We acknowledge that developing a theory in a simplified framework and then testing it experimentally in a more complicated setup has its flaws. The point of our paper is not to say that the mechanisms behind linear models extend to non-linear models but to illustrate that grokking is nevertheless observed if we use the insights obtained theoretically on linear models. We'll try to make this point clearer in the final version of the paper.

• We acknowledge that grokking can be caused by a number of factors. This is the point of our work, even if we focus on regularization. Other work (Grokfast, etc) focuses on other factors to amplify/reduce the Grokking delay, and we will discuss this in detail in the related works section.

Methods And Evaluation Criteria, Theoretical Claims

The point of our work is not that our theoretical results on linear models trivially extend to non-linear models. We simply show that the mechanism behind grokking (and regularization) goes beyond the $\ell_2$ norm: we show it theoretically and empirically in specific settings and just empirically in others (and leave the theory in such settings as future work).

Essential References Not Discussed

Thanks for the references, we will discuss this in the related work section. These works differ from ours in that the first focuses on double-descent (which is not the point of our work) and the other on other strategies to accelerate grokking (we focus on different types of regularization and data selection).

Questions For Authors

We excluded Transformer results on modular arithmetic as the paper was already too long. On Transformer, $\ell_1$ and $\ell_*$ still affect the delay between memorization and generalization.

References

[1] Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit (https://arxiv.org/abs/2207.08799)

[2] A Tale of Two Circuits: Grokking as Competition of Sparse and Dense Subnetworks (https://arxiv.org/abs/2303.11873)

[3] Grokking modular arithmetic (https://arxiv.org/abs/2301.02679)

We thank the reviewer for their time and insightful comments. We appreciate the effort taken to evaluate our work and provide constructive feedback. Below, we address each of the points raised.

Weaknesses

We understand that the paper is poorly organized. That is why, after submission, we completely rewrote it as follows:

• All the proofs have been simplified, and we've included the proof sketch for each main theorem. For example, here's the main idea behind Theorem 2.1. We first show (using standard results from compressed sensing literature) that, given a matrix $\tilde{\mathbf{X}} \in \mathbb{R}^{N \times n}$ and a vector $\mathbf{b}^* \in \mathbb{R}^{n}$ of sparsity $s$ ($\|\mathbf{b}^*\|_0 \le s$), we can write, for any vector $\mathbf{b} \in \mathbb{R}^{n}$, $\|\mathbf{b} - \mathbf{b}^*\|_2 \le C_1 \| \tilde{\mathbf{X}} \left(\mathbf{b} - \mathbf{b}^*\right)\|_2 + C_2 | \|\mathbf{b}\|_1 - \|\mathbf{b}^*\|_1|$ where $C_1$ and $C_2$ are constants that depend only on $s$ and $\tilde{\mathbf{X}}$. After a short amortization phase ($t \le t_1$), the term $\| \tilde{\mathbf{X}} \left(\mathbf{b}^{(t)} - \mathbf{b}^*\right)\|_2$ vanish (or becomes proportional to the noise if there is any), since we have $\tilde{\mathbf{X}} \mathbf{b}^{(t)} \approx \mathbf{y}^* = \tilde{\mathbf{X}} \mathbf{b}^* + \xi$. Then comes a second phase, during which $\|\mathbf{b}^{(t)}\|_1$ converges to $\|\mathbf{b}^*\|_1$ at $t_2$: our contribution is to show that $t_2 - t_1$ is inversely proportional to $\alpha$ (the learning step) and $\beta_1$ (the regularization strength).
• We made the body of the paper self-contained insofar as any contribution mentioned in the abstract and introduction is proven in the body of the paper theoretically and/or empirically (grokking time, effect of data selection on grokking, limits of $\ell_2$, and grokking without understanding, ...). So, the results in the appendix this time are just extensions of those in the main text.
• With the rewrite, the appendix has also been reduced from $\mathcal{O}(70)$ to less than $50$ pages.
• We have left the part on matrix factorization in the main text since the mechanics behind the generalization phase, in this case, differs from sparse recovery. In sparse recovery problems, we take into account only the iterate $\mathbf{b}^{(t)}$ (especially its $\ell_1$ norm), whereas, in matrix factorization, we take into account not only the singular value matrix $\Sigma^{(t)}$ (especially its $\ell_1$ norm) of the iterate $\mathbf{A}^{(t)}$, but also its singular vectors $\mathbf{U}^{(t)}$ (left) and $\mathbf{V}^{(t)}$ (right), which makes the dynamic in the second case completely different from the first case. Also, in a matrix factorization problem, replacing $\ell_*$ (the nuclear norm regularization) by $\ell_{*/2}$ does not induce grokking (at least in the shallow case used in the theory).

Thanks for the suggestion about figures and parameters of interest. We will define a parameter $\tilde{\alpha} = \alpha \beta$, which is the main parameter in the theory ($\alpha$ is the learning rate and $\beta$ is the regularization).

What is the overlap between $\mathbf{b}^{(t_1)}$ and $\mathbf{b}^{(t_2)}$? What changes more, the norm or the direction? Once memorization occurs ($\mathbf{X}\mathbf{b}^{(t_1)}=\mathbf{X}\mathbf{b}^*+ \xi$), the direction of $\mathbf{b}^{(t)}$ that explains the data does not need to change drastically. The main change from $t_1$ to $t_2$ is the magnitude of $\mathbf{b}^{(t)}$ in $\ell_1$—i.e., it transitions from something closer to the least‐square solution’s norm $\|\hat{\mathbf{b}}\|_1$ down to the teacher’s norm $\|\mathbf{b}^*\|_1$ (while still maintaining some form of memorization). So, $\mathbf{b}^{(t_1)}$ and $\mathbf{b}^{(t_2)}$ are closely aligned in direction; the main difference is that the $\ell_1$ norm shrinks from $\|\hat{\mathbf{b}}\|_1$ to $\|\mathbf{b}^*\|_1$. Hence, the norm changes more than the direction.

Can you comment on the fact that the test loss is always monotonic in your settings?

The monotonic decrease in test loss is due to how we choose the hyperparameters, mainly the learning rate and the regularization strengths. In the case of sparse recovery, for example, the test loss is controlled mainly by $\|**b**^{(t)}\|_1$, even more so after memorization, which takes only a few steps. After this memorization, $\|**b**^{(t)}\|_1$ converge very slowly towards $\|**b**^*\|_1$. When $\alpha$ and/or $\beta_1$ are larger than the values used in the paper, there are indeed a lot of oscillations in test loss. The same reasoning applies to experiments outside the linear scope (e.g., figure 5). We chose the learning rate and regularization strength to illustrate grokking. Larger values show oscillations (e.g., the Slingshot Mechanism [1]).

[1] The Slingshot Mechanism: An Empirical Study of Adaptive Optimizers and the Grokking Phenomenon (https://arxiv.org/abs/2206.04817)