Facts in Stats: Impacts of Pretraining Diversity on Language Model Generalization

TL;DR

This work tackles how pretraining diversity shapes language model generalization by introducing a minimal synthetic testbed that simultaneously models a statistical learning stream and a factual recall stream. It formalizes data generation with $N$ templates and $K$ facts, controlled by diversity $DIV$ and contextual structures MC1PosN, MCNPos1, MCNPosN, to dissect how context and paraphrase variation influence ID versus OOD performance. The key findings show that higher contextual diversity can slow ID factual recall, but the effect on OOD generalization depends on the contextual structure; with sufficient training, high diversity can yield robust generalization, while low diversity often impedes OOD recall. The authors localize optimization bottlenecks to embedding and unembedding layers and attention, demonstrating that targeted module interventions can restore or improve generalization, thereby providing a controlled, low-cost framework for studying learning dynamics in language models.

Abstract

Language models are pretrained on sequences that blend statistical regularities (making text fluent) with factual associations between specific tokens (knowledge of facts). While recent work suggests that the variability of their interaction, such as paraphrases of factual associations, critically determines generalization ability, we lack a systematic analysis of these impacts. This paper introduces a flexible synthetic testbed that combines a statistical stream of generic tokens with an abstract factual stream of source-target token pairs, enabling fine-grained control over their interaction. The design enables the independent control of diversity nature by manipulating stream composition (contextual structure) and the diversity level by varying which statistical streams each fact appears in. Through controlled experiments, we find that while higher contextual diversity delays in-distribution (ID) factual accuracy, its impact on out-of-distribution (OOD) factual generalization depends critically on contextual structure. In some cases, OOD performance follows the same trend as ID, but in others, diversity becomes essential for non-trivial factual recall. Even when low diversity prohibits factual recall, optimal diversity levels depend on training duration. Beyond factual recall failures, we identify structures where statistical generalization fails independently, and others where both capabilities degrade. This shows how the interplay between contextual design and diversity level impacts different generalization aspects. Further, through a series of controlled interventions on the model components, we trace the OOD failures to distinct optimization bottlenecks, highlighting the importance of the embedding and unembedding layers. Our synthetic framework allows us to isolate effects that would be confounded in large-scale studies, offering a controlled testbed for future investigations.

Figures (15)

• Figure 1: Proposed testbed and summary of key findings.(A) Sequences are generated by combining a factual stream, a set of atomic pairs, and a statistical stream of templates, each a statistical distribution over sequences with two fact placeholder positions. Final sequences insert a fact into a sequence sampled from a template (Sec. \ref{['sec:setup']}) . (B) We control two key data properties. (Left) Contextual structure: The degree of statistical and (fact) positional variation across templates. We explore three structures: Stats Var (varied statistics, fixed positions), Pos Var (fixed statistics, varied positions), and Stats-Pos Var (varied on both axes). (Right) Diversity level: The number of unique templates each fact is paired with in the training corpus. The table shows an exposure matrix, where dots indicate template-fact pairs used for training (Sec. \ref{['sec:div']}-\ref{['sec:templates']}) . (C) Low diversity impairs out-of-distribution generalization, but the specific aspects affected (factual recall vs. statistical pattern generation) depend on contextual structure. At high diversity, all structures enable both types of generalization, though with varying efficiency (colored boxes: darker means more training iterations needed with higher $\texttt{DIV}$) (Sec. \ref{['sec:results']}) . (D) While training on low-diversity data impairs statistical and factual generalization, when the embeddings or unembeddings from a high-diversity model are "patched" in, the remaining modules can still learn a generalizing solution even when trained on that same low-diversity data (Sec. \ref{['sec:bottleneck']}) .
• Figure 2: Generalization performance for varying diversity levels and context structures. Heatmaps show (a)${\textbf{Loss}_{\rm{stat}}}$, (b)${\textbf{Acc}_{\rm{pos}}}$ and (c)${\textbf{Acc}_{\rm{fact}}}$ over training iterations (horizontal axis) and diversity level (vertical axis) for the three contextual structures defined in Sec. \ref{['sec:templates']}: (top) $\textsc{MC10Pos1}$ with only statistical variations, (middle) $\textsc{MC1Pos10}$ with only positional variations and (bottom) $\textsc{MC10Pos10}$ with both types of variations across templates. Each metric is displayed for in‑distribution (ID, left column) and out‑of‑distribution (OOD, right column) fact–template pairs. Brighter colors denote higher loss for ${\textbf{Loss}_{\rm{stat}}}$ and higher accuracy for ${\textbf{Acc}_{\rm{pos}}}$ and ${\textbf{Acc}_{\rm{fact}}}$. Increasing diversity can slow down convergence, whereas very low diversity leads to OOD failure on one or multiple metrics. See Sec. \ref{['sec:evaluation_metrics']} for metrics and Sec. \ref{['sec:results']} for details.
• Figure 3: Diversity pays off with time. Factual accuracy versus diversity level after $3$K, $4$K and $16$K iterations. Depending on the template structure, very low diversity can remain unrecoverable; no amount of additional training can restore OOD performance. High-diversity runs begin with lower accuracy but continue to improve until they surpass low-diversity models, demonstrating that diversity incurs an initial cost yet yields long-term benefits. Each bold curve is the average over three individual runs for OOD (solid) and ID (dashed) data.
• Figure 4: The complexity of one stream impacts learning of the other. Speed of learning the statistical, position, and factual metrics as data complexity is varied in a high-diversity setting, by varying (top) the MC order, and (bottom) the generic vocabulary size $V_\mathcal{D}$, number of facts $K$, and sequence length $T$. See Sec.\ref{['sec:param_ablation']} for discussion and App.\ref{['app:loss_metrics']} for metrics.
• Figure 5: Probing the Quality of Last-Layer Features. Learning curves for the unembedding layer ($\mathbf{U}$) retrained on last-layer features ($\Phi$) from models pre-trained on high- or low-diversity data. When using features from a low-diversity model, retraining $\mathbf{U}$ even on high-diversity data shows little improvement over the baseline (blue). Conversely, retraining $\mathbf{U}$ on low-diversity data leads to successful generalization when using features from a high-diversity model (green). However, these high-quality features degrade if they are not frozen and are updated along with $\mathbf{U}$ on the low-diversity data (red). In the legend, ${\text{❄}}$ denotes that the parameters are frozen, and the subscript indicates the diversity level ($\texttt{DIV}$) on which those parameters were originally trained. See Sec. \ref{['sec:prob']} for details.