Grokking：神經網絡為何在記住資料很久後才突然泛化

Grokking: Why Neural Networks Suddenly Generalise Long After Memorising

Train a small neural network on modular arithmetic and watch the curves. Training accuracy hits 100 percent quickly. Test accuracy sits near chance for a long time. Then, after thousands more steps with no change in the training data, test accuracy jumps to near 100 percent as well. Nothing in the data changed. Nothing in the task changed. The network grokked.

Grokking is the delayed, sudden transition from memorisation to generalisation. It was first noticed clearly on small algorithmic tasks, but the mechanism behind it - a slow competition between a memorising solution and a generalising solution, decided by regularisation and optimisation dynamics - matters far beyond toy arithmetic. It is a warning against judging a model too early, and a window into how networks build internal structure.

The setup: a task that can be memorised or understood

The classic grokking setup is modular addition. Inputs are pairs (a, b), the target is (a + b) mod p, for some prime p. The training set contains only a fraction of all possible pairs, and the test set contains the rest. A model that memorises the training pairs will get training accuracy right and test accuracy wrong. A model that learns the actual addition rule will get both right.

The point of choosing such a clean task is that generalisation is unambiguous. There is a correct rule, it is simple, and it is either learned or it is not. In messier real tasks, memorisation and generalisation blend, which makes the transition harder to see but no less real.

A typical training objective adds weight decay to the usual loss:

L(θ)=Ldata(θ)+λ∣θ∣2

Here theta is the vector of parameters, L_data is the training loss, and lambda controls the strength of weight decay. That small second term turns out to be central to grokking, as we will see.

What the curves actually look like

In the early phase, the network memorises. It fits the training pairs by building a lookup-like mapping. Training loss falls to near zero. Test loss stays high. To a practitioner watching only the first part of training, the model looks like a failure: it has overfitted, and further training looks pointless.

In the middle phase, almost nothing visible happens on the test curve. Training accuracy stays at 100 percent. Test accuracy stays low. But inside the network, something is changing slowly. The weights are drifting, under the pressure of weight decay, toward a simpler configuration that also fits the training data.

In the late phase, test accuracy rises sharply, often over a relatively small number of steps compared with the long plateau before it. Training accuracy does not change, because it was already perfect. The model has not learned new data. It has reorganised what it already knew.

That shape - long plateau, then sharp jump - is the signature of grokking. It is not gradual improvement. It is a phase transition in the learned solution.

Two solutions, one loss landscape

Think of the network as having (at least) two kinds of solution available.

The memorising solution fits each training pair individually. It needs many effective parameters, because each pair is handled more or less separately. Its weight norm tends to be large: memorisation is expensive in norm.

The generalising solution implements the underlying rule. For modular addition, that rule can be represented with trigonometric features: the network can learn to map a and b to rotations, add the angles, and read off the result. That solution is compact. Its weight norm tends to be smaller, because one shared mechanism replaces many separate memories.

Early in training, gradient descent finds the memorising solution first because it is easier to reach. Fitting training points directly is a straightforward direction for the gradient. Building a clean trigonometric circuit takes longer: it requires coordinating many weights into a coherent structure, and the gradient signal for that structure is weaker while memorisation already gives near-zero loss.

But once training loss is near zero, the data term of the loss stops pushing hard. Weight decay keeps pushing, slowly, toward smaller norm. The memorising solution, with its large norm, is penalised more than the generalising solution. Over many steps, the network drifts from the expensive solution to the cheap one, while staying on the manifold of near-zero training loss. When the drift crosses a threshold, the generalising circuit becomes strong enough to dominate predictions on unseen pairs, and test accuracy jumps.

In short: optimisation finds the fast solution first, regularisation selects the simple solution eventually, and grokking is the time between those two events.

Why weight decay matters so much

Remove weight decay, or make it too small, and grokking may never happen, or may take far longer. The memorising solution has no reason to move. With strong weight decay, the pressure toward low norm is stronger, and grokking tends to happen earlier, though too much decay can also prevent fitting in the first place.

This also explains why grokking time is so sensitive to hyperparameters. Learning rate, initialisation scale, weight decay, and the fraction of data used for training all shift the balance between the two solutions and the speed of the drift. A small change can turn a run that groks in a few thousand steps into one that never groks within the training budget. That sensitivity is not noise to be averaged away. It is evidence that grokking is a competition with a sharp threshold, not a smooth learning curve.

A related quantity is the effective complexity of the solution. One way to formalise the intuition is to compare solutions by a norm or by a description length: the generalising solution is simpler in the sense that it can be described by a shared rule rather than a list of cases. Weight decay is a crude proxy for that simplicity, but in this setting it is a surprisingly effective one.

What is happening inside the network?

For modular addition, researchers have reverse-engineered the generalising circuit. The network learns embeddings that behave like Fourier features: each input is mapped to a set of sine and cosine components, the components are combined in a way that corresponds to adding angles, and a readout picks out the correct residue.

The memorising phase does not build that circuit. It builds something closer to a table. The transition to generalisation corresponds to the Fourier-like structure emerging and taking over, while the table-like components shrink under weight decay.

You do not need to follow every detail of that circuit to take the practical lesson. The lesson is that test accuracy can be flat while internal structure is changing in a way that matters enormously. A flat test curve is not proof that nothing is being learned. It may mean that the learned structure is not yet strong enough to change predictions, even though it is growing steadily underneath.

This is a useful corrective to the habit of early stopping at the first sign of overfitting. Early stopping is often right, but grokking shows a case where it would be exactly wrong: stopping at the plateau would freeze the memorising solution in place and discard a generalising solution that was on its way.

A minimal example you can run

The following PyTorch sketch trains a small model on modular addition. It is deliberately small: a two-layer network on one-hot inputs is enough to show the pattern on a reduced problem, and a small transformer shows it more cleanly on the full task. Treat this as a starting point for experiment, not as a tuned recipe, because grokking is hyperparameter sensitive.

import torch
import torch.nn as nn

p = 97
torch.manual_seed(0)

# All pairs (a, b), target (a + b) mod p
a = torch.arange(p).repeat_interleave(p)
b = torch.arange(p).repeat(p)
y = (a + b) % p
perm = torch.randperm(p * p)
train_idx = perm[:int(0.5 * p * p)]

class SmallMLP(nn.Module):
    def __init__(self, p, hidden=128):
        super().__init__()
        self.emb_a = nn.Embedding(p, hidden)
        self.emb_b = nn.Embedding(p, hidden)
        self.fc = nn.Sequential(
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, p),
        )
    def forward(self, a, b):
        h = self.emb_a(a) + self.emb_b(b)
        return self.fc(h)

model = SmallMLP(p)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1.0)
loss_fn = nn.CrossEntropyLoss()

for step in range(20000):
    idx = train_idx[torch.randint(len(train_idx), (512,))]
    logits = model(a[idx], b[idx])
    loss = loss_fn(logits, y[idx])
    opt.zero_grad()
    loss.backward()
    opt.step()
    if step % 2000 == 0:
        with torch.no_grad():
            train_acc = (model(a[train_idx], b[train_idx]).argmax(-1) == y[train_idx]).float().mean()
            test_mask = torch.ones(p * p, dtype=torch.bool)
            test_mask[train_idx] = False
            test_acc = (model(a[test_mask], b[test_mask]).argmax(-1) == y[test_mask]).float().mean()
        print(f"step={step} train_acc={train_acc:.3f} test_acc={test_acc:.3f}")

If you run a version of this, do not expect grokking on a fixed schedule. Vary weight decay and the training fraction, and log test accuracy over a long horizon. The interesting runs are the ones where test accuracy stays low for a long stretch and then jumps. If test accuracy rises immediately, you are likely in a regime where the generalising solution is found first, which is good for practice but does not show grokking.

Does grokking happen on real tasks?

This is the right question, and the honest answer is: not usually in this clean form, but the underlying mechanism shows up in related ways.

Real datasets are noisy, tasks are less cleanly separable into memorise versus understand, and training budgets, regularisation, and data scale all differ. You should not expect a single sharp jump in test accuracy on a large language model in the same way you see it on modular addition.

What does carry over is the idea of delayed generalisation through simplification. Several related phenomena point the same way:

• Longer training with strong regularisation can improve generalisation after training accuracy has saturated, because the solution keeps simplifying.
• Representations can become cleaner and more structured late in training, even when accuracy metrics look flat.
• Models can switch from a shortcut solution to a more robust solution when the shortcut is penalised, for example by weight decay, data augmentation, or a change in the training distribution.

In other words, grokking is an extreme, legible case of a broader pattern: the solution found first is not necessarily the solution that survives regularisation, and the difference between them may only appear after a long delay.

Limitations and common pitfalls

Do not generalise from one toy task. Modular addition is chosen because it is clean. Most real tasks do not have a single compact rule waiting to be found, and memorisation may be partly necessary, for example for facts that cannot be derived from a rule.

Do not assume longer training always helps. Grokking is a case where it does, but in many settings longer training after saturation just overfits further or wastes compute. The difference is whether a simpler, equally fitting solution exists and whether your regularisation pushes toward it.

Weight decay is not magic. It helps here because the generalising solution has smaller norm. In other settings, the robust solution may not be the smallest-norm one, and weight decay may hurt. Test the mechanism, do not just copy the hyperparameter.

A flat test curve is ambiguous. It can mean the model is stuck memorising, or it can mean a generalising structure is growing underneath. You cannot tell from the curve alone. Inspect representations, probe the model on structured test cases, or vary regularisation to see which explanation fits.

Grokking time is fragile. Reporting that a model groks without reporting weight decay, initialisation, training fraction, and optimiser settings is incomplete. Small changes move the transition by orders of magnitude, so reproducibility needs the full setup.

Practical takeaways

• If training accuracy saturates and test accuracy is flat, do not immediately conclude the run is done. Check whether a simpler solution could exist, and whether regularisation is still pushing the weights.
• Log internal structure, not just accuracy, on long runs: weight norms, representation geometry, and performance on structured probes can move before test accuracy does.
• Treat weight decay and other regularisers as part of the learning dynamics, not just as a way to prevent overfitting. They decide which of several fitting solutions survives.
• Be careful with early stopping on algorithmic or highly structured tasks. The best solution may arrive long after the training loss has bottomed out.
• When you see a sudden jump in generalisation, ask what changed underneath: a new circuit, a norm collapse, or a shift in representation. The jump is a symptom, and the mechanism is the interesting part.

Grokking is a reminder that learning is not just about reducing loss. It is about which solution, among the many that achieve low loss, the training process ends up with, and how long it takes to get there. On modular arithmetic, that process is visible as a long plateau followed by a sudden jump. On real tasks, it is quieter and messier, but the same question applies: is your model memorising the training set, or has it found a rule that will survive new data? Sometimes you have to wait a surprisingly long time for the answer.

Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial