Learning to Read Out: Unembedding Dynamics in Language Model Pretraining

Learning to Read Out. Unembedding Dynamics in Language Model Pretraining. By Matteo He and five others, from the University of Cambridge, Flower Labs, and one other institution. Published on the fourteenth of October. Its code repository is worked out: about one hundred and forty source files, with a readme and a licence.

When does a language model acquire a capability? Developmental studies usually answer by tracking when the model starts producing the correct token. But that prediction passes through the unembedding matrix, the readout, and the readout is itself still being learned during pretraining. The claim is that a model's hidden states can encode task information before its own readout expresses it. Timelines read from outputs therefore mix together two things: what the model encodes, and what its readout is able to express. A model might know the answer internally and still fail to say it, simply because the last layer that turns hidden states into tokens has not caught up. Dating the capability by the output would then date the readout, not the knowledge.

To watch the readout directly, the authors introduce parameter-trajectory crosscoding. It fits one sparse dictionary to every token's readout row, at every checkpoint. Unlike crosscoders applied to activations, it needs nothing but the checkpoint weights. Each token has a row in the readout, and the method follows all of those rows across training with one shared dictionary, so changes in the readout can be compared from checkpoint to checkpoint.

Applied to Pythia, it locates a narrow early window, between zero point four and one point one percent of the way through pretraining. In that window the readout reorganises, and WordNet word categories become decodable from its rows. By that point, linear probes on the hidden states already decode task information that the model's own outputs do not yet show. The inside is ahead of the outside: the information is there to be decoded before the model puts it into its predictions.

The second line of evidence is an intervention. To test whether such gaps depend on how fast the readout learns, the authors pretrain small models from scratch that differ only in the readout's learning rate. Raising that rate makes the readout develop earlier and the gap close sooner, with little effect on the final loss. Speed up the readout, and the behaviour shows up earlier, even though what the model ends up with barely changes.

The claim concerns early training, which is also where many developmental transitions are dated. The readout window sits in Pythia's pretraining, and the learning-rate test uses small models trained from scratch. And the conclusion is that a behavioral date can reflect when the readout learns to express task information: a possibility for any given milestone, not a rule for every one.

The implication the authors draw is a correction to how milestones are read. When a capability first appears in a model's outputs early in training, that date may mark when the readout learned to express the information, not when the model learned it. Watching outputs alone, the two are indistinguishable.