🤖 machine learning

Weight Decay Regimes in Grokking Transformers: Cheap Online Diagnostics

This paper demonstrates that weight decay serves as a critical control parameter governing the transition between memorization, generalization (grokking), and collapse in transformers trained on modular arithmetic, and introduces two computationally efficient online diagnostics based on attention activations to track these dynamics across various model scales and architectures.

Original authors: Lucky Verma

Published 2026-05-21

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Lucky Verma

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 teaching a group of students (the "attention heads" inside a computer brain) to solve a tricky math puzzle: modular arithmetic (like figuring out what time it is on a clock that only has 97 hours).

For a long time, these students seem to just be memorizing the answers by rote. They get perfect scores on the practice test but fail the real exam because they haven't actually learned the rules of the game. Then, suddenly, something magical happens: they "grok" the concept. They stop memorizing and start understanding, allowing them to solve any new problem instantly.

This paper is a field guide to understanding when and how this sudden "aha!" moment happens, and what happens if you push the students too hard or not hard enough.

Here is the breakdown of the findings using simple analogies:

1. The "Volume Knob" of Learning (Weight Decay)

The researchers discovered that a setting called Weight Decay acts like a master volume knob for the learning process. It controls the "personality" of the training:

Too Quiet (Low Weight Decay): The students are too relaxed. They just memorize the answers and never figure out the underlying logic. They stay stuck in "memorization mode."
Just Right (Medium Weight Decay): This is the "Goldilocks" zone. The students go through a specific two-step dance:
1. Sync Phase: First, all the students start thinking exactly the same way (they synchronize). They all agree on the basic rules.
2. Differentiation Phase: Then, they start to specialize. Some students become experts at addition, others at multiplication, etc. They stop being clones and become a diverse team. This is when the "grokking" happens.
Too Loud (High Weight Decay): If you turn the knob up too high, the students get overwhelmed. They all collapse into a single, confused state where they stop learning entirely.

2. The "Cheap Thermometers" (Online Diagnostics)

Usually, to know if a student is about to have an "aha!" moment, you have to wait until the end of the semester and give a final exam. That takes a long time and costs a lot of money (computing power).

The authors invented two cheap, real-time thermometers that you can check every few minutes while the students are still in class:

The "Group Hug" Meter (Cosine Similarity): This measures how much the students are agreeing with each other. When they start "grokking," they first hug (agree) tightly, and then slowly drift apart as they specialize.
The "Chaos" Meter (Entropy Standard Deviation): This measures how much the students are disagreeing or varying their thoughts. When the "aha!" moment is about to happen, this meter spikes, showing that the students are breaking out of their synchronized state to find their own unique roles.

These tools let researchers see the "phase change" happening in real-time without waiting for the final exam.

3. The "Critical Tipping Point"

The paper found a very specific number (a "tipping point") for the volume knob.

If the setting is below 0.0158, the students almost never learn the rules; they just memorize.
If the setting is above 0.0158, they almost always learn the rules.
The time it takes for them to learn gets faster the higher you turn the knob (up to a point), following a predictable mathematical curve.

4. The "Late-Stage Crash" (Anti-Grokking)

There is a twist. If you train these students for way too long (20,000 steps instead of the usual few thousand), some of them might suddenly forget everything and go back to guessing randomly. The paper calls this "anti-grokking."

Good News: This doesn't happen to everyone. It depends on the "seed" (the random starting conditions). Some students are fragile and crash; others are tough and stay smart forever.
The Cause: It seems to be caused by the students losing their unique roles and becoming too similar again, but in a broken way.

5. Does this work for other "Brains"?

The researchers tested if this "Volume Knob" rule works for different types of computer brains, not just the standard ones (Transformers).

They tried it on MLPs (simple feed-forward networks), LSTMs (old-school memory networks), and Mamba (a newer, efficient architecture).
Result: Yes! The "Volume Knob" still works. All of them have a "too quiet," "just right," and "too loud" zone. However, the exact number on the knob needed to hit the "just right" zone is different for each type of brain.

Summary of What They Don't Claim

They are not saying this applies to giant language models (like the ones that write essays or chat with you) yet. They only tested small math models.
They are not claiming to have found the ultimate "law of physics" for all AI. They found a specific pattern for this specific math puzzle and are careful to say it might be different for other tasks.
They are not saying this is a medical cure or a biological fact, even though they used words like "heartbeat" and "synchronization" as metaphors.

In a nutshell: The paper shows that by turning a single dial (Weight Decay), you can control whether a computer brain memorizes, learns deeply, or crashes. They also gave us two simple tools to watch this process happen in real-time, proving that learning often happens in two distinct stages: first, everyone agrees; second, everyone specializes.

Technical Summary: Weight Decay Regimes in Grokking Transformers

Problem Statement
Grokking refers to the phenomenon where transformers trained on algorithmic tasks (specifically modular arithmetic) exhibit a sharp transition from memorization to generalization long after training loss saturates. While mechanistic studies have identified the formation of internal circuits (e.g., Fourier representations) and induction heads as drivers of this transition, existing literature lacks a "cheap online diagnostic" to track these dynamics in real-time. Furthermore, the role of weight decay ( $\lambda$ ) as a control parameter for these regimes has not been systematically mapped across model scales, nor has the specific contribution of attention-head coordination been isolated from weight magnitude effects.

Methodology
The authors train decoder-only transformers on modular arithmetic tasks (mod+, mod-, mod $\times$ , mod $\div$ ) using the AdamW optimizer. The study employs a dense sweep of weight decay values and varies model scales (from 0.82M to 85M parameters) and architectures (Transformers, MLPs, LSTMs, Mamba).

Key methodological components include:

Online Order Parameters: Two scalar diagnostics are computed at every training step using only attention activations:
- Mean Pairwise Attention-Head Cosine Similarity ( $\bar{s}(t)$ ): Measures the coordination (synchronization) among attention heads.
- Entropy Standard Deviation ( $\sigma_H(t)$ ): Measures the dispersion of Shannon entropy across heads, tracking differentiation.
Finite-Size Scaling: The authors fit a logistic function to the probability of grokking ( $P(\text{grok})$ ) versus weight decay to estimate a critical threshold $\lambda_c$ . They also fit a power law to the "time-to-grok" ( $t_{\text{grok}}$ ) to determine the critical exponent $\nu$ .
Causal Interventions: To distinguish between head-pattern structure and weight magnitude, the authors perform paired experiments where attention heads are re-initialized versus where weight norms are clipped at the peak of the differentiation phase.
Formal Verification: Diagnostic identities (bounds and algebraic properties) are machine-checked in Lean 4 to ensure mathematical well-formedness.

Key Contributions

Quantitative Critical Threshold: The study identifies a sharp weight-decay critical threshold for the canonical 4-layer, 8-head transformer on mod+ tasks: $\lambda_c = 0.0158$ (95% CI [0.0109, 0.0200]).
Regime Diagram: A two-axis ( $\lambda$ $λ$ , $N$ $N$ ) map reveals three distinct regimes:
- Memorization: $\lambda < \lambda_c$ (near-zero grokking).
- Developmental Grokking: $\lambda \ge \lambda_c$ (reaching ~100% grok rate, with time-to-grok decreasing as $\lambda$ increases).
- Collapse: $\lambda = 10$ (all heads converge to identical patterns, $\bar{s}=1.0$ ).
Two-Phase Dynamics: The training trajectory is characterized by:
- Phase 1 (Synchronization): Heads converge ( $\bar{s}$ rises from ~0.93 to 0.995) coinciding with the onset of grokking.
- Phase 2 (Differentiation): Heads diverge ( $\bar{s}$ drops to ~0.88, $\sigma_H$ increases) while test accuracy remains high.
Per-Head Dimension as Amplitude Modulator: The amplitude of differentiation ( $\sigma_H$ ) is empirically dominated by the per-head dimension ( $d/H$ ), saturating monotonically as $d/H$ increases.
Cross-Architecture Scope: The weight-decay controlled transition is observed in non-attention architectures (MLP, LSTM, Mamba), though the critical threshold $\lambda_c$ shifts significantly (spanning roughly an order of magnitude) depending on the architecture.

Results

Scaling Exponent: The time-to-grok follows a power law $t_{\text{grok}} \propto (\lambda - \lambda_c)^{-\nu}$ with an empirical exponent $\nu = 0.757$ (95% CI [0.725, 0.799]). The authors explicitly note that this value does not match standard reference exponents (e.g., $\nu=1/2$ or 3D Ising $\nu \approx 0.63$ ) under their current grid density, deferring universality-class identification to future work.
Long-Horizon Fragility: At extended horizons (20,000 epochs), a seed-dependent "anti-grokking" collapse is observed in some trajectories, characterized by a five-stage pattern ending in accuracy decay. This is identified as a fragility rather than a universal cycle.
Causal Isolation: Re-initializing attention heads significantly reduces the Phase 2 amplitude ( $\sigma_H$ ) compared to controls, whereas weight-norm clipping does not. This isolates the effect to head-pattern structure rather than weight magnitude.
Multi-Task Replication: The weight-decay control pattern holds across four modular operations and two model scales, confirming the robustness of the regime map beyond addition.

Significance and Claims
The paper claims to provide a cheap, online-computable diagnostic framework for tracking grokking dynamics using only attention activations, complementing more expensive loss-landscape diagnostics like the refined local learning coefficient (rLLC). It establishes weight decay as a primary empirical control parameter that separates memorization, developmental grokking, and collapse regimes.

The authors are modest in their scope:

They claim no universality class for the observed exponent $\nu \approx 0.76$ , noting that tested reference exponents fall outside their confidence intervals.
They do not claim that weight decay is necessary for generalization in all architectures, only that it acts as a control parameter within standard attention architectures.
They explicitly defer causal mechanism claims regarding why differentiation occurs to follow-up work, focusing instead on when and how it manifests.
The findings are scoped to modular arithmetic in small transformer attention models; claims regarding large language models or broader universality are out of scope.

The work positions itself as a phenomenological and diagnostic contribution, providing a "regime map" and "order parameters" that complement concurrent theoretical work on finite-size scaling, Lyapunov contraction, and glass relaxation.