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
- Zhengzhong Jin (ORCID: https://orcid.org/0009-0007-3408-212X)
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