Learning to Read Out
Unembedding Dynamics in Language Model Pretraining
2026
When does a language model acquire information, and when can it use that information to predict the next word?
A model’s output layer turns its internal information into scores for possible next words. I studied how that layer changes during training, using probes to check what information was already present and trajectory crosscoders to track how the output weights developed.
In subject–verb agreement tests, grammatical information was detectable in hidden states hundreds of training steps before the model used it reliably in its predictions. The gap remained on harder sentences at the end of training. Predictions alone can therefore give a misleading date for when this information became available.
Explore how readout features develop during training →
Method & scope
I tracked snapshots of the output weights during pretraining and fitted trajectory crosscoders to study how they changed. I tested the resulting features with readout swaps, linear probes, and ablations. On harder examples involving relative clauses, the gap between linear decodability and the model’s own readout persisted at convergence.
Using fixed hidden states from Pythia-1B at step 1000, removing eight localized features from the reconstructed final readout reduced agreement accuracy on held out examples from 90.0% to 48.7%. Matched controls stayed at 90.0%. In control runs with 31 million parameters and a single random seed, each 4× increase in the readout learning rate halved the training step at which reorganization occurred. Final validation losses agreed within 0.06 nats.
The public release includes reusable analysis tools, examples that run on a CPU, guided notebooks, and a manifest linking experiments to figures. The paper was accepted to NeurIPS 2026.
