HeadMap: Locating and Enhancing Knowledge Circuits in LLMs
Authors: Xuehao Wang
Liyuan Wang
Binghuai Lin
Yu Zhang
Venue: ICLR 2025 Poster (2025)
Keywords: Large Language Model
Parameter-Efficient Fine-Tuning
Interpretability
Primary area: other topics in machine learning (i.e., none of the above)

Abstract: Large language models (LLMs), through pretraining on extensive corpora, encompass rich semantic knowledge and exhibit the potential for efficient adaptation to diverse downstream tasks. However, the intrinsic mechanisms underlying LLMs remain unexplored, limiting the efficacy of applying these models to downstream tasks. In this paper, we explore the intrinsic mechanisms of LLMs from the perspective of knowledge circuits. Specifically, considering layer dependencies, we propose a layer-conditioned locating algorithm to identify a series of attention heads, which is a knowledge circuit of some tasks. Experiments demonstrate that simply masking a small portion of attention heads in the knowledge circuit can significantly reduce the model's ability to make correct predictions. This suggests that the knowledge flow within the knowledge circuit plays a critical role when the model makes a correct prediction. Inspired by this observation, we propose a novel parameter-efficient fine-tuning method called HeadMap, which maps the activations of these critical heads in the located knowledge circuit to the residual stream by two linear layers, thus enhancing knowledge flow from the knowledge circuit in the residual stream. Extensive experiments conducted on diverse datasets demonstrate the efficiency and efficacy of the proposed method. Our code is available at https://github.com/XuehaoWangFi/HeadMap.

Published as a conference paper at ICLR 2025
HEADMAP: LOCATING AND ENHANCING KNOWL- EDGE CIRCUITS IN LLMS
Xuchao Wang! 122. Liyuan Wang", Binghuai Lin", Yu Zhang
Soothern University of Science and Technology
Tencent Technology Co., Ltd, China
(xuehacwangfi, summersumue Sonmersunshinne366 1 inbinghuaa, yu. zhang. ust bost)eee (egmaiml com
ABSTRACT
Large language models (LLMs), through pretraining on extensive corpora, encom- pass rich semantic knowledge and exhibit the potential for efficient adaptation to diverse downstream tasks. However, the intrinsic mechanisms underlying LLMs remain unexplored, limiting the efficacy of applying these models to downstream tasks. In this paper, we explore the intrinsic mechanisms of LLMs from the per- spective of knowledge circuits. Specifically, considering layer dependencies, we propose a layer-conditioned locating algorithm to identify a series of attention heads, which is knowledge circuit of some tasks. Experiments demonstrate that simply masking a small portion of attention heads the knowledge circuit can significantly reduce the model's ability to make correct predictions. This suggests that the knowledge flow within the knowledge circuit plays a critical role when the model makes a correct prediction. Inspired by this observation, we propose a novel parameter-efficient fine-tuning method called HeadMap, which maps the activations of these critical heads in the located knowledge circuit to the resid- ual stream by two linear layers, thus enhancing knowledge flow from the knowl- edge circuit the residual stream. Extensive experiments conducted on diverse datasets demonstrate the efficiency and efficacy of the proposed method. Our code is available hhtts:///itthhb.. http://giiiuu..ooom/e.oaoooooeet Xuehao Wanghii/aa///HeaaMdMM
INTRODUCTION
Large language models (LLMs) have demonstrated strong language understanding capabilities (Rad-
ford et al.) 2018 2019 Kenton & Toutanova 2019 Lewis al., 2020; Brown et al. 2020) by pertaining on large corpora and can achieve excellent performance on downstream tasks with fine-
tuning (Houlsby al. 2019; Li Liang] 2021 Hu al. 2021 Chen et al. 2024 Wei et al. 2024 Yu et al. 2024). Consequently, LLMs are applied to wide range of fields, which creates growing need understand model behavior. However, due to the complex nonlinear interactions within these models, understanding their intrinsic mechanisms remains significant challenge.
To understand the intrinsic mechanisms of models, previous works (Meng et al. 2022 Wang et al. 2022 Zhhng et al. 2024) analyze the roles of different components (e.g.. Multi-Hes Head Atten- tion (MHA) blocks and Ford-Forward Network (FFN) blocks) using the causal mediation analysis well-constructed sentences. In particular, Zhang et al. (2024) proposes precise fine-tuning of key attention heads based on their findings to improve the performance of LLMs mathematical computation tasks. While effective, these studies focus on specific tasks and sentence formats, limiting their scalability.
(CMA) method (PEARL 2001 Vig et al. 2020)
To address this limitation, analyze the mechanisms of LLMs from the perspective of multi-head attention. Through experiments, we find that masking a single attention head, randomly masking different heads across various layers, employing simple greedy masking does significantly impact the model's 's ability to make correct predictions. This indicates the presence of complex structure LLMs that affects the model's 's ability to make accurate predictions, and that the behavior
Work done during Xuehao Wang was visiting Tencent as a research intern. Corresponding author.

---

Published as a conference paper at ICLR 2025
of attention heads is influenced by those in earlier layers. Based on these findings, we propose a layer-conditioned locating algorithm, which accounts for the dependency relationships between different layers, to identify important attention heads layer by layer. Using this method, we locate a series of attention heads crucial for the ability to make correct predictions on a specific task. We refer this series of attention heads across all layers as the knowledge eircuit Yaao al. 2024).
In addition to enhancing the understanding of LLMs behavior, since the forward pass of LLMs residual stream where each component of LLMs reads inputs from it and writes outputs to it (Wang et al. [022 Todd et al. 2024 Zhang et al. 2024), the knowledge circuit identified by the proposed layer- conyit-ondiiionall locating algorithm could help improve the performance of LLMs on downstream tasks by enhancing the influence of the knowledge flow in it. To achieve that, we propose a novel parameter-effiieent fine-tuning (PEFT) method called HeadMap. Instead of updating all the attention heads as did in conventional PEFT methods, the proposed HeadMap method only transforms the output of attention heads in the identified knowledge circuit via up and down projection and adds the transformed output into the residual stream.helping LLMs better adapt to downstream tasks. The number of learnable parameters in HeadMap is small, since the number of attention heads in the knowledge cireuit is not large, leading to its parameter efficiency. Through extensive evaluations on diverse datasets, we demonstrate that the proposed HeadMap method achieves comparable performance to the baseline with fewer parameters. Additionally, the HeadMap method is complementary to other PEFT methods (e.g., LoRA (Hu et al. 2021) ) and combining them could significantly enhance the performance.
can be considered
The contributions of this paper include the following points: (i) We propose the layer-conditioned locating algorithm identify the existence of knowledge circuits in LLMs that significantly influ- ence the ability to make correct predictions; (ii) Based on knowledge circuits, we proposed the HeadMap method to enhance the knowledge circuits in the residual stream during the forward pass of LLMs. aiding the model in better adapting to downstream tasks; (iii Experiments on extensive datasets demonstrate the parameter efficiency and effectiveness of the proposed HeadMap method.
RELATED WORKS
Mechanistic Interpretability of Large Language Models. As LLMs play increasingly important
roles across various domains (Obermeyer et al. (019 Rudin 2019: Bender et al. 202T). understand- ing the underlying mechanisms behind their varying performance on different tasks has become a growing research focus (Vig et al.) 2020) Meng et al. 2022]] Wang et al., 2022 Todd et al. 2024] Zhang et al. 2024). To investigate the underlying mechanisms of LLMs, Vig et al. (2020) employ causal mediation analysis (CMA) (PEARL 2001) to interpret the instinct mechanism of LLM. Fol- lowing (Vig et al., 2020), [Meng et al. (2022) apply CMA method to identify the crucial activations for the model's factual predictions, while Todd et al. (2024) employ CMA in the In-Context Learn- ing (ICL) scenario and reveal that some attention heads LLMs transmit compact representations corresponding to the task. To enhance the analysis method proposed in (Vig et al. 2020). Wang al.) (2022) propose the path patching method to analyze the role of attention heads on the indirect object identification task. Inspired by (Vig et al., 2020), Zhang al. (0224 construct dataset for math- ematical computation tasks and employ a path patching method to identify that only a small subset of attention heads plays a crucial role in this task. Yao al. (2024) focus on factual recall tasks and ablate the special edge in the computation graph of the language model to identify the knowledge circuits. Although those studies analyze the intrinsic mechanisms of LLMs from various perspec- tives, their analytical methods heavily rely on well-constructed data and are limited to specific tasks, restricting their application to broader scenarios. In contrast, the proposed layer-conditioned locat- ing algorithm, simple yet efficient method, can identify the knowledge circuit that plays a crucial role making correct predictions for various tasks.
Parameter-Efficient Fine-Tuning, Due to the increasing size of LLMs, fully fine-tuning them for each downstream task has become increasingly difficult. To address this challenge, various PEFT methods have been proposed. Existing parameter-efficient methods are typically categorized into adapter- based methods, prompt tuning methods, and low-rank adaptation methods. Specifically, adapter- based methods (Houlsby et al. 2019 Mahabadi et al. 202 I Karimi Mahabadi et al. 2021]
insert compact modules between transtormer layers. Prompt tuning methods (Li & Liang 2021
Lester et al. 2021 Razdaibiedina et al. 2023 Chen al. 2024 Jiang et al. 2023 add trainable
2

---

Published as a conference paper at ICLR 2025
tokens prefixes input intermediate sequences. Low-rank adaptation methods (Hu et al. 2021 Kopiczko et al., 2023] Wu et al. 2023] Ding et al. 2023 Valipour et al. 2023] Li et al. 2024
Liu et al. 2024; Zhuang et al. 2024) introduces trainable low-rank matrices approximate weight
updates. For example, LoRA represents the weight update matrix as the product of two low-rank matrices and DoRA (Liu al.
2024) decomposes pre-trained weights into magnitude and direction
while fine-tuning direction by LoRA. Different from those methods, the proposed MeadMap method built the critical role of knowledge circuits in LLMs and only makes updates the attention heads in the knowledge circuits.
3 PRELIMINARY
Transformers. In this work, we mainly focus on the autoregressive transformer model, which is adopted by most LLMs. Given a sequence of input token embeddings X xx..x2..... ,xV] NNxd where N the number of tokens the input and d is the dimension of embeddings. For the model with L transformer layers, each layer consists of an MHA block and an FFN block. Formally, the hidden state x of i-th token layer calculated
a! + m'.
where af and m denote the output of the MHA block and the FFN block for the i-th token at layer I, respectively. The calculation of a and m are introduced in the following.
Each MHA block consists of H attention heads, which allow for attending to different represen- tation subspaces at different positions in the sequence (Vaswani 2017). For an individual head h in layer 1, it can be parameterized by three matrices: W Wi.. W E RR # Formally, for the input X
K
V
[x|-|. Xx
x layer the output for attention head calculated as: hh Q
h f(x
x 1 -
x -
(2)
where /(Q.K,
QKT Va/R
and denotes the softmax function. Thus, the output A'
[a), al,, ... anyl of the MHA block formulated as
A' = [h' h42 h W.'
33)
where W E Rddd is a learnable output matrix.
FFN block consists of a up projection matrix W up E RRdxdd and a down projection matrix Wi down Rdd xd where d' d. For token i, the output of the FFN block layer is:
m = <(xx + w! up) down
W ww
where (()) is an activation function.
Next token prediction loss. Given LLM M, the input sequence X, and the target sequence Y Y2.. with N, tokens, the next token prediction loss is defined as
N log M(,
(5)
No
4 FROM MULTI-HEAD ATTENTION TO KNOWLEDGE CIRCUIT
It is generally accepted that different attention heads in multi-head attention modules focus on differ- ent parts the input and different subspaces of feature representations (Vaswani 2017). Therefore, previous works (Todd et al. 2024 Jiang et al. 2024) suggest that small number of attention heads in LLMs can comprehend the reasoning required by tasks and guide the model to make appropriate predictions for diverse ICL tasks. However, it remains unclear whether this phenomenon persists in more challenging tasks, e.g., commonsense reasoning tasks. In this section, we empirically validate this phenomenon commonsense reasoning tasks and present some empirical findings.
3

---

Published as a conference paper at ICLR 2025
26
12-
Kinker
Single monely
26
10
275
2.33
10
Laye
Layer
(a)
(b)
(c)
Randen
Simple Gredyy
1.3
35
20
23
10
Leser
1ayer
(d)
(c)
(f)
Figure 1: Figures (a) and (d) show the average loss on 64 samples from the SIQA and HellaSwag
datasets, respectively, when masking a single attention head. Figures (b) and (e) show the average loss on 64 samples from the SIQA and HellaSwag datasets, respectively, when applying the proposed layer-conditiones locating algorithm. Figures (c) and (f) show the average loss w.r.t. different layers of LLMs on the SIQA and HellaSwag datasets, respectively.
Model. In this section, empirical studies are conducted the LLaMA2-7B model (Touroon et al, 2023). LLaMA2-7B utilizes a decoder-only transformer architecture and is trained on large text corpus, thus it offers competitive performance across various natural language processing tasks.
Task and Datasets. We conduct experiments on commonsense reasoning tasks from the popular evaluation framework im-evaluation--hane harness Gao al. (2024) that have garnered widespread at- tention. Those tasks are designed to challenge models beyond mere pattern recognition, requiring them to apply real-world knowledge make inferences. Due the complexity of the task inputs, performing causal inference analysis through input sequence modifications to identify underlying knowledge circuits, explored in previous works (Wang et al. 2022; Zhang et al. 2024; Meng et al. 2022), presents challenges. The commonsense reasoning tasks contain eight datasets, includ- ing six question-ant -anteinnnnn datasets (i.e., BoolQ (Clark al. al.] 2019 PIQA (Bisk
2020), SIQA (Sap et al. 2019), ARC-e (Clark et al., (018), ARC-c (Clark et al. 2018). and OBQA (Mihaylov al.] 2018)). context completion dataset (i.e., HellaSwag (Zellers et al. 2019 ) and fill-in-the- blank dataset (i.e., WinoGrande (Sakaguchi et al. (222T). Following (Todd et al.) 2024 Jiang al., 2024), our analysis focuses on samples that LLMs can predict successfully and utilizes such 64 samples with the lowest next token prediction loss each dataset to conduct experiments.
4.1 EXPERIMENT 1: LOCATING CRITICAL ATTENTION HEADS
Critical attention heads are considered those who change their outputs and then have a substantial influence on the LLMs' ability to make correct predictions for certain task. Therefore, a straight- forward method to locate critical attention heads is to mask their outputs and observe whether causes the LLM to make incorrect predictions. That is, to evaluate the importance of attention head layer I, we set h loss is much larger than the loss without the masking, this attention head may be critical.
and calculate the training loss (i.e. next token prediction loss). If this
Results. Based on this simple idea, we conduct experiments on commonsense reasoning tasks. Fig- ure I(a) and I(d) illustrate the effect of masking each head in each layer for LLaMA2-7B SIQA and HellaSwag datasets. The overall results are shown in Figure 4 of Appendix A.I According to the results, see that masking single attention head does not cause significant changes to
4

---

Published as a conference paper at ICLR 2025
Algorithm Layer- Condi--onddi Locating Algorithm
Require: Number of selected attention heads of each layer, set of selected samples D, model M with L layers and H heads at each layer.
Initialize HI <
2: for < to L do
3:
for 4- to H do
4:
H - ((1,
5: Lossi,1 <-
6:
for (X, Y) E D do
7:
LossLA f loss,, + L(MH, Y) >. Calculate the loss when mask the heads in H end for H H ((1,)) end for
8:
9:
11:
for E arg.top Saa.tooo_Klosss K(lssst n) do
H h(]
Select heads with higher influence of loss
13:
end for end for
15:
return IHT
the loss of L.Ms.. For example, the highest average loss change, the difference between the aver- age loss after and before masking the SIQA and HellaSwag datasets is only +3.02 and +2.58, respectively. This suggests that for complex tasks, the loss of LLMs not driven by the output of single attention head. Inspired by (Elhage et al. 202T), the knowledge flow in the LLMs can be considered as a residual stream, while MHA and FFN blocks can read and write knowledge stream. Therefore, we hypothesize that within LLMs, there exists knowledge circuit composed of multiple attention heads across different layers, which aids the model in making correct predictions for specific tasks by reading and writing the important knowledge in the residual stream.
this
4.2 EXPERIMENT 2: SIMPLE GREEDY LOCATING ALGORITHM
To have a better understanding the knowledge circuit, we consider the forward pass of the LLM computation graph (Wang al. 2022 Zhang et al. 2024). Each node in this computation graph denotes a computation component edge between different nodes denotes the hidden state, which is calculated from the starting node and is the input of the ending node. Therefore, the knowledge circuit can be viewed as a subgraph of the computational graph in LLMs and facilitates the critical knowledge flow from input to output.
LLM (i.e., attention heads and FFN blocks) and the
Based on this perspective, natural approach to locating knowledge circuit is to make a simple greedy locating. This greedy locating method selects the top-n most critical attention heads in each layer, where the importance of each attention head is calculated as in Section [.] and constructs the knowledge circuit as the subgraph the computation graph consisting of selected attention heads.
Results. Here we set to be and hence the knowledge circuit consists of 8 x 32 attention heads. The green lines in Figures [[c]]and (ff show the loss change curve when attention heads are masked layer by layer using the aforementioned greedy locating method. The overall results are shown in Figure [6 of Appendix AA.3 It is observed that while the initial layers show significant impact, masking attention heads after the 5-th layer does not lead to substantial changes. This contradicts previous findings (Jin et al.] (2024). which finds that attention heads in the later layers still play a significant role. plausible explanation is that this greedy locating method overlooks the inter-layer dependencies within L.Mss. Specifically, when critical attention heads in earlier layers are masked, the attention heads in later layers may behave differently, since some attention heads that seem impactful individually might merely reinforce the knowledge provided by earlier critical attention heads, and their influence disappears once those earlier heads are masked.
4.3 EXPERIMENT 3:
-OEERR-CNDIIONEDD LOCATING ALGORITHM
To address this issue, we propose a layer-conditioned locating algorithm to locate knowledge cir- cuits by taking layer dependencies into consideration. Specifically, the algorithm performs greedy
5

---

Published as a conference paper at ICLR 2025
search layer by layer, with the loss calculation for each masking attention head depending the attention heads masked in the previous layers. The detailed process is shown in Algorithm whose time complexity is ((HL) with H as the number of attention heads per layer in LLMs and L as the total number of layers. Similar to the greedy locating algorithm, we set n = 8 and evaluate the layer-conditioned locating algorithm on LLaMA2-7B.
Results. Figures l(b) and [(e) show the average loss when masking the corresponding attention heads, with critical attention heads in previous layers already masked. The overall results shown in Figure of Appendix A.2 Figures [ and [ show the trend of the average loss w.r.t the
number masked layers for selecting heads per layer randomly by the greedy locating algorithm and by layer-conditioned locating algorithm, respectively. The overall results are shown in Figure [ of Appendix A..3] It can be seen that the knowledge circuit selected by the layer-conditioned locating algorithm has a greater impact on the training loss compared to that selected randomly or by the greedy locating algorithm.
4.4 FINDINGS
Based the above experimental results and analyses, we have the following findings.
Finding 1. Masking some attention heads (i.e. the knowledge circuit) significantly impairs the large language model's ability to make accurate predictions, indicating that the knowledge written and read by the knowledge circuit in the residual stream is crucial for correct predictions.
Finding 2. Masking redundant attention heads in some cases reduces loss, suggesting that knowl- edge written and read by certain redundant attention heads in the residual stream may interfere with the model's accurate predictions.
Finding 3. For a given task, a knowledge circuit impacts the ability of LLMs to make correct predic- tions, and knowledge circuits varies across tasks. This suggests that LLMs use different knowledge and patterns for different tasks, and there is no single knowledge circuit that can handle all tasks.
5 METHODOLOGG
Findings 1 and 3 suggest, there exists crucial subset of attention heads, referred to as the knowl- edge circuit, that are essential for making the right predictions on the corresponding task. These knowledge circuits enable the model to make accurate predictions by reading and writing important knowledge the residual stream of the L.M.. However, Finding suggests the presence of numer- redundant attention heads in the model, whose outputs interfere with the knowledge flow in the knowledge circuit, thus leading to incorrect predictions. Inspired by these findings, we propose a novel parameter-efficient fine-tuning method, HeadMap. This method enhances the knowledge flow from knowledge circuits in the residual stream, aiding large language models in making accurate predictions.
Due to the FFN block reading and writing on both the residual stream and MHA output, we propose two variants of the HeadMap method: the first one enhances the residual stream directly (Section 5.1) and the other enhances in the MHA block (Section (..2). Specifically, the former bypasses the FFN block and Wo, directly writing the knowledge from the knowledge circuit into the residual stream, while the latter writes the knowledge into the MHA outputs, thereby affecting both the residual stream and the operation of the FFN block.
5.1 HEADMAP FOR HIDDEN STATES
The forward pass of LLMs can be viewed a long residual stream, where all attention blocks and FFN blocks read from and write to Thus, at the l-th layer, to enhance the influence of the knowledge circuit located by the faydr-conditioned locating algorithm, we map the outputs of the selected attention heads and add them to the hidden state X Formally, the output of layer is calculated
x X'= + +' M' + Map(h''> hl
(6)
where M' = imf m myy denotes the output of the FFN block in layer I, Map(.) denotes a neural network, Hi denotes the set containing the selected attention heads in the knowledge circuits,

---

Published as a conference paper at ICLR 2025
is the number of selected attention heads per layer, and (1) E H. E [1. n). To reduce the number of trainable parameters in Map(-), we use a bottleneck architecture where we first down- project the outputs of the selected attention heads into r-dimension and then up-project them into d. Thus, the output can be formulated as
X'= A + M Up (Down([h!e 3m h()*)))
5.2 HEADMAP FOR OUTPUTS OF MHA BLOCKS
In addition to adding directly the residual stream, another way to enhance the knowledge circuit is to place MeadMap within the MHA. This involves directly adding the mapped output to the output A' of the MHA block at the 1-th layer, resulting in the following output as
A' 11 h4
4.8 W + Up (Down [h h"(*)))
(8)
Since use two linear mappings here, this method is similar to selecting {Wos (1, perform LoRA fine-tuning and does not introduce inference delay.
HI) to
5.3 IMPLEMENTED DETIALS
Initialization. The parameters of down projection layers are initialized as O, while the parameters of up-project layers are initialized by Kaiming uniform. Therefore, in the begining of the training process, the HeadMap method does not change the output of LLMs.
Parameter complexity. We provide a parameter complexity analysis of the proposed HeadMap method, where the two variants have the same number of parameters. To fine-tune the large lan- guage model with hidden dimension d, the proposed HeadMap method selects n attention heads at each layer. Thus, the HeadMap method introduces O ((n x + 1) x + (r + 1) x d) trainable parameters each layer, where is the dimension after down-projection and r < d.
6 EXPERIMENT
In this section, we empirically evaluate the proposed HeadMap method.
6.1 EXPERIMENTS ON COMMONSENSE REASONING TASKS
In this section, we conduct experiments commonsense reasoning tasks.
Datasets.
As introduced before, the commonsense reasoning tasks comprise eight datasets, in- cluding BoolQ, PIQA, SIQA, ARC-e, ARC-c, OBQA, HellaSwag, and WinoGrande. Each dataset contains a training and testing set. We use the accuracy metric to measure the performance.
Baselines. ILLMM2-7B (Touvron et al. 2023) and ILLMM3-8B (Meta 2024 used to be LLMs. We compare the proposed HeadMap method with LoRA (Hu et al.) 2021) and DoRA (Liu et al. 2024). LoRA and DoRA are used Wq. WK. Wy. Wo. W ups and W Wdown each layer. Addi- tionally, we explore combining HeadMap with LoRA and DoRA by applying LoRA or DoRA to fine-tune Wv. Wups and W down while implementing HeadMap.
Setups. For commonsense reasoning datasets, we fine-tune and evaluate models on each single dataset. The batch size set to 16, and the AdamW optimizer is used. We use the linear learning rate scheduler and train all methods for 2 epochs with 100 warmup steps. The learning rate is set to 0.0002 for LoRA and DoRA, while set to 0.001 for HeadMap. We set rank 32 for all methods and = 64 for LoRA and DoRA. We select 8 attention heads for each MHA block and use 64 samples in each dataset to locate the knowledge circuit.
The experimental results commonsense reasoning datasets shown in Table It can be observed that by enhancing the impact of critical heads, both variants of the HeadMap method achieve comparable results to the baselines while introducing only one-tenth of the parameters. Ad- ditionally, we find that HeadMap and MoRA-based methods (i.e., LoRA and DoRA) are complemen- tary combining LoRA and D with HeadMap achieves better results with fewer parameters.
Results.

---

Published as a conference paper at ICLR 2025
Table 1: Accuracy (%) eight commonsense reasoning datasets using LLaMA2-7B and LLaMA33 8B. The best result in each comparison group is in bold. HсаdMaриня denotes the variant of MeadMap for outputs of MHA blocks as introduced in Section 5.2
and HeadMapHs denotes the
variant of MeadMap for hidden states in Section 5.1
Model
Params (%)
BoolQ
PIQA
SIQA
HellaSwag
WinoGiande
ARC-e
ARC-c
OBQA
Avg.
LAMA227B
LoRA
0.95
63,52
81.34
78.36
93.78
81.93
79.71
54.78
77.40
76.38
0.96
64.25
78.30
82.95
78.20
HeadMapsmt
0.08
63,94
81.19
79.07
92.74
79.40
80.30
57.00
78.60
76.78
LoRA
0.66
65.38
84.66
79.58
80.58
81.65
80.20
77.76
DoRA (V.
0.67
65,93
84.06
79.79
93.67
81.93
82.03
57.00
81.00
78.18
HeadMapus
0.08
63.73
83.30
79.73
92.63
78.37
80,30
55.38
78.20
76.46
LoRA (V.
0.66
66.85
84.06
79.84
93.59
82.08
82.37
56.31
80,60
78.21
DoRA (V.U,
0.67
67.16
84.28
80.14
93.59
82,00
81,73
56.57
$1.00
78.31
LoRA
0.80
76.30
86.07
80.60
96.01
86.82
90,70
78.67
83,40
84.82
DoRA
76.73
86.40
80.60
96.13
86,42
91.37
78.75
87.60
HeadMapsara
0.07
75.29
89.17
81.47
95.64
85,24
92.00
78.75
88,40
85.74
+ LoRA (V. D)
0.60
76.30
90.48
81.27
96.16
87.13
92.05
80.46
88.00
86.48
DoRA (V,
0.61
76.67
89.83
81.68
96.15
87.92
92.13
79.69
87.80
86.51
HeadMapus
0.07
75.35
89.23 28
81.53
95.53
85,40
92.17
78.92
87.80
85. 75
LoRA D)
0.60
76.24
89.39
81.52
87.29
91.84
81.06
88.60
86.56
DoRA (V, U, D)
0.61
76.48
90.04
81.73
96.44
87.29
92.30
79.78
87.00
86.38
Table 2: Accuracy (%) on five NLU datasets using LLaMA2-7B and LLaMA3-8B.
The best result
in each comparison group is bold. HeadMap is for outputs MHA blocks. HeadMapus for hidden states.
Method
Params (9)
CoLAA
MRPC
RIE
SST-2
WNLI
Avg.
LIaMA2-7B
LoRA
0.95
49.20
70.34
87.00
96.10
39.15
72.36
HеadMaрaии
01.88
69.12
82.67
95.87
LoRA (V, D)
0.66
59.92
82.11
87.36
96.90
56.34
76.53
HeadMapus
0.08
47.63
77.70
85.92
95.53
54.93
72.34
LoRA (V,
0.66
64.41
80.39
87.36
96.44
59.15
77.55
LLaMA3-8B
LoRA
0.80
65.48
85.29
88.81
96.10
70.42
81.22
HeadMapsmi LoRA (V. U.
0..07
62.75
85.78
88.09
95.76
71.83
80.84
0.60
67.82
87.25
90.25
96.33
80.28
84.39
HeadMapus
0.07
85.29
88.81
95.99
LoRA (V, D)
0.60
68.77
87.75
89.89
96.44
77.46
84.06
Moreover, on LLaMA2-7B. combining HcadMap for hidden states with DoRA achieves the best per- formance on average (+ (+1.29% compared to DoRA) and some tasks (i.e. BoolQ, SIQA, and OBQA), while HeadMap for hidden states with LoRA achieves the best performance compared to DoRA) and several tasks (i.e. HellaSwag, Arc-c, and OBQA) on LLaMA3-8B.
average (+1.06%
6.2 EXPERIMENTS ON NATURAL LANGUAGE UNDERSTANDING TASKS
In this section, we conduct experiments on Natural Language Understanding (NLU) tasks.
Datasets. For NLU tasks, we compare baselines on five datasets, including CoLA (Warstadt al. 2018), MRPC (Dolan & Brocketi (005), RTE (Wang et al. 2019), SST-2 (Socher et al. 2013)). and WNLI Wang et al. 2019). We use Matthew's correlation for CoLA and accuracy for other tasks to measure the performance of each method.
Baselines. For NLU tasks, we compare the proposed HeadMap method with LoRA on LLaMA2-7B and LLLMA3-8B. LoRA are used on Wo, Wx. Wv. Wo. W and Wdown at each layer. Moreover, we combine HeadMap with LoRA.
Setups. Similar to commonsense reasoning datasets, we fine-tune and evaluate models on each single dataset. The batch size is set to 16, and the AdamW optimizer is used. We use the linear learning rate scheduler and train all methods for 2 epochs with 100 warmup steps. For the learning rate each method, we perform the grid search in {0.0001, 0.0002, 0.0004, 0.0008, 0.001 on each 64 for LoRA. We select 8 attention heads for
dataset. We set rank each MHA block. We select 8 attention heads for each MHA block and use 64 samples with the lowest loss in each dataset locate the knowledge eircuit.
32 for all methods and
8

---

Published as a conference paper at ICLR 2025
Table Accuracy (%) on eight commonsense reasoning datasets using LLaMA2-7B and LLaMA3- 8B for different selection strategies. MeadMap is for outputs of MHA blocks. HeadMapHs is for hidden states.
Method
BoolQ
PIQA
SIQA
HellaSwag
WinoGrande
ARC-c
ARC-c
OBQA
Avg.
LLaMA2-7B
HeadMapxmn (Random)
62.63
82.48
78.25
91.88
78.14
78.07
54.78
75.00
75.15
HeadMapsaнA
63.94
83.19
79.07
92.74
79.40
80.30
57.00
78.60
76.78
Headdapps (Random)
63.27
83.57
78.30
92.51
76.64
78.79
53.41
75.40
75.24
HeadMapas
63.73
83.30
79.73
92.63
78.37
80.30
55,38
78.20
76.46
--MA3388
Mapsaaa (Random)
74.46
87.92
80.50
84.05
90.95
77.90
84.74
HeudMapxmn
75.29
89.17
81.47
95.64
85.24
92.00
78.75
88.40
85.74
HeadMapus (Random)
74.34
88.25
80.30
95.40
83.82
91.08
77.30
86.80
84.66
HeaddMapas
75.35
89.28
81.53
95.53
85,40
92.17
78.92
87,80
85,75
Table Accuracy (%) on eight commonsense reasoning datasets using LLaMA2-7B and LLaMA3- 8B for different number of selected heads. HeadMap is for outputs of MHA blocks. HeadMapus is for hidden states.
Method
# Params (%)
BoolQ
PIQA
SIQA
HellaSwag
WinoGrrande
ARC-c
ARC-c
OBQA
Avg.
HeadMapana (4)
0.07
64.50
82.10
79.93
92.59
78.93
79.29
55.38
78.20
76.37
Head HeadMapaus (8)
0.08
63.94
83.19
79.07
92,74
79.40
80.30
57.00
78.60
76.78
HcadMapaиs (16)
64.28
83.41
79.32
79.50
56.06
79.00
76.77 74
HeadMapus (4)
0.07
64.04
82.92
78.81
92.56
79.01
79.67
54.18
79.20
76.30
HeadMapus (8)
63.73
79.73
92.63
80,30
55,38
78.20
HeadMapus (16)
0.10
64.28
83.68
78.97
92.80
78.85
79.63
55.72
77.60
76.44
LLaMA3-8B
HcadMapииa
75.44
88.63
81.68
95.68
84,93
92.30
79.30
87.80
85.72
HeadMapsns (8)
0.07
75.29
89.17
81.47
95.64
85.24
92.00
78.75
88.40
85.74
HcadMapwnи (16)
74.31
89.23
81.83
95.86
86.11
90.91
79.10
87.00
HeadMapus (4)
0.06
75.05
88.36
81.17
95.35
85.08
91.75
79.86
87.20
85.48
HeadMapus (8)
75.35
89.28
81.53
95.53
85.40
92.17
78.92
87.80
85.75
HcadMapus
75.08
81.63
95.76
85.63
92.13
79.44
87.60
85.75
The experimental results on diverse NLU datasets are shown in Table 2 As can be seen, HeadMap achieves comparable performance with fewer parameters. Additionally, compared to directly only applying LoRA for fine-tuning, combining it with HeadMap significantly enhances performance. Specifically, HeadMap for hidden states has much better performance on average (45.19% compared to LoRA) and on the CoLA, RTE, and SST-2 datasets with LLaMA2-7B as the LLM. For LLaMA3-8B, HeadMap for outputs of MHA blocks achieves the best results on average (+2.84% compared to LoRA) and on the CoLA. MRPC, and SST-2 datasets.
Results.
6.3 ABLATION STUDY
Effect of layer-conditioned locating algorithm. Here, we conduct ablation study on the proposed layer-conditioned locating algorithm. Following the hyperparameters setup in Section we com- pared with a random selection algorithm on commonsense reasoning datasets. According to the results shown in Table 3
that attention heads selected by the layer-conditioned locating algorithm outperforms random selection across all datasets and models, demonstrating the effective- ness of the proposed layer-conditioned locating algorithm.
Sensitivity to the number of selected heads per layer. Here, we investigate the impact of the number of selected attention heads to the fine-tuning performance with the same settings of other hyperparameters as Section 6..1 According to the results shown in Table we can observe that the performance of the proposed HeadMap method remains stable the number of heads across the two LLMs, which indicates that critical attention heads in each layer are limited and adding more does not lead performance improvement.
6.4 ANALYSIS
Overlapping knowledge circuits in different datasets. In this experiment, we analyze the overlapping of attention heads in knowledge circuits identified in different datasets by adopting the

---

Published as a conference paper at ICLR 2025
1.0
0.32
250
0.37
5:0
0.3
0:32
0.31
0.32
0.32
5.25
200
0.25
0:31
0.28
0.3
1.0
0:35
0.25
3.32
0.52
0:33
1-0
0.47
0.32
0.55
0.31
RAC
0.32
0:35
50
0880
0.3
2:28
0:34
9:36
0.35
0.32
8-46
1.0
HelaSSaap
ARCi
(b)
Figure 2: (a) The overlap of knowledge circuits between different tasks LLaMA2-7B. (b) Distribution of attention head selection frequency.
32
20
20
20
25
30
20
25
(a)
(b)
Figure 3: The visualization of located knowledge circuits for (a) Arc-e and (b) WinoGrande on LLaMA2-7B.
same setups as Section 6.1 datasets a and b as
Specifically, we calculate the overlap rate of knowledge circuits between
LxH
The overlap rate between different tasks on LLaMA2-7B and LLaMA3-SB are shown in Figures 2(a) and 8 respectively. The visualization of knowledge circuits for the Arc-easy and WinoGrande datasets on LLaMA2-7B is shown Figure 3 It can be observed that for most tasks, the overlap rate of their identified knowledge circuits is low. This suggests that the knowledge circuits tend to vary significantly due to differing knowledge requirements for different tasks. However, some tasks show higher overlap rates. For example, when selecting only 8 attention heads per layer, the overlap rate between the knowledge circuits in the ARC-e and WinoGrande datasets reaches 0.47, reflecting the similar knowledge and behaviors needed for model reasoning these tasks.
The number of times for attention heads being selected. In Figure 2(b) we record the number of times being selected for each attention head over the eight commonsense reasoning datasets. It can be seen that only a small number of attention heads are located by many different tasks, indicating that most attention heads play crucial role only a few specific tasks.
CONCLUSION
In this paper, we propose the layer-conditioned locating algorithm identify knowledge circuits containing series of attention heads that play crucial role in a specific task. After locating the knowledge circuit, we propose the HeadMap method to enhance the influence of the knowledge circuit on the residual stream, allowing the model to better adapt to downstream tasks. Extension experiments demonstrate that HeadMap achieves comparable performance with few parameters. Moreover, HeadMap and LoRA-based methods are complementary. and their combination effec- tively improves model performance. In our future work, we are interested in applying HeadMap to more applications.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
This work is supported by National Key R&D Program of China 20222D0160300 and NSFC key grant under grant no. 62136005.
REFERENCES
Emily M Bender, Timnit Gebru, Angelina McMillan- Major, and Shmargaret Shmitchell. On the dangers of stochastic parrots: Can language models be too big? In Proceedings of the 2021 ACM conference on fairness, accountability, and transparency, pp. 610-623, 2021.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language. In Thirty-Fourte AAAI Conference Artificial Intelligence, 2020.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:187771901,, 2020.
Yang Chen, Shuai Fu, and Yu Zhang. Mopd: of distillation for vision-lange Aision-anguaage models, 2024. URL wttpttttp...........t... : /tt///xiivve. wttp://///abb///////// 9087
Christopher Clark, Kenton Lee, Ming- Weig Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. BoolQ: Exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Lingnistics: Human Language Technologies, Volume 1 (Long and Short Papers), pp. 2924-2936, Minneapolis, Minnesota, June 2019. Association for Computational Linguisties. doi: 10.18653/ v1/N19-1300. URL wtt.:::ttts...m : //tcc///cc ac antho conttolloge logy. org /N1 9-1100-1000
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge.
1oXYv:18033.555571.. 2018.
Ning Ding, Xingtai Lv, Qiaosen Wang. Yulin Chen, Bowen Zhou, Zhiyuan Liu, and Maosong Sun. Sparse low-rank adaptation of pre-trained language models. In The 2023 Conference on Empirical Methods in Natural Language Processing, 2023.
William B. Dolan and Chris Brockett. Automatically constructing corpus sentential paraphrases. In Proceedings of the Third International Workshop on Paraphrasing (IWP2005), 2005. URL
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. http::/taanstommeerccicciitppppp /ttp://02///ammewo///neee. htmmll...
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Fos- ter, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, h, Haonan Li, Kyle McDonell, Niklas Muen- nighoff, Chris Ociepa. Jason Phang. Laria Reynolds. Hailey Schoelkopf, Aviya Skowron, Lin- tang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 07 2024. URL 12608602
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, An- drea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficcent transfer learning for nlp. In International conference machine learning, pp. 2790-2799. PMLR, 2019.
Edward J Hu, Yelong Shen. Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv: 10c.2/00.09665 2021.
11

---

Published as a conference paper at ICLR 2025
Gangwei Jiang, Zhaoyi Li, Caigao Jiang, Siqiao Xue, Jun Zhou, Linqi Song, Defu Lian, and Ying Wei. Interpretable catastrophic forgetting of large language model fine-tuning via instruction vector. arXiv preprint arXiv:2406.12227, 2024.
Weisen Jiang, Yu Zhang, and James Kwok. Effective structured prompting by meta-learning and representative verbalizer. In International Conference on Machine Learning, pp. 15186-15199. PMLR, 2023.
Zhuoran Jin, Pengfei Cao, Hongbang Yuan, Yubo Chen, Jiexin Xu, Huaijun Li, Xiaojian Jiang, Kang Liu, and Jun Zhao. Cutting off the head ends the conflict: A mechanism for interpreting and mitigating knowledge conflicts in language models. arXiv preprint arXiv:2402.18/36 2024.
Rabeeh Karimi Mahabadi, James Henderson, and Sebastian Ruder. Compacter: Efficient low-rank hypercomplex adapter layers. Advances in Neural Information Processing Systems, 34:1022 1035, 2021.
Jacob Devlin Ming-Wei Chang Kenton and Lee Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings NAACLHHLT, pp. 4171- 4186, 2019.
Dawid Jan Kopiczko, Tijmen Blankevoort, and Yuki M Asano. Elora: Efficient low-rank adaptation with random matrices. In The Twelfth International Conference Learning Representations, 2023.
Brian Lester, Rami AI-Rfou, and Noah Constant, The power of scale for parameter-efficiet prompt tuning. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Pro- cessing, pp. 3045-3059, 2021.
Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Veselin Stoyanov, and Luke Zettlemoyer. Bart: Denoising seneene-to-sequenee pre- training for natural language generation. translation, and comprehension. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 7871-7880, 2020.
Dengchun Li, Yingzi Ma, Naizheng Wang, Zhengmao Ye, Zhiyuan Cheng, Yinghao Tang, Yan Zhang, Lei Duan, Jie Zuo, Cal Yang, and Mingjie Tang. Mixlora: Enhancing large language models fine-tuning with lora-based mixture of experts, 2024.
Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 1 Itth International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pp. 4582 4597, 2021.
Shih- Yang Liu, Chien-Yt Yi Wang, Hongxu Yin, Pavlo Molchanov, Yu-Chiang Frank Wang, Kwang- Ting Cheng, and Min-Hung Chen. Dora: Weight-decomposed low-rank adaptation. arXiv preprint ar arXiv:2402.0953 2024.
Rabeeh Karimi Mahabadi, Sebastian Ruder, Mostafa Dehghani, and James Henderson. Parameter- efficient multi-task fine-tuning for transformers via shared hypemetworks. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguisties and the I 11th International 2021.
Joint Conference on Natural Language Processing ( (Volume 1: Long Papers), pp.
Kevin Meng. David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt. Advances Neural Information Processing Systems, 35:17359-17372 2022.
AI Meta. Introducing meta llma 3: The most capable openly available I1m to date. Meta AI, 2024.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? new dataset for open book question answering. In EMNLP, 2018.
Ziad Obermeyer, Brian Powers, Christine Vogeli, and Sendhil Mullainathan. Dissecting racial bias in an algorithm used to manage the health of populations. Science, 366(6464).447-455 2019.
J PEARL. Direct and indirect effects. In Proc. of the 17th Conference on Uncertainty in Artificial Intelligence, 200/, pp. 4114111111 411-420, 2001.
12

---

Published as a conference paper at ICLR 2025
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language under- standing by generative pre- vre-training. 2018.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
Anastasiaa Razdaibiedina, Yuning Mao, Madian Khabsa, Mike Lewis, Rui Hou, Jimmy Ba, and Amjad Almahairi. Residual prompt tuning: improving prompt tuning with residual reparameteri- zation. In The sl Annual Meeting Thehe The Association For Computational Linguistics, 2023.
Cynthia Rudin. Stop explaining black box machine learning models for high stakes decisions and use interpretable models instead. Nature machine intelligence, 165):266-2155 2019.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adver- sarial winograd schema challenge scale. Communications of the ACM, 64(9)999-106, 2021.
Maarten Sap, Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi. Socialiqa: Common- sense reasoning about social interactions. arXiv preprint 1rXiv:1904.09725. 2019.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning. Andrew Ng, sentiment
and Christopher Potts. Recursive deep models for semantic compositionality treebank. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing, pp. 1631-1642, Seattle, Washington, USA,, October 2013. Association for Computa- 1170
tional Linguistics. URL
Eric Todd, Millicent Li, Arnab Sharma, Aaron Mueller, Byron C Wallace, and David Bau. Function vectors in large language models. In International Conference on Learning Representations. ICLR, 2024.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint 2rXiv:2307.09288, 2023.
Mojtaba Valipour, Mehdi Rezagholizadeh, Ivan Kobyzev, and Ali Ghodsi. Dylora: Parameter- efficient tuning of pre-trained models using dynamic search-free low-rank adaptation. In Pro- ceedings of the 17th Conference of the European Chapter of the Association for Computational Linguistics, pp. 3266-3279, 2023.
Vaswani. Attention all you need. Advances Neural Information Processing Systems, 2017.
Jesse Vig, Sebastian Gehrmann, Yonatan Belinkov, Sharon Qian, Daniel Nevo, Yaron Singer, and Stuart Shieber. Investigating gender bias language models using causal mediation analysis. Advances in neural information processing systems, 13:12888-1240) 2020.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. In International Conference on Learning Representations, 2019.
Kevin Ro Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Inter- pretability in the wild: a circuit for indirect object identification gpt-2 small. In The Eleventh International Conference Learning Representations, 2022.
Alex Warstadt, Amanpreet Singh, and Samuel R Bowman. Neural network acceptability judgments. arXiv preprint arXiv: 000.11005..7111 2018
Yanbin Wei, Shuai Fu, Weisen Jiang, Zejian Zhang, Zhixiong Zeng, Qi Wu, James Kwok, and Yu Zhang. GITA: Graph to visual and textual integration for lasion-language graph reasoning. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL net htt.:/oormmmmdssdoo111
https:://penns httpr//rreeeeeeeeeeooooooooooooo
Xun Wu, Shaohan Huang, and Furu Wei. Mole: Mixture of lora experts. In The Twelfth International Conference Learning Representations, 2023.
13

---

Published as a conference paper at ICLR 2025
Yunzhi Yao, Ningyu Zhang. Zekun Xi, Mengru Wang. Ziwen Xu, Shumin Deng, and Huajun Chen. Knowledge circuits in pretrained transformers. In The eiirty- Annual Conference on Neu- ral Information Processing Systems, 2024. URL pettps::/openzeviiee net wet./oruu/mdd YVXzZNxcac
Longhui Yu, Weisen Jiang, Han Shi, Jincheng YU, Zhengying Liu, Yu Zhang, James Kwok, Zhen- guo Li, Adrian Weller, and Weiyang Liu. Metamath: Bootstrap your own mathematical questions for large language models. In The Twelfth International Conference on Learning Representations, 2024. httpttttss//pneeeeeet
ew net/ forum? GENNOhhNNDR
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can ma- chine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.
Wei Zhang, Chaoqun Wan, Yonggang Zhang, Yiu-ming Cheung, Xinmei Tian, Xu Shen, and Jieping Ye. Interpreting and improving large language models in arithmetic calculation. In Forty-first International Conference on Machine Learning, 2024.
Zhan Zhuang, Yulong Zhang, Xuehao Wang, Jiangang Lu, Ying Wei, and Yu Zhang. Time- varying loRA: Towards effective cross-domain fine-tuning diffusion models. In The Thirty- eighth Annual Conference Neural Information Processing Systems, 2024. URL https:
7/ //opereeviww.re wttp://etttttooomieeeeooooeooee
14

---

Published as a conference paper at ICLR 2025
A ADDITIONAL EXPERIMENTAL RESULTS
A.1 LOCATING THE CRITICAL ATTENTION HEAD
(a) Arc-easy.
Arc-challenge.
(c) BoolQ.
(d) OBQA.
(e) PIQA.
WinoGrande.
Figure 4:
When masking a single attention head, the average loss of 64 samples from each dataset.
The average loss model predictions when a single attention head is masked is shown in Figure It can be observed that masking single attention head does not have significant impact across different datasets.
A.2 LAYER LANDIICDND LOCATING ALGORITHM
(a) Are-easy.
(b) Arc-challenge.
(c) BoolQ.
1-
(d) OBQA.
(c) PIQA.
(1) WinoGrande.
Figure 5: When applying the proposed layer-conditioned locating algorithm, the average loss of 64 samples from each dataset.
Figure 5 illustrates the average loss model predictions when single attention head is masked in the current layer, after previous layers have been masked, using the layer-conditioned locating algorithm. As can be seen, the proposed method can find some attention heads that have much larger impact on model loss compared to masking single attention head. This indicates that the proposed method effectively identifies a series of attention heads, termed knowledge circuits, that significantly influence the training loss of large language models.
A.3 AVERAGE LOSS CURVES OF DIFFERENT LOCATING METHODS
Figure shows the loss curves of different location methods across various datasets w.n. layer of L.LMs. can be observed that, compared random selection and simple greedy methods,
15

---

Published as a conference paper at ICLR 2025
Uppr only
121
(a) Arc-easy.
Are-challenge.
(c) BoolQ.
Uaper oney cotfonatttttttt
Comder coddy
(d) OBQA.
(e) PIQA.
WinoGrande.
Figure 6: The average loss w.r.t. layer of LLMs all datasets.
proposed layer-conditioned locating algorithm accurately identifies the attention heads with the greatest impact on model training loss.
A.4 AVERAGE LOSS CURVES ON ALL SAMPLES AND SELECTED SAMPLES OF EACH DATASET
(a) Arc-easy.
(b) Arc-challenge.
(c) BoolQ.
12
1400
(d) OBQA.
(e) PIQA.
(1)
WinoGrande.
Figure 7: The average loss w.r.t. layer of LLMs on all samples of datasets.
To demonstrate that using samples with minimal loss effectively represents the model's ability to handle the corresponding task, we plotted the loss across layers for all samples and for those with minimal loss, as shown in Figure As can be seen, the curves align closely, indicating that the loss variations in samples with the smallest loss reflect those of all samples. This suggests that using samples with the smallest loss to identify knowledge circuits an efficient and effective approach.
A.5 SIMILARITY OF KNOWLEDGE CIRCUITS
Figure 3 shows the similarity of knowledge circuits identified across different datasets on LLaMA3- 8B. Most tasks exhibit very low similarity, while few related tasks, such ARC-e and ARC-c, have higher similarity. This highlights the potential of using knowledge circuits to analyze the behavior of similar tasks.
16

---

Published as a conference paper at ICLR 2025
BoolQ
1.0
0.31
0.31
0.29
0.26
0.34
0.27
0.31
MOA
0.31
1.0
0.29
0.35
0.3
0.39
0.35
0.33
SIQA
HelloSwag
0.29
D.35
0.38
1.0
0.39
0.41
0.45
0.42
GindGrande
0.26
0.3
0.36
0.39
1.0
0.36
0.47
0.38
ARC-8
ARC-c
0.27
0.35
0.38
0.45
0.47
0.52
1.0
0.51
OBOA
0.31
0.33
0.38
0.42
0.38
0.46
0.51
1.0
BoolQ
PIQA
SIQA
HellaSwag
WineGrande
ARC-a
ARC-C
OBQA
Figure 8: The similarity of selected heads between different tasks LLaMA3-SB.
A.6 PATTERN OF KNOWLEDGE CIRCUITS
(a) SIQA.
(b) Arc-challenge.
(c) BoolQ.
(d) OBQA.
(e) PIQA.
(f) HellaSwag.
Figure 9: The located knowledge circuit of each dataset.
Figure presents the patterns of knowledge circuits located by the layer-conditioned locating algo- rithm across different datasets. While some datasets exhibit similar patterns, the majority show sig- nificant variations. This motivates to search for distinct knowledge circuits in different datasets.
B IMPLEMENTED DETAILS
presents the learning rates of LoRA, HeadMap, and their combination across different NLU datasets. For HeadMap combined with the LoRA method, the former values in parentheses indicate the learning rate for HeadMap and the latter is for LoRA.
C ABLATION STUDY
Transferability. In order to explore whether the knowledge circuit has transferability. we enhance the knowledge circuit identified in the WinoGrande dataset different datasets. The results pre-
17

---

LoRA Published as a conference paper at ICLR 2025
Table 5: Learning rate of each baseline and each dataset.
Method
CoLA
MRPC
RTE
SST-2
WNLI
LLaMA2-7B
2e-4
8e-4
8e-4
2e-4
le-3
HeadMap
1e-3
8e4
8e-4
le.4
1e-3
LoRA (V, U,
(le-4, 4e-4)
(Ie-4, 8e-4)
(le-3, 8e-4)
(Ie-3, 2e-4)
(8e-4, 4e-4)
HeadMapus
8e-4
8e-4
le-3
2e-4
8e-4
LoRA (V, D)
(8e-4, 4c-4)
(Ie-4, le-3)
(le-3, 8e-4)
(2e-4, 2c-4)
(2e-4, le-3)
L.aMA3-8B
LoRA
2e-4
8e-4
4e-4
2e-4
2e-4
HeadMap
4e-4
8c-4
8c-4
lc-4
le-3
LoRA (V, D)
(Ie-4, 4c-4) 8e-4
(2e-4, 2c-4)
4c-4) (le-4,
(Ie-4, 2c-4)
(8e-4, 8c-4)
HeadMaphs
le-3
8e-4
2e-4
4e-4
LoRA (V, D)
(4e-4, le-4)
(Ie-3, 2e-4)
(2e-4, 4e-4
(4e-4, le-4)
(8e-4, 8c.4)
Table 6: Accuracy (%) eight commonsense reasoning datasets using LLaMA2-7B and LLaMA3- 8B for different selection strategies. HeadMap is for outputs of MHA blocks. HeadMapıs for hidden states.
Method
HellaSwag
WinoCrandd
Avg.
LLaMA2-7B
HaadMappanx (Redundant)
62.35
77.74
74.86
HeadMapxии (Random)
62.63
82.48
78.25
91 88
78.14
78.07
54.78
75.00
75.15
HeudMaapes (Transfer from Winogrande)
64.04
82.86
79,06
92.53
79.40
79.46
56.40
77.40
76.39
HeadMMapma
63.94
83.19
79.07
92.74
79.40
80.30
57.00
78.60
76.78
HeadMapus (Redundant)
63.03
82.37
76.92
91.37
78.30
77.99
53.42
75.80
74.90
HeadlMaps (Random)
63.27
83.57
78.30
92.51
76.64
78.79
53.41
75.40
75.24
HeadMapus (Transter from Winogrande)
64.53
82.97
78.92
92.58
78.37
79.76
54.95
77.60
76.21
HeadMapus
79.73
92.63
78.37
80.30
78.20
76.46
Table 7: Time consumption of layer-condittoneed locating algorithm on each dataset.
BoolQ
PIQA
SIQA
HellaSwag WinoGrande
ARC-e
ARC-c
OBQA
Cost Time (min)
3.58
22.13
8.62
22.2
6.53
16.68
17.30
11.10
sented in Table It can be observed that compared to random selection, the transferred knowledge circuit demonstrates better performance. This suggests that some of the attention heads within the knowledge circuit contain general knowledge, not merely task-specific or dataset-speeific knowl- edge. Hence, the knowledge circuits located by the layer-conditioned locating algorithm are capable of generalizing to other tasks or datasets.
Rebundant. We conduct additional ablation studies redundant heads. Specifically, domly select 8 attention heads from redundant heads in each layer and enhance these attention heads. The results shown in Table 6 It be observed that compared to fully random selection in each layer, the performance worse when selecting and enhancing attention heads from redun- dant heads. This indicates that the assistance provided by these redundant heads for the model's adaptation to downstream tasks is more limited.
D TIME CONSUMPTION OF LAYER-CONDIINN CONDITIONED LOCATING
The time consumption of layer-conditioned locating is acceptable, as only a few samples require inference. This is further supported by the time required for knowledge circuit localization on each dataset, as shown in Table 7 Moreover, the localization procedure needs to be conducted only once per dataset, or even across multiple datasets, since experiments demonstrate that the located knowledge circuits are transferable.
18


Reviews:
[{"review_id": "oBrh6qQHWB", "reviewer": "Reviewer_Kvw1", "summary": "This paper represents another method of task-specific training and is one of the parameter-efficient fine-tuning approaches, referred to as HeadMap. The authors discovered that for certain tasks, specific attention heads are particularly influential; masking these heads significantly decreases performance. They propose a layer-conditioned locating algorithm to identify knowledge circuits in LLMs that greatly impact predictive accuracy. Based on this, they suggest training focused on these knowledge circuits, where only a small number of parameters are updated. The results are complementary to those of LoRA, and together they yield improved outcomes.", "questions": "Is the code available? How can we verify the results?", "limitations": "", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "It complements LoRA-type fine-tuning with utilizing significantly fewer parameters. When combined, they enhance accuracy for specific tasks and could be a valuable method for various applications.", "weaknesses": "Compared to LoRA, this approach is more challenging to utilize in practice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dO7oVz1MNl", "reviewer": "Reviewer_pHFG", "summary": "This paper attempts to explore the intrinsic mechanism of LLMs with the concept of knowledge circuits. Specifically, the termed knowledge circuits refer to the attention heads that are more important to specific tasks in all attention heads across all transformer layers. To locate the knowledge circuits, the paper explores direct strategy, simple greedy strategy and layer-conditioned strategy. The underlying assumption is that if the attention head is important, masking this attention head would cause significant increase of the task loss. After locating knowledge circuits, the paper fine-tunes the attention heads therein with LoRA-like Map modules. Experimental results evaluate the effectiveness of this design.", "questions": "1. In Figure 1 (a) and (b), it has different value ranges of attention visualization between the results on SIQA and HellaSwag. What may be the reason of that? Why there is only a light point in Figure 1(d) on Layer-0 and Head-25?\n2. What are the impacts of other attention heads beside the knowledge circuits? Can these attention heads be ablated to further enhance the accuracy and efficiency?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Knowledge circuit is an interesting concept that can help understand the mechanism of LLM performance on specific tasks.\n2. The paper proposes three strategies to locate the knowledge circuits, and two strategies to fine-tune the model with knowledge circuits, which can be inspirable for future works on knowledge circuits.\n3. Experimental results demonstrate the effectiveness of the proposed model.", "weaknesses": "1. The use of locating knowledge circuits is somewhat straightforward, and the concept of knowledge circuit was proposed be existing studies [1]. Actually, the reference [1] is missing in this paper, and it requires to make discussion with [1].\n2. The proposed method may have less efficiency and less generality. The knowledge circuit has to be located specifically for specific tasks, and then be adopted to enhance the model on the target task. The method requires to detect the 64 samples with lowest losses, and then conduct layer-conditioned locating by masking each attention head in each layer one-by-one. This procedure would have high time-consume. Besides, the knowledge circuit may hardly be transferred on different tasks. This would cause the limitation of the proposed method in application. The generality is also concerned whether the method can be used for different tasks.\n3. The selected samples with lowest losses may also be dubious. The 64 samples with lowest losses may be the 64 most easy samples for each task. Therefore, the located knowledge circuit may not be optimal.\n4. The performance improvements are somewhat marginal. This also affects the contribution of the proposed method.\n\n[1] Knowledge Circuits in Pretrained Transformers. NeurIPS 24.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lzfi4Is4xz", "reviewer": "Reviewer_37tq", "summary": "This paper introduces a novel method called HeadMap, aimed at identifying and enhancing knowledge circuits within large language models (LLMs). The authors propose a layer-conditioned locating algorithm that identifies critical attention heads layer by layer, which are essential for making accurate predictions on specific tasks. Through extensive experiments on various commonsense reasoning datasets, the authors demonstrate the effectiveness of the HeadMap method in improving model performance while maintaining parameter efficiency. Additionally, the paper explores the overlap of knowledge circuits across different datasets, revealing the varying knowledge requirements for different tasks.", "questions": "Refer to Weaknesses.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1.\tThe proposed layer-conditioned locating algorithm and the HeadMap method provide a fresh perspective on understanding and optimizing LLMs, particularly in identifying and leveraging critical attention heads, showcasing innovation.\n2.\tHeadMap outperforms random selection and simple greedy methods, and the HeadMap method achieves comparable performance to baseline models while using fewer parameters", "weaknesses": "1.\tWhile the paper shows that masking certain attention heads affects performance, it does not provide in-depth mechanistic insights into why these specific heads are critical.\n2.\tThe layer-conditioned locating algorithm may introduce bias in selecting attention heads based on the specific datasets used. If the algorithm is overly tuned to the characteristics of these datasets, it may not generalize well to other tasks or datasets\n3.\tThe experiments primarily focus on commonsense reasoning tasks. How does it perform in tasks like generative tasks, domain-specific applications?\n4.\tLack of Ablation Studies on Redundant Heads", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3CAtJHYZAe", "reviewer": "Reviewer_mNPV", "summary": "This paper ‌explores‌ the intrinsic mechanisms underlying the attention heads of LLMs, and ‌proposes‌ an algorithm to detect the important heads that ‌play‌ a critical role in LLMs. Based on these observations, the authors further ‌propose‌ a new PEFT method that only fine-tunes these important heads, and ‌verify‌ its effectiveness via extensive experiments.", "questions": "Please see the weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The overall writing is good. This paper is easy to follow.\n\nThe analysis about the  attention heads is comprehensive and reasonable.\n\nThe experimental results demonstrate the state-of-the-art performance, and the ablation study empirically proves the effectiveness of the proposed method.", "weaknesses": "The concept of the knowledge circuit is interesting, but the algorithm to find such a circuit is overly greedy. From my perspective, the selected heads are only layer-wise optimal. It seems this problem can be ‌solved‌ by dynamic programming to find a more optimal result, and ‌wouldn't‌ cost much. And then, the knowledge circuit can become a real \"circuit\".\n\nThe improvement of adopting HeadMap is not significant in Table 1, ‌especially‌ when using Llama3-8B as the LLM.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)