Progress measures for grokking via mechanistic interpretability — part 1/6

Cleaned-text code points [0, 15546). Bibliography omitted here; source.md includes it.

# Progress measures for grokking via mechanistic interpretability

###### Abstract

Neural networks often exhibit emergent behavior, where qualitatively new capabilities arise from scaling up the amount of parameters, training data, or training steps. One approach to understanding emergence is to find continuous progress measures that underlie the seemingly discontinuous qualitative changes. We argue 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 performing ablations in Fourier space. Based on this understanding, we define progress measures that allow us to study the dynamics of training and split training into three continuous phases: memorization, circuit formation, and cleanup. Our results show that grokking, rather than being a sudden shift, arises from the gradual amplification of structured mechanisms encoded in the weights, followed by the later removal of memorizing components.

## 1 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 to a number of breakthroughs, via capabilities such as in-context learning ( Radford et al., 2019 ; Brown et al., 2020 ) and chain-of-thought prompting ( Wei et al., 2022b ) . However, it also poses risks: Pan et al. (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 it is abrupt, as in the case of reward hacking, chain-of-thought reasoning, or 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 are causally linked to 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, we introduce a different approach to uncovering hidden progress measures: via mechanistic explanations . 1 1 1 Interactive versions of figures, as well as the code to reproduce our results, are available at https://neelnanda.io/grokking-paper . A mechanistic explanation aims to reverse engineer the mechanisms of the network, generally by identifying the circuits ( Cammarata et al., 2020 ; Elhage et al., 2021 ) within a model that implement a behavior. Using such explanations, we study grokking , where models abruptly transition to a generalizing solution after a large number of training steps, despite initially overfitting ( Power et al., 2022 ) . Specifically, we study modular addition, where a model takes inputs a , b ∈ { 0 , … , P − 1 } a,b\in\{0,\ldots,P-1\} for some prime P P and predicts their sum c c mod P P . Small transformers trained with weight decay on this task consistently exhibit grokking (Figure 2 , Appendix C.2 ).

We reverse engineer the weights of these transformers and find that they perform this task by mapping the inputs onto a circle and performing addition on the circle. Specifically, we show that the embedding matrix maps the inputs a , b a,b to sines and cosines at a sparse set of key frequencies w k w_{k} . The attention and MLP layers then combine these using trigonometric identities to compute the sine and cosine of w k ​ ( a + b ) w_{k}(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 L W_{L} is well approximated by a 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 2 polynomials of trigonometric functions of a single frequency; and (4) ablating key frequencies used by the model reduces performance 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 modular addition task— restricted 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 formation , 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 a 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 qualitative behaviors as they are scaled up or trained longer ( Ganguli et al., 2022 ; Wei et al., 2022a ) . McGrath et al. (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 be 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 use 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 as an implicit regularizer.

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 —metrics 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 mechanistic intepretability to discover progress measures empirically.

## 3 Setup and Background

We train transformers to perform addition mod P P . The input to the model is of the form “ a ​ b = a\ b\ \!= ”, where a a and b b are encoded as P P -dimensional one-hot vectors, and = = is a special token above which we read the output c c . In our mainline experiment, we take P = 113 P=113 and use a one-layer ReLU transformer, token embeddings with d = 128 d=128 , learned positional embeddings, 4 4 attention heads of dimension d / 4 = 32 d/4=32 , and n = 512 n=512 hidden units in 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 113\cdot 113 pairs of numbers mod P P ). We use full batch gradient descent using the AdamW optimizer ( Loshchilov & Hutter, 2017 ) with learning rate γ = 0.001 \gamma=0.001 and weight decay parameter λ = 1 \lambda=1 . We perform 40,000 40,000 epochs of training. As there are only 113 ⋅ 113 113\cdot 113 possible pairs, we evaluate test loss and accuracy on all pairs of 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 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, we follow the conventions and notations laid out in Elhage et al. (2021) . We focus on the d × p d\times p embedding matrix W E W_{E} , the d × n d\times n output matrix of the MLP layer W o ​ u ​ t W_{out} , and the P × d P\times d unembedding matrix W U W_{U} . 2 2 2 We ignore the embedding and unembedding of the ‘ = = ’ token for simplicity. Let Logits ​ ( a , b ) \textrm{Logits}(a,b) denote the logit vector on inputs a , b a,b , and M ​ L ​ P ​ ( a , b ) MLP(a,b) denote the MLP activations. Empirically, our networks do not significantly use the skip connection around the MLP (Appendix A.1 ), so Logits ​ ( a , b ) ≈ W U ​ W o ​ u ​ t ​ MLP ​ ( a , b ) \textrm{Logits}(a,b)\approx W_{U}W_{out}\textrm{MLP}(a,b) . We therefore also study the P × n P\times n neuron-logit map W L = W U ​ W o ​ u ​ t W_{L}=W_{U}W_{out} .

### 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 , b a,b map these to sin ⁡ ( w k ​ a ) \sin({w_{k}a}) , cos ⁡ ( w k ​ a ) \cos(w_{k}a) , sin ⁡ ( w k ​ b ) \sin(w_{k}b) , and cos ⁡ ( w k ​ b ) \cos({w_{k}b}) using the embedding matrix, for various frequencies w k = 2 ​ k ​ π P , k ∈ ℕ w_{k}=\frac{2k\pi}{P},k\in\mathbb{N} .

• Compute cos ⁡ ( w k ​ ( a + b ) ) \cos\left(w_{k}(a+b)\right) and sin ⁡ ( w k ​ ( a + b ) ) \sin\left(w_{k}(a+b)\right) using the trigonometric identities: cos ⁡ ( w k ​ ( a + b ) ) \displaystyle\cos\left(w_{k}(a+b)\right) = cos ⁡ ( w k ​ a ) ​ cos ⁡ ( w k ​ a ) − sin ⁡ ( w k ​ a ) ​ sin ⁡ ( w k ​ b ) \displaystyle=\cos\left(w_{k}a\right)\cos\left(w_{k}a\right)-\sin\left(w_{k}a\right)\sin\left(w_{k}b\right) sin ⁡ ( w k ​ ( a + b ) ) \displaystyle\sin\left(w_{k}(a+b)\right) = sin ⁡ ( w k ​ a ) ​ cos ⁡ ( w k ​ b ) + cos ⁡ ( w k ​ a ) ​ sin ⁡ ( w k ​ b ) \displaystyle=\sin(w_{k}a)\cos\left(w_{k}b\right)+\cos\left(w_{k}a\right)\sin\left(w_{k}b\right) In our networks, this is computed in the attention and MLP layers.

• For each output logit c c , compute cos ⁡ ( w k ​ ( a + b − c ) ) \cos\left(w_{k}(a+b-c)\right) using the trigonometric identity: cos ⁡ ( w k ​ ( a + b − c ) ) \displaystyle\cos\left(w_{k}(a+b-c)\right) = cos ⁡ ( w k ​ ( a + b ) ) ​ cos ⁡ ( w k ​ c ) + sin ⁡ ( w k ​ ( a + b ) ) ​ sin ⁡ ( w k ​ c ) . \displaystyle=\cos\left(w_{k}(a+b)\right)\cos\left(w_{k}c\right)+\sin\left(w_{k}(a+b)\right)\sin\left(w_{k}c\right). (1) This is a linear function of the already-computed values cos ⁡ ( w k ​ ( a + b ) ) \cos(w_{k}(a+b)) , sin ⁡ ( w k ​ ( a + b ) ) \sin(w_{k}(a+b)) and is implemented in the product of the output and unembedding matrices W L W_{L} .

• The unembedding matrix also adds together cos ⁡ ( w k ​ ( a + b − c ) ) \cos\left(w_{k}({a+b-c})\right) for the various k k s. This causes the cosine waves to constructively interfere at c ∗ = a + b mod p c^{*}=a+b\mod p (giving c ∗ c^{*} a large logit), and destructively interfere everywhere else (thus giving small logits to other c c s).

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, we describe four lines of evidence that our transformers are using the Fourier multiplication 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 number 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.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 W L W_{L} is (approximately) rank 10, where each direction corresponds to the cosine or sine of one of 5 key frequencies. Projecting the MLP activations onto the components of W L W_{L} approximately produces multiples of the functions cos ⁡ ( w k ​ ( a + b ) ) \cos\left(w_{k}(a+b)\right) and sin ⁡ ( w k ​ ( a + b ) ) \sin\left(w_{k}(a+b)\right) , 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 are well-approximated by degree- 2 2 polynomials of sines and cosines at a single frequency. Moreover, the corresponding direction in W L W_{L} 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 to 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 a function of the input tokens.

Periodicity in the embeddings. We start by examining the embeddings. We apply a Fourier transform along the input dimension of the embedding matrix W E W_{E} then compute the ℓ 2 \ell_{2} -norm along the other dimension; results are shown in Figure 3 . We plot only the components for the first 56 frequencies, as the norm of the components for frequencies k k and P − k P-k are symmetric. The embedding matrix W E W_{E} is sparse in the Fourier basis–it only has significant nonnegligible norm at 6 frequencies. Of these frequencies, only 5 appear to be used significantly in later parts of the model (corresponding to k ∈ { 14 , 35 , 41 , 42 , 52 } k\in\{14,35,41,42,52\} ). We dub these the key frequencies of the model.