How do language models learn facts? Dynamics, curricula and hallucinations — part 3/5

Cleaned-text code points [31137, 46983). Bibliography omitted here; source.md includes it.

This analysis reveals that, in a simple but relevant setting, attention-based recall circuits emerge before the formation of associative memories in feedforward layers. We hypothesize that this occurs as established circuits amplify the correlation between the inputs and backpropagated errors received by feedforward layers. Prior to circuit formation, task learning remains possible (e.g., by artificially increasing name token values), but progress is slow, performance plateaus, and generalization likely suffers. Phenomena like grokking ( Power et al., 2022 ; Nanda et al., 2023a ) may indicate that learning dynamics initially find this shortcut, before abandoning it due to regularization. The proposed credit-assignment argument for plateau formation generalizes to other architectures, as evidenced by ablations in which attention is replaced with recurrence (Figure C ). Decoupling the token-mixing role of attention from other computations, as advocated by Elhage et al. (2021) for mechanistic understanding of trained Transformers, also proves powerful for analyzing learning dynamics. This perspective offers a promising avenue for future investigations into neural network learning dynamics.

##### On the role of non-uniformity and connections to developmental psychology.

A key finding of this work is the precise analysis of how imbalances in the training distribution enable faster escape from performance plateaus. While not explicitly connected to the formation of attention patterns, similar phenomena have been observed in diverse settings ( Charton and Kempe, 2024 ; Park et al., 2024 ) , though sometimes at the cost of limited generalization. Our data scheduling results, by adapting the data to the network’s current learning phase, offer a promising path towards accelerated learning without sacrificing generalization. Intriguingly, these strategies mirror the implicit curriculum experienced by developing infants ( Smith et al., 2018 ) , arising from factors like limited mobility or frequent exposure to familiar faces, and similar strategies have proven successful in reinforcement learning ( Baranes and Oudeyer, 2013 ; Stooke et al., 2021 ) . As argued earlier in the discussion, repetition strengthens the signal-to-noise ratio in correlations between events (tokens in our case), facilitating faster identification of key relationships and accelerating learning. Crucially, however, (progressive) diversification remains essential for optimal generalization in later stages of learning. Our findings could therefore serve as a building block for a statistical theory of development.

### Limitations

In our study, we train relatively small language models on a synthetic factual recall task to investigate knowledge acquisition dynamics. While this approach enables precise experimental control and mechanistic insights, it presents several important limitations.

##### Scale and architecture.

The 44 44 M parameter models we train are substantially smaller than the billion-parameter models used in practice. The learning dynamics we observe may manifest differently at larger scales or with different architectural choices. However, our ablations across model sizes, architectures (including recurrent variants), and hyperparameters suggest the core phenomena are robust to these variations.

##### Data.

The synthetic factual recall task, while designed to capture essential properties of knowledge acquisition, only represents one specific type of knowledge language models acquire during pretraining. Different types of knowledge could interact in sophisticated ways that are not modeled in our task. This may create extra redundancy that language models could leverage to accelerate learning. Nevertheless, our fine-tuning results already align with practically observed inefficiencies of fine-tuning for integrating new knowledge, suggesting some degree of transferability.

Despite these limitations, we believe our work provides a necessary foundation for understanding knowledge acquisition in neural language models. The mechanistic insights we derive offer concrete hypotheses that future work can validate in more realistic scenarios. We hope these findings will guide the design of more effective training strategies and data scheduling approaches for large-scale language models.

### Acknowledgements

The authors thank Nino Scherrer, Katja Filippova, Carter Blum, Jay McClelland and Arslan Chaudhry for insightful discussions.

\nobibliography *

## Appendix

### Appendix A Related work

#### A.1 Associative memories and factual knowledge of neural networks

Associative memories have been extensively studied in neuroscience. The foundational experimental work on conditioning by Pavlov (1927) sparked extensive theoretical research into the computational mechanism underlying associative memories. This led to breakthrough developments like Hebb’s rule ( Hebb, 1949 ) and Hopfield networks ( Amari, 1972 ; Hopfield, 1982 ) , which later proved instrumental in the emergence of deep learning.

In recent years, as language models have grown increasingly powerful, research has begun examining them through the lens of associative memories. Our work directly builds on this line of research. Petroni et al. (2019) first proposed that masked language models like BERT ( Devlin et al., 2018 ) , could encode relational knowledge within their weights. This analysis was subsequently extended to autoregressive Transformer-based language models: Geva et al. (2021) demonstrated that feed-forward layers act as associative key-value memories, Meng et al. (2022) showed how this stored knowledge could be located and edited, Geva et al. (2023) ; Nanda et al. (2023b) revealed the functional mechanisms underlying memory recall, and Allen-Zhu and Li (2023) identified conditions under which these general mechanisms are developed. Beyond these mechanistic insights, associative memories have provided a framework for evaluating the knowledge capacity of language models. This includes empirical scaling laws ( Allen-Zhu and Li, 2024 ) and theoretical analyses ( Nichani et al., 2024 ) which follow from a long history of capacity analyses of Hopfield networks and related models ( McEliece et al., 1987 , e.g.,) .

While the above research primarily examines parametric memories stored in network weights, recent years have also seen growing interest in the in-context associative recall capabilities of sequential modeling networks. Notable examples include the induction head in Transformers Olsson et al. (2022) and comparative studies showing that the primary difference between attention-based and attention-free language models can be attributed to their different in-context associative recall capabilities ( Arora et al., 2024 ) .

#### A.2 Learning dynamics of neural networks

The study of neural network learning dynamics has deep roots in early connectionist research ( Hinton, 1986 ; McClelland, 1995 ; Baldi and Hornik, 1995 , e.g.,) . It has been particularly influential in the early days of deep learning. For example, the seminal work of Saxe et al. (2014) provided fundamental insights into the role of depth in neural network training.

The recent demonstration of in-context learning capabilities in large-scale models Brown et al. (2020) has sparked renewed interest in understanding the underlying dynamics. A significant contribution to this understanding came from Olsson et al. (2022) , who established a causal link between the development of induction heads and a phase transition that enables in-context learning. This mechanistic understanding has been complemented by research on the influence of training data distribution on learning trajectories and implemented algorithms. Notably, Chan et al. (2022) and Singh et al. (2024) revealed how different data distributions guide networks through distinct algorithmic solutions, either facilitating in-context learning or not. The generality of these findings was later confirmed by Park et al. (2024) , who reproduced them using a simplified training distribution.

Parallel to these developments, Nguyen (2024) showed that Transformer-based language models’ predictions can be effectively approximated by N − N- gram statistics, with the optimal N N increasing throughout training. Our work observes similar dynamics during the transition from bigram to trigram predictions during the plateau phase. However, a key distinction lies in the level of abstraction: while Nguyen (2024) ’s N − N- gram associations occur directly at the token level, our work observes these sequential dependencies at a higher level of abstraction.

Finally, and most related to ours, is the study of factual knowledge acquisition in large language models from Chang et al. (2024) . Their key finding is that factual knowledge acquisition occurs through accumulating small probability increases when the model encounters the same knowledge, followed by gradual forgetting when not being exposed to it. These results are consistent with ours, although we find forgetting to be more pronounced, likely because of larger training distribution shifts. Additionally, while their findings were obtained by slightly altering the training distribution of a large language model, our work goes one step further by studying in depth the impact of training distributions on learning speed.

### Appendix B Experimental setup

#### B.1 Rationale behind our design choices

Our goal is to identify a synthetic setting, as simple as possible, where a model exhibits knowledge of its training data (that is a flexible usage of it, cf. Section 1.1 , not mere memorization) using mechanisms similar to large language models. This allows precise control over the data distribution and increases the likelihood that our findings generalize to large language models.

We build upon the synthetic biography dataset and analysis of Allen-Zhu and Li (2023) , as well as the mechanistic interpretability studies of Geva et al. (2023) and Nanda et al. (2023b) . Geva et al. (2023) and Nanda et al. (2023b) find that pretrained large language models encode individual-specific information in their residual stream as soon as the name of the individual is encountered. Allen-Zhu and Li (2023) observe such a behavior when training models on biographies whose order is permuted every new occurrence, along with having multiple biographies for “celebrities” (see their Q-probing analysis). We hypothesize that textual variation within the biography distribution is key to achieve this behavior. Therefore, our setup ensures unique biographies by introducing 25 distinct templates per attribute type (see next section) and permuting their order. Importantly, biographies are resampled for each new sequence. Figure A demonstrates the critical importance of random ordering and the need for some textual diversity within individual biographies. While we have not ablated the total number of templates per attribute type, we believe that the overall textual diversity of the distribution is necessary for the learning of robust features. In our experiments, the individual’s name is repeated in every sentence to facilitate knowledge development, though we believe later occurrences could be replaced with pronouns without significantly affecting results, as ( Allen-Zhu and Li, 2023 ) found the benefits of full name repetition primarily in the non-permuted and not so diverse regime.

#### B.2 Details about the biography generation process

Our data generation process largely follows Allen-Zhu and Li (2023) , with a few important differences regarding how templates are generated and manipulated.

Prior to training, we generate a population of individuals, each with a full name (first, middle, and last) and six attributes: birth place, birth date, university, major, company, and current location. These are generated as follows: • Full name. First and middle names are sampled from a list of 900, and last names from a list of 1,000, resulting in 810 million potential unique names. We ensure that each individual has a unique full name. These names are among the most commonly used worldwide.

• Birth place and current location. Sampled from a list of the 800 largest cities worldwide. Unlike Allen-Zhu and Li (2023) , we do not correlate current location with the company.

• University. Sampled from the 200 largest universities worldwide.

• Major. Sampled from a list of 100 majors.

• Company. Sampled from the 500 companies with largest valuations.

• Birth date. Year, month, and day are sampled independently (between 1900-2100, January-December, and 1-31, respectively). This allows unrealistic dates (e.g., February 31st), but should not significantly impact tokenization.

These values are chosen to reasonably approximate the token distribution a large language model might encounter during pre-training.

For the sake of our analysis, the template generation requires extra care, as we want all the information needed to predict the attribute value and as we want to evaluate the model on sentences it has never seen. Such considerations were not needed in Allen-Zhu and Li (2023) , so we had to adapt their setup. In our dataset, templates are generated prior to training with the assistance of a large language model, using the prompting scheme in Figure B . This yields 25 25 distinct templates per attribute type. As discussed in Section 1.2 , we (manually) ensure that the individual’s name and information identifying the attribute type precede the attribute value, allowing a model with perfect knowledge to achieve zero loss. For each individual, we pick 20 20 templates for training and keep 5 5 for evaluation, ensuring the model encounters novel template-individual combinations during evaluation, thus measuring knowledge rather than memorization. Additionally, we introduce special tokens for tags (like name or birth date) in the templates and replace these tags by the desired content only after tokenization. This way, we can make sure that, after this manipulation, the tokens coming from names or attribute values directly appear in the token sequence (e.g. "1990." will be tokenized as "[1990][.]" and not "[1990.]" as it would usually be). Such a precaution facilitates analysis as we can be sure that each of the name or attribute value tokens only contain this information, and not some unrelated semantic information about the sentence.

Biography generation, for both training and evaluation, is summarized in Figure 1 and involves two steps: 1. Individual sampling. By default, we sample individuals uniformly at random from the population. This distribution is modified in some experiments (Section 3.1 ) and can be time-step dependent (Section 1.1 and Appendix F.5 ). Here are the detail of these different distributions: • Inverse power law / Zipf distribution. The i i -th individual is sampled with probability proportional to i − α i^{-\alpha} with α \alpha an hyperparameter. Recall that we get the uniform distribution when α = 0 \alpha=0 and, when α = 1 \alpha=1 , we get the Zipf distribution.

• “Celebrities” distribution. A subset of the individuals of a size n_celebrities is oversampled is sampled weight_celebrities times more frequently larger than individuals outside the group. Here, n_celebrities and weight_celebrities are the two hyperparameters we vary.

• “Warm-up” distribution. Training begins on a subset of indiv_warmup individuals for epochs_warmup epochs. An epoch is here defined as the number of training steps required to see as many biographies as there are individuals in the entire population (number of individuals divided by batch size steps). Once the warm-up is over, we train on the full population. In both phases, individuals within the relevant group are uniformly sampled.