Learning to Read Out: Unembedding Dynamics in Language Model Pretraining
Abstract
Hidden states encode task information before the readout expresses it, so behavioral timelines reflect readout learning rather than model acquisition, and faster readout training closes this gap earlier.
When does a language model acquire a capability? Developmental studies usually answer this by tracking when models start producing the correct token. But that prediction passes through the unembedding matrix, or readout, which is itself still being learned during pretraining. We show that a model's hidden states can encode task information before its own readout expresses it, so timelines read from outputs conflate what the model encodes with what its readout expresses. To track the readout directly, we introduce parameter-trajectory crosscoding, which fits one sparse dictionary to every token's readout row at every checkpoint. Unlike crosscoders on activations, it needs only checkpoint weights. We use it to locate a narrow early window, 0.4-1.1% of the way through Pythia's pretraining, in which the readout reorganizes and WordNet word categories become decodable from its rows. By then, linear probes on the hidden states already decode task information that the model's own outputs do not yet show. To test whether such gaps depend on how fast the readout learns, we pretrain small models from scratch that differ only in the readout's learning rate. We find that raising it makes the readout develop earlier and the gap close sooner, with little effect on final loss. Together, these results imply that early in training, where many developmental transitions are dated, a behavioral date can reflect when the readout learns to express task information and not when the model learns it.
Community
We introduce parameter-trajectory crosscoding, which fits one sparse dictionary to every token's readout (unembedding) row at every checkpoint, using only checkpoint weights. On Pythia, the readout reorganizes in a narrow early window, 0.4–1.1% of the way through pretraining. By then, linear probes on hidden states already decode subject number before the model's own readout expresses it, so a behavioral date can reflect when the readout learns rather than when the model does. In small models pretrained from scratch, raising only the readout's learning rate makes the readout develop earlier and closes the gap sooner.
Dictionaries (Pythia-160M, 1B and 6.9B; OLMo-2-7B) and code are released, and you can browse every feature in the atlas: https://learning-to-read-out.pages.dev/atlas/
Get this paper in your agent:
hf papers read 2610.21025 Don't have the latest CLI?
curl -LsSf https://hf.co/cli/install.sh | bash