Mechanistic analysis demonstrates that early query–key parameter dynamics control grokking delays in one-layer Transformers, indicating that early attention weights dictate generalization fate.
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.
No takes yet. Share an insight, caveat, or question.
Zhengzhong Jin (2026) studied this question.
Synapse has enriched 5 closely related papers on similar clinical questions. Consider them for comparative context: