Grokking: A Network That Memorizes, Stalls, Then Generalizes

I am reading some "old" AI papers at the moment to increase my knowledge/understanding of the field and build some basis so I can follow up more easily on the recent research. One of those papers was Power's 2022 paper "Grokking: Generalization beyond overfitting on small algorithmic datasets", Grokking is a bit of a slang way to say "delayed generalization".

The paper discusses the phenomenon where neural networks achieve sudden, perfect generalization on algorithmic tasks long after they have already completely overfitted and memorized the training data, and the authors show that it takes time to reach the "Aha!" moment where the model has finally generalize on unseen data while on the same time may have already overfit in train data. They also show that regularization is needed to help the model eventually discover the unterlying patterns with techniques like weight decay

In order to fully grasp the idea I thought was important to try to reproduce the paper with a small example/experiment I could do myself, end to end, on a laptop. The task was a silly clock that wraps but instead of 12 positions it has 82, dont ask me why 82, I just randomly picked that, so the model has to learn (a + b) mod p for p = 82. So this clock if it's at position 70 and you add 50, you don't land on 120, you wrap around to 38.

There are 82 x 82 = 6,724 possible (a, b) pairs. I generated all of them, shuffled, and split 50/50. The model never sees the correct answer for the held-out half during training. The only way it scores well there is by inferring the actual rule, not by pattern-matching against problems it's already seen.

The model

It's small on purpose: an embedding table plus a two-layer MLP.

class ClockMathBrain(nn.Module):
    def __init__(self, clock_size, representation_size):
        super().__init__()
        self.number_to_representation = nn.Embedding(clock_size, representation_size)
        self.mixer = nn.Sequential(
            nn.Linear(representation_size * 2, representation_size * 4),
            nn.ReLU(),
            nn.Linear(representation_size * 4, representation_size * 4),
            nn.ReLU(),
            nn.Linear(representation_size * 4, clock_size),
        )

    def forward(self, a, b):
        representation = torch.cat(
            [self.number_to_representation(a), self.number_to_representation(b)],
            dim=-1,
        )
        return self.mixer(representation)

Each number from 0 to 81 gets its own 128-dimensional embedding. The embeddings for a and b are concatenated and pushed through the Multilayer Perceptron, which outputs a score for each of the 82 possible answers. Nobody tells it about wraparound or addition. It has to recover both from the examples. I used ReLU as activation function between the layers, since is the most common used nowdays.

The one setting that matters: weight decay

The optimizer is AdamW with weight_decay=1.0, a much higher value than you'd typically use in a production model. This is the whole experiment. Weight decay penalizes the total size of the model's parameters, tracked as an L2 norm:

$$\text{parameter L2 norm} = \sqrt{\sum_i \|W_i\|_2^2}$$

where each $W_i$ is one weight or bias tensor. Every training step, weight decay quietly pulls this number down unless the loss gradient pushes back hard enough to justify keeping it up. In effect, the model isn't just fitting the practice data, it's also being charged rent on every parameter it uses to do so.

I trained for 3,000 steps, checking the practice score, the held-out (quiz) score, and the parameter L2 norm every 100 steps.

What happened

The run breaks into three distinct phases.

Steps 0-100: memorization. Practice score hits 1.00 almost immediately. Quiz score sits at 0.00. Classic overfitting: the model is looking up answers, not deriving them.

Steps 100-1,200: the plateau. Practice score stays pinned at 1.00. Quiz score doesn't move. If you stopped training here, you'd conclude the model had learned nothing generalizable and call it a day. Twelve hundred steps of watching a flat line is exactly why grokking is easy to miss in ordinary training runs. Nobody waits around after a metric flatlines.

Steps 1,200-2,400: the click. Quiz score breaks from 0.00, and climbs the rest of the way to 1.00 by step 2,400:

Nothing in the training loop changed between step 1,200 and step 1,300. The trigger was internal.

The weight norm is the clock

Here's the part that made this worth writing up instead of just running once and moving on. Plot the parameter L2 norm under the same timeline:

The norm falls the entire time, from about 106 right after memorization to 49.6 by the end. There's no kink in that curve when the quiz score takes off around step 1,300 to 1,900. Generalization doesn't happen at some special weight norm value. It happens partway down a smooth, monotonic decline that started long before the click and continues after it.

That's consistent with what two 2026 papers argue directly: The Weight Norm Sets the Grokking Timescale and What Does the Weight Norm Control in Grokking? both claim the norm's rate of decay predicts when the generalizing solution becomes reachable, not that grokking fires at some threshold value. While the model is memorizing, weight decay is silently shrinking a solution that overfits toward one that also generalizes, and the moment the trajectory crosses into the generalizing region is what looks, from the outside, like a sudden click.

Turning off the rule

So I ran the obvious ablation: set weight_decay=0.0 and retrained from scratch. Same architecture, same 6,724 problems, same 3,000 steps, only the weight decay coefficient changed.

Quiz score: 0.00 at step 0, 0.00 at step 3,000. It never moves. The parameter L2 norm does the opposite of the weight-decay run, it climbs from 104 to 130, because nothing is charging the model rent for growing its weights to fit the practice set more confidently. Without the constraint pulling the norm down, the model has no pressure to find the smaller, simpler solution that happens to generalize. It just gets better and better at memorizing, forever.

That's the cheapest, most direct evidence in this whole exercise that weight decay isn't a minor regularization nicety here. It's the mechanism.

Why bother with a toy clock

An 82-element addition table is silly on its own. Nobody needs a neural network to compute modular arithmetic. The reason this is worth building is that the same qualitative signature, instant memorization, a long flat stretch, then a sudden jump to generalization, has been observed in training runs on much larger, real models, where a full run costs real time and money and you can't afford to just retrain it five times to poke at one hyperparameter. Having an honest, minute-long reproduction on a laptop means you can test a claim like "the weight norm's decay rate predicts the click" by actually changing the decay rate and watching what happens, instead of reading someone else's plot and hoping it generalizes to your setup.

Sources:

• Power et al., Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets (2022), the original paper.
• The Weight Norm Sets the Grokking Timescale: A Causal Delay Law (2026)
• What Does the Weight Norm Control in Grokking? (2026)
• Muon Optimizer Accelerates Grokking (2025)