Paper: Fine-Tuning Enhances Existing Mechanisms: A Case Study on Entity Tracking (ICLR-2024)

Reviewer comment (Experiments):
The methods employed in this paper are limited to a single type of model on an entity tracking dataset. Moreover, the entity tracking mechanism is likely just one of many contributing factors, raising questions about the generalizability of their claims.

Author response:
Although we have focused on the specific task of entity tracking, the methods employed and introduced in this work are generic and can be easily applied to other tasks and models, by generating relevant data. Both path patching and DCM have already been applied to other tasks like indirect object identification and variable binding in arithmetic expression (Wang et. al (2022), Davies et. al (2023)). Additionally, CMAP can also be employed among various kinds of models that share the same dimensionality.
- While we have meticulously analyzed the impact of fine-tuning on the entity tracking mechanism, we acknowledge that understanding whether such mechanism invariance is universal would require further studies involving additional tasks and models. *We have updated our discussion section to clarify this limitation.*

Reviewer comment (Novelty):
I personally do not see the significance of the entity tracking problem to be particularly high.

Author response:
We understand not all researchers might be interested in the same topics, but there is a significant community of researchers who are interested in both entity tracking and name binding: it has been a significant area of study in not just AI, but also in Linguistics and Cognitive Neuroscience (Karttunen et. al (1976), Heim et. al (1983), Nieuwland et. al (2006), Kamp et. al (2010), Marcus et. al (2018)). 
- For instance, Kamp et. al (2010) proposed multiple structures for discourse representation, with the fundamental atomic steps as binding names with their attributes defined in the context, making it crucial for understanding the discourse capabilities of LLMs. 
- Additionally, Marcus et. al (2018) argued that representing abstractions, instantiating variables with instances, and applying operations to those variables are indispensable to the human mind, which is often ignored in AI research.
- While we are among the first to study the internal mechanism for the entity tracking task, multiple researchers have previously investigated such capabilities in various deep neural networks, particularly using probing classifiers (Li et. al (2021), Li et. al (2022), Kim et. al (2023)). This interest is current and shared by others; for example, there is another anonymous submission to ICLR (https://openreview.net/forum?id=zb3b6oKO77) that examines the name-binding capabilities from another perspective, indicating the current interest within the research community in this topic.
 - *We have also updated our manuscript to highlight these points to provide context on the significance of the entity tracking task.*

Reviewer comment (Experiments):
The experiment results can not lead to the claim that the original model and its fine-tuned versions implement entity tracking with the “same” circuit: (1) the fine-tuned model, Goat-7B, reaches an accuracy of 82% while the circuit identified in Llama-7B reaches an accuracy of 68%; this considerable performance gap indicates that Goat-7B’s circuit may be different from Llama-7B’s circuit, although these two circuits may have considerable overlap. (2) It is necessary to explore the mentioned overlap. For example, could authors apply Path Patching to Vicuna-7B and Goat-7B to extract their circuits and compute the overlap between their circuits and Llama-7B’s circuit?

Author response:
**The experiment results can not lead to the claim that the original model and its fine-tuned versions implement entity tracking with the “same” circuit: (1) the fine-tuned model, Goat-7B, reaches an accuracy of 82% while the circuit identified in Llama-7B reaches an accuracy of 68%; this considerable performance gap indicates that Goat-7B’s circuit may be different from Llama-7B’s circuit, although these two circuits may have considerable overlap. (2) It is necessary to explore the mentioned overlap. For example, could authors apply Path Patching to Vicuna-7B and Goat-7B to extract their circuits and compute the overlap between their circuits and Llama-7B’s circuit?**

-   We acknowledge that while the performance of the circuit identified in Llama-7B on Goat-7B is satisfactory, its faithfulness score is slightly lower than that of the circuit on Llama-7B, suggesting that the circuit employed by Goat-7B might have a few additional components. Hence, we applied path patching (as suggested) on the Goat-7B model to recover its circuit.
    
-   We found that similar to the circuit identified in Llama-7B, the Goat-7B circuit comprises of four sets of attention heads, each exhibiting a similar information flow mechanism to that of the Llama-7B circuit. Using the minimality score, we pruned out redundant heads, making it a circuit with 76 attention heads, comparable to the previously identified Llama-7B circuit consisting of 75 heads.
    
-   We evaluate the performance of the Goat-7B circuit across models and compare its performance to the circuit found in Llama-7B (see the following table):
    

	1.  Not only the performance of the Goat and Llama circuits in Goat-7B is similar (0.70 vs. 0.68), but their faithfulness scores are also comparable (0.85 vs. 0.83), indicating that the two circuits have roughly the same performance.
    
	2.  The identified Goat circuit also performs well on other models. For instance, the Goat circuit could achieve the performance of the entire Vicuna-7B model, indicating that it is generalizable across models, similar to the Llama circuit.
    
	3.  Both Goat and Llama circuits exhibit comparable performance across various models, suggesting that identifying a circuit in one model and applying it to another can yield results as effective as identifying the circuit directly within the target model itself.

| Model      | Full-Model Accuracy | Goat Circuit Accuracy | Llama Circuit Accuracy | Random Circuit Accuracy   | Goat Circuit Faithfulness | Llama Circuit Faithfulness |
|------------|------------|--------------|---------------|-------------------|--------------|---------------|
| Llama-7B   |  0.66        |  0.57          |  0.60           |  0.005   |  0.86          |  0.91           |
| Vicuna-7B  |  0.67        |  0.67          |  0.65           |  0.0342  |  1.00          |  0.96           |
| Goat-7B    |  0.82        |  0.70          |  0.68           |  0.0044 |  0.85          |  0.83           |


-   To understand it thoroughly, we computed the overlap between the attention heads of each of the four groups of heads in both circuits and found that there is a significant overlap between the heads of each of the four groups, as shown in the following table. We observe that more than half of the attention heads are shared between both circuits. We also found that most heads with high causal impact on the final output are shared between both circuits, based on Fig. 6 (Appendix D). 

| Head Group(s)     | #Heads in Llama Circuit | #Heads in Goat Circuit | Intersection | Precision   | Recall |
|------------|------------|--------------|---------------|-------------------|--------------|
| A   |  40        |  34          |  18           |  0.53    |  0.45          | 
| B  |  9        |  14          |  6           |  0.42  |  0.67          |
| C   |  22        |  21          |  15           |  0.71  |  0.68          |
| D   |  4        |  7          |  3           |  0.43  |  0.75          |
| A + B   |  49        |  48          |  28           |  0.58  |  0.57          |


- The findings from the previous table and Fig. 6 suggest that Llama-7B and Goat-7B have roughly the same circuit, with overlapping components that have a high causal impact on the final output. While some of the components are unique, their overall impact is minimal.
    
-   Please refer to Appendix D for more details.

Reviewer comment (Evaluation):
As mentioned in (Wang et al., 2022) [1], faithfulness is not sufficient to prescribe which circuits explain the behavior well. Why do not authors show the completeness and minimality scores, similar to (Wang et al., 2022) [1]?

Author response:
**As mentioned in (Wang et al., 2022) [1], faithfulness is not sufficient to prescribe which circuits explain the behavior well. Why do not authors show the completeness and minimality scores, similar to (Wang et al., 2022) [1]?**

**Minimality**
-   Please note that we do use the minimality criterion from Wang et al. (2022), but instead of reporting the scores we are using it as an optimization criterion; A drawback of path patching is its inability to determine the optimal number of components to incorporate into the circuit. Consequently, there is a high probability that the initially assumed number of components constituting the circuit may contain redundancy. To address this challenge, we employ the minimality criterion to systematically remove redundant attention heads from our initial circuit; unlike Wang et al. (2022), who simply reported the minimality scores without explaining how did they decide on the number of heads to include in the circuit. 
    
-   Specifically, we remove heads that contribute less than 0.5% of the functionality defined by their subcircuit (after removing respective subset K), thus reducing from 100 to 75 attention heads (as described in sec 4.2 in the paper).
    
-   Further, we are investigating the Llama-7B model which is significantly larger than GPT-2 small. As a result, we had to come up with a greedy approach to compute the subset K for each head in the circuit when computing their minimality score, instead of directly using all the heads in a specific group.

**Completeness**
-  Following your suggestion, we have added an evaluation with the completeness score from Wang et al. (2022) to Appendix B of the paper. We used the following two sampling methods, proposed in the original paper:
	1.  *Random*: We uniformly sample a group of circuit components.
	2.  *Circuit subgroups*: We define the subset to be one of four groups of the circuit: A, B, C, and D.

- Results are reported in the following table and Fig. 5 (Appendix B). For the random sampling methods, we report the mean and standard deviation over 20 random subsets.

|           | Full model                      | Circuit                       | Completeness                 |   
|-----------|-----------------------|-----------------------|----------------------|
| Random         |  0.166 ± 0.028                 |  0.068 ± 0.023                  |  0.098 ± 0.037                 |
| Group A        |  0.34                            |  0.004                          |  0.336                         |
| Group B        |  0.152                           |  0.132                          |  0.012                         |
| Group C        |  0.146                          |  0.144                          |  0.002                         |
| Group D        |  0.612                         |  0.424                          |  0.188                         |


-   When we measure completeness over the circuit, we obtain 0.09, which indicates that, for randomly selected components that have an impact on the task, only about 9% of the performance gap is unexplained by the circuit itself.  For groups B and C, the measurement is even lower, indicating highly complete identification of heads within those groups.  But for groups A and D we obtain 0.34 and 0.19, suggesting a lower level of completeness. Although the performance of the current circuit is quite good, this result suggests that including more heads in groups A and D can improve it further.
    
-   Please refer to Appendix B for more details.

Reviewer comment (Experiments):
What is the motivation to divide heads into four groups (A, B, C, D) instead of three or five? In other words, since authors iteratively identify groups of heads with high direct effects on each other using the path patching score, why did authors end with four groups instead of other numbers?

Author response:
**What is the motivation to divide heads into four groups (A, B, C, D) instead of three or five? In other words, since authors iteratively identify groups of heads with high direct effects on each other using the path patching score, why did authors end with four groups instead of other numbers?**

-   During path patching, we observed that the heads at the correct object token and its next position, which could be influencing heads in Group D, had a minimal effect on the final output. This observation is also supported by our faithfulness measure, where we systematically ablated all the heads except those in the circuit. Despite this, we were able to recover most of the model's performance. Consequently, we decided not to include heads after Group D.
    
-   We applied the path patching algorithm for our iterations for the entity tracking task on Llama-7B. For each iteration, we found a set of heads that either directly or indirectly influenced the final output of the model. Therefore, we divided all the attention heads into four groups A, B, C, and D.

Reviewer comment (Writing):
When establishing the Desiderata for identifying circuit functionality, the connection between the three tasks and their corresponding abilities remains ambiguously articulated. A clearer exposition of these relationships would have enhanced the clarity and rigor of the study.

Author response:
To identify the functionality of various circuit components, we conceptualized and defined three desiderata:
 
 - *Object desideratum*: It is used to identify circuit components that are encoding the value of correct object information in their output. Consequently, when the output of these components is patched from the counterfactual to the original run, the final output of the original run changes to the correct object of the counterfactual example, as shown in Fig. 2(a).
 - *Label desideratum*: It is used to identify circuit components that are encoding the query box label value information. Hence, when their output is patched from the counterfactual to the original run, the final output of the original run changes to the object associated with the query box label, of the counterfactual, in the original example, as shown in Fig. 2(b).
 - *Position desideratum*: It is used to identify circuit components that encode the positional information of the correct object, i.e. when they are patched from counterfactual to the original run, the final output of the original run changes to the object in the original example which is at the same location as the correct object of the counterfactual example, as shown in Fig. 2(c).

For each of the three desiderata, we train a binary mask over the model components to identify the circuit components encoding the corresponding vital information to accomplish the task. *We will elucidate these details in the next version of the manuscript.*

Reviewer comment (Theory):
The section detailing the use of CMAP to patch activations from Goat-7B to Llama-7B lacks clarity, particularly when validating the impact of QKV on the model's performance. The rationale behind why patching the QK-circuit of the Value Fetcher and the value vector of the Position Transmitter heads results in the most significant enhancement is not well-explained.

Author response:
As described in Elhage et. al (2021), the QK (query-key) circuit and OV (output-value) circuit are two independent computations of an attention head. The QK circuit computes the attention pattern and the OV circuit computes the output.

We observe a significant performance gain when the QK-circuit of Value Fetcher heads are patched from Goat-7B to Llama-7B (Fig. 4(b)) which indicates that their capability to attend to the correct object using the positional information forwarded by the Position Transmitter heads has improved during fine-tuning. The computation of the QK-circuit depends on 1) Weight matrices (W_Q & W_K) and 2) their inputs, which in the case of the Value Fetcher heads are primarily constructed from the output of the Position Transmitter heads. Since we observe only a marginal performance increment when patching the output of the Position Transmitter heads across models (Fig. 4(a), purple bar), we attribute the improved capability of Value Fetcher heads to the transformation in the W_Q and W_K matrices of these heads, i.e. fine-tuning has refined these matrices, making them more adept at leveraging positional information to attend to the correct object token. On the other hand, in Position Transmitter heads, the value vector (hence OV-circuit) is responsible (Fig. 4(c)) for the increase of performance, suggesting that their contribution to the performance gain can be attributed to the information encoded in their outputs, which we know from DCM experiments is primarily positional information.

*We will elucidate these details in the next version of the manuscript.*