# An interactive introduction to grokking and mechanistic interpretability

Our write up largely agrees with [@Quintin Pope](https://www.alignmentforum.org/users/quintin-pope?mention=user)'s [summary](https://www.alignmentforum.org/posts/GpSzShaaf8po4rcmA/qapr-5-grokking-is-maybe-not-that-big-a-deal), with the addition of training trajectory visualizations and an explanation of the MLP construction that solves modular addition.  

A meta note that didn't make it into the article — with so many people looking into this problem over the last 18 months, I'm surprised this construction took so long to find. The modular addition task with a 1-layer MLP is about as simple as you can get![^9gn5hx1hv4p]

Scaling mechanistic interpretability up to more complex tasks/models seems worth continuing to try, but I'm less sure extracting crisp explanations will be possible.[^dlczoq2as9d] Even if we "solve" [superposition](https://transformer-circuits.pub/2022/toy_model/index.html), figuring the construction here — where there's no superposition in the generalizing model — wasn't trivial. 

gif/twitter summary 
====================

If we train a MLP to solve modular addition, the generalizing phase has suggestive periodic patterns.

![](https://39669.cdn.cke-cs.com/rQvD3VnunXZu34m86e5f/animations/6f0906b2e835393c352fc88d395a61386da4a855ab8c60f8.gif)

To figure out why the model generalizes, we first look at task where we know the generalizing solution — sparse parity. You can see the model generalizing as weight decay prunes spurious connections.

![](https://39669.cdn.cke-cs.com/rQvD3VnunXZu34m86e5f/animations/1b4c67f67defe8c1912981317a3604e7eb9b4b70491cbe85.gif)

One point from the [Omnigrok paper](https://transformer-circuits.pub/2022/toy_model/index.html) I hadn't internalized before training lots of models: grokking only happens when hyper-parameters are just right.

We can make other weird things happen too, like AdamW oscillating between low train loss and low weights.

![](https://39669.cdn.cke-cs.com/rQvD3VnunXZu34m86e5f/animations/e2ef5dd4fa86505d0be98f21241de94fb005cb7e0b753657.gif)

To understand how a MLP solves modular addition, we train a much smaller model with a circular input embedding baked in.

![](https://39669.cdn.cke-cs.com/rQvD3VnunXZu34m86e5f/animations/76f0a790a898b197c8b41eaa8247e721d9036f1d72afd29d.gif)

Following [@Neel Nanda](https://www.alignmentforum.org/users/neel-nanda-1?mention=user) and applying a discrete Fourier transform, we see larger models trained from scratch use the same star trick!  
![](https://39669.cdn.cke-cs.com/rQvD3VnunXZu34m86e5f/animations/d1e4fc62667bd8ea00dbf28569f5a52e0342327ec09ae093.gif)

Finally, we show what the stars are doing and prove that they work:

![](https://39669.cdn.cke-cs.com/rQvD3VnunXZu34m86e5f/animations/1bdab2f5f3ab9e8c3e075bd019d4372701bab3c71d629733.gif)

Our ReLU activation has a small error, but it's close enough to the exact solution — an x² activation suggested in [Grokking modular arithmetic](https://arxiv.org/abs/2301.02679) — for the model to patch everything up w/ constructive interference.

[^9gn5hx1hv4p]: And there are still open question: why are the frequencies with >5 neurons lopsided? Why does factoring \(\textbf{W}_{\text{input}}\) not do that same thing as factoring \(\textbf{W}_{\text{output}}\)? 

[^dlczoq2as9d]: Also see The Hydra Effect: Emergent Self-repair in Language Model Computations