💻 computer science

Grokking the Discrete Logarithm: Scaling Across Model Capacity and Mechanistic Analysis

This paper presents the first systematic demonstration that transformers can exhibit "grokking"—a delayed transition from memorization to algorithmic generalization—on the full variable-base discrete logarithm problem, revealing that this phenomenon scales non-monotonically with model capacity and is driven by the emergence of multiplicative-character and Fourier-aligned representations.

Original authors: David Vesterlund

Published 2026-10-13

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: David Vesterlund

Original paper licensed under CC BY 4.0 (https://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the vast landscape of artificial intelligence, researchers are constantly trying to understand how machines learn. A particularly puzzling behavior, known as "grokking," has recently captured their attention. Imagine a student who memorizes every single answer in a textbook for a test. They can recite the answers perfectly, yet if asked a slightly different question, they fail completely. Grokking is the moment, often occurring long after the student has already memorized the material, when they suddenly stop just reciting facts and begin to truly understand the underlying rules. This shift from rote memorization to genuine generalization happens abruptly, sometimes after thousands of hours of practice, leaving scientists to wonder how and why the machine's brain suddenly reorganizes itself to find the logic behind the data.

This phenomenon was first noticed in simple math tasks, but a new study by David Vesterlund takes it into much more difficult territory. The researcher set out to see if a computer model could grok a problem that is notoriously hard for humans and machines alike: the discrete logarithm. In simple terms, this is a reverse-engineering puzzle. If you know a starting number and a final number, the task is to figure out exactly how many times you had to multiply the starting number by itself to reach the final one. While the forward calculation is easy, working backward is incredibly difficult, especially when the numbers are large and the rules are strict. This specific type of math problem is the foundation of modern digital security, making it a serious test for any learning system. Vesterlund wanted to see if a neural network, a type of AI designed to mimic the human brain, could eventually crack this code not by memorizing every possible answer, but by discovering the hidden mathematical structure that connects them.

To investigate this, the researcher built a series of computer models, ranging from small to very large, and fed them a massive dataset of these reverse-engineering puzzles. The models were trained on a specific prime number, 113, which defines the rules of the game. The smallest successful model had about 427,000 adjustable settings, a number that represents the complexity of its internal brain. When training began, the model behaved exactly as expected for a machine learning system: it quickly memorized the answers it had seen before, achieving near-perfect scores on the practice problems. However, when tested on new, unseen puzzles, it performed no better than random guessing. For a very long time, the model seemed stuck in this state of memorization, unable to generalize.

Then, after thousands of additional training steps, something dramatic happened. The model suddenly began to solve the new puzzles correctly, jumping from random guessing to nearly perfect accuracy. This was the moment of grokking. The delay was striking; the model required roughly three hundred times more training steps to achieve true understanding than it did to simply memorize the answers. To understand if this was a fluke or a fundamental property of these systems, the researcher tested twelve different models of varying sizes, from tiny ones with 90,000 settings to massive ones with over 14 million. The results showed that grokking is not limited to a specific size. It appeared in the medium-sized models and again in the largest ones, provided they were given enough practice data. Interestingly, there was a middle range where the models failed to learn, suggesting that the relationship between the size of the brain and the amount of data needed is complex and not a simple straight line.

To figure out what was happening inside the model's mind during this transformation, the researcher looked closely at the internal signals the model used to make decisions. They found that as the model moved from memorization to understanding, its internal language changed. It stopped relying on simple, surface-level patterns and began organizing its knowledge around a specific type of mathematical rhythm, similar to how sound waves can be broken down into distinct frequencies. By using a technique that searches for simple equations to describe these internal signals, the researcher discovered that the model had spontaneously learned to use trigonometric patterns—waves of sine and cosine—to represent the solution. This was a profound finding because the model was never explicitly taught to use these waves; it discovered them on its own as the most efficient way to solve the puzzle.

The study also included rigorous checks to ensure the model wasn't just finding a shortcut specific to the numbers used. The researcher tested the models on a different version of the problem where the rules were slightly more complex and the numbers were prime, removing any easy shortcuts. The models still grokked, proving they had learned a general rule rather than a specific trick. Furthermore, when the models were tested on generators they had never seen before, they still succeeded, showing they had truly learned the underlying logic of the system. The research suggests that the key to this sudden leap in understanding was a specific training technique that gradually increased the pressure to simplify the model's internal connections, forcing it to discard memorized details and keep only the essential structure.

This work establishes the discrete logarithm as a powerful new tool for studying how artificial intelligence learns. It demonstrates that even for difficult, inverse problems where no simple rule exists, neural networks can eventually transition from rote memorization to deep structural understanding. The findings confirm that this transition is not a magical event but a measurable process involving the reorganization of internal representations into efficient, wave-like patterns. While the study was conducted on relatively small numbers and does not solve the problem of breaking real-world encryption, it provides a clear window into the mechanics of learning. It shows that given enough time and the right conditions, these systems can move beyond memorizing the past to understanding the rules that govern the future.

Technical Summary: Grokking the Discrete Logarithm

Problem Definition

This study investigates the phenomenon of grokking—a delayed transition from memorization to algorithmic generalization in neural networks—on the full variable-base finite-field discrete logarithm problem (DLP). Unlike prior grokking research focused on forward algebraic operations (e.g., modular addition or multiplication), the DLP is an inverse mapping: given a generator $g$ and an element $h$ in a multiplicative group $\mathbb{F}_p^*$ , the model must find $x$ such that $g^x \equiv h \pmod p$ .

The task is defined as a classification problem over $p-1$ classes. The input consists of a generator $g$ (drawn uniformly from primitive roots) and an element $h$ , tokenized as $[BOS, g, SEP, h, EOS]$ . The study primarily focuses on the prime $p=113$ , where the total domain size is 5,376 unique triples $(g, h, x)$ .

Methodology

Model Architecture and Training

The author employs a GrokkingTransformer, a standard pre-norm encoder with GELU activations. The primary anchor model (M04) consists of 2 layers, $d_{model}=128$ , and 4 attention heads, totaling approximately 427k parameters.

Optimizer: AdamW with a constant learning rate ( $\eta=10^{-3}$ ).
Weight Decay (WD) Schedule: A critical component is a progressive weight decay ramp. WD starts at 0.05 and increases by 0.05 every 1,000 epochs only after the model achieves >99% training accuracy. It caps at 0.30 (scaled down for larger models based on parameter count).
Scaling Sweep: A 12-model grid (M01–M12) spans from 90k to 14.4M parameters ( $d_{model}$ from 56 to 768). Training fractions vary from 33% to 60% across the sweep to compensate for capacity differences.

Mechanistic Analysis

To understand the internal representations, the author utilizes:

Group Structure Probes: Testing consistency with algebraic identities (composition, cyclic shift, inverse).
Fourier Spectral Analysis: Projecting hidden states onto multiplicative character bases (aligned with discrete logs) and additive character bases.
Symbolic Regression (SR): Using PySR to fit equations to hidden layer activations. Two operator sets are used:
- Free SR: Standard operators without modular arithmetic knowledge.
- Algebra-aware SR: Includes precomputed modular features (e.g., $\cos(2\pi x/q)$ ).

Controls

Prime-Order Control: Experiments conducted in a subgroup of order 113 within $\mathbb{F}_{227}^*$ to rule out shortcuts based on the smooth factorization of $p-1$ (112).
Disjoint-Generator Split: Training and testing on non-overlapping sets of generators to test generalization beyond seen generators.

Key Results

1. Discovery of DLP Grokking

The anchor model (M04) exhibits a classic grokking trajectory:

Memorization Phase: The model reaches >99% training accuracy at epoch 90 but remains at chance-level test accuracy (~15%) for approximately 22,000 epochs.
Grokking Transition: After a prolonged plateau, test accuracy abruptly jumps to 99.8% by epoch 37,481.
Delay Ratio: The grokking delay ( $G = T_{90}/T_{mem}$ ) is approximately 298, indicating the model requires ~300 times more steps to generalize than to memorize.

2. Scaling Across Model Capacity

A 12-model sweep reveals a non-monotonic relationship between capacity and grokking:

Underfitting/Transition: Smaller models (M01–M03) fail to generalize. M04 (427k) succeeds.
Capacity Valley: Models M05 (657k) and M06 (1.09M) fail to grok at the initial 33% training fraction. However, a local experiment shows M05 does grok when the training fraction is increased to 40%, suggesting the failure was due to insufficient data rather than inherent model incapacity.
Large Models: Models M07 through M12 (1.64M to 14.4M parameters) successfully grok, with the time to grok decreasing as model size increases.

3. Training Dynamics

The transition involves 18 distinct "slingshot" events, where training accuracy temporarily collapses due to the interaction between AdamW and weight decay. These events precede the final generalization jump. Post-grokking, the weight norm decreases significantly (from ~148 to ~30), indicating a shift to a lower-norm solution.

4. Mechanistic Signatures

Algebraic Alignment: During the grokking transition, the model's internal consistency with group-theoretic identities (composition, cyclic shift) rises in tandem with test accuracy.
Spectral Sparsity: The hidden representation becomes sparse in the multiplicative character basis (Fourier harmonics of the discrete log), while the additive basis remains diffuse.
Symbolic Regression:
- Algebra-aware SR preferentially recovers trigonometric features aligned with multiplicative characters ( $\cos(2\pi x/q)$ ).
- Crucially, Free SR (without modular knowledge) independently recovers valid Fourier harmonics (e.g., frequencies corresponding to $\pi/2$ and $\pi$ ), confirming the structure is intrinsic to the learned representation and not an artifact of the operator set.

5. Structural Controls

Prime-Order DLP: The phenomenon persists in a prime-order subgroup (order 113), proving grokking does not rely on the smooth factorization of the group order.
Disjoint Generators: The model generalizes to generators unseen during training, achieving 100% test accuracy in the disjoint-generator split.

Significance and Claims

The paper claims to provide the first systematic demonstration of grokking on the full variable-base discrete logarithm mapping. Its primary contributions are:

Benchmarking: Establishing DLP as a compact, nontrivial benchmark for studying the interplay between memorization, representation formation, and generalization.
Scaling Laws: Demonstrating that grokking is not restricted to small models but reappears across a wide range of neural scales (up to 14.4M parameters), though it is sensitive to the ratio of model capacity to training data.
Mechanistic Insight: Providing evidence that successful generalization on this inverse problem corresponds to the emergence of a sparse, multiplicative-character-aligned Fourier representation.
Robustness: Showing that the phenomenon is not an artifact of specific group structures (smooth vs. prime order) or specific generators seen during training.

The author remains modest regarding practical cryptanalysis, noting that the experiments use small finite groups ( $p=113$ ) and do not claim to solve DLP for large cryptographic primes. Instead, the work serves as a controlled environment to study the "slingshot mechanism" and the formation of algorithmic circuits in deep learning.