Pretraining Data Statistics Shape the Phases of Learning Entity Comparison in Language Models

How does data shape language model (LM) behavior throughout pretraining? We investigate this question through a case study on entity comparison, e.g., Between France and Brazil, which country is larger?. We begin with controlled experiments in which we train small LMs (124M parameters) on mixtures of natural text from pretraining corpora and synthetic data from entity comparison tasks. We identify three distinct phases of learning: (1) an early phase where the LM selects entities by frequency, (2) a middle phase where the LM selects entities by position in a prompt (first vs. last), and (3) a late phase where the LM selects the entity that is the correct answer to the question. We show that the emergence of these three phases is controlled by statistical properties of the training data. With small amounts of task-specific synthetic data, we observe only the first two phases and the model fails to learn the task; with large amounts, the model jumps directly from the frequency-based heuristic to solving the task correctly. Moreover, if we modify the frequency of entities in data from a naturally occurring Zipfian distribution (a small number of entities are very common and the vast majority are rare) to a uniform distribution, the first phase disappears and the model learns the task more quickly. Finally, we find the same three phases of learning in the pretraining of open-sourced OLMo models. Together, our findings demonstrate that properties of pretraining data are causal drivers of heuristic learning and show that small-scale synthetic experiments can predict training dynamics at larger scales.