How AI Rediscovers the Language of Math

We know how every "neuron" that makes up an AI works, but this is akin to claiming we understand the grandeur of a cathedral just because we know the properties of a brick. When tens of billions of neurons are connected in ways too complex for us to fathom, and something we call "intelligence" emerges from a simple task like "predicting the next word," our comprehension begins to collapse.

Are these behemoths merely memorizing, or are they genuinely learning and reasoning?

A peculiar phenomenon known as "Grokking" offers us a perfect window into the depths of the AI "mind."

An Accidental Discovery: The "Grokking" Phenomenon

The story begins with a 2021 study from OpenAI. Researchers were training a very small Transformer model to perform a seemingly simple task: modular arithmetic.

You can think of modular arithmetic like a clock. On a 12-hour clock, 11 o'clock plus 2 hours is 1 o'clock, not 13, because time "wraps around" after hitting 12. The researchers set up similar rules for the model; for instance, in a modulo 5 system, the result of 4 + 2 is not 6, but 6 mod 5 = 1.

Initially, the results were unremarkable. The model quickly learned to give 100% correct answers on the training data, but its performance on test data it had never seen before was abysmal. It seemed to be yet another familiar tale of a model that could only resort to rote memorization.

However, the truly compelling chapter unfolded after a researcher returned from vacation. He discovered that after an unexpectedly vast number of training steps, the model seemed to experience an "epiphany." Its performance suddenly leaped, achieving near-perfect accuracy on both the training and test sets.

The model had perfectly memorized all the training answers after just a few hundred steps, then seemed to stagnate for thousands more, only to suddenly grasp the essence of the problem. How did this delayed "epiphany" happen?

The researchers named this phenomenon "Grokking." The term originates from a Martian word coined by science fiction writer Robert A. Heinlein in his 1961 novel Stranger in a Strange Land, meaning "to understand something so thoroughly, intuitively, and completely that you merge with it." The name aptly describes the model's leap from mechanical memorization to true understanding.

Delving into the "Black Box": An Elegant Feat of Mathematical Reasoning

The discovery of "Grokking" was stunning, but the mechanism behind it is what truly reveals the elegance of AI's "thought." A year after the OpenAI paper was published, a team led by researcher Neil Nanda conducted a "mechanistic interpretability" analysis on a single-layer Transformer model that exhibited Grokking, attempting to open this "black box."

The model they studied was tasked with addition modulo 113. As the researchers peered into the model's internal states—its neuron activations—during training, they found astonishing patterns.

When the input X was fixed and Y was iterated from 0 to 112, the outputs of certain neurons inside the model showed smooth, periodic structures resembling sinusoidal waves. When they visualized the outputs of different neurons as X-Y coordinates, these data points formed perfect circles.

The circular structure formed by the model's internal activations

source: Welch Labs / Neil Nandanda et al.

These waves and circles strongly suggested that the model was using trigonometric functions (sine and cosine) to represent the input numbers. It was as if the model had invented its own "clock": each input number was no longer an isolated symbol but was mapped to a specific position on a circle. Using the Discrete Fourier Transform, the researchers confirmed this hypothesis and precisely calculated the frequencies of these waves.

But simply encoding inputs into a circular structure isn't enough. The model's core task is to compute X+Y. It cannot simply add the two numbers, as it only receives abstract input patterns representing them.

The solution lay hidden deeper within the Multi-Layer Perceptron (MLP). The researchers found that some neurons in the MLP had learned to compute cos(x) * cos(y), while others had learned to compute sin(x) * sin(y).

And here, the most breathtaking part of the explanation took the stage.

When the neurons in the next layer combined these two results with specific weighted sums—for example, setting the weight for cos(x) * cos(y) to 1 and the weight for sin(x) * sin(y) to -1—their combination, cos(x)cos(y) - sin(x)sin(y), produced a marvelous interference effect, forming diagonal ripples.

Diagonal ripples produced by the trigonometric identity

source: Welch Labs / Neil Nandanda et al.

This is none other than the Sum of Angles Identity, deeply buried in mathematics textbooks, and its result is precisely equal to cos(x+y).

The model, without being explicitly told, had independently "rediscovered" this mathematical law that took humanity centuries to formulate. It used this law to transform the trigonometric representations of two separate inputs into the trigonometric representation of their sum, thereby brilliantly completing the addition operation.

Ultimately, the model's output layer learned to recognize these patterns formed by cos(x+y). For instance, the neuron responsible for outputting the answer "7" would be maximally activated for all input combinations where X+Y equals 7 (or equals 7 modulo 113).

The model's output layer response for X+Y=7

source: Welch Labs / Neil Nandanda et al.

Why the Delayed Epiphany? The Long Gestation from Memory to Wisdom

Now we can answer the initial question: why does Grokking happen with a delay?

In the early stages of training, the model takes a "shortcut": rote memorization of the answers in the training data. This is an inefficient but quick way to reduce training error. However, during the seemingly flat training plateau, the model was not idle.

Using a new metric called "Excluded Loss," researchers discovered that while the model's performance was stagnating, its internals were slowly but surely building the elegant trigonometric structure described above. This structure is far more efficient and generalizable than rote memorization, but it takes much longer to construct.

"Excluded Loss" metric reveals the model is still learning during the plateau phase

source: Welch Labs / Neil Nandanda et al.

The moment "Grokking" occurs is not when this trigonometric structure is first completed, but when it becomes strong and reliable enough for the model to completely discard the "memorization" crutch it previously relied on. It is a phase transition from an inefficient memorization strategy to an efficient algorithmic one—a moment where quantitative change leads to a qualitative leap.

From Toy Models to Claude: The Frontier of Interpretability

You might ask, this is just a toy model solving a simple math problem. Is it helpful for understanding behemoths like ChatGPT or Claude?

The answer is an emphatic yes. This "mechanistic interpretability" research method is being applied to the most advanced large language models. For instance, a team at Anthropic recently studied how Claude 3.5 Haiku decides when to insert a line break. They found that Haiku represents the character count of the current line and the total line width in a high-dimensional space (a six-dimensional manifold) with a structure strikingly similar to the circles we saw in the toy model. By rotating and comparing these high-dimensional representations in the attention heads (an operation known as the Q-K twist), the model can precisely calculate how much space is left before the end of the line and thus decide where to elegantly break it.

This suggests that the principle of "computation through geometric structures" discovered in the Grokking case may not be an exception, but rather a universal, intrinsic logic that AI models employ to solve a wide variety of problems.

Conclusion: Summoning Ghosts, Not Building Machines

As Andrej Karpathy, a prominent researcher in the AI field, commented, training large language models is "more like summoning a ghost than building an animal intelligence." The story of Grokking perfectly illustrates this point.

We communicate with models using human language, but this layer of interaction is just a thin veneer. Beneath it lies "extremely complex patterns" and an "alien" geometry that we struggle to intuitively grasp. The model's independent discovery of the trigonometric identity to solve modular arithmetic showcases its profound "understanding," yet it also feels alienating due to its non-human approach.

As our knowledge of the inner workings of AI deepens, will these intelligences feel more familiar to us, or more like "ghosts"? This is perhaps one of the most thought-provoking questions in the tech world for the next decade.