🤖 machine learning

Feature Repulsion and Spectral Lock-in: An Empirical Study of Two-Layer Network Grokking

This empirical study validates Tian's (2025) theoretical feature repulsion mechanism in modular addition networks, demonstrating that while the predicted sign structure of feature interactions holds robustly across activation functions, the resulting spectral signatures in parameter updates are strictly dependent on the activation derivative, with quadratic activations enabling a detectable rank-2 spectral lock-in during grokking that ReLU activations fail to produce.

Original authors: Yongzhong Xu

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yongzhong Xu

Original paper licensed under CC BY 4.0 (http://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

Imagine you are watching a student take a very difficult math test. At first, they are just memorizing the answers by rote (Stage 1). Then, they start to understand the underlying patterns, but they are still confused and making mistakes (Stage 2). Suddenly, at a specific moment, they have a "lightbulb" moment where they stop guessing and start getting perfect scores on questions they've never seen before. This sudden leap from memorization to true understanding is called "Grokking."

This paper investigates how and when that lightbulb moment happens inside a simple computer brain (a neural network). The researchers are testing a specific theory proposed by a scientist named Tian, which suggests that during the learning process, the computer's internal "ideas" (features) push each other away if they are too similar, forcing the brain to organize itself efficiently.

Here is the breakdown of their findings using simple analogies:

1. The "Repulsion" Rule (The Theory)

Tian's theory says that when two internal ideas in the computer brain become too similar, a mathematical force pushes them apart. Think of it like two magnets with the same pole facing each other; they repel.

What they tested: They checked if this "pushing apart" actually happens in real training.
The Result: Yes, it happens. It works perfectly whether the computer brain uses a "smooth" activation function (like squaring numbers, $x^2$ ) or a "bumpy" one (like ReLU).
The Metaphor: Imagine a crowded dance floor. If two dancers start moving in the exact same way, a rule kicks in that forces them to move to different spots. This rule works no matter what kind of music (activation function) is playing.

2. The "Spectral Lock-in" (The Observable Signal)

The researchers wanted to know: Can we see a specific signal on a computer screen that tells us the "lightbulb moment" is about to happen?
They looked at the "update spectrum"—essentially, how the computer's weights (the knobs it turns to learn) are changing over time. They were looking for a specific pattern: Rank-2 Lock-in.

The Analogy: Imagine the computer is trying to find the best path through a maze.
- Before the lightbulb: The computer is spinning its wheels, trying 100 different directions at once (chaos).
- The Lock-in: Suddenly, it realizes it only needs two specific directions to solve the maze. All other directions collapse into noise. The computer "locks in" to just those two paths.
The Result: This "locking in" signal is only visible when the computer uses the "smooth" activation ( $x^2$ $x^{2}$ ).
- With $x^2$ : The signal is loud and clear. The computer collapses its many directions down to exactly two.
- With ReLU: The signal disappears. The computer doesn't lock into two directions; it stays spread out, dominated by just one main direction.

3. The Big Discovery: "The Rule vs. The Result"

This is the most important part of the paper. The researchers found a split between the mechanism (the rule) and the signature (the visible result).

The Mechanism (The Repulsion): The rule that "similar ideas push each other apart" is universal. It happens with both types of computer brains ( $x^2$ and ReLU).
The Signature (The Lock-in): The visible sign that the computer has "figured it out" (the Rank-2 lock-in) is specific to the $x^2$ brain.

The Metaphor:
Imagine two different types of cars (one with a smooth engine, one with a bumpy engine) driving up a hill.

The Rule: Both cars have a rule that says, "If the wheels slip, the brakes must engage." This rule works for both cars.
The Signature: However, only the smooth car makes a distinct clicking sound when the brakes engage. The bumpy car engages the brakes too, but it just makes a humming noise.
The Lesson: If you only listen for the clicking sound, you might think the bumpy car isn't braking. But it is! It's just that the "sound" (the spectral signature) depends on the engine type, even though the "braking rule" (the repulsion) is the same.

4. Practical Takeaways for the "Driver"

The paper also offers some advice for people trying to detect this "lightbulb moment" in real-time:

Window Size Matters: To see the "clicking sound" (the lock-in), you have to look at the data over a specific time window (about 20 to 30 steps). If you look too quickly (5 steps), you get false alarms. If you look too slowly, you miss the moment.
Timing: The "clicking sound" happens right before the computer starts getting perfect test scores. It's a reliable early warning system, but only for the smooth ( $x^2$ ) type of computer.
The "Slow" Mode: If you slow down the learning process (by reducing a parameter called $\eta$ ), the "lightbulb moment" happens much later, but the math still holds up. The "clicking sound" is just quieter and takes longer to appear.

Summary

The paper proves that the "repulsion" rule Tian proposed is real and happens in all cases. However, the specific visual signal we use to detect it (the "Rank-2 Lock-in") is a trick of the light that only works for certain types of computer brains. If you are using a different type of brain (like ReLU), you have to look for a different signal, because the "lock-in" happens differently, even though the underlying logic is the same.

Technical Summary: Feature Repulsion and Spectral Lock-in

Problem Statement

Grokking—the phenomenon where neural networks abruptly transition from memorization to generalization long after training loss has plateaued—has been explained through various lenses, including mechanistic interpretability and implicit regularization. Tian [2025] provided a principled framework decomposing grokking dynamics in two-layer networks into three stages, culminating in an "Interactive feature learning" stage (Stage III). Within this stage, Tian's Theorem 6 posits a "repulsion theorem": similar hidden features repel each other due to the off-diagonal structure of the matrix $B = (\tilde{F}^\top \tilde{F} + \eta I)^{-1}$ .

However, Tian's framework left two empirical questions unanswered:

Timing: At what specific point in training does this repulsion mechanism become empirically observable?
Observability: Does this mechanism manifest as a measurable, online signature in quantities practitioners can compute without expensive offline diagnostics?

This paper addresses these questions by testing Theorem 6 and a candidate spectral observable on Tian's exact modular-addition setup.

Methodology

The study replicates Tian's [2025] modular-addition setup ( $M=71$ classes, $K=2048$ hidden units, training fraction $p \approx 0.40$ ) using Adam optimization with weight decay $\eta$ . The authors employ two complementary tests across multiple random seeds:

Direct Verification of Theorem 6:
- The authors compute $B$ exactly using the Woodbury identity to reduce computational complexity.
- At deterministic checkpoints, they identify the top-200 most-similar feature pairs based on cosine similarity.
- They verify the sign rule: $\text{sgn}(B_{j\ell}) = -\text{sgn}(\tilde{f}_j^\top P_{\eta, -j\ell} \tilde{f}_\ell)$ , where $P$ is a projection matrix excluding columns $j$ and $\ell$ .
- This verification is performed across different activation functions ( $\sigma(x)=x^2$ and $\sigma(x)=\text{ReLU}$ ).
Spectral Signature Detection:
- The authors monitor the rolling-window Gram matrix of parameter updates ( $\Delta W$ ).
- They track the eigengap ratio $\sigma_2/\sigma_3$ (the ratio of the second to third eigenvalues).
- A slope-based detector identifies the "lock-in" moment when this ratio enters a sustained rise, hypothesizing that repulsion consolidates redundant features into a low-rank structure (specifically rank-2).
- Sensitivity analyses are conducted on window size ( $W$ ) and weight decay ( $\eta$ ).

Key Results

1. Verification of Theorem 6 (Activation-General)

The sign rule of Theorem 6 holds empirically across different activation functions, confirming the mechanism is activation-general.

$\sigma(x) = x^2$ : The empirical sign-match for the top-200 similar pairs rises from 0.865 at epoch 50 to a tight saturation of 0.985 by epoch 300. Saturation ( $\ge 0.95$ ) occurs consistently at epoch 175 across seeds.
$\sigma(x) = \text{ReLU}$ : The sign rule also holds, saturating even faster (1.000 by epoch 500).
Conclusion: The repulsion mechanism, governed by the structure of $\tilde{F}^\top \tilde{F}$ , operates regardless of whether the activation is a power function or ReLU.

2. Spectral Lock-in (Activation-Specific)

While the mechanism is general, the parameter-update spectral signature is activation-specific.

$\sigma(x) = x^2$ (Focused Memorization): The slope-based detector fires in 15/15 grokking seeds at epoch 174. There is a massive magnitude separation ( $229\times$ ) in the late-stage $\sigma_2/\sigma_3$ ratio between grokking seeds (median $\approx 300$ ) and control seeds ( $\eta=0$ , median $\approx 1.31$ ). The spectrum collapses to rank-2, indicating features consolidate onto two persistent directions.
$\sigma(x) = \text{ReLU}$ (Spreading Memorization): The detector fires in 0/15 grokking seeds. The magnitude separation collapses to $1.4\times$ , and the spectrum remains rank-1 dominated ( $\sigma_1 \gg \sigma_2 \approx \sigma_3$ ).
Conclusion: The rank-2 lock-in signature only emerges in the "focused memorization" regime (power activations) where features collapse onto sharp peaks. In the "spreading memorization" regime (ReLU), features remain distributed, preventing the formation of a distinct rank-2 spectral signature.

3. Methodological Controls and Scaling

Window Size Sensitivity: The detector requires a window size $W \in \{20, 30\}$ . Smaller windows ( $W \le 10$ ) produce false positives in control conditions.
$\eta$ Scaling: An extended run at $\eta = 10^{-5}$ confirms Tian's $1/\eta$ scaling prediction for grokking time (grokking at epoch 1527). However, the lock-in magnitude ( $\sigma_2/\sigma_3$ ) is significantly reduced ( $\approx 25$ vs $\approx 300$ ), suggesting the rank-2 structure is incompletely developed at this slow timescale.
Level-Metric Detector ( $\rho_{tian}$ ): A simpler metric based on the activation Gram's off-diagonal structure predicts grokking onset with positive lead time for $\sigma=x^2$ but fails to generalize to ReLU or different $(M, p)$ configurations, marking feature dynamics initiation rather than successful grokking.

Significance and Claims

The paper draws a critical structure–mechanism distinction:

The Mechanism (General): Tian's Theorem 6 governs the off-diagonal sign structure of $B$ via properties of the feature Gram $\tilde{F}^\top \tilde{F}$ . This mechanism is robust and independent of the specific activation function $\sigma$ .
The Signature (Specific): The observable spectral signature in parameter updates depends on how the repulsion translates into weight updates, which is mediated by the derivative $\sigma'$ $σ^{'}$ .
- Power activations ( $\sigma=x^2$ ) produce focused features that consolidate onto two directions, yielding a rank-2 spectral lock-in.
- ReLU produces spreading features that remain rank-1 dominated, lacking the spectral lock-in signature despite the presence of the underlying repulsion mechanism.

The authors connect this empirical dissociation to Tian's Theorem 5, which formally distinguishes between focused and spreading memorization regimes. The paper concludes that while the repulsion mechanism is universal in Stage III, the spectral approach to grokking diagnostics is activation-dependent. Practitioners cannot assume a single spectral metric (like $\sigma_2/\sigma_3$ ) will detect grokking across all activation functions; the choice of diagnostic must align with the specific memorization regime induced by the activation function.