Learning to Read Out: Unembedding Dynamics in Language Model Pretraining

AI summary

Overview

Research area: Mechanistic interpretability and language-model pretraining dynamics — specifically, how the unembedding matrix (the "readout") develops during pretraining and how that development distorts the dating of when capabilities emerge.

Technical level: Advanced. The paper assumes familiarity with sparse autoencoders, crosscoders, logit lens, linear probing, and transformer internals.

Scope: The paper introduces parameter-trajectory crosscoding to track readout development across 32 public checkpoints of Pythia-160M/1B/6.9B and OLMo-2-7B, shows that hidden states can encode task information before the readout expresses it, and tests via controlled 31M-parameter pretraining runs whether the readout's learning rate sets that timing.

What This Paper Is About

Developmental studies of language models usually date a capability by when the model starts producing the correct token. But that prediction passes through the unembedding matrix, which is itself still being learned. The authors show that a model's hidden states can already encode task information while its own readout does not yet express it, so output-based timelines conflate what the model encodes with what its readout expresses. Their goal is to track the readout directly, locate when it reorganizes, and test whether its learning speed controls when available information becomes natively expressed.

Key Contributions

1. Problem formalization. The paper formalizes a confound in behavioral dating: a logit change decomposes into a hidden-state term, a readout term, and their interaction, and logits alone cannot separate them. This ambiguity also affects logit-lens comparisons, vocabulary projections, and attention-head scores computed through the readout. Diagnostics are proposed to separate the two sources.

2. Method: parameter-trajectory crosscoding. A crosscoder fit to readout weights rather than activations — one sparse dictionary fitted to every token's readout row at every checkpoint, with a shared code per token and checkpoint-specific decoders. Unlike activation-based crosscoders, it needs only checkpoint weights (no activation corpus, no labels) and covers the full vocabulary rather than only tokens a corpus elicits.

3. Readout dynamics. Identification of a narrow early window, spanning 0.4–1.1% of Pythia's pretraining (between step 512 and step 1.6k), in which the readout reorganizes and WordNet word categories become decodable from its rows.

4. Availability versus expression. Evidence that Pythia-1B and 6.9B hidden states encode subject number before the model's own readout or any intermediate readout expresses it, plus a controlled pretraining intervention showing that raising only the readout's learning rate makes the readout develop earlier and closes the gap sooner, with little effect on final loss.

Main Findings

• A single early reorganization window in Pythia. The reorganization score peaks between step 512 and step 1.6k in Pythia-160M, 1B, and 6.9B — 0.4% to 1.1% of training — a window containing the end of learning-rate warmup at step 1430. It stands out against circular-shift nulls with p = 0.005 at 160M and 1B and p = 0.04 at 6.9B. Of 21 Pythia dictionaries fitted, 20 selected the same window and the remaining one selected the adjacent window.

• OLMo-2-7B reorganizes at a similar fraction of training but not significantly. Its score peaks between step 9k and step 28k, 1–3% of the way through OLMo's first training stage, but does not separate from the nulls (p = 0.12).

• Features do not share one growth pattern. In Pythia-160M and 1B some features are early-decaying, some persistent, and some late-emerging; nearly all Pythia-6.9B features persist, and most OLMo-2-7B features emerge late. A Pythia-1B dictionary fitted only from step 256 onward assigns 0.02% of features to the early-decaying profile, against 24% in the main dictionary.

• Lexical organization develops alongside structural change. In Pythia-160M the median WordNet probe accuracy across 35 categories rises from 0.499 at initialization to 0.874 at the final checkpoint; all 35 categories become decodable by step 1000, and 33 have their largest gain between steps 128 and 1000. In Pythia-1B the median category becomes decodable at step 512, with onsets from step 128 to step 2000. A lemma-split control lowers final median accuracy to 0.841 with 34 of 35 categories decodable by step 1000.

• Availability precedes native expression. On Pythia-6.9B simple subject–verb agreement, the logistic probe approaches ceiling by step 512 while native expression and even the best swept intermediate readout remain at chance.

• Later readouts decode early states better than their own. On Wikipedia text in nine languages, the step-1000 readout lowers the aggregate next-token loss of Pythia-1B step-256 and step-512 states by 1.50 and 0.15 nats. The final readout raises early task accuracy from 0.800 to 0.925 on Pythia-1B step-1000 SVA states, but this falls to 0.566 after orthogonal Procrustes alignment.

• A few readout directions carry the agreement margin. Eight features carry the SVA margin on step-1000 states under the reconstructed final readout: ablating them lowers accuracy from 0.900 to 0.487, keeping only them gives 0.938, and ablating matched random features leaves 0.931.

• Readout learning rate shifts timing. Across the 31M runs, raising the readout learning-rate multiplier moved the median feature peak step from 128 to 64 to 32 (one checkpoint per fourfold increase), while halving warmup from 1430 to 715 steps left it at baseline. Gaps between availability and native expression close by step 1024 in every run except the 0.25× run, which closes by step 2048.

• Final loss is largely unaffected. At 0.27B tokens the 4× run's validation loss is 1.20 nats below baseline and 2.08 nats below the 0.25× run, while final validation losses agree within 0.06 nats.

• Dictionary fidelity. Across all four models the dictionaries explain 0.808–0.920 of the variance in readout trajectories; each readout row uses 1.17–2.26% of features at each checkpoint, and no feature is dead.

Methodology in Plain English

The authors train a sparse dictionary on the rows of the unembedding matrix rather than on activations. Each vocabulary token has one shared code across all checkpoints, but each checkpoint has its own decoder directions and its own gating threshold. The gate compares a feature's contribution norm to a learned threshold, which lets a feature switch on or off over training while keeping the same index — so a single feature can be followed from early to late checkpoints. The training objective combines reconstruction error with a saturating firing-rate penalty that discourages broadly used features.

Features are then described by their decoder-norm trajectories over checkpoints, normalized by their own maximum (so shapes, not heights, are compared) and classified into lifecycle profiles such as early-decaying, persistent, or late-emerging. A "reorganization score" combines six kinds of change between adjacent checkpoints into a geometric mean and is tested against circular-shift nulls that shuffle each measure's timing independently.

To separate what a hidden state encodes from what the readout expresses, the authors use two probes of the same fixed state: a readout swap, which decodes a checkpoint's hidden state with another checkpoint's readout and measures the margin between a correct answer and a foil, and a linear probe, which reads the state directly without any readout. They compare native expression (own readout), an oracle sweep over all intermediate readouts, and probe accuracy. Finally, they train four 31M-parameter models from a shared initialization on the first 10B tokens of Pythia's data, differing only in the readout's learning-rate multiplier (0.25×, 1×, 4×) or in warmup length, to test causal influence.

Why This Matters

Impact on research. Any developmental analysis that reads progress off logits — surprisal curves, logit lens, translation-accuracy dates — depends on a readout that is still learning. The paper argues that early behavioral dates can reflect when the readout learned to express information rather than when the model learned it, and recommends reporting native expression, a probe with its own initialization baseline, and a readout sweep on the same states. Because readout change concentrates at 0.4–1.1% of Pythia training, close to where the induction-head phase change is dated (roughly 1–2%), the confound may touch landmark developmental findings.

Real-world applications:

• Evaluating early checkpoints and small proxy runs. Recipe comparisons made on partially trained models can favor a faster readout even when final losses agree, so benchmark rankings from short runs may mislead.
• Vocabulary extension and domain adaptation. Adding tokens to a pretrained vocabulary creates untrained readout rows; giving those rows their own learning rate may speed adaptation.
• Capability forecasting for safety. Internal progress measures may anticipate behavioral transitions, and separating availability from expression could sharpen forecasts of safety-relevant capabilities.
• Pretraining diagnostics. The reorganization score gives a weights-only signal for when the readout is unstable, without needing an activation corpus.

Industry relevance. Teams comparing training recipes, scaling laws, or checkpoint-selection strategies on partial runs should account for readout speed as a confound in early evaluation, particularly since the intervention alters early loss substantially while leaving final loss within 0.06 nats.

Future Directions

1. Test whether recipe rankings from probes or swapped readouts match final rankings. The authors suggest this as a direct check on a possible fix for early-checkpoint evaluation.
2. Extend the readout learning-rate result beyond one 31M setting. The evidence rests on a single configuration with one run per condition; the readout holds about 42% of parameters in those 31M models versus about 10% in Pythia-1B, so larger models remain untested.
3. Train newly added vocabulary rows with their own learning rate to see whether this speeds adaptation when a pretrained model's vocabulary is extended.
4. Measure weight drift in every layer to determine whether each component, like the readout, settles on its own schedule.
5. Explain the different feature trajectories within the early window. The authors speculate that early-decaying features may record a unigram-frequency stage and late-emerging features a later context-using stage, but this remains open.

Target Audience

Interpretability researchers studying training dynamics, sparse dictionary learning, and feature evolution across checkpoints; pretraining practitioners who compare recipes or select checkpoints using early or small-scale runs; and evaluation researchers who date capabilities from model outputs. The paper is less suited to beginners, since it presupposes comfort with crosscoders, logit lens, linear probing, and transformer weight structure. Readers interested in capability forecasting and in the distinction between encoded and expressed knowledge will find the framing most useful.

Authors’ 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.