Early Query–Key Learning Determines the Grokking Delay in Modular Addition

This deposit contains two documents and their LaTeX sources, reporting a causallocalization of the grokking delay in a one-layer Transformer trained on modularaddition (c = a + b mod 113). The four-page brief (note.pdf) reports the headline evidence. The fullmanuscript (main_preprint.pdf) develops the mechanistic account of thememorization and generalization computations and of the handoff between them.latex_source.zip holds the compilable sources for both. Headline findings: - The early window determines the fate. Freezing W_Q and W_K at initialization yields early crossing (validation accuracy above 0.9 at steps 1375-1925, 10/10 seeds). Freezing them only after the first 200 training steps leaves the network on a long structured-lookup plateau (censored 10/10 seeds). Later query-key plasticity is neither necessary for the fast branch nor sufficient to release the early slow state.- The outcome follows the query-key values, not the body. In a 2x2 state-factorial transplant at t=200, the two mismatched arms invert with the query-key block in both directions and in every seed: trained query-key into a would-be-fast body gives the slow fate, and a freshly initialized query-key block into the native body gives the fast fate.- The effective content is an operand-scoring code. The installed code's seed-specific assignment is causally dispensable, while its amplitude grades the delay length; the plateau is maintained jointly by the query-key values and the rest of the network. Relation to prior work. Yildirim (2026, arXiv:2603.05228) shows that replacinglearned, data-dependent attention routing with uniform aggregation removes thedelay on Z_p cyclic arithmetic, and that a bounded spherical topology does thesame. This deposit supplies the complementary direction: what native trainingwrites into the query-key block before the delay becomes visible, and whetherthat state is causally portable. We deliberately stop short of claiming a closedbidirectional account, since the two sets of interventions act at differentlevels (a wholesale replacement of the attention mechanism versus aparameter-level transplant within one training regime). Scope. One task family (c = a + b mod 113, "a b =" format) and a one-layerpre-norm Transformer; the primary regime exempts the query-key block from weightdecay. The early-window, transplant, identity-maintenance and amplitude resultsare ten-seed; the geometric and rank-1-sufficiency characterizations come from athree-seed discovery set and are labelled as such in both documents. Code and data. The reproduction package - per-claim run and verificationscripts, the archived per-seed results, and a manifest of sha256 hashes for everyfile - is public at https://github.com/keyiadiannao/frozen-QK-reduce-delay, withits own archived release under DOI 10.5281/zenodo.22710044. Its ERRATA.md recordstwo known errors in sealed aggregate files, one claim statement that over-readits own result, one omitted data file, the retained invalid runs, and thefindings of a pre-release audit together with their resolution. AI tools were used extensively during code development, experimental iteration,analysis assistance and manuscript preparation; the author designed the studies,verified all reported results, and is responsible for them.

Authors

Publication Details

Journal
Zenodo (CERN European Organization for Nuclear Research)
Published
2026-09-11
DOI
https://doi.org/10.5281/zenodo.22710058
Primary Topic
Machine Learning in Materials Science
Type
preprint
Controls
|||
ALL TIME
JAN
FEB
MAR
APR
MAY
JUN
JUL
AUG
SEP
OCT
preprint

Early Query–Key Learning Determines the Grokking Delay in Modular Addition

Zhengzhong Jin
Zenodo (CERN European Organization for Nuclear Research)
Machine Learning in Materials Science
preprint

Early Query–Key Learning Determines the Grokking Delay in Modular Addition

Zhengzhong Jin
preprint en

Abstract

This deposit contains two documents and their LaTeX sources, reporting a causallocalization of the grokking delay in a one-layer Transformer trained on modularaddition (c = a + b mod 113). The four-page brief (note.pdf) reports the headline evidence. The fullmanuscript (main_preprint.pdf) develops the mechanistic account of thememorization and generalization computations and of the handoff between them.latex_source.zip holds the compilable sources for both. Headline findings: - The early window determines the fate. Freezing W_Q and W_K at initialization yields early crossing (validation accuracy above 0.9 at steps 1375-1925, 10/10 seeds). Freezing them only after the first 200 training steps leaves the network on a long structured-lookup plateau (censored 10/10 seeds). Later query-key plasticity is neither necessary for the fast branch nor sufficient to release the early slow state.- The outcome follows the query-key values, not the body. In a 2x2 state-factorial transplant at t=200, the two mismatched arms invert with the query-key block in both directions and in every seed: trained query-key into a would-be-fast body gives the slow fate, and a freshly initialized query-key block into the native body gives the fast fate.- The effective content is an operand-scoring code. The installed code's seed-specific assignment is causally dispensable, while its amplitude grades the delay length; the plateau is maintained jointly by the query-key values and the rest of the network. Relation to prior work. Yildirim (2026, arXiv:2603.05228) shows that replacinglearned, data-dependent attention routing with uniform aggregation removes thedelay on Z_p cyclic arithmetic, and that a bounded spherical topology does thesame. This deposit supplies the complementary direction: what native trainingwrites into the query-key block before the delay becomes visible, and whetherthat state is causally portable. We deliberately stop short of claiming a closedbidirectional account, since the two sets of interventions act at differentlevels (a wholesale replacement of the attention mechanism versus aparameter-level transplant within one training regime). Scope. One task family (c = a + b mod 113, "a b =" format) and a one-layerpre-norm Transformer; the primary regime exempts the query-key block from weightdecay. The early-window, transplant, identity-maintenance and amplitude resultsare ten-seed; the geometric and rank-1-sufficiency characterizations come from athree-seed discovery set and are labelled as such in both documents. Code and data. The reproduction package - per-claim run and verificationscripts, the archived per-seed results, and a manifest of sha256 hashes for everyfile - is public at https://github.com/keyiadiannao/frozen-QK-reduce-delay, withits own archived release under DOI 10.5281/zenodo.22710044. Its ERRATA.md recordstwo known errors in sealed aggregate files, one claim statement that over-readits own result, one omitted data file, the retained invalid runs, and thefindings of a pre-release audit together with their resolution. AI tools were used extensively during code development, experimental iteration,analysis assistance and manuscript preparation; the author designed the studies,verified all reported results, and is responsible for them.

Zenodo (CERN European Organization for Nuclear Research)
Machine Learning in Materials Science
AI Navigator

Ask Laika to Summarize, Analyze, and Connect papers live on the map.

Summarize Papers & Methodologies

Extract key findings, datasets, and comparative methods across publications.

Benchmark Rankings & Visual Analytics

Rank top research institutions, authors, funders, topics, and journals by Field-Weighted Citation Impact (FWCI) and paper volume with instant charts.

Connect Distant Disciplines

Bridge topological clusters on the map to find hidden collaborative intersections.