To Grok or not to Grok: Disentangling Generalization and Memorization on Corrupted Algorithmic Datasets
Authors: Darshil Doshi
Aritra Das
Tianyu He
Andrey Gromov
Venue: ICLR 2024 poster (2024)
Keywords: Interpretability
Grokking
Label noise
Generalization
Memorization
Representations
Modular Arithmetic
Primary area: visualization or interpretation of learned representations
TL;DR: We study interpretable models+task where generalizing and memorizing representations are distinguishable. Training with corrupted data, we isolate the performance on corrupted and uncorrupted data; and explain the effect of explicit regularization.

Abstract: Robust generalization is a major challenge in deep learning, particularly when the number of trainable parameters is very large. In general, it is very difficult to know if the network has memorized a particular set of examples or understood the underlying rule (or both). Motivated by this challenge, we study an interpretable model where generalizing representations are understood analytically, and are easily distinguishable from the memorizing ones. Namely, we consider multi-layer perceptron (MLP) and Transformer architectures trained on modular arithmetic tasks, where ($\xi \cdot 100\\%$) of labels are corrupted (*i.e.* some results of the modular operations in the training set are incorrect). We show that (i) it is possible for the network to memorize the corrupted labels *and* achieve $100\\%$ generalization at the same time; (ii) the memorizing neurons can be identified and pruned, lowering the accuracy on corrupted data and improving the accuracy on uncorrupted data; (iii) regularization methods such as weight decay, dropout and BatchNorm force the network to ignore the corrupted data during optimization, and achieve $100\\%$ accuracy on the uncorrupted dataset; and (iv) the effect of these regularization methods is ("mechanistically") interpretable: weight decay and dropout force all the neurons to learn generalizing representations, while BatchNorm de-amplifies the output of memorizing neurons and amplifies the output of the generalizing ones. Finally, we show that in the presence of regularization, the training dynamics involves two consecutive stages: first, the network undergoes *grokking* dynamics reaching high train *and* test accuracy; second, it unlearns the memorizing representations, where the train accuracy suddenly jumps from $100\\%$ to $100 (1-\xi)\\%$.

Published as a conference paper at ICLR 2024
TO GROK OR NOT TO GROK: DISENTANGLING GENER- ALIZATION AND MEMORIZATION ON CORRUPTED AL- GORITHMIC DATASETS
Darshil Doshi b.
Aritra Das b.*
Tiayu He b,
Andrey Gromov a.a,b
(ddoshi, aritrad, tianyuh]eumd.eem
groomvvnndeetaa.co com
ABSTRACT
Robust generalization is major challenge deep learning, particularly when the number of trainable parameters is very large. In general, it is very difficult know if the network has memorized particular set of examples or understood the underlying rule (or both). Motivated by this challenge, we study an interpretable model where generalizing representations are understood analytically, and are eas- ily distinguishable from the memorizing ones. Namely, we consider multi multi-layer perceptron (MLP) and Transformer architectures trained on modular arithmetic tasks, where (E 100%) of labels are corrupted (i.e. some results of the modular operations in the training set are incorrect). We show that (i) it is possible for the network to memorize the corrupted labels and achieve 100% generalization at the same time; (ii) the memorizing neurons can be identified and pruned, lowering the accuracy on corrupted data and improving the accuracy on uncorrupted data; (iii) regularization methods such as weight decay, dropout and BatchNorm force the network to ignore the corrupted data during optimization, and achieve 100% ac- curacy the uncorrupted dataset; and (iv) the effect of these regularization meth- ods is ("mechanistically") interpretable: weight decay and dropout force all the neurons learn generalizing representations, while BatchNorm de-amplifies the output of memorizing neurons and amplifies the output of the generalizing ones. Finally, we show that in the presence of regularization, the training dynamics in- volves two consecutive stages: first, the network undergoes grokking dynamics reaching high train and test accuracy; second, it unlearns the memorizing repre- sentations, where the train accuracy suddenly jumps from 100% to 100(1 <)%
INTRODUCTION
The astounding progress of deep learning in the last decade has been facilitated by massive, high- quality datasets. Annotated real-world datasets inevitably contain noisy labels, due to biases of annotation schemes (Paolacci et al. 2010 Cothey| 2004) or inherent ambiguity (Beyer et al. 2020). key challenge in training large models to prevent overfitting the noisy data and attain robust generalization performance, On the other hand, in large models, it is possible for memorization and generalization to coexist (Zhang et al.) 2017: 202T). By and large, the tussle between memorization and generalization, especially in the presence of label corruption, remains poorly understood.
In generative language models the problem of memorization is even more nuanced. On the one hand, some factual knowledge is critical for the language models to produce accurate information. On the other hand, verbatim memorization the training data is generally unwanted due to privacy concerns (Carlini et al. 202T). The full scope of the conditions that affect memorization in large
"Condensed Matter Theory Center, University of Maryland, College Park DDeparmeent of Physics, University of Maryland, College Park
'Meta AI
soonespn author
"These authors contributed equally
Code reproduce results is available httptp/////ggtttbb.mmo/////

---

Published as a conference paper at ICLR 2024
Coexistence
Partial Inversion
Full Inversion
100-
100
100
accuracy (%)
75-
75-
75
50
50
50
train test
train
train
25
25
test
25
test
0-
250
500
750
1000
200
400
200
400
steps
steps
steps
(a)
(b)
(c)
Figure 1: Training curves in various phases. All plots are made for networks trained with data- fraction a = 0.5 and corruption-fraction E = 0.35. (a) No regularization: Coexistence of general- ization and memorization - both train and test accuracies are high. (b)(c) Adding weight decay: The network generalizes on the test data but does not memorize the corrupted training data, resulting in negative generalization gap! Remarkably, the network predicts the "true" labels for the corrupted examples. We term these phases Partial Inversion and Full Inversion, based on the degree of mem- orization of corrupted data. ("inversion" refers to test accuracy being higher than train accuracy.)
language models (LLMs) is still under investigation (Carlini al. 2022
Tirumala et al. 2022). In
particular, larger models as well as repeated data (Hernandez et al. 2022 favor memorization.
In this work, we take an approach to disentangling generalization and memorization motivated by theoretical physics. Namely, we focus on simple, analytically tractable, ("mechanisticallyy inter- pretable models trained on algorithmic datasets. More concretely, we study two-layer Multilayer Perceptrons (MLP) trained modular arithmetic tasks. In this case, the network is trained to learn a rule, such as = (m + n n)pp = m + n mod p from a number of examples. These tasks are phrased as classification problems. We then force the network memorize & 100% of the exam- ples by corrupting the labels, iee. injecting data points where (m + %Pp is mapped to 2' * 2. In this setting, the generalizing representations are understood qualitatively Nanda et al. (2023); Liu et al. (2022) as well as quantitatively, thanks to the availability of an analytic solution for the rep- resentations described in /romov (2023). We empirically show that sufficiently large networks can memorize the corrupted examples while achieving nearly perfect generalization on the test set. The neurons responsible for memorization can be identified using inverse participation ratio (IPR) and pruned away, leading to perfect accuracy on the un-corrupted dataset. We further show that regu- larization dramatically affects the network's ability to memorize the corrupted examples. Weight decay, Dropout and BatchNorm all prevent memorization, but in different ways: Weight decay and Dropout eliminate the memorizing neurons by converting them into generalizing ones, while Batch- Norm de- amplifies the signal coming from the memorizing neurons without eliminating them.
1.1 PRELIMINARIES
Grokking modular arithmetic. We consider the modular addition task (m + n) %.p It can be learned by two-layer MIP. Explicitly, the network function takes form
fmm n) W (em + e,,)) Wo e.. Ve,)
Here, e, en E Iº are one.hot encoded numbers m, n. "@" denotes concatenation of vectors (em @ en TP.. Win E NR and W E IRPxN are the first and second layer weight matrices, respectively. ø is the activation function. Win is decomposed into two N x P blocks: U, V E R.XP, U, V serve as embedding vectors for m. respectively. f(m, n) E RP is the network- output on one example pair (m, n2). The targets are one encoded answers e(m+n) V.'
The dataset consists of p² examples pairs; from which we randomly pick ap² examples for training and use the other (1 - a)p examples as a test set. We will refer to a as the data fraction.
This setup exhibits grokking: delayed and sudden occurrence of generalization, long after memo-
rization (Power et al. 2022 Nanda al. 2023 Liu et al.) 2022 Gromov 2023).
Throughout the text, networks are trained with Full-batch AdamW and MSE loss, unless otherwise stated.
2

---

Published as a conference paper at ICLR 2024
100
1.00
0.2
before grokking after grokking
500
before grokking
accuracy (%)
75
0.1
of neurons
400
after grakking
50
0.50
0
0.0
300
train acc
200
25
acc
0.25
-0.1
100
25
50
avg
75
100
0.00
-0.
100
8.0
0.2
0.4
0.6
0.8
1.0
steps
IPR
(a)
(b)
(c)
Figure 2: Grokking the modular arithmetic task over Egy with 2-layer MLP, trained with Adam W. (a) Sharp transition test accuracy, long after overfitting. IPR == Ex [IPRg] monotonically in-
creases over time, indicating periodic representations. (b) Example row vector (Uk.) before and
after grokking. Generalization is achieved through periodic weights (equation 2 (c) Histogram of per-neuron IPRs before and after grokking the distribution shifts to high IPR.
Periodic weights. Choosing quadratic activation function x(x) = x.² makes the problem analyti-
cally solvable [Gromov] 2023). Upon grokking, the trained network weights are qualitatively similar to the analytical solution presented in Gromov (2023) The following analytical expression for the network weights leads to 100% generalization (for sufficiently large width)
2T (k) +
Vkj (2)
2TT jo(k) + +
(1)
Uki
COS
COS
N
p
p
2 N
2m qo (k)
W pk
COS
p
(2)
where (lk) denotes a random permutation in Sn reflecting the permutation symmetry of the hidden neurons. The phase of( and (0) are uniformly i..d sampled between (-T, x]. We refer the reader to Appendix D for further details about the analytical solution.
Inverse participation ratio. To characterize the generalizing representations quantitatively, utilize a quantity familiar from the physics of localization: the inverse participation ratio (IPR). Its role is detect periodicity in the weight matrix.
Let U V.. (.x.) denote the Ath row(column) vectors of the weights U, (WW. Consider the
discrete Fourier transforms of these vectors, denoted by Uk., Vz-(Wi.) The Fourier decompositions of the periodic rows(columns) of these weights are highly localized. We leverage this to quantify the similarity of trained weights to the analytic solution (equation 2). To that end, we define the Inverse
Participation Ratio (IPR) (Girvin & Yang 2019 Pastor-Satori & Castellano 2016 for these vectors:
Gromov 2023
||UAo:14 1
IPR(
IPR(") =
IPR (w)
VV ;
W.
where || denotes the -norm of the vector. One can readily see from equation 3 that IPR E [1/p, 1]. with higher values for periodic vectors and lower values for non non-periodic ones. useful to quantify the periodicity of each neuron in the hidden layer of the network (indexed by k). We define per-neuron IPR by averaging over the weight-vectors connected to the neuron:
1 IPR + IPR (v) + IPR("))
IPRg = 3
(4)
Since trained networks generalize via the periodic features, a larger population of high-IPR neu- rons leads to better generalization. To quantify the overall similarity of the trained network to the analytical solution (equation) 2 we average IPR, over all hidden neurons: IPR Ex [IPRR].
PPeriodicity of weights is approximate in trained networks. 'In general IPR is defined (1) U/e We set a common choice.
3

---

Published as a conference paper at ICLR 2024
Label corruption. To introduce label corruption, we choose a fraction & of the training examples, and replace the true labels with random labels. The corrupted labels are generated from a uniform random distribution over all the labels. This is often called symunetric label noise in literature (Song et al. 2022). This type of label corruption does not introduce label asymmetry.
1.2 RELATED WORKS
Grokking Grokking was first reported by Power et al. 2022) for Transformers trained modular arithmetic datasets. [Luu et al. (2022) 2023) presented explanations for grokking in terms of quality of representation learning as well weight-norms. They also found grokking- -iie behaviours other datasets. Nanda et al. 2023); Zhong et al. (2023) reverse- enverse-nggineeredd grokked Transformer models and revealed the underlying learned algorithm. Barrk et al. (2022) observed grokking on the sparse- parity problem, which they attributed to the amplification of the so-called Fourier gap. Merrill et al. (203)) examined grokking in the ppase-pprrtty setup and attributed it to the emergence of sparse subnetworks. Zunkovië & Ilievski (022)) discuss solvable models for grokking in teacher-student setup. Notsaawo et al. (2023) investigated the loss landscape at early training time and used to predict the occurrence of grokking. Varma al. (2023) demonstrated new by investigating the critical amount of data needed for grokking.
behaviours
Label Noise Zhang et al. (2017) showed that neural networks can easily fit random labels. A large body of work has been dedicated to developing techniques for robust training with label noise: ex- plicit regularization
al.) al.] 2012 [offe & Szegedy| (015). noise-robust loss functions (Ghosh)
et al. 2017), sample selection (Katharopoulos & Fleuret 2018 8:
Paul et al., (021), noise modeling
and dedicated architectures (Xiao et al. 2015 Han et al.
2018
Yao et al.) 2019) etc. We refer the
reader to recent reviews by [Sng al. (022)) Liang et al. (2022) for comprchensive overview.
Regularization and Pruning Morcos al.] (2018) investigated the interplay between label cor- ruption and regularization in vision tasks. They showed that BatchNorm (loffe & Szegedy 2015) and Dropout (Hinton et al. 2012) have qualitatively different effects. Dropout effectively reduces the network-size but does not discourage over-reliance on a few special directions in weights space; while BatchNorm spreads the generalizing features over many dimensions. Bartoldson et al. (2020) found that there is trade-off between stability and generalization when pruning trained networks.
2 GROKKING WITH LABEL NOISE
Grokking MLP on modular addition dataset is remarkably robust to label corruption. Even without explicit regularization, the model can generalize tion. In many cases, the network surprisingly manages to "correct" some of the corrupted examples, resulting in Inversion (i.e. test accuracy training accuracy). We emphasise that this is in stark contrast to the common belief that grokking requires explicit regularization. Adding regularization makes grokking more robust to label corruption, with stronger Inversion. To quantify these be- haviours, we observe the training and test performance of the network with varying sizes of training dataset, amount of label corruption and regularization. We summarize our findings in the form of empirical phase diagrams (Figure 3 featuring the following phases. More detailed phase diagrams along with training and test performances are presented in Appendix L
100% accuracy with sizable label comup-
Coexistence. In the absence of explicit regularization the network generalizes on test data and memorizes the corrupted training data. In other words, both train and test accuracy are close to 100% despite label corruption.
Partial Inversion. In this phase, the network generalizes on the test data but only memorizes a fraction of the corrupted training data. Remarkably, the network predicts the "true" labels on the remaining corrupted examples. other words, the network corrects a fraction the corrupted data. Consequently, we see < 100% train accuracy but 1ear-100% test accuracy, resulting in a negative generalization gap (Figure [b)) We term this phenomenon Partial Inversion; where "inversion" refers to the test accuracy being higher than train accuracy. Remarkably, partial inversion occurs even in the absence of any explicit regularization, but only when there is ample training data (left- most panels in Figure g a.))).
4

---

Published as a conference paper at ICLR 2024
wd = 0.0
wd = 5.0
wd 15.0
wd 20.0
0.8
corruption frac
Memorization
Memorization
Memorization
Memorization
0.6
Coexistence
P..
0.4
Partial Inversion
Full Inversion
Full Inversion
Forgetting
0.2
Coexistence
0.2
0,4
0.6
0.8
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
0.2
0,4
0.6
0.8
data frac
data frac
data frac
data
frac
(a) Regularization with weight decay. (wd denotes the weight decay value)
dropout prob.
0.0
dropout prob.
0.3
dropout prob.
0.6
dropout prob.
0.9
corruption frac
Memorization
Memorization
Memorization
Confusion
Coexistence
P.I.
Partial
Inversion
Full
Full Inversion
Full Inversion
Coexistence
Inversion
0.6 data frac
data frac
data frac
data frac
(b) Regularization with Dropout.
Figure 3: Modular Addition phase diagrams with various regularization methods. A larger data fraction leads to more "correct" examples, leading to higher corruption- robuutiss Increasing reg- ularization, in the form of weight decay dropout, enhances robustness label corruption and facilitates better generalization.
Full Inversion. Upon adding regularization, the network often generalizes on the test set but does memorize any of the corrupted examples. This results in training accuracy plateauing close to 100(1 - % (E is the corruption fraction); while the test accuracy is at 100% (Figure [(c). Increasing regularization enhances this phase (middle columns in Figure a.))).
Memorization. With very high label corruption and/or very low training data cata-fraction the network memorizes all of the training data (corrupted as well as non-corrupted), but does not generalize. We close 100% accuracy on training data, but random guessing accuracy on test data in this phase.
Forgetting. With very high weight decay, the network performs poorly on both training and test data (right-most panel in Figure 3 a)). However, upon examination of the training curves (Figure M()) we find that the accuracies plummet after the network has generalized! This performance drop result of the norms of weights collapsing arbitrarily small values (Figure {())). Forgetting exclusively occurs grokked networks with activation functions degree higher than (e.g. we use d(x) 2²2 and with adaptive optimizers. (See Appendix for detailed discussion.)
With very high dropout the network becomes effectively narrow, resulting in reduced capacity to memorize. Consequently, even in the absence of grokking, the network does not manage memorize all the corrupted labels (right-most panel in Figure b)).
Confusion.
2.1 GENERALIZING AND MEMORIZING SUB-NETWORKS
The generalizing features in our setup are "mechanisticall"" interpretable (equation As a result, leverage measures of periodicity such IPR to quantitatively characterize various phases presented above. also allows us to analyze the effect of various regularization schemes. To that end, state the following hypothesis and extensively examine it with experiments.

---

Published as a conference paper at ICLR 2024
Partial Inversion
Full Inversion
Forgetting Confusion
Coexistence
Memorization
Memorizes corrupted train data
Yes
Partially
No
Yes
No
Generalizes
Yes
Yes
Yes
No
No
on test data
Table 1: Qualitative characterization of phases.
Coexistence (wd 0.0)
Parita Inversion (wd =.00
Full Inversion (wd = =1.0)
300
300
300
of neurons
100
100
100
8.0
0.2
0.4
0.6
0.8
0.0
0.2
0.4
0.6
8.0
0.2
0.4
0.8
1.0
IPR
IPR
IPR
(a) Regularization with weight decay. (wd denotes the weight decay value)
Coexistence (dp = 0.0)
Parital Inversion (dp 0.3)
Full Inversion (dp 0.6)
300
300
300
# of neurons
200
200
200
100
8.0
0.2
0.4
0.6
10
0
0.2
0.4
1.0
8.0
0.2
0.4
1.0
IPR
IPR
IPR
(b) Regularization with Dropout. (dp denotes the Dropout probability)
Figure 4: Distribution of per-neuron IPR for trained networks in various phases. Coexistence phase has bimodal IPR distribution, where the high and low IPR neurons facilitate generalization and memorization, respectively. Regularization with weight decay or Dropout shifts the IPR distribution towards higher values Generalizing neurons get more populous compared memorizing ones; resulting in more robust generalization and Inversion. All plots are made for networks trained with data-fraction a = 0.5 and corruption-fraction 0.35; with various regularization strengths.
Hypothesis 2.1. For a two-layer MLP with quadratic activation trained on (one-hot encoded) mod- ular addition datasets with label corruption, the high IPR neurons facilitate generalization via feature-lerrning whereas the low IPR neurons cause memorization of corrupted training data.
To test Hypothesis 2.1 we study the distribution of per-neuron IPR various phases, with differ- ent levels of regularization. In the Coexistence phase, the network simultaneously generalizes and memorizes the corrupted train data. From Hypothesis 2.1 we expect the presence of both general- izing and memorizing neurons. Indeed, we find bi-modal distribution of IPRs in trained networks (left column of Figure 4. Remarkably, the network not only distinguishes between the tasks of feature-leaming and memorization, it also assigns distinct and independent sub-networks to each.
As mentioned the previous section, adding regularization in the form of weight decay or Dropout makes the generalization more robust and facilitates Inversion. Since these phases see decreased memorization of corrupted labels, we expect to see a decrease in the number of memorizing neu- and corresponding increase in the generalizing ones. Indeed, observe distribution-shi towards higher IPR in the Partial Inversion phase (middle column of Figure 4. the Full Inversion that memorizing neurons almost entirely phase (induced by weight decay or dropout),

---

Published as conference paper at ICLR 2024
withouth BatchNorm
BatchNorm
[BatchNorm weightj
300
300
of neurons
of neurons
0.04
200
200
0.03
* 100
100
0.02
0.0
0.2
0.8
1.0
8.0
0.4
0.6
0.8
1.0
0.00
0.25
0.50
0.75
1.00
IPR
IPR
IPR
Figure 5: IPR distributions of networks with/without BatchNorm; the correlation between Batch- Norm weights and IPR. The IPR distributions in the models do not exhibit significant differ- ences. However, BatchNorm helps generalization by assigning higher weights (7k) to high IPR neurons compared low IPR neurons. Both models trained data fraction 0.65 and cor- ruption fraction ç = 0.2, with batch-size 64.
Pruning Memorizing (low IPR)
Pruning Generalizing (high
1.00
100
accuracy (96)
80
accuracy (9)
80
training test set
corrupted training data uncorrupted training 100 E1%
40
8.0
0.2
0.6
8.0
0.1 0.2 pruned
fraction of neurons pruned (from lowest IPR)
fraction (from highest IPR)
Figure 6: Progressively pruning neurons from trained model (Coexistence phase), based on IPR. (data fraction a 0.5, corruption fraction g = 0.35). (Left) Pruning out neurons starting from the lowest IPR. The accuracy corrupted training data decreases, providing evidence that low IPR neurons are responsible for memorization. Test accuracy as well as accuracy on uncorrupted training data remains high. (Right) Pruning neurons starting from the highest IPR. Test accuracy decreases providing evidence that high IPR neurons are responsible for generalization. Accuracy on corrupted training data remains high.
eliminated in favour of generalizing ones (right column of Figure 4). This provides mechanism by which the network recovers the "true" labels even for the corrupted data - in the absence of memo- rizing sub-network, its predictions are guided by the underlying rule and corrupted samples.
BatchNorm also has a regularizing effect on training (Figure 9), albeit, through a qualitatively dis- tinct mechanism. In the network (Equation (T). we apply BatchNorm after the activation function (post-BN):
f(m, n) W BN (6 (U em + Ven))
(5)
BatchNorm layer does does not affect the IPR distribution of the neurons significantly the low IPR neurons persist even in the Full Inversion phase (Figure Instead, BatchNorm adjusts its own weights de-amplify the low-IPR (memorizing) neurons, blasing the network towards generalizing representations. This can be seen from the strong correlation between IPR, and the corresponding BatchNorm weights in Figure [irrihht). BatchNorm effectively prunes the network and learns to do this directly from the data! Note that this effect of BatchNorm acts in conjunction with the implicit regularization from the mini-batch noise in this case. We find that decreasing the batch-size enhances Inversion even without explicit regularization (Figure I0).
To further isolate the effect individual neurons on performance, we perform the two complemen- tary pruning experiments (Figure 6) Starting from trained network in the Coexistence phase, gradually prune out neurons, one-at-a-tin time, the highest IPR neuron, while keeping track of training and test accuracies. To distinguish network's ability to memorize, we also track the accuracies on corrupted and uncorrupted parts the training dataset. Before pruning, the network has 10ar-100% accuracy both train and test data. In
(i) starting from the lowest IPR neuron (ii) starting from

---

Published as a conference paper at ICLR 2024
10.0
wd 50.0
Coexistence (wd 10.0)
Full Inversion (wd 50.0)
250
corruption frac
o neuron as
200
Memorization
Memorization
150
100
PL.
Full
Conxistence
Inversion
0.6
0.8
0.2
0.4
0.6
8.
IPR
(a) I-layer Transformer.
wd 5.0
wd
20.0
Coexistence (wd 5.0)
Full Inversion (wd 20.0)
500
corruption frac
of neurons
400-
Memorization
Memorization
300
200
RL
Coevistence
Full Inversion
100
0.2
0.4
0.6 0.8
0.2
0.4
0.6
0.8
8.0
0.2
0.4
0.6
0.8
1.00.0
0.2
0.4
0.6
0.8
1.0
data frac
data frac
IPR
IPR
(b) 3-layer ReLU MLP.
Figure 7: Modular Addition phase diagrams and IPR histograms with different architectures and weight decay values (denoted by "wd"). The IPR histograms plotted for neurons connected to the input (embedding) layers.
(i) (Figure lleet)), we see a monotonic decrease in accuracy on corrupted examples as more and more low IPR neurons are pruned. The accuracy on test data as well ucorrupted training data remains high. The resulting (overall) training accuracy platcaus around 100(1 - 4),, corresponding to the fraction of uncorrupted examples. In case (ii) (Figure (right)) more high IPR neurons pruned, we see a decline in test accuracy (as well as the accuracy on the uncorrupted training data). The network retains its ability memorize, the accuracy corrupted training data remains high.
We emphasize that, although each generalizing neuron learns the useful periodic features, the overall computation in the network is emergent: it is performed collectively by all generalizing neurons.
2.2 GENERAL ARCHITECTURES
this subsection, we aim to check (i) we can quantify memorizaation/geeeniiaaaioo capabilities of general architectures using IPR distribution of the neurons; (ii) our understanding of regularization can be generalized to such cases.
Hypothesis 2.2. For any network that achieves non-trivial generalization accuracy the modular addition dataset with label corruption, operation can be decomposed into two steps: (I) The input (embedding) layer with row-wise periodic weight maps the data onto periodic features; (II) The remainder of the network nonlinearly implements trigonometric operations over the periodic neurons and maps them the output predictions. We hypothesize that step (I) depends only the dataset characteristics; and hence, is a universal property across various network architectures.
To verify Hypothesis 2.2 and better understand the impact of regularization, we conduct further experiments on 1 1-layer Transformers and 3-layer ReL.U MI.Ps. Our findings, illustrated the first two columns of Figure 7 and paralleling the observations in Figure a), indicate that increased weight decay promotes generalization over memorization. IPR histograms in the last two columns of Figure 7 further confirm that weight decay encourages the correct features and corroborates the assertions made in Hypothesis 2.2
For these general models, Dropout mostly leads extended confusion phase, while BatchNorm only facilitates coexistence. We observe weight decay to be essential for these models to have the full inversion phase. We refer the reader to Appendix for in-depth discussion.
8

---

Published as a conference paper at ICLR 2024
3 EFFECT OF REGULARIZATION
In this section, we summarize the effect of various regularization techniques and compare them with existing literature.
3.1 WEIGHT DECAY
As shown Figure B(a, increasing weight decay has two notable effects: (i) Generalization with less data: The phase boundary between the Memorization and Coexistence phases shifts toward lower data-fraction. Weight decay encourages generalization even when training data is scarce; in line with classical ideas (Bartleu 1998). (ii) Emergence of the Full Inversion phase: In this phase, weight decay prevents memorization and helps the network "correct" the labels for corrupted data points. This effect further elucidated in Figure For fixed amount of data and label corrup- tion, increased weight decay results in diminished proportion of low IPR (memorizing) neurons in favor of high IPR (generalizing) ones.
3.2 DROPOUT
is believed that dropout prevents complex co-adaptations between neurons in different layers, en- couraging each neuron learn useful features (Hinton al. 2012 Srivastava al.) (014)). This intuition becomes more precise in our setup. The generalizing solution constitutes of periodic neu- rons, each corresponding a weight-vector of given "frequency" (l+). Such solution requires neurons with different frequencies (see Equation 6) In wide network (N = 500 in our setup), there can be multiple neurons with the same frequency Due this redundancy, generalizing solu- tions are more robust to dropped neurons; whereas overfitting solutions are not. As a result, training with Dropout gets attracted such periodic solutions.
CONCLUSION
We have investigated the interplay between memorization and generalization in two-layer neural networks trained on algorithmic datasets. The memorization was induced by creating dataset with corrupted labels, forcing the network to memorize incorrect examples and learn the rule at the same time. We have found that the memorizing neurons can be explicitly identified and pruned away leading to perfect generalization.
Next we showed that various regularization methods such as weight decay, Dropout and BatchNorm encourage the network ignore the corrupted labels and lead to perfect generalization. We have leveraged the interpretability of the representations learnt the modular addition task to quanti- tatively characterize the effects of these regularization methods. Namely, we have found that the weight decay and dropout act in similar way by reducing the number of memorizing neurons; while BatchNorm de-amplifies the output of regularizing neurons without eliminating them, effec- tively pruning the network.
5 LIMITATIONS AND FUTURE WORK
Although the trained models in our setup are completely interpretable by means of analytical so- lutions, the training dynamics that lead to these solutions remains an open question. Solving the dynamics should also shed more light the effects of regularization and label corruption.
While our work is limited to the exactly solvable models and algorithmic datasets, paves the way quantitative investigation of generalization and memorization in more realistic settings.
$ "oore precisely, it only requires [p/2] frequencics for real weights. 'The analytical solution becomes exact only large N; with sub-lcading corrections of order 1/VN.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
A.G.'s work at the University of Maryland was supported in part by NSF CAREER Award DMR- 2045181, Sloan Foundation and the Laboratory for Physical Sciences through the Condensed Matter Theory Center. The authors acknowledge the University of Maryland supercomputing resources (http///ppcc.umd..dee wm....cmmmddd............ edu) made available for conducting the research reported in this paper.
REFERENCES
Boaz Barak, Benjamin L. Edelman, Surbhi Goel, Sham M. Kakade, eran malach, and Cyril Zhang- Hidden progress in deep learning: SGD learns parities near the computational limit. In Al- ice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neu- ral Information Processing Systems, 2022. URL /teeps:///penreview... net htwp:r/rwwwww.eeoooootmotm formmmir Neww-eww-i1 im
P.L. Bartlett. The sample complexity of pattern classification with neural networks: the size of the weights more important than the size of the network. IEEE Transactions on Information Theory, 44(2):525-533 536, 1998. doi: 10.1 109/10911/00 5026
Brian R. Bartoldson, Ari S. Morcos, Adrian Barbu, and Gordon Erlebacher. The generalization- stability tradeoff in neural network pruning, 2020.
Lucas Beyer, Olivier J. Hénaff, Alexander Kolesnikov, Xiaohua Zhai, and Aäron van den Oord. Are we done with imagenet?, 2020.
Nicholas Carlini, Florian Tramer, Eric Wallace, Matthew Jagielski, Ariel Herbert-Vose Katherine Lee, Adam Roberts, Tom Brown, Dawn Song, Ulfar Erlingsson, al. Extracting training data from large language models. In 30th USENIX Security Symposium (USENIX Security 21), pp- 2633-2650, 2021.
Nicholas Carlini, Daphne Ippolito, Matthew Jagielski, Katherine Lee, Florian Tramer, and Chiyuan Zhang. Quantifying memorization across neural language models. arXiv preprint arXiv:2202.07644 07646, 2022.
Viv Cothey, Web-crawling reliability. Journal of the American Society for Information Science and Technology, 55(14): 1238, 2004. doi: https:///doi.og///0.00///sssisi 20078. URL //onnineeibbraaaeo wiley com/ htt//doi//bba 1ttp://...1///a..20002
Aritra Ghosh, Himanshu Kumar, and P Shanti Sastry. Robust loss functions under label noise for deep neural networks. In Proceedings of the AAAI conference on artificial intelligence, volume 31, 2017.
Steven M. Girvin and Kun Yang. Modern Condensed Matter Physics. Cambridge University Press, 2019. doi: 10.1017/7781116 131466649.
Andrey Gromov.
Grokking modular arithmetic, 2023. URL https://axiv.org/eee
2301. 02679
Bo Han, Jiangchao Yao. Gang Niu, Mingyuan Zhou, Ivor Tsang, Ya Zhang, and Masashi Sugiyama. Masking: A new perspective of noisy supervision. In Bengio. H. Wallach, H. Larochelle, K. Grauman, N. Aesa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31, 2018.
Danny Hernandez, Tom Brown, Tom Conerly, Nova DasSarma, Dawn Drain, Sheer El-Showk, Nel- Elhage, Zac Hatfiel-Dodds, Tom Henighan, Tristan Hume, al. Scaling laws and inter- pretability of learning from repeated data. arXiv preprint arXiv:2255. /048, 2022.
Geoffrey E. Hinton, Nitish Srivastava, Alex Krizhevsky, Ilya Sutskever, and Ruslan R. Salakhutdi- nov. Improving neural networks by preventing co-adaptation of feature detectors, 2012.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift, 2015.
10

---

Published as a conference paper at ICLR 2024
Arthur Jacot, François Ged, Berfin Şimşek, Clément Hongler, and Franck Gabriel. soddle-to-saddee dynamics in deep linear networks: Small initialization training. symmetry, and sparsity, 2022.
Angelos Katharopoulos and Francois Fleuret. Not all samples are created equal: Deep learning with importance sampling. In Proceedings the 35th International Conference Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 2525-25344 PMLR, 10-15 Jul 2018.
Xuefeng Liang, Xingyu Liu, and Longshan Yao. Review-a survey of learning from noisy labels. ECS Sensors Plus, 1(2):02140) jun 2022. doi: 10.1149/7754222264477551 URL https:// wtw...org....11..
dx ex. www.orgg/11... 1tt.1911.119212185
2726 2/c755
Ziming Liu, Ouail Kitouni, Niklas S Nolte, Eric Michaud, Max Tegmark, and Mike Williams. To- wards understanding grokking: An effective theory of representation learning. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neu- ral Information Processing Systems, volume 35, pp. 34651-34663. Curran Associates, Inc., 2022. URL ltttps:/ wtwp:eeedddngg..uusssss.o.moo www.neuuiiss... co/pper hatpaaaerr_ieeepappapeeeee fitess/appp/e///22// pdf
Ziming Liu, Eric Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data.
In The Eleventh International Conference on Learning Representations, 2023. URL https:
//opezeeview Sepeereeveewnnt//Eeee oet.
William Merrill, Nikolaos Tsilivis, and Aman Shukla. A tale of two circuits: Grokking as competi- tion of sparse and dense subnetworks, 2023.
Ari S. Morcos, David G. T. Barrett, Neil C. Rabinowitz, and Matthew Botvinick. On the importance of single directions for generalization, 2018.
Neel Nanda, Lawrence Chan, Tom Lieberum. Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability.. 2023.
Pascal Jr. Tikeng Notsawo, Hattie Zhou, Mohammad Pezeshki, Irina Rish, and Guillaume Dumas. Predicting grokking long before it happens: A look into the loss landscape of models which grok, 2023.
Gabriele Paolacci, Jesse Chandler, and Panagiotis G. Ipeirotis. Running experiments on ama- zon mechanical turk. Judgment and Decision Making, 5(5):411-419, 2010. doi: 10.1017/
$1930297500002255.
Romualdo Pastor-Satorras and Claudio Castellano. Distinct types of eigenvector localization in networks. Scientific Reports, 6(1):18847, 2016. doi: 10.1038/srepl 8847. WRLLhttps://ddd
htr//00.038//ree.
Mansheej Paul, Surya Ganguli, and Gintare Karolina Dziugaite. Deep learning on a data diet: Finding important examples early in training. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances Neural Information Processing Systems, volume 34, pp. 10596-20601 Curran Associates, Inc., 2021.
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gener- alization beyond overfitting on small algorithmic datasets, 2022.
Hwanjun Song, Minseok Kim, Dongmin Park, Yooju Shin, and Jae-Gil Lee. Learning from noisy labels with deep neural networks: A survey, IEEE Transactions Neural Networks and Learning Systems, pp. 1-19, 2022. doi: 10. L09/TNNLSS.022.1112222
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: A simple way to prevent neural networks from overfitting. Journal of Machine Learning Research, 15(56):192-11958 2014. URL http: http://jmrr..rrrr harg//aperss/v1.
ari Priivsttvvall 4a .html
Kushal Tirumala, Aram Markosyan, Luke Zettlemoyer, and Armen Aghajanyan. Memorization without overfitting: Analyzing the training dynamics of large language models. Advances in Newral Information Processing Systems, 35:38274-38290, 2022.

---

Published as a conference paper at ICLR 2024
Vikrant Varma, Rohin Shah, Zachary Kenton, János Kramár, and Ramana Kumar. Explaining grokking through circuit efficiency, 2023.
Tong Xiao, Tian Xia, Yi Yang, Chang Huang, and Xiaogang Wang. Learning from massive noisy labeled data for image classification. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2015.
Jiangchao Yao, Jiajie Wang, Ivor W. Tsang, Ya Zhang, Jun Sun, Chengqi Zhang, and Rui Zhang. Deep learning from noisy image labels with quality embedding. IEEE Transactions on Image 2019. doi: 10.1 10.11/9//1 10.110999120011.28799999 2.277993.
Processing,
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. International Conference on Learning Rep- resentations, 2017. httplttttt://prreeeeeee wet.wwiwwwwweeeeeeeeeee..ee.eeee /
Chiyuan Zhang, Samy Bengio. Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107- 115, 2021.
Ziqian Zhong, Ziming Liu, Max Tegmark, and Jacob Andreas. The clock and the pizza: Two stories in mechanistic explanation of neural networks, 2023.
Bojan Zunkovič and Enej Ilievski. Grokking phase transitions in learning local rules with gradient descent, 2022.
12

---

Published as a conference paper at ICLR 2024
A FURTHER EXPERIMENTAL DETAILS
In every plot presented, unless explicitly stated otherwise, we use the following default setting:
Default setting A 2-layer fully connected network with quadratic activation function and width N 500, initial weights sampled from (0, (16N) trained on 97 Modular Addition dataset. The networks are trained for 2000 optimization steps, using full-batch, MSE loss, AdamW optimizer, learning rate 17 0.01 and B (0.9, 0.98).
A.1 FIGURES MAIN TEXT
Figure
Default setting with 0.5, { = 0.35 and (a) wd = 0 (b) wd 5 and (c) wd 15.
Figure Default setting with a = 0.5, { = 0.0, wd = 5.0.
Figure Each plot is scanned over 17 different data fractions ranging from 0.1 to 0.9 and 19 noise levels { from 0.0 to 0.9 using the default setting.
Figures and
Default setting with 0.5, 5 = 0.35.
Figure 2-layer fully connected network with quadratic activation function and width 500, with and without BatchNorm. We train for 8000 steps, with MSE loss, Adam optimizer with minibatch of size 64, and learning rate r = 0.005. Data fraction a 0.65, corruption fraction =
General setting for minibatch training: We train the network with batch size 256 for 2000 steps. For batch sizes 8 and 64, we scale the number of total steps to keep the number of epochs roughly equal. (In other words, all the networks see the same amount of data.) The learning rate is scaled as 0.01 / batch size/256.
Figure 7 All networks are trained with MSE loss using AdamW optimizer with learning rate 7j 0.001 and (3), = (0.9, 0.98). (a) Encoder-only Transformer, embedding dimension dembed = 128 with ReLU activation, trained for 3000 steps with 3 random seeds; (b) Width N 500 trained for 5000 steps with 3 random seeds.
B DETAILS OF RESOURCES USED
Figure requires 72 hours on a single NVIDIA RTX 3090 GPU.
Figure used 60 hours 1/7 NVIDIA A100 GPU.
All the remaining figures except Figure 7 in the main text are either further experiments performed on the data from the Figure 3xxuureeexprrmment
very quick to perform any machine.
We used 300 GPU hours for early-stage exploration and making phase diagrams the appendix.
13

---

Published as a conference paper at ICLR 2024
C QUANTITATIVE CHARACTERIZATION OF PHASES
All phase diagrams are plotted using the empirical accuracies of trained networks.
Trained networks with 2 90% test accuracies get classified into Coexistence, Partial Inversion and Full Inversion based on their training accuracy (specifically, their ability to memorize corrupted ex- amples). Coexistence phase has 90% training accuracy, since it memorizes most of the corrupted training examples. Partial Inversion phase has < 90% but > 100(1.05 5)% training accuracy, where & is the corruption fraction. Note that we subtract Ç from 1.05 instead of 1.00 to account for the 2% random guessing accuracy well statistical fluctuations around this value. Full Inversion phase has 100(1.05 <)% training aceuracy, since it does not memorize any of the corrupted examples. Memorization phase has < 90% test accuracy but 2 90% training accuracy, indicating that the network memorizes all the training data, including the corrupted and uncor- rupted examples. Forgetting confusion phase has < 90% test well as training accuracies. the network neither generalizes nor manages to memorize all the training examples. We distinguish Forgetting from confusion using the training curve. In the Forgetring phase the grokked network loses its performance due to steadily decreasing weight-norms. Whereas the confusion phase, the network never reaches high performance.
Partial
Full Inversion
Forgetting Confusion
Coexistence
Memorization
Inversion
Test accuracy
90%
90%
90%
90%
90%
Training accuracy
< 90% and
90%
100(1.05 8)%
90%
90%
2 100(1.05 - 9
Table 2: Quantitative characterization of phases.
14

---

Published as a conference paper at ICLR 2024
D ANALYTICAL SOLUTION
We show that the network with periodic weights stated in Equation 2 indeed the analytical solution for the modular arithmetic task. The network output given by
Sa(m, n)
W (Ukm Vkn)²
2
o(kig - (0 + +
cos
N
p
km
2π o (k)m + +06%)) + cos
2n o
(k)) + ob.²) ))
cos
p
N
COS
2T
o (2 -q) +0k (n)
0x
(1)
N
4
p
1=1
+
COS
2T
o (k) + +4) +30 + +
(p)
4
2T
+
4
CO8
o (k) 9) -0/ + Qk
2T
(2) + 9) + (w) oh + 36k (1)
+
cos
4
2TT =())(mm (nn)(++n-pp)
oos
2
1
2T
-m) + +) 20// +20
2
2T
(k) 4) -20/
cos
2
1
2T
(m))m 9) 20/
+
2T
[ok)q ((xh) )(()) }
cos
(6)
We have highlighted the term that will give us the desired output with a box. Note that the desired term is the only one that does not have additive phases in the argument. Recall that the phases (() and of") randomly chosen uniformly i..d. sampled between (- rT.. Consequently. as N becomes large, all other terms will vanish due to random phase approximation. The only surviving term will be the boxed term. We can write this boxed term a more suggestive form to make the analytical solution apparent.
N Σ 1=1
(2)0(((m++ + 1)) x 8 8m + -)))
2T
fm n)
N
where we have defined the modular Kronecker Delta function (P()) Kronecker Delta function up to integer multiples the modular base p.
8 (z) = [1 0 otherwise
=
(r E Z)
(8)
where Z denotes the set of all integers. Note that 5m(m + 9) the desired one_hot encoded labels for the Modular Arithmetic task, by definition. Thus our network output with periodic weights indeed solution.
15

---

Published as a conference paper at ICLR 2024
E
FORGETTING PHASE
100
train test
accuracy (96)
weight norms
10
[IVI
[|W]
50
25
01
2000
4000
2000
4000
steps
steps
(a)
(b)
Figure 8: Forgetting phase (data fraction a 0.5, corruption fraction { 0.35, weight decay 20). (a) Training curves. (b) Frobenius norms of network weights with training.
We saw in Figure (a) (right-most panel, wd=20.0) that extremely high weight decay leads For- getting phase. This phase has random guessing accuracies on both training and test datasets at the end of training. Naively, this would seem similar the Confusion phase. However, Upon close inspection of their training dynamics, we find rich behaviour. Namely, the network initially under- goes dynamics similar to Full Inversion, i.e., the test accuracy reaches 100% while the train accuracy plateaus around 100(1 ()%. However, both accuracies subsequently decay down to random guess- ing. This curious phenomenon explained by the effect of the high weight decay on the network weights. We notice that the Frobenius norms of the weights of the network steadily decay down to arbitrarily small values.
This curious behaviour only occurs when we use activation functions of degree higher than I (e.g. ((x) = x² in this case). (More precisely, we mean that the first term in the series expansion should be of degree at least 2.) This can be intuitively explained by examining the gradients carefully. With quadratic activation functions and MSE loss, gradients take the following form
Oc
4
N
Wox (URb + Vkn)² 6 n
Wa (Uab + Van)
aUab
DDp
(b,))ED
Oc
4
Σ
Σ
Σ M=
Wor gh (Ukm + - 9) Wga (Uam + Vab) (9)
(b
VV.b
DDp
(m.b)ED
OC
4
Σ
Σ
N
Wok (Ukm + Vkn) (m + 9) (Uam + Vdn)²
aW
D
(m.))D
k=
where D is the training dataset and [D] is its size.
For non-generalizing solutions, the term inside the bracket is of degree O in weights, due to the delta function term. But close to generalizing solutions, the term inside the bracket becomes effec- tively degree 3 in weights, since the delta function term cancels with the network output. Conse- quently, around the generalizing solution, the weight gradients have a higher degree dependence on the weights themselves. This makes such solutions more prone to the collapse due high weight decay. This serves heuristic justification to why only grokked networks seemed to be prone to Forgetting, while memorizing networks seem to be immune.
If we use vanilla Gradient Descent optimizer, the updates would plateau once the balance between gradient updates and weight decay contribution reached. But this not the case in Adaptive optimizers. Indeed, we find that this phenomenon exclusively occurs with Adaptive optimizers (e.g. Adam in this case).
We believe that it is possible to get further insight into this phase by closely examining the opti- mization landscape and network dynamics. Since this behaviour is exclusive to grokked networks, understanding this phase would lead to further understanding of grokking dynamics. Alternatively, a better understanding of the dynamics of grokking should give us insights into this curious phe- nomenon. We defer this analysis for future work and welcome examination by readers.
16

---

Published as a conference paper at ICLR 2024
F COMPARISON OF BATCHNORM, LAYERNORM AND MINIBATCH w/O NORMALIZATION
Here examine the effects of BatchNorm, LayerNorm and mini-batch noise. We find that gen- eralization and Inversion are indeed enhanced by decreasing the batch-size. This demonstrates that implicit regularization from mini-batch noise favours simpler solutions in our setup.
Additionally, upon comparing Figure 9 Figure and Figure I1 we find enhanced Inversion in cases with normalization layers (especially BatchNorm). For BatchNorm, we provide explanation in the main text, with the the BatchNorm weights de-amplifying memorizing neurons in favour of generalizing ones. For LayerNorm, however, we do not find such effect.
batch size 256
batch size 64
batch size 32
batch size 8
0.8
corruption frac
Memorization
Memorization
Memorization
Memorization
0.6
0.4
Coexistence
Coexistence
P.I.
Coesistence
P.I
P.L
Full
Full
Full Inversion
Full Inversion
0.2
Inversion
Inversion
0.2
0.4
0.6
0.8
0
0.4
0.6
0.8
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
data frac
data frac
data frac
data frac
Figure 9: Phase diagram with BatchNorm BNN
Full batch
Batch size
256
Batch size
0.8
corruption frac
Memorization
Memorization
Memorization
0.4
Partial Inversion
Coexistence
PI.
P.I.
Full
0.2
Coexistence
F.I
C.E.
Inversion
0.2
0.4
0.6
0.8
0.2
0.8
0.2
0.4
0.6
0.8
data frac
data frac
data
frac
Figure 10: Phase diagram with mini-brr bathh training, without any explicit regularization normal-
ization layers.
Full batch
Batch size 256
Batch size
corruption frac
Memorization
Memorization
Memorization
0.6
0.4
Partial Inversion
Coexistence
P.I.
P.I.
Full Inversion
Full
0.2
Coexistence
C.E.
Inversion
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
data frac
data
frac
data
frac
Figure 11: Phase diagram with LayerNorm (post-LNN
17

---

Published as a conference paper at ICLR 2024
G TRANSFORMERS AND DEEP MLPS
In this section, we present (i) more weight decay experiments on different configurations of Trans- formers; (ii) the Effect of Dropout on Transformers; (iii) the Effect of BatchNorm on 3-layer ReLU MLPs.
For reasons that do fully understand, 3-layer ReLU MLPs with even 0.05 Dropout probabil- ities failed to memorize the dataset for any data fraction a > 0.3 and noise level & 2 0.1, even after an extensive hyperparameter search. We leave the discussion about this for the future.
G. EFFECT OF DEPTH NUMBER OF HEADS WRIGHT TYING
To further support our Hypothesis 2.2
we conducted further experiments by varying the config-
uration of Transformer models. Figure T2
suggests that our Hypothesis holds irrespective of the
Transformer configurations.
1-layer
1-layer Weight Tying
4-layer
4-layer Weight Tying
neurons
100
8.0
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
1.0
0.0
0.2
0.6
0.8
1.0 1.0
0.0
0.2
0.4 0.6
0.8
IPR
IPR
IPR
(a) Transformer with heads and embedding dimension 128.
1-layer
1-layer Weight Tying
4-layer
4-layer + Weight Tying
neurons
100-
50-
8.0
1.0
IPR
IPR
IPR
IPR
(b) Transformer with heads embedding dimension 128.
Figure 12: The IPR histograms of the Embedding layer for Transformers with different configura- tions. We see that depth does not affect the IPR distribution, whereas more heads encourages high IPR embeddings and weight tying discourages it. All models are trained with corruption fraction = 0.8 and corruption fraction { 0.2.
G.2 DROPOUT PHASE DIAGRAM FOR TRANSFORMERS
We conduct further experiments check the effect of Dropout on Transformers. From Figure 13 we see that with small dropout probability, model managed have a coexistence phase with an average higher IPR for neurons connected to the embedding layer. Notably, any dropout probability larger than 0.1 does not help.
G.3 EFFECT OF PARAMETERIZATION FOR 3-LAYER MLPS
As we mentioned Hypothesis 2.2 the input layer should have high IPR columns irrespective of the details of the network. In this section, show that different parameterization of the network does not break our Hypothesis.
In particular,
consider two different parameterizations: (i) The standard parameterization with weights sampled from (0, 2/faa_ifnnn (ii) The parameterization with weights sampled from (0 2/ 2/N2/33, which corresponds to network parameterized a more lazy training regime (Jccot al. 2022
18

---

Published as a conference paper at ICLR 2024
drepout prob. 0.0
dropout 0.05
Memorization
Coexistence 0.05)
200
corruption frac
Confusion
of neurons
Confusion
Memorizatiin
Memorization
100
50
C.E.
0.6
0.4
0.6
8.0
(a) 1-layer Transformer.
Figure 13: Modular Addition phase diagrams and IPR histograms with different dropout probabil- ities (denoted by "dp")) The IPR histograms are plotted for neurons connected to the embedding layer. We select points trained with data fraction 0.9 and corruption fraction & = 0.1.
Our results in Figure 14 show that as the network gets into lazier regime, the row-wise IPR dis- tribution of the input layer weight barely shifts, whereas the column- wise IPR distribution of the output layer weight has larger values. We interpret this shift as happening because the network with lazier parameterization fails to utilize the hidden layers to decode the periodic features embedded by input layers.
wd = 20.0
Input
Output
500
corruption frac
of neurons
400
Memorization
300
FRRt Final
Finel
200
P.I.
Full Inversion
100
0.2
0.4
0.6
0.8
8.0
0.2
0.4
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
data frac
IPR
IPR
(a) MLPs standard paramcterization.
20.0
Input
Output
500
corruption frac
of neurons
400
Memorization
300
Tinal
200
PJ.
Full Inversion
100
0.2
0.4
0.6
0.8
8.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
data frac
IPR
IPR
(b) MLPs in lazier parameterization.
Figure 14: Phase diagrams and the IPR histograms of the columns of input layer weight and rows of the output layer weights for 3-layer MLPs with different parameterizations. Selected models trained with data fraction = 0.8 and corruption fraction ę = 0.2.
G.4 BATCHNORM FOR 3-LAYER RELU MLPS
This subsection tests the correlation between BatchNorm weights and high IPR neurons for 3-layer ReLU MLPs. We found that BatchNorm for 3-layer MLPs only has coexistence phases and can perfectly memorize and generalize for small corruption fractions &, we showed representative training curve in Figure T5 We find that with BatchNorm, the 3-layer ML.P tends to have high IPR
19

---

Published as a conference paper at ICLR 2024
neurons close to the output layer instead of having high IPR for close the imput layer. We believe that failure to remove the very low IPR neurons connected to the input layer why this network is always in the coexistence phase.
500
[BatchNorm weight
0.150
Input
Output
* of neurons
400
0.125
300
0.100
200
0.075
100
0.050
8.0
0.2
0.6
0.4
0.6
0.02 0.0
0.2
IPR
IPR
IPR.
(a) 3- 3-layer ReLU MLP, trained with batch size B = 64.
Figure 15: IPR of 3-layer MLP in coexistence phase, trained with data fraction = 0.8 and corruption fraction & 0.1. We plotted row-wise IPR for input layer weight and column-wise IPR for output layer weight; the correlation between column- wise IPR of output layer weights with the absolute value of BatchNorm weights right before it. The correlation between the BatchNorm weights and IPR values we pointed out in the main text still holds. However, we did not observe strong correlation between row-wise IPR of the input layer and the BatchNorm weight next to
H LABEL CORRUPTION ON GENERAL DATASETS
We train a 4-layer ReLU MLP on the MNIST dataset with label corruption with weight decay and dropout (CrossEntropy loss). Weight decay indeed enhanses generalization and leads to Full Inversion. Very high weight decay leads confusion. Dropout, on the other hand, not as effective. Increasing dropout hurts both test and train accuracies in this case.
The generalizing features are harder to quantify for MNIST we leave this analysis for future work.
wd = 0.0
wd 20.0
wd 40.0
wd 80.0
0.8
Confusion
Confusion
corruption frac
0.6
Memorization
Confusion
0.4
Full Inversion
Full Inversion
0.2
20000
40000
20000 N train
40000
20000 train
40000
20000 N train
40000
N ttariin
(a) Weight decay
0.00
dp 0.30
dp 0.60
dp 0.90
0.8
corruption frac
Confusion
0.6
Memorization
Confusion
Confusion
0.4
Memorization
0.2
Coexistence
20000
40000
20000 N_ N_train
40000
20000
40000
20000 40000
N_ trariin
N N_taain
N_train
(b) Dropout
Figure 16: 4-layer ReLU MLP trained on MNIST dataset with label corruption.
20

---

Published as a conference paper at ICLR 2024
I ADDITIONAL PRUNING CURVES
I.1 ACCURACIES
In Figure 6 showed pruning low and high IPR neurons in the Coexistence phase. Here we show the extended version of the pruning experiments. In contrast to the plots in Figure In Figure 17 we plot longer x-axes we show accuracies with fraction of pruned neurons ranging from 0.0 1.0. We also perform the pruning experiment for Partial Inversion and Full Inversion phases. Pruning affects generalization and memorization differently in different phases because the proportion of high and low IPR neurons in trained networks depends on the phase.
Pruning Memorizing (low IPR)
Pruning Generalizing (high IPR)
100
accuracy (36
80
80
training set
test set
corrupted training data uncorrupted training data
40
10011 - 51%
20
8.0
0.8 fraction of neurons pruned (from lowest IPR)
0.2
0.4
0.6
8.0
0.2
0.6
0.8
1.0
fraction of neurons pruned (from highest (PR)
(a) Coexistence phase (weight decay = 0.0). (Left) Pruning from low IPR neurons retains performance on test data and uncorrupted training data; while decreases performance corrupted training data down 100 (1 - 4)%. (Right) Pruning high IPR neurons decreases performance on data; while retaining performance on training data. Note that accuracy on corrupted data is retained longer than on uncorrupted data.
Pruning Memorizing (low IPR) neurons
Pruning Generalizing (high IPR) neurons
100
accuracy (9))
(96
80
training set
test set
60
60
corrupted training data uncorrupted training data 10011 - 51%
20
20
8.0
8.0
fraction of neurons pruned (from lowest IPR)
fraction neurons pruned (from highest IPR)
(b) Partial Inversion phase (weight decay 5.0). (Left) Pruning from low IPR neurons retains performance on test data and uncorrupted training data; while decreases performance on corrupted training data down to 100 (1 1)% (Right) Pruning high IPR neurons decreases performance performance on corrupted training data and retains uncorrupted training data. In the absence of periodic neurons, the (memorizing) low-IPR neurons have an increased influence on network predictions, which leads to this increase in memorization.
test data. It improves the
Pruning Memorizing (low IPR) neurons
Pruning Generalizing (high IPR) neurons
accuracy (%)
80
training set
test set corrupted training data uncorrupted training data 100/1 E1%
accuracy
40
20
20
8.0
0.8
8.0
0.2
fraction of neurons pruned (from lowest (PR)
fraction of neurons pruned (from highest (PR)
(c) Full Inversion phase decay= 15.0). (Left) Pruning from low IPR ncurons retains performance on test data and uncoruupted training data; while decreases performance corrupted training data 100 (1 E)%. (Right) Pruning high IPR neurons daes not improve memorization significantly, except for the small increase accuracy on corrupted training data towards the end. This because fraction of the neurons that have low IPR is very small in this case.
Figure 17: Pruning experiments in all various phases (N 500, 0.5, 0.35). Pruning affects generalization and memorization differently in different phases because the proportion of high and low IPR neurons in trained networks depends on the phase.
21

---

Published as a conference paper at ICLR 2024
1.2 LOSSES
Here we show how the train and test losses are affected by pruning. Train loss is lowered by, both, memorizing (low IPR) and generalizing (high IPR) neurons. This leads monotonic behaviors in all the cases. Test loss is lowered by periodic neurons, while increased by memorizing neurons. An effect that sits on top of these is the effective cancellation of sub-leading terms in the analytical solution Equation 66). Since the accuracies only depend on the largest logit, these sub-leading terms have little effect on them (at sufficient width). However, losses are significantly affected by them. With zero small weight decay values, this cancellation is achieved by low IPR neurons. Thus, in these cases, low IPR neurons serve the dual purpose of memorization as well as canceling the sub- leading terms. With large weight decay values, this cancellation is achieved by a different algorithm, where "secondary peaks" appear in the Fourier spectrum of the periodic neurons. We leave an in- depth analysis of these algorithmic biases for future work.
Pruning Memorizing (low (PR) neurons
Pruning Generalizing (high IPA) neurons
0.020
MSE l00g
INSE 1699
20.015
training set
0.010
300010
0.005
0.005
0.6
0.8
1.0
fraction neurons pruned (from lowest (PR)
fraction neurons pruned (from highest IPR)
(a) Coexistence phase (weight decay = 0.0). (Left) Pruning from low neurons (almost) monotonically increases the train loss. The test loss has two competing factors: lower memorization and fewer neurons that cancel the sub-Icading terms. As a result, we see non-monotonic behaviour in the first half. (Right) Pruning high IPR neurons increases both train and test losses. We non-monotonic behaviour in the second half since the cancellation sub-Icading terms irrelevant in the absenee periodic neurons. Pruning General lizng (high (PR) neurons
Pruning Memarizing (low (PR) neurons
0.010
MSE loss
MSE loss
trail traiining set test set
0.008
0.008
0.006
0.4
0.6
0.8
1.0
fraction pruned (from (PR)
fraction pruned (from highest (PR)
(b) Partial Inversion phase (weight decay 5.0). (Left) Pruning from low IPR neurons: Initially, the removal of memorizing neurons causes the train loss increase and the test loss to decrease. In the latter half, once we start removing the periodic neurons, both train and test losses increase. (Right) Pruning from high IPR neurons: Initially, both train and test losses increase due to the removal of periodic neurons. In the latter half, the removal of the memorizing neurons causes the train and test losses to increase and decrease, respectively. Pruning Memorizing (low neurons
Pruning Generalizing (high neurons
0.010
0,010
MSE loss
trai ning set
0.009
MSE
6.0099
test set
0.008
0.0
0.2
0.4
0.5
0.8
I.0
0.0
0.2
0.4
0.6
0.8
1.0
fraction of neurons pruned (from lawest (PR)
fraction of neurons pruned (from highest (PR)
(c) Full Inversion phase (weight decay= = 15.0). (Left) Pruning from low IPR neurons monotonically increases the train loss, since is affected by removal of memorizing and periodic neurons. On the other hand, test loss briefly remains constant (during the removal of the low IPR neurons), followed by monotonic increase. (Right) Removing high IPR neurons monotonically increases both train and test losses.
Figure 18: Pruning experiments in all various phases (N 500, 0.5, 0.35). Pruming affeets generalization and memorization differently in different phases because the proportion of high and low IPR neurons in trained networks depends the phase.
22

---

Published as a conference paper at ICLR 2024
J ADDITIONAL TRAINING CURVES
Coexistence
Partial Inversion
Full Inversion
Confusion
100
100
100
100
% accuracy
50-
50
50-
train test
train test
train
test
test
0
500 steps
1000
260
400
280
40o
1000
steps
steps
steps
(a)
(b)
(c)
(d)
Figure 19: Training curves with dropout (a) dp=0.0 (b) dp=0.3 (c) dp=0.6 (d) dp=0.9
Coexistence
Partial Inversion
Full Inversion
100
100
100
%6 accuracy
75
75
50
train
train
train
25
25
25
test
0-
1000 2000
3000
4000
500
1000 steps
1500
2000
500 1000
1500 steps
2000
steps
(a)
(b)
(c)
Figure 20:
Training curves in various phases with BatchNorm. All networks are trained with cor- ruption fraction & = 0.3 and batch-size = 64; with variable data fractions. (a) Coexistence (a = 0.4)
(b) Partial Inversion (a 0.5) (c) Full Inversion (a 0.8)
Coexistence
Partial Inversion
Full Inversion
100
100
100
% accuracy
75
50-
50-
50
train
train test
train
25
test
25
25
test
3000
4000
500
1000
1500
2000
500
1500 steps
1000
2000
steps
steps
(a)
(b)
(c)
Figure 21: Training curves in various phases with LayerNorm. All networks are trained with cor- ruption fraction & = 0.3 and batch-size = 64; with variable data fractions. (a) Coexistence (a 0.4)
(b) Partial Inversion (a = 0.5) (c) Full Inversion (a = 0.8)
Coexistence
Partial Inversion
Full Inversion
100
100
accuracy
75
50
50
50
train
train
train
25
test
test
25
test
5k
1000
1500
500 1000
2000
steps
steps
steps
(a)
(b)
(c)
Figure 22: Training curves in various phases without any normalization, but with mini-batches.
All networks are trained with corruption fraction & 0.3 and batch-size 64; with variable data fractions. (a) Coexistence (a- = 0.4) (b) Partial Inversion (a 0.5) (c) Full Inversion (a = 0.8)
23

---

Published as a conference paper at ICLR 2024
K
OUTPUT LOGITS
Here we show how the output logits of trained networks look like in various phases (i.e. different amount weight decay) corrupted training examples. In all cases, we see two peaks in logits, corresponding the the corrupted and the "correct" labels. Their relative value determines the degree of memorization, and hence, the phase.
labal
0.8-
cornuptio
comrsttd labal
corrapted
conrupted label
loglt logit value
0.6
0.4
0.2
0.0
Nout 1o
99 11
A0 1/
MAMM MMli
50
75
100
25
50
75
100
0
50 75
100
25
75 100
output Index
putput_ Index
putput Index
output, Index
(a) Coexistence phase (wd=0.0). The logits corresponding to corrupted label have higher values than the correct ones, which explains the memorization of corrupted data in this phase.
label
0.3
corrupted
comppted
corrupted
comuupted
logit legit value
0.2
0.1
Ma
MM M
0.0
16
40
50
75
100
25
75
100
50
75
100
25
75
100
output Index
output
output Index
output
(b) Partial Inversion phase (wd=5.0). The logits corresponding to cornupted labels have similar values to the correct ones, which explains partial memorization of corrupted data in this phase.
0.2
connupted
comupted label
corrupted
comnupted label
logit value
0.1
W/
Am Ma
0.0
0
50
75
100
25
75
100
0
50
75
100
25
50
75
100
output_ index
output_ index
output index
outputt index
(c) Full Inversion phase (wd=15.0). The logits corresponding to corrupted labels are heavily suppressed com- pared to the correct ones, which explains the absence of memorization of corrupted data in this phase.
Figure 23: Network output logits on 4 different (corrupted) training examples examples, in various phases (N = 500, a 0.5. g = 0.35). Blue dots correspond to the logit value for the "correct" labels. Red dots correspond the logit value for the corrupted labels.
24

---

Published as a conference paper at ICLR 2024
L
DETAILED PHASE DIAGRAMS
Here we replot the phase diagrams in Figure 3 together with train and test accuracies.
wd
0.0
wd 5.0
wd 15.0
wd 20.0
1.0
corruption frac
0.8
acc
acc
acc
acc
0.6
0.8
0.4
0.6
test acc
test acc
test acc
test acc.
corruption frac
0.4
0.0
corruption frac
0.8
Memorization
Memorization
Memorization
Memorization
0.6
Coexistence
P.
0.4
Partial
Full Inversion
Inversion
Full
Forgetting
0.2
Coexistence
Inversion
0.4 data frac
0.4
0.8
data frac
data frac
data frac
Figure 24: Modular Addition with weight decay. The first two rows are train and accuracies.
dp 0.00
dp 0.30
dp 0.60
dp 0.90
1.0
corruption frac
0.8
acc
acc
acc
acc
0.6-
0.8
0.4
test acc
test acc
test acc
test acc
corruption frac
0.0
corruption frac
0.8
Memorization
Memorization
Memorization
Confusion
0.6
Coodistence
0.4
Partial
Inversion
Full
Full
Full Inversion
0.2.
Coexistence
Inversion
Inversion
0.8
0.4
0.8
0.4
0.8
data frac
data frac
data frac
data frac
Figure 25: Modular Addition with Dropout. The first two rows are train and test accuracies.
25

---

Published as a conference paper at ICLR 2024
M CHOICE OF LOSS FUNCTION
CrossEntropy loss is known to be sensitive to label noise (Ghosh et al. 2017). By adding regular- ization, we recover similar phase diagrams. The requirement for regularization is much greater for CrossEntropy loss compared to MSE. We find that CrossEntropy training is very sensitive to weight decay, with Inversion occuring for narrow band of weight decay values, beyond which find Confusion. We do not observe Forgetting phase in this case. See Figure 26 With Dropout, we find that CrossEntropy training requires very high dropout probabilities grok and/or attain Inversion.
wd = 0.0
wd = 2.0
wd 5.0
wd = 10.0
0.8
corruption frac
Memorization
0.6
Memorization
Memorization
Confusion
0.4
P.I.
Full Inversion
0.2
Coexistence
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
data frac
data frac
data
frac
data
frac
(a) Weight decay
dropout prob.
0.0
dropout prob. 0.3
dropout prob. 0.6
dropout prob. 0,9
0.8
corruption frac
Memorization
0.6
Memorization
Memorization
Memorization
0.4
PLL
Full Inversion
0.2
C.E.
Coexistence
C.E.
0.2
0.2
0.2
0.8
0.2
0.6
0.8
data frac
data frac
data frac
data frac
(b) Dropout
Figure 26: Modular Addition trained with CrossEntropy Loss. Robustness to label corruption is weaker compared to MSE.
26

---

Published as a conference paper at ICLR 2024
N CHOICE OF ACTIVATION FUNCTION
To check if the story depends on the activation function, we did the modular addition experiments for 2-layer ReLU network, see Figures 27 and 28 The requirement for regularization for grokking is higher compared to quadratic activation function. We noticed that the phase diagram for weight decay behaves similarly to quadratic activation functions, except for the disappearance of the For- getting phase. Dropout does not seem as effective for ReLU networks, with the network entering Confusion phase for even moderate dropout probabilites.
wd = 0.0
wd = 5.0
wd 15.0
wd = 50.0
0.8
corruption frac
Conf-
0.6
Memorization
usion
Memorization
Confusion
Memorization
Memorization
0.4
Full Inversion
Full Inversion
0.2
F.I.
F.I.
0.2
0.4
0.6
0.8
0
0.4
0.6
0.8
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
data frac
data frac
data
frac
data
frac
Figure 27: Phase diagram with weight decay, ReLU.
dropout prob. 0.0
dropout prob. 0.3
dropout prob. 0.6
dropout prob. 0.9
0.8
corruption frac
Memorization
Memorization
Memorization
Confusion
Confusion
Confusion
P.I.
C.E.
F.I.
data frac
data frac
data frac
data frac
Figure 28: Phase diagram with Dropout, ReLU.
27

---

Published as a conference paper at ICLR 2024
O
EFFECT OF OF WIDTH
wd = 0.0
wd = 5.0
wd 15.0
wd = 20.0
0.8
corruption frac
Memorization
Memorization
Memorization
Memorization
Pl.
0.4
Coexistence
PI.
Full
Full
Forgetting
0.2
Coexistence
Inversion
Inversion
0.6 data frac
0.8
0.2
0.4
0.6
0.8
0.2
0.4
0.6 0.8 data frac
data frac
data frac
Figure 29: Phase diagram for width N = 5000 fully connected network with quadratic activation. Regularization with weight decay. Compared N 500 network in Figure Ba), the network generalizes marginally poorly. Note: the absence of Partial Inversion phase and smaller Coexistence phase with wd 0.0. This can be attributed to the increased capacity to memorize due to abundance
of Neurons.
wd = 0.0
wd = 10.0
wd 15.0
wd 50.0
corruption frac
Memorization
Memorization
Memorization
0.6
Memorization
0.4
Coexistence
PL
Coexistence
Full Inversion
0.2
Coexistence
FI1
data frac
data frac
data frac
data frac
Figure 30: Phase diagram for width N 5000 fully connected network with ReLU activation. Regularization with weight decay. In this case, the network performance is drastically improved compared to N 500 (Figure(27) The confusion phase is largely eliminated, due to the increased capacity of the network, result of abundance of neurons.
P OTHER MODULAR ARITHMETIC DATASETS
In this section, check if the phase diagram story is unique modular addition. We repeat experiments with modular multiplication (z (mn) %p = mn mod p) and obtain almost identical
phase diagrams modular addition (see Figure 31.
5.0
wd 15.0
wd 20.0
corruption frac
Memorization
Memorization
Memorization
Memorization
0.6
Coexistence
0.4
Partial Inversion
Full Inversion
Full Inversion
Forgetting
0.2
Coexistence
0.6
0.8
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.4 0.6 data frac
0.8
data frac
data frac
data frac
Figure 31: Phase diagram for Modular Multiplication, weight decay. All other settings are identical to Figure Note that the phase diagram is identical to Modular Arithmetic case representative of the similarity the algorithms learnt by the network.
28

---

Published as a conference paper at ICLR 2024
Q HOW MANY NEURONS CAN BE PRUNED WHILE RETAINING GENERALIZATION?
We showed in the pruning experiments (Figures 6, 17) that significant proportion of neurons (start- ing from low IPR) can be pruned out without decreasing test error. In Figure 32 we show how many neurons can pruned from a trained network various widths. We train networks with widths be- tween 500 and 10000; with learning rate 0.005 and weight decay 1.0, with AdamW optimizer and MSE loss. We use corrupted modular addition data, with data fraction a 0.5 and corruption fraction ç = 0.35. We repeat the experiment 40 times, averaging over different random seeds.
% of pruned neurons
100
98
of neurons
94
70
90
60
88
10a
104
10a
10
network width (N)
network width (N)
(a) Number of neurons required after pruning (from low IPR) to retain generalization. We see sub-linear increase in the required neurons with width. shaded region denoted the standard error over 40 runs.
(b) % of neurons that can be pruned (from low IPR) while retaining generalization.
Figure 32: Pruning as many low IPR neurons as possible while retaining generalization, for various network widths (a = 0.5, & = 0.35). The curves are averaged over 40 random seeds. Note that the x-axes on both plots have logarithmic scale.
29


Reviews:
[{"review_id": "NkjLqOUj38", "reviewer": "Reviewer_kDnz", "summary": "This paper investigates the grokking phenomenon on algorithmic datasets with label corruptions.  In particular, a two-layer MLP is fitted to a modular arithmetic dataset with varying degrees of label corruptions, weight decay, and dropout rates. As a measure of generalization, the authors propose to use the inverse participation ratio, which quantifies the periodicity in the parameters. Experiments show that (i) the model can perfectly memorize the mislabeled examples while achieving near-perfect test accuracy, (ii) the model learns to \"correct\" some or all mislabeled examples depending on the level of regularization, (iii) the memorizing/generalizing neurons are identifiable.", "questions": "- Interestingly, even in the case of full inversion without batch normalization, low IPR neurons do not disappear. Why does this occur? \n- 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?\n- Multi-modality in Figure 4 alone does not imply low/high IPR neurons memorizing/generalizing as there is no evidence that low IPR neurons are indeed connected to memorization. I think this becomes clear only in the second to the last paragraph of section 2.\n- _In trained networks, a larger population of high-IPR neurons leads to better generalization_ requires citation or explanation.\n- _Choosing quadratic activation function ... makes the problem analytically solvable._ requires citation.\n- _It’s role ..._ <--- typo\n- _In the previous Section_ <--- S should be small", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is written clearly and the results are discussed well. The specific idea of training under label corruption is interesting, which seems to allow for studying memorization on the neuron level. I find some results important (e.g. simultaneous memorization and generalization and sub-network identification) while some are not as significant (e.g. the impact of weight decay).", "weaknesses": "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.\n\n- 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.\n\n- 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.\n\n- 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pnFGjrLCdB", "reviewer": "Reviewer_j4fN", "summary": "Grokking is a recent empirical phenomenon where training achieves zero error long before neural network models generalize. People have identified this phenomenon to explore memorization vs. generalization and acquire mechanistic understanding of training dynamics of neural networks.\n\nIn this paper, the authors present randomization experiments where a fraction of labels are corrupted, i.e., replaced by random labels. A simple arithmetic task, namely modular addition, is used to train a small 2-MLP network. According to levels of memorization and generalization, the trained neural networks are categorized into several regimes: (1) memorization (2) coexistence (3) partial inversion (4) full inversion (5) forgetting. Each regime corresponds to an explanation of the inner mechanism.", "questions": "See the above sections.\n\nIn addition, is there any scenario where we believe Partial Inversion or Full Inversion may be observed in practice?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "Overall, I find this paper very clear and interesting. The authors focus on a recent phenomenon, and by experimenting under idealized data and networks, phase transitions are identified, each with meaningful interpretations.\n\nSome of the key contributions include:\n1. Randomization experiments reminiscent of [1], which helps understand generalization of neural nets in the presence of label noise.\n2. A clear phase transition phenomenon and separates different data generating regimes into interpretable regions. This may be helpful for understanding the inner workings of neural networks and training dynamics.\n3. Clear measurements that quantify the behavior of the trained network.\n4. Some insights into regularization techniques such as batchNorm and weight decay.\n\n[1] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. In International Conference on Learning Rep- resentations, 2017.", "weaknesses": "This paper follows a line of papers that study grokking, and provides a refined analysis of the phenomenon instead of proposing more general principles. \n\n1. I feel that compared with the initial phenomenon identified in earlier papers, this paper is in a way less innovative.\n2. The scope of the analysis is limited, as toy datasets and toy neural nets are studied. It is unclear whether these analyses can generalize to practical settings.\n3. This paper does not contain a detailed study of the training dynamics, how the weights behave, or some theoretical analyses. It is unclear whether this paper has impact on the theory community.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qW9xBqp5YJ", "reviewer": "Reviewer_4cwR", "summary": "This paper studies the interplay between memorization and generalization in a controlled setting of synthetic task, where analytical solution exists and can be found by neural network optimizers. It demonstrates interesting behavior where generalization and memorization co-exists, and identified subnetworks that are responsible for each behavior. It further studies how popular regularization techniques impact the interplay between memorization and generalization.", "questions": "1. 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.)?\n\n2. 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?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This paper presented a simple synthetic machine learning setting where memorization and generalization can be studied in a very controlled manner. It further demonstrate that memorization and generalization could co-exist, and when it happens, there are clear distinction of subnetworks (at neuron level) that are responsible for each aspect. Moreover, by using the inverse participation ration (IPR), those neurons can be easily identified. This allows the authors to conduct further analysis on how different regularization techniques impact generalization and memorization from the perspective of memorization neurons. For example, it is shown that batch normalization operates in a very different manner from dropout and weight decay. Although the task of study is completely synthetic, it is still valuable to have such a controlled task that demonstrate interesting neural network behaviors such as grokking and memorization/generalization.", "weaknesses": "While some of the results are interesting, this paper is not strong in novelty and depth.\n\n1. The main synthetic task studied in this paper is from previous papers, including the construction of the analytical solution and the matching neural network architecture that could realize such solutions. The main technique to distinguish memorization and generalization neurons, the inverse participation ration (IPR), is also from previous literature. While this paper does has some interesting observations when analyzing the memorization-generalization behaviors, I think it could benefit significantly from more novel analytical or algorithmic contributions.\n\n2. 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.\n\n3. The experiments could be improved by more in-depth analysis. See below for a few examples.\n\n    1. The paper shows regularizers help improve the generalization and demonstrated the \"inversion\" behaviors. However, given the emphasis on \"grok\" from the motivations, I would expect more in-depth analysis and insights to how those regularizers impact the *learning dynamics* that lead to the end result of \"inversion\" or \"coexistence\".\n\n    2. 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? \n\n    3. The paper studies the two-layer neural networks with quadratic activation function because analytical solution could be admitted with this architecture. However, most of the empirical studies (e.g. the phase diagrams) does not rely on this property. I think it is very helpful to see experiments with more realistic deep neural networks and diverse neural network architectures to see if they demonstrate consistent behaviors, and if not, to analyze what contribute to the different behaviors. There are some experiments in the appendix but only with small variations of the activation function or the width, and without in-depth analysis.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)