Learning to Read Out: Unembedding Dynamics in Language Model Pretraining
Organizations: Department of Computer Science and Technology, University of Cambridge · Flower Labs · Isomorphic Labs
Abstract
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.
Figures & tables
| Task | Example prompt | Correct | Foil | |||
| SVA | simple | The keys to the cabinet | are | is | 320 | 320 |
| across PP | The judge by the street | appears | appear | 2,000 | – | |
| subject RC | The author that admires the senators | likes | like | 1,200 | – | |
| object RC | The author that the senators admire | likes | like | 1,200 | – | |
| IOI | availability | Lisa and Jane went to the store. Jane waited. Lisa gave the ring to | Jane | Lisa | 2,000 | – |
| localization | Ryan and Adam went to the store. Ryan gave the pen to | Adam | Ryan | – | 462 | |
Appendix figures & tables37 assets
Supplementary material from the paper’s appendix.
Appendix
| Control | Comparison |
|---|---|
| Step-0 hidden-state probe | Distinctions decodable before training versus gains in decodability during training. |
| Shuffled/random labels | Probe predictability under targets unrelated to the task assignment. |
| Prompt-token probe | Hidden-state predictions versus information in token presence/counts. |
| Gaussian-feature probe | Hidden-state predictions versus fitting an equally sized unrelated feature matrix. |
| WordNet matched controls | Lexical labels versus frequency, length, row norm, spacing, and token-class metadata or matched token sets. |
| Answer-row geometry | Differences in contrast scores attributable to candidate norms and token-frequency patterns. |
| Suite | Reason not used |
|---|---|
| OLMo 7B, April 2024 [ Groeneveld et al., 2024 ] | Sparser coverage before step 1000 than OLMo-2 |
| Olmo 3 [ Team Olmo et al., 2025 ] | No stage-1 checkpoint between steps 0 and 1000 |
| LLM360 Amber [ Liu et al., 2024 ] | Checkpoints indexed ckpt_000 to ckpt_359 instead of by training step |
| LLM360 Crystal [ Liu et al., 2024 ] | Phase-specific revisions, custom loading, and a code and data mixture |
| Cerebras-GPT Intermediate [ Dey et al., 2023 ] | No documented early step schedule, and some maximal update parametrization ( P) variants need custom loading |
| OpenLLaMA previews [ Geng and Liu, 2023 ] | Sparse previews that begin after step 1000 (3B from step 2.5k) |
| Hyperparameter | Pythia-1B | Pythia-160M | Pythia-6.9B | Pythia-6.9B | OLMo-2-7B |
|---|---|---|---|---|---|
| 24,576 | 24,576 | 32,768 | 32,768 | 32,768 | |
| Expansion factor | 12.0 | 32.0 | 8.0 | 8.0 | 8.0 |
| Sparsity coefficient | 0.3 | 0.3 | 0.6 | 0.3 | 0.3 |
| Optimizer | Adam | Adam | Adam | Adam | SparseAdam |
| Learning rate | |||||
| JumpReLU LR factor | 0.2 | 0.1 | 0.4 | 0.3 | 0.3 |
| Run family | Worker class | GPU memory | Wall clock | GPU-hours |
|---|---|---|---|---|
| Pythia-160M checkpoint-spanning fits | 1–3 A100-class GPUs | 40–80 GB/GPU | 4–8 h | 4–24 |
| Pythia-1B capacity sweep | 1 A100-80GB GPU | 80 GB | up to 12 h | up to 12 |
| Pythia-6.9B checkpoint-spanning fit, per setting | 4 A100/H100-class GPUs | 40–80 GB/GPU | 4–8 h | 16–32 |
| OLMo-2-7B checkpoint-spanning fit | 4 A100/H100-class GPUs | 40–80 GB/GPU | 4–8 h | 16–32 |
| 31M pretraining runs (four compared) | 1 A40 GPU per run | 48 GB | about 21 h per run | about 85 in total |
| Probe, null, baseline, and figure jobs | CPU workstation or 1 GPU | 24–80 GB if GPU | minutes–4 h | up to 4 |
| Model | Matrix / Family | Exp. | Seed | Params | EV | Mean | % Act. | Dead | ||
|---|---|---|---|---|---|---|---|---|---|---|
| 160M | cross-snap-32 | 8 192 | 10.7 | 0–4 | 403 M | 0.776 | 203.6–203.9 | 2.49 | 0.000 | |
| 160M | cross-snap-32 d | 16 384 | 21.3 | 0 | 806 M | 0.780 | 104.9 | 0.64 | 0.000 | |
| 160M | cross-snap-32 | 24 576 | 32 | 0–2 | 1.21 B | 0.920 | 286.4–287.2 | 1.17 | 0.000 | |
| 160M | final-snap SAE | 6 144 | 8 | — | 9.5 M | 0.758 | 864.2 | 14.07 | 0.000 | |
| 160M | final-snap SAE | 8 192 | 10.7 | — | 12.6 M | 0.793 | 1 088 | 13.28 | 0.000 | |
| 160M | final-snap SAE | 16 384 | 21.3 | — | 25.2 M | 0.867 | 1 921 | 11.72 | 0.000 |
| Profile | Rule | Meaning |
|---|---|---|
| persistent | or | sustained across training |
| transitional | , , , , and | rare localized mid-training peak |
| early-decaying | , , and | strong early, weaker late |
| late-emerging | , , and | weak early, stronger late |
| mixed/ambiguous | no previous rule matched | residual category |
| Model | Selected window | Same window | Other windows (resamples) | |
|---|---|---|---|---|
| Pythia-160M | 512–1.6k | 0.005 | 370 | 1k–3.2k (30) |
| Pythia-1B | 512–1.6k | 0.005 | 283 | 1k–3.2k (117) |
| Pythia-6.9B | 512–1.6k | 0.04 | 222 | 9k–28k (134) , 14k–44k (33) , 1k–3.2k (10) , 8k–25k (1) |
| OLMo-2-7B | 9k–28k | 0.12 | 399 | 900–2.8k (1) |
| Model | Width | Seeds | Early-decaying | Persistent | Late-emerging | Window | ||
|---|---|---|---|---|---|---|---|---|
| Pythia-160M | 8,192 | 0.3 | 5 | 45.6–46.1 | 0.5–0.9 | 52.5–53.0 | 512–1.6k | 0.005 |
| Pythia-160M † | 16,384 | 0.3 | 1 | 7.8 | 80.9 | 11.2 | 512–1.6k | 0.005 |
| Pythia-160M ∗ | 24,576 | 0.3 | 3 | 42.9–43.1 | 6.9 | 49.5–49.7 | 512–1.6k | 0.005 |
| Pythia-160M | 8,192 | 0.4 | 1 | 43.2 | 1.2 | 55.0 | 512–1.6k | 0.005 |
| Pythia-160M | 8,192 | 1.0 | 1 | 27.8 | 43.4 | 28.3 | 512–1.6k | 0.005 |
| Pythia-160M | 8,192 | 1.2 | 1 | 22.3 | 57.5 | 19.7 | 512–1.6k | 0.005 |
| Task | Template | Answers | Set | |
|---|---|---|---|---|
| Simple SVA | The {noun} {prepositional phrase} | are / is | A, L | 320 |
| Across-PP SVA | The {head} {preposition} the {attractor} | agreeing verb / other form | A | 2000 |
| Subject-RC SVA | The {head} that {clause verb} the {attractor} | agreeing main verb / other form | A | 1200 |
| Object-RC SVA | The {head} that the {attractor} {clause verb} | agreeing main verb / other form | A | 1200 |
| IOI | {A} and {B} went to the store. {R} waited. {S} gave the {object} to | {R} / {S} | A | 2000 |
| IOI | {A} and {B} went to the store. {A} gave the {object} to | {B} / {A} | L | 462 |
| family | unreach | (1B) | regime | ||||||
|---|---|---|---|---|---|---|---|---|---|
| SVA simple | 320 | 0.60 | 1.00 | 1.00 | 1.00 | 0.00 | – | 0.00 | transient |
| SVA across-PP | 2000 | 0.43 | 1.00 | 0.97 | 0.98 | 0.03 | 63% | 0.03 | transient |
| SVA subject-RC | 1200 | 0.73 | 1.00 | 0.75 | 0.79 | 0.25 | 85% | 0.21 | persistent |
| SVA object-RC | 1200 | 0.73 | 1.00 | 0.91 | 0.93 | 0.10 | 72% | 0.13 | persistent |
| IOI | 2000 | 0.92 | 1.00 | 1.00 | 1.00 | 0.00 | – | 0.00 | transient |
| numeric | 2000 | 0.86 | 1.00 | 0.63 | 0.82 | 0.37 | 47% | 0.55 | persistent |
| family | example prompt (truncated) | / | rank1 | |||
|---|---|---|---|---|---|---|
| BLiMP anaphor | Katherine can’t help | herself / himself | 1455 | 0.99 | 0.99 | 0.02 |
| BLiMP det-noun | Raymond is selling this | sketch / sketches | 838 | 0.92 | 0.94 | 0.03 |
| BLiMP NPI | Teresa had | not / probably | 656 | 0.98 | 0.97 | 0.15 |
| copy control | valley valley anchor anchor zebra | zebra / jacket | 2000 | 0.98 | 0.98 | 0.01 |
| fixed-token control | mirror thunder cookie music island | the / , | 2000 | 0.98 | 0.90 | 0.00 |
| FV antonym | selfish:altruistic / note:silence / known… | perfect / socialist | 2000 | 1.00 | 1.00 | 0.00 |
| step | step | final | |||||||
|---|---|---|---|---|---|---|---|---|---|
| model | family | ||||||||
| 6.9B | SVA simple | 320 | 0.60 | 0.93 | 0.50 | 0.49 | 1.00 | 1.00 | 1.00 |
| 6.9B | SVA across-PP | 2000 | 0.43 | 0.95 | 0.49 | 0.48 | 1.00 | 0.90 | 0.97 |
| 6.9B | BLiMP anaphor | 1455 | 0.60 | 0.63 | 0.48 | 0.48 | 0.85 | 0.90 | 0.99 |
| 1B | SVA simple | 320 | 0.45 | 0.97 | 0.50 | 0.49 | 1.00 | 1.00 | 1.00 |
| 1B | SVA across-PP | 2000 | 0.53 | 0.98 | 0.48 | 0.48 | 1.00 | 0.93 | 0.97 |
| Language | Targets | Share (%) | Gain | Gain | Gain | |||
| Aggregate | 173,229 | 100.0 | 1000 | 1.503 [1.455, 1.548] | 1000 | 0.147 [0.141, 0.153] | 1000 | 0 |
| English | 3,939 | 2.3 | 512 | 0.048 [0.026, 0.070] | 1000 | 0.064 [0.039, 0.089] | 2000 † | 0.002 [ 0.021, 0.026] |
| Russian | 10,639 | 6.1 | 512 | 1.033 [1.007, 1.060] | 1000 | 0.114 [0.090, 0.139] | 1000 † | 0 |
| Chinese | 25,330 | 14.6 | 1000 | 1.251 [1.176, 1.305] | 1000 | 0.134 [0.123, 0.146] | 1000 | 0 |
| Japanese | 20,698 | 11.9 | 2000 | 1.148 [1.090, 1.195] | 1000 | 0.192 [0.178, 0.207] | 2000 | 0.013 [0.002, 0.024] |
| ablate | preserve-only | hidden projection | ||||||
|---|---|---|---|---|---|---|---|---|
| Family | orig. | recon | top-8 | sign | top-8 | sign | native | proj. |
| Subject–verb agreement | 0.925 | 0.900 | 0.487 | 0.931 | 0.938 | 0.463 | 0.781 | 0.500 |
| Greater-than completion | 0.529 ∗ | 0.512 | 0.084 | 0.500 | 0.894 | 0.507 | 0.601 | 0.497 |
| Indirect object id. (IOI) | 0.455 ∗ | 0.429 | 0.394 | 0.424 | 0.481 | 0.420 | 0.377 | 0.338 |
| Relational facts | 0.269 ∗ | 0.231 | 0.154 | 0.231 | 0.923 | 0.308 | 0.308 | 0.154 |
| Condition | EV | val. loss | persist. | early-dec. | late-em. | mixed | med. peak | turnover step | |
|---|---|---|---|---|---|---|---|---|---|
| wu_lr_0p25 | 0.905 | 193.9 | 3.405 | 0.875 | 0.000 | 0.093 | 0.032 | 128 | 2048 |
| baseline | 0.891 | 177.2 | 3.381 | 0.403 | 0.219 | 0.125 | 0.252 | 64 | 1024 |
| warmup_short | 0.890 | 175.5 | 3.350 | 0.011 | 0.442 | 0.153 | 0.394 | 64 | 1024 |
| wu_lr_4x | 0.874 | 168.5 | 3.369 | 0.003 | 0.675 | 0.131 | 0.191 | 32 | 512 |
| Family | Analyzed object | Input data | Temporal correspondence | What the comparison isolates |
|---|---|---|---|---|
| Lens methods | Hidden states in vocabulary space | Text and forward passes | Vocabulary coordinates; native or learned readout | Predictions under the chosen readout |
| Vocabulary geometry | Input/output embedding rows | Weights; lexical annotations | Token identities or geometric statistics | Changes in geometry and lexical association |
| Independent SAEs | Activations per checkpoint | Activation corpus | Separate indices; post-hoc matching | Features within each fit |
| Activation crosscoders | Activations across checkpoints | Activation corpus | Shared sparse indices | Representation change; feature use via interventions |
| Parameter methods | Decomposed or sparse weights | Weights and inputs | Typically one model state | Parameter components and circuits |
| This work | Readout row trajectories | Weights; text for tests | Token identities and shared sparse indices | Readout change; swaps and ablations at fixed states |
Language Model Pretraining
Language model pretraining teaches models language patterns and knowledge from large text corpora before they are adapted to specific tasks.
See all 440 papers