LLM Circuit Analyses Are Consistent Across Training and Scale

Thanks for your thorough review! We respond below, abbreviating your comments for space reasons.

Weaknesses

the methods introduced are not scalable to large datasets of auto-discovered circuits

We disagree in part about the scalability of our methods. EAP-IG is efficient, and can now be used with models around 7B parameters in size (unfortunately, this was not the case when writing the paper). The fact that small models are increasingly capable means that studying e.g. OLMo-1B or 7B could show us how models learn even rather complex tasks. However, we do agree that studying 156 checkpoints for each model was compute-intensive, and some parts (copy suppression evaluation) scale poorly.

the motivation of fine-tuning/continual training is not quite applicable

We agree that narrow fine-tuning doesn’t not quite fit as a motivation, and will revise this. However, continuous training is often done on a wider distribution of text (e.g., a large variety of user queries) and there may also be cases where researchers would want to look at circuits at intermediate checkpoints of a model before it has finished training.

some results are not novel and merely reproduce earlier work

Although our induction head findings do reproduce findings in Olsson et al., we primarily mention them in passing; the bulk of our paper consists of novel findings: other component types, performance emergence, algorithmic stability. We will try to focus on the more novel parts of our work.

Questions

what components are the circuits over?

The circuits are composed of attention heads and MLPs (plus the inputs and logits). We will clarify this!

It would be good to also have a graph showing Jaccard similarity of the current circuit vs. the final circuit

Thanks for the suggestion, implemented in Figures 2/3 of the response PDF. We see that in Figures 2/3 there are still many fluctuations compared with the final circuits which indicates that component swapping occurs during training. We can also observe a rising trend, indicating that the circuit grows increasingly similar to the final circuit during training. Therefore, the model both slowly drifts towards the final circuit, even as components swap during training.

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.

Apologies for the confusion in Section 4.2! To clarify, there are two criteria for including a head in a group: whether it meets a component score threshold, and whether it is important to model performance as measured by the path patching causal intervention. The latter intervention is roughly what EAP-IG approximates, so it’s actually a stronger test than checking circuit membership; see our response to dE5s for more details. Given this procedure, we start Figure 4 at the 30B token mark (10% of training) because this is where name-movers (NMHs) first appear as part of the IOI circuits. S2-inhibition and other heads are found via their impact on NMHs, so without NMHs, there’s nothing to measure / plot. Before the 30B token mark, most IOI circuits in these models consist simply of a single copy suppression head. Gradual emergence of relevant component types is better seen in Figures 2 and 3 of the main text.

Line 162 says we're going to learn why bigger models do worse

Our apologies for the confusion! Large models do not learn any faster than smaller models do (above a certain threshold), because such models do not develop the responsible heads any faster than smaller models do. We don’t say that explicitly right now, but we will add that to the paper.

Perhaps they could both be plotted on the same axes in the appendix with different line patterns to tell the different colors apart?

We agree that plotting things in one plot, with the same axis would help, but we didn’t have the space in our response PDF. However, we plotted the sum of component effects for in-circuit heads using the same x-axis as the behavioral plot. In this new plot, you can see clearly how greater-than components arise at or just before task behavior emerges, at 2x$10^9$ tokens. For IOI, copy suppression, the initial part of the circuit, arises at 4x$10^9$ tokens, while the name movers form a bit later, at 7 to 8x$10^9$ tokens. We will include the plot that you suggest in the appendix.

in Figure 1, is there a reason some of it is logit diff and some of it is prob diff?

For tasks with one in/correct answer, we use logit diff; for tasks with multiple, we use prob diff, aligning with the original metrics for these tasks. Logit diff doesn’t quite work for multi-answer tasks; we can’t sum over the logits of all in/correct answers as with prob diff. We could use prob diff for all tasks, but some argue against this (Heimersheim and Nanda, 2024).

The authors should ideally cite The Quantization Model of Neural Scaling

Thank you for pointing out that connection! We’ll add that citation.

Maybe you can use the weight of edges in the DAG to assign nodes importance and do some sort of importance-weighted jaccard?

We have included related plots in the response PDF (Figures 2 and 3). For now, instead of plotting importance-weighted Jaccard similarity (JS) for node sets, we do the same for edge sets, as EAP-IG gives us only edge importance; we think computing node importance from edge importance would give similar results. We compute the Jaccard similarity between intermediate checkpoints and final circuit. We can see that the JS slowly climbs and remains high at the end. It aligns with the conclusion that the model both slowly drifts towards the final circuit along with components swapping during training. To see specific attention head switching, we refer to Figure 3 (main text).

Thanks for your helpful review! We agree that many of the issues you bring up are important, and have attempted to address them below.

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.

We agree that it remains to be seen whether more complex circuits, or a broader set of circuits in general show the same stability over training and scale. Currently, however, only a very small set of circuits have yet been identified, and finding new circuits was somewhat beyond the scope of our paper. We hope that our work will be useful to follow-on investigations of future circuits and their potential for consistency across these dimensions. We do not claim that all circuits will retain this consistency, and it remains to be seen what happens in much, much larger models; however, we do suggest that this provides some evidence that circuits found at specific points and model sizes can provide information that holds to at least some degree in across these dimensions.

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 phenomena could have significantly improved the paper.

We also agree that investigating load balancing or token-dependent formation of components would have added to the paper; however, we felt such investigations likely would deserve their own dedicated projects (likely involving significant amounts of model training to get more granular checkpoints) and research output, and as such decided the scope of this paper would be best limited to an initial investigation of a wide set of phenomena. Those two topics are certainly worth investigation, however, and we hope follow-up work will address them!

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 …”

Thanks for catching these! We will fix them.

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.

To provide more detail the question on Jaccard similarity and circuit stability, we have added a set of graphs (Figures 2 and 3) that show circuit Jaccard similarity at each checkpoint to the final circuit. We do observe periods of stability in some models and tasks, but in other cases we see gradual or spiky transitions towards the final circuit. We also note that we conducted EAP-IG on a number of seed variations of the Pythia models, finding a variety of patterns; in some models, the circuits were quite stable for long periods, while others showing oscillating sharp changes between checkpoints. Unfortunately, we didn’t have room to include these plots, but we could include them in the appendix.

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?

In Section 5, we compute the Pearson correlation between circuit sizes and model sizes. The Pearson correlation is r = 0.72 for IOI and SVA, 0.9 for Greater-Than, 0.6 for Gender Pronoun.

We do find some differences in circuits across model scales. As we cover in Section 4, the algorithmic structure remains similar, but as shown in Section 5, the circuits are larger in larger models. In the case of the IOI task, some components demonstrate more replication than others; for example, copy suppression heads and S-inhibition heads remain few in number across circuits in different model scales, while induction heads and name-mover heads increase significantly. On balance, less-volatile head types seem to increase the most in number, which does indeed support your suggestion that this is related to circuit stability in larger models. Regardless, though, all identified components seem to be present in all models, and seemed to be similarly important to their circuits for IOI.

Thanks for your review and helpful suggestions! We respond to your questions (which contain your stated weaknesses), below.

1.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)?

Here are additional details about our algorithmic consistency experiment, which we will add to the appendix. At each checkpoint, heads were selected for ablation as follows:

1. We ablated all heads upstream of a target or set of targets (e.g., the final logits or the set of name-mover heads) one at a time via path patching to determine their effect on the logit difference metric (either directly or through an intermediate node); this mirrors Wang et al. (2023).

2. We tested each of these heads for their target function (e.g., copy score for name-mover heads). We tested the set of heads that both A. had a component function score over a specific threshold (10%) and B. had a negative effect on logit difference when patched. For S2-inhibition heads, we tested whether ablating positional signal while keeping the token signal A. reduced logit difference, B. reduced NMH attention to the IO, and C. increased NMH attention to S1. See appendix D for the component tests for induction heads and duplicate token heads. All of these tests also come from Wang et al. (2023)

1.2 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.

You are right to question our use of the term direct in this paragraph. Upon re-reading this, we see that we inadvertently overloaded the term and should not have used “direct” here. The procedure is as you described.

1.3 “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)?

Taking the notation you used in the previous question, if we take the numerator as the effect of ablating heads H (e.g., all S-Inhibition heads) on Y through intermediate heads R (e.g. NHMs), the denominator is the effect on Y of ablating heads G (all heads upstream of R, and includes H as a subset) through intermediate heads R.

2. 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.

We apologize for the lack of clarity here and agree that this section might have been better combined with Section 4. We aimed to quantify the shifting set of nodes that comprised each task circuit, and to illustrate that these shifts are not model-dependent. In most cases, shifts are gradual—Jaccard similarity is often between 0.6 and 0.9—but nevertheless real. We have added plots in our response PDF (Figures 1 and 2) that illustrate similarity to the final circuit. These show gradual changes result in a significant shift between the circuits at the beginning and end of training, despite our algorithmic consistency results.

3. 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).

This is a good suggestion, and one that we incorporated into Figure 3 of our response PDF. The figure now shows the sum of head effects across all in-circuit heads at a given checkpoint. Some trends remain the same as previously: the timestep at which components emerge is still tightly coupled to the emergence of task behavior. However, the degree to which the sum is constant over time varies across components; we believe that this reflects not only variation in the total component effects, but how much sense it makes to sum each component’s score.

Thanks for your insightful review! These are good points, which we answer below.

Weaknesses

1.1 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.

Indeed, this could be a concern. In our PDF response, we plotted the sum of all component effects from all heads in the circuit (Figure 1). We see that the emergence of these components comes right as (or just before) the emergence of task behavior as well, indicating that our previous results are not due to the inclusion of non-circuit heads.

1.2 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.

It’s true that Hanna et al. (2023) study MLPs and their neurons. However, unlike the induction and successor heads on which we focus, these aren’t known to perform a generalizable function that is reused across tasks. We could track the MLPs that connect to the logits in the circuit, and which neurons therein matter most, but we wouldn’t gain any insights about how higher-level model abilities develop. For these reasons, we excluded them.

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.

We agree that the phenomenon of components emerging at similar token counts across model scales deserves more exploration. A deeper look into this phenomenon seemed beyond the scope of our paper, but we hope it will be the subject of future work. We hypothesize that the formation of the other head types may occur in a phase change similar to the induction head phase change observed in Olsson et al. (2022), though model retraining for the purpose of obtaining more granular checkpoints would be needed to look in further detail at what is happening as these heads develop. This by itself would not explain why this happens, however, and we hope to see further investigations with additional approaches.

3. Section 4 investigates only the IOI task, which raises concerns regarding the generalizability of its results.

It is true that running a similar analysis on a broader set of circuits would strengthen our claim of algorithmic consistency. Unfortunately, few circuits with a clear and quantifiable algorithm beyond IOI have been identified to-date; even quantifying the relatively well-characterized Greater-Than is challenging. For this reason, we used IOI as the subject of our analysis. We do not claim that the property of generalizability over time and scale will apply to all circuits, especially more complex ones, but we hope this is investigated further.

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).

We have also hypothesized that there may well be a connection between the number of gradient updates and the magnitude of changes to model weights required for models to learn a task. This isn’t quite in the scope of the current work, but we agree this is an interesting question.

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.

This is true—we run the circuit at all checkpoints, but at early checkpoints, there’s little model behavior to localize. The circuits discovered before model behavior emerges are not very meaningful; though they attain 80% of the model’s original performance, that performance is near 0. These pre-performance “circuits” consist of fairly stochastically-changing, hard-to-interpret sets of components that contribute only to the extremely small positive and negative variances in the task performance metrics. Due to this, we don’t base our conclusions on these early circuits; we will stress this in the text.

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.

Wang et al. (2023) indeed propose a completeness metric, which involves comparing model and circuit performance while ablating a set of components from both. However, this has flaws. For one, it’s challenging to choose components; Wang et al. propose three methods for doing this, which yield different results, and are not all compatible with our needs. More importantly, these ablations are computationally expensive, and not feasible to perform over many checkpoints. As a result, this metric has not been widely adopted, and we find most of the proposed alternatives (e.g. faithfulness of the circuit's complement) inadequate. We can add some discussion of this issue.