cs.LGAug 26, 2026
SaveMapping the Emergence of Regularization-Driven Dynamics in Grokking
Organizations: School of Artificial Intelligence, University of Chinese Academy of Sciences
Abstract
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 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 ) and ending epochs before the baseline first reaches test accuracy (epoch ). (b) For each , all perturbed branches start from the same full training state . 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 . (d) WD-response maps for and . This run contains pulse start times and perturbed branches. Early in the plateau, the responses show no stable ordering; around – , 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 ( ) for three example Task 1 runs; insets show the corresponding baseline training and test accuracy. (b) Mean dose–response linearity across 30 Task 1 runs (left) and baseline relative weight norm (right). (c,d) WD-response maps ( ) 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 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 intervention grid. For both the baseline and each perturbed branch, the generalization checkpoint is defined as the first time test accuracy reaches . (a) Test-loss barrier between the perturbed and baseline generalization checkpoints. We evaluate test loss at 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 of each perturbed branch relative to the baseline. Negative values indicate earlier generalization, and positive values indicate later generalization.
Figure 4: Transient perturbations reveal a similar reorganization of the generalization-time response. A representative run from the -regularized grokking setting of Junior et al. [1] is shown. (a) Baseline training and test accuracy. (b) Generalization-time shifts under transient changes in 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 |
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 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 , , and 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 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: 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 .