A model splits a word's two meanings early in pretraining: two heads carry 14M, while 410M spreads the job across steadier ones
Synopsis
During pretraining of Pythia language models, word-sense disambiguation improves in sharp steps around 1,000 to 2,000 training steps, two layer-3 attention heads in 14M track that improvement and their ablation measurably hurts it, while 410M spreads the function across more heads that hold up better under position and part-of-speech perturbations.
Interpretation
The ability to separate the senses of an ambiguous word takes shape early in pretraining rather than climbing steadily. Prior work on attention-head specialization centered on induction heads, which track repeated token sequences; lexical ambiguity, where one word form carries entangled meanings, had rarely been used to characterize attention specialization. Across final checkpoints of the Pythia suite from 14M to 12B, 14M reached an R² of 0.15 and the best model, 6.9B, reached 0.598; 410M reached 92% of 12B's performance with about 3% of its parameters, and 14M had already approached its final level by about 4.2B tokens, roughly 1.4% into pretraining, with marked shifts at steps 1,000 and 2,000.
The heads that carry disambiguation are not the same set in 14M and 410M, and the larger model's heads depend less on the particular sentence frame. In 14M only heads (3,1) and (3,2) in layer 3 strengthened in step with disambiguation performance, and head (3,2)'s attention to the disambiguating word sat about 3.7 standard deviations above the model's mean attention to that word. When an inserted string separated the disambiguating cue from the target noun, only (3,1) in 14M stayed stable while (3,2) shifted to the token immediately preceding the target; under part-of-speech changes, the 14M heads attending to disambiguating verbs and those attending to disambiguating nouns were nearly disjoint. In 410M, heads including (1,14), (4,7), (1,3), (2,3), (6,10) and (1,13) kept attending to disambiguating cues across both perturbations.
Ablating these heads directly degrades disambiguation, but the size of the damage depends on model scale and on which layer's representations are read out. Earlier causal tests of attention heads targeted induction heads and in-context learning; disambiguation had not been systematically ablated this way. Setting a target head's query and key matrices to zeros, or reverting them to their step-1 values, significantly increased ΔR² and reduced the fraction of intact R² in 14M, with effects concentrated in layer 3; in 410M the effect was present but very small when averaged across layers, absent at the best-performing layer 24, and by the end of training indistinguishable from ablating control heads.
Perspective
The work examines English noun disambiguation in the Pythia suite across pretraining checkpoints, using 504 minimal sentence pairs from RAW-C, of which 318 senses were polysemous and 186 homonymous; causal ablation was run only on 14M and 410M, and functional assessment only on the original RAW-C pairs. The authors propose collecting human judgments for the modified stimuli, extending the full analysis to large Pythia models beyond 14M and 410M or to other model families with public checkpoints, and using ambiguity datasets in languages other than English.
The authors note that candidate heads in 14M did not generalize across sentence frames or parts of speech, so they may not be generalized disambiguation heads and may participate in noun disambiguation specifically; in 410M the effect of ablating a single head was very small, which may reflect overlap among head functions. Replication across nine random seeds showed highly consistent developmental trajectories of disambiguation performance, but the layer hosting the key heads was bimodal across seeds, falling in layer 3 for roughly half and layer 4 for the other half. The authors also leave open whether these heads are selective for ambiguous target words or participate in contextualizing any target word.