首页 > AI前沿 > Flat Loss, Evolving Features: Grokking in Modular Addition

Flat Loss, Evolving Features: Grokking in Modular Addition

arXiv机器学习 2026-10-16 03:43 3 阅读 查看原文

Can a network keep learning while its training cross-entropy stays nearly flat at a positive level? We show that it can.

Modular Addition and Grokking

On modular addition, a standard testbed for grokking, ReLU networks trained with supported-Polar hidden updates, normalized output-layer updates and weight decay can keep 100% training accuracy while their hidden features change and held-out accuracy rises from zero to above 90%, at relatively flat training cross-entropy.

Correct classification need not minimize cross-entropy, which still depends on prediction confidence.

Finite-Window Theorems

We prove this coexistence for one-hidden-layer networks in two finite-window theorems from constructed initializations: one gives a constant gain in the Fourier fraction of the hidden features, and the other holds at arbitrarily low positive loss, where a readout refitted on the training data generalizes before the network.

Experiments

Experiments from random initialization show the same coexistence on a window fixed in advance, and in deeper networks held-out accuracy also rises over post-hoc windows with a net increase in training cross-entropy.

Deep Networks and Last-Layer Rate

With a last-layer rate set per training fraction and head decay scaled with the prime, deeper networks at p=47 and 61 meet a criterion on post-hoc windows, fixed in advance, in nearly all fresh runs.

Transformer Training

A one-block transformer trained with the same updates and a slowly decaying output gain keeps every training pair correct at every logged checkpoint while its training cross-entropy rises and its held-out accuracy approaches one; with two blocks this holds through the held-out rise, after which training pairs begin to flip.

A plateau in training loss need not be a plateau in learning.

Code: https://github.com/akashkumar-d/flat-loss-evolving-features