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

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

In this work, we use mechanistic interpretability to define progress measures for small transformers trained on a modular addition task. We find that the transformers embed the input onto rotations in ℝ 2 \mathbb{R}^{2} and compose the rotations using trigonometric identities to compute a + b mod 113 a+b\mod 113 . Using our reverse-engineered 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 a proof of concept for using mechanistic interpretability to understand emergent behavior.

Larger models and realistic tasks . In this work, we studied the behavior of small transformers on a simple algorithmic task, solved with a 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 metrics are specific to small networks on one particular algorithmic task. Methods for automating the analysis and finding task-independent progress measures seem necessary to 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 increase 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 us 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.

## Reproducibility Statement

An annotated Colab notebook containing the code to replicate our results, including download instructions for model checkpoints, is available at https://neelnanda.io/grokking-paper .

## Author Contributions

Neel Nanda was 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 in Section 4 . He also discovered the restricted and excluded loss progress measures and that grokking in mainline model could be 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 a 1L Transformer on 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 hyper-parameters.

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 to thank Adam Gleave and Chengcheng 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 et 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 to 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 Anthropic 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.

## Appendix 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 it in our specific case.

We denote our hyperparameters as follows: d v ​ o ​ c ​ a ​ b = 113 d_{vocab}=113 is the size of the input and output spaces (treating ‘ = = ’ separately), d m ​ o ​ d ​ e ​ l = 128 d_{model}=128 is the width of the residual stream (i.e. embedding size), d h ​ e ​ a ​ d = 32 d_{head}=32 is the size of query, key and value vectors for a single attention head, and d m ​ l ​ p = 512 d_{mlp}=512 is the number of neurons.

We denote the parameters as follows: W E W_{E} (embedding layer); W p ​ o ​ s W_{pos} (positional embedding); W Q j W_{Q}^{j} (queries), W K j W_{K}^{j} (keys), W V j W_{V}^{j} (values), W O j W_{O}^{j} (attention output) (the 4 weight matrices of head j j in the attention layer); W i ​ n W_{in} and b i ​ n b_{in} for the input linear map of the MLP layer; W o ​ u ​ t W_{out} and b o ​ u ​ t b_{out} for the output linear map of the MLP layer; and W U W_{U} (unembedding layer). Note that we do not have biases in our embedding, attention layer or unembedding, and we do not tie the matrices for the embedding/unembedding layers.

We now describe the mathematical structure of our 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 to the final token. We use t i t_{i} to denote the token in position i i (as a one-hot encoded vector), p i p_{i} to denote the i i th positional embedding, x i ( 0 ) x^{(0)}_{i} to denote the initial residual stream on token with index i i , A ( i ) A^{(i)} to denote the attention scores from = = to all previous tokens from head i i , x ( 1 ) x^{(1)} to 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, x ( 2 ) x^{(2)} the final residual stream on the final token, Logits the logits on the final token.

The logits are calculated via the following equations:

x i ( 0 ) \displaystyle x_{i}^{(0)} = W E ​ t i + p i \displaystyle=W_{E}t_{i}+p_{i} A j \displaystyle A^{j} = softmax ​ ( x ( 0 ) T ​ W K j T ​ W Q j ​ x 2 ( 0 ) ) \displaystyle=\textrm{softmax}(x^{(0)^{T}}W_{K}^{j^{T}}W_{Q}^{j}x^{(0)}_{2}) x ( 1 ) \displaystyle x^{(1)} = [ ∑ j W O j ​ W V j ​ ( x ( 0 ) ⋅ A j ) ] + x 2 ( 0 ) \displaystyle=[\sum_{j}W_{O}^{j}W_{V}^{j}(x^{(0)}\cdot A^{j})]+x^{(0)}_{2} MLP = ReLU ​ ( W i ​ n ​ x ( 1 ) ) \displaystyle=\textrm{ReLU}(W_{in}x^{(1)}) x ( 2 ) \displaystyle x^{(2)} = W o ​ u ​ t ​ N + x ( 1 ) = W o ​ u ​ t ​ ReLU ​ ( W i ​ n ​ x ( 1 ) ) + x ( 1 ) \displaystyle=W_{out}N+x^{(1)}=W_{out}\textrm{ReLU}(W_{in}x^{(1)})+x^{(1)} Logits = W U ​ x ( 2 ) \displaystyle=W_{U}x^{(2)}

As in Elhage et al. (2021) , we refer to the term W O j ​ W V j ​ ( x ( 0 ) ) W_{O}^{j}W_{V}^{j}(x^{(0)}) as the OV circuit for head j 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 0.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, if 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 − 7 2.4\cdot 10^{-7} to 9.12 ⋅ 10 − 7 9.12\cdot 10^{-7} and 7.25 ⋅ 10 − 7 7.25\cdot 10^{-7} respectively. This is a significant increase in loss, but from such a small baseline that we can still ignore it 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: zero 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. a sigmoid over the difference. And x 2 ( 0 ) x^{(0)}_{2} is constant, as it is independent of x x and y y , and the embedding and positional embedding of ‘ = = ’ are fixed. So A 0 j = σ ⁡ ( x 2 ( 0 ) T ​ W Q j T ​ W K ​ ( x 0 ( 0 ) − x 1 ( 0 ) ) ) A^{j}_{0}=\sigma\left({x^{(0)}_{2}}^{T}W_{Q}^{j^{T}}W_{K}(x^{(0)}_{0}-x^{(0)}_{1})\right) (and A 1 j = 1 − A 0 j A^{j}_{1}=1-A^{j}_{0} )

A consequence of the second observation is that Logits ≈ W U ​ W o ​ u ​ t ​ MLP \textrm{Logits}\approx W_{U}W_{out}\textrm{MLP} , which we denote as W L = W U ​ W o ​ u ​ t W_{L}=W_{U}W_{out} . From the perspective of the network, W L W_{L} is the meaningful matrix, not either of its constituents, since they compose linearly.

## Appendix 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 a function whose value at x = 0 mod 113 x=0\mod 113 is significantly larger than any other x x .

For example, consider the function f 14 ​ ( x ) = cos ⁡ ( 2 ​ π ⋅ 14 113 ​ x ) f_{14}(x)=\cos\left(\frac{2\pi\cdot 14}{113}x\right) . This function has period 113 and is maximized at x = 0 mod 113 x=0\mod 113 . However, other values of x x cause this function to be close to 1: f 14 ​ ( 8 ) = f 14 ​ ( 105 ) = 0.998 f_{14}(8)=f_{14}(105)=0.998 , f 14 ​ ( 16 ) = f 14 ​ ( 89 ) = 0.994 f_{14}(16)=f_{14}(89)=0.994 , etc.

Now consider f 35 ​ ( x ) = cos ⁡ ( 2 ​ π ⋅ 35 113 ​ x ) f_{35}(x)=\cos\left(\frac{2\pi\cdot 35}{113}x\right) . While this function also has period 113 and is maximized at x = 0 mod 113 x=0\mod 113 , it turns out that f 35 ​ ( 8 ) = f 35 ​ ( 105 ) = − 0.990 f_{35}(8)=f_{35}(105)=-0.990 . This means that by adding together f 14 f_{14} and f 35 f_{35} , we end up with a function that is not close to 1 at x = 8 mod 113 x=8\mod 113 . Similarly, while f 35 ​ ( 16 ) = 0.961 f_{35}(16)=0.961 , f 52 ​ ( 16 ) = − 0.56 f_{52}(16)=-0.56 , and so adding a third frequency reduces the peak at x = 16 mod 113 x=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 .

## Appendix C Supporting evidence for mechanistic analysis of modular arithmetic networks

### C.1 Further analysis of the specific training run discussed in the paper

In this section, we provide additional evidence relating to the mainline model.

#### 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 a for all 4 attention heads, as a heatmap over the inputs a a and b b , as this is a scalar for each head. We observe a striking periodicity and further that heads 1 and 3 represent the same frequency while heads 0 and 2 are different.

As shown in Appendix A.1 , the attention paid from ‘ = = ’ to itself is negligible, so A 0 j = 1 − A 1 j A^{j}_{0}=1-A^{j}_{1} and it suffices to plot attention to a a .

#### C.1.2 Approximating attention heads with Sines and Cosines

Attention heads approximately compute degree- 2 2 polynomials of a single frequency or are used to amplify W E W_{E} . In order to compute terms like cos ⁡ ( w k ​ ( a + b ) ) \cos\left(w_{k}(a+b)\right) , the model needs to compute the product of the sine and cosine embeddings output by W E W_{E} . As the attention heads are approximately 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 a single frequency w k w_{k} . 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 a linear function actually improves performance). Thus the attention weights multiply with the OV output to create degree- 2 2 polynomials of the frequency w k w_{k} , as would be needed for 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 of a single frequency.

The periodicity of the attention heads has a striking form— A 0 j A^{j}_{0} is well approximated by 0.5 + α j ​ ( cos ⁡ ( w k ​ a ) − cos ⁡ ( w k ​ b ) ) + β j ​ ( sin ⁡ ( w k ​ a ) − sin ⁡ ( w k ​ b ) ) 0.5+\alpha^{j}(\cos(w_{k}a)-\cos(w_{k}b))+\beta^{j}(\sin(w_{k}a)-\sin(w_{k}b)) , for some frequency w k w_{k} and constants α j \alpha^{j} and β j \beta^{j} (which may differ for each head). Note further that this simplifies to 0.5 + γ ⁡ ( cos ⁡ ( w k ​ ( a + θ ) ) − cos ⁡ ( w k ​ ( b + θ ) ) ) 0.5+\gamma(\cos(w_{k}(a+\theta))-\cos(w_{k}(b+\theta))) for some constants γ \gamma and θ \theta . We show the coefficients and fraction of variance explained in Table 1

##### Mechanistic Analysis of Attention Patterns.

We can further mechanistically analyse how the model achieves this form. The following is a high-level sketch of what is going on:

First, note that the attention score on position 0 0 and head j j is just a lookup table on the input token a a (of size P P ). To see why, note that A 0 j = m ​ x 0 ( 0 ) T ​ W K j T ​ W Q j ​ x 2 ( 0 ) A^{j}_{0}={mx^{(0)}_{0}}^{T}W_{K}^{j^{T}}W_{Q}^{j}x^{(0)}_{2} . x 2 ( 0 ) x^{(0)}_{2} is constant since the token is always ‘ = = ’ and x 0 ( 0 ) = W E ​ t 0 + p 0 x^{(0)}_{0}=W_{E}t_{0}+p_{0} . So this reduces to t 0 ⋅ C j + D t_{0}\cdot C_{j}+D for some constant vector C j = W E T ​ W K j T ​ W Q j ​ x 2 ( 0 ) ∈ ℝ p C_{j}=W_{E}^{T}W_{K}^{j^{T}}W_{Q}^{j}x^{(0)}_{2}\in\mathbb{R}^{p} and some scalar D = p 0 T ​ W K j T ​ W Q j ​ x 2 ( 0 ) D=p_{0}^{T}{W_{K}^{j}}^{T}W_{Q}^{j}x^{(0)}_{2} . As t 0 t_{0} is one-hot encoded, this is just a lookup table, which we may instead denote as C j ​ [ a ] C_{j}[a]