The most suitable audience for reviewing this material are Machine Learning Researchers, AI Theorists, and Specialists in Mechanistic Interpretability.
Abstract
This analysis details the phenomenon of "grokking" within transformer models, focusing on a single-layer transformer trained on modular arithmetic. Grokking is defined as sudden generalization that occurs long after a model has achieved perfect memorization (overfitting) of its training data. Mechanistic interpretability techniques, notably Fourier analysis and linear probing, reveal that this generalization phase corresponds to the model abandoning rote memorization and robustly implementing the modular addition operation via trigonometric identities. The model learns to encode inputs as sine and cosine components within its embedding space, leverages the identity $\cos(x)\cos(y) - \sin(x)\sin(y) = \cos(x+y)$ within its Multi-Layer Perceptron (MLP) block, and develops geometric structures (diagonal symmetries) to compute the required sum. A specialized metric, "excluded loss," confirms that structural learning progresses actively during the period where conventional accuracy plateaus. This work serves as a cornerstone for understanding core circuit mechanisms in transformers, paralleled by recent findings demonstrating 6-dimensional manifolds used for arithmetic in large-scale models like Claude Haiku.
Grokking Modular Arithmetic: A Mechanistic Interpretability Analysis
0:58 Complexity and Understanding: The video positions the single-layer transformer solving modular arithmetic via "grokking" as the most complex AI model mechanism currently understood in full detail, contrasting it with the black-box nature of large-scale models.
2:39 Modular Addition Setup: The model is trained on the operation $X+Y \pmod{M}$. A toy example uses $M=5$. The complexity is bounded, and the input is processed using sparse one-hot encoded vectors (e.g., 6 tokens for modulo 5, including the equals sign).
6:52 Accidental Discovery: Grokking was discovered when an OpenAI researcher accidentally left a small model training; performance initially overfit the training set rapidly (high training accuracy, poor test accuracy), then remained dormant for thousands of steps before suddenly achieving perfect generalization on the test set.
8:49 Mechanistic Analysis: Subsequent analysis, led by Neel Nanda, utilized a single-layer transformer trained on modulo 113 arithmetic. The model memorized the data around 140 training steps and generalized after 7,000 steps.
11:10 Structural Emergence: Examination of Multi-Layer Perceptron (MLP) activations (specifically the second layer) during the grokking phase reveals sine wave-like curves and circular "loop shapes" in neuron correlation plots, suggesting the network is learning a Fourier domain representation.
13:57 Fourier Components: Discrete Fourier Transform (DFT) analysis of the activation patterns identifies key frequencies (e.g., $8\pi/113$ and $6\pi/113$) corresponding to the learned waves.
15:13 Linear Probes and Signal Retrieval: The model learns to compute the sine and cosine components of the input variables ($x$ and $y$) by taking sparse weighted sums (linear probes) across its embedding vectors. This effectively extracts clean sinusoidal signals that are initially spread across multiple vector components.
16:59 Clock Analogy and Summation: Modular addition is analogized to an analog clock. The core computational challenge is combining the variables themselves ($x+y$), not merely summing their component representations (e.g., $\cos(x) + \cos(y)$).
20:43 Product Term Identification: 2D Fourier analysis of neuron outputs (varying both $X$ and $Y$ inputs) reveals that the strongest frequency component corresponds to the product of input cosines: $\cos(x) \times \cos(y)$.
21:42 Diagonal Symmetry in Deeper Layers: Neurons in the final MLP layer exhibit diagonal wave crests. These neurons fire maximally for combinations of inputs $(x, y)$ that sum to a constant value (e.g., $x+y = 65$), confirming that the model has learned the summation operation. This applies modularly, as $66+112 = 178 \equiv 65 \pmod{113}$.
23:19 Trigonometric Identity Implementation: The network successfully implements the trigonometric sum identity $\cos(x+y) = \cos(x)\cos(y) - \sin(x)\sin(y)$. This is achieved by combining the output surfaces of earlier neurons, applying weights of $+1$ to the $\cos(x)\cos(y)$ term and $-1$ to the $\sin(x)\sin(y)$ term, resulting in the desired diagonal symmetry surface ($\cos(x+y)$).
26:38 Excluded Loss Metric: To track learning during the plateau phase, the "excluded loss" metric was introduced. By measuring loss after removing the information corresponding to the key Fourier frequencies, researchers demonstrated that this metric slowly increases during the period of flat accuracy, confirming the model is actively constructing the structural solution and relying more heavily on the frequency-domain representation.
28:51 Cleanup Phase: Grokking occurs just after the structural phase is complete, during a "cleanup phase" where the model actively removes the previously memorized examples.
29:56 Extensibility to Large Models: The principles of mechanistic interpretability are shown to apply to full-size models; Anthropic researchers found that Claude 3.5 Haiku uses a 6-dimensional manifold to arithmetically track character count and line length, employing a "QK twist" mechanism in the attention blocks to calculate distance from the end of the line.
32:26 Final Takeaway: The detailed understanding of grokking demonstrates an intelligence mechanism that is highly efficient, mathematically elegant, and fundamentally "alien" or "ghost-like" compared to human cognitive processes.