Facts in Stats: Impacts of Pretraining Diversity on Language Model Generalization
Abstract
Language models are pretrained on sequences that blend statistical regularities, which make text fluent, with factual associations between specific tokens, which encode knowledge of facts. While recent works suggest that generalization depends critically on the interaction of these two streams, such as the diversity of the contexts in which facts appear, these effects remain difficult to study systematically. 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, such as their composition into a sequence (contextual structure) or the level of context diversity carrying the facts at training time. Through controlled experiments, we find that higher contextual diversity delays in-distribution factual learning, while low diversity can harm out-of-distribution generalization in ways that depend on the contextual structure. As a result, the optimal diversity level depends on the training budget. Beyond factual recall failures, we also identify failures in statistical generalization as we study how the interplay between contextual design and diversity level impacts different aspects of generalization. Furthermore, through a series of controlled interventions on the model components, we trace failure in different aspects of generalization to distinct optimization bottlenecks, highlighting the importance of the embedding and unembedding layers. Overall, our synthetic framework allows us to isolate effects that would be confounded in large-scale studies, offering a controlled testbed for future investigations.
Lay Summary
Language models are trained on large collections of text by learning to predict the next word. From this simple objective, they learn both regular patterns of language, such as grammar, and factual knowledge. In this work, we want to understand when a model can separate these sources of information and use a fact correctly in a new context while still following the linguistic rules. Studying this directly in natural language is difficult, as facts and linguistic patterns are deeply entangled and hard to track separately. Instead, we build an abstract dataset where facts are represented as pre-defined word pairs, similar to how Paris and France are related, and the surrounding context is generated by simple statistical rules. The simplicity then lets us control how many different contexts each fact appears in during training, similar to controlling how many paraphrases of a fact the model sees. We find that the context diversity creates a trade-off in the model's performance. Too little diversity makes the model tie facts to the few contexts where they appeared, so it fails to use them in new contexts later at test time. More diversity helps generalization in the long run, but can slow learning early in training, making reliable learning more expensive. By testing different parts of the model, we also find that different components contribute to factual and statistical generalization in different ways. Overall, our work shows how a small controlled framework can reveal learning behaviors that are hard and expensive to isolate in large language models and can guide future studies for more realistic settings.