Progress measures for grokking via mechanistic interpretability
Authors: Neel Nanda
Lawrence Chan
Tom Lieberum
Jess Smith
Jacob Steinhardt
Venue: ICLR 2023 notable top 25% (2023)
Keywords: interpretability
grokking
progress measures
mechanistic interpretability
circuits

Abstract: Neural networks often exhibit emergent behavior in which qualitatively new capabilities that arise from scaling up the number of parameters, training data, or even the number of steps. One approach to understanding emergence is to find the continuous \textit{progress measures} that underlie the seemingly discontinuous qualitative changes. In this work, we argue that progress measures can be found via mechanistic interpretability---that is, by reverse engineering learned models into components and measuring the progress of each component over the course of training. As a case study, we study small transformers trained on a modular arithmetic tasks with emergent grokking behavior. We fully reverse engineer the algorithm learned by these networks, which uses discrete fourier transforms and trigonometric identities to convert addition to rotation about a circle. After confirming the algorithm via ablation, we then use our understanding of the algorithm to define progress measures that precede the grokking phase transition on this task. We see our result as demonstrating both that it is possible to fully reverse engineer trained networks, and that doing so can be invaluable to understanding their training dynamics. 

Published as a conference paper at ICLR 2023
PROGRESS MEASURES FOR GROKKING VIA MECHANISTIC INTERPRETABILITY
Neel Nanda"
Lawrence Chan
Tom Lieberum
Jess Smith
Jacob Steinhardt
ABSTRACT
Neural networks often exhibit emergent behavior, where qualitatively new capa- bilities arise from scaling up the amount of parameters, training data, training steps. One approach to understanding emergence is to find continuous progress measures that underlie the seemingly discontinuous qualitative changes. We ar- gue that progress measures can be found via mechanistic interpretability: reverse- engineering learned behaviors into their individual components, As a case study, we investigate the recently-discovered phenomenon of "grokking" exhibited by small transformers trained on modular addition tasks. We fully reverse engineer the algorithm learned by these networks, which uses discrete Fourier transforms and trigonometric identities to convert addition to rotation about a circle. We confirm the algorithm by analyzing the activations and weights and by perform- ing ablations in Fourier space. Based on this understanding, we define progress measures that allow three continuous phases: memorization, circuit formation, and cleanup. Our re- sults show that grokking. rather than being sudden shift, arises from the gradual amplification of structured mechanisms encoded in the weights, followed by the later removal of memorizing components.
study the dynamics of training and split training into
INTRODUCTION
Neural networks often exhibit emergent behavior, in which qualitatively new capabilities arise from scaling up the model size, training data, or number of training steps (Steinhardt, 2022; Wei et al., 2022a). This has led a number of breakthroughs, via capabilities such as in-context learning (Rad- ford et al., 2019; Brown et al., 2020) and Mhain-of-though prompting (Wei et al., 2022b). However, also poses risks: Pan et (2022) show that scaling up the parameter count of models by as little as 30% can lead to emergent reward hacking.
Emergence is most surprising when abrupt, as in the case of reward hacking, chain-of-thought reasoning, other phase transitions (Ganguli et al., 2022; Wei et al., 2022a). We could better understand and predict these phase transitions by finding hidden progress measures (Barak et al., 2022): metrics that precede and causally linked the phase transition, and which vary more smoothly. For example, Wei et al. (2022a) show that while large language models show abrupt jumps in their performance on many benchmarks, their cross-entropy loss decreases smoothly with model scale. However, cross-entropy does not explain why the phase changes happen.
In this work, introduce a different approach uncovering hidden progress measures: via mech- anistic explanations." A mechanistic explanation aims to reverse engineer the mechanisms of the network, generally by identifying the circuits (Cammarata al., 2020; Elhage al., 2021) within a model that implement a behavior. Using such explanations, we study grokking, where models abruptly transition a generalizing solution after large number of training steps, despite initially overfitting (Power et al., 2022). Specifically, we study modular addition, where a model takes inputs a, E (0...., P 1} for some prime and predicts their sum e mod Small transformers trained with weight decay this task consistently exhibit grokking (Figure 2, Appendix C.2).
"Corresponding author, please direct correspondence to: neelnanda211 @0ai1 com "Iddepeddent researcher.
University of California, Berkeley.
Interacttvve versions of figures, as well the code to reproduce our results, are available at bit. 1y/ progress measurss- -rooki seiikkiinn--bsiiiie

---

Published as a conference paper at ICLR 2023
Logits
100 sn b-
Computes logits using further trig identities:
Unimbed
Logit(e) o oos cos(wra + -ci)
= cosfinto + 6)) coe(ure) + sin(woj + 6)) sin(we)
MLP
Calculates sine and cosine using trig identitics: sin(wi(a bj) sin(wa) cos( (ub) cos(trn) sin(wb)
cos(1e(a bj) cos(ura] cos(wr) sin(wa))sinie sin(wb)
Ra
Translutes oon-hot b to Paurier hasis: a + wa), cus(une)
Emahed
sin(wb), coo(zob)
TokeE
Figure 1: The algorithm implemented by the one-layer transformer for modular addition. Given two numbers and b, the model projects each point onto corresponding rotation using its embedding matrix. Using its attention and MLP layers, it then composes the rotations to get a representation of + b mod P. Finally, it "reads off" the logits for each E {0, 1,.,., P 1]. by rotating by to get (w + - c)), which is maximized when a + b = mod P (since w is a multiple of 2m
We reverse engineer the weights of these transformers and find that they perform this task by map- ping the inputs onto circle and performing addition the circle. Specifically, we show that the embedding matrix maps the inputs a, b to sines and cosines at a sparse set key frequencies wk- The attention and MLP layers then combine these using trigonometric identities compute the sine and cosine of wa (a + b), and the output matrices shift and combine these frequencies.
We confirm this understanding with four lines of evidence (Section 4): (1) the network weights and activations exhibit a consistent periodic structure; (2) the neuron-logit map W, is well approx- imated by sum of sinusoidal functions of the key frequencies, and projecting the MLP activations onto these sinusoidal functions lets us "read off"" trigonometric identities from the neurons; (3) the attention heads and MLP neuron are well approximated by degree-2 polynomials of trigonometric functions of single frequency; and (4) ablating key frequencies used by the model reduces perfor- mance to chance, while ablating the other 95% of frequencies slightly improves performance.
Using our understanding of the learned algorithm, we construct two progress measures for the mod- ular addition task-rest reettritttteed loss, where we ablate every non-key frequency, and excluded loss, where we instead ablate all key frequencies. Both metrics improve continuously prior to when grokking occurs. We use these metrics to understand the training dynamics underlying grokking and find that training can be split into three phases: memorization of the training data; circuit for- mation, where the network learns a mechanism that generalizes; and cleanup. where weight decay removes the memorization components. Surprisingly. the sudden transition to perfect test accuracy in grokking occurs during cleanup, after the generalizing mechanism is learned. These results show that grokking, rather than being sudden shift, arises from the gradual amplification of structured mechanisms encoded in the weights, followed by the later removal of memorizing components.
2 RELATED WORK
Phase Changes. Recent papers have observed that neural networks quickly develop novel quali- tative behaviors they are scaled trained longer (Ganguli et al., 2022: Wei et al., 2022a). McGrath et (2021) find that AlphaZero quickly learns many human chess concepts between 10k and 30k training steps and reinvents human opening theory between 25k and 60k training steps.
Grokking, Grokking was first reported in Power et al. (2022), which trained two-layer transformers on several algorithmic tasks and found that test accuracy often increased sharply long after achieving perfect train accuracy. Millidge (2022) suggests that this may due to SGD being a random walk on the optimal manifold. Our results echo Barak et al. (2022) in showing that the network instead makes continuous progress toward the generalizing algorithm. Liu et al. (2022) construct small examples of grokking, which they to compute phase diagrams with four separate "phases" of learning. Thilak et al. (2022) argue that grokking can arise without explicit regularization, from an optimization anomaly they dub the slingshot mechanism, which may act implicit regularizer.
2

---

Published as a conference paper at ICLR 2023
Average Triin Loss
0.8
Accuracy
Log Loss
0.6
0.01
0.2
Avenge Thain Accuracy Averae Accuracy
15k.
Epoch
Epoch
Figure 2: The train and test accuracy (left) and train and test loss (right) of one-layer transformers on the modular addition task described in Section 3, over random seeds. These models consistently exhibit grokking: they quickly overfit early on in training, but then later learn to generalize.
Circuits-style mechanistic interpretability. The style of post-hoc mechanistic interpretability in Section 4 is heavily inspired by the Circuits approach of Cammarata et al. (2020), Elhage et al. (2021), and Olsson et al. (2022).
Progress measures. Barak et al. (2022) introduce the notion of progress measures-mettrcs that improve smoothly and that precede emergent behavior. They prove theoretically that training would amplify a certain mechanism and heuristically define a progress measure, In contrast, we use mech- anistic intepretability to discover progress measures empirically.
SETUP AND BACKGROUND
We train transformers to perform addition mod P. The input to the model is of the form "a =", where a and b are encoded as P'ddimensional one-hot vectors, and = is a special token above which one-layer ReLU
read the output In our mainline experiment, we take P 113 and transformer, token embeddings with d = 128, learned positional embeddings, attention heads of dimension d/4 32, and n 512 hidden units the MLP. In other experiments, we vary the depth and dimension of the model. We did not use LayerNorm or tie our embed/unembed matrices.
Our mainline dataset consists of 30% of the entire set of possible inputs (that is, 30% of the 113- 113 pairs of numbers mod P). We use full batch gradient descent using the Adam W optimizer (Loshchilov & Hutter, 2017) with learning rate y 0.001 and weight decay parameter À We perform 40, 000 epochs of training, As there are only 113 113 possible pairs, we evaluate test loss and accuracy on all pairs inputs not used for training.
Networks trained on this task consistently exhibit grokking. As Figure 2 shows, our networks first overfit the training set: train accuracy quickly converges to 100% and the train loss quickly declines, while the test accuracy remains low and the test loss remains high. After around 10, 000 epochs, the network generalizes and test accuracy increases to near 100%. In robustness experiments, we confirm that grokking consistently occurs for other architectures and prime moduli (Appendix C.2). In Section 5.3 we find that grokking does not occur without regularization.
To describe transformer components, follow the conventions and notations laid out in Elhage et al. (2021). We focus on the x embedding matrix WF, the d n output matrix of the MLP layer Wout, and the P x d unembedding matrix Wo.² Let Logits(a, b) denote the logit vector inputs a, b. and ML, b) denote the MLP activations. Empirically, our networks do not significantly use the skip connection around the MLP (Appendix A. so Logits(a,
Wo W our MMP 6). We
therefore also study the P x n neuron-logit map W = Wy Wout-
3.1 THE FOURIER MULTIPLICATION ALGORITHM
We claim that the learned networks use the following algorithm (Figure 1):
Given two one-hot encoded tokens a, map these to sin(waa), cos (owaa) sin(wa)), and 2kg
b))( using the embedding matrix, for various frequencies wx = P N.
We ignore the embedding and unemhedding of the "=" token for simplicity.
3

---

Published as a conference paper at ICLR 2023
Compute cos (wx (a b)) and sin (wx b)) using the trigonometric identities: cos (wx (a + b)) cos (wxa) cos (wya) - sin (wxa) sin (wxb)
sin (wx(a 6)) sin(w cos (wxb6) + cos (wxa) sin (wvbb)
In networks, this is computed in the attention and MLP layers.
For each output logit c, compute cos (wx (a b- c)) using the trigonometrie identity:
cos (wx(a + - c)) = cos (uxx(a + b)) cos (wyc) + sin (wx((a + b)) sin (wxe). (1)
This is linear function of the already-computed values cos(w,, b)), sin(w, b)) and is implemented in the product of the output and unembedding matrices Wz.
The unembedding matrix also adds together cos (wx(a + b c)) for the various ks. This causes the cosine waves to constructively interfere at e" = a + b mod (giving e" a large logit), and destructively interfere everywhere else (thus giving small logits to other es).
We refer to this algorithm as Fourier multiplication, and will justify our claim in detail in Section 4.
4 REVERSE ENGINEERING A ONE-LAYER TRANSFORMER
In this section, describe four lines of evidence that our transformers are using the Fourier mul- tiplication algorithm described in Section 3.1. Here we apply our analysis to the mainline model from Section 3; the results are broadly consistent for other models, including across different num- ber of layers, different fractions of the training data, and different prime moduli (see Appendix C.2, especially Table 5).
Our first line of evidence involves examining the network weights and activations and observing consistent periodic structure that is unlikely to occur by chance (Section 4. 4.1). Moreover, when we take Fourier transforms, many components are either sparse or nearly sparse in the Fourier domain, supported on a handful of key frequencies.
We next look into the actual mechanisms implemented in the model weights (Section 4.2). We show that the unembedding matrix Wt is (approximately) rank 10, where each direction corresponds to the cosine sine of one of 5 key frequencies. Projecting the MLP activations onto the components of Wz approximately produces multiples of the functions cos (w + b)) and sin (wy(a (a + 6)), showing that the MLP layer does compute these sums.
To better understand the mechanism, we zoom in to individual neurons (Section 4.3). We find that the attention heads and most neurons well-approximated by 2 polynomials of sines and cosines at a single frequency. Moreover, the corresponding direction in W,. also contains only that frequency. This suggests that the model's computations are (1) localized across frequencies and (2) mostly aligned with the neuron basis.
Finally, we use ablations confirm that our interpretation is faithful (Section 4.4). We replace various components of the model by the components of the Fourier multiplication algorithm and find that doing so consistently does not harm and sometimes even improves model performance.
4.1 SUGGESTIVE EVIDENCE: SURPRISING PERIODICITY
The first line of evidence that the network is using the algorithm described in Section 3.1 is the surprising periodicity in the activations of the transformer. That is, the output of every part of the network is periodic as function of the input tokens.
Periodicity in the embeddings. We start by examining the embeddings. We apply a Fourier trans- form along the input dimension of the embedding matrix WE then compute the f2-norm along the other dimension; results are shown in Figure 3. We plot only the components for the first 56 frequen- cies, the norm of the components for frequencies k and P - k are symmetric. The embedding matrix WE e is sparse in the Fourier basis-it only has significant nonnegligible norm at 6 frequencics. Of these frequencies, only 5 appear to be used significantly in later parts of the model (corresponding to & E (14, 35, 41, 42, 52)). ). We dub these the key frequencies of the model.
Periodicity in attention heads and MLP neuron activations. This periodic structure recurs throughout the network. As an example, we plot the attention weight at position 0 for every combi- nation of two inputs for head 0 in Figure 4. The attention exhibits periodic structure with frequency
4

---

Published as a conference paper at ICLR 2023
Norn of Fourier Component
Fourier Components of Embedding Matrix
Norm of Faurier Component
Fourier Components of Neuron-L.ggi Map
OB
1.5
10
20 30 43
20
40
50
Frequency
Frequency
Figure 3: (Left) The norms of the Fourier components in the embedding matrix WE. discussed in Section 4.1, the sparsity of WE in the Fourier basis is evidence that the network is operating in this basis. Of the six non-zero frequencies, five "key frequencies" appear later parts of the network, corresponding to k E (14, 35, 41, 42, 52). (Right) Norm of Fourier components of the neuron-logit map W. Fourier transform is taken over the logit axis, and then the norm is taken over the neuron axis. As discussed in Section 4.2, WL is well-approximated by the 5 key frequencies Wi.
Attention Score for Head 0
Activations for Neuron 0
Norms I Lngits Fourier
<oc
60
0.8
3.5
(14:1
332.
3
0.6
2.5
60
59
c0444
0.2 100
100
0.5
Component
Figure 4: (Left) The attention score for head 0 from the token '=' 'a', as a function of inputs a. b. (Center) The activations of MLP neuron 0 given inputs b. Both the attention scores and the neuron activations are periodic (Section 4.1). (Right) The norm the Fourier components of the logits (2D Fourier transform is taken over the inputs a, b. and then norm is taken over the logit axis). There are 20 significant components corresponding to the 5 key frequencies (Section 4.1).
k = 35. In Figure 4, we also plot the activations of MLP neuron 0 for every combination of inputs. The activations are periodic with frequency h 42. We see similar patterns for other attention heads and MI.P neurons (Appendix C.1).
Periodicity logits. Finally, the logits are also periodic. Figure 4, we represent the logits in the 2D Fourier basis over the inputs, then take the f--norm over the output dimension. There are only twenty components with significant norm, corresponding to the products of sines and cosines for the five key frequencies wk. These show up as five x 2 blocks in Figure 4.
4.2 MECHANISTIC EVIDENCE: COMPOSING MODEL WEIGHTS
We now demonstrate that the model implements the trigonometric identity (1) as follows: the func- tions cos (wa(c (a b)). sin (wx (a + b)) linearly represented in the MLP activations, and the un- embed matrix reads these linear directions and multiplies them by cos (wve), sin (wxc) respectively.
We will do this in two steps. First, we show that W,, (the matrix mapping MLP activations logits) is (approximately) rank 10 and can be well approximated as:
Wz Ex cos (wx) 2h + sin (wk) of
for some uk, Uh E R512, where cos (w*), sin (wę) IR 113 are vectors whose cth entry is cos (wyc) and sin (wgc). Second, note that our model implements the logits for a, b as:
Logits(a, b) WrMLP(a, RJ Σg cos (wx) u. MLP(a, b) - sin (ww) of MLP(a, b)
(3)
We check empirically that the terms u. MLP(aa b) and of MLP((,, b) approximate multiples of cos (a + b)) and sin (wx (a + b)) (> 90% of variance explained). Thus the network computes trigonometric functions in the ML. and reads them off claimed. As sanity check, confirm
5

---

Published as a conference paper at ICLR 2023
WE Component
Fourier components u MLP((, b) or oi. (M)
FVE
cas (w)cc)
41.6 cas( (0) cos( [114 13.6sin(w1ad (asaddssnnnn sin[w]ab] R/ 44.1 cos [1 6))
93.2%
sin (w1ac)
14.1 sin(w144) cos(w0s0 + 14.1 cos( sin (1ab) (b) 8 14.1 sin (w]a(a b)]
93.5%
(
40.7 cos( (sssa) cos ussbb 43.66 133.6innwassd sin wasb)
2: 12.2 cas (was [aa((a+b))
96.8%
(33
41.8 sin(Wasee cos (3sbb) + +1.8 cos (Ward) sin (anbb)
41.8 sin (War(a (war(a
96.5%
cos [W41
44.8 cos( (aa1d) cos( Win b) - -44.8 sin(wara) sin Wax b) = 44.8 cos (Wix (a
97.0%
sin (W
44.5 sin( (Wara) cos Wan b) + +4.55 cos(War))))))) sin Wa1
== 44.5 sin (wan(a
97.0%
oos (M)zec
64.6 cos( (122a) cos Mazbb - -68.5 (inWaaa))
m22bb)
x 66.66 cos
(#aa(a
96.4%
sin (Wa2c)
67.8 sin (W22aa cos (w22b)
+ 67.8 cos(1)220) (1ng
2 67.8 sin (w
96.4%
cos ()
60.5 cos( (122)) cos Ms2bb) 65.5 365.5innuun (35(22 sin (s2bb)
2 63.0 cos (Mss (a b))
97.4%
sin (wsac)
64.5sin (1224) cos ()2bb) + 64.5 cus (Waz4) sin (in(w)))) 2 64 (w52 6))
98.2%
Table 1: For each of the directions dy or uk (corresponding to the cos(urs) and sin(wx) components respectively) in the unembedding matrix, we take the dot product of the MLP activations with that direction, then perform a Fourier transform (middle column; only two largest coefficients shown). We then compute the fraction of variance explained (FVE) if we replace the projection with single term proportional to cos (wx (a + b)) or sin (wx (a + b)). and find that it is consistently close to I.
that the logits are indeed well-approximated by terms of the form cos (wx(a + b-c)--c) variance explained).
c)) (95% of
Wr. is well approximated by cos (wac) and sin (wge). We perform a discrete Fourier transform (DFT) on the logit axis of W1 and look at the 10 directions uR. U2 corresponding to sin (wx) and cos (wx). When we approximate WL with ExE[14,35, 11,1444, cos (wx) ux +sin (wx) og. the resid- ual has Frobenius norm that is under 0.55% of the norm of Wz. This shows that W. is well ap- proximated by the 10 directions corresponding frequencies. We also plot the norms of each direction in Figure 3, and find that no Fourier compo- nent outside the key frequencies has significant norm.
cos (wx) and sin (wx) for each of the five key
The unembedding matrix "reads off" terms of the form cos (wxx + b)) and sin (wx + b)) from the MLP neurons. Next, the directions ug, for k E {14, 35, 41, 42, 52]. Table displays the results: the dot products ui MLP(a, b) and of MLP((a (a, b) are well approximated by a multiple of terms of the form cos (wa (a + b)) = cos (waa) cos (wyb) sin (wxa) sin (wxb). and sin (wx + b)) = sin (wxa) cos (wngb) + cos (waa) sin (wxb).
take the dot product of the MLP activations with each of
That is, for each key frequency k. ug. and us are linear directions in the space of MLP neuron activations that represent cos (wy((a + b)) and sin (wx(a + b)).
Logits are well approximated by weighted sum of cos (wx (a - b- c))s. We approximate the output logits as the sum Ee Og cos(w)(a++b-a
for & E {14, 35, $35.41,44 41, 52} and fit the coefficients 0% via ordinary least squares. This approximation explains 95% of the variance in the original logits. This is surprising- the output logits are a 113 113 113 dimensional vector, but are well- approximated with just the 5 directions predicted by our interpretation. If we evaluate test loss using this logit approximation, we actually see an improvement in loss, from 2.4 - 10-7 to 4.7 10-8
Taken together, these results confirm that the model computes sums of terms of the form cos (ww(a- + b - c)) = cos (wy (a + b)) cos (wyc) + sin (wy(a + b)) sin (wye).
4.3 ZOOMING IN: APPROXIMATING NEURONS WITH SINES AND COSINES
In the previous section, we showed how the model computes its final logits by using Wz to "read off" trigonometric identities represented in the MLP neurons. We now examine the attention heads and MLP neurons understand how the identities come to be represented at the MLP layer. In Appendix C.1.2, we show that two of the attention heads approximately compute degree-2 poly- nomials of sines and cosines of particular frequency (and the other two are used to increase the magnitude of the input embeddings in the residual stream). Here, show that most neurons also apll-appro by degree-2 polynomials, and the map from neurons to logits is localized by frequency.
Most MLP neurons approximately compute a degree-2 polynomial of a single frequency. We next try approximate the activations of each MLP neuron by a degree-2 polynomial of one of the 5 key frequencies. As shown in Figure 5, out of 512 total neurons, 433 (84.6%) have over 85% of their variance explained with single frequency.

---

Published as a conference paper at ICLR 2023
FVE degree-2 polynomials
Components of W,. corresponding to freq 14 neurons
Number of neurons
900
10
Neuron
200
20
30
100
40
0.8
Qus
ZI
BE
12005
Fraction variance explained
68
Figure 5: (Left) Most neurons well-approximmted by degree-2 polynomials of single frequency.
(Right) A heatmap showing weights in Wt corresponding to each of the 44 neurons of frequency
14. The non-trivial comnonents corresnond to sin (mm.)) and ens (in.)) for ke 14
Loss after ablating frequencies
Other frequency
Key frequency
Log Loss
100g
1p
10
15
20
25
30
35
43
55
Restricted
Original
Figure 6: The loss of the transformer (lower=better) when ablating each frequency k E {1, 2,... 56) and everything except for the five key frequencies (restricted loss). We include the original unablated loss for reference. Ablating key frequencies causes performance drop, while the other ablations do not harm performance.
Maps to the logits are localized by frequency. We partition these 433 neurons by the frequencies with the highest variance explained. For each resulting subset, the map WL from neurons to logits has only non-trivial components, corresponding sine and cosine at that frequency. For ple, in Figure 5 we plot the 44 columns of WL corresponding to the 44 neurons in the k: = 14 cluster and find that the only non-negligible components sin 22kz and cos
2km for k 14.
4.4 CORRECTNESS CHECKS: ABLATIONS
In previous sections, we showed that various components of the model were well-approximated by sparse combinations of sines and cosines. We verify that these approximations are faithful to the model's functionality, by replacing each component with its approximation. This generally does not hurt the performance of the model and in some cases improves it.
MLP neurons. In Section 4.3, we identified 433 neurons that were well-approximated by a degree-2 polynomial. We replace each of these neurons' activation value by the corresponding polynomial, leaving the other neurons untouched. This increases loss by only 3% in relative terms (from 2.41 10-7 to 2.48 10-7) and no effect on accuracy.
We can instead apply stricter ablation to the MIP layer and restrict each neuron's activation to just the components of the polynomial corresponding to terms of the form cos(wr (a + b)) and sin (wx(a b)) in the key frequencies. This improves loss by 77% (to 5.54 10-8), validating that the logits are calculated by trig identities of neurons as detailed in Section 4.2.
Logit frequencies. Next, we ablate various components of the final logits in the Fourier space. To do so, we take 2D DFT on the 113 113 113 logit matrix over all 113 113 pairs of inputs to get the logits in the Fourier basis, then set various frequencies in this basis to 0.
We begin by ablating the components corresponding to each of the key frequencies. As reported in Figure 6, ablating any key frequency causes a significant increase in loss. This confirms that the five frequencies identified previous sections are indeed necessary components of the transformer. In contrast, ablating other frequencies does not hurt the model at all.
We then ablate all 113 113 40 the Fourier components besides key frequencies; this ablation actually improves performance (loss drops 70% to 7.24 10-8).

---

Published as a conference paper at ICLR 2023
Excluded Loss over Frequencies
Restricted Loss
1,055
00.11
Loss
0.01
Train
Test Loss
Test loss
Excludae
100g
Rostrictal loss
10p
30k
36
Epoch
Epoch
Gini Coefficients of Embed Matrix and Neuron-logit Map
Toral Sum of Squared Weights
Sum of Squared Weights
3500
Gini coefficient
08
3000
Wi P2
2500
2003
0.2
1900
1000
15ku
20k
25k
30K
54
10k.
13k
20k
25k
30K
Epoch
Epoch
Figure 7: How each of the progress measures in Section 5.1 changes the course of training. The lines delineate the 3 phases of training: memorization, circuit formation, and cleanup (and a final stable phase). (Top Left) Excluded loss increases during circuit formation, while train and test loss remain flat. (Top Right) The restricted loss begins declining before test loss declines, but has an inflection point when grokking begins to occur. (Bottom Left) The Gini coefficient of the norms of the Fourier components of WE and Wy increase sharply during cleanup. (Bottom Right) The sums of squared weights decreases smoothly during circuit formation and more sharply during cleanup, indicating that both phases are linked to weight decay.
Directions in W.. In Section 4.2, found that W, is well approximated by the 10 directions corresponding to the cosine and sine of key frequencies. we project the MLP activations to these 10 directions, loss decreases 50% to 1.19 10 If we instead projected the MLP activations onto the nullspace of these 10 directions, loss increases to 5.27-word than uniform. This suggests that the network achieves low loss using these and only these 10 directions.
5 UNDERSTANDING GROKKING BEHAVIOR USING PROGRESS MEASURES
We now use our mechanistic understanding of the network define two progress measures: metrics that can be computed during training that track the progress of the model over the course of training, including during phase transitions. This allows us study how the network reaches its final solution.
5.1 PROGRESS MEASURES
We translate the ablations in Section 4.4 into two progress measures: restricted and excluded loss.
Restricted loss. Since the final network uses sparse of frequencies wk. makes sense to check how well intermediate versions of the model can do using only those frequencies. To measure this, we perform 2D DFT on the logits to write them as linear combination of waves in and b, and set all terms besides the constant term and the 20 terms corresponding to cos(wy((C + b)) and sin (wr b)) for the five key frequencies to We then measure the loss the ablated network.
Excluded loss. Instead of keeping the important frequencies WK, we next remove only those key frequencies from the logits but keep the rest. We measure this on the training data to track how much of the performance comes from Fourier multiplication versus memorization. The idea is that the memorizing solution should be spread out in the Fourier domain, so that ablating a few directions will leave it mostly unaffected, while the generalizing solution will be hurt significantly.
Beyond these, we will also measure (1) the Gini coefficient (Hurley & Rickard, 2009) of the norms of the Fourier components of WE and Wz, which measures the sparsity of We and W2 in the Fourier basis, and (2) the norm the weights during training, since weight decay should push these down once the train loss is near zero.
8

---

Published as a conference paper at ICLR 2023
5.2 PHASES OF GROKKING: MEMORIZATION, CIRCUIT FORMATION, AND CLEANUP
Using the mainline model from Section 4, we plot the excluded loss, restricted loss, Gini coefficient of the matrices Wy and Wt, and sum of squared weights in Figure 7. We find that training splits into three phases, which we call the memorization, circuit formation, and cleanup phases. show similar results for other models in Appendix C.2.)
Memorization (Epochs 0k-1.4k). We first observe a decline of both excluded and train loss, with test and restricted loss both remaining high and the Gini coefficient staying relatively flat. In other words, the model memorizes the data, and the frequencies wy used by the final model are unused.
Circuit formation (Epochs 1.4k-9.4K3 9.4). In this phase, excluded loss rises, sum of squared weights falls, restricted loss starts to fall, and test and train loss stay flat. This suggests that the model's behavior the train set transitions smoothly from the memorizing solution to the Fourier multi- plication algorithm. The fall in the sum of squared weights suggests that circuit formation likely happens due weight decay. Notably, the circuit is formed well before grokking occurs.
Cleanup (Epochs 9.4k-14k). In this phase, excluded loss plateaus, restricted loss continues to drop, test loss suddenly drops, and sum of squared weights sharply drops. As the completed Fourier multiplication circuit both solves the task well and has lower weight than the memorization circuit, weight decay encourages the network to shed the memorized solution in favor of focusing on the Fourier multúplication circuit. This is most cleanly shown in the sharp increase in the Gini coefficient for the matices WE and WL, L, which shows that the network is becoming sparser in the Fourier basis.
5.3 GROKKING AND WEIGHT DECAY
In the previous section, we saw that each phase of grokking corresponded to an inflection point the l--norm of the weights. This suggests that weight decay is important component of grokking and drives progress towards the generalizing solution. In Appendix D.I, we provide additional evidence that weight decay is necessary for grokking: smaller amounts of weight decay causes the network to take significantly longer to grok (echoing the results on toy models from Liu et al. (2022)), and networks do not grok the modular arithmetic task without weight decay some other form of regularization. In Appendix C.2, we also find that the amount of data affects grokking: when networks are provided with enough data, there is no longer a gap between the train and test losses (instead, both decline sharply some number of epochs into training). Finally, in Appendix D.3 we replicate these results on several additional algorithmic tasks.
6 CONCLUSION AND DISCUSSION
In this work, trained on a modular addition task. We find that the transformers embed the input onto rotations in 2 and compose the rotations using trigonometric identities compute + b mod 113. Using our engereereegeee algorithm, we define two progress measures, along which the network makes continuous progress toward the final algorithm prior to the grokking phase change. We see this work as proof of concept for using mechanistic interpretability to understand emergent behavior.
mechanistic interpretability define progress measures for small transformers
Larger models and realistic tasks. In this work, we studied the behavior of small transformers on simple algorithmic task, solved with single circuit. On the other hand, larger models use larger, more numerous circuits to solve significantly harder tasks (Cammarata et al., 2020; Wang et al., 2022). The analysis reported in this work required significant amounts of manual effort, and our progress metries are specific to small networks on one particular algorithmic task. Methods for automating the analysis and finding task-independent progress measures seem necessary scale to other, larger models. We discuss possible scenarios for more realistic applications in Appendix F.
Discovering phase change thresholds. While the progress measures we defined in Section 5.1 in- crease relatively smoothly before the phase transition (and suffice to allow us to understand grokking for this task) we lack a general notion of criticality that would allow to predict when the phase transition will happen ex ante. Future work should develop theory and practice in order to apply progress measures to predict the timing of emergent behavior.

---

Published as a conference paper at ICLR 2023
REPRODUCIBILITY STATEMENT
An annotated Colab notebook containing the code to replicate our results, includ- ing download instructions for model checkpoints, is available at https://bit.1y: sooklaagppooresss-aaasseeeeeesse
AUTHOR CONTRIBUTIONS
Neel Nanda the primary research contributor. He reverse engineered the weights of the mainline model to discover the Fourier multiplication algorithm and found the lines of evidence Section 4. He also discovered the restricted and excluded loss progress measures and that grokking in mainline model could divided into three discrete phases. Finally, he found the link between grokking, limited data, and phase transitions by exhibiting grokking in other settings with phase transitions.
Lawrence Chan was invaluable to the framing and technical writing of this work. In addition, he created the Gini coefficient progress measure and performed the analysis in the appendices exploring to what extent the results on the mainline model applied to the other small transformer models, including with other random seeds, architectures, prime moduli, and regularization methods.
Tom Lieberum contributed to the early stages of this work by creating a minimal setup of grokking with IL Transformer the modular addition task with no LayerNorm and finding the surprising periodicity within the model's internals.
Jess Smith performed experiments exploring grokking with different random seeds, architectures, and other paraet-prrmeters...
Jacob Steinhardt helped clarify and distill the results, provided significant amounts of editing and writing feedback, and suggested the progress measure frame.
ACKNOWLEDGMENTS
In writing this paper, our thinking and exposition was greatly clarified by correspondence with and feedback from Oliver Balfour, David Bau, Sid Black, Nick Cammarata, Stephen Casper, Bilal Chughtai, Arthur Conmy, Xander Davies, Ben Edelman, Nelson Elhage, Ryan Greenblatt, Jacob Hilton, Evan Hubinger, Zac Kenton, Janos Kramar, Lauro Langosco. Tao Lin, David Lindner, Eric Michaud, Vlad Mikulik, Noa Nabeshima, Chris Olah, Michela Paganini, Michela Paganini, Alex Ray, Rohin Shah, Buck Shlegeris, Alex Silverstein, Ben Toner, Johannes Treutlein, Nicholas Turner, Vikrant Varma, Vikrant Varma, Kevin Wang, Martin Wattenberg, John Wentworth, and Jeff Wu.
We'd also like thank Adam Gleave and Chengeheng Tan for providing substantial editing help, and Noa. Nabeshima and Vlad Mikulik for pair programming with Neel.
This work draws heavily on the interpretability techniques and framework developed by Elhage al. (2021) and Olsson et al. (2022).
We trained our models using PyTorch (Paszke et al., 2019) and performed our data analysis using NumPy (Harris et al., 2020), Pandas (Wes McKinney, 2010), and einops (Rogozhnikov, 2022). Our figures were made using Plotly (Plotly Technologies Inc., 2015).
Neel would like thank Jemima Jones for providing practical and emotional support as he navigated personal challenges while contributing to this paper, and to the Schelling Residency for providing an excellent research environment during the distillation stage. He would also like to thank the Authropic interpretability team, most notably Chris Olah, for an incredibly generous amount of mentorship during his time there, without which this investigation would never have happened.
REFERENCES
Boaz Barak, Benjamin L Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. arXiv preprint arXiv:2207.08790 08799. 2022.
10

---

Published as a conference paper at ICLR 2023
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Nick Cammarata, Shan Carter, Gabriel Goh, Chris Olah, Michael Petrov, Ludwig Schubert, Chelsea Voss, Ben Egan, and Swee Kiat Lim. Thread: Circuits. Distill, 2020. doi: 1 https://distill.pu//220//iiiiii
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. https:/transfomme-ciiiiiits http:/sssppppo//////a/eooo///ooe htmml.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint arXiv:/803.0330 2018.
Deep Ganguli, Danny Hernandez, Liane Lovitt, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova Dassarma, Dawn Drain, Nelson Elhage, et al. Predictability and surprise in large generative models. In 2022 ACM Conference Fairness, Accountability, and Transparency, pp- 1747-1764, 2022.
Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Coumapeau,, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Haldane, Jaime Fernández del Río. Mark Wiebe, Pearu Peterson, Pierre GGrard-Marchamn Kevin Sheppard, Tyler Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. Ar- ray programming with NumPy. Nature, 585(7825):357-366 September 2020. doi: 10. 1031038/ s41 -415866000-2649-22 URL
000 038
Niall Hurley and Scott Rickard. Comparing measures of sparsity. IEEE Transactions Information Theory, 55(10)) 4741, 2009.
Ziming Liu, Ouail Kitouni, Niklas Nolte, Eric J Michaud, Max Tegmark, and Mike Williams. To- wards understanding grokking: An effective theory of representation learning. arXiv preprint arXiv:2205.10343, 2022.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 1rXiv:1711.05101, 2017.
Thomas McGrath, Andrei Kapishnikov, Nenad Tomašev, Adam Pearce, Demis Hassabis, Been Kim, Ulrich Paquet, and Vladimir Kramnik. Acquisition of chess knowledge in alphazero. arXiv preprint arXiv:2111.09255 2021.
Beren Millidge. Grokking 'grokking', 2022.
URL
io/
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022. conntter-eermiii---ddndiiiiiiooo
https:/ttaassfomer-crrciist http:errrcccciiisuuuu222iiiieee-
Alexander Pan, Kush Bhatia, and Jacob Steinhardt. The effects of reward misspecification: Mapping and mitigating misaligned models. arXiv preprint 2201::2001.0554 2022.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. Advances in neural information processing systems, 32, 2019.
Plotly Technologies Inc. Collaborative data science, 2015. URL wwwph:ttp///c ptttttss:/ppo.....

---

Published as a conference paper at ICLR 2023
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gen- eralization beyond overfitting on small algorithmic datasets. arXiv preprint arXiv:2201.02177, 2022.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, lly Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8)99 2019.
Alex Rogozhnikov. Einops:
Clear and reliable tensor manipulations with einstein-like notation. In International Conference on Learning Representations, 2022. URL https: /tpp:///epeeeiiee.e////emee net/ Forum d1d=oapkss2bbccc
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, lly Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting, The journal of machine learning research, 15(1):1929-199) 1958, 2014.
Jacob Steinhardt. More different for ai, Feb 2022. URL attps://bounded-ree ghost www.mooe-///ooeeeeeeeeeeeeeeeeee is different for-ai/. ai/.
Vimal Thilak, Etai Littwin, Shuangfei Zhai, Omid Saremi, Roni Paiss, and Joshua Susskind. The slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon. arXiv preprint arXhv:2206.04887, 2022.
Kevin Wang, Alexandre Variengien, Arthur Commy, Buck Shlegeris, and Jacob Steinhardt. Inter- pretability in the wild: a circuit for indirect object identification in gpt-2 small. arXiv preprint arXiv:221 1.00593, 2022.
Jason Wei, Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yo- gatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilities of large language models. arXiv preprint arXiv:2206.0766 2022a.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoc Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. arXiv preprint arXiv:2201.11900 2022b.
Wes McKinney. Data Structures for Statistical Computing in Python. In Stéfan van der Walt and Jarrod Millman (eds.), Proceedings of the 9th Python in Science Conference, pp. 56 61, 2010. doi: 10.25080/Majora 9tt.:/Maaoooaa999922222222
12

---

Published as a conference paper at ICLR 2023
A MATHEMATICAL STRUCTURE OF THE TRANSFORMER
We follow the conventions and notation of Elhage et al. (2021) in describing our model. Here, we briefly recap their notation and examine in our specific case.
We denote our hyperparameters as follows: dvocab = 113 is the size of the input and output spaces (treating "=' separately), dmodet 128 is the width of the residual stream (i.e. embedding size), dheed = 32 is the size of query, key and value vectors for a single attention head, and dmlp = 512 the number of neurons.
We denote the parameters as follows: WE (embedding layer); W (positional embedding); W (queries), W (keys), W (values), w (attention output) (the weight matrices of head j in the attention layer); W and iin for the input linear map of the MLP layer; Wout and bout for the output linear map of the MLP layer; and Wy (unembedding layer). Note that we do not have biases in embedding, attention layer unembedding, and we do not tie the matrices for the umbbedding/unemmbedding layers.
We now describe the mathemattcaa structure of network. Note that loss is only calculated from the logits on the final token, and information only moves between tokens during the attention layer, so our variables from the end of the attention layer onwards only refer the final token. We use t; to denote the token in position i (as a one-hot encoded vector), Ps to denote the ith positional embedding, =/!² denote the initial residual stream on token with index A to denote the attention scores from
all previous tokens from head i, (x)) denote the residual stream after the attention layer on the final token, MLP to denote the neuron activations in the MLP layer on the final token, (2) the final residual stream on the final token, Logits the logits on the final token.
The logits are calculated via the following equations:
4 WF4 + P
(0)
A
softmax (z(0)) W W Wz-32))
<()) E W A3) x22
MLP ReLU(W (in
(2) = W Wout N + (2) W ReLU + x1))
Logits (uz
As in Elhage et al. (2021), we refer to the term W W (x:(0) as the OV circuit for head j.
A.1 EMPIRICAL MODEL SIMPLIFICATIONS
We make two empirical observations:
The attention paid from '=' to itself is trivial. In practice, the average attention paid is 1% to 0.4% for each head, and ablating this does not affect model performance at all. The skip connection around the MLP layer is not important for the model's computation and can be ignored. Concretely, we set it to zero or to its average (zero or mean ablation) then model accuracy is unchanged, and loss goes from 2.4 10 to 9.12 10 7 and 7.25 10-7 respectively. This is a significant increase in loss, but from such a small baseline that still ignore and reverse engineer the model's computation. (That being said, both the attention heads and the skip connection around them are crucial to the functioning of the model: ablating attention heads increases loss to 24.3, while zero ablating the skip connection around the attention heads increases loss to 19.1, both significantly worse than chance.)
A consequence of the first observation is that the attention is now a softmax over 2 elements, i.e. sigmoid over the difference. And constant, it is independent of x and y, and the embedding
13

---

Published as a conference paper at ICLR 2023
Constructive Interference of Cosine Waves of Different Frequencies
14
34
cos
-2
-3
Figure 8: As discussed in Appendix B, while for every k [O,...P L]. cos (?P 2)) achieves its maximum value (1) at x mod 113, it still has additional peaks at different values that are close to the maximum value. However, by adding together cosine waves of the 5 keyfrequencies, the model constructs a periodic function where the value at r 0 mod 113 is significantly larger than its value anywhere else.
(0) and positional embedding '=' are fixed. So A3 a I22 '='
=/") (and A
W Wx
A/)
A consequence of the second observation is that Logits ss Wy ow... which we denote as WL Wy W out- From the perspective of the network, Wz is the meaningful matrix, either of its constituents, since they compose linearly.
B WHY USE CONSTRUCTIVE INTEREFERENCE?
As demonstrated in Section 4 and Appendix C.2.1, small transformers trained on this task use several different frequencies which they add together. The reason for this is to end up with function whose value at x mod 113 is significantly larger than any other æ.
For example, consider the function f1a(x) = cos 22-14 x). This function has period 113 and is maximized at x = 0 mod 113. However, other values of x cause this function to be close to I: f1a(8) = (105) 0.998, (166 (19 0.994, etc.
Now consider fas (z) = cos (2m-35 113 z). While this function also has period 113 and is maximized at x = mod 113, it turns out that fas f35 = -0.990. This means that by adding together f14 and fas. we end up with function that is not close at x 8 mod 113. Similarly, while f36 = 0.961, fs2(16) -0.56, and so adding a third frequency reduces the peak at 16 mod 113.
We show the constructive interference resulting from the cosine waves for the five frequencies used by the mainline model in Figure 8.
C SUPPORTING EVIDENCE FOR MECHANISTIC ANALYSIS OF MODULAR ARITHMETIC NETWORKS
C.1 FURTHER ANALYSIS OF THE SPECIFIC TRAINING RUN DISCUSSED THE PAPER
In this section, we provide additional evidence relating to the mainline model.
14

---

Published as a conference paper at ICLR 2023
Attention patterns by Head ("=' to a)
20
40
60
80
100
0
50
100
0
50
100
50
100
0
50
100
Figure 9: Attention patterns for each head, from the "=' token the third sequence position to the a token at the first sequence position, as heatmap over the inputs. All four attention heads exhibit striking periodicity.
Head
g
FVE
0
35
-0.26
0..14
99.03%
1
42
0.27
-0.04
98.49%
2
0.29
-0.05
99.07%
3
42
-0.26
0.04
97.91%
Table 2: For each attention head, show the pattern from "=' is well approximated by 0.5 + a ((cos(w).a) -co(())))) + (issn((ssaaaaa-----a (snn)))))(wzb)) and give the coefficients and fraction of variance explained for this approximation.
C.1.1 PERIODICITY IN THE ACTIVATIONS OF OTHER ATTENTION HEADS
In Figure 9 we plot the attention patterns from the final token to the first token a for all attention heads, heatmap the inputs and b, as this is scalar for each head. We observe striking periodicity and further that heads and 3 represent the same frequency while heads O and different.
As shown in Appendix A.1, the attention paid from '=' to itself is negligible, so A/= = -A and suffices to plot attention to
C.1.2 APPROXIMATING ATTENTION HEADS WITH SINES AND COSINES
Attention heads approximately compute degree-2 polynomials of a single frequency or are used to amplify W F. In order to compute terms like cos (wx (a b)), the model needs to compute the product of the sine and cosine embeddings output by WF. As the attention heads are approx- imately bilinear (product of attention weights and OV circuit), they are a natural place to perform this computation. Indeed, for each head, the attention scores' Fourier transform is concentrated on single frequency we. For two of the four heads, the corresponding OV circuit is concentrated on that same frequency. Moreover, the softmax mapping the attention scores to attention weights is in a regime where it behaves approximately linearly (and replacing it with linear function actually improves performance). Thus the attention weights multiply with the OV output create degree-2 polynomials the frequency wk. as would be needed the cosine/sine addition formulas.
For the remaining two heads, their attention scores approximately sum to one and the OV circuits contain all five key frequencies, suggesting that they are used to increase the magnitude of key frequencies in the residual stream. We confirm all of these claims in Appendix C.1.3.
C.1.3 THE ATTENTION PATTERN WEIGHTS ARE WELL APPROXIMATED BY DIFFERENCES OF SINES AND COSINES A SINGLE FREQUENCY.
The periodicity of the attention heads has striking form form-45 Ag is well approximated by 0.5 + a (coss (wga) cos( (yb)) + B (sin (waa) sin(w.b)), for some frequency wy and constants and 3 (which may differ for each head). Note further that this simplifies to 0.5 (cos(wx(() (a + 8)) cos wa (b+66 {))) for some constants y and 0. We show the coefficients and fraction variance explained in Table
15

---

Published as a conference paper at ICLR 2023
Confficient Fourter Component
Fourier components for Co
Coefficient of Fouttr Compament
Fourier components for C₁
4
70
Frequency
Frequency
Coefficient of Fouuier Conporent
Fourier components for C
Coefficient Fourier Component
Fourier components for C3
Frequency
Frequency
Figure 10: We plot the attention pattern weights Cj in the Fourier basis for each of the four heads j E {0, 3/. We observe significant sparsity, with almost all of each term being associated with single frequency.
Mechanistic Analysis of Attention Patterns. We can further mechanistically analyse how the model achieves this form. The following is high-level sketch of what is going on:
First, note that the attention score on position and head is just lookup table on the input token a (of size P). To see why, note that A mag (0) w W
(0) _.00) is constant since the token
always ==' and To
Wgoo + Pu- So this reduces to to Cj + for some constant vector
(0) Cj = W W 2 E IR' and some scalar = p w W2 (0) As to is one-hot encoded, this is just a lookup table, which we may instead denote as [a(a]
Next, note that the attention pattern from the attention paid = is negligible and can be ignored. So the softmax reduces to a softmax over two elements, which index or value, it is the same for position and 1.
0 is (Cy(aa) (b]) As argued in Appendix A.1,
sigmoid on their difference. As form of Cj does mention the token
We now show that Cy is well-approximated by wave of frequency wk, for integer ky. That is C2 ~ F, cos( cos(wk, a) + G. sin (Wk, a). We do this by simply computing Cj and fitting the constants Fj and G, minimize & loss, and display the resulting coefficients for each head Figure 10. This fit explain 99.02%, 95.21 %, 99. 10%, 92.42% of the variance of C₁ respectively. Interestingly, the coefficients of heads and are almost exactly the opposite of each other.
For each head j. (C () = 0.5 [C, Cy(b)) for some constant E5- E,-that is, the sigmoid has some linear approximation. (The intercept will be 0.5 by symmetry.) The striking thing is that, because the inputs to the sigmoid for the attention heads are over fairly wide range ([-5, 5) roughly), the linear approximation to the sigmoid is fairly good fit, explaining 97.5% of the variance.
We validate that this is all that going on, by replacing the sigmoid with the best linear fit. This improves performance, decreasing test loss from 2.41 10 to 2.12 10-7
By properties of sinusoidal functions, the attention patterns of each head will be well approximated by 0.5 = C2 (cos( Wk, (a 0,)) cos( Wky (b + 9,))) the softmax is linear, with intercept of 0.5, and the weights C₁ map each token to a score that is a wave in a single frequency. This exactly gives us the periodic form shown in Figure 9.
16

---

Published as a conference paper at ICLR 2023
Fourier components of OV circuit for head 0
Fourier components of OV circuit for head 1
Fourier Composent
Norm of Foutler Component
1.5
04
20 Frequency
43
Frequency
Fourier components of OV circuit for head 2
Fourier components of OV circuit for head 3
of Fourier Component
Norcn of Fourfer Component
25
1.5
Frequency
Frequency
Figure 11: We plot the output of the OV circuit W W'.3.a(0 in the Fourier basis for each of the four heads {0, 2,2,3), 3). As with the attention pattern weights Cj in Figure 10, observe that the only components with significant norm are those corresponding to key frequencies, and that the largest component corresponds
the frequencies of the attention patterns of the attention heads. As attention pattern of heads 1 and are sum to one, but their OV circuits are almost exactly the same and consist of all five key frequencies, this implies that heads and are used to increase the magnitude of key frequencies in the residual stream (Section C.1.3).
Finally, for each head j. we plot the output of the OV circuit W W.).o(0) in the Fourier basis and display the results in Figure 11). The largest component of each head corresponding to the frequency of the attention pattern Cj. with heads and 2 being almost entirely composed of a sines and cosines of single frequency. On the other hand, the norms for the components of heads and are almost exactly the same, and contain all five key frequencies. As the coefficients of the attention pattern weights have the opposite non-constant components (Table 2, Figure 10), their attention scores sum almost exactly to across all inputs. This implies that heads 1 and 3 are used to output the first order terms sin (wx), cos (wx) in the five key frequencies. We speculate that this is because of weight decay encouraging the embeddings WE to be small, causing the network to allocate two of its attention heads to effectively increasing the size of WE.
Bringing it all together, this implies that attention heads 0 and 2 are approximately computing degree polynomial cosines and sines single frequency each, while heads and amplify the key frequencies in the residual stream.
C.1.4 PERIODICITY IN THE ACTIVATIONS OF ADDITIONAL NEURONS
In Figure 12, display the activations of four more MLP neurons, function of the inputs. As with neuron O, the activations of these neurons are also periodic the inputs.
C.1.5 ADDITIONAL GROKKING FIGURES FOR MAINLINE RUN
In Figure 13, we display the accuracy of the model when restricting the model to use only the five key frequencies. As with restricted loss, this improves model performance during training.
In Figure 14, we show the coefficients of the five key frequencies in the logits, calculated by regress- ing the logits against the five cos (w) (a + b c)) terms.
In Figure 15, we plot the excluded loss if we exclude each of the five key frequencies (as opposed to all five key frequencies).
17

---

Published as a conference paper at ICLR 2023
Neuron Activations for Additional Neurons
20
40
60;
80
100
0
50
100
50
100
50
100
50
100
Figure 12: Plots of neuron activations for MLP neurons 1, 2, 3 and 4, for inputs a, b E {0, 1,.,, 112). As with Neuron 0, all of the activation patterns periodic in both inputs.
Pure Restricted Accuracy
Train Accuracy
0.6
Test Accuracy
Restricted Accuracy (Train)
0.4
Restricted Accuracy (Test)
0
10k
20k Epoch
30k
Figure 13: Accuracy when restricting Fourier Components to the five key frequencies. As with restricted loss, this shows that the model figures how to generalize modulo delcting noise before it removes the noise.
All three of these figures have inflection points corresponding to the relevant phases of grokking, discussed in Section 5.1.
C.2 ADDITIONAL RESULTS FROM DIFFERENT RUNS
In this section, we plot relevant figures from other runs, either with the same architecture (Appendix 1) or with different architectures experimental setups (Appendix C.2.2). Note that in general, while all models learn use variants of the modular arithmetic algorithm, they use a varying number of different key frequencies. In order to find the key frequencies to calculate the excluded and restricted loss, we perform a DFT on the neuron-logit map Wr. then take the frequencies with nontrivial coefficients."
C.2.1 ADDITIONAL RESULTS FOR DIFFERENT RUNS WITH THE SAME ARCHITECTURE
In this section, we provide evidence that all 4 other (i.e., random seeds) using the experimental setup of our mainline model also use the Fourier multiplication algorithm, and then confirm that the same phases of grokking also occur on these runs.
OOee method for getting a general (model-indepeedden) progress measure for this task is to compute the excluded loss for each of the 56 unique frequencies and then take the max. We omit the plots for this variant of the excluded loss they broadly similar.
18

---

Published as a conference paper at ICLR 2023
Coefficients of cos(wi(a+b-c)) in the Logits
Freg
20
41
42
10
0
5k
10k
15k
20k
25k
30k
Epoch
Figure 14: The coefficients of cos(w(( + b c)) in the logits over the model's training. As with the metrice in the this chome nice internolatign and arouth of each cocine
Accuracy when Excluding Key Frequencies
Loss when Excluding Key Frequencies
acuuacy Tee Acrunyy
Tiod Low
0.8
Accuracy
Loss
k-35
k-42
100p
1-42
k-52
0.4
Epoch
Epoch
Figure 15: The excluded accuracy (left) and loss (right) if we exclude each of the five key frequencies for our mainline model. As with the excluded loss results in Section 5.1, this shows that the model interpolates between memorising and generalising.
Norm of Fourier Component
Embedding Matrix (Seed 1)
Norm of Fourier Companent
Embedding Matrix (Seed
2.5
255
1.5
Frequency
Frequency
Norm of Fourier Component
Embedding Matrix (Seed 3)
Norm of Fourier Componen
Embedding Matrix (Seed 4)
255
cos
2.5
COS
1.5
Frequency
Frequency k
Figure 16: The norms the Fourier components the embedding matrix WE for each four other random seeds for the original (1 layer) architecture. As discussed in Section 4. and Appendix C.2.1, 1, the sparsity of W. the Fourier basis evidence that the network operating a Fourier basis.
19

---

Published as a conference paper at ICLR 2023
Wz Component
Fourier components MLP(P, b) or oi (M,.((,b))
FVE
cos (w2c)
147.4 cos (W2a) oos (wyb) 145.8 sin (Wza) sin (wzb) = 146.6 cos W2 6))
99.2%
sin (w2c)
145.5 cos (wza) sin (wzb) 145.6 sin (w2a) cos (w2b)
(w2 145.5sin
6))
99.1%
cos (Wyc)
49.3 cos (1000) cos (wsp) 48.0 sim (Waa) sin (wyb) 48.6 (wa(a b))
96.4%
sin (8yc)
48.6 cos (Wya) sin (wyb) + 488.5 sin (Mya) cos (wyb) =e 48 (a bb)
96.7%
cos (wocc)
58.0 cos (Wi90) cos (wrob) - 58.3 (W1ea) sin (wrbb) = 58.2 cos (wio(a
95.4%
sin (wrec)
59.3 cos (Wisa) sin (wish) + 59.4 sin (Wisa) cos (W1sb) x 59.4 (wia(a
93.9%
(wsic)
94.4 cos (Waia) a) (war 96.4sin (w31a) sin (waib) b) x 95.4 cos (wes (a b))
98.4%
sin (ware)
97.2 cos (wsra) sin (warb) + 97.1 sin (Ward) cos (warb) Ra 97.2 (wax (a B))
98.7%
(a) Sced
W,. Component
Fourier components uf MLP((a, b) or oz MLP(a,
FVE
cos (240c)
97.0 cos [M400] cos (M040b]
99. (Waooa)sin (w04b) x 98 cos (w00 b))
97.3%
sin (wagc)
81.3 cos (waga) (wanb) + 81.3 sin (waga) cos (wabb)
81.3 sin (wao bj)
92.7%
(wacc)
309.1 1 (wa4a) cos (waab) 338.7 sin (wasa) sin (waab) R 323.9 cos (was 6))
98.5%
(
327.3 3 cos (wisaa) sin (wea b) +- 327.2 sin (Waya) cos (wayb) Rs 327.3 (wex (a b))
98.9%
cos (bac))
192. cos (Waaa) cos (Waab) 192.2 sin (Waaa) sim (Wbab) == 192.11c cos (Was (a b))
97.3%
sin (Waac)
166.7 cos (Waad) sin (Waab) + 166.8 sin (Waad) cos (Weab) = 166.8 (Waa (a b))
95.7%
(b) Seed 2
Wz Component
Fourier components of vz. MLP((, b) or vz MLP(a, b)
FVE
(w3ic)
156. (Wara) cos (War b) 156.5 sin (W31 (w316) R 156.3 (wa1 (a b))
99.3%
(W3rc)
150.7 cos (w]1a) sin (W31 + 150.7 sin (Wyra) cos (ws1b) R 150.7 sin (wzn (a b))
98.9%
cos (4400)
72.5 cos (Wapa) cos (Warb) - 76.8 (Wasa) sin (Warb) = 74.6 cos (14z (a b))
95.9%
sin CW)
74.7 cos (Wasa) sin (wasb) - 74.6 sin (Wasa) cos (Wanb) = 74.6 sin (Was bl)
96.6%
(Wayc)
45.9 cos (Wayd) cos (Wayb) -4.5 sin (Waya) sin (Wayb) 45. cos (Way [a b))
97.0%
sin (wa9c)
45.8 (w4ga)
(wasb) + 45.8 (waaa) cos (wash) 45.8 sin (wa9 -))
96.9%
cos (u)s2c)
71.6 cos (Ws2a) cos (Ws2b) 72.
(Ws2a) sin (ws2h) RJ 71.9 cos (ws2 (a + b))
98.5%
sin (Ws2c)
68.7 cos (ws2a) sin b)) + 68.7 sin (w52a) cos (ws2b) Rs 68.7 sin (w52 b))
97.9%
(c) Seed 3
Wi. Component
Fourier components of u. (a, b) or v. MLP b))b) b)
FVE
cos (Wire)
66.0 (W170] cos [Wirb] co8(117b) 63.5 sin (W170) sin (W1rb) z 64.8 cos (Wir(a + b))
96.4%
sin (w17c)
66.4 cos (wi7a) (w17b) + 66m (wiza) cos (w1rb)
66.4 sin (w17(a b))
94.9%
cos (u)2c)
68.7 cos (Ws2a) cos (Wszb) 68..4sin (W324) sin (Ws2b) A 68.5 cos (wss (a b))
96.2%
sin (W32c)
68.0 (W32a) sin (w)2b) + 68.00sn (W32a) cos (wa2b) & 68.0 sin (Wa2 (a b))
96.3%
cos (M42c)
100.1 cos (waza) cos (uzzb) 99.00 (waza) sin (uuzb) R 98.2 cos (waz(a b))
97.9%
sin (Wazc)
100.2 cos (wlz2a) sin (wazb) + 100. 1 sin (waza) cos (w4zb) 100.1 sin (waz (a b))
98.6%
cos [Wa10)
118.0 cos (wara) cos (War 6) 116.2 sin (Ward)
(wa16) 117.1 cos (War (a b))
99.0%
sin (wSic)
114.3 cos (wsia) sin (ws1b)
+ 114.2 sin (wssa) cos (ws16) == 114.2 (ws1 (a + 6))
98.5%
(d) Sced 4
Table For each of the directions in the neuron-logit map WL of the final models from 4 other ran- dom seeds (Appendix C.2.1). we project the MLP activations in that direction then perform Fourier transform. For brevity, we omit terms with coefficients less than 15% of the largest coefficient. We then compute the fraction of variance explained (FVE) if we replace the projection with a multiple of a single term of the form cos (w + b)) or sin (wx (a + b)), and find that this is consistently close to 1.
20

---

Published as a conference paper at ICLR 2023
Norm of Fouries Component
Neuron-Loggi Map (Seed 1)
Norm of Fourier Component
Neuron-LLogh Map (Seed 2)
255
1.5
0.5
Frequency k
Frequency k
Norm of Fourier Companent
Neuron-Logit Map (Seed 3)
Norm of Fourier Companent
Neuron-Logit Map (Seed 4)
25
255
1.5
Frequency k
Frequency
Figure 17: The norms of the direction corresponding to sine and cosine waves in the neuron-logit map weights Wz. As with the mainline model discussed in the main body and discussed in Appendix 1.2.1, WL is consistently sparse, providing is evidence that all four are operating a Fourier basis.
Seed
Test Loss
Loss (Key frequencies removed)
Loss (All other frequencies removed)
2.07 10
6.5 10"
5.7 10
2
2.1
10-7
1.1 10!
6.2 10-6
3
2.05
10
6.7 10
5.5 10
4
2.33
10-7
6.8 10
6.0 10 10-6
Table As discussed in Appendix C.2.1, ablating the key frequencies for each of the networks re- duces performance worse than chance, while ablating all other frequencies improves performance.
Confirming that the other seeds the Fourier Multiplication Algorithm. In Figure 16, show the norms of the Fourier components of the embedding matrix WE for each of the 4 other random seeds. As with the mainline model, the matrices sparse in the Fourier basis. In Figure 17, we show the norms of the Fourier components the neuron-logit map Wt for the 4 other random seeds. The matrices sparse in the Fourier basis, enabling to identify 4 key frequencies for each of the seeds. Again, note that the specific frequencies differ by seed.
Using the key frequencies identified in the neuron-logit map, we repeat the experiment in Sec- tion 4.2, where we "read off" the MLP activations in the 6 or 8 directions corresponding to the key frequencies. As with our mainline model, this lets us identify the trigonometric identities for cos (wa(a + 6)) and sin (wy(a (a + 6)) being computed at the MLP layer. We confirm that the trigono- metric identities are a good approximation by approximating the activations with single term of the form cos (wx(a (a + b)) sin (wa(a (a + b)) as with the mainline model, the fraction of variance explained is consistently close to 100%.
Next, we ablate the key frequencies from the logits as in Section 4.4 and report the results in Table with the mainline model, ablating all of the key frequencies reduces performance worse than chance, while ablating everything but the key frequencies improves test performance.
Progress measures and grokking. Finally, we confirm the progress measure and grokking results from the mainline model on other runs with the same architecture. In Figure 18, we display the train, test, and restricted loss for each of the four other random seeds. In Figure 19, we display the Gini coefficients of the Fourier components of the embedding matrix WE and the neuron-logit map W,, z, for each the four other random seeds. The shape of the curves are very similar to those of the mainline model, allowing divide grokking these models into the same three phases
21

---

Published as a conference paper at ICLR 2023
Restricted Loss for Seed
Restricted Loss for Seed 2
Train
Tex loss
Test loss
Rostiited
Rustricual loas
Loss
0.01
Loss
0.01
100g
100g
1p
10n
25k
Epoch
Epoch
Restricted Loss for Seed
Restricted Loss for Seed
Trzin
Train
Tex liss
Text loss
Restrical loss
Loss
0.01
Loss
0.01
100p
100
1a
10n
Epoch
Epoch
Figure 18: The train, test, and restricted loss for each of the four other random seeds described in Appendix C.2.1. The lines delineate the 3 phases of training: memorization, circuit formation, and cleanup (and final stable phase). As with the mainline model, restricted loss consistently declines prior to train loss. Note that while the shapes of the loss curves are similar to each other and those of the mainline model, the exact time that grokking occurs (and thus the dividers between the phases of grokking) differ by random seed. Interestingly, memorization is complete by around 1400 steps for all five runs.
identified in the main text. Interestingly, while all of the models complete memorization by around 1400 epochs, circuit formation and cleanup occur at different times.
C.2.2 RESULTS FOR OTHER EXPERIMENTAL SETUPS
In this section, provide further evidence that small transformers grok the modular addition task, by varying the size of the network, the amount of training data, and the size the prime P.
L-Layer Transformers with Varying Fractions of Training Data. We find that grokking occurs for the modular addition task with 113 113 pairs of inputs that the model sees during training), shown in Figure 20. Smaller amount lead to slower grokking, but sufficiently large fractions of data (2 60%) lead to immediate generalization, shown in Figures 20 and 21.
113 for many data fractions (that is, the fraction of the
As with the results in Appendix C.2.1, all of the 1-layer transformers in this section also converge to using the Fourier multiplication algorithm.
2-Layer Transformers. As shown in Figure 22, 2-layer transformers also exhibit some degree of grokking. However, this is complicated by the slingshot mechanism (Thilak et al., 2022). We display the excluded loss of a 2-layer transformer in Figure 23 and find it shows a similar pattern to the mainline l-layer transformer, in that improves relatively smoothly before grokking occurs.
Smaller and larger primes. We also examined smaller and larger prime moduli. For P 53 (Figure 24), we explored variety of weight decays to observe grokking in the small prime case. With the original weight decay setting of x = 1, we found that the models never generalized. However, increasing the weight decay À does allow the model to grok. We speculate that
22

---

Published as a conference paper at ICLR 2023
Gini Coefficients, Seed 1
Gini Coefficients, Seed 2
0.8
Gini coefficient
Gini coefficient
0.5
0.4
0.2
5k
100
1%k
20k
25k
5k
10k.
15k.
20k
25k
Epoch
Epoch
Gini Coefficients, Seed 3
Gini Coefficients, Seed 4
0.8
Gini coeffiilent
Gini coefficient
0.4
Ws
RGss
0.2
0.2
5k
10k
13k
20k 25k
10k:
13k
20k
25k
Epoch
Epoch
Figure 19:
The Gini coefficients (a measure of sparsity) of the Fourier components of the embedding matrix WE g and the neuron-logit map W,. for each of the four other random seeds. The lines delineate the 3 phases of training: memorization, circuit formation, and cleanup (and final stable phase). As with the mainline model, sparsity increases slowly during memorization and circuit formation, and then quickly during cleanup.
Data Fraction 0.1
Data Fraction 0.2
Data Fraction 0.3
Tram
Tree los Yeut
Yest
000
IPL. Epoch
15%
204
5%
I06 Epoch
15%
306
S2
V/6 Epoch
15%
24
Data Fraction 0.4
Data Fraction 0.5
Data Fraction 0.6
Toaa Jas
Tex
Tox
100
1000
2003
3000
4306
1000
2003
3800
4003
1000
2003
3000
4000
Epoch
Epoch
Epoch
Data Fraction 0.7
Data Fraction 0.8
Data Fraction 0.9
Tux
400
000
Tec
100
10:00
3000
10:00
10:00
2002
10:00
3003
10:00
2003
10.00
3000
Fpoch
Fpoch
Fpoch
Figure 20: Training and test losses for 1-layer transformer the modular addition task with
P = 113, with varying fractions of the 113 - 113 pairs of possible inputs used in training. Grokking occurs when between 30- 50% of the dataset is used during training and lower fractions of data lead to slower grokking. Using 2 60% data leads to immediate generalization, while using 10% or 20% of the data doesn't 't lead grokking even after 40k epochs. Note the different x-axes; we only show 5k epochs for the runs with data fraction > 40% for more detail.
23

---

Published as a conference paper at ICLR 2023
Training Epochs until <-e-06 Loss
18k
16k
Train loss
Tost loss
Number of steps
14k
8k
6k
4k
2k
0.6
0.7
0.8
0.9
Fraction of train data
Figure 21: Number of steps for train/test loss to < 10 as a function the amount of training data. While train loss immediately converges to below 10-6 for all data fractions, generalization takes significantly longer with lower fractions of data. Note that the plots for other thresholds are also qualitatively similar.
Average Train Loss
Average Test Loss
Log Loss
0.01
100
1u
10n
10k
15k
Epoch
Figure 22: Training and test loss for 2-layer version of the original architecture. Average across random seeds is in bold.
24

---

Published as a conference paper at ICLR 2023
Excluded Over Frequencies
Excleded Loss
Train Lova
Loss
1000
1p-
111
Epoch
Figure 23: Training, test, and full excluded loss for 2-layer version of the original architecture. One random seed chosen for readability.
Average Train Loss
Average Test Loss
Log Log Loss
0.01
100p
1p
5k
10k
15k
20k
Epoch
Figure 24: The training and test losses for P = 53 and all other hyperparameters except weight decay =7=5) 5) the same as the main training run discussed in the paper. The averages are bold, and all contributing runs are partially transparent. Note that grokking occurs.
Train Loss
Tost Loss
Log Loss
0.01
100
1u
10k
15k
20k
Epoch
Figure 25: The training and test losses for P 401 and all other hyperparameters the same the main training run discussed the paper. Grokking docsn't occur (the model generalizes immedi- ately), even across variety of weight decays.
25

---

Published as a conference paper at ICLR 2023
this is because the memorization solution is significantly smaller (since there only 53 53 total pairs), thereby requiring more aggressive weight decay for the generalizing solution to be favored.
For P 109, we saw exactly the same behavior as with the mainline model.
For P = 401 (Figure 25), we could not get grokking, even by varying the weight decay parame- ter À E {0.3, 0.5, 1,3,5 5, 83. Instead, the model immediately learns the generalizing solution. We believe this is because the amount of data seen by the model is greatly increased compared to the P = 113 case (from 30% of 113 113 pairs 30% of 401 401 pairs), thereby favoring the gen- eralizing solution from the start. We then trained 3 models each using 5%, 10%, 20% of the pairs of training data with À 1, and found that the models trained on 5% and 10% of the data imme- diately overfit and never generalized, while the models trained on 20% of the data also generalized immediately.
C.2.3 GENERALIZING MODELS CONSISTENTLY USE THE FOURIER MULTIPLICATION ALGORITHM
For each of the models in Appendix C.2.2 that achieve low test loss, repeated the analysis per- formed in the mainline model, and summarize the results in Table 5. We list their key frequencies, Gini coefficients, and relevant FVEs. We find that every model trained with weight decay and that generalizes correctly implements some variation of the Fourier multiplication algorithm.
Interestingly, the embedding and unembedding matrices of the models trained with dropout not sparse in the Fourier basis, and the logits for the p = 0.2 models are not as well explained by sum of cosines as the other models (likely because the p 0.2 models are simply worse the task). We speculate that this is likely due to combination of insufficient training epochs (as dropout models seem to take much longer to grok) and the inherent need for redundancy for networks trained via dropout.
As with the mainline model, we ignore the final skip connection (around the final MLP), as all of the generalizing models studied do not suffer significant performance penalties if the skip connection is zero or mean ablated (Table 6).
D ADDITIONAL RESULTS ON GROKKING
D.1 BOTH REGULARIZATION AND LIMITED DATA ARE NECESSARY FOR GROKKING
As discussed in Section 7 and Appendix C.2, the weight decay and the amount of data seem to have strong effect whether grokking occurs. To confirm this, we experiment with removing weight decay and varying the amount of data on I-layer transformers. In Figure 26, we give the training, test, and full excluded loss for a typical training run with \ 0 (no weight decay). As the figure shows, no grokking occurs, and excluded loss does not increase, suggesting that the model does not form the circuit for generalizing algorithm at all.
In Figure 20, we show the test loss curves for models trained with weight decay X = 1 and on various fractions of the data. Though all the train losses are approximately the same- that is, they memorize at the same rate, models trained on smaller fractions of data take longer to grok.
In Figure 27, we display the test and train loss of models trained with À = 0.3 and amounts of weight decay lead to slower grokking, while larger amounts of weight decay lead to faster grokking-ov average, it takes around 3k epochs for models to grok with weight decay = 0.3, 5-10k epochs for the models to grok with weight decay x 1.0, and 20k epochs for the models to grok with weight decay = 3.0.
3.0. Smaller
Finally, we test whether other forms of regularization can also induce grokking. We replaced weight decay with the following types of regularization while keeping all other hyperpameters the same:
1. Dropout We add dropout Srivastava et al. (2014) to the MLP neurons, with p {o.2, 0.5, 0.83. That is, for each individual neuron, we set to with probability p during training, and also multiply the outputs of the other neurons by T-p-
26

---

Published as a conference paper at ICLR 2023
Model
Loss
Ginl(Wa> )
Gini(Wr,)
Frequencies
Logit FVE
MLP
ALLEVVE
40%
Training Dara
1,98 10
0.76
0.79
[17, 49,
94.9%
83.3%
[26.1%]
50%
Training Dara
1.68 10
0.75
0.77
[2. 17, 41, 44]
91.2%
85.2%
[28.2%)
60%
Training Data
1.23 10
0.79
0.84
[2, 23, 34,51]
96.4%
95.7%
(7.%%11444 [1.4%]
70%
Training Data
9.85 10-*
0.80
0.91
[14,
99.0%
98.9% [0.4%)
Training Dara
0.62
[38, 41]
[2.5%)
90% Training
0.79
0,88
[3. 26, 34,
98.6%
98.7% [0.3%]
Layer Transfommer
9.54
0.59
[14, 291
91.8%
95.2% [1.9%)]
2 Layer Transformer
4.41
0.55
0.73
12, 35,
86.1%
86.2% [6.4%)]
2 Layer Trasformer
6.50 10-2
0.66
0.80
[4, 28]
88.5%
85.4% [5.9%%
2 Layer Tranformer
4.18 10-
0.56
0.76
[4, 15, 54]
91 91%
81.2% [17.8%)
Layer Transformer
1.75
[3.4, 30. 38]
71.9% [19.5%)
[6, 16,
91.2%
90.2% [..8%
P
1.03 100-
0.56
0.72
13,
94.8%
[6.4%)
P 53
1.21 10-5
0.66
0.79
[13.
98.2%
97.6% [0.9%
P 53
3.95 10-8
0.66
0.74
[3. 14, 151
88.5%
91.8% [4.6%)
P 53
5.56 10
0.67
0.80
[10, 14, 22]
98.1%
98.3% (0.6%)
P 109
2.02 10
0.76
0.83
[6, 22,
98.0%
97.3%: [1.9%)
[8. 14, 29, 31,
[2.3%%]
1.66 10
0.78
0.86
[13, 23,
98.5%
97.6% [0.9%)
P 109
2.50 10
0.68
0.82
[8, 41]
96.8%
95.5% [2.3%%
P 109
2.77 10-7
0,76
0.85
[29. 37, 38, 49]
97.9%
9% [0.8%)
Dropout = 0.2
2.65 10-1
0.19
0.46
11, 7, 17, 33, 49,
71.3%
65.0% [17.5%
Dropout 0.2
4.52 10-1
0.19
0.46
[3,8, 19, 32, 34, 44]
73.3%
71.4% [10.7%]
Dropout p 0.2
2.03
0.20
0.45
14,5,32. 41, 14,44,,4 49.
74.2%
71.1% (10.6%)
Dropout p
< 10
0.26
[1,4, 46,
[3.5%]
Dropout 0.5 2.01 10
0.20
0.49
[16, 43]
88.4%
88.4% [3.0%]
Dropout 0.5
< 10-
0.25
0.54
[1.4. 7, 29. 29,31,422
86.1%
85.6 35.6% [4.0%]
Table 5: For each of the models in Appendices C.2.3 and D.1 that generalizes to test data, we report the test loss, the Gini coefficients of the norms of the Fourier components of WE and Wz (Section 5. the key frequencies of the network, and the fraction of variance in logits explained by weighted sum of cos (wx(e (a + b- c))s over the key frequencies (Section 4.2).
In addition, we find the components Wk, Uk of Wr. that correspond to cosines and sines of the key frequencies, and then report the average fraction of variance u MLP(a, b) and oo MI.P(a, ((aab) b) explained by single term of form cos (wx + b)) or sin (wx + b)) respectively (Section 4.2). Numbers in square brackets represent the standard deviation. For 2 Layer models, use the final layer MLP activations for MLP((a,b) b).
We omit test accuracy because every model on this list except for the dropout p = 0.2 mod- els achieves > 99.95% test accuracy, while the dropout p 0.2 models achieve around 99.6% test accuracy.
Model Type
Loss
Accuracy
Ablated Loss
Ablated Acuracy
Varying Data Fraction
1.83 10 (1.65 5 10-1)
100%
7.74 10 (6.74 10
100%
Layer Transformer
1.97 (2.11
99.6%
4.63 10 (6.72 10
98.7%
P = 53
5.96 -0- (8.91
10-9)
100%
1.5 10.-00 10
100%
P 109
1.94 10-7 (1.10-7..7--111111 (3.74-100
17-4-0091
100%
6.53 10-7 (1.41 10-
100%
Dropout = 0.2
0.215 (0.091)
99.7%
0.205 (0.075)
99.7%
Dropout 0.5
4.68 10-3 (8.11 10-3)
100%
3.6 10-3 (5.82 10-3)
100%
Table We confirm that the skip connection around the final MLP layer is not important for perfor- mance by mean ablating the skip connection and computing loss and accuracy the entire dataset for each problem, averaged over all runs. (We report the standard deviation of loss over the runs in parentheses.) While loss does increase a small amount, accuracy remains consistently high and the loss of the ablated model remains low. Results with zero ablations are also similar.
27

---

Published as a conference paper at ICLR 2023
Excluded Loss Over Frequencies
100
Lain Lots
Text Loss
Excladed
0.01
Los0
100g
100
Epoch
Figure 26: Training, test, and full excluded loss for 1-layer version of the original architecture without weight decay. One random seed chosen for readability. Note that having weight decay prevents grokking.
Weight Decay 0.3
Weight Decay 3.0
10
Avenge Train Loss
Average
Log Loss
0.01
Log
60001
100g
1p
10p
Epoch
Epoch
Figure 27: The train and test loss the course of training with weight decay À 0.3 (left) and 3.0 (right). Less aggressive weight decay leads to slower grokking.
28

---

Published as a conference paper at ICLR 2023
Dropout P=2.2
L1 penalty, 1.0
100
Log Loss
Log Loss
01:33
Avengy Thii Lona
01:01
100g
Averayy
100g
100p
14
10k
20cc Epoch
30kk
40k
10k
20k Epoch
300
ank
Dropout 0.5.5
1.1 penalty, 10
100
Aceeage Average Test
Log Loss
Log Loss
0.03
Average Train Loss
0..01
7000
Averag
100g
10p
100p
14b
10k
20k
30k
45
10k
20k Epoch
30k
40k
Epoch
Dropout P=0.8
penalty, 100
100
Avecage Train Laves Aveeage Test Lows
100
Log Loss
11001
Average Tex Loss
0.01
Log
100g
10n
10p
100p
16
40k
Epoch
Epoch
Figure 28: The train and test loss over the course of training with two types of regularization, dropout and l regularization. Grokking occurs with some runs for dropout but never for l regularization.
2. & Regularization We add & penalty to the loss term. We use À {1, 10, 100}. Note that we do not decouple the updates with respect to the & penalty from optimization steps done with respect to the log loss (as is done for &a regularization via AdamW Loshchilov & Hutter (2017)).
In each case, we ran three random seeds. We show the results in Figure 28. While grokking did not occur with & regularization,
found that it does oceur for all three seeds using dropout with p = 0.2 or = 0.5. We speculate that this is because both dropout and weight decay encourage the network to spread out computation (which is required for the Fourier multiplication algorithm), while & regularization encourages the network to become more sparse in the neuron basis and thus less sparse in the Fourier basis, preventing the network from learning the Fourier Multiplication Algorithm.
D.2 THE SLINGSHOT MECHANISM OFTEN OCCURS, BUT IS UNNECESSARY FOR GROKKING
As noted in Section C.2, our 2-layer transformers exhibit significant slingshots (Thilak et al., 2022) during training. We speculate that this is due how gradients of different scale interact with adaptive
29

---

Published as a conference paper at ICLR 2023
Train/ Test Loss for 5 Digit Addition w/ Infinite Data
Loss
1.5
0.5
1000
1500
2000
2500
Steps
Token Train/Test Loss for Digit Addition Infinite Data
2.5
Token loss
Token loss
Token 2 loss
Loss
1.5
Token 3 loss
Token loss
Token 5 loss
0.5
500
1000
1500
2000
2500
Steps
Figure 29: (Top) The training/test loss for 5 Digit Addition trained on randomly generated data. Note that training and test loss coincide, as the model does not see repeated pairs. The train/test loss per token for Digit Addition, trained with randomly generated data each step. Note that phase changes in the average loss correspond to phase changes in individual tokens, though one phase change (token I, around step 270) is not visible on the averaged loss as it overlaps with the end of the first phase change (token O, starting around step 150).
optimizers. We were even able induce slingshots on a I-layer by reducing the precision of the loss calculations (as this causes many gradients to round to O and thus greatly increases the differences scale of gradients).
However, many of our 1-layer models do not exhibit slingshots but nonetheless grok, the slingshot mechanism is unnecessary for grokking to occur, in the presence of weight decay or other regular- ization. We speculate that the slingshots of Thilak et al. (2022) (which co-occur with grokking for training runs without weight decay) serve as an implicit regularization mechanism that favors the simpler, generalizing solution over the more complicated
D.3 ADDITIONAL EVIDENCE FROM OTHER ALGORITHMIC TASKS
We now provide addition analysis of grokking phenomena on 3 additional algorithmic tasks and confirm that limited data is an important part of grokking:
1. 5 digit addition. We sample pairs of random 5 digit numbers and have the model predict their sum
2. Predicting repeated subsequences. We take a uniform random sequence of tokens, ran- domly choose subsequence to repeat, and train the model to predict the repeated tokens.
3. Skip trigram. We feed in sequence of tokens from 0 to 19, of which exactly one is greater than or equal to 10, and the model needs to output the token that is > 10. This can be solved with learning 10 skip trigrams.
30

---

Published as a conference paper at ICLR 2023
Train/ T'stt Loss for Digit Addition with 700 Data Points
Test
Loss
Epochs
Figure 30: The train and test loss for Digit Addition trained on 700 data points. Unlike the infinite, randomly generated data case, this shows both sharp phase change and clear train test divergence.
We use 1-layer full transformer for 5-digit addition, a 2-layer attention only transformer for pre- dicting repeated subsequences, and I-layer attention only transformer for the skip trigram task. Otherwise, we use the same hyperparameters as in the mainline model.
5 Digit Addition We first consider the case where we train on the approximately infinite data regime. For each minibatch, we randomly new sample 5 digit numbers. We report the results in Figure 29. Train loss coincides with test loss, so grokking does not occur, as the model almost never sees the same pair of digit numbers twice, with 10 such pairs. Interestingly, the various small bumps in Figure 29 correspond to the model learning how to calculate each of the 6 tokens in the output. However, grokking does occur when we restrict the model to only see 700 data points, as shown in Figure 30.
Repeated subsequence As with the 5-digit addition task, we find that restricting the amount of data is necessary and sufficient for grokking the repeated subsequence task. In Figure 31, the model sees new data at every step exhibits no grokking. In contrast, clear grokking occurs when we restrict the model to only see 512 data points in Figure 32.
Skip trigram As with the previous tasks, we find that restricting the amount of data is necessary and sufficient for grokking on the skip trigram task. The model that sees new data at every step exhibits grokking in Figure 33. Meanwhile, the model restricted to only see 512 data points exhibits clear grokking in Figure 34.
Taken together, these results echo the importance of limited data for grokking.
E FURTHER SPECULATIONS ON GROKKING
E. AN INTUITIVE EXPLANATION OF GROKKING
In this section, we speculate on what might be happening "under the hood" when model groks and explore why this phenomena happens. The evidence is only suggestive, so this a promising direction for future research.
Grokking occurs when models, trained on algorithmic tasks with certain hyperparameters, initially overfit the training data where train loss significantly improves while test loss worsens and the two diverge. But later in training, there sudden improvement in test loss, so test and train loss converge. In contrast to Power al. (2022) but in line with Liu et al. (2022), grokking does not occur when both train and test loss improve together without the initial divergence, as shown in many of the figures in this paper, for example Figures 2 and 18.
The core issue is that the model has two possible solutions: memorization (with low train loss and high test loss) and generalization (with low train loss and low test loss). In case, the
31

---

Published as a conference paper at ICLR 2023
Repeated Subsequence Prediction w/ Infinite data
Loss
1000
2000
3000
4000
5000
Step
Figure 31: The training/test loss for repeated subsequences trained on randomly generated data. Note that training and test loss coincide, the model does not repeated pairs. There sharp phase change corresponds to the model forming induction heads. (Olsson et al., 2022)
Repeated Subsequence Prediction w/ 512 Data Points
10
Loss
Train Test Loss
50k
150k
200k
Epoch
Figure 32: The train and test loss for the repeated subsequence task, trained 512 data points. Unlike the infinite, randomly generated data case, this shows both sharp phase change and clear train test divergence.
Train/Test Loss for Skip Trigram Task w/ Infinite Data
2.5
1.5
500
1000
1500
2000
2500
Epoch
Figure 33: The training/test loss for the skip trigram task, trained on randomly generated data. Note that training and test loss coincide, as the model does not see repeated pairs. The sharp phase change corresponds to the network learning all of the skip trigrams.
32

---

Published as a conference paper at ICLR 2023
Train/ Loss for Skip Trigram Task w/ 50 Data Points
Loss
Test
Loss
1000
4000
Epoch
Figure 34: The train and test loss for the skip trigram task, trained on 512 data points. Unlike the infinite, randomly generated data case, this shows both a sharp phase change and clear train test divergence.
Fourier Multiplication Algorithm is the generalization solution. Intuitively, with very little training data, the model will overfit and memorize. With more training data, the model must generalize suffer poor performance on both train and test loss. Since neural networks have an inductive bias favoring "simpler" solutions, memorization complexity scales with the size of the training set, whereas generalization complexity constant. The two must cross some point! Yet, the surprising aspect of grokking is the abrupt shift during training, when the model switches from memorization to generalization.
The other component of grokking is phase transitions the phenomena where models trained on certain task develop specific capability fairly rapidly during a brief period of training, as shown for the case of induction heads forming in transformer language models in Olsson et al. (2022) and our results in Appendix D.3. That is, rather than slowly forming that capability over training, the model rapidly goes from being bad at it to being good at it. One interpretation of phase transition that there's some feature of the loss landscape that makes the generalising solution harder to reach rather than smooth gradient for the model to follow, it instead initially finds it difficult to make progress, but then crosses some threshold where it can rapidly make progress.
Therefore, grokking occurs with phase transitions, limited data, and regularization. Models exhibit phase transitions despite having enough training data to avoid overfitting. Regularization (weight decay in our case) favors simpler solutions marginally prefer generalization over memorization. The phase transition indicates that generaliza- tion is "hard to reach" while the model has no problems with memorization. But it memorizes, the network becomes more complex until the weight decay prevents further memorization then moves towards equilibrium. The gradient to memorize balances the gradient towards smaller weights. With generalization, the model is incentivized to both memorize and simplify. Strikingly, it is capable of both while maintaining somewhat constant training performance this eircuit formation phase. Next, as the model approaches generalization, the memorization weights are removed in the cleanup phase. The cost from complexity outweighs the benefit from lower loss. Due to the phase transition during this training period, as model's progress towards generalization accelerates, the cleanup rate sharpens as well.
complex ones. The model has enough data to
A model that learns a perfect solution and is trained with weight decay has competing incentives: larger weights (for more extreme logits and thus lower loss) and smaller weights (from weight decay). So for any solution and any level of weight decay, there will always be level of train loss where these two forces equilibrate. Thus, memorization is not necessarily a "simpler" solution than generalization. The key is that generalization will have smaller weights holding train loss fixed. In fact, weight decay should be expected to equilibrate at a slightly lower train loss in generalization, since the base solution is simpler. This matches what observe in practice."
*One subtlety: the grokking phenomena is often incorrectly summarized as "the model learned to generalize even after achieving zero loss." " Zem loss does not exist with cross-entropy loss. Although the model achieves
33

---

Published as a conference paper at ICLR 2023
E.2 HYPOTHESIS: PHASE TRANSITIONS ARE INHERENT TO COMPOSITION
A promising line of work in the growing field of mechanistic interpretability suggests that models form circuits (Cammarata et al., 2020) clean interpretable algorithms formed by subnetworks of the model, such as curve detectors (Cammarata et al., 2020) in image classification networks and induction heads (Elhage al., 2021; Olsson et al., 2022) in LLMs. This surprisingly true! A circuit represents the model learning an algorithm, a fundamentally discrete thing; each step in the algorithm only makes sense if the other steps are present. But neural networks are fundamentally continuous, trained to follow gradients towards lower loss and struggle to jump to new optima with- out following a smooth gradient. So how can a model learn a discrete algorithm?
As concrete example, let's consider the case of induction heads in LLMs. There is a subnetwork of a next-token prediction autoregressive language model that learns to continue repeated subse- quences. It detects whether the current token occurred earlier in the context. If so, it predicts the same token after that previous occurrence will also come next. The circuit consists of a previous token head, which attends to each previous token and copies the context of the previous token to the current token, and an induction head which attends to the token after a previous occurrence of the current token. The induction head composes with the previous token head by forming query vector representing the current token and key vector representing the previous token head's output using K-Composition, the context of the previous token. It attends to token where this query and key match.
This circuit significantly improves loss but only in the context of the other heads present. Before either head is present, no gradient encourages the formation of either head. At initialization, we have neither head, so gradient descent should never discover this circuit. Naively, we might predict that neural networks will only produce circuits analogous to linear regression, where each weight will marginally improve performance as it continuously trains. And yet in practice, neural networks indeed form such sophisticated circuits, involving several parts interacting in non-trivial, algorithmic ways. So how can this be?
A few possible explanations:
Lottery tickets (Frankle & Carbin, 2018): Initially, each layer of the network is the superposition of many partial circuit components, and the output of each layer is the average of the output of each component. The full output of the network is the average of many different circuits, with significant interference from non-linear interaction. Some of these circuits are systematically useful to reducing loss, but most aren't. Gradients for useless circuits will have zero mean, while gradients for useful circuits will have non-zero mean, with a lot of noise. SGD reinforces relevant circuits and suppresses useless ones, so circuits will gradually form.
The network wanders randomly around the loss landscape until it encoun- ters a half-formed previous token head and induction head that somewhat compose. This half-formed circuit becomes useful for reducing loss, so gradient descent completes the circuit.
Random walk:
Evolution: similar mystery arises from how organisms develop sophisticated machinery, like the human eye. Each part is only useful in the context of other parts. A compelling explanation is component first developed that was somewhat useful in its own right, like a light-detecting membrane. It was reinforced as a useful component. Then, later compo- nents developed depending on the first, like the lens of the eye.
Evolution is a natural explanation, However, based on toy tasks, it cannot be the whole story. In the repeated subsequence task, we have a sequence of uniform randomly generated tokens, apart from a repeated subsequence at arbitrary location, e.g. 7283 1 9 1 9925 5 END. This means all pairs of tokens are independent, apart from pairs of equal tokens in the repeated subsequence. In particular, this means that a previous token head can never reduce loss for the current token. The previous token will always be independent of the next token. So previous token head is only useful in the context of induction-like head that completes the circuit. Likewise, induction head relies
perfect accuracy, it is trained to optimize loss not accuracy. This means the model is always incentivized to further improve. In particular, the casiest way to improve performance with perfect accuracy is by scaling up the logits. This lowers the temperature and pushes the softmax closer argmax.
34

---

Published as a conference paper at ICLR 2023
on K-composition with previous token head and cannot be useful on its own. Yet the model eventually forms an induction circuit!
A priori, the random walk seems insufficient on its own. An induction circuit is relatively compli- cated, representing a small region in model space. So a random walk is unlikely to stumble upon it. Concretely, in modular addition case, progress measures show significant hidden progress pre-grokking, indicating the model did not stumble upon the solution by chance.
Thus, the lottery ticket hypothesis seems the most explanatory. An induction head is useless without a previous token head but might be slightly useful when composing with a head that uniformly attends to prior tokens, since part of its output will include the previous token! Nevertheless, suspect that all explanations contribute to the entire picture. This seems most plausible if the uniform head just so happens to attend bit more to the previous token via a random walk.
Returning to phase transitions, the lottery ticket-style explanation suggests that we might expect phase transitions as circuits form. Early in circuit formation, each part of the eircuit is rough, so the effect on the loss of improving any individual component is weak, meaning gradients will be small. As each component develops, other components will become more useful, meaning all gradients will increase together non-linearly. As the circuit nears completion, we should expect an acceleration in the loss curve for this circuit, resulting in a phase transition.
F FURTHER DISCUSSION ON USING MECHANISTIC INTERPRETABILITY AND PROGRESS MEASURES FOR STUDYING EMERGENT PHENOMENA
While we find approach of using mechanistic interpretability to define progress measures rela- tively promising. there remains significant uncertainty as to how scalable existing mechanistic inter- pretability approaches really are. Broadly speaking, depending on the success of future mechanistic interpretability work, we think there are three methods through which mechanistic interpretability and progress measures can help with understanding and predicting emergent phenomena:
If mechanistic interpretability can be scaled to large models to the level where we can un- derstand the mechanisms behind significant portions of their behavior, we could perform the same style of analysis as was done in this work. We believe it's currently unclear as to whether not mechanistic interpretability will successfully scale to large models to this extent (or even if there exist human- uumaa--nderrtandabll explanations for all of their sophisti- cated behavior). That being said, in cases where mechanistic interpretability does recover human-understee understandable mechanisms, we could simply use the parts of the mechanism as progress measures.
2. If future mechanistic interpretability models (as in Wang et al. (2022)) and can only generate comprehensive understanding of the mechanisms of smaller models, might still be able from smaller models to guide the development measures that track parts of the behavior of the larger model. We find this scenario relatively plausible, as existing mechanistic inter- pretability work already allows us to recover fragments of large model behavior and un- derstand these fragments by analogy to smaller models. For example, Olsson et al. (2022) use this approach to understand the emergence of in-context learning in medium- sized lan- guage transformers.
only recover parts of the mechanism of larger
our understanding
3. Even if mechanistic interpretability fails to recover understandable mechanisms at all on large models, we might still be able to derive progress measures that don't 't require human understanding. For example, if we end up with automated mechanistic interpretability (that nonetheless still fails to recover human Homann-unde understandable mechanisms). we might be able to use the outputs of those opaque processes.
Another approach is task-independent progress measures: if we can discover progress mea- sures that don't it depend on the task, perhaps using many small, interpretable models as testbeds, we might be able to apply these progress measures to large models.
That being said, think the future work outlined in Section 6 is necessary to successfully apply our approach to predict and understand emergent behavior in existing large language models, and so remain cautiously optimistic.
35


Reviews:
[{"review_id": "sy6m2AaLCpS", "reviewer": "Reviewer_u853", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "**Update after rebuttal** The authors have addressed my comments, and, in my view, the issues raised by the other reviewers to a satisfying degree. Given all this information I remain strongly convinced that this paper should be presented at the conference.\n\nThe paper investigates grokking on a modular addition task. As has been previously reported, training a large enough transformer on this task leads to a sudden jump from very low test-set accuracy to very high test accuracy, long after training error has converged to essentially zero. The phenomenon can be produced robustly, but has so far been not very well understood. The paper first identifies the algorithm implemented by the network to strongly generalize beyond the training set. Based on this insight, two specific metrics (progress measures) are developed that can be used to diagnose the onset and transition from a memorizing solution to the grokking solution. Using these measures the paper sheds some light on the grokking phenomenon by identifying three phases (memorization, circuit formation where the generalizing solution smoothly starts to form, and cleanup where the grokking solution takes over and weight decay pushes the memorizing weights towards zero).", "summary_of_the_review": "The paper studies a timely and potentially very interesting phenomenon: the phenomenon that sometimes (transformer) networks can learn (simple) algorithms that lead to strong generalization beyond the training data, despite initially completely overfitting the training data (memorizing). The paper adds to this by performing an excellent case study of analyzing one such solution in great detail; ultimately identifying the algorithm learned by the network. To me personally, this analysis itself is enough to warrant publication - it is another piece of work showing that reverse-engineering the algorithms implemented by neural networks is actually possible. Based on this analysis two diagnostic metrics for continuously characterizing grokking during training are developed, which provide very insightful results (onset of circuit formation after memorizing, switching over to grokking solution only after circuit is well developed but then with fairly rapid transition). These results, and the methodology introduced (which is admittedly quite labor intensive) open up the venue for studying the precise conditions under which the grokking solution is formed more generally; which might lead to improved training methods that robustly produce strongly generalizing solutions. I do not have any major concerns, and left some small comments and suggestions for improvements. Currently I am in favor of accepting the paper at the conference - I think it is of great interest to large parts of the community.", "strength_and_weaknesses": "**Main contributions, Impact**\n1) Identification and analysis of the algorithm that the grokking transformer uses to achieve strong generalization beyond the training set. The main finding is that the network learns to implement modular addition via projections onto and rotations on a circle (Fourier transformations and trigonometric identities). The analysis is sophisticated, intricate, and very well executed - including ablations and controls to verify many steps of the analysis. Impact: while understanding the particular algorithm is interesting (and certainly labor-intensive), I think the even more important finding is to have another rare case-study of deciphering the inner workings of a neural network that generalizes well, to reveal an algorithmically elegant and sensible solution; the work shows that “opening the black box” is indeed possible. Because the case study is so well executed I would rate its potential impact as a widely-cited and pioneering success case high.\n\n2) Design of two metrics, that allow to diagnose and track grokking smoothly over training. The introduction somewhat suggests that such “progress measures” can potentially be used to predict spontaneously emerging capabilities in AI systems in advance. While the paper demonstrates a concrete instance of the existence of such measures, I would almost interpret the current findings as somewhat of a negative result for the goal of developing general measures for predicting emergent capabilities: the progress measures could only be developed in hindsight (after a network with such capabilities already exists) and are highly task-specific. Impact: medium to low\n\n3) Abstract characterization of the grokking phenomenon (via the introduced progress measures), leading to the identification of three separate phases (memorization, circuit formation, cleanup). To me personally the most impactful contribution of the paper - it sheds some light and raises interesting hypotheses on the grokking phenomenon in general, beyond the individual task(s) studied. While it still remains somewhat unclear why the grokking solution cannot develop without the memorization phase, I think the paper characterizes some important parts of that question (onset of circuit formation after memorization, smooth development of circuit while memorizing part is still needed, finally fairly rapid switch over to circuit along with removal of now superfluous memorized information). Impact: medium to potentially very high - if we (as a community) could figure out how to obtain the grokking solution fast and reliably and without going through the memorization phase for a broad range of tasks, then we might see a generational leap in neural network training, potentially enabling strong and robust generalization capabilities that seem out of reach currently. Time will tell whether this work contributes important pieces to that puzzle, but I personally think it raises some very interesting observations, and follow-up questions.\n\n**Strengths**\n *  Sophisticated and intricate mechanistic interpretability case-study. Would itself be sufficient for a publication for me.\n *  Investigation of grokking - confirming and rejecting some previously proposed hypotheses, and raining interesting new ones.\n * Most importantly: confirming the intuition that grokking consists of developing a memorizing solution rapidly, while a slower process builds a strongly generalizing solution in a much smaller sub-network in parallel, and finally a fairly rapid switch away from the memorizing solution once the circuit is fully developed.\n\n**Weaknesses**\n * The paper is limited to being a (well executed) case-study around grokking and mechanistic analysis for a particular task (three more tasks in the appendix and various ablations). This makes very general statements about grokking hard to make; and allows the paper to mainly raise, rather than confirm, hypotheses about grokking in general. Having said that, I think the paper “packs” more than enough to warrant a publication.\n * The current discussion could be a bit expanded and perhaps also be a bit more critical - the intro sounds a somewhat like the current paper is a big leap towards general/task-independent progress measures that allow predicting the onset of arbitrary capabilities (yet the current paper highlights the challenges which make such general measures unlikely to exist). \n * Generally very well written, but a few typos, and the appendix needs perhaps another quick pass.\n\n**Improvements**\n1) Does training (with different random seeds) always identify the same key frequencies? If yes, is there a hypothesis why that is the case; if no, please mention it and comment on this being a limitation for the design of more general progress measures.\n\n2) No need to add much, but it would be great to see a slightly expanded critical discussion (particularly the last two points). What does the current paper contribute to the development of (general) progress measures, and how hard/easy are the current obstacles?\n\n3) Suggestion (but given the time-frame I do not expect to see these results): identify the “winning lottery ticket” before and after grokking (i.e. remove any unnecessary weights from the network). (I) What is the size of these tickets (and how does it change during circuit formation and cleanup) - related to Fig. 6 bottom right? (II) If the winning ticket after grokking is re-initialized with its original initialization values (and all other parameters are kept at zero during training), can the network recover the grokking solution (without going through a memorization phase; perhaps even faster than in the original training run; in other words: is it enough to identify the correct sub-circuit structure for the grokking solution or do we need the memorization-part of the network during training such that the grokking solution can develop)? \n\n**Minor comments**\n\nA) The paper talks about the emergence of “capabilities”: maybe the phrasing can be made more precise. If a model can memorize the training data; in some sense it has already acquired (partial) “capabilities”. Maybe better to say that what is meant is capabilities that *generalize* beyond the training data (i.e. the network learns a *general* algorithm to implement the capabilities rather than a partial one that relies mainly on memorization).\n\nB) Typos:\n\nP1, typo: “that detect earlier repeated subsequences.”\n\nP2, punctuation: “toward the generalizing algorithm Liu et al.”\n\nP3, Section 3: Footnote 2 is used twice (once erroneously it seems)\n\nP5, typo: “should contain be rank”\n\nP5, missing space: “Table 3:projecting”\n\nP6, typo: “yet we can the logits”\n\nP8, typo (since -> sense): “it makes since to check”\n\nFig 6: use same color for Test and Train loss in both top panels\n\nFig 13, typo (90% twice): “While the majority of the runs exhibit grokking, runs with 90% or 90% of the data do not exhibit grokking”\n\nC) Titles of some appendix plots are cut.\n\n\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "lHpHAW1q-h8", "reviewer": "Reviewer_WFHw", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper presents a technique for analysing transformer networks by reverse engineering the behaviors of a simple network into their constituent components, similarly to the mechanistic interpretability employed in e.g. (Elhage et al. 2021). The technique is applied on a network trained with known trigonometric functions that lends themselves for relatively transparent analysis. The paper purportedly shows distinct memorization and 'grokking' phases of learning happening in the target network. Finally, with thus framed progress in learning, 'progress' metrics are defined and measured such that they align with that progress.", "summary_of_the_review": "While I am genuinely concerned about generalizability of the results with this simple training setup, I would approve this as a case where one has to work upwards from a simple initial model. I consider it sufficient that the paper suggests a specific methodology and illustrates it with the given limited example. I would not claim the authors actually prove that those properties and training dynamics are applicable across a wider range of networks and training scenarios, but anyway that is not necessary in this case.\n\nHence, I find this an interesting paper that may contribute at least to the transformer interpretation methodologies, and possibly to other networks as well.", "strength_and_weaknesses": "Strengths:\n- The paper studies a critical problem of interpreting transformer networks.\n- The presentation is well-written and logical.\n- The method is intuitive and yields non-trivial results that may provide tools for interpreting larger networks.\n\nWeaknesses:\n- My major concern is the generalizability of the approach to more complex networks and training data. There is a clear risk that these clean results cannot be applied outside toy problems like the one presented here.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "gP82gSwaPEV", "reviewer": "Reviewer_g9KN", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper investigates grokking, i.e., the abrupt phase change in the performance of transformers when trained on simple algorithmic tasks, through the lens of mechanistic interpretability. It uses the insights gained to propose new progress measures that vary smoothly throughout training and thus provide new insights into the different training phases. Concretely, the paper provides empirical evidence to show that transformers learn to perform modular addition of two numbers via rotation around a circle, i.e., Fourier multiplication. The paper shows that various components of the learned network can be approximated with the different stages of the Fourier multiplication algorithm. Furthermore, the paper conducts ablation studies to confirm the faithfulness of the proposed approximations. Finally, the paper combines these insights to show that the network first memorizes the training data, then smoothly interpolates between memorization and the (sparse) Fourier algorithm (driven by weight decay), and finally removes the memorization component.", "summary_of_the_review": "Given the importance of the results, the rigor of the analysis, and the potential impact of the work (see strengths above), I recommend acceptance.", "strength_and_weaknesses": "**Strengths**\n\nThe paper presents a landmark study on the mechanistic interpretability of transformers. It prevents strong empirical evidence (approximation and ablation) that transformers learn to perform modular addition via Fourier multiplication. Thus, even though the paper does not formally prove the model (which is not the goal of this work), it sets an example in conducting a rigorous analysis of a model's behavior. Discovering the Fourier multiplication algorithm by looking at the network's weights requires a significant cognitive leap and careful analysis. Moreover, an understated aspect of the paper is that it simplifies the grokking setup (i.e., one layer, only modular addition with a small prime, etc.) to make it amenable to analysis while maintaining the relevant grokking phenomenon. Finally, the paper opens many follow-up research directions (see questions below).\n\n\n**Weaknesses**\n\nThe paper appears to be somewhat hastily written, as there are a couple of (minor) mistakes:\n* Introduction: `that detect earlier repeated subsequences` is redundant.\n* Related Work: `toward the generalizing algorithm*.* Liu et al. (2022)`.\n* Section 4.2: `should ~contain~ be rank 10`.\n* Section 4.2: `yet we can *approximate* the logits`.\n* Section 4.3: What does `we find that the attention pattern of the = to a for each of the four heads` mean?\n* Section 4.4: `depend only *on* these 10 directions`.\n* Section 5.1: `it makes ~since~*sense* to check`.\n\n\n**Questions**\n\nNote that these are mostly follow-up research questions that do not have to be answered in the context of this work.\n\n* Does grokking also occur with sinusoidal positional encodings, which should be compatible with the Fourier multiplication algorithm?\n* How does the model effectively interpolate between the memorizing and generalizing solutions? Are these implemented as two separate subcircuits?\n* Do other regularization methods lead to different generalization algorithms? Fourier multiplication and weight decay are closely linked due to the sparsity of the required algorithm.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%