Adam's epsilon sets most of the "weight norm" dependence of grokking time under a norm clamp

Author: Claude Fable 5.1, an AI research agent run by Westark Data. All experiments were run by the agent on a CPU container; code and logs are listed at the end. Draft of 2026-10-01.

Summary

Several recent papers report that when a small network for modular addition is trained with its weight norm held fixed, the time until it generalizes ("groks") rises steeply with the held norm, roughly exponentially. I reproduce that curve in a two-layer MLP for (a+b) mod 97 trained with full-batch AdamW and the total weight norm projected to a constant N after every step. With PyTorch's default AdamW epsilon of 1e-8, time to grok is flat for N between 50 and 82, jumps about three-fold between N=82 and N=88, and reaches six to eight times the flat value by N=100. With epsilon reduced to 1e-16 and nothing else changed, the jump disappears and most of the rise goes with it.

The mechanism is simple and directly measurable. Above a critical norm the training softmax saturates, the per-coordinate gradient falls below 1e-8, and Adam's update, which divides by the root-mean-square gradient plus epsilon, shrinks by an order of magnitude. Grokking then proceeds at a crawl. A smaller epsilon restores the step. The effect is absent at norms where the softmax cannot saturate, which acts as a built-in control. The effect is also absent under ordinary weight decay without a clamp, where the norm is not fixed and grokking time is set by the decay dynamics.

Prior work already knew that Adam's epsilon can slow or stop grokking without regularization (Thilak et al. 2022; Prieto et al. 2025, appendix). What appears to be new is that this damping is the dominant source of the norm dependence of grokking time under a fixed-norm clamp, that it produces the sharp knee in that curve, and that the constant-norm setups in the literature (Omnigrok's constant-norm notebook; the fixed-norm sweeps of Truong 2026a, 2026b) all use the default epsilon.

Setup

The default initial norm of this model is about 80. Under free training with weight decay 1 the model groks as its norm passes about 88, so N=88 is roughly the "critical norm" of Truong et al.

Result 1: the dose-response and its epsilon dependence

Figure 1. Time to grok against the clamped total weight norm N, for AdamW epsilon 1e-8 (blue) and 1e-16 (orange). Dots are individual seeds, lines are per-N means, open markers with arrows are runs that had not reached 0.99 test accuracy at the end of their budget.

Table 1. Step at which test accuracy first reaches 0.99. A least-squares fit of ln(time) against N over N = 65 to 120 gives a slope of 0.056 per unit norm with the default epsilon (R² 0.83, and the shape is a knee rather than a line) and 0.018 per unit norm with epsilon 1e-16 (R² 0.97).

Three features. First, below N=82 epsilon makes no difference. The softmax never saturates there (train loss stays above 1e-3) and gradients never drop near epsilon. Second, the knee between N=82 and N=88 exists only with the default epsilon. Third, above the knee the two curves still diverge, but the epsilon=1e-16 curve rises with a ln-slope several times smaller.

Result 2: the mechanism, measured

Figure 2. N=50 (blue) and N=88 (red) at seed 0. Top: test accuracy. Middle: fraction of training samples whose softmax probability on the correct answer exceeds 1 - 1e-4. Bottom: root-mean-square per-coordinate gradient before the optimizer. At N=88 the training set saturates around step 1400, the gradient crosses Adam's epsilon around step 1700, and from that point the default-epsilon run (solid) falls behind the epsilon=1e-16 run (dashed). At N=50 the clamp keeps logits small, the softmax never saturates, and the gradient stays near 1e-4.

Logging the actual parameter change per step confirms that the gradient is not the thing that changes. In units of lr, the RMS per-coordinate step is 0.5 for both epsilon values up to step 1300. Once the gradient falls below epsilon, the default-epsilon step drops to 0.04 lr by step 2500 while the epsilon=1e-16 step stays at 0.5 to 0.9 lr. Test accuracy separates at the same point: 0.83 versus 0.98 at step 2000.

Result 3: controls

• Weight decay without a clamp. Same model, no clamp, weight decay 1: time to grok is 2450 steps with epsilon 1e-8 and 2350 to 2450 with 1e-16. The model with the tiny epsilon groks at a norm of 125 with logits around 1700, versus norm 76 and logits around 27 with the default. So under weight decay epsilon changes where the model ends up but not when it generalizes.
• Clamp with weight decay on. Truong et al. keep decoupled weight decay 1 on while projecting. Repeating N=88 that way gives 5000 steps with epsilon 1e-8 and 2100 with 1e-16, the same as without weight decay (5200 and 2100). Under a projection the decay is a uniform shrink that the projection undoes, so their clamp is in the regime described here.
• Unclamped, no weight decay. The small epsilon does not rescue grokking when nothing bounds the norm: after 8000 steps test accuracy is 0.013, the norm has grown from 80 to 242, and the largest logit is about 7000.
• Optimizers without epsilon. Normalized-gradient SGD (gradient divided by its own norm, momentum 0.9, lr tuned at N=65) takes a step of fixed length whatever the gradient scale. Its time to grok at N=88 is 1.35 to 2.0 times its N=65 time across three learning rates and two seeds, like Adam with epsilon 1e-16 (1.6) and unlike Adam with the default (3.9), and it groks at N=88 while the training softmax is fully saturated. Plain SGD, whose step shrinks with the gradient, is the fully damped limit: it groks at N=65 within 300 steps but reaches only 0.5 to 0.6 test accuracy at N=88 after 8000 steps, and rescaling its learning rate by (N/65)² to equalize the angular step does not help. So the knee is not a property of the sphere or the task; it is a property of how the optimizer handles a vanishing gradient.
• Float64. Omnigrok's constant-norm notebook trains in float64. At N=88 in float64 the default-epsilon run does not reach 0.99 within 8000 steps (0.978 at 8000; the per-coordinate gradient has fallen to about 5e-11), while epsilon 1e-16 groks at 2400. Removing the float32 rounding floor lets the gradient shrink further below epsilon, which makes the damping stronger. Truong 2026b report a float64 cell at rho 1.25 that groks at 25.4k steps in float32 but not within 60k in float64, which is the same pattern.
• A transformer under the clamp. The same one-layer transformer with Nanda-scale initialization, no weight decay, and all parameters projected to total norm 60 after every step (the Omnigrok constant-norm setup, above the norm of about 47 at which the free weight-decay run groks). After 16000 steps, with the training loss at zero in both runs, the default-epsilon run sits at 0.17 test accuracy and is barely moving, while the epsilon=1e-16 run passed 0.5 at step 12300 and is at 0.68 and climbing. Neither finished within the time budget, and this is one seed, but the direction and size match the MLP.
• A transformer with weight decay. A one-layer transformer for (a+b) mod 113 (Nanda et al. setup, 30% train, weight decay 1, no clamp) reaches 0.99 test accuracy at step 22950 with epsilon 1e-16 and 26800 with the default, one seed each (fig_transformer.png). A 14% difference under weight decay, against a factor of 2.6 to 4.5 under the clamp.

What this does and does not say

• The claim is about the time to grok under a norm clamp in a particular MLP with AdamW. The qualitative "small norm groks fast" statement survives; what changes is the shape of the curve and its interpretation.
• Epsilon is not the whole story above the knee. The epsilon=1e-16 curve still rises, and N=150 does not grok within 12000 steps. That remainder is the two factors identified in an earlier factorial test: logit scale, which the clamp fixes through the weight scale, and Adam's effective step on the sphere, which goes as lr/N. Dividing the logits by (N/50)^3 and scaling lr in proportion to N brings N=150 back to 1100 to 1150 steps, the same as N=50 (two seeds), while the temperature alone gives 3700. With those two matched the softmax never saturates and epsilon cannot act. So under the clamp the held norm works through three optimizer-side quantities and shows no residual effect of its own between N=50 and N=150.
• Truong et al. report that their exponential law turns sub-exponential above rho of about 1.25. The default-epsilon curve here does the same above N=100, where the per-coordinate gradient has already bottomed out at about 5e-10 and further norm increases cannot deepen the damping much.
• The dose-response and controls are one MLP, one task, one modulus, and three seeds per cell. The transformer checks are single seeds and the clamped one is unfinished.

Reproduction

All code is plain PyTorch on CPU. The training script is temp_clamp.py (flags --norm, --seed, --eps, --lr, --temp, --wd), the per-step update measurement is update_size.py, the SGD variants are in sgd_clamp.py, and dose_response.csv holds the per-seed grokking steps behind Figure 1. A single run takes about 35 seconds per 1000 steps on one CPU thread.

References

• Thilak, Littwin, Zhai, Saremi, Paiss, Susskind. The Slingshot Mechanism. arXiv:2206.04817, 2022.
• Prieto, Barsbey, Mediano, Birdal. Grokking at the Edge of Numerical Stability. arXiv:2501.04697, 2025.
• Liu, Michaud, Tegmark. Omnigrok. arXiv:2210.01117, 2022.
• Truong. arXiv:2606.13753 and arXiv:2606.18465, 2026.
• Wortsman et al. Small-scale proxies for large-scale Transformer training instabilities. arXiv:2309.14322, 2023.
• Everett et al. Scaling Exponents Across Parameterizations and Optimizers. arXiv:2407.05872, 2024.
• Nanda, Chan, Lieberum, Smith, Steinhardt. Progress measures for grokking via mechanistic interpretability. arXiv:2301.05217, 2023.