For overparameterized neural networks, many solutions can fit the training data equally well while behaving very differently on unseen samples. Grokking separates training fit from visible generalization, providing a window for studying how this selection develops during training. We sweep short, fixed-duration weight decay (WD) perturbations across the pre-generalization plateau and measure how they shift later generalization time. Across three grokking tasks, these shifts are unordered early in the plateau but later form a stable dose ordering before visible generalization, with stronger WD increases leading to earlier generalization and stronger WD decreases leading to later generalization. Test-loss barriers between perturbed and baseline generalization checkpoints collapse toward zero while the ordered timing effects persist. A similar response reorganization is observed under ℓ1 regularization in the grokking setting of Junior et al. (2025). Drawing on Waddington's developmental landscape as an analogy, we call this combination of increasingly constrained solution selection and persistent dose-ordered timing shifts the canalization of grokking solution selection. Together, our response maps and loss-barrier measurements reveal a dynamical reorganization before visible generalization that is consistent with the theoretical picture of regularization-driven motion along a stable slow manifold (Boursier et al., 2025).
Figures & tables
Figure 1: Weight-decay pulse interventions and construction of the WD-response map. An example run from the Task 1 parity-matching task is shown. (a) Baseline training and test accuracy. The horizontal bar at the bottom marks the scan range of pulse start times, beginning at full training fit (epoch 375 ) and ending 1,000 epochs before the baseline first reaches 95% test accuracy (epoch 8625 ). (b) For each t0 , all perturbed branches start from the same full training state St0 . Weight decay is changed for τ epochs and then restored to its baseline value. (c) Generalization-time responses to different WD perturbation magnitudes at a fixed pulse start t0 . (d) WD-response maps for α=0.75 and α=0.90 . This run contains 331 pulse start times and 3,310 perturbed branches. Early in the plateau, the responses show no stable ordering; around t0≈5000 – 5500 , they begin to form a persistent dose-dependent structure, while the baseline is still on the chance-level plateau shown in panel (a), before any visible generalization. Intervention and scan settings are given in Appendix C .
Figure 2: WD responses develop a stable dose-dependent structure before visible generalization, and this pattern is reproduced across initializations and tasks. (a) WD-response curves ( α=0.90 ) for three example Task 1 runs; insets show the corresponding baseline training and test accuracy. (b) Mean dose–response linearity Rcorr2 across 30 Task 1 runs (left) and baseline relative weight norm (right). (c,d) WD-response maps ( α=0.90 ) for example runs from Task 2 and Task 3. Task and training settings are given in Appendix A ; dose–response and weight-norm analyses in Appendix B ; and intervention settings in Appendix C . Figures 10 – 12 show the signed ΔT0.95 response maps for all 70 analyzed runs, allowing the direction, dose ordering, and temporal persistence of the responses to be inspected run by run.
Figure 3: Test-loss barriers collapse while ordered WD timing responses persist. Each column shows one example run from a task (the Task 1 run is the same as in Figure 1 ), and the two rows share the same (t0,Δλ) intervention grid. For both the baseline and each perturbed branch, the generalization checkpoint is defined as the first time test accuracy reaches α=0.95 . (a) Test-loss barrier between the perturbed and baseline generalization checkpoints. We evaluate test loss at 51 equally spaced interpolation points along the straight line in parameter space connecting the two checkpoints, and define the barrier as the maximum excess above the linear interpolation of the endpoint losses (see Appendix C for the definition). Gray cells indicate perturbed branches that do not reach the target accuracy within the training limit. (b) Generalization-time shift ΔT0.95 of each perturbed branch relative to the baseline. Negative values indicate earlier generalization, and positive values indicate later generalization.
Figure 4: Transient ℓ1 perturbations reveal a similar reorganization of the generalization-time response. A representative run from the ℓ1 -regularized grokking setting of Junior et al. [1] is shown. (a) Baseline training and test accuracy. (b) Generalization-time shifts under transient changes in ℓ1 regularization strength. (c,d) Test- and training-cross-entropy barriers between perturbed and baseline branches at matched generalization checkpoints. The test-loss barrier remains close to zero throughout most of the scan, whereas the training-loss barrier retains a clear late-stage collapse similar to that observed in the WD experiments.
Figure 5: Waddington’s epigenetic landscape. Reproduced from Fig. 1 of Mitteroecker and Stansfield [29] under CC BY 4.0. Developing cells are depicted as balls rolling down valleys that become deeper and narrower over time, making perturbations increasingly less able to change their eventual direction. This landscape provides an intuitive analogy for the canalization of solution selection described here.
Figure 6: Baseline test loss (top) and global weight norm (bottom) across Tasks1̃–3. Task 1 shows little sustained norm decrease until test loss begins to improve, whereas Tasks 2 and 3 exhibit clearer norm contraction during late training. All quantities are evaluated every 25 epochs.
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
Task
Runs
Pulse starts
Avg. per run
Branches
Task 1 (parity match)
30
8,119
271
81,190
Task 2 (sparse parity)
24
8,975
374
89,750
Task 3 (factored modular addition)
16
6,832
427
68,320
Total
70
23,926
—
239,260
Appendix
Table 1: Scale of the WD-pulse intervention analysis across tasks. Each pulse start is evaluated with 10 WD perturbation magnitudes, giving 10 intervention branches per pulse start.
Figure 7: All baseline generalization trajectories included in the main analysis. We show 30 runs for Task 1, 24 runs for Task 2, and 16 runs for Task 3. Thin lines show baseline test accuracy for individual initializations, thick lines show the median across runs, and dashed lines indicate chance level.
Figure 8: Projected update-direction trajectories after early and ordering-onset WD pulses. Columns correspond to Task 1 seeds 2, 4, and 9; rows show early and ordering-onset pulses. Early pulses begin 500 epochs before the corresponding ordering-onset pulses. Black dashed curves show the baseline; blue and orange curves show branches with decreased and increased WD, respectively. Trajectories are displayed from the pulse end, with circles, triangles, squares, and diamonds marking the pulse end and 500 , 1,000 , and 2,000 epochs afterward. Each panel reports the number of training epochs shown after the pulse ends. Both panels within a run share a jointly fitted PCA basis, coordinate range, and axis aspect ratio; axis labels report the explained-variance ratios. Early branches show clearer deviations from the baseline, while ordering-onset branches follow largely overlapping paths in the displayed projection.
Figure 9: Existing grokking progress measures provide little advance signal during the Task 1 plateau. Each column corresponds to one clear-grokking run (seeds 2, 4, 9, and 21). (a) Nanda-style restricted/excluded decomposition. Restricted test BCE remains near chance level throughout the plateau and begins to decrease only about 50 epochs before the held-out test BCE begins to decrease (black curve); excluded train BCE continues to decrease and rises again near the generalization transition. The blue curve shows test accuracy (right axis). (b) Barak-style parameter and feature amplification. Parameter movement reaches a stable level early in training and changes little thereafter; the Fourier energy fraction of the joint-parity mode in the fourth hidden layer remains near zero during the plateau and rises near the generalization transition as held-out test BCE decreases.
Figure 10: WD responses and test-loss barriers across 30 independent Task 1 initializations. Run-selection criteria are described in Appendix A .
Figure 11: WD responses and test-loss barriers across 24 independent Task 2 initializations. Run-selection criteria are described in Appendix A .
Figure 12: WD responses and test-loss barriers across 16 independent Task 3 initializations, corresponding to seeds 1–16.
Figure 13: ℓ1 responses and loss barriers across three independent initializations. Rows correspond to seeds 14, 16, and 18, from top to bottom. Columns show generalization-time shifts, test-loss barriers, and training-loss barriers, respectively. Experimental settings and intervention details are provided in Appendix D .
Delayed generalization, or grokking, remains poorly understood despite extensive empirical study. We identify an exactly solvable late-time relaxation mechanism for grokking in linear models trained with full-batch heavy-ball optimization and weight decay, together with a locally quadratic extension to nonlinear neural networks. Our analysis reveals a distinguished population-active component of the empirical null space, which we call the grokking subspace. Along this subspace, the training predictions remain unchanged, leaving weight decay as the sole restoring force and giving rise to a slow dissipative relaxation governed by an exact discrete-time and continuous-time law. We show that only this subspace contributes to the slow asymptotic decay of the population risk and derive explicit iteration-scale predictions for the grokking time, recovering the familiar (1−β)/(ηλ) scaling in the weak-regularization regime. The theory further predicts distinct effects of optimizer choice, distinguishing coupled L2 regularization from decoupled weight decay, and yields causal predictions for interventions that modify the grokking component. We verify all theoretical identities without fitted parameters in a synthetic model where every subspace and relaxation rate is computable in closed form. We further observe genuine delayed generalization in modular addition, where the measured delay follows the predicted scaling and the late-time relaxation agrees closely with the theoretical clock.
Taeyoung Kim
Center for AI and Natural Sciences Korea Institute for Advanced Study Seoul 02455, Republic of Korea
Delayed generalization (\ie~grokking) refers to the phenomenon in which a neural network fits its training data early in training but only begins to generalize after a prolonged delay, often through an abrupt transition. Despite extensive empirical study, its underlying mechanism remains poorly understood. In this work, we first theoretically characterize a shell--core topological configuration of the reachable solution space induced by Adam's optimization dynamics with weight-shrinkage regularization, supported by empirical evidence. This optimization-induced topological configuration gives rise to grokking. In model's parameter space, random initialization solutions concentrate on a thin outer spherical shell, enclosing another spherical shell of memorization solutions, which in turn contains a core corresponding to the generalization solutions. Leveraging stopping-time theory, we then analyze the geometry of this topological configuration and the solution transition time at which optimization trajectories escape the memorization manifold and first reach the boundary of the generalization manifold. Our theoretical analysis derives grokking scaling laws for the learning rate, batch size, and ℓ2 regularization coefficient, which are further validated through experiments and shown to recover results from prior literature.
Róisín Luo, Christian Gagné, Jonas Ngnawé +2
Research Ireland – Centre for Research Training in AI · University of Galway · Universit´e Laval +2
Understanding generalization remains a central challenge in machine learning because it requires jointly considering data, architecture, and training dynamics. In this paper, we develop a theoretical framework that characterizes how these factors jointly shape generalization performance throughout training. More precisely, we study a broad class of neural networks trained under the ℓ2 loss by gradient descent (GD) with weight decay, and prove the convergence of GD to a neighbourhood of the global minimizers of the empirical loss. By partitioning the space based on the input data, we then decompose the population error into data error, optimization error, and prediction variation error, and bound them separately. In particular, for the prediction variation error, which measures the oscillations of the learned function, we propose (local) approximate homogeneity and derive explicit cellwise and layerwise bounds for its evolution along the training trajectory. These bounds yield two important implications: a necessary condition of improved generalization explains differences in layerwise generalization behavior; a sufficient condition describes delayed generalization and provides a theoretical characterization of grokking.
Yuqing Wang, Ioannis G. Kevrekidis, Mikhail Belkin
Johns Hopkins University · University of California San Diego