LLM Circuit Analyses Are Consistent Across Training and Scale

Summary

This paper examines how a few common circuits (IOI, greater than, verb agreement, gendered pronoun prediction) develop at different model scales and timings in training. The main findings are that these circuits develop at the same time across different model scales, that once they develop they do not disappear, and that although individual components of the circuits can change once developing the overall circuit structure mostly remains the same once it develops.

Strengths

- The "sharing" behavior, where one attention head stops being suited for a particular purpose and another one starts being used for that purpose during training, is extremely interesting and to my knowledge a novel finding. - The authors are the first to my knowledge to study how the structures of circuits change over training. - The authors develop novel metrics that allow them to test if a given model implements a certain type of circuit. - The authors contribute to universality by showing that previously studied circuits occurs in Pythia models of many scales.

Weaknesses

- Only a few simple, well known circuits are analyzed. Significantly, the methods introduced are not scalable to large datasets of auto-discovered circuits. - The motivation of fine-tuning/continual training is not quite applicable, since the work in this paper studies only pretraining checkpoints (which is fundamentally different from something like RLHF or fine-tuning on a narrow distribution). - As the authors acknowledge, some results are not novel and merely reproduce earlier work (e.g. when induction heads arise).

Questions

- Maybe I missed it, but what components are the circuits over? Are they MLPs + attention heads? Stating this very clearly could help, as well as pictures of the circuits in the appendix. - It would be good to also have a graph showing Jaccard similarity of the current circuit vs. the final circuit, as it is hard to tell if the circuit is slowly drifting across nodes or if it stays mostly the same with a few heads switching back and forth. - I'm confused about section 4.2. Are you rerunning the circuit discover algorithm at each checkpoint to get candidate heads? Or are you just running on all possible heads? How do you know the circuit uses the head if it’s the latter? I would also like to see Figure 4 at all tokens, including before the heads develop; since the lines are just horizontal it’s not clear to me that such heads don’t always exist, it'd be great to see a steep increase when the circuit develops. - Line 162 says we're going to learn why bigger models do worse (an interesting finding in and of itself!) but I didn't see where this was touched on. - One large claim (line 194) is that good performance on the circuits emerges at the same time as the heads emerge, but as these are on two different plots (figure 1 and figure 2) with different axes, this is hard to tell. Perhaps they could both be plotted on the same axes in the appendix with different line patterns to tell the different colors apart? - Nit: in Figure 1, is there a reason some of it is logit diff and some of it is prob diff? - The authors should ideally cite The Quantization Model of Neural Scaling (the hypothesis that language models learn skills ("quanta") in order of importance), since their work supports the Quanta hypothesis in a very interesting way (the circuits arise at the same time across models). - For weighted jaccard similarity, it would be good to “see” attention heads switching. Maybe you can use the weight of edges in the DAG to assign nodes importance and do some sort of importance-weighted jaccard? As is it is hard to tell what’s changing in the circuits, much of the variance in the graph could just be from the size of the circuit changing as opposed to the more interesting "sharing". I would be happy to raise my score if some of these questions/concerns were addressed.

Rating

5

Confidence

4

Soundness

4

Presentation

3

Contribution

3

Limitations

Yes.

Summary

This paper studies the stability of circuits in language models during pre-training and across different model sizes. Specifically, it examines the Pythia model suite and a selection of four simple tasks with known circuits: indirect object identification (IOI), subject-verb agreement, greater-than, and gendered pronouns. The analysis involves multiple steps, each repeated across different model sizes (70m - 2.8b) and pre-training checkpoints: 1. Evaluate the ability of the language model to solve each of these tasks. 2. Evaluate the emergence of specific attention heads, which previous work has established as important for each of these tasks. 3. Narrow down on the IOI task to evaluate whether the three core algorithmic steps, again known from previous work, are consistent across training and scale. To this end, the authors use path patching to ablate the connection between the components involved in each of these steps. 4. Evaluate the consistency of the circuits during training by looking at component and edge overlap. To identify these circuits, the authors use edge attribution patching. Overall, the results reveal a significant level of consistency of circuits during training and across scale: important attention heads tend to start emerging after roughly the same number of tokens, although at different paces; the effect of important heads in the IOI circuit is somewhat consistent once these components emerge; and there is significant circuit component overlap across checkpoints.

Strengths

- The paper investigates an interesting question: to what extent might existing circuit analysis results generalise across training and scale? Understanding the training dynamics of circuits is important for the field of (mechanistic) interpretability. - The authors employ a variety of analysis techniques, including path patching and edge attribution patching, to systematically and causally test their hypotheses. - Despite focusing on a narrow set of tasks, the authors identify potentially general motifs, such as “load balancing”. These insights open up avenues for future work to develop a more fundamental understanding of the stability of circuits across training and scale.

Weaknesses

- While the results demonstrate a significant degree of consistency and stability of circuits across training, the focus on a small number of simple tasks provides limited insight into whether circuits for other tasks are consistent as well. - The results reveal several training dynamics that are left unexplored. For example, load balancing or the observation that many attention heads emerge after roughly the same number could have been explored in more detail. I believe a detailed investigation of one of these phenomenons could have significantly improved the paper. Minor issues: - The title of the paper (“Stability and Generalizability of Language Model Mechanisms Across Training and Scale”) does not match the title in OpenReview (“LLM Circuit Analyses Are Consistent Across Training and Scale”). - Typo in L67: “… we can verify that they are …”

Questions

1. In Figure 5, you focused on the Jaccard similarity with *all* previous checkpoints. This makes it hard to evaluate whether circuits had phases of high consistency with only the previous checkpoint. Did you observe any periods of significant circuit stability during training? This would be interesting to see as the authors of latent state models of training dynamics [1] suggest phase transitions between different algorithmic solutions. 2. You mention that “Circuits in larger models require more components, with circuit sizes positively correlating with model scale” in your contributions, but I don’t see this discussed in the following sections. Have you been able to study how these circuits differ across scale? Do we find duplications of the same components? How does this relate to the higher stability of circuits in larger models? [1] M. Y. Hu, A. Chen, N. Saphra, and K. Cho, ‘Latent State Models of Training Dynamics’, arXiv [cs.LG]. 2024.

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

2

Limitations

The selection of tasks limits the generalisation of their findings, as previous studies suggested that various behaviours of language models emerge at scale [2] or qualitatively change across scale [3]. Both of these suggest that circuits still fundamentally change for more complex tasks. However, I believe that this is mostly addressed in the limitations section of the paper. [2] J. Wei et al., ‘Emergent Abilities of Large Language Models’, arXiv [cs.CL]. 2022. [3] J. Wei et al., ‘Larger Language Models Do In-Context Learning Differently’, arXiv [cs.CL]. 2023.

Summary

This study explores the emergence and evolution of internal mechanisms in language models of varied sizes during the training process. Specifically, it examines simple tasks such as IOI, Greater-than, Gender Pronoun, and Subject-verb agreement using Pythia models. The findings indicate that models of different sizes tend to learn these tasks after a similar number of token counts during training. Moreover, while individual components of the models may change in functionality, the overall algorithm implemented by the models remains consistent throughout the training process. Lastly, the study identifies that once a circuit emerges, it generally remains stable thereafter.

Strengths

1. This work is highly relevant for two key reasons: - Most mechanistic interpretability studies do not examine the internal mechanisms of models throughout the training and fine-tuning processes. As a result, they fail to offer a comprehensive understanding of how and when a model learns its mechanisms and how these mechanisms evolve after their creation. - Unlike many existing studies, this work analyzes multiple tasks across various models of different sizes. 2. The results from Section 4 indicate that, although some components of the model change their functionality, the overall algorithm used by the model to solve simple tasks, such as IOI, remains consistent. This finding aligns with previous results from [1], which state that even though fine-tuned models have larger circuits, their mechanisms remain consistent, even for more complex tasks like entity tracking. 3. The presentation is clear, but it does require the reader to have some background in mechanistic interpretability. Additionally, there are a few other pertinent works on mechanistically understanding the impact of fine-tuning that the authors should cite [2, 3]. [1] Prakash et al, "Fine-Tuning Enhances Existing Mechanisms: A Case Study on Entity Tracking", 2024. [2] Jain et al, “Mechanistically analyzing the effects of fine-tuning on procedurally defined tasks”, 2023. [3] Lee et al, “A Mechanistic Understanding of Alignment Algorithms: A Case Study on DPO and Toxicity”, 2024.

Weaknesses

1. I’m unsure about the results of section 3.2 stating that emergence of circuit components that are involved in the internal computation of the tasks at the same time as the model behavior performance, suggests that the former is responsible for the emergence of latter, because of the following reasons: - All the model heads are evaluated, rather than only those involved in the circuit performing the task. It is possible that some heads exhibit certain behaviors without actually being part of the circuit. Thus, concluding that their occurrence is responsible for the model's performance could be misleading. - Furthermore, the analysis is primarily conducted for a single task, IOI, rather than all four tasks mentioned earlier. Although some of the heads studied are involved in the Greater-than tasks, MLP neurons, which have been shown to be part of the circuit, are not analyzed. 2. While it is interesting to note that individual circuit components emerge simultaneously with the model's behavior, it still does not explain why these components emerge after similar token counts in models of varied scales. 3. Section 4 investigates only the IOI task, which raises concerns regarding the generalizability of its results.

Questions

1. The finding that models of varied sizes learn a task after a similar number of tokens is intriguing and somewhat counterintuitive. This makes me wonder if analyzing how gradient updates modify model weights and comparing these changes across models could provide a better understanding of how language models learn (potentially in future works). 2. Section 3.2 states that to validate the importance of four mentioned types of attention heads, a circuit is discovered for each model at each checkpoint. I’m unsure of what a circuit discovery algorithm will discover for early checkpoints where the model does not even have behavioral ability to perform the task. So, it’s surprising to me that authors were still able to identify circuits and functionality of their components which is consistent with existing discovered circuits in the literature. I would like to see evaluation results of these circuits. 3. Section 2.2 mentions that there is no definitive method for verifying the entirety of the identified circuit. This needs more explanation, particularly regarding why metrics like completeness proposed in [4] are considered inadequate for this purpose. [4] Wang et al, “Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small”, 2022.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

4

Limitations

The primary limitation of this work is the use of simple tasks such as IOI, Greater-than, Subject-verb, and Gender Pronoun for analysis. Some of the results may not be applicable to more complex tasks. Additionally, relying solely on Pythia model suits trained using the same training data poses a risk to the generalizability of the findings. However, despite these limitations, the results are insightful and should be valuable for the mechanistic interpretability community and beyond.

Summary

This paper presents a set of analyses on the dynamics with which internal language models’ mechanisms emerge and change during training. The four mechanisms studied are internal circuits that the model implements to carry out four simple tasks: indirect object identification (IOI), gendered pronoun, greater-than, and subject-verb agreement. These circuits have been identified by previous work and shown to be implemented by a set of specialized attention heads that display an interpretable behavior (e.g., induction heads, which attend to a previous occurrence of a substring that is repeated in the input). The objects of the study are 7 models of the Pythia family, with sizes ranging from 70M to 12B parameters. The authors identify the components that implement each circuit using a technique called edge attribution patching with integrated gradients. In the first analysis, the authors show that the ability of the models to carry out the four tasks emerges roughly at the same point during training (in terms of # of tokens seen by the model). At the same point during training, a subset of the models’ attention heads are observed to start implementing the specialized behaviors used to implement the tasks studied, leading to the conclusion that these specialized heads are responsible for the emergence of the model performance. In a second analysis, the authors study the dynamics which is how different sets of heads in Pythia 160M display the specialized behaviors relative to the four tasks considered. One observation is that the functional behavior of some heads decreases over time. This is surprising as this decrease is not reflected by the model’s performance. Copy suppression might be a possible explanation for this phenomenon. Additionally, the paper includes an analysis of the stability of the algorithm implemented by the Pythia models on IOI, suggesting that the algorithm does not undergo significant changes after emerging during training. Finally, the paper analyzes the changes in the actual components constituting the four circuits in Pythia 70M-2.8B.

Strengths

- The analysis of circuit dynamics during language model pre-training is novel. - The results about the emergence (and disappearance) of specialized heads, and about the models’ algorithmic stability during training are informative and are likely to be appreciated by the (mechanistic) interpretability community.

Weaknesses

- Parts of the paper might benefit from additional details and clarification (see questions 1.1-1.3). - The conclusions to be drawn from some of the results are not completely clear (see question 2).

Questions

1. In Section 4.2, the authors measure, through path patching, the effect of a node in the computational graph in Figure 4A, on the model’s accuracy when performing IOI. 1. As a reader, I would appreciate some additional details about the experimental procedure here: what are the exact heads that are being ablated? How are they selected for each model (e.g., the 5 heads with the highest CPSA score)? 1. It is not clear to me why this effect is termed “direct.” As illustrated in Appendix B, path patching involves intervening on a node H (e.g., S-Inhibition heads), measuring the change produced by the intervention on a second node R that depends on H (e.g., NMH), and finally how the change in R an outcome variable Y (in this case, the model’s output). Through this procedure, what is being measured is the *indirect* effect of H on Y that is mediated by R. 1. “For each step, our metric measures this direct effect, divided by the sum of the direct effects of ablating each edge with the same endpoint” (lines 254-255). I am confused about the denominator in this operation: what are the effects that are being summed here? Also, do the results look similar when unnormalized (i.e., when measuring the absolute effect)? 1. I am not sure whether there’s any clear conclusion that we can draw from the results presented in Section 5. It is true that the EWMA-JS for Pythia 70M seems to have a higher variance, but besides this, I struggle to see a clear trend in the measurements for the other models. 1. If you quantified the cumulative score for each type of head (e.g., induction score) over time for the top k heads, would it be constant after the point during training at which the model learns the task? This would confirm your hypothesis for which specialized heads are replaced by other heads as they lose their functional behavior during training (although my guess would be that such cumulative scores might be decreasing over time).

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Limitations seem to be adequately discussed.