The Long Delay to Arithmetic Generalization: When Learned Representations Outrun Behavior

Laura Gomezjurado Gonzalez

cs.LGMar 30, 2026

The Long Delay to Arithmetic Generalization: When Learned Representations Outrun Behavior

Authors: Laura Gomezjurado Gonzalez

Organizations: Stanford University

Abstract

Grokking in transformers trained on algorithmic tasks is characterized by a long delay between training-set fit and abrupt generalization, but the source of that delay remains poorly understood. In encoder-decoder arithmetic models, we argue that this delay reflects limited access to already learned structure rather than failure to acquire that structure in the first place. We study one-step Collatz prediction and find that the encoder organizes parity and residue structure within the first few thousand training steps, while output accuracy remains near chance for tens of thousands more. Causal interventions support the decoder bottleneck hypothesis. Transplanting a trained encoder into a fresh model accelerates grokking by 2.75 times, while transplanting a trained decoder actively hurts. Freezing a converged encoder and retraining only the decoder eliminates the plateau entirely and yields 97.6% accuracy, compared to 86.1% for joint training. What makes the decoder's job harder or easier depends on numeral representation. Across 15 bases, those whose factorization aligns with the Collatz map's arithmetic (e.g., base 24) reach 99.8% accuracy, while binary fails completely because its representations collapse and never recover. The choice of base acts as an inductive bias that controls how much local digit structure the decoder can exploit, producing large differences in learnability from the same underlying task.

Figures & tables

Figure 1: The encoder becomes informative long before accuracy improves. (A) Sequence accuracy on one-step Collatz prediction in base 8 remains low until a late jump around step 44k, with even-branch examples learned earlier than odd-branch examples. (B) All six encoder-layer parity probes (

n\bmod 2

) exceed

99.7\%

accuracy from step 2k onward (left axis), while sequence-level output accuracy (dashed, right axis) rises only after the grokking transition. The persistent gap between the two is the shadow-knowledge gap . Parity is linearly decodable throughout the encoder long before it appears in the outputs.

Figure 2: Encoder transplant accelerates grokking, whereas decoder transplant does not. Lines show mean accuracy across

3

random seeds; shaded ribbons denote

\pm 1

standard deviation. Top: overall accuracy. A trained, frozen encoder paired with a fresh decoder ( enc. transplant ) reaches

70\%

accuracy

2.75\times

earlier than joint training from scratch and converges to a higher final accuracy. By contrast, a trained, frozen decoder paired with a fresh encoder ( dec. transplant ) declines over training. Bottom: the same asymmetry is largest on the odd branch (

3n{+}1

), where readout is hardest.

Figure 3: Modular structure becomes linearly decodable early in the encoder. (A) Probe accuracy at layer $5$ for targets modulo $2$ , $4$ , $8$ , and $16$ over training (base 8, log-scale steps). All probes approach ceiling by about step 2k, with coarser targets learned slightly earlier. (B) Probe accuracy at step 2k across encoder layers and modular targets. Accuracy is near ceiling across layers, though modulo $16$ is slightly weaker in earlier layers.

Figure 4: Rewinding only the decoder largely removes the grokking plateau. Curves show mean accuracy across

3

random seeds; shaded ribbons denote

\pm 1

standard deviation. (A) Starting from a converged model, we freeze the encoder, rewind the decoder to its step-2k weights, and resume training. The rewound model improves immediately, whereas training from scratch remains on a prolonged plateau before grokking. (B) The early-training zoom shows this gain on both even and odd branches in the rewound condition, while the scratch baseline stays near chance. Because the two conditions share the same converged encoder, this difference isolates decoder initialization.

Sequence accuracy (%)
Base $b$	All	Even	Odd	Base $b$	All	Even	Odd	Base $b$	All	Even	Odd	Base $b$	All	Even	Odd
2	0.0	0.0	0.0	8	97.3	99.7	94.9	16	93.6	99.8	87.3	32	96.3	99.9	92.7
3	45.8	33.4	58.3	9	92.1	89.1	95.2	18	99.6	99.7	99.5	36	97.1	99.6	94.7
4	99.2	99.8	98.7	10	99.0	99.9	98.1	24	99.8	99.9	99.6	48	93.7	97.4	90.0
6	99.7	99.4	100.0	12	99.4	99.9	99.0	27	92.4	98.8	85.9

Table 1: Numeral-base sweep. Accuracy varies with the numeral base, near ceiling for bases divisible by

2

and

3

, and zero for binary. Held-out exact sequence-level accuracy at step

500{,}000

, with

n_{\mathrm{train}}=1{,}000

integers sampled per step and evaluation on

5{,}000

unseen integers. “All” is overall exact-match accuracy on the full digit sequence of

T(n)

; “Even” and “Odd” restrict evaluation to even and odd inputs. The base-

2

row reports the final post-collapse checkpoint, after an earlier memorization phase.

Figure 5: Probe accuracy saturates before output accuracy moves, inside a pretrained

1.4

B decoder-only model, on two arithmetic tasks. Left: Pythia-

1.4

B fine-tuned on base-

8

Collatz. Linear probes for

T(n)\bmod 8

and

n\bmod 16

read

1.00

at step

0

, the pretrained checkpoint before any task gradient, while sequence-level exact-match is exactly

0.000

; the gap closes over the first

\sim 500

fine-tuning steps. Right: the same model fine-tuned on

4

-digit addition. The units-digit carry probe holds in the

0.78

–

0.87

band from step

0

, while exact-match stays below

0.14

for

\sim 2{,}000

steps. The probed feature is task-specific (residues for Collatz, carry for addition); the gap is task-general.

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: The same integer can induce different decoder readout problems across bases. The integer

n=80

is shown in bases 8, 10, and 2. While the least-significant digit stays fixed at the right edge, the number of preceding tokens varies with the base. Numeral base therefore changes both sequence length and access to low-order arithmetic cues.

Condition	Seed	Overall acc.	$95\%$ CI	Notes
Encoder transplant	primary	(main text)		trajectory in Fig. 2
Encoder transplant	alt. 1	54.4%	[53.0, 55.7]	data seed 10
Encoder transplant	alt. 2	93.6%	[92.9, 94.2]	data seed 11
Scratch baseline	primary	(main text)
Scratch baseline	alt. 1	94.4%	[93.7, 95.0]	data seed 10
Scratch baseline	alt. 2	96.4%	[95.9, 97.0]	data seed 11

Appendix

Table 2: Per-seed endpoints for causal replications. Overall sequence-level accuracy on the same

5{,}000

-example held-out set with exact

95\%

Clopper–Pearson intervals (binomial). “Primary” denotes the main-text run; alternate rows use training data seeds 10 and 11. All alternate rows are evaluated at step 200,000 (same training budget as the primary runs).

Seed	Transplant (src= $24$ , tgt= $2$ )	Scratch (tgt= $2$ )
$0$	$0.807$	$0.477$
$7$	$0.766$	$0.388$
$13$	$0.699$	$0.306$
$17$	$0.743$	$0.381$
$23$	$0.721$	$0.520$
$29$	$0.770$	$0.117$

Appendix

Table 3: Six-seed cross-base transplant. Final exact-match accuracy at step

300{,}000

for encoder transplant (source base

24

, target base

2

) versus the matched scratch baseline. The transplant-to-scratch ratio is

2.06

overall and

4.4

on the odd branch, with non-overlapping 95% intervals on both metrics.

Figure 7: Six-seed encoder transplant on the largest base-gap configuration (base-24 source, base-2 target). (A) Per-seed final exact-match accuracy at step 300,000. The transplant arm exceeds the scratch baseline on every seed; its mean (

0.751

) is

2.06\times

the scratch mean (

0.365

), with non-overlapping 95% intervals. (B) Aggregated mean accuracy with

\pm 1

standard deviation. The separation is largest on the odd branch (ratio

\sim 4.4

), where the scratch base-2 odd-branch mean is

0.119

. Transplant variance is

\sim 4\times

smaller than scratch variance.

Figure 8: Per-seed training trajectories for the three transplant conditions. Each panel shows the three individual training runs (solid, dashed, dotted lines) with their mean (thicker solid line), for overall sequence-level accuracy over 200k steps. Left (encoder transplant). All three seeds show rapid early progress with tight agreement, so the speedup in Fig. 2 is reproducible across seeds rather than a single-run artifact. Center (scratch baseline). Seeds agree on final accuracy but differ in the step at which the grokking transition occurs, producing the wider ribbon in the main figure. Right (decoder transplant). All three seeds decline or stall throughout training, so the asymmetry between encoder and decoder transplant holds across seeds.

Figure 9: Per-seed training trajectories for decoder rewind and scratch baseline. Each panel shows three individual seeds with their mean, for overall accuracy over 200k steps. Left (decoder rewind). Two of three seeds improve rapidly from early training with no visible plateau; seed 2 (data seed 10) is slower but still improves monotonically. Right (scratch baseline). Seeds agree closely on the grokking transition window and final accuracy, consistent with the tighter ribbon in Fig. 4 . The contrast between rewind and scratch is reproducible across seeds, even where individual seed endpoints vary.

Figure 10: The encoder resolves residue structure coarse-to-fine. Conditional probe accuracy measures how much information the final encoder layer carries about each level of the 2-adic filtration, given that all coarser levels are known. (A) Probes

k_{2}

(mod 4

\mid

mod 2) and

k_{3}

(mod 8

\mid

mod 4) saturate above 0.99 from step 2k, while

k_{4}

(mod 16

\mid

mod 8) starts at 0.676 and rises slowly, peaking near 0.96 around step 130k. The ordering is sequential coarse-to-fine rather than all-at-once. (B) The

k_{4}

trajectory has the same shape across two seeds, both showing the slow rise and late peak, so the pattern is reproducible rather than an artifact of a single run.

Base $b$	Exact-match	Digit-level (micro)	Exact @ len=4	Exact @ len=5
2	0.00	0.51	—	—
4	0.98	0.996	0.95	0.98
8	0.93	0.98	0.99	0.95
10	0.96	0.99	0.99	0.99
16	0.91	0.97	0.91	—
24	0.996	0.999	1.00	—

Appendix

Table 4: Length-normalized base sweep. Exact sequence-level accuracy, digit-level micro accuracy, and exact-match restricted to fixed output-length buckets at the converged checkpoint. The substantive ordering (composite-with-2

>

prime-power-of-2

>

odd-composite

>

binary) is preserved under the length-clean metric; bins with too few held-out integers are marked “—”.

Seed	Run name	Peak	Layer	Step	Head	L1 plateau peak	L1 @ converg.
42	factor_sweep_base8	$\mathbf{0.89}$	1	66,000	$h_{1}$	0.89	0.19
0	grok_base8_seed0	$\mathbf{0.95}$	5	60,000	$h_{3}$	0.64	—
1	grok_base8_seed1	0.47	4	70,000	$h_{0}$	0.45	0.44
2	grok_base8_seed2	$\mathbf{0.97}$	3	70,000	$h_{6}$	0.45	0.22

Appendix

Table 5: Decoder cross-attention head specialization during the plateau, across four base-8 seeds. Top-head share is the fraction of mean cross-attention to the last input-position digit captured by the single dominant head (out of 8), conditioned on odd-parity input. Plateau region

=

steps in

[25\mathrm{k},80\mathrm{k}]

. The convergence column reports the layer-1 share at the latest available checkpoint (500,000 for three seeds, 342,000 for seed 0).

Figure 11: Decoder depth and carry exposure shape decoder learning. (A) Odd-branch accuracy over 500k steps for decoder depths 1, 2, 4, and 6, with the encoder fixed at 6 layers. Performance is non-monotonic: depth 4 performs best at convergence, depth 1 learns fastest early and finishes close behind, and depths 2 and 6 lag throughout. (B) Even (dashed) and odd (solid) accuracy under three training distributions. Carry-stratified training preserves odd accuracy but reduces even accuracy to 26.9%, while short-carry training preserves even accuracy but caps odd accuracy at 38.3%, showing that deep-carry examples are necessary for odd-branch generalization.

	$k=3$ (standard)		$k=-1$
Base $b$	$H^{\mathrm{even}}$	$H^{\mathrm{odd}}$	$H^{\mathrm{even}}$	$H^{\mathrm{odd}}$
2	1.00	0.00	1.00	0.00
3	0.00	0.00	0.00	0.00
4	1.00	0.00	1.00	0.00
6	1.00	0.00	1.00	0.00
8	1.00	0.00	1.00	0.00

Appendix

Table 6: Local predictability is identical across Collatz family members. Branchwise conditional entropy

H^{(r)}

for the standard (

k=3

) and

k=-1

maps. The two columns are numerically identical: odd-branch entropy is zero for every base, even-branch entropy is near-maximal for even-factor bases and zero for purely-odd bases (3, 9, 27). This is not a coincidence. It is a consequence of

H^{(r)}

depending only on the even branch (

n/2

, independent of

k

) and on whether the odd branch resolves within a two-digit window (which it does for any fixed

k

in any base).

Base $b$	$k=3$ acc. (%)	$k=-1$ acc. (%)
4	99.2	82.0
8	97.3	94.0
10	99.0	95.1
16	93.6	86.4
18	99.6	99.2
27	92.4	91.3

Appendix

Table 7: Base-dependent learnability ordering is preserved for

k=-1

. Final sequence-level accuracy for the standard (

k=3

, step 500,000) and

k=-1

Collatz maps. The

k=-1

values for bases 6, 9, 12, and 24 are taken at their latest checkpoint (

\sim

210k–218k steps); all other

k=-1

runs are at step 500,000. The relative ordering is largely preserved: highly-composite even-factor bases (18, 32, 36, 48) remain among the easiest, and the purely-odd base 27 remains harder than most even-factor bases in both conditions.

$d_{\mathrm{model}}$ \ depth	$2$ ( $1{+}1$ )	$4$ ( $2{+}2$ )	$6$ ( $3{+}3$ )
$128$	$0.968$	$0.858$	$0.863$
$256$	$0.959$	$0.853$	$0.918$
$512$	$0.938$	$0.936$	$0.967$

Appendix

Table 8: Architecture sweep, final exact-match. Encoder–decoder transformers at three widths and three total depths, base-

8

Collatz,

200{,}000

training steps. The shadow-knowledge gap appears in every cell; per-cell probe and output trajectories are released with the supplementary code.

Figure 12: Decoder-only baseline at base 8, step 200k. Linear probes recover residue and last-digit features at

\geq 0.95

across layers

0

–

8

, then collapse in the final layers, while sequence-level output exact-match is

0.000

. The shadow-knowledge gap is wider than in the encoder–decoder setting.

Figure 13: Probe-vs-output dissociation inside Pythia-

160

M (base-

8

Collatz). (A) Best-layer linear probes for

T(n)\bmod 8

and

n\bmod 16

are saturated (

\geq 0.98

) from step

0

, while output exact-match rises from

0.000

at step

0

\geq 0.97

by step

10{,}000

. (B) The

T(n)\bmod 8

probe resolved by layer: every layer from

3

12

carries the feature at

\geq 0.96

, while the embedding layer (

0

) sits at

0.755

; output is overlaid for reference. This is the shadow-knowledge gap of Fig. 1 reproduced in a pretrained

160

M decoder-only LM, at the

\sim\!2{,}000

-step checkpoint resolution of this run.

Figure 14: Probe-vs-output dissociation inside Pythia-

1.4

B (base-

8

Collatz). (A) Best-layer linear probes for

T(n)\bmod 8

and

n\bmod 16

are saturated at

1.00

from step

0

(the pretrained checkpoint, before any task gradient), while output exact-match climbs from

0.000

at step

0

\geq 0.97

by step

\sim\!500

, with the dense early-step schedule resolving the gap-closing into a smooth curve. (B) The

T(n)\bmod 8

probe resolved by layer: every layer from

4

24

carries the feature at

\geq 0.94

throughout, while the embedding layer (

0

) sits at

0.755

; output is overlaid for reference. The shadow-knowledge gap of Fig. 1 , reproduced inside a pretrained

1.4

B-parameter decoder-only LM at fine-grained early-step resolution.

Figure 15: Causal probe-direction erasure inside Pythia-

1.4

B at step

200

of the Collatz fine-tune. Probe-direction erasures collapse accuracy at layer

4

(

\Delta\geq+0.67

for

T(n)\bmod 8

, last_digit_b

8

, and

n\bmod 16

), while same-rank random subspaces cost essentially nothing. The dependence concentrates at exactly the layer where the pretrained probes saturate, and disperses with depth.

Figure 16: Erasing the encoder’s parity direction selectively impairs performance during the plateau. (A) At each checkpoint, we remove the learned parity direction from encoder hidden states at inference time and compare performance with the unmodified model. The erased model underperforms the baseline most strongly before grokking, even though parity probe accuracy is already near ceiling. (B) The effect of parity erasure, measured as baseline minus erased accuracy, peaks during the plateau and drops after grokking.

Erased feature	Subspace rank	Probe acc. before erasure	Accuracy drop after erasure
mod $2$	$1$	$1.000$	$0.000$
mod $4$	$3$	$0.997$	$0.000$
mod $8$	$7$	$0.999$	$0.000$
mod $16$	$15$	$0.999$	$0.001$
last digit	$7$	$0.999$	$0.000$
magnitude	$7$	$0.987$	$\mathbf{0.141}$

Appendix

Table 9: Full erasure hierarchy at convergence (step

492

k, base

8

). Every low-order residue subspace is non-load-bearing at convergence; the magnitude subspace is the single load-bearing axis. Random-direction controls of equal rank cost

\leq 0.005

in every case.

Figure 17: Accuracy drop after projecting out the probe subspace for each encoder-side feature, resolved across training. The five residue features (mod 2/4/8/16, last digit) peak during the plateau and decay to

\Delta\leq 0.001

by convergence; the shaded band traces their min–max envelope. The magnitude erasure rises monotonically and dominates at convergence (

\Delta=0.141

). The crossover at

\sim\!5{\times}10^{4}

steps marks in time the transition from the simple-cue readout to the distributed readout.

Grokking

Grokking is delayed generalization in which neural networks memorize training data before abruptly learning a rule that works on unseen examples.

See all 77 papers

Early-Warning Signals of Grokking via Loss-Landscape Geometry
Feb 19, 2026Yongzhong XuGrokking Neural Network Generalization
Structure-Specific Representational Priors Causally Control the Grokking Delay
Jul 5, 2026Gunner Levi HoweTransformer Grokking
Grokking Is Conditional and Fragile: A Fully-Tractable, Multi-Seed Study at 12K Parameters
Jul 6, 2026Yoshiyuki OotaniPhase Transitions ML Reproducibility
Independent Researcher

Feb 19, 2026J/K move · Enter open · S save

Early-Warning Signals of Grokking via Loss-Landscape Geometry

Yongzhong Xu

Grokking Neural Network Generalization Loss Landscape Geometry

Open paper AskPDF arXiv

Grokking -- the abrupt transition from memorization to generalization after prolonged training -- has been linked to confinement on low-dimensional execution manifolds in modular arithmetic. Whether this mechanism extends beyond arithmetic remains open. We study two sequence-learning benchmarks: SCAN compositional generalization and Dyck-1 depth prediction. Across both tasks and a wide range of learning rates, the commutator defect -- a curvature measure derived from non-commuting gradient updates -- rises well before generalization, with lead times following a superlinear power law (alpha approximately 1.18 for SCAN, approximately 1.13 for Dyck), consistent with prior results on modular arithmetic. Weight-space PCA reveals that spectral concentration is not a universal precursor; the commutator defect is. Causal interventions demonstrate a mechanistic role: amplifying non-commutativity accelerates grokking (roughly 32% on SCAN, roughly 50% on Dyck), while suppressing orthogonal gradient flow delays or prevents it. The three task families form a spectrum of causal sensitivity -- modular arithmetic is rigid, Dyck is responsive, SCAN is intermediate -- yet suppression delays or prevents grokking in all cases, establishing necessity as a universal finding. These results identify the commutator defect as a robust, architecture-agnostic, causally implicated early-warning signal for delayed generalization in transformers.

Abstract

Grokking

Explore similar work

Early-Warning Signals of Grokking via Loss-Landscape Geometry

Structure-Specific Representational Priors Causally Control the Grokking Delay

Grokking Is Conditional and Fragile: A Fully-Tractable, Multi-Seed Study at 12K Parameters