PaperlayerView exact revision on arXiv

Source-linked AI summary

Beyond Embedding Transfer: Component Roles in Grokking Transfer and Stability

Zeyu Jia

arXiv:2609.18078v1cs.LG

TL;DR

The paper asks which transferred components accelerate algorithmic generalization and whether those gains persist under continued optimization. It evaluates cross-operator transfer and targeted carrier protection, finding distinct component roles in acquisition and stability, with protection bounded by architecture and task.

  • Problem

    Prior embedding-transfer work leaves open whether internal attention/MLP weights add computational utility and how warm-started models evolve under continued optimization.

    Source: §1 · pp. 1–2

  • Method

    The study separates early efficacy from post-reach stability while testing component transfer and representation-carrier protection across modular-arithmetic and S5 tasks.

    Source: §2.1 · p. 2, §1 · p. 2, §2.3 · p. 4

  • Results

    Internal-block transfer adds acquisition benefits beyond E + U, while freezing E + U suppresses post-reach relapse and online shielding reduces drawdown in modular tasks.

    Source: §1 · p. 2, §3.1 · p. 5, §5.2 · p. 7

  • Takeaways & Limitations

    Acquisition acceleration and retention stability have distinct component roles, and the effectiveness of shielding has empirical architectural and task boundaries.

    Source: §1 · p. 2, §8 · p. 10

    Shielding provides incomplete protection in two-layer models and no confirmed benefit in non-abelian S5.

    Source: §8 · p. 10

Abstract

from arXiv · show

Warm-start transfer can make algorithmic tasks generalize rapidly, yet it is unclear which model components provide the gain and whether that gain remains stable under continued optimization. We study cross-operator transfer on modular arithmetic and separate efficacy (early velocity) from stability (post-reach drawdown). In a scale-matched 108-run battery across 12 seed blocks (96-run 2^3 factorial plus 12-run scale control), transferring internal attention/MLP weights (B) alongside token embeddings and readout (E+U) improves early accuracy by 5.46 pp (Holm p=0.0039) and cuts confirmation latency by 558 steps (Holm p=0.0088). While readout plus internal-block transfer satisfies the pre-specified +/-500-step latency equivalence criterion in 1-layer models (TOST p=0.0011, though Full is faster in 11/12 paired seeds), a prospective 2-layer replication confirms the internal-block advantage (12/12 seeds, +704.67 integral units, p=4.88x10^-4) while revealing an architectural boundary: omitting donor embeddings falls 4475.6 units below Full, outside the +/-250-unit margin. Continued target training frequently triggers severe post-grokking relapse. Freezing transferred representation carriers (E, U) nearly eliminates offline relapse (19.40% -> 0.07%, Holm p=0.005859). Online validation-triggered gating slashes True Max Drawdown from 22.06% to 0.60% on 2a+b (p=0.000488), with prospective confirmations extending protection across affine, nonlinear quadratic, and 2-layer targets (10.94-23.47 pp reductions), distinguishing continual stabilization from static early stopping. In non-abelian S_5, unshielded transfer surges transiently (95.4% peak), but a prospective shielding cohort yields no confirmed benefit (+0.15 +/- 1.14 pp). These results establish a component-level dissociation between transfer acceleration and trajectory stability, and expose the empirical boundaries of parameter shielding.

1 INTRODUCTION

Grokking delays held-out generalization despite rapid memorization, and prior embedding-transfer work leaves open whether internal weights add computational utility and whether gains persist during continued training. The paper separates acquisition speed from trajectory stability and identifies distinct component roles across these questions.

1 INTRODUCTION

  • Prior work links delayed modular-arithmetic generalization to gradual formation of structured Fourier multiplication circuits and circular representation manifolds in embeddings and attention weights.The passage cites Nanda et al. (2023), Gromov (2023), and Varma et al. (2023).Source: §1 · p. 1
  • Earlier embedding-transfer results accelerate algorithmic generalization, but leave internal attention/MLP weight utility and post-transfer evolution under continued optimization open.The cited prior study is Xu et al. (2025).Source: §1 · p. 1
  • Full transfer reaches 90% accuracy by step 50 versus 250–500 steps for representation conditions, while freezing E + U reduces post-peak relapse from 19.40% to 0.07%.The caption reports Holm p = 0.005859 for the relapse reduction.Source: Figure 1 · p. 2
  • The study asks whether internal attention/MLP matrices add computational acceleration beyond representation carriers, why warm-started models relapse, and whether parameter-level protection can suppress relapse without degrading early velocity or terminal accuracy.It also asks where such protections break down.Source: §1 · p. 2
  • Internal attention/MLP weights (B) contribute to acquisition velocity, while shielding representation carriers controls retention stability; acquisition’s carrier requirements depend on architecture depth.The contributions report +5.46 pp early accuracy and +704.67 units in a two-layer replication, while omitting donor embeddings in two layers produces a −4475.6-unit deficit.Source: §1 · p. 2

2 PROBLEM FORMULATION & DUAL-TRACK PROTOCOL

The protocol studies cross-operator transfer in modular arithmetic and S5, separating early learning efficacy from sustained confirmation and post-reach stability. It scale-matches transferred tensor directions and specifies the metrics and inference procedures used to evaluate these outcomes.

2.1 TASK DEFINITIONS & ARCHITECTURE

  • The tasks pair modular addition with an affine modular target and S5 multiplication with conjugation, using pre-LN Transformers that transfer embeddings (E), readout (U), and/or internal attention/MLP weights (B).Full transfer moves E, U, and B; positional embeddings, LayerNorm parameters, and MLP biases retain cold-initialized values.Source: §2.1 · pp. 2–3

2.2 PER-TENSOR SCALE MATCHING (DECOUPLING DIRECTION FROM MAGNITUDE)

  • Transferred tensors retain donor direction but are rescaled to the native cold initialization’s Frobenius norm, isolating direction from overall tensor magnitude.Donor checkpoints qualify after validation accuracy remains at least 90% for a continuous 1,000-step persistence window.Source: §2.2 · p. 3

2.3 DUAL-TRACK EVALUATION METRIC PROTOCOL

  • Efficacy is measured by first validation reach at 90%, pre-5k test-accuracy integral, and full-trajectory mean test accuracy; test accuracy is for reporting, not triggering or model selection.Validation accuracy and held-out test accuracy are tracked separately.Source: §2.3 · p. 3
  • Stability uses 21 consecutive validation evaluations at or above 90% to confirm grokking, with confirmation latency truncated at 30,000 steps and post-peak test-accuracy dip as the primary offline relapse metric.Online interventions use validation-triggered branches and evaluate post-trigger stability over a specified finite horizon.Source: §2.3 · p. 3
  • Online triggering requires two consecutive validation evaluations above threshold at a 50-step cadence, using q = 0.95 for modular arithmetic and q = 0.70 for S5.Source: §2.3 · p. 3
  • The protocol reports full transfer’s +5.46 pp early-accuracy advantage over E + U and evaluates post-trigger stability using True Max Drawdown.The captioned protocol specifies post-trigger horizons of 1,500 steps for modular cohorts and 3,000 for S5.Source: §2.3 · p. 4

2.4 STATISTICAL INFERENCE AND FAMILY-WISE ERROR RATE CONTROL

  • Paired contrasts use exact sign-flip permutation tests with Holm-Bonferroni correction; U + B versus Full latency equivalence uses paired TOST against a prespecified ±500-step margin.The passage also specifies BCa bootstrap confidence intervals and separate inference families for efficacy and stability.Source: §2.4 · p. 4

3 WHAT INFORMATION IS TRANSFERRED? CARRIER LOCALIZATION AND SYNERGY

A scale-matched factorial study finds that internal-block transfer adds acquisition benefits beyond E + U, while isolated components show asymmetries and U + B meets the prespecified mean-latency equivalence criterion to Full in one-layer models.

3.1 INTERNAL WEIGHT MATRICES PROVIDE SIGNIFICANT INCREMENTAL ACCELERATION

  • Full transfer increases pre-5k integral by +273.0 units, a +5.46 percentage-point average test-accuracy gain, and reduces confirmation latency by +558.3 steps versus E + U.Both incremental advantages are significant under family-wise error control.Source: Table 1 · p. 5, §3.1 · p. 5

3.3 BOUNDED LATENCY EQUIVALENCE OF U + B VS. FULL

  • U + B meets the prespecified ±500-step average-latency equivalence criterion to Full (mean difference −16.7 steps; TOST p = 0.001089), although Full is faster in 11 of 12 paired seeds.The conclusion is operational average-latency equivalence, not computational parity or identical per-seed trajectories.Source: Table 1 · p. 5, §3.3 · p. 6

3.2 FACTORIAL SYNERGY: READOUT POWER AND INTERNAL INDEPENDENCE

  • Readout-only transfer confirms in 2,600 steps versus 4,220.8 for embedding-only transfer, while B alone has 0/12 confirmations within 30,000 steps and does not significantly improve on cold start.For B versus cold, the exact sign-flip test gives p = 0.5000; 10 of 12 blocks are mutually censored.Source: §3.2 · p. 5
  • Normalized transfer of learned embedding magnitudes severely hinders downstream optimization, validating the scale-matching protocol; separately, a negative descriptive three-way interaction (−3,141.9 integral units, p = 0.0005) is interpreted as ceiling-compressed diminishing returns.Source: §3.2 · p. 6
  • The descriptive three-way interaction is −3,141.9 integral units (p = 0.0005), interpreted as ceiling-compressed diminishing returns because I5k is bounded above by 5,000.The passage does not interpret this interaction as a distinct mechanism.Source: §3.2 · p. 6

4 WHY DOES GENERALIZATION RELAPSE? BOUNDED MECHANISTIC DIAGNOSTICS

The paper probes relapse through state replacement, freezing interventions, displacement analysis, and linear readout refitting, while emphasizing that these diagnostics do not establish a universal causal mechanism.

4 WHY DOES GENERALIZATION RELAPSE? BOUNDED MECHANISTIC DIAGNOSTICS

  • On one 25.24-point relapse checkpoint pair, restoring pre-relapse B raises trough accuracy from 74.76% to 99.85%, whereas restoring E + U reaches 81.45%.The authors interpret this checkpoint-specific state swap as locating the functional deficit primarily in internal blocks.Source: §4 · p. 6
  • Freezing either E + U or B prevents the impending 50-step dip on the probed trajectory; at least 88.4% of internal-block displacement is tangential, and only its tangential component reproduces low accuracy.These interventions implicate joint co-adaptation and tangential internal-block motion on this trajectory.Source: §4 · p. 6
  • Across four selected relapse events, linear readout refitting improves trough accuracy by 8.37–11.03 points, but trough probes remain less accurate than pre-relapse probes at every tested calibration budget.The authors conclude that trough representations retain recoverable target information but are harder to decode linearly; cross-initialization differences remain unresolved.Source: §4 · p. 6
  • The prospective gating table compares online carrier freezing with unshielded AdamW and trigger-time early stopping across one-layer, affine, quadratic, and two-layer cohorts.Trigger Snapshot has zero post-trigger drawdown by construction because optimization stops at the trigger.Source: Table 2 · p. 7

5 TARGETED PROTECTION: OFFLINE SHIELDING AND ONLINE DYNAMIC GATING

Freezing representation carriers suppresses post-reach relapse, while online validation-triggered freezing stabilizes continued optimization. Protection extends across affine, quadratic, and two-layer targets, but its residual drawdown and task-specific boundary remain important.

5.1 INTERVENTION FORMULATION

  • Carrier freezing bundles gradient suppression with stopping weight-decay erosion, while positional embeddings, LayerNorms, and internal blocks remain trainable.The tested alternatives are soft carrier learning-rate reduction and temporary warmup freezing.Source: §5.1 · p. 7

5.2 CONFIRMATORY BATTERY RESULTS (N = 10 INDEPENDENT BLOCKS 230–239, 60 RUNS)

  • 19.40% to 0.07% post-peak relapse dip: permanently freezing E+U nearly eliminates offline relapse.The reduction is +19.33 percentage points (Holm p = 0.005859), with positive contrasts in all 10 independent blocks.Source: §5.2 · p. 7
  • Neither 0.1× carrier learning rate nor 500-step temporary freezing significantly suppresses relapse under the tested schedules.Mean dips were 20.80% and 14.12%, respectively; delayed relapse remained within the 10,000-step horizon.Source: §5.2 · p. 7

5.3 RESOLVING THE 4-QUADRANT LIFECYCLE: TAKEOFF AND STABILITY SYNERGIES

  • Full transfer plus carrier freezing preserves 50-step takeoff while reducing relapse to 0.02% in an exploratory 10-run extension.The confirmatory findings distinguish internal-parameter transfer's early takeoff benefit from E+U freezing's stability benefit.Source: §5.3 · p. 8

5.4 ONLINE CAUSAL PROTECTION VIA DYNAMIC VALIDATION GATING

  • 22.06% to 0.60% True Max Drawdown: online carrier freezing sharply reduces post-trigger collapse on 2a+b.The reduction is +21.45 percentage points across 12/12 concordant pairs; the intervention freezes WE and WU after validation accuracy reaches at least 95% for two consecutive evaluations.Source: §5.4 · p. 8

5.5 PROSPECTIVE CONFIRMATION ACROSS TASKS AND ARCHITECTURE DEPTHS

  • Online carrier freezing reduces True Max Drawdown by 23.47 percentage points on the affine target, but the cohort mean remains 5.09%.The prospectively tested ≤2.0% cohort-mean drawdown boundary was not met.Source: §5.5 · p. 8
  • Online carrier freezing reduces True Max Drawdown by 10.94 percentage points on the nonlinear quadratic target, with 12/12 positive wins.The caption reports minimal terminal cost for this confirmation.Source: §5.5 · p. 8, Figure 3 · p. 9
  • In two-layer targets, carrier shielding cuts True Max Drawdown from 80.58% to 65.23%, leaving substantial residual drawdown.The reduction is 15.35 percentage points across 11/12 concordant wins.Source: §5.5 · p. 9

6 GENERALITY BOUNDARIES: NON-ABELIAN GROUPS AND LAYER DEPTH

The tests identify boundaries for transfer and shielding beyond the one-layer affine setting. Internal blocks accelerate two-layer acquisition, but donor embeddings remain necessary; prospective shielding shows no confirmed benefit in non-abelian S5.

6.1 NON-ABELIAN CONJUGATION: TRANSIENT TRANSFER AND EMPIRICAL BOUNDARY

  • In S5 conjugation transfer, accuracy briefly peaks at 95.39% and 84.62% versus 5.83% for cold start, but none of four runs persistently confirms.The prospective cohort triggered in 9/12 blocks; shielding showed no confirmed drawdown benefit on triggered pairs (∆D = +0.15 ± 1.14 pp, p = 0.894531).Source: §6.1 · p. 9

6.2 TWO-LAYER COMPONENT ROLES: ACQUISITION AND EMBEDDINGS

  • The two-layer component test used 12 fresh blocks and compared cold, Full, E+U, and U+B target conditions.The donors were two-layer Transformers, and transfer used cold Frobenius scale matching.Source: §6.2 · pp. 9–10
  • Internal blocks add +704.67 I5k units beyond E+U in two-layer models, with 12/12 wins.The compared I5k values are 4600.9 ± 31.3 for internal blocks plus carriers and 3896.3 ± 73.4 for E+U.Source: §6.2 · p. 10
  • U+B without donor embeddings falls 4475.63 I5k units below Full, outside the ±250-unit equivalence margin.The condition was untriggered in 10/12 blocks.Source: §6.2 · p. 10

7 RELATED WORK

The related work situates the study among accounts of grokking circuits, embedding transfer, and parameter freezing, while distinguishing its cross-operator analysis of acquisition and continued-optimization stability.

7 RELATED WORK

  • Unlike single-task embedding-transfer and circuit-unlearning studies, this work separates internal computation from carriers across operators and examines dynamic gating during continued optimization.The passage relates grokking to trigonometric circuits and lazy-to-rich transitions, and notes that scale matching isolates directional transfer.Source: §7 · p. 10

8 DISCUSSION, LIMITATIONS, AND CONCLUSION

The conclusion presents component-level differences between acquisition and continual stability, while noting that their generality beyond the tested architectures and domains remains unresolved.

8 DISCUSSION, LIMITATIONS, AND CONCLUSION

  • Whether component specialization generalizes to deeper architectures and non-algorithmic domains remains open.The conclusion also reports that online carrier freezing stabilizes continued optimization beyond the trigger, unlike static early stopping.Source: §8 · p. 10

A DIAGNOSTIC LINEAR PROBING SPECIFICATION, OPTIMIZATION CONVERGENCE CHECKS, AND FULL ACCOUNTING

The appendix specifies an exploratory linear-probing investigation and defines the embedding-subspace alignment measure used in Figure 4(b).

A DIAGNOSTIC LINEAR PROBING SPECIFICATION, OPTIMIZATION CONVERGENCE CHECKS, AND FULL ACCOUNTING

  • The appendix supplies the technical specification, convergence checks, trajectory-event accounting, and numerical results for the diagnostic linear-probing investigation.The alignment measure is described in the diagnostic-figure subsection.Source: A DIAGNOSTIC LINEAR PROBING… · p. 12, §A.1 · p. 12

A.1 DIAGNOSTIC FIGURES

  • Alignment in Figure 4(b) is the mean canonical principal-angle cosine between the leading k = 12 left singular subspaces of current and donor token embeddings.The subspaces are spanned by the top-k left singular vectors; the metric averages singular values of their inner projection matrix.Source: §A.1 · p. 12

A.2 PROTOCOL SPECIFICATION

This appendix specifies a test-conditioned probe protocol for examining relapse trajectories and representation decodability. Its development cohort includes deep relapses in all three broken scratch runs.

A.2 PROTOCOL SPECIFICATION

  • The exploratory cohort comprises 12 trajectories across scratch, transfer EU, and transfer full initialization conditions.Transfer EU initializes token embeddings and readout from donors while internal blocks and LayerNorms remain cold; transfer full also transfers internal attention and MLP weights.Source: §A.2 · p. 12
  • Probe representations are extracted at the final sequence position after final LayerNorm and before the readout, with logits checked against the model’s forward pass.The check ensures the probed representations match the operational readout input.Source: §A.2 · p. 12
  • Relapse events are tracked from the running peak after validation accuracy reaches 95% for two consecutive checks, with troughs requiring a dip of at least 10 pp and accuracy below 90%.Recovery is the first subsequent evaluation at or above 95%.Source: §A.2 · p. 12
  • Figure 4’s single checkpoint-pair swap recovers 99.85% accuracy by restoring internal blocks, versus 81.45% by restoring E + U.These exploratory observations localize the contemporaneous functional deficit but do not identify a general cause.Source: Figure 4 · p. 13
  • Trough probes improve with larger calibration budgets but remain below paired pre-relapse probes in the four test-conditioned relapse events.The finite-budget deficit neither establishes irreversible information loss nor excludes a better decoder.Source: Figure 5 · p. 13
  • Probe fitting uses nested calibration subsets of 256, 512, 1,024, or 2,554 training examples, with penalties selected by validation cross-entropy.Because event times are selected using test accuracy, the analysis is test-conditioned rather than globally blind.Source: §A.2 · p. 13

A.3 OPTIMIZATION CONVERGENCE CHECKS

  • Of 288 candidate probe fits, 286 pass the gradient-norm gate, and all 48 validation-selected probes pass.The two exceptions were not selected; both occurred at λ = 1 and N = 2554.Source: §A.3 · p. 13

A.4 FULL TRAJECTORY EVENT ACCOUNTING

  • 18.37, 13.86, and 29.95 pp were the relapse dips in the three broken scratch trajectories.All three had qualified relapses, challenging the assumption that scratch training converges smoothly.Source: §A.4 · p. 13

A.5 FULL LINEAR PROBE DECODABILITY AND PAIRED DEFICITS

This section reports trajectory-event accounting and probe decodability across checkpoint stages and calibration budgets, while defining the paired trough deficit and re-fit gain.

A.5 FULL LINEAR PROBE DECODABILITY AND PAIRED DEFICITS

  • All three broken scratch runs in the 12-run cohort suffered deep qualified relapses, with dips of at least 10 pp and test accuracy below 90%.The table’s event accounting covers breakthrough, pre-relapse peak, trough, and recovery.Source: Table 3 · p. 14
  • Table 4 reports probe test accuracies across four test-conditioned relapse events and calibration budgets of 256, 512, 1,024, and 2,554 examples.The penalty is selected by validation loss; the supplied passage does not provide the table’s cell values.Source: §A.5 · p. 13, Table 4 · p. 14
  • The paired trough deficit is trough-probe accuracy minus pre-relapse-probe accuracy, while re-fit gain compares the full-sample trough probe with the unadapted trough model.The parameter census assigns E to token embeddings, U to the readout, and B to internal blocks.Source: Table 4 · p. 14, Table 5 · p. 14

B MODEL ARCHITECTURE, PARAMETER KEYS, AND INITIALIZATION DETAILS

The appendix inventories the GrokTransformer parameters and specifies which tensors transfer, including the scale-matching convention and components that remain cold.

B.1 ARCHITECTURAL CENSUS AND TENSOR INVENTORY

  • The 1-layer pre-LayerNorm Transformer has dmodel = 128, four attention heads, a 512-unit MLP, and 227,712 parameters.The MLP has two layers and GELU activations.Source: §B.1 · p. 15
  • The 2-layer model contains 425,472 parameters, including 422,400 transferable 2D weights, mapped layer-by-layer when internal blocks transfer.Its two identical blocks each contain 197,760 parameters.Source: §B.1 · p. 15

B.2 EXACT MODULE CONVENTIONS AND TRANSFER RULES

  • Internal-block transfer includes fused attention QKV and output-projection matrices, while LayerNorm parameters, MLP biases, and positional embeddings retain cold initialization.The appendix specifies the attention projections and cold-initialized vectors as module conventions.Source: §B.2 · p. 15
  • Transferred 2D tensors inherit donor direction while matching the target seed’s cold-initialization Frobenius norm.The embed-orig-scale control instead copies donor token embeddings without normalization.Source: §B.2 · p. 15

Related papers

Loading 2609.18078v1…