P ROGRESS MEASURES FOR GROKKING VIA MECHANISTIC INTERPRETABILITY
Anonymous authors
Paper under double-blind review
A BSTRACT
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: reverseengineering 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 a sudden shift, arises from the combination of a gradual amplification of structured mechanisms encoded in the weights and the removal of components from the network that perform memorization.
I NTRODUCTION
Neural networks often exhibit emergent behavior, in which qualitatively new capabilities, often generalizing beyond the training data, arise from scaling up model size, training data, or the 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-ofthought 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 doesn’t explain why the phase changes happen.
In this work we introduce a different approach to uncovering hidden progress measures: via mechanistic explanations .
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.
We instantiate this strategy to study grokking , where overparameterized models abruptly transition to a generalizing solution after initially overfitting.
Specifically, we study modular addition, where a transformer takes inputs a, b ∈ { 0 , . .
, P − 1 } for some prime
P and predicts their sum c mod P .
Small transformers trained with weight decay on this task exhibit grokking (Appendix C.2).
We reverse engineer the weights of these transformers and find that they perform this task by projecting the inputs onto a circle and performing addition on the circle.
Specifically, we show that the embedding matrix maps the inputs a, b to sines and cosines at a sparse set of key frequencies w k .
The attention and MLP layers then combine these using trigonometric identities to compute the sine and cosine of w k ( a + b ) , and the output matrices shift and combine these frequencies.
We confirm this understanding with several lines of evidence.
For example, we show that the embedding/unembedding matrices consist of sines and cosines at key frequencies.
Most MLP neurons use a single key frequency and the activations and logits of the model are well approximated by components of the interpreted algorithm.
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 algorithm implemented by these transformers, we construct two measures of progress during training on this task— restricted loss , where we ablate every non-key frequency, and excluded loss , which measures performance with key frequencies removed.
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 of the mechanism that implements the generalizing algorithm; and cleanup , where weight decay removes the transformer’s memorizing algorithm.
Surprisingly, the sudden transition to perfect test accuracy in grokking occurs during cleanup, after the generalizing mechanism is learned.
While significant work remains for using this approach on larger networks and more realistic tasks, we believe that our results show that using mechanistic interpretability to define progress measures can further our understanding of emergent phenomena.
R ELATED W ORK
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.
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 construct 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.
Mechanistic interpretability and progress measures.
There is a growing literature that seeks to mechanistically understand neural networks.
The style of post-hoc mechanistic interpretability in Section 4 is heavily inspired by the Circuits approach of Cammarata
Elhage
(2021), and Olsson et
Meng et al.
(2022) use patching and restoration to locate the components in GPT-2 that encode factual knowledge.
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.
S ETUP AND B ACKGROUND
We train transformers to perform addition mod P .
The input to the model is of the form “ a b = ”, where a and b are encoded as P -dimensional one-hot vectors, and = is a special token above which we read the output c .
In our mainline experiment, we take P = 113 and use a one-layer ReLU transformer, token embeddings with d = 128 , learned positional embeddings, 4 attention heads of dimension d/ 4 = 32 , and 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 such problems (that is, 30% of the 113 ·pairs of numbers mod P ).
We use full batch gradient descent, weight decay (Hanson & Pratt, 1988) and the AdamW optimizer (Loshchilov & Hutter, 2017) with learning rate γ = 0 .
001 and weight decay parameter λ
= 1 .
We perform 40 , 000 epochs of training.
As there are only 113 · 113 possible pairs, we can evaluate the model’s performance on all points in the distribution.
Networks trained on this task consistently exhibit grokking (Power et al., 2022).
As Figure 2 shows, the 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 as test accuracy also increases to near 100%.
To denote our transformer, we follow the conventions and notations laid out in Elhage et al.
For our purposes, we focus on the embedding matrix W E (of shape d × P ), the second matrix in the
MLP layer W out (of shape d × n ), and the unembedding matrix W U (of shape P × d ).
Empirically, the transformer does not use the skip connection around the MLP layer as a significant part of its computation (Appendix A.1), so L ≈ W U W out N , where L are the output logits and N is the neuron activations (which are in turn a function of the embeddings and the attention heads).
We therefore also focus on the composition of W U and W out , W L = W U W out (of shape P × n ).
3. T HE F OURIER MULTIPLICATION ALGORITHM
We claim that the learned networks use the following algorithm:
• Given two one-hot encoded tokens a, b , for various k ∈ N , map these to sin( w k a ) , cos( w k a ) , sin( w k b ) , and cos( w k b ) using the embedding matrix, where w k = 2 kπ P .
• Compute cos ( w k ( a + b )) and sin ( w k ( a + b )) using the trigonometric identities: cos ( w k ( a + b ))
=
cos ( w k a ) cos ( w k a ) − sin ( w k a ) sin ( w k b ) sin ( w k ( a + b ))
=
sin( w k a ) cos ( w k b ) + cos ( w k a ) sin ( w k b )
In our networks, this is computed in the attention and MLP layers.
• For each output logit c , compute cos ( w k ( a + b − c )) using the trigonometric identity: cos ( w k ( a + b − c ))
=
cos ( w k ( a + b )) cos ( w k c )
+ sin ( w k ( a + b )) sin ( w k c ) .
This is a linear function of the already-computed values cos( w k ( a + b )) , sin( w k ( a + b )) and is implemented in the product of the output and unembedding matrices W L .
• The unembedding matrix also adds together cos ( w k ( a + b − c )) for the various k s. This causes the cosine waves to constructively interfere at c ∗ = a + b mod p (giving c ∗ a large logit), and destructively interfere everywhere else (and thus small logits for other c s).
We will justify this in detail in Section 4.
R EVERSE 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 (as we confirm in Appendix C.2).
Our first line of evidence involves examining the network weights and activations and observing consistent periodic structure that was 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 is 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 approximately produces multiples of the functions cos ( w k ( a + b )) and sin ( w k ( a + b )) , 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 most neurons are well-approximated by sines and cosines at a single frequency.
Moreover, the corresponding direction in W L also contains only that frequency.
This suggests that the models’ 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 algorithm we found above and find that doing so does not harm and sometimes improves model performance.
4. S UGGESTIVE 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 then compute the ℓ 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 and P − k are symmetric.
The embedding matrix W E is sparse in the Fourier basis–it only has significant nonnegligible norm at 6 frequencies.
Of these frequencies, 5 frequencies are the only frequencies used by other parts of the model (corresponding to k ∈ { 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 combination of two inputs for head 0 in Figure 4.
The attention exhibits a periodic structure with frequency
[Figure]
Fourier Components of Embedding Matrix
[Figure]
Fourier Components of Neuron
[Figure]
Logit
[Figure] [Figure]
k = 35 .
Other heads (Appendix C.1) also exhibit periodicity, always with one of five key frequencies.
In Figure 4, we plot the activations of neuron 0, given inputs a and b .
The activations are periodic with frequency k = 42 .
We see similar patterns for other MLP neurons (Appendix C.1).
Periodicity in logits.
Finally, the logits also are periodic.
In Figure 4, we represent the logits in the 2D Fourier basis over the inputs then take the ℓ 2 -norm of each component.
There are only twenty components with significant norm, corresponding to the products of sines and cosines for the five key frequencies w k .
These show up as five 2 × 2 blocks in Figure 4.
4. M ECHANISTIC E VIDENCE : C OMPOSING M ODEL W EIGHTS
We now demonstrate that the model computes its logits using the trigonometric identity for cos ( w k ( a + b − c )) .
First, we’ll show that the logits are well approximated (95% of variance explained) by terms of the form cos ( w k ( a + b − c )) .
Second, we’ll show that W L is rank 10 and can be written as:
[Equation]
for some u k , v k ∈ R .
Third, note that our model implements the logits for a, b as:
[Equation]
We’ll show that the projections of MLP ( a, b ) in the cos ( w k c ) and sin ( w k c ) directions are multiples of cos ( w k ( a + b )) and cos ( w k ( a + b )) respectively.
Taken together, these three steps demonstrate that the network computes its logits using the corresponding trigonometric identity.
Logits are well approximated by a weighted sum of cos ( w k ( a + b − c ))
s. First we approximate the output logits as a sum of the form (cid:80) k α k cos( w k ( a + b − c )) for k ∈ { 14 , 35 , 41 , 42 , 52 } and fit the coefficients α k via a linear regression.
This approximation explains 95% of the variance in the original logits.
This is surprising—naively, the output logits are a 113 · 113 · 113 dimensional vector,
[Table]
but it’s 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 .
The W L is well approximated by cos ( w k c ) and sin ( w k c ) .
We perform a discrete Fourier transform (DFT) on the logit axis of W L and look at the 10 directions corresponding to sin ( w k c ) and cos ( w k c ) .
When we approximate W L with these 10 directions, the residual has Frobenius norm that is under 0 .
55% of the norm of W L .
The unembedding matrix “reads off” terms of the form cos ( w k ( a + b )) and sin ( w k ( a + b )) from the MLP neurons.
Finally, we take the dot product of the MLP activations with each of the 10 directions in W L .
Table 1 displays the results in the Fourier basis: projecting the MLP activations onto the cos ( w k c ) direction only has non-trivial terms of the form cos ( w k a ) cos ( w k b ) and sin ( w k a ) sin ( w k b ) , and projecting onto the sin ( w k c ) direction gives us terms of the form sin ( w k a ) cos ( w k b ) and sin ( w k a )
cos ( w k b ) , with approximately equal coefficients.
4. Z OOMING I N : A PPROXIMATING N EURONS WITH S INES AND C OSINES
In the previous section, we showed how the model computes its final logits by examining the weights of W L .
Now, we examine the role that individual neurons play in the model, by approximating their behavior on all 113 · 113 pairs of possible inputs in the input distribution.
Attention heads approximately compute products of sinusoidal functions.
In order to compute terms like cos ( w k ( a + b )) , the model needs to multiply the results of the embedding matrix together.
One natural place to perform this computation is the attention heads.
By performing a DFT on the scores of the 4 attention heads from the third sequence position (‘ = ’ token) to the first sequence position (‘ a ’ token), we find the scores are well approximated by 0 .
5+ C (cos ( w k a ) − cos ( w k b ))
+
D (sin ( w k a ) − sin ( w k b )) for a single frequency w k , where C, D are constants (Appendix C.1.1).
As the token embeddings are also sinusoidal functions of the input, the output of the attention heads is approximately a quadratic form of sinusoidal functions, as is necessary for the identified algorithm.
Most MLP neurons are well approximated by a quadratic form of sinusoidal functions of a single frequency.
If our interpretation is correct, then we expect that the MLP neurons will compute quadratic forms on sin ( w k a ) , sin ( w k b ) , cos ( w k a ) , cos ( w k b ) of a single frequency w k .
To check this, we attempt to approximate every MLP neuron with such a quadratic form for each of thefrequencies identified in Section 4.1.
Out of a total of 512 neurons, 433 ( 84 .
6% ) have over 85% variance explained with some frequency.
Figure 5 displays the histogram.
Maps to the logits are localized by frequency.
We next partition these 433 neurons by the frequencies with the highest variance explained.
In each of these clusters, the weights of the map W L from the neurons to the logits are only non-trivial for two components in the Fourier basis—corresponding to sine and cosine functions with the corresponding frequency.
We plot the W L weights of theneurons in the k = 14 cluster in the Fourier basis in Figure 5.
As the figure shows, the nonnegligible components are sin (cid:0) 2 kπP (cid:1) and cos (cid:0) 2 kπP (cid:1) for k = 14 , consistent with the hypothesis that these neurons only use the k = 14 sine and cosine waves.
[Figure] [Figure]
4. C ORRECTNESS 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 quadratic form.
We replace each of these neurons’ activation value by the corresponding quadratic form, 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 has no effect on accuracy.
We can instead apply a subtly different ablation to the MLP layer and replace its activations with cos( w k ( a + b )) and sin( w k ( 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 a 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 pairs of frequencies in this basis to 0.
We begin by ablating the components corresponding to each of the key frequencies.
As reported in
Table 2, ablating any key frequency causes a significant increase in loss.
This confirms that the five frequencies identified in previous sections are indeed necessary components of the transformer.
In contrast, ablating other frequencies does not hurt the model at all.
We ablated three randomly selected other frequencies and observed no impact on performance (Table 2).
Finally, we ablated all 113 · 113 − 40 of the Fourier components besides key frequencies; this improved performance.
Insofar as other frequencies are used, they appear to serve as noise that degrades performance.
Directions in W L .
In Section 4.2, we found that W L is well approximated by the 10 directions corresponding to the cosine and sine of key frequencies.
If we project the MLP activations to these 10 directions, loss decreases 50% to 1 .
19 · 10 − 7 .
If we instead projected the MLP activations onto the nullspace of these 10 directions, loss increases to 5 .
27 —worse than uniform.
This tells us that the logits depend on these and only these 10 directions.
U NDERSTANDING GROKKING BEHAVIOR USING PROGRESS MEASURES
We now use our mechanistic understanding of the network to define two progress measures : metrics that can be computed during training that track the progress of the model over the course of training.
This allows us to better understand how the network reaches its final solution.
5. P ROGRESS MEASURES
We translate the ablations in Section 4.4 into progress measures: restricted loss and excluded loss.
Restricted loss.
Since the final network uses a sparse set of frequencies w k , it makes sense to check how well intermediate versions of the model can do using only those frequencies.
To measure this, we perform a 2D DFT on the logits to write them as a linear combination of waves in a and b ,
[Table]
and set all terms besides the constant term and the 10 terms corresponding to cos( w k ( a + b )) and sin( w k ( a + b )) for the five key frequencies to 0 .
We then measure the loss of the ablated network.
Excluded loss.
The excluded loss is opposite of the restricted loss.
Instead of keeping the important frequencies w k , we remove only those key frequencies from the logits but keep the rest.
Intuitively, this metric tracks how much the model’s performance on the training data comes from the Fourier multiplication algorithm versus memorization.
This exploits the fact that the Fourier multiplication algorithm produces sparse logits in predictable dimensions, while memorization does not (and thus will be mostly unaffected by ablating a few directions).
So the ablation will only damage the component of training performance coming from the Fourier multiplication algorithm
Beyond these measures, we will also measure 1) the Gini coefficient (Hurley & Rickard, 2009) of the norms of the Fourier components of W E and W L , which measures the sparsity of W E and W L in the Fourier basis, and 2) the ℓ 2 -norm of the weights during training, since weight decay should push these down once the train loss is near zero.
5. P HASES OF GROKKING : MEMORIZATION , CIRCUIT FORMATION , AND CLEANUP
We use these progress measures to study how the model evolves throughout training.
We use the same mainline model from Section 4 and plot the excluded loss, restricted loss, ratio of test loss to restricted loss, and sum of squared weights in Figure 6.
We find that training splits into three partially overlapping phases, which we call the memorization, circuit formation, and cleanup phases.
(We 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 is relatively flat.
In other words, the model memorizes the data, and the frequencies w k used by the final model are unused.
Circuit formation (Epochs 1.4k–9.4k).
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 on the train set transitions smoothly from the memorizing solution to the Fourier multiplication algorithm.
The fall in the sum of squared weights suggests that circuit formation likely happens due to weight decay.
Notably, the circuit is formed well before grokking occurs.
Cleanup (Epochs 9.4k–14k).
In this phase, the restricted loss continues to drop, test loss suddenly drops, sum of squared weights sharply drops, and the ratio of test to restricted loss plateaus then sharply decreases.
As the Fourier multiplication circuit both solves the task well and has lower weight than the memorization circuits, weight decay encourages the network to shed the memorized solution in favor of focusing on the Fourier multiplication circuit.
This is most cleanly shown in the sharp increase in the Gini coefficient of the Fourier norms of both the embed and unembed matices
W E and W L , suggesting that the network is becoming sparser in the Fourier basis.
5. G ROKKING AND W EIGHT D ECAY
In the previous section, we saw that each phase of grokking corresponded to an inflection point in the ℓ 2 -norm of the weights.
This suggests that weight decay is an important component of grokking and drives progress towards the generalizing solution.
In this section, we support the hypothesis via ablation—in the absence of weight decay, grokking fails to occur.
In Appendix C.2 and D, we provide additional evidence showing that both weight decay and the right amount of data are necessary for grokking.
We find that 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))
[Figure]
All Frequencies
[Figure]
and that our networks do not grok on the modular arithmetic task without weight decay.
We also find that when the networks are provided with enough data, there is no longer a gap between the training and test losses (instead, both decline sharply some number of epochs into training).
Finally, in Section D.3 we replicate this result on several additional simple algorithmic tasks.
C ONCLUSION AND DISCUSSION
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 R 2 and compose the rotations using trigonometric identities to compute 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.
In Appendix F, we provide three different scenarios where this approach can be used to understand emergent phenomena in general.
Below, we outline areas of future work:
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 many more circuits implemented on orders of magnitude more neurons to solve significantly harder tasks.
Automated mechanistic analysis and task-independent progress measures .
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 for this approach to scale to other, larger models.
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 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.
R EPRODUCIBILITY S TATEMENT
An annotated Colab notebook containing the code to replicate our results, including download instructions for model checkpoints, is available at https://bit.ly/ grokking-progress-measures-website .
A M ATHEMATICAL S TRUCTURE OF THE T RANSFORMER
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 vocab = 113 is the size of the input and output spaces (treating ‘ = ’ separately), d model = 128 is the width of the residual stream (i.e. embedding size), d head = 32 is the size of query, key and value vectors for a single attention head, and d mlp =is the number of neurons.
We denote the parameters as follows: W E (embedding layer); W pos (positional embedding); W jQ (queries), W jK (keys), W jV (values), W jO (attention output) (the 4 weight matrices of head j in the attention layer); W in and b in for the input linear map of the MLP layer; W out and b out for the output linear map of the MLP layer; and 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 to denote the token in position i (as a one-hot encoded vector), p i to denote the i th positional embedding, x (0) i to denote the initial residual stream on token with index i , A ( i ) to denote the attention scores from = to all previous tokens from head i , x (1) to denote the residual stream after the attention layer on the final token, N to denote the neuron activations in the MLP layer on the final token, x (2) the final residual stream on the final token, L the logits on the final token.
The logits are calculated via the following equations:
[Equation]
A. E MPIRICAL M ODEL S IMPLIFICATIONS
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.
• 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 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 we can still ignore it and reverse engineer the model’s computation.
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 (0)2 is constant, as it is independent of x and y , and the embedding and positional embedding of ‘ = ’ are fixed.
So A j 0 = σ ( x (0)2 T W j T Q W K ( x (0)0 − x (0)1 ) (and A j 1 = 1 − A j 0 )
A consequence of the second observation is that L ≈ W
U W out N , which we denote as W L =
W U W out .
From the perspective of the network, W L is the meaningful matrix, not either of its constituents, since they compose linearly.
B W HY 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 is significantly larger than any other x .
For example, consider the function f 14 ( x )
= cos (cid:0) 2 π · 14 113 x (cid:1) .
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 1: f 14 (8) = f 14 (105) = 0 .
998 , f 14 (16) =
f 14 (105) = 0 . 994 , etc.
Now consider f 35 ( x ) =
cos (cid:0) 2 π · 35 113 x (cid:1) .
While this function also has period 113 and is maximized at x = 0 mod 113 , it turns out that f 35 (8) = f 35 (105) = − 0 .
This means that by adding together f 14 and f 35 , we end up with a function that is not close to 1 at x = 8 mod 113 .
Similarly, while f 35 (16) = 0 .
961 , f 52 (16) = − 0 .
56 , and so adding a third frequency reduces the peak at 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 7.
C S UPPORTING EVIDENCE FOR MECHANISTIC ANALYSIS OF MODULAR ARITHMETIC NETWORKS C.
F URTHER ANALYSIS OF THE SPECIFIC TRAINING RUN DISCUSSED IN THE
PAPER
In this section, we provide additional evidence relating to the mainline model.
C.1. P ERIODICITY IN THE ACTIVATIONS OF OTHER ATTENTION HEADS
In Figure 8 we plot the attention patterns from the final token ‘ = ’ to the first token a for allattention heads, as a heatmap over the inputs a and 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 andare different.
As shown in Appendix A.1, the attention paid from ‘ = ’ to itself is negligible,
so A j 0 = 1 − A j
and it suffices to plot attention to a .
Attention patterns by Head ('=' to a)
a a a a
[Figure] [Figure]
[Table]
C.1. T HE 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 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 )) , for some frequency w k and constants α j and β j (which may differ for each head).
Note further that this simplifies to 0 .
5 + γ (cos( w k ( a + θ )) − cos( w k ( b + θ ))) for some constants γ and θ .
We show the coefficients and fraction of variance explained in Table
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 and head j is just a lookup table on the input token a (of size P ).
To see why, note that A j 0 = mx (0)0 T W jKT W jQ x (0)2 .
x (0)2 is constant since the token is always ‘ = ’ and x (0) =
W E t 0 + p 0 .
So this reduces to t 0 · C j + D for some constant vector
C j = W TE W j KT W j Q x (0) ∈
R p and some scalar D = p T 0 W j KT W j Q x (0)2 .
As t 0 is one-hot encoded, this is just a lookup table, which we may instead denote as C j [ a ]
Next, note that the attention pattern from = → 0 is σ ( C j [ a ] − C j [ b ]) .
As argued in Appendix A.1, the attention paid =
→ = is negligible and can be ignored.
So the softmax reduces to a softmax over two elements, which is a sigmoid on their difference.
As form of C j does not mention the token index or value, it is the same for position 0 and 1 .
We now show that C j is well-approximated by a wave of frequency w k j for some integer k j .
That is,
C j [ a ] ≈ F j cos( w k j a )+ G j sin( w k j a ) .
We do this by simply computing C j and fitting the constants
F j and G j to minimize ℓ 2 loss, and display the resulting coefficients for each head in Figure 9.
This fit explain 99.02%, 95.21%, 99.10%, 92.42% of the variance of C j respectively.
For each head j , σ ( C j [ a ] − C j [ b ]) ≈ 0 .
5+ E j
( C j [ a ] − C j [ b ]) for some constant E j —that is, the sigmoid has some linear approximationm.
It is always the case that we can take a linear approximation to scalar function, by a 1D Linear Regression.
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 a fairly wide range ( [ − 5 , 5] roughly), the linear approximation to the sigmoid is a fairly good fit, explaining 97.5% of the variance.
[Figure] [Figure]
ff i
c i e n t
o f F ou r i e r C
o m pon e n t
[Figure]
− −  sin cos
Fourier components
[Figure]
C
[Figure]
Frequency
[Figure]
C o e ff i
c i e n t
o f F ou r i e r C
o m pon e n t
[Figure]
−
[Figure]
e ff i
c i e n t
o f F ou r i e r C
o m pon e n t
[Figure] [Figure]
Fourier components
[Figure]
Frequency
[Figure]
C o
[Figure]
ff i
c i e n t
o f F ou r i e r C
o m pon e n t
[Figure]
As with Neuron 0, all of the activation patterns are periodic in inputs.
We validate that this is all that is going on, by replacing the sigmoid with the best linear fit.
This improves loss from 2 .
41 · 10 − 7 to 2 .
12 · 10 − 7 .
Bringing it all back together, we see that the attention patterns of each head will be well approximated by 0 .
5 ± C j (cos( w k j ( a + θ j )) − cos( w k j ( b + θ j )))
- the softmax is linear, with an intercept of 0 . 5 , and the weights C j map each token to a score that is a wave in a single frequency.
This exactly gives us the periodic form shown in Figure 8.
C.1. P ERIODICITY IN THE ACTIVATIONS OF ADDITIONAL NEURONS
In Figure 10, we display the activations of four more MLP neurons, as a function of the inputs.
As with neuron 0, the activations of these neurons are also periodic in the inputs.
[Figure] [Figure]
[Figure] [Figure]
C.1. A DDITIONAL GROKKING FIGURES FOR MAINLINE RUN
In Figure 11, 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 12, we show the coefficients of the five key frequencies in the logits, calculated by regressing the logits against the five cos ( w k ( a + b − c )) terms.
In Figure 13, we plot the excluded loss if we exclude each of the five key frequencies (as opposed to all five key frequencies).
All three of these figures have inflection points corresponding to the relevant phases of grokking, discussed in Section 5.1.
C. A DDITIONAL RESULTS FROM DIFFERENT RUNS
In this section, we plot relevant figures from other runs, either with the same architecture (Appendix
C.2.1) or with different architectures or experimental setups (Appendix C.2.2).
Note that in general, while all models learn to 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
[Figure] [Figure]
ou r i
[Figure]
r
C
[Figure]
pon
[Figure]
ou r i
[Figure]
r
[Figure]
pon
[Figure]
W E for each of four other random seeds for the original (1 layer) architecture.
As discussed in Section 4.1 and Appendix
C.2.1, the sparsity of W E in the Fourier basis is evidence that the network is operating in a Fourier basis.
restricted loss, we perform a DFT on the neuron-logit map W L , then take the frequencies with nontrivial coefficients.
C.2. A DDITIONAL RESULTS FOR DIFFERENT RUNS WITH THE SAME ARCHITECTURE
In this section, we provide evidence that all 4 other runs 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.
Confirming that the other seeds use the Fourier Multiplication Algorithm.
In Figure 14, we show the norms of the Fourier components of the embedding matrix W E for each of the 4 other random seeds.
As with the mainline model, the matrices are sparse in the Fourier basis.
In Figure 15, we show the norms of the Fourier components of the neuron-logit map W L for the 4 other random
[Table]
[Table] [Table]
[Table] [Table]
b )) , and find that this is consistently close to 1.
[Figure]
i
[Figure]
r
[Figure]
pon
[Figure]
[Table]
the networks reduces performance to worse than chance, while ablating all other frequencies improves performance.
seeds.
The matrices are sparse in the Fourier basis, enabling us to identify 3 or 4 key frequencies for each of the seeds.
Again, note that these are different frequencies per seed.
Using the key frequencies identified in the neuron-logit map, we repeat the experiment in Section 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 ( w k ( a + b )) and sin ( w k ( a + b )) being computed at the MLP layer.
We confirm that the trigonometric identities are a good approximation by approximating the activations with a single term of the form cos ( w k ( a + b )) or sin ( w k ( a + b )) —as with the mainline model, the fraction of variance explained is consistently close to 100%.
Finally, we ablate the key frequencies from the logits as in Section 4.4 and report the results in Table 5.
As with the mainline model, ablating all of the key frequencies reduces performance to 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 16, we display the train, test, and restricted loss for each of the four other random seeds.
In Figure 17, we display the Gini coefficients of the Fourier components of the embedding matrix W E and the neuron-logit map W L for each of the four other random seeds.
The shape of the curves are very similar to those of the mainline model, allowing us to divide grokking on these models into the same three phases 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.
[Figure] [Figure]
G i n i c
o e ff i
c i e n t
[Figure]
i
[Figure]
i
[Figure]
i
[Figure]
C.2. R ESULTS FOR OTHER EXPERIMENTAL SETUPS
In this section, we provide further evidence that small transformers grok on the modular addition task, by varying the size of the network, the amount of training data, and the size of the prime P .
1-Layer Transformers with Varying Fractions of Training Data.
We find that grokking occurs for the modular addition task with P = 113 for many data fractions (that is, the fraction of the 113 · 113 pairs of inputs that the model sees during training), as shown in Figure 18.
Smaller amount lead to slower grokking, but sufficiently large fractions of data ( ≥ 85% ) lead to immediate generalization, as shown in Figure 18.
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 20, 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 21 and find it shows a similar pattern to the mainline 1-layer transformer, in that it improves relatively smoothly before grokking occurs.
Smaller and larger primes.
We also examined smaller and larger prime moduli.
For P =(Figure 22), we explored a variety of weight decays to observe grokking in the small prime case.
With the original weight decay setting of λ = 1 , we found that the models never generalized.
However, increasing the weight decay to λ = 5 does allow the model to grok.
We speculate that this is because the memorization solution is significantly smaller (since there are only 53 · 53 total pairs), thereby requiring more aggressive weight decay for the generalizing solution to be favored.
For P = 401 (Figure 23), we could not get grokking, even by varying the weight decay parameter.
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 to 30% of 401 · 401 pairs), thereby favoring the generalizing solution from the start.
[Figure]
[Figure] [Figure]
5k
[Figure]
Number of steps to achieve <1e-2 loss Fraction of
[Figure]
data
[Figure]
r o f s t
e p s
D A DDITIONAL RESULTS ON GROKKING D. B OTH WEIGHT DECAY AND LIMITED DATA ARE NECESSARY FOR GROKKING
As discussed in Section 6 and Appendix C.2, the weight decay and the amount of data seem to have a strong effect on whether grokking occurs.
To confirm this, we experiment with removing weight decay and varying the amount of data on 1-layer transformers.
In Figure 24, we give the training, test, and full excluded loss for a typical training run with λ =
(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 18, we show the test loss curves for models trained 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
. D. T HE 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 to how gradients of different scale interact with adaptive
[Figure]
Loss for 5 Digit Addition w/ Infinite Data
[Figure]
Per Token Train/Test Loss for 5 Digit Addition w/ Infinite Data
[Figure]
optimizers.
We were even able to induce slingshots on a 1-layer by reducing the precision of the loss calculations (as this causes many gradients to round to 0 and thus greatly increases the differences in scale of gradients).
However, as 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.
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. A DDITIONAL 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, randomly choose a subsequence to repeat, and train the model to predict the repeated tokens.
Skip trigram.
We feed in a 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.
[Figure]
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 26.
Train loss coincides with test loss, so grokking does not occur, as the model almost never sees the same pair of 5 digit numbers twice, with 10 10 such pairs.
Interestingly, the various small bumps in Figure 26 correspond to the model learning how to calculate each of the 6 tokens in the output.
Figure ??
shows the breakdown of the loss per token.
However, grokking does occur when we restrict the model to only see 700 data points, as shown in Figure 26.
Repeated subsequence As with the 5-digit addition task, we find that restricting the amount of data is necessary and sufficient for grokking on the repeated subsequence task.
In Figure ?? , 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 28.
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 no grokking in Figure 29.
Meanwhile, the model restricted to only see 512 data points exhibits clear grokking in Figure 30.
Taken together, these results echo the importance of limited data for grokking.
E F URTHER SPECULATIONS ON GROKKING
E.  A N INTUITIVE EXPLANATION OF GROKKING
In this section, we speculate on what might be happening “under the hood” when a 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 is a sudden improvement in test loss, so test and train loss converge.
In contrast to Power et 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 16.
The core issue is that the model has two possible solutions: memorization (with low train loss and high test loss) and a generalization (with low train loss and low test loss).
In our case, the
[Figure]
[Figure]
[Figure]
[Figure]
/Test
Loss for Skip Trigram Task w/ 50 Data Points [Figure] 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 or 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 is constant.
The two must cross at 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 a certain task develop a 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.
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 a phase transition is that there’s some feature of the loss landscape that makes the generalising solution harder to reach - rather than a 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 over complex ones.
The model has enough data to marginally prefer generalization over memorization.
The phase transition indicates that generalization is ”hard to reach” while the model has no problems with memorization.
But as 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 a somewhat constant training performance in this circuit 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.
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 a 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 we observe in practice.
E. H YPOTHESIS : P HASE T RANSITIONS 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 et al., 2021;
Olsson et al., 2022) in LLMs.
This is 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 without following a smooth gradient.
So how can a model learn a discrete algorithm?
As a 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 subsequences.
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 a query vector representing the current token and a key vector representing the previous token head’s output using K-Composition, the context of the previous token.
It attends to a 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.
• Random walk: The network wanders randomly around the loss landscape until it encounters 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.
• Evolution: A 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 a 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 components developed depending on the first, like the lens of the eye.
Evolution is a natural explanation, However, based on our 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 an arbitrary location, e.g. 7 2 8 3 1 9 3 8 3 1 9 9 2 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 a previous token head is only useful in the context of an induction-like head that completes the circuit.
Likewise, an 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 easiest way to improve performance with perfect accuracy is by scaling up on K-composition with a previous token head and so 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 complicated, representing a small region in model space.
So a random walk is unlikely to stumble upon it.
Concretely, in our 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, we suspect that all explanations contribute to the entire picture.
This seems most plausible if the uniform head just so happens to attend a 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 circuit 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 F URTHER DISCUSSION ON USING MECHANISTIC INTERPRETABILITY AND PROGRESS MEASURES FOR STUDYING EMERGENT PHENOMENA
While we find approach of using mechanistic interpretability to define progress measures relatively promising, there remains significant uncertainty as to how scalable existing mechanistic interpretability 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 understand 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 or not mechanistic interpretability will successfully scale to large models to this extent (or even if there exist human-understandable explanations for all of their sophisticated behavior).
That being said, in cases where mechanistic interpretability does recover human-understandable mechanisms, we could simply use the parts of the mechanism as progress measures.
If future mechanistic interpretability can only recover parts of the mechanism of larger models and can only generate comprehensive understanding of the mechanisms of smaller models, we might still be able to use our understanding 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 interpretability work already allows us to recover fragments of large model behavior and understand 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 language transformers.
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 require human understanding.
For example, if we end up with automated mechanistic interpretability (that nonetheless still fails to recover human-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 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, we think the future work outlined in 6 is necessary to successfully apply our approach to predict and understand emergent behavior in existing large language models, and so remain cautiously optimistic.