Flat Loss, Evolving Features: Grokking in Modular Addition
Abstract
Can a network keep learning while its training cross-entropy stays nearly flat at a positive level? We show that it can. 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. 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 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. 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. 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.
Related
- Evolving LLM-Generated Features for Interpretable Classification
- Grokking or Glitching? How Low-Precision Drives Slingshot Loss Spikes
- Grokking through the Lens of Minimum-Norm Interpolation
- Algorithmic Information Dynamics of Learning: A Certified, Differentiable Complexity Controller for Grokking
- A Simple Proof of I. Chowla's Theorem on Addition of Residue Classes
BibTeX
@misc{kumar2026flat,
title = {{Flat Loss, Evolving Features: Grokking in Modular Addition}},
author = {Akash Kumar and Karthik Sankaranarayanan},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.22593},
eprint = {2610.22593},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.22593},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Flat Loss, Evolving Features: Grokking in Modular Addition
AU - Akash Kumar
AU - Karthik Sankaranarayanan
PY - 2026
DA - 2026/10/15
PB - arXiv
DO - 10.48550/arxiv.2610.22593
UR - https://arxiv.org/abs/2610.22593
AB - Can a network keep learning while its training cross-entropy stays nearly flat at a positive level? We show that it can. 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. 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 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. 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. 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.
ER -