Generalization Dynamics of LM Pre-training
Jiaxin Wen, Zhengxuan Wu, Dawn Song, Lijie Chen
UC Berkeley / Stanford University, Google DeepMind
cs.CL
2026-09-27
A toy suite shows pre-training is not a stable climb: OLMo3-32B drops 81% to 0% in 40B tokens, and a 4.5T checkpoint beats 4.9T on GPQA after math SFT (36.3% vs 29.8%).
The usual mental model of pre-training is a smooth climb. Loss falls, standard benchmarks rise, and the model grows from a parrot that copies shallow patterns into a system that extracts the actual task. That picture rests on pre-training loss and downstream scores. Both look calm.
Berkeley and Stanford built a cheap eval suite that baits those shortcuts, then walked every public checkpoint of OLMo3 and Apertus. While loss stays flat, the model still snaps between parrot-like and generalizing computation. They call this mode-hopping. A hop can fit inside a few tens of billions of tokens, and it can reverse.
The models are dense LMs with public data and intermediate checkpoints: OLMo3 7B and 32B, Apertus 8B and 70B. Training already sits at 9x to 90x the Chinchilla-optimal token budget, so undertraining is not an excuse. The analysis keeps only the general pre-training stage, dropping mid-training and long-context stages so data sampling cannot explain the swings.
The suite is meant to be cheap: zero-shot or few-shot. Each prompt plants a tempting shallow pattern in context, then the test item breaks it. Scoring compares probabilities on the correct and incorrect answer spans, not graded samples, which shuts out generic instruction-following noise. Every number is a mean over 4 seeds.
Six prompting evals, one bait each:
Two fine-tuning tests sit on top. Out-of-context reasoning fine-tunes on I/O of an anonymized Python function or city and asks the model to name it. Emergent misalignment fine-tunes on insecure code and measures probability on misaligned choices for broader user queries.
The cleanest trace is Successive Answer, the answer+1 task. OLMo3-32B is at 81% after 2.17T tokens, 0% at 2.19T, and 81.7% at 2.21T. The window is about 40B tokens.
The same hopping shows up across flipped labels, repeated answers, truthiness, System 1, multi-hop personas, and out-of-context reasoning after fine-tuning. For emergent misalignment, ΔP(misaligned choice) crosses zero more than once: the same insecure-code fine-tune makes some checkpoints worse and others slightly safer. From 2T tokens on, OLMo3-32B names the Hitler persona at about 90%, while multi-hop accuracy on neighboring checkpoints swings between 20% and 95%. The 7B model locks more often; letter counting sits near 0% for the whole run. Apertus 8B and 70B trace the same family of curves. Scaling makes generalization more frequent. Larger models still hop, just on harder items.
Several nulls fail. On 10 ordinary datasets covering sentiment, topic, math word problems, and knowledge QA, the same checkpoints are flat. Soft probability P(correct) - P(incorrect) still hops, and so does accuracy versus pre-training FLOPs, so this is not a hard-accuracy mirage. One extra pre-training step barely moves the suite probabilities, even at learning rate 1e-2 and small batch size, which does not match edge-of-stability oscillation. Averaging K=5 consecutive checkpoints along the oscillating curve only damps the hops. It does not remove them.
Hops are not synchronized across datasets. Under Flipped Answer, SST-2 and IMDB correlate at only 0.43; sentiment versus topic is near zero. Paraphrases of the same dataset correlate strongly. Larger models show higher average correlation. Complexity proxies on activations and gradients (RankMe, participation ratio, Fisher trace, sharpness, gradient cosine), after picking the best layer, average 0.35 to 0.56. A random baseline already hits 0.4 under that selection. Well-generalizing checkpoints on the same layer can be high-rank or low-rank.
| Setting | Metric | 4.5T ckpt | 4.9T ckpt |
| Math SFT, then GPQA | accuracy | 36.3% | 29.8% |
| General post-training + prefilling | alignment robustness | 53% | 21% |
The 4.5T checkpoint generalizes more consistently on the toy suite, and it beats 4.9T after post-training. It also beats the final pre-training and mid-training checkpoints. Further pre-training or mid-training lifts in-distribution scores. A small continued-pretraining study that picks windows which originally pushed the model toward pattern-matching or toward generalization can pin later dynamics in the chosen direction. Random data keeps hopping.
The last checkpoint is not guaranteed to carry the strongest generalization prior. The suite is cheap and already useful for picking an intermediate point that post-training can actually use: +6.5 points on GPQA after math SFT, and prefilling robustness 53% versus 21%.
The working story is capacity allocation. In a capacity-bounded model, shallow circuits learned early compete with later generalizable ones, and the data in each window may decide the winner. Training longer does not automatically stabilize the winner. At matched FLOPs, larger models show no generalization edge.
Teams that already dump intermediate checkpoints can run this suite today. Teams that do not should stop treating the final point as the default. The paper flags data mixture, optimizers, and architecture as the places where generalizable circuits might overwrite shallow ones more steadily.
The evals are intentionally tiny. Persona QA has 5 test questions per figure, sampled 200 times each. Multi-hop items about historical figures are a narrow stand-in for persona representations. Cross-dataset correlation is low, so one or two toy tasks are not a universal generalization dashboard.
The data-selection study is a small preliminary run, not a 32B train-from-scratch. Weight averaging tried only K=5. Complexity metrics used best-layer selection, so 0.35 to 0.56 correlation is barely above a 0.4 random baseline.
Only OLMo3 and Apertus are covered. Closed models are untested. On the fine-tuning side, ΔP crossing zero is still a conjecture: the representations that support fine-tuning generalization may break and recover at different steps. The hoped-for picture is a capacity-bounded model that keeps rewriting circuits toward generalization. Small models stay stuck in shallow circuits. Large models still fall back.