💻 computer science

Mechanistic Interpretability of Discrete Logarithm Grokking: How Representation Determines the Internal Algorithm

This paper demonstrates that the input representation of the discrete logarithm problem fundamentally dictates the internal algorithm discovered by neural networks, revealing that algebraic alignment with the problem's true factorization is a prerequisite for rapid generalization (grokking) and that the necessary interaction-rich computational structures form during the memorization phase long before performance improves.

Original authors: David Vesterlund

Published 2026-10-05

📖 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 quiet corners of computer science, researchers are trying to understand how artificial intelligence actually thinks. For years, we have known that these systems can learn to solve complex puzzles, but the process often looks like magic: a model might fail repeatedly for a long time, suddenly memorize the answers, and then, after a strange pause, begin to generalize and solve new problems it has never seen. This phenomenon, where a machine seems to "grok" a concept long after it has simply memorized the data, has become a central mystery. To crack this code, scientists often use mathematical problems that have a clear, hidden structure, much like a lock with a specific key. One such problem is the discrete logarithm, a task that asks a computer to reverse a multiplication process to find a hidden number. While this sounds abstract, it is the kind of logical puzzle that reveals the inner machinery of neural networks. By watching how a computer learns to solve this specific riddle, researchers can see exactly what internal strategies the machine invents and whether the way we show it the problem changes the solution it finds.

A new study by David Vesterlund takes this investigation a step further by asking a simple but profound question: does the way we present information to a computer change the way it thinks? To find out, the researcher trained artificial intelligence models to solve the discrete logarithm problem using two very different methods of showing the data. In the first method, the model was given the raw numbers, just as a human might see them on a page, with no special hints about the underlying structure. In the second method, the researcher broke the problem down into its natural mathematical parts before showing it to the model, effectively handing the computer a pre-sorted toolkit. The results revealed that these two approaches led the models to discover completely different internal algorithms. The model trained on the raw data developed a complex, tangled web of connections where different parts of the problem were deeply intertwined. In contrast, the model given the pre-sorted, structured data developed a clean, separated organization where the different parts of the problem were handled independently.

Perhaps the most surprising discovery was the timing of this learning. The researchers found that the complex, tangled structure in the raw-data model began to form long before the model actually started to get the answers right. For thousands of steps, the model's internal connections were rearranging themselves into this intricate pattern while its test scores remained stuck at a low level. It was only much later, after this internal structure was fully built, that the model suddenly "clicked" and began to solve the problem correctly. This suggests that the machine was building the necessary machinery for understanding long before it could demonstrate that understanding. The study also showed that this internal architecture is not just a side effect but is essential for the task. When the researchers surgically removed specific parts of the model's internal math, the system failed completely, proving that these specific patterns were the actual engine driving the solution.

The study also corrected a misunderstanding from previous research regarding how to test these models. A common technique involved scrambling the labels of the data to see if the model relied on the specific order of numbers. The new analysis proved that this scrambling method was actually a trick; because the model learns its own dictionary of symbols, simply renaming the numbers did not actually change the underlying mathematical structure it was seeing. The model was still solving the same structured problem, just with different names. This finding was crucial because it cleared the way for a more honest test. The researchers then created a hierarchy of new tests, some of which kept the mathematical structure intact and others that deliberately broke it. They found that models could only learn to solve the problem quickly if the input preserved the true mathematical "partition" of the problem. If the input scrambled the relationship between the parts of the equation, the model would memorize the training data perfectly but would never learn to solve new problems, even after being trained for a very long time.

Ultimately, this work suggests that the input representation does not just speed up or slow down learning; it fundamentally determines the algorithm the machine discovers. When the problem is presented in a way that respects its natural mathematical divisions, the machine finds a simple, efficient path to the solution. When the problem is presented in a raw, unstructured way, the machine is forced to build a much more complex internal world to bridge the gaps, and it takes much longer to realize how to use that complexity. The study confirms that for these artificial systems, the path to understanding is paved by the structure of the information they are given. It is not merely a matter of how much data they see, but how that data is arranged. The machine does not just learn the answer; it learns a specific way of thinking that is dictated by the shape of the question itself.

Technical Summary: Mechanistic Interpretability of Discrete Logarithm Grokking

1. Problem Statement

This research investigates how the input representation of an algebraic learning problem dictates the internal algorithm discovered by a neural network. The study focuses on the Discrete Logarithm Problem (DLP) within the context of grokking—the phenomenon where a model memorizes training data long before suddenly generalizing to test data.

The core question is not merely if grokking occurs, but what internal mechanism the network constructs and how this mechanism is shaped by whether the input representation exposes or disrupts the underlying algebraic structure (specifically the Chinese Remainder Theorem, or CRT, factorization of the cyclic group).

2. Methodology and Experimental Setup

2.1 Task and Data

Problem: Solve $h = g^x \pmod p$ for $x$ , where $p=113$ (a prime) and the group order is $N = p-1 = 112 = 7 \times 16$ .
CRT Decomposition: The problem allows a natural factorization into $x \pmod 7$ and $x \pmod{16}$ .
Dataset: All 5,376 triples $(g, g^x, x)$ generated from 48 generators. Split: 30% training, 70% testing.
Model: A 2-layer Transformer (427K parameters, 128 hidden dimension, 4 heads).

2.2 Input Representations

The study compares several input encodings to test the necessity of algebraic alignment:

RAW (R0): Standard sequence $[BOS, g, SEP, h, EOS]$ . No explicit factorization.
TRUE CRT (R1): Explicitly factorized input $[BOS, g_b, SEP, g_a, SEP, h_b, SEP, h_a, EOS]$ , exposing the subgroup projections.
SCRAMBLED: Applies a single fixed permutation to all CRT positions.
Negative Controls (Hierarchy of Algebraic Destruction):
- PERMUTED CRT (R2): Independent permutations per factor (preserves partition, destroys ordering).
- ROLE-DECOUPLED CRT: Different permutations for generator vs. result positions (preserves partition, destroys shared conjugacy).
- GLOBAL RANDOM BIJECTION (R3): Random bijection destroying the group law and factor partition.
- ORDER-MATCHED (R3O) & MIXED RADIX (R4): Variations that preserve information but destroy algebraic structure.

2.3 Analysis Techniques

2D CRT Fourier Decomposition: Hidden activations are decomposed into Fourier modes: $E_A$ (factor A only), $E_B$ (factor B only), and $E_{AB}$ (cross-factor interaction). The interaction ratio $R_{int} = E_{AB} / (E_A + E_B + E_{AB})$ quantifies representational separability.
Temporal Analysis: Tracking Fourier energy and test accuracy across 32 checkpoints (0–65K epochs) to determine when internal structures form relative to generalization.
Causal Ablation: Top-k ablation of Fourier subspaces ( $U_A, U_B, U_{AB}$ ) with 100 rank-matched and variance-matched random controls to verify causal necessity.
Algebraic Structure Scores: Metrics including homomorphism score ( $S_{hom}$ ) and adjusted Rand index (ARI) to quantify how well a representation preserves the group law and factor partition.

3. Key Contributions

Divergent Internal Algorithms: Demonstrated that RAW and CRT representations lead to qualitatively different internal organizations. RAW develops a non-separable, interaction-rich structure ( $E_{AB} \approx 0.65$ ), while CRT develops a separable structure ( $E_{AB} \approx 0.01$ ), despite solving the same task.
Temporal Formation of Interaction: Showed that the interaction-rich structure in RAW models forms during the memorization plateau (epochs 15K–30K), approximately 20K epochs before visible generalization (grokking at ~50K).
SCRAMBLED Audit: Proved that the previously studied "SCRAMBLED" representation is exactly conjugate to TRUE CRT at initialization. Because the model uses learned categorical embeddings, a global token permutation is absorbed by the embedding table, meaning SCRAMBLED preserves all algebraic structure ( $S_{hom}=1.0$ ) and does not test algebraic destruction.
Hierarchy of Negative Controls: Established that preservation of the true CRT factor partition is the critical requirement for rapid generalization. Representations preserving the partition (TRUE CRT, PERMUTED CRT, ROLE-DECOUPLED CRT) all grok, while those destroying the partition (R3, R3O, R4) fail to grok even after 100K epochs.
Causal Necessity of Fourier Modes: Confirmed via top-k ablation that A-modes, B-modes, and AB-interaction modes are all causally necessary for computation ( $p < 0.01$ ).

4. Key Results

4.1 Representation Determines Algorithm

RAW Model: Develops a complex internal algorithm relying on cross-factor interactions. The interaction energy ( $E_{AB}$ ) rises from 35% to 52% during the plateau while test accuracy remains flat (~30%).
CRT Model: Develops a separable algorithm. The model discovers the component information ( $x \pmod 7$ and $x \pmod{16}$ ) early but learns to recombine them later.
Conclusion: The input representation does not just change the speed of learning; it fundamentally changes the internal algorithm (separable vs. interaction-rich).

4.2 The Necessity of Factor Partition Preservation

Fast Grokking: TRUE CRT (~~20K epochs), PERMUTED CRT (~~8K), and ROLE-DECOUPLED CRT (~26K) all achieve near-perfect generalization.
Failure to Grok: GLOBAL RANDOM BIJECTION (R3), ORDER-MATCHED (R3O), and MIXED RADIX (R4) fail to generalize ( $T_{grok} > 100K$ ) despite full memorization and maximum weight decay.
Implication: Merely factorizing the input into two coordinates is insufficient. The representation must preserve the true algebraic partition (the specific mapping of group elements to the CRT factors).

4.3 The "SCRAMBLED" Misconception

The study corrects a potential misinterpretation in prior literature. Because SCRAMBLED is conjugate to CRT, its rapid grokking does not imply that algebraic alignment is unnecessary. It simply confirms that the model can learn the algorithm when the algebraic structure is preserved via vocabulary relabeling.

4.4 Causal Verification

Ablating the top-k Fourier subspaces causes catastrophic performance drops ( $\Delta \approx -0.89$ to $-0.96$), whereas random controls show negligible impact. This confirms that the identified Fourier modes are not just correlated features but are causally required for the computation.

5. Significance and Claims

The paper claims that input representation is a primary determinant of the internal algorithm a neural network discovers.

Mechanistic Insight: Grokking is not a monolithic phenomenon; the path to generalization depends on whether the network must construct complex interaction modes (RAW) or can rely on pre-exposed separable structures (CRT).
Algebraic Alignment: Rapid generalization in algebraic tasks requires the input representation to preserve the true factor partition of the underlying group structure. Simple information preservation or arbitrary factorization is insufficient.
Methodological Correction: The study highlights a critical pitfall in mechanistic interpretability: "scrambling" inputs via token permutations may be vacuous if the model uses learned embeddings, as the permutation can be absorbed without destroying the underlying algebraic logic.
Temporal Dynamics: The internal structure required for generalization (specifically cross-factor interactions in RAW) forms well before the model exhibits generalization capabilities, suggesting that the "plateau" is a phase of structural construction rather than just memorization.

The author maintains a modest scope, noting limitations such as the use of a single seed for temporal analysis and the current restriction to $p=113$ , framing these as directions for future work rather than current conclusions.