Paper: To Grok or not to Grok: Disentangling Generalization and Memorization on Corrupted Algorithmic Datasets (ICLR-2024)

Reviewer comment (Evaluation):
The title and motivation put a lot of emphasis on grokking, yet there are no in-depth analysis of the grokking behaviors presented in the paper. I think the paper could be improved if it could show that the analytical tools used in this paper could bring new insights to our understanding of the grokking behaviors.

Author response:
## New insights on grokking

Grokking on modular arithmetic datasets has been observed in multitudes of setups with different architectures, optimizers, regularization schemes (or lack thereof), loss functions, etc. These setups, despite having very different dynamics, feature similar grokking behaviours. Consequently, we focus our attention on the common denominator, namely the dataset. We study how changing/corrupting the data affects these behaviors and also present new behaviors emerging out of this examination. As we mention in the Novelty section, on modular arithmetic datasets, label corruption doesn't just introduce conflicting examples -- it leads to a new modular arithmetic task. Thus, robust generalization on the target task with significant label corruption points to a bias towards a certain class of modular tasks, possibly ones that can be learnt with an algorithm similar to Hypotheses 2.1, 2.2. We believe that these insights are valuable for future investigations on grokking.

Another important aspect of grokking is the tussle between memorization and generalization. With that in mind, we demonstrate that one can tune the generalization/memorization performance by tuning label corruption and regularization. We also characterize these behaviors using phase diagrams and interpret them using detailed IPR results. Moreover, our analysis also generalizes to various network architectures. All in all, we think that our experiments and tools will significantly advance the understanding of the subject.

Reviewer comment (Evaluation):
In the case of co-existence, are there any consistent patterns on which neurons would become memorization and which would become generalization? Are they mostly in the lower layer or the upper layer? Do they have different learning dynamics?

Author response:
## Patterns of generalizing/memorizing neurons in the *coexistence* phase

We can predict the occurrence of generalizing neurons before the grokking transition occurs by keeping track of per-neuron IPR during training. Furthermore, one can predict these behaviors during early training times by tracking the per-neuron IPR in the gradient updates -- the increase in IPR of gradients precedes that in weights.

However, predicting such behaviors *at initialization* is difficult using IPR. As such, we do not find a consistent correlation between the Fourier spectrum of the weight vectors at initialization and their periodicity after training. We view the occurrence of periodic features and the partition of neurons into generalizing and memorizing neurons as emergent collective phenomena -- a comprehensive theory of grokking dynamics may give us further insight into the question. 

> Are they mostly in the lower layer or the upper layer?

We are unsure about what the reviewer means by this question. We urge the reviewer to clarify their question so that we can give an accurate response.

To avoid any possible confusion, we clarify the point about periodic neurons. In our 2-layer setup, *all* the weight matrices have periodic features. Specifically, using notation from the paper, the $k^{th}$ row of $U$, $k^{th}$ row of $V$ and $k^{th}$ column of $W$ matrices connect to a given $k^{th}$ (hidden) neuron. All three of these row/column vectors are found to be periodic with identical frequency. The per-(hidden)-neuron IPR histograms presented in the paper average the IPR over these three vectors. We refer the reviewer to Equation 4 and the surrounding text for further clarification. It can also be seen from the analytical solution that this has to be the case.

Reviewer comment (Experiments):
The paper uses p=97. Why such a value is chosen? Does the behavior critically depend on the actual value of p (odd, even, prime, very small, very large, etc.)?

Author response:
## Choice of p

Neural networks can generalize on modular addition tasks for all natural numbers $p$, given a large enough width. Furthermore, we find this to be true even on more general modular arithmetic tasks (such as $(mn)\\%p$ and $(m^2+n^2)\\%p$, etc.). 
However, formally, general modular arithmetic tasks are only well-defined over finite fields -- which require the modular base to be a prime number. We find that network training circumvents this subtlety in many cases.

The size of the dataset also depends on this choice -- the number of total examples is $p^2$. Hence, a very small $p$ (such as $2,3,5$) would prevent learning. On the other hand, a very large $p$, while learnable, would demand excessive computing resources. The prime $p=97$ is the choice from the original paper on Grokking by Power et al. (2022). This choice strikes a reasonable balance between the aforementioned factors, but is in no way unique.

Reviewer comment (Experiments):
In Fig.6, the test accuracy remains almost perfect even after 60% of the neurons being pruned. How does the ratio (60%) depend on various aspects of the underlying task? Does it remain constant when the network size changes?

Author response:
## How many neurons can be pruned without losing performance?

We would first like to clarify that perfect generalization in the pruning experiment (Figure 5) remains **even beyond** pruning $60\\%$ of the low IPR neurons. Figure 6(left) showed only up to $70\\%$ pruning to explain the (sole) effect of pruning out low IPR neurons. We have added the extended version of this pruning experiment in Appendix I.1 (Figure 17(a)), where one can see that it is possible to prune $\\sim 88\\%$ of the neurons without losing test performance. We find that for $N=500$, the number of remaining neurons after maximal pruning is $\\sim 60$ ($N$ is the network width).

According to the analytical solution (Equation 2), $\\lceil p/2 \\rceil$ neurons with different frequencies are required for the network to generalize correctly. Empirically, in real networks, the periodicity in neurons is approximate and the neuronal frequencies are not always unique. Consequently, real networks require more than $\\lceil p/2 \\rceil$ neurons for proper generalization. This is reflected in the aforementioned number $60$ being higher than $\\lceil 97/2 \\rceil=49$. 

We also empirically find that wider networks are more prone to such effects, and hence require more neurons for generalization. We have added this experiment in Appendix Q (Figure 32). As such, we see a monotonic, albeit sub-linear, increase in the number of neurons required with increasing network width. Note that the ratio of prune-able neurons *increases* substantially due to the sub-linearity with width. For N=10000, we find that $\\sim 99\\%$ of the neurons can be pruned while retaining generalization. In this case, the network requires $\sim 100$ neurons to maintain generalization.

Moreover, the number of prune-able neurons also weakly depends on the phase in which the network lies (coexistence/partial-inversion/full-inversion.) Intuitively, due to the scarcity of periodic neurons in the coexistence phase, the network seems more economical, compared to the inversion phases. These comparisons are explained in Figure 17 of Appendix I.1 of the revised manuscript.

Reviewer comment (Experiments):
Is there any scenario where we believe Partial Inversion or Full Inversion may be observed in practice?

Author response:
- Yes, we do observe partial as well as full inversion on Modular Arithmetic tasks learned with the Transformer architecture as well as deeper MLP networks (Section 2.2  and Appendix G). Even in these general settings, we can relate periodic weights and IPR to generalization.
We also observe inversion in corrupted image tasks such as MNIST (Appendix H). For further details please refer to the Global Response.

Reviewer comment (Experiments):
Recently, there have been a series of papers on grokking, most of which rely on empirical analysis of training dynamics or final networks. Although such findings could be thought-provoking, they often come with the risk that findings/questions may apply to the particular setting considered in the work. I'm not sure whether the analysis and findings in this work can be extrapolated to other setups, e.g., different architectures or datasets. This is why I'm closer to a rejection but would be happy to discuss it with the authors and other reviewers.

Author response:
For clarity, we re-emphasize the two different avenues of generalizing our results; namely, (i) general architectures, and (ii) general datasets. 

On the general architecture side, as we mentioned in the General Response, we have shown that our results extend to Transformers and deeper MLPs (new Section 2.2 and Appendix G in the revised manuscript).

On the general dataset front, we agree that analyses using IPR cannot be carried over verbatim to general datasets like ImageNet. However, we do see that image datasets like MNIST have similar phases and inversion phenomena as modular arithmetic datasets (Appendix H). Thus, they feature the same phenomena, albeit with a different (unknown) algorithmic bias. The challenge for future research is to identify this bias, which would allow a way to find these sub-networks for general datasets. To that end, we believe that our methodology, i.e., introducing incremental label corruption to the dataset and then finding the invariant features in various phases, could be a useful method for understanding memorization and generalization in a broader setting.

Reviewer comment (Novelty):
The paper analyzes the learning under label corruptions in the same vein as previous works. As shown in Nanda et al. (2023), the algorithm implemented by the algorithm after the grokking phase, roughly speaking, boils down to a handful of periodic functions, which is further confirmed by Gromov (2023) in terms of an analytical solution. Based on this observation, the authors propose to use IPR as a measure of "how periodic a neural representation is", which serves as a proxy for whether a neuron generalizes or memorizes. Then, the effect of pruning individual neurons is investigated in the same spirit as Nanda et al. (2023). All in all, on the procedural side, the only novelties seem to be the use of IPR and label corruption.

Author response:
We first clarify that the word "simpler" means completely different things in the paper by Nanda et al. compared to ours. 

- In their work (Appendix E.1), they speculated a vague notion of simplicity: "memorization complexity scales with the size of the training set, whereas generalization complexity is constant. The two must cross at some point!" 
They conjectured that there are only two solutions, namely, a memorizing solution (complex) and a generalizing solution (simpler). Both correspond to the *same* task, and regularization would bias the network towards the simpler one. 

- Whereas the point we are trying to convey is that, given some training data, there are an exponentially large number of solutions and they all correspond to *some* modular arithmetic polynomial. Then, simplicity refers to the polynomial that the network learns upon training.

To clarify this idea, let's view the learning modular arithmetic task as "completing the $p\times p$ table"; where the entries of the table correspond to the map from $(m,n)\rightarrow z$. The network has to complete the table given a fraction of the entries, a.k.a. the training data.

First, let's look at the uncorrupted training set. In this setup, every possible completion of the table corresponds to a resulting polynomial, the "target polynomial" being one of them. The fact that the network can generalize, already shows a preference for the simple target polynomial compared to other completions. This is known from the very first paper on grokking; even before the work by Nanda et al. Moreover, this bias remains even in the absence of any regularization.

Whereas in our case, a significant portion (as high as $50\%$) of the training data (i.e. pre-filled table) has corrupted labels. Since *any* table corresponds to a polynomial in modular arithmetic tasks, the model could have recognized the data as some uncorrupted training data from another polynomial. The fact that regularization can help the network overcome the penalty of making "wrong" predictions on the training set (*inversion* phase) and finding the correct simpler polynomial is highly non-trivial. This inductive bias is different and much stronger than the uncorrupted case.

To summarize, our findings on corrupted datasets points to an alternative viewpoint to the inductive bias on modular arithmetic tasks. We believe that future research should focus on explaining the (strong) bias towards a certain small set of modular polynomials amongst the astronomically large set. This may prove crucial in explaining why certain modular polynomials are un-grokkable; for example $(m^3 + mn^2 + n) \% p$ [Power et al. (2022)].

Reviewer comment (Experiments):
I find the finding that there are separate generalizing/memorizing sub-networks (or rather collections of neurons) interesting. Yet, the identification of these sub-networks is demonstrated on this toy algorithmic problem whose solution is known to be periodic. Consequently, I'm not sure if this finding translates into more general problems with non-periodic subnetworks that generalize.

Author response:
On the general dataset front, we agree that analyses using IPR cannot be carried over verbatim to general datasets like ImageNet. However, we do see that image datasets like MNIST have similar phases and inversion phenomena as modular arithmetic datasets (Appendix H). Thus, they feature the same phenomena, albeit with a different (unknown) algorithmic bias. The challenge for future research is to identify this bias, which would allow a way to find these sub-networks for general datasets. To that end, we believe that our methodology, i.e., introducing incremental label corruption to the dataset and then finding the invariant features in various phases, could be a useful method for understanding memorization and generalization in a broader setting.

Reviewer comment (Novelty):
I do not understand the main takeaway of sections 3.1 and 3.2. Is there any new finding or just a confirmation of well-known results?

Author response:
In the previous version of the manuscript, we dedicated Section 3 to summarize the effect of various regularization techniques and compare these effects with existing literature.
In the revised version, considering the reviewers' suggestions, we have shortened this discussion in favor of the results for general architectures -- Transformers and deeper MLPs.

Reviewer comment (Evaluation):
Likewise, batch normalization seems to downweight the low IPR activations but does not zero them out, meaning that memorization still occurs. How do the authors interpret this?

Author response:
Indeed, BatchNorm does not zero-out the low IPR neurons, but merely downweights them. As a result, the regularization effect of BatchNorm is less substantial compared to weight decay and dropout. This can be seen by comparing the phase diagrams in the Figures 4, 8. We emphasize that the regularizing effect of BatchNorm, albeit weaker, is noteworthy due to its interpretability.

Similar to the logits in Figure 23, we also see multiple peaks in the presence of BatchNorm. The suppression of low IPR neurons with BatchNorm suppresses the peak corresponding to the corrupted label in favor of the correct one, which explains the regularizing effect of BatchNorm.

Reviewer comment (Writing):
_It’s role ..._ <--- typo

Author response:
We thank the reviewer for pointing out the typos and additional citations. We have made these corrections in the revised manuscript.