Progress measures for grokking via mechanistic interpretability — part 2/6
Cleaned-text code points [15546, 31275). Bibliography omitted here; source.md includes it.
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 k = 35 k=35 . In Figure 4 , we also plot the activations of MLP neuron 0 for every combination of inputs. The activations are periodic with frequency k = 42 k=42 . We see similar patterns for other attention heads and MLP neurons (Appendix C.1 ).
Periodicity in logits. Finally, the logits are also periodic. In Figure 4 , we represent the logits in the 2D Fourier basis over the inputs, then take the ℓ 2 \ell_{2} -norm over the output dimension. There are only twenty components with significant norm, corresponding to the products of sines and cosines for the five key frequencies w k w_{k} . These show up as five 2 × 2 2\times 2 blocks in Figure 4 .
### 4.2 Mechanistic Evidence: Composing Model Weights
We now demonstrate that the model implements the trigonometric identity ( 1 ) as follows: the functions cos ( w k ( a + b ) ) \cos\left(w_{k}(a+b)\right) , sin ( w k ( a + b ) ) \sin\left(w_{k}(a+b)\right) are linearly represented in the MLP activations, and the unembed matrix reads these linear directions and multiplies them by cos ( w k c ) \cos\left(w_{k}c\right) , sin ( w k c ) \sin\left(w_{k}c\right) respectively.
We will do this in two steps. First, we show that W L W_{L} (the matrix mapping MLP activations to logits) is (approximately) rank 10 and can be well approximated as: W L = ∑ k ∈ { 14 , 35 , 41 , 42 , 52 } cos ( w k ) u k T + sin ( w k ) v k T \displaystyle W_{L}={\sum}_{k\in\{14,35,41,42,52\}}\cos\left(w_{k}\right)u_{k}^{T}+\sin\left(w_{k}\right)v_{k}^{T} (2) for some u k , v k ∈ ℝ 512 u_{k},v_{k}\in\mathbb{R}^{512} , where cos ( w k ) , sin ( w k ) ∈ ℝ 113 \cos\left(w_{k}\right),\sin\left(w_{k}\right)\in\mathbb{R}^{113} are vectors whose c c th entry is cos ( w k c ) \cos\left(w_{k}c\right) and sin ( w k c ) \sin\left(w_{k}c\right) . Second, note that our model implements the logits for a , b a,b as: Logits ( a , b ) = W L MLP ( a , b ) ≈ ∑ k cos ( w k ) u k T MLP ( a , b ) + sin ( w k ) v k T MLP ( a , b ) \displaystyle\textrm{Logits}(a,b)=W_{L}\textrm{MLP}(a,b)\approx{\sum}_{k}\cos\left(w_{k}\right)u_{k}^{T}\textrm{MLP}(a,b)+\sin\left(w_{k}\right)v_{k}^{T}\textrm{MLP}(a,b) (3) We check empirically that the terms u k T MLP ( a , b ) u_{k}^{T}\textrm{MLP}(a,b) and v k T MLP ( a , b ) v_{k}^{T}\textrm{MLP}(a,b) are approximate multiples of 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) ( > 90 % >\!90\% of variance explained). Thus the network computes trigonometric functions in the MLP and reads them off as claimed. As a sanity check, we confirm that the logits are indeed well-approximated by terms of the form cos ( w k ( a + b − c ) ) \cos\left(w_{k}(a+b-c)\right) (95% of variance explained).
W L W_{L} is well approximated by cos ( w k c ) \cos\left(w_{k}c\right) and sin ( w k c ) \sin\left(w_{k}c\right) . We perform a discrete Fourier transform (DFT) on the logit axis of W L W_{L} and look at the 10 directions u k , v k u_{k},v_{k} corresponding to sin ( w k ) \sin\left(w_{k}\right) and cos ( w k ) \cos\left(w_{k}\right) . When we approximate W L W_{L} with ∑ k ∈ { 14 , 35 , 41 , 42 , 52 } cos ( w k ) u k T + sin ( w k ) v k T {\sum}_{k\in\{14,35,41,42,52\}}\cos\left(w_{k}\right)u_{k}^{T}+\sin\left(w_{k}\right)v_{k}^{T} , the residual has Frobenius norm that is under 0.55 % 0.55\% of the norm of W L W_{L} . This shows that W L W_{L} is well approximated by the 10 directions corresponding to cos ( w k ) \cos\left(w_{k}\right) and sin ( w k ) \sin\left(w_{k}\right) for each of the five key frequencies. We also plot the norms of each direction in Figure 3 , and find that no Fourier component outside the 5 5 key frequencies has significant norm.
The unembedding matrix “reads off” terms of the form 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) from the MLP neurons. Next, we take the dot product of the MLP activations with each of the directions u k , v k u_{k},v_{k} for k ∈ { 14 , 35 , 41 , 42 , 52 } k\in\{14,35,41,42,52\} . Table 1 displays the results: the dot products u k T MLP ( a , b ) u_{k}^{T}\textrm{MLP}(a,b) and v k T MLP ( a , b ) v_{k}^{T}\textrm{MLP}(a,b) are well approximated by a multiple of terms of the form cos ( w k ( a + b ) ) \displaystyle\cos\left(w_{k}(a+b)\right) = cos ( w k a ) cos ( w k b ) − sin ( w k a ) sin ( w k b ) , and \displaystyle=\cos\left(w_{k}a\right)\cos\left(w_{k}b\right)-\sin\left(w_{k}a\right)\sin\left(w_{k}b\right)\textrm{, and } 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\left(w_{k}a\right)\cos\left(w_{k}b\right)+\cos\left(w_{k}a\right)\sin\left(w_{k}b\right). That is, for each key frequency k k , u k u_{k} and v k v_{k} are linear directions in the space of MLP neuron activations that represent 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) .
Logits are well approximated by a weighted sum of cos ( w k ( a + b − c ) ) \cos\left(w_{k}(a+b-c)\right) s. We approximate the output logits as the sum ∑ k α k cos ( w k ( a + b − c ) ) \sum_{k}\alpha_{k}\cos(w_{k}(a+b-c)) for k ∈ { 14 , 35 , 41 , 42 , 52 } k\in\{14,35,41,42,52\} and fit the coefficients α k \alpha_{k} via ordinary least squares. This approximation explains 95 % 95\% of the variance in the original logits. This is surprising—the output logits are a 113 ⋅ 113 ⋅ 113 113\cdot 113\cdot 113 dimensional vector, but are well-approximated with just the 5 directions predicted by our interpretation. If we evaluate test loss using this logit approximation, we actually see an improvement in loss, from 2.4 ⋅ 10 − 7 2.4\cdot 10^{-7} to 4.7 ⋅ 10 − 8 4.7\cdot 10^{-8} .
Taken together, these results confirm that the model computes sums of terms of the form cos ( w k ( a + b − c ) ) = cos ( w k ( a + b ) ) cos ( w k c ) + sin ( w k ( a + b ) ) sin ( w k c ) \cos\left(w_{k}(a+b-c)\right)=\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) .
### 4.3 Zooming In: Approximating Neurons with Sines and Cosines
In the previous section, we showed how the model computes its final logits by using W L W_{L} to “read off” trigonometric identities represented in the MLP neurons. We now examine the attention heads and MLP neurons to understand how the identities come to be represented at the MLP layer. In Appendix C.1.2 , we show that two of the attention heads approximately compute degree- 2 2 polynomials of sines and cosines of a particular frequency (and the other two are used to increase the magnitude of the input embeddings in the residual stream). Here, we show that most neurons are also well-approximated by degree- 2 2 polynomials, and the map from neurons to logits is localized by frequency.
Most MLP neurons approximately compute a degree- 2 2 polynomial of a single frequency. We next try to approximate the activations of each MLP neuron by a degree- 2 2 polynomial of one of the 5 key frequencies. As shown in Figure 5 , out of 512 total neurons, 433 ( 84.6 % 84.6\% ) have over 85 % 85\% of their variance explained with a single frequency.
Maps to the logits are localized by frequency. We partition these 433 neurons by the frequencies with the highest variance explained. For each resulting subset, the map W L W_{L} from neurons to logits has only two non-trivial components, corresponding to sine and cosine at that frequency. For example, in Figure 5 we plot the 44 columns of W L W_{L} corresponding to the 44 neurons in the k = 14 k=14 cluster and find that the only non-negligible components are sin ( 2 k π P ) \sin\left(\frac{2k\pi}{P}\right) and cos ( 2 k π P ) \cos\left(\frac{2k\pi}{P}\right) for k = 14 k=14 .
### 4.4 Correctness checks: ablations
In previous sections, we showed that various components of the model were well-approximated by sparse combinations of sines and cosines. We verify that these approximations are faithful to the model’s functionality, by replacing each component with its approximation. This generally does not hurt the performance of the model and in some cases improves it.
MLP neurons. In Section 4.3 , we identified 433 neurons that were well-approximated by a degree- 2 2 polynomial. We replace each of these neurons’ activation value by the corresponding polynomial, leaving the other neurons untouched. This increases loss by only 3 % 3\% in relative terms (from 2.41 ⋅ 10 − 7 2.41\cdot 10^{-7} to 2.48 ⋅ 10 − 7 2.48\cdot 10^{-7} ) and has no effect on accuracy.
We can instead apply a stricter ablation to the MLP layer and restrict each neuron’s activation to just the components of the polynomial corresponding to terms of the form cos ( w k ( a + b ) ) \cos(w_{k}(a+b)) and sin ( w k ( a + b ) ) \sin(w_{k}(a+b)) in the key frequencies. This improves loss by 77% (to 5.54 ⋅ 10 − 8 5.54\cdot 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 113\cdot 113\cdot 113 logit matrix over all 113 ⋅ 113 113\cdot 113 pairs of inputs to get the logits in the Fourier basis, then set various frequencies in this basis to 0.
We begin by ablating the components corresponding to each of the key frequencies. As reported in Figure 6 , ablating any key frequency causes a significant increase in loss. This confirms that the five frequencies identified in previous sections are indeed necessary components of the transformer. In contrast, ablating other frequencies does not hurt the model at all.
We then ablate all 113 ⋅ 113 − 40 113\cdot 113-40 of the Fourier components besides key frequencies; this ablation actually improves performance (loss drops 70 % 70\% to 7.24 ⋅ 10 − 8 7.24\cdot 10^{-8} ).
Directions in W L W_{L} . In Section 4.2 , we found that W L 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 1.19\cdot 10^{-7} . If we instead projected the MLP activations onto the nullspace of these 10 directions, loss increases to 5.27 5.27 —worse than uniform. This suggests that the network achieves low loss using these and only these 10 directions.
## 5 Understanding grokking behavior using progress measures
We now use our mechanistic understanding of the network to define two progress measures : metrics that can be computed during training that track the progress of the model over the course of training, including during phase transitions. This allows us to study how the network reaches its final solution.
### 5.1 Progress measures
We translate the ablations in Section 4.4 into two progress measures: restricted and excluded loss.
Restricted loss. Since the final network uses a sparse set of frequencies w k 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 a and b b , and set all terms besides the constant term and the 20 20 terms corresponding to cos ( w k ( a + b ) ) \cos(w_{k}(a+b)) and sin ( w k ( a + b ) ) \sin(w_{k}(a+b)) for the five key frequencies to 0 0 . We then measure the loss of the ablated network.
Excluded loss. Instead of keeping the important frequencies w k w_{k} , we next remove only those key frequencies from the logits but keep the rest. We measure this on the training data to track how much of the performance comes from Fourier multiplication versus memorization. The idea is that the memorizing solution should be spread out in the Fourier domain, so that ablating a few directions will leave it mostly unaffected, while the generalizing solution will be hurt significantly.
Beyond these, we will also measure (1) the Gini coefficient ( Hurley & Rickard, 2009 ) of the norms of the Fourier components of W E W_{E} and W L W_{L} , which measures the sparsity of W E W_{E} and W L W_{L} in the Fourier basis, and (2) the ℓ 2 \ell_{2} -norm of the weights during training, since weight decay should push these down once the train loss is near zero.
### 5.2 Phases of grokking: memorization, circuit formation, and cleanup
Using the mainline model from Section 4 , we plot the excluded loss, restricted loss, Gini coefficient of the matrices W U W_{U} and W L W_{L} , and sum of squared weights in Figure 7 . We find that training splits into three phases, which we call the memorization, circuit formation, and cleanup phases. (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 staying relatively flat. In other words, the model memorizes the data, and the frequencies w k 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, excluded loss plateaus, restricted loss continues to drop, test loss suddenly drops, and sum of squared weights sharply drops. As the completed Fourier multiplication circuit both solves the task well and has lower weight than the memorization circuit, weight decay encourages the network to shed the memorized solution in favor of focusing on the Fourier multiplication circuit. This is most cleanly shown in the sharp increase in the Gini coefficient for the matices W E W_{E} and W L W_{L} , which shows that the network is becoming sparser in the Fourier basis.
### 5.3 Grokking and Weight Decay
In the previous section, we saw that each phase of grokking corresponded to an inflection point in the ℓ 2 \ell_{2} -norm of the weights. This suggests that weight decay is an important component of grokking and drives progress towards the generalizing solution. In Appendix D.1 , we provide additional evidence that weight decay is necessary for grokking: smaller amounts of weight decay causes the network to take significantly longer to grok (echoing the results on toy models from Liu et al. (2022) ), and our networks do not grok on the modular arithmetic task without weight decay or some other form of regularization. In Appendix C.2 , we also find that the amount of data affects grokking: when networks are provided with enough data, there is no longer a gap between the train and test losses (instead, both decline sharply some number of epochs into training). Finally, in Appendix D.3 we replicate these results on several additional algorithmic tasks.
## 6 Conclusion and discussion