Generalization Dynamics of LM Pre-training

· ArXivIQ ·

9 min read Original article ↗

Authors: Jiaxin Wen, Zhengxuan Wu, Dawn Song, Lijie Chen
Affiliations: UC Berkeley, Stanford University, Google DeepMind
Paper: https://arxiv.org/abs/2609.33150
Code: N/A
Model: N/A

WHAT was done?
The authors probe whether language models mature smoothly from pattern-matching heuristics into generalizable computation by tracking dense intermediate pre-training checkpoints across open-weight models using a lightweight, diagnostic evaluation suite.

WHAT was found?
Across scales spanning 9x to 90x the Chinchilla-optimal budget, models do not stably retain reasoning capabilities: on the successive arithmetic task, OLMo3-32B hits 81% accuracy at 2.17T tokens, collapses to 0% at 2.19T tokens, and rebounds to 81.7% at 2.21T tokens. Selecting an intermediate 4.5T checkpoint over the 4.9T checkpoint via this behavioral diagnostic yields 36.3% vs. 29.8% accuracy on GPQA after math fine-tuning, and increases resilience against prefilling jailbreaks from 21% to 53%.

WHY it matters?
Selecting the final checkpoint from a long pre-training run often inherits degraded reasoning and safety representations, though the diagnostic signals show high variance across distinct domains.

Standard monitoring during pre-training relies on validation loss and aggregate benchmarks. Both metrics present an orderly picture. Cross-entropy loss falls smoothly, while downstream benchmarks show monotonic progression or sharp emergence as characterized by Wei et al. (2022).

Mechanistic frameworks such as The Implicit Curriculum Hypothesis suggest that skills emerge in stable, predictable orders. In contrast, this paper asks a sharper behavioral question: when a model faces a prompt where a superficial pattern conflicts with the true underlying computation, does it execute the rule or follow the heuristic?

By evaluating intermediate checkpoints of OLMo3 (7B and 32B) and Apertus (8B and 70B), the authors uncover a phenomenon they term mode-hopping. Models repeatedly alternate between parrot-like pattern matching and generalizable reasoning throughout pre-training.

This oscillation is completely invisible on standard benchmarks like MMLU or HellaSwag, as shown in Figure 9.

Figure 9

To isolate mode-hopping from instruction-following variance, the authors build a diagnostic evaluation suite spanning six task categories, summarized in Table 1 and detailed in Table 3. Each task sets up a strong, tempting local heuristic in the context window or training prior that directly contradicts the correct output. Instead of open generation, evaluation measures log-likelihoods over predefined target tokens.

Table 1
Table 3

To see the mechanics of mode-hopping in action, consider the arithmetic successive answer evaluation from Table 5. The prompt contains in-context demonstrations that follow an arithmetic sequence while incrementing the target output: Q: 8 - 7 = ? -> 1, Q: 1 + 1 = ? -> 2, Q: 192 - 189 = ? -> 3. The query prompt then tests an arithmetic operation whose real result breaks the incrementing pattern: Q: 68 - 60 = ?

Table 5

A model operating in a pattern-matching regime relies on sequence continuation circuits like successor heads, emitting 4. A model executing in-context task induction activates functional representations, outputting 8.

When tracking OLMo3-32B across pre-training tokens in Figure 4, the model alternates between these two circuits. It jumps to 81% accuracy at 2.17T tokens, completely collapses to 0% at 2.19T tokens by emitting 4, and then jumps back to 81.7% at 2.21T tokens.

Figure 4

The same dynamic appears across all six evaluations. In label-flipped sentiment classification (Figure 2), models oscillate between following flipped in-context exemplars and reverting to memorized semantic priors.

Figure 2

In Cognitive Reflection Tests (Figure 6), models alternate between intuitive System 1 answers (such as predicting $0.10 for the classic bat-and-ball problem) and calculated System 2 answers ($0.05).

Figure 6

In multi-hop persona identification (Figure 7), models track isolated biographical facts stably, yet their ability to integrate those facts across multiple hops swings between 20% and 95% accuracy across neighboring checkpoints.

Figure 7

A natural reaction is to dismiss these oscillations as artifacts of hard accuracy thresholds, measurement noise, or optimization instability. The authors rule out these explanations through targeted counter-experiments:

  • Soft Metric Evaluation: Replacing discrete accuracy with the soft probability margin P(correct) - P(incorrect) does not smooth the curves. As shown in Figure 26 and Figure 28, the probability margins swing repeatedly between positive and negative values.

    Figure 26
    Figure 28
  • Edge of Stability: Under the edge-of-stability regime documented by Cohen et al. (2021), loss oscillations occur over single update steps. The authors test the local stability of mode-hopping in Figure 10 by loading intermediate Adam optimizer states and taking single gradient steps across various batch sizes and learning rates up to 10-2. The output probabilities on the suite do not change, demonstrating that each behavioral mode is locally stable.

    Figure 10
  • Weight Averaging: Applying model soups (Wortsman et al., 2022) by averaging the weights of K = 5 consecutive checkpoints along an oscillation wave fails to resolve the instability (Figure 11). Averaging dampens the variance slightly but preserves the underlying mode switches.

Figure 11

I find the local stability analysis in Figure 10 to be the most convincing piece of evidence in the paper. It establishes that mode-hopping is not gradient noise or high-learning-rate turbulence. Instead, it represents genuine transitions between distinct, multi-step attractors in parameter space.

However, I find the authors’ dismissal of checkpoint averaging based solely on uniform averaging of K = 5 consecutive checkpoints somewhat premature. Testing exponential moving averages or greedy soups with non-uniform weights would have offered a more thorough check of weight-space interpolation.

Why do these transitions occur? The authors frame mode-hopping as a capacity-allocation contest between shallow heuristics and generalizable circuits (Gu et al., 2025). Early in training, simplicity bias drives models to learn low-complexity pattern-matching heuristics. Later, generalizable circuits emerge, as explored in mechanistic studies of sparse attention patterns (Baherwani et al., 2026).

In capacity-bounded networks, these circuits compete for the residual stream. Depending on the statistical composition of documents within a given training window, shallow circuits can suppress general circuits, or vice versa.

To evaluate whether representation complexity predicts which mode dominates, the authors analyze geometric and spectral metrics across layers. The effective rank is captured via RankMe over the normalized eigenvalue spectrum pᵢ = λᵢ/∑ⱼ λⱼ of the activation Gram matrix: RankMe = exp(-∑ᵢ pᵢlog pᵢ)

To quantify directional alignment, they compute the mean absolute cosine similarity between per-example loss gradients gᵢ across N test examples:

\(|\cos| = \frac{2}{N(N-1)}\sum_{i<j} \frac{|\langle g_i, g_j \rangle|}{\|g_i\| \|g_j\|}\)

The results in Figure 38 and Figure 39 complicate standard intuitions. A widespread hypothesis in representation learning suggests that generalizable solutions are simpler, which would imply lower activation rank or sharper compression (Conklin et al., 2026).

Yet the authors find that RankMe yields both strongly positive correlations (ρ+ ≈ 0.7) and strongly negative correlations (ρ- ≈ -0.6) depending entirely on which layer is measured. As illustrated in Figure 40, well-generalized checkpoints alternate between high-rank and low-rank representations across identical layers. Generalizable solutions are not uniformly simpler; they merely allocate capacity differently across network depth.

The practical consequences of mode-hopping become apparent when selecting checkpoints for post-training. Standard practice typically selects the final pre-training or mid-training checkpoint. However, if the final checkpoint lands in a shallow heuristic valley, it transfers poorly to downstream tasks.

The authors contrast an OLMo3-32B checkpoint trained on 4.5T tokens (which scores high on the toy suite) with a later checkpoint trained on 4.9T tokens (which scores low, as shown in Figure 17).

Figure 17

Both checkpoints undergo identical mathematical supervised fine-tuning using the data regime of Ren et al. (2026). As recorded in Figure 13a, the 4.5T checkpoint generalizes substantially better to out-of-distribution reasoning on GPQA (36.3% vs. 29.8%).

Figure 13

When aligned using safety data from STAR-1 and evaluated against prefilling attacks (Qi et al., 2025), the 4.5T checkpoint achieves 53% safety retention compared to 21% for the 4.9T checkpoint (Figure 13b). Inspecting Figure 13 reveals that the 4.5T model also outperforms the final pre-training and mid-training checkpoints. Extra pre-training tokens can actively degrade reasoning foundations if training terminates during a shallow-circuit collapse.

The demonstration that the 4.5T checkpoint beats the final 5.15T and mid-training checkpoints on GPQA is an important empirical finding. However, evaluating only two hand-selected checkpoints leaves open the concern that this specific pair was selected post-hoc rather than discovered through an automated stopping criterion.

Furthermore, the authors run a controlled intervention in Figure 14. They categorize pre-training data slices based on whether they previously triggered generalization or pattern matching. Continued pre-training on selected generalizing data stabilizes performance on the successive answer task, whereas pattern-heavy data drives the model into a persistent heuristic state.

Figure 14

Despite these findings, two key limitations weaken the broader reach of the claims:

First, mode-hopping lacks cross-task consistency. As shown in Figure 12 and the correlation matrix in Figure 15, the rank correlation of mode-hopping between different dataset pairs is low. On flipped answer tasks, the correlation between SST-2 and IMDB is modest (r = 0.43). The correlation between sentiment classification and Yahoo topic classification is near zero.

Figure 12
Figure 15

A checkpoint that reasons well on arithmetic might simultaneously collapse on label flipping. This limits the diagnostic utility of the suite: selecting a checkpoint that excels across one subset of tasks does not guarantee global generalization across all domains.

Second, the data selection experiment in Figure 14 is a retrospective proof of concept. The authors identify generalizing data by observing its historical effect on a single task in a prior run. They do not provide an unsupervised or forward-looking method to classify raw pre-training text as generalizing without first running checkpoint evals.

I would put the following question to the authors: Given that cross-dataset correlations are near zero between distinct domains, how can one construct a balanced, holistic score from this diagnostic suite to ensure that selecting for arithmetic generalization does not select against factual persona reasoning?

This work demonstrates that pre-training convergence is far less monolithic than standard loss tracking suggests. Practitioners training foundation models should not assume that the final step yields the best reasoning base.

Running lightweight behavioral conflict tests throughout pre-training can identify checkpoints that retain generalizable reasoning circuits before downstream fine-tuning. Future research must focus on dynamic data mixtures and optimizers designed to prevent shallow heuristics from suppressing generalizable circuits during long training runs.

Discussion about this post

Ready for more?