Progress measures for grokking via mechanistic interpretability — part 4/6
Cleaned-text code points [46874, 62497). Bibliography omitted here; source.md includes it.
Next, note that the attention pattern from = → 0 =\to 0 is σ ( C j [ a ] − C j [ b ] ) \sigma(C_{j}[a]-C_{j}[b]) . As argued in Appendix A.1 , the attention paid = → = =\to= 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 C_{j} does not mention the token index or value, it is the same for position 0 0 and 1 1 .
We now show that C j C_{j} is well-approximated by a wave of frequency w k j w_{k_{j}} for some integer k j k_{j} . That is, C j [ a ] ≈ F j cos ( w k j a ) + G j sin ( w k j a ) C_{j}[a]\approx F_{j}\cos(w_{k_{j}}a)+G_{j}\sin(w_{k_{j}}a) . We do this by simply computing C j C_{j} and fitting the constants F j F_{j} and G j G_{j} to minimize ℓ 2 \ell_{2} loss, and display the resulting coefficients for each head in Figure 10 . This fit explain 99.02%, 95.21%, 99.10%, 92.42% of the variance of C j C_{j} respectively. Interestingly, the coefficients of heads 1 1 and 3 3 are almost exactly the opposite of each other.
For each head j j , σ ( C j [ a ] − C j [ b ] ) ≈ 0.5 + E j ( C j [ a ] − C j [ b ] ) \sigma(C_{j}[a]-C_{j}[b])\approx 0.5+E_{j}(C_{j}[a]-C_{j}[b]) for some constant E j E_{j} —that is, the sigmoid has some linear approximation. (The intercept will be 0.5 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 ] [-5,5] roughly), the linear approximation to the sigmoid is a fairly good fit, explaining 97.5% of the variance.
We validate that this is all that is going on, by replacing the sigmoid with the best linear fit. This improves performance, decreasing test loss from 2.41 ⋅ 10 − 7 2.41\cdot 10^{-7} to 2.12 ⋅ 10 − 7 2.12\cdot 10^{-7} .
By properties of sinusoidal functions, 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 ) ) ) 0.5\pm C_{j}(\cos(w_{k_{j}}(a+\theta_{j}))-\cos(w_{k_{j}}(b+\theta_{j}))) - the softmax is linear, with an intercept of 0.5 0.5 , and the weights C j 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 9 .
Finally, for each head j j , we plot the output of the OV circuit W O j W V j x ( 0 ) W_{O}^{j}W_{V}^{j}x^{(0)} in the Fourier basis and display the results in Figure 11 ). The largest component of each head corresponding to the frequency of the attention pattern C j C_{j} , with heads 0 and 2 being almost entirely composed of a sines and cosines of a single frequency. On the other hand, the norms for the components of heads 1 1 and 3 3 are almost exactly the same, and contain all five key frequencies. As the coefficients of the attention pattern weights have the opposite non-constant components (Table 2 , Figure 10 ), their attention scores sum almost exactly to 1 across all inputs. This implies that heads 1 1 and 3 3 are used to output the first order terms sin ( w k ) , cos ( w k ) \sin\left(w_{k}\right),\cos\left(w_{k}\right) in the five key frequencies. We speculate that this is because of weight decay encouraging the embeddings W E W_{E} to be small, causing the network to allocate two of its attention heads to effectively increasing the size of W E W_{E} .
Bringing it all together, this implies that attention heads 0 0 and 2 2 are approximately computing a degree 2 polynomial of cosines and sines of a single frequency each, while heads 1 1 and 3 3 amplify the key frequencies in the residual stream.
#### C.1.4 Periodicity in the activations of additional neurons
In Figure 12 , 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.
#### C.1.5 Additional grokking figures for mainline run
In Figure 13 , we display the accuracy of the model when restricting the model to use only the five key frequencies. As with restricted loss, this improves model performance during training.
In Figure 14 , we show the coefficients of the five key frequencies in the logits, calculated by regressing the logits against the five cos ( w k ( a + b − c ) ) \cos\left(w_{k}(a+b-c)\right) terms.
In Figure 15 , we plot the excluded loss if we exclude each of the five key frequencies (as opposed to all five key frequencies).
All three of these figures have inflection points corresponding to the relevant phases of grokking, discussed in Section 5.1 .
### C.2 Additional results from different runs
In this section, we plot relevant figures from other runs, either with the same architecture (Appendix 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 restricted loss, we perform a DFT on the neuron-logit map W L W_{L} , then take the frequencies with nontrivial coefficients. 3 3 3 One method for getting a general (model-independent) progress measure for this task is to compute the excluded loss for each of the 56 unique frequencies and then take the max. We omit the plots for this variant of the excluded loss as they are broadly similar.
#### C.2.1 Additional results for different runs with the same architecture
In this section, we provide evidence that all 4 other runs (i.e., random seeds) using the experimental setup of our mainline model also use the Fourier multiplication algorithm, and then confirm that the same phases of grokking also occur on these runs.
Confirming that the other seeds use the Fourier Multiplication Algorithm. In Figure 16 , we show the norms of the Fourier components of the embedding matrix W E 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 17 , we show the norms of the Fourier components of the neuron-logit map W L W_{L} for the 4 other random 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 the specific frequencies differ by 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 ) ) \cos\left(w_{k}(a+b)\right) and sin ( w k ( a + b ) ) \sin\left(w_{k}(a+b)\right) 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 ) ) \cos\left(w_{k}(a+b)\right) or sin ( w k ( a + b ) ) \sin\left(w_{k}(a+b)\right) —as with the mainline model, the fraction of variance explained is consistently close to 100%.
Next, we ablate the key frequencies from the logits as in Section 4.4 and report the results in Table 4 . 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 18 , we display the train, test, and restricted loss for each of the four other random seeds. In Figure 19 , we display the Gini coefficients of the Fourier components of the embedding matrix W E W_{E} and the neuron-logit map W L 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.
#### C.2.2 Results 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 P .
##### 1-Layer Transformers with Varying Fractions of Training Data.
We find that grokking occurs for the modular addition task with P = 113 P=113 for many data fractions (that is, the fraction of the 113 ⋅ 113 113\cdot 113 pairs of inputs that the model sees during training), as shown in Figure 20 . Smaller amount lead to slower grokking, but sufficiently large fractions of data ( ≥ 60 % \geq 60\% ) lead to immediate generalization, as shown in Figures 20 and 21 .
As with the results in Appendix C.2.1 , all of the 1-layer transformers in this section also converge to using the Fourier multiplication algorithm.
##### 2-Layer Transformers.
As shown in Figure 22 , 2-layer transformers also exhibit some degree of grokking. However, this is complicated by the slingshot mechanism ( Thilak et al., 2022 ) . We display the excluded loss of a 2-layer transformer in Figure 23 and find it shows a similar pattern to the mainline 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 = 53 P=53 (Figure 24 ), we explored a variety of weight decays to observe grokking in the small prime case. With the original weight decay setting of λ = 1 \lambda=1 , we found that the models never generalized. However, increasing the weight decay to λ = 5 \lambda=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 53\cdot 53 total pairs), thereby requiring more aggressive weight decay for the generalizing solution to be favored.
For P = 109 P=109 , we saw exactly the same behavior as with the mainline model.
For P = 401 P=401 (Figure 25 ), we could not get grokking, even by varying the weight decay parameter λ ∈ { 0.3 , 0.5 , 1 , 3 , 5 , 8 } \lambda\in\{0.3,0.5,1,3,5,8\} . 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 P=113 case (from 30% of 113 ⋅ 113 113\cdot 113 pairs to 30% of 401 ⋅ 401 401\cdot 401 pairs), thereby favoring the generalizing solution from the start. We then trained 3 models each using 5 % 5\% , 10 % 10\% , 20 % 20\% of the pairs of training data with λ = 1 \lambda=1 , and found that the models trained on 5 % 5\% and 10 % 10\% of the data immediately overfit and never generalized, while the models trained on 20 % 20\% of the data also generalized immediately.
#### C.2.3 Generalizing models consistently use the Fourier Multiplication Algorithm
For each of the models in Appendix C.2.2 that achieve low test loss, we repeated the analysis performed in the mainline model, and summarize the results in Table 5 . We list their key frequencies, Gini coefficients, and relevant FVEs. We find that every model trained with weight decay and that generalizes correctly implements some variation of the Fourier multiplication algorithm.
Interestingly, the embedding and unembedding matrices of the models trained with dropout are not sparse in the Fourier basis, and the logits for the p = 0.2 p=0.2 models are not as well explained by a sum of cosines as the other models (likely because the p = 0.2 p=0.2 models are simply worse at the task). We speculate that this is likely due to a combination of insufficient training epochs (as dropout models seem to take much longer to grok) and the inherent need for redundancy for networks trained via dropout.
As with the mainline model, we ignore the final skip connection (around the final MLP), as all of the generalizing models studied do not suffer significant performance penalties if the skip connection is zero or mean ablated (Table 6 ).
## Appendix D Additional results on grokking
### D.1 Both regularization and limited data are necessary for grokking
As discussed in Section 7 and Appendix C.2 , the weight decay and the amount of data seem to have 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 26 , we give the training, test, and full excluded loss for a typical training run with λ = 0 \lambda=0 (no weight decay). As the figure shows, no grokking occurs, and excluded loss does not increase, suggesting that the model does not form the circuit for generalizing algorithm at all.
In Figure 20 , we show the test loss curves for models trained with weight decay λ = 1 \lambda=1 and on various fractions of the data. Though all the train losses are approximately the same—that is, they memorize at the same rate, models trained on smaller fractions of data take longer to grok.
In Figure 27 , we display the test and train loss of models trained with λ = 0.3 \lambda=0.3 and λ = 3.0 \lambda=3.0 . Smaller amounts of weight decay lead to slower grokking, while larger amounts of weight decay lead to faster grokking—on average, it takes around 3k epochs for models to grok with weight decay λ = 0.3 \lambda=0.3 , 5-10k epochs for the models to grok with weight decay λ = 1.0 \lambda=1.0 , and 20k epochs for the models to grok with weight decay λ = 3.0 \lambda=3.0 .
Finally, we test whether other forms of regularization can also induce grokking. We replaced weight decay with the following types of regularization while keeping all other hyperpameters the same: 1. Dropout We add dropout Srivastava et al. (2014) to the MLP neurons, with p ∈ { 0.2 , 0.5 , 0.8 } p\in\{0.2,0.5,0.8\} . That is, for each individual neuron, we set it to 0 0 with probability p p during training, and also multiply the outputs of the other neurons by 1 1 − p \frac{1}{1-p} .
2. ℓ 1 \ell_{1} Regularization We add an ℓ 1 \ell_{1} penalty to the loss term. We use λ ∈ { 1 , 10 , 100 } \lambda\in\{1,10,100\} . Note that we do not decouple the updates with respect to the ℓ 1 \ell_{1} penalty from optimization steps done with respect to the log loss (as is done for ℓ 2 \ell_{2} regularization via AdamW Loshchilov & Hutter (2017) ).
In each case, we ran three random seeds. We show the results in Figure 28 . While grokking did not occur with ℓ 1 \ell_{1} regularization, we found that it does occur for all three seeds using dropout with p = 0.2 p=0.2 or p = 0.5 p=0.5 . We speculate that this is because both dropout and weight decay encourage the network to spread out computation (which is required for the Fourier multiplication algorithm), while ℓ 1 \ell_{1} regularization encourages the network to become more sparse in the neuron basis and thus less sparse in the Fourier basis, preventing the network from learning the Fourier Multiplication Algorithm.
### D.2 The slingshot mechanism often occurs, but is unnecessary for grokking
As noted in Section C.2 , our 2-layer transformers exhibit significant slingshots ( Thilak et al., 2022 ) during training. We speculate that this is due to how gradients of different scale interact with adaptive 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).