The Atlas of In-Context Learning: How Attention Heads Shape In-Context Retrieval Augmentation
Authors: Patrick Kahardipraja
Reduan Achtibat
Thomas Wiegand
Wojciech Samek
Sebastian Lapuschkin
Venue: NeurIPS 2025 poster ()
Keywords: feature attribution
in-context learning
function vectors
retrieval-augmented language models
transformers
explainable AI
mechanistic Interpretability
Primary area: social_and_economic_aspects_of_machine_learning
TL;DR: A method to identify specialized attention heads for in-context retrieval augmented LMs, with analyses of their roles and an application use case to trace knowledge sources.

Abstract: Large language models are able to exploit in-context learning to access external knowledge beyond their training data through retrieval-augmentation. While promising, its inner workings remain unclear. In this work, we shed light on the mechanism of in-context retrieval augmentation for question answering by viewing a prompt as a composition of informational components. We propose an attribution-based method to identify specialized attention heads, revealing in-context heads that comprehend instructions and retrieve relevant contextual information, and parametric heads that store entities' relational knowledge. To better understand their roles, we extract function vectors and modify their attention weights to show how they can influence the answer generation process. Finally, we leverage the gained insights to trace the sources of knowledge used during inference, paving the way towards more safe and transparent language models.

The Atlas of In-Context Learning: How Attention Heads Shape In-Context Retrieval Augmentation
Patrick Kahardipraja
Reduan Achtibat
Thomas Wiegand 1,2,3
Wojeiech Samek
Sehastian Lapuschkin
Depaarmeent of Artificial Intelligence, Fraunhofer Heinrich Hertz Institute
DDpartment of Electrical Engineering and Computer Science, Technische Universität Berlin
BIFOLD Berlin Institute for the Foundations of Learning and Data *Centre of eXplainable Artificial Intelligence, Technological University Dublin (firstname lastname] (aattname)@hhhl @hhhi fraunhofer.de
Abstract
Large language models able exploit in-context learning to access exter- nal knowledge beyond their training data through cerrevallaugmennatioo.. While promising, its inner workings remain unclear. In this work, shed light the mechanism in-context retrieval augmentation for question answering by viewing prompt as a composition of informational components. We propose an attribution- based method to identify specialized attention heads, revealing in-context heads that comprehend instructions and retrieve relevant contextual information, and parametric heads that store entities' relational knowledge. To better understand their roles, extract function vectors and modify their attention weights to show how they can influence the answer generation process. Finally, we leverage the gained insights to trace the sources of knowledge used during inference, paving the way towards more safe and transparent language models.
1 Introduction
Many, not most language tasks can be framed sequence to sequence problem 61 69 is integral to how modern Large Language Models (LLMs) operate, as they are able to approximate relations between an input and an output sequence not only as continuation of text, but also as response to a stimulus [64]. In a sense, input prompts serve a query to search and induce function(s) in vast, high-dimensional latent space, where the corresponding process can be cast as question
This view
answering [50] or instruction following [56] 80)
This capability brought forth with the introduction of in-context learning (ICL) [12 that enables LLMs to adapt to new tasks with few demonstrations at inference time, without additional fine-tuning. Previous work has investigated ICL from various perspectives, including its relation induction heads that can replicate token patterns during prediction [54), the ability to compress attention heads to function vectors (FVs) representing specific task [31 74). and how it can emerge when transformers (75) implicitly learn to perform gradient-basse based optimization 4 [7]. Besides meta- learning, ICL can be used for aerrivaalaaugmentaaie
[62]. where external knowledge from web
retrieval corpora [1] 29 33
or dialogue information [66 88 92 is given instead of input-output
pairs to ground LLMs during generation. However, the mechanism behind ICL for knowledge
retrieval is not yet fully understood. In this work, we aim to shed light on this question.
"quual contribution.
Corresponding authors.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Inprut Prompt:
Esecute instructions snother prompe
Mike Tason a firelighter and Guramedig/Mith the New Yora City Fire Department.
patai instrodian
Mangaret Mitchell was born Gengial She is novelist
00 What has he wolked so? He boxer firefohter
Tosk heads mncade instractions What occupation?
10
bi Change retrieved answer object modilyy al encion weichi
Layer ID
Mike Tysen a alfreefightrrjand and aparamediil [- parumedic
Retrieval heuis copy, ansserss frrrighter
20
Overwrite entitiee attributes
attritutes
Tyson Inoo last taken
Renumetrio Rutude encode attribucee
the occopation Albert Binstall'
birthultee:3 June 1968
occupation:
30
boxer
birthplace: Brooklyn
5
10 15
20
25
30
Head
Figure 1: Functional map of in-context and parametric heads in ---ma-3.1-BB-Insiii They are surprisingly well-structured and operate on the input prompt at various levels, with in-context heads processing information the prompt. including instruction comprehension and retrieval operations and parametric heads that encode relational knowledge. In-context heads can specialize to task heads to parse instructions (blue) or retrieval heads for verbatim copying (green). Together with parametric heads, they affect the answer generation process through function vectors that they transport (a, c) or their attention weights (b). Our relevance analysis (bar plot) shows that mnstrcctoonfollowee capabilities emerge in middle layers, while answer retrieval occurs in later layers. Details in C.T
On question answering tasks, we show that by viewing a prompt as a composition of informational components, certain attention heads perform various operations on the prompt at different stages of inference and layers (Figure D Our method identifies two groups of heads based on their functions: parametric heads that encode relational knowledge 27 57) and in-context heads responsible for processing information in the prompt. Further, as in-context heads need to understand which prompt components to process and how, we hypothesize that they specialize to fill their respective roles.
Our analysis shows that in-context heads can indeed execute specialized functions such instruction comprehension and retrieval of relevant contextual information. To investigate this further, we curate controlled biography dataset with entities extracted from Wikidata (78). Remarkably, we find that through compressing them to FVs or modifying their weights, both in-context and parametric heads can induce specific, targeted functions.
Building on these insights, we probe for sources of knowledge used during retrieval-augment question answering and show where it is localized within the input context. Our attempt shows promising results, serving as an initial step towards more transparent aetrieval-augmeeee generation (RAG) systems. Overall, our contributions can be summarized as follows:
We describe an attribution- method to determine attention heads that play a key role during in-context retrieval augmentation for question answering, revealing that they operate on the prompt in a distinct manner. Our method can thus be used to reliably categorize attention heads into in-context and parametric heads.
We analyze how in-context heads specialize in reading instructions and retrieving informa- tion, mapping their location across model layers. Additionally, demonstrate the influence of in-context and parametric heads on the answer generation process, by compacting them into function vectors or modifying their attention weights.
We present preliminary results on enabling causal tracing of source information for retrieval- augmented LMs, suggesting fruitful directions for interpretability of RAG systems.
2 Related Work
Retrieval Augmentation Retrieval-augmenter language models (RALMs) 40 45
address in-
herent limitations of LMs by providing external knowledge sources. Despite this advantage, issues such as discrepancies between contextual and parametric knowledge may occur 48 87 Some

---

works have studied mechanisms behind knowledge preference in RALMS ST 55 91], but they focus on simple domains. On the other hand, those that explored more complex domains mostly only Closely related to our work is that of Jin al. [37], where the authors also discovered in-context and parametric heads. Compared to them, we show that interactions among in in-context heads are subtly more complex, since they can specialize into task or retrieval heads. Our approach also allows to demonstrate their roles in shaping the model's representation along with parametric heads, by transforming them into FVs or modifying their respective weights.
analyzed RALMs' behavior at the output level [4
71 85
Another disadvantage of RALMs is that they cannot guarantee faithful answer attribution to con- textual passages [25], which necessitates a shift to interpretability. Recent efforts address this include leveraging contrastive gradient attribution [60] 89), fitting surrogate models [16], and treating attention weights as features [17]. In relation to this, we train a probe based on retrieval heads to track for knowledge provenance.
The Role of Attention Attention mechanisms have been previously observed to encode many kinds of information. Clark et al. [15] showed that they correspond well to linguistic properties such as syntax and coreference. Similarly, Voita et al. (76) found that syntactic heads play an important role in machine translation models. In relation to world knowledge, Geva et al. [27] proposed that certain heads store factual associations and demonstrated how they extract an attribute of a given subject- relation pair. Interestingly, attention also appears to encode sense of "truthful" directions [46]. With the exception of Voita et al. [76). the above works make use of attention weights, which might not fully capture the model's prediction [10 35 81. Our work can be seen as an attempt to reconcile both perspectives: analyses based on attention weights and feature attribution methods (9]
In- In-Context Learning Numerous works have studied ICL since its introduction. Liu et al. 47 studied what constitutes a good example for demonstrations. Dai al. (20) suggested that ICL can be understood as an implicit fine-tuning. ICL is a general phenomenon, although it is commonly assumed to be unique to autoregressive models [63]. At the component level, ICL is primarily associated with induction heads [21]54] However, recent findings showed that certain heads can also be compressed form FVs that represent a demonstrated task [31 74). Yin and Steinhardt [90] investigated the connection between these heads and induction heads, showing that they are distinct from each other and how models' ICI. performance is mainly driven by the former. ICL. can also be viewed as a mixture of meta-learning and retrieval (22. In that regard, we study the latter perspective to understand its mechanism as a specific instantiation of ICL. with a focus on the retrieval augmentation paradigm.
3 Background and Preliminaries
The self-attention mechanism in transformers poses a challenge in understanding which heads actually contribute during in-context retrieval augmentation, and how they process various components in the prompt. This is mainly due to the fact that information from different tokens gets increasingly mixed
as layers go deeper and how several attention heads may implement redundant functions [79]. A natural option is to analyze attention weights, as they are an inherent part of a model's computation. However, attention can be unfaithful [34], which questions its validity as an explanation [10 65).
This problem is further exacerbated by "attention sinks" [3 84] - a phenomenon where heads heavily attend the first token and obscure the weights of other tokens in the sequence.
An alternative would be to use feature attribution methods [9, as they trace the contribution of each input element to the model's output prediction. Propagation-base based feature attribution 88 67 68 especially takes the entire computation path into account, which can be used to characterize attention heads 766 or identify latent concepts []]. Furthermore, feature attribution is able to estimate causal and thus enables to observe how
relations 26 e.g., to automate circuit discovery [8 22
30 32
specific attention head affects model's prediction.
In this section, we provide a description of AunLRP 2] on which our method based, due to its superior performance and efficiency in transformer architectures compared to other attribution methods. We also provide an overview of the multi-head attention mechanism in transformers, which
Here, the term answer attribution means the use of external documents to support the generated response. which is different from feature attribution used throughout this work to describe interpretability techniques.
3

---

we leverage through AttnLRP to identify both in in-context and parametric heads ($5). Additionally, we analyze the specialization of in-context heads, show causal roles of the identified heads ($6), and use this information for reliable and efficient source tracking of facts in rerrieval-augmented I.Ms ($7).
3.1 Layer-wise Relevance Propagation
Feature attribution methods aim to quantify the contribution of input features the overall activation of an output y in linear but also highly non-linear systems. We define a function R that maps the input activations to relevance scores indicating their causal effect model's 's output logit y:
R: R-R.V R R.A.
(x. y) R(x y)-
In principle, any feature attribution method can be employed for R. though trade-off's between faithful- ness and computational efficiency must be carefully considered. Perturbation-based approaches [49] typically offer high faithfulness but incur exponential computational costs, as each ablated latent component requires at least one dedicated forward pass [30]. In contrast, gradient-based methods [67] are computationally more efficient, requiring only a single backward pass, which makes them well suited for large-scale interpretability studies. However, they are susceptible to noisy gradients, which are distorted by non-linear components such layer normalization [6]. To address these limitations, we adopt AttnLRP 22, an extension of Layer-wise Relevance Propagation (LRP) [ designed for transformer architectures. As a backpropagation--usee technique, AttnLRP propagates relevance scores from model output to its inputs in a layer-wise manner and can be implemented efficiently via modified gradient computation in a single backward pass. Importantly, it incorporates stabilization procedures for non-linear operations, thereby improving the faithfulness of relevance distributions compared to standard gradient- or other decomposition- based methods 77].
Relevance scores produced by (AP can be either positive or negative. Positive relevance indi- cates an amplifying effect on the final model logit y. whereas negative relevance reflects an inhibitory influence. Without loss of generality, we focus our analysis on Signal-amplifiin components by considering only positive relevance scores. Formally, we define:
R* (xly) max (R(xly).0
(I)
This yields clearer separation between in-context and parametric heads in the subsequent analysis.
3.2 Attention Mechanism
While the original formulation of the multi-head attention mechanism 755 concisely summarizes the parallel computation of attention heads, our goal is to isolate their individual contributions. To this end, we reformulate the equations make the influence of each head more explicit [2]] 23). Let X = (X1,... xs) E d denote the matrix of hidden token representations for sequence of
length S with dimension d, and suppose our model employs H parallel heads, each of dimension d, = d H. Then, the computation of the multi-head attention layer can be reformulated into H complementary operations, where each head produces an intermediate attention output z E Rda:
zf
Σ Alt (W xx))
where A is the attention weight of token attending to token j, and W dd xd is the pead value projection. The final output is obtained by multiplying the intermediate output of each head with their corresponding output projection matrix W E Rdxd, followed by summing:
H
(3)
h=1
We leverage the multi-head attention mechanism in transformers through the lens of AttnLRP to identify both in-context and parametric heads in and how in-context heads specialize in {6
4 Experimental Setup
Models We use instruction-tuned I.IMss due to their increased capability on question answering (QA) tasks in our preliminary experiments: --sttrrttttt (28). - Instrruct--O.. (36),

---

and Gemma-2-9B-1 (72). We apply AttLLR based on their huggingface implementations [2] For the rest of this work, we refer each model by their family prefix.
Datasets To perform our analyses, we use two popular open-domain QA datasets: NQ-Swap 48 and TriviaQA (TQA) 38] NQ-Swap is derived from Natural Questions 44 QA dataset with questions collected from real Google search queries and answer spans annotated in Wikipedia articles. TQA contains trivia questions with answers sourced from the web. Both datasets are based on the MRQA 2019 shared task version (24]
Similar to Petroni al. [88 consider different types of contextual information to see how they affect in-context and parametric heads. We use oracle contexts as they are always relevant to the question and contain the true answer. In addition, we use counterfactual contexts as they contain information that is usually not seen during pre-training and fine-tuning stages, thus forcing models to rely on the text to answer the question correctly. Oracle context is often not available; therefore also use Dense Passage Retriever (DPR) (39] with a Wikipedia dump from December 2018 as our retrieval corpus. For simplicity, we only select the top one retrieved document. We show results for oracle and counterfactual contexts in the main paper and retrieved DPR contexts in Appendix B
Inspired by Allen-Zhu and Li 66. we build a human biography datasets to allow to better understand the characteristic of in-context and parametric heads and conduct controlled experiments. Using Wikidata [78). collect profiles for random 4.255 notable individuals containing their date of birth, birth place, educational institute attended, and occupation. We concatenate the attributes of each individual in random order form biographical entry and ask Llama-3.1-80 --8--Instrue to paraphrase it. See Appendix A for more details
5 Localization of In-Context and Parametric Heads
In retrieval-augiee -ugmeente generation, L.MS are faced with the option to generate responses by using a wealth of knowledge they learn during pre-training or by relying on contextual information provided in the prompt through ICL. Here, we categorize attention heads that are responsible for both capabilities.
Method We aim to identify the sets of in-context heads Hexx and parametric heads Hparam as depicted in Figure T We define in-context heads those that mainly contribute to the model's 's pre- diction during RAG by using contextual information, whereas parametric heads primarily contribute upon reliance on internal knowledge. We hypothesize that each head type contributes maximally under a specific condition while having minimal influence in others, i.e., in-context heads are expected to contribute the most open-book settings and the least in closed-book settings, and vice versa. We analyze questions with counterfactual contexts, forcing retrieval to produce a counterfactual prediction Yer that disagrees with the parametric answer. Conversely, we also focus on closed-boce -book settings where contextual information is minimized, to identify parametric heads and reduce the chance that in-context heads contribute. We restrict our analysis to instances where gold reference answer Mgoodd is predicted, to ensure that relevance attribution reflects genuine parametric behavior.
We use Att to quantify the contribution of each attention head h to the prediction by summing the positive relevance scores assigned to its latent output z" across its dimension d, and over all i token positions, when explaining the targeted prediction yt, which can be either a gold reference answer Ugoid or a counterfactual output Ver, depending on the setting:
$ dn
RM()
R (z)
wt)s E
To contrast heads across settings, we compute a difference score D representing their average contribution in open-book versus closed-book conditions for all N₂, heads in the model:
D Exx (R*())) Excs (Ugola))
1, 2,....NNp)
We then identify the most distinctive heads for each behavior by selecting the top 100 heads (around 10%-15% of total heads) with the highest positive and lowest negative difference scores:
Hes {argsort Jesc (D)) 100
Hparam = {argsort 201 (10
100
(6)
*Our implementation is publicly available htt.ttss: http:///////gggoooooaaaaaaaaa/aa
5

---

AttuLRP (In (In-conlext)
AMtaLRP (Parametric)
AWR
Ramdom
NQ-Swap (Counterfact)
NQ-Swap (Oracle)
NQ-Swap (Closed-book)
TQA (Oracle)
TQA (Closed-book)
100
80-
Recall
60-
60-
20-
20
40 Go
80 100
20 40
so 100
20 60 80 100
20 40 GO 80
20 40 so 100 Heads
Num Heads Ablated
Heads Ablated
Num Heads Ablated
Num Heads A Ablated
Figure 2: Recall analysis for Llama 3.1 when either in-context parametric heads ablated. Removing identified in-context heads noticeably affects the model's 's performance in open-book QA across various configurations. Conversely, removal of identified parametric heads most strongly affects the model's s closed-book QA capabilities. Compared to Wu et al. [83] that only yield AWR (retrieval) heads, our method allows obtain both in-context and parametric heads.
Experiments To ensure that the identified in-context and parametric heads play a role in QA tasks, we ablate both sets of heads and measure performance drops in settings where they are expected to be mostly active (open- and closed-book book, respectively). We also measure if the removal of in-context heads affects the models' capabilities to answer in closed-bood setting and vice versa, since this informs to what extent both sets of heads are dependent on each other. Furthermore, we want to know if the identified in-context and parametric heads can generalize to other datasets. To test this, compute the score of both heads only over NQ-Swap and reuse the same sets of heads on TQA. To evaluate the aforementioned criteria, report recall 3 as instruction-tune -annudteon---eed models tend to produce verbose responses. As baselines, we select random heads, and also adopt the Attention Weight Recall (AWR) method based on attention maps' activations, described in 83
Results We show results for Llama 3.1 here and other models in Appendix Figure shows how the recall score evolves when either heads identified as in-context parametric heads are ablated. We observe that the removal of 20 heads (100 heads) reduces the performance by 13.86%-63..8 144.26%-68.66%% for open- and closed-book. settings across different configurations, indicating the causal influence these heads have on the answers' correctness. Moreover, the performance drops on TQA hold even though the heads are computed NQ-Swap, showing that the identified in-context and parametric heads are transferrable to other datasets.
We compare in-context heads identified with our method against AWR heads, and find that the removal of 20 in-context heads results in a roughly similar reduction recall as removing 100 AWR heads. Furthermore, ablating in-context heads yield more drastic performance decrease compared to the removal AWR heads, suggesting that our method is more suitable than those based on attention scores alone study heads that contribute to response generation. Ablating randomly chosen heads barely affects the model's 's ability to answer correctly.
We examine whether in-context and parametric heads are independent of each other. As expected, ablating parametric heads has little influence to the model's performance our open-book setting. Interestingly, this leads to slight performance increase on NQ-Swap with counterfactual contexts, which suggest that the ablation forces the model to rely more on the given context instead of its own parametric knowledge. Surprisingly, ablating in-context heads in the closed-book setting incurs non-negligible performance reduction. This is likely due to the influence in-context heads have when processing the input prompt. We explore this in S
5.1 Resolving Knowledge Conflicts
We further test whether targeted ablation of parametric and in-context heads also helps to resolve conflicts between contextual and parametric knowledge sources. On NQ-Swap with counterfactual contexts, we selectively ablate heads from each category. Removing in-context heads forces the model rely more heavily on parametric memory, leading to substantial gains in recall with respect to the original, con-counterfactual answers. Conversely, ablating parametric heads promotes stronger

---

grounding in the provided context as shown before, improving recall for counterfactual open-book answers. We present the full results in Appendix
book
6 Functional Roles of In-Context and Parametric Heads
Given that ablating in-context heads yields a non-negligible drop in closed-book QA performance, where external documents are available, we posit that in-context heads not only process the context but also interpret the intensional frame - the semantic structure imposed by the instruction itself (4). In the counterfactual example below, the intensional frame (the question prompt) is shown in italics, the object instance in bold, and two equally plausible answers in color:
"[Mike Tyson was firefighter from 1980 1984 with the New York City Fire Department... Q: What has Mike Tyson worked as? A: boxer / firefighter"
To answer correctly, the model must map the intensional frame onto the knowledge triple (s, r, o*). where s is the subject ("Mike Tyson"), "), r is the semantic relation (the predicate specified by the question, here "has worked as"), and o* is the object (the yet to be determined answer, "boxer" or "firefighter"). Depending on where the answer resides, o" may be retrieved from the model's paramet- ric memory (o") or from the context (o9). By treating (s, r, r.)") as the complete task specification, we analyze how in-context and parametric heads specialize both to comprehend the intensional frame and to retrieve the object needed to generate the correct answer.
6.1 Disentangling the Functionality of In-Context Heads
1.2
Other Heads
1.0
Task Heads
soridwe (In-Contes Score)
Retrieval Heads
0.8
0.6
0.4
0.2
0.0
-0.2
-0.4
0
200 400
600
800
1000
0
10
Sorted Index
Count
Figure 3: Sorted in-context scores for 1024 heads of Llama 3.1, comparing open-book and closed- book settings via score D. Positive scores indicate in-context behavior, while negative scores reflect parametric behavior. Retrieval heads (green) and task heads (blue) are predominantly high-scoring in-context heads. See Appendix Figure for other models.
Our goal to identify heads specialized in processing the intensional frame and those specialized in retrieving the answer object from the context. Inspired by the work of (1, which demonstrates that relevance is effective for separating functional components in latent space, we measure how much relevance of attention head is assigned to the question and retrieved answer tokens.
frame
Method For each head h, we compute the total relevance attributed to the attention weight Ab, when explaining the logit output yt- Since relevance flows backwards from the output to the input, our goal is to obtain relevance at the input level of each layer. Given that each head transfers information from the key at position j to the query at position i, we aggregate this backward relevance all possible query positions i to obtain a single relevance score for the source token at key position j:
Σ
R
(Ais wt)
Here, of represents the total relevance assigned at head h to token j when contributing to logit yr- Next, we aggregate the relevance scores separately for two sets of token positions within the context:
7

---

Table 1: Zero-shot recall scores for task, parametric, and retrieval heads.
Pater Higgs was born on 29 May 1922 in inspire fature penerrations of acientist and peweard 0:A4 which in inst lution e was be PPA
Mila Kunis was born in the picturesque city Chernivtsi, place that of Hollywond's most beloved stars University of Calliooriaa
Neweastle upon Tyme
pasition
Models
Husk
H.
Hist
Llama 3.1 (random)
18.00
6.68
15. 15.94
dal Exrrationn and Insertion Task PVs
FVs/ / Attn Weight
94.75
38.84
93.45
Albert Was born
Q What the acmupation of Mike Iyson?
Mistral v0.3 (random)
9.50
12.95
8.56
physleist
FVs/ Attn Weight
88.50
44.04
97.03
06 Extraction and Iisection of Paramettee TV
Gemma 2 (random)
7.50
6.79
3.89
Figure Extraction and insertion of task and parametric FVs. The induced generation is highlighted in italic.
FVs/ / Attn Weight
88.00
34.77
87.36
the intensional-frant tokens, denoted as j E Jusk, which comprise the question token positions, and the answer object tokens, denoted as j E Jree, which represent positions of the retrieved object.
Prask
Σ
Piet
Σ
(8)
JEJua
JE
Finally, to obtain the sets of specialized task and retrieval heads, we rank heads by their aggregated relevance and select the top K, a hyperparameter determined separately for each experiment.
H.K = (argsort(Pinskaas I H = (argsort(Pffa)dee ) n=
K
K
(9)
1=1'
Results We compute the task relevance score Plack and the retrieval relevance score Phcc over NQ- Swap with counterfactual contexts to minimize influences of parametric heads, and aggregate their distributions across the model layers. In Figure I we observe that Pisk initially increases in the early layers where few parametric heads are located, suggesting that early parametric heads enrich the question with relational knowledge. The relevance peaks in the middle layers, where in-context heads dominate, aligning with the transition to a more context-dependent reasoning. In contrast, the retrieval relevance score Phe peaks in deeper layers, reflecting the point where the model extracts the final answer object o Figure 3 further illustrates the sorted average difference D between open-book and closed-book settings for all heads, alongside the top 40 task heads Hask and retrieval heads H We observe that the highest-scoring in-context heads are primarily composed of retrieval and task heads, emphasizing their role for retrieval augmented generation.
6.2 Causal Effects of In-Context and Parametric Heads
An important question whether the heads we identify truly reflect their assigned functionalities. We examine this under controlled conditions and investigate their causal effects on the answer generation.
Experiments We conjecture that task heads Hask encode the intensional frame (s, r, o²) and that parametric heads H Hparrm contain information of the subject which depending on the training data may or may not include of. On the other hand, retrieval heads Hact search for 0f, allowing them to copy any tokens from the context verbatim. without being restricted only plausible answers. For task and parametric heads, we compute FV for each head and insert them into various settings to trigger the execution of their functions. Following Wu al. 83), we opt for a needle-in-s in-a-haystack (NIAH) setting for retrieval heads and determine their ability to retrieve relevant information from the context by modifying the attention weights. To isolate these behaviors, on the biography dataset ($4) and measure recall [3]. For comparison, we also consider random heads for FV extraction and attention modification. See Appendix Cfor additional results and details, including the selection of hyperparameter K.
conduct our analysis
Task Heads We demonstrate that task heads encode intensional frames. In zero-shot manner, extract task FVs from each head in Hask for four questions relating to all recorded attributes from the biography dataset. Then, we insert them to another biographical entry without question at the final token position, and also for all subsequent token generations (Figure top). We examine whether
A FV can be defined as a sum of averaged heads outputs over a task 774 or computed individually [31 Following the latter, consider FV to be output of task parametric head scaled by scalar a.
8

---

they reliably induce responses aligned with the original question. In Table I (left), we show that applying FVs in a zero-shot manner allows all models respond accordingly wrt. the intensional frame, yielding an average improvement of 78.75 75 points over random heads.
Parametric Heads Parametric heads contain relational knowledge. To show this, first select random attribute of an individual and convert it to cloze-style statement. Then, we extract FVs from Hparam. which are inserted We observe if the generated response contains information of the original entity conditioned on the intensional frame. For simplicity, we restrict extraction to cases where the closed- book answer is correct wrt. gold reference. We see that in Table [ (middle), adding parametric FVs allows all models recover the original attributes significandly, with an increase of 30.41 points compared to random.
question prompt of another unrelated individual (Figure 4 bottom).
Retrieval Heads We assess retrieval heads' ability to copy verbatim by using famous poem titles needles, inserted at random position in the biographical entries. At the last token of the entry and for the following generations, we increase the attention weights of all heads in Hra on every token of the needle to force the model to copy. Our results (Table I right) show drastic increase of 83.15 points over the random baseline, indicating that retrieval heads are indeed able to perform copy operations independent of the token position.
6.2.1 How Does This Generalize to Other In-Context Tasks Beyond QA?
To assess whether our method generalizes beyond question answering, we conduct a preliminary investigation on machine translation using the OPUS Europarl dataset (31 Directly transferring the task heads identified in the QA experiments proved ineffective, indicating that translation relies on different functional subspace. Nonetheless, by applying our disentanglement procedure of $ to the intensional frame of translation prompt, we are able to identify compact set of 15 attention heads that robustly encode translation behavior independent source or target language. Patching these heads induces zero-shot translation across multiple language pairs, yielding BLEU scores comparable to explicit prompting whereas patching random heads fails to induce any translation behavior. Complete results and experimental details are provided in Appendix F
7
Source Tracking
Fine- Grained Detection of Knowledge Source
Causal Source
Source-Tracking Tracking of Retrieved Documents
overlay RAG context
Question: Where does llama origlinte from?
Retrieval Heads' attention maps at last query position
Answer: Llama from...
Wh
0.5
South America
0.4
Africa
Head2
10
South Americcs
Body
W/
UMP Oevunie
Vostorday; an opon souroe language model was reloused dvannc open research It was trained on public datasets, such as Common Cravi and Gititub reponitories, and is available four purameter sizes 78, 13B, 33B, and 65B. This model, built debuted on Pebruary 23, 2023. It has demionstrated impressive performance, with the 198 version surpussing GPT-3, despite its smaller size. Distributed under the 3 Hinnns, provides researebers with full aesesa the madel's wrights.
Meta
LLaMA (Large Lan gunguage Model
framework, which
Figure 5: (a) When asked "Where does Ilama originate from?", the retrieval-head probe classifies "South America" and "Africa" as parametric, while "Meta" as contextual. The UMAP projection of retrieval head activations displays the linear probe's 's decision boundary (dashed line) separating parametric from contextual clusters. (b) The weighted aggregation of retrieval head attention maps at the final query position is superimposed on the document to pinpoint the retrieved source span.
Our experiment in $6.2 demonstrates that retrieval heads reliably perform verbatim copying of text spans when their corresponding attention maps focus on the retrieved tokens. As can be seen on Figure we now aim to investigate if we can (i) detect when retrieval heads initiate the copying process for the first answer token (i.e., whether token is derived from external contexts rather than from the model parameters), and (ii) accurately localize its position within that context using the attention maps. To this end, we train linear probe on NQ-Swap with counterfactual contexts. Each

---

retrieval head's output at the last token's 's position zó is decoded via logit lens [53] converting each head's 's activation into a score for token & N using the model's unembedding matrix We E [R and layer normalization LN(-):
L 1)) = LN(WW) d3) Woll E
(10)
where O E RRddd is the head's output projection and Worl the row of the unembedding matrix corresponding to token t. As such, L(z's t) computes how strongly head writes token t into the residual stream. In Appendix Figure 8 we illustrate histograms of the logit lens scores.
Next, we train probe via linear regression, yielding the weights {wn, Je For source localization, we aggregate each head's attention map using its weight and logit-lens score, and predict the source token index as
k arg max
Σ wh L (z's ()) Al...
(11)
For details, please refer to Appendix D Additionally, we use standard AttnLRP backward pass from the model output to compute an input heatmap baseline for comparison.
In Table the retrieval-head probe achieves an ROC AUC of at least 94%, reliably distinguish- ing contextual from parametric predictions and thus confirms linearly separable representation of the retrieval task. A promising direction for future research is to leverage the probe's 's abil- ity to distinguish between parametric and con- textual predictions. enabling dynamic control over the model's token selection. This approach could reduce hallucinations by explicitly guid- ing the model to prioritize context over paramet- ric memory when appropriate. In addition, each model attains a top-1 localization accuracy of at least 84%. In Appendix Figure(9) 99] we illustrate heatmaps of the aggregated attention maps superimposed on the input, highlighting the positions of the predicted tokens. While AttLLRP outperforms the probe, it requires an additional backward pass increasing computational cost, while the probe only requires attention maps computed during the forward pass.
Table 2: Performance of the retrieval- hhed probe across models.
Localization
Models
ROC AUC
Attention
AttLIRP
Llama
95%
97%
98%
Mistral v0.3
98%
96%
99%
Gemma 2
94%
84%
96%
Conclusion
We propose a method to explore the inner workings of ICL for retrieval-uggmeentttonn, revealing in-context and parametric heads that operate the input prompt in distinct manner and find that in-context heads can specialize into either task or retrieval heads, depending on whether they encode intensional frames or retrieve relevant information. We study the roles of the identified heads by converting them into FVs or modifying their weights, showing how they can affect the generation process. Finally, we present probe precisely and efficiently track for knowledge provenance, opening up a path towards more interpretable retrieval- auggeeted LMs.
Limitations We focus our investigation on attention heads since they are primarily associated with ICL. However, how they interact with components in MLP modules e,g., knowledge neurons [19 to induce functions remains an open question. Our analyses are also mainly centered on QA (with minor part on MT). It would be interesting to see if similar mechanisms arise in other tasks. We find that several heads exhibit inhibitory effects, and that some intermediate heads are not exclusively specialized for either in-context or parametric roles both of which warrant further investigation. In addition, there is possibility for redundant heads [79]. which are yet to be uncovered. We leave this avenue for future work.
Broader Impacts Our research enhances trust in retrieval- auggmentdd LMs by elucidating the mechanisms through which they access and use external knowledge. Furthermore, it enables precise source attribution, allowing users to trace the origins of the information leveraged in response generation. However, we caution against its potential for misuse, such as using the identified heads to induce malicious behavior.
10

---

Acknowledgements
We thank the anonymous reviewers for their critical reading of our manuscript and their insightful comments and suggestions.
References
[1] R. Achtibat, M. Dreyer, Eisenbraun, S. Bosse, T. Wiegand, W. Samek, and Lapuschkin. From attribution maps to human-understandable explanations through concept relevance prop- agation. Nature Machine Intelligence, 100):000-1019, Sep 2023. ISSN 2522-5839. doi: 10. (008//422556023-0000
URL
(038/542256.2 001442266022-0071111 023-00711-8
[2] R. Achtibat, S. M. V. Hatefi, M. Dreyer, A. Jain, T. Wiegand, S. Lapuschkin, and W. Samek. At- mLRP: Attention-aware layer-wise relevance propagation for transformers. In R. Salakhutdinov, Z. Kolter, K. Heller, A. Weller, N. Oliver, J. Scarlett, and F. Berkenkamp, editors, Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 135-168. PMLR, 21 Jul 2024.
[3] V. Adlakha, BehnamGhader, X. H. Lu, N. Meade, and Reddy. Evaluating correctness and faithfulness of instruction-folloiing models for question answering. Transactions of the Association for Computational Linguistics, 12:681 1999, 2024. doi: 10.1162/tacl. a 00667.
URL htap:::::///claahhoollll// wrg.y/rrg/22//// wttp://2..aaal11....2
[4] E. Akyürek, D. Schuurmans, J. Andreas, Ma, and D. Zhou. What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations, 2023. URL http:: h/tter:///perrvvee..o www.:/ww.eeeffooorrooo 1d=0gOX4H8yN44I
[5] A. Ali, T. Schnake, O. Eberle, G. Montavon, K.-R. Müller, and L. Wolf. XAI for transformers: Better explanations through conservative propagation. In K. Chaudhuri. S. Jegelka, L. Song, C. Szepesvari, G. Niu, and S. Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages www.:mees...m/ee.mm..m..
435 451. PMLR, 17-23 Jul 2022. html
mlr wwwpp//ssseee/// v12// ali22a
[6] Z. Allen-Zhu and Y. Li. Physics of language models: Part 3.1, knowledge storage and extraction, 2024. URL BtRLhttps:/// http:::://xxxvv.rrpppp wttp://..rrrgba//////// 14316.
[7] L. Arras, B. Puri, P. Kahardipraja. S. Lapuschkin, and W. Samek. A close look decomposition- based xai-methods for transformer language models, 2025. /ttpps:///// orwpxxiiw.rrraeeeoo / 2502. 15886
[8] S. Bach, A. Binder, G. Montavon, F. Klauschen, K.-R. Müller, and W. Samek. On wixel-wise explanations for non-linear classifier decisions by layer-wise relevance propagation. PLOS ONE, 10(7):1-66 07 2015. doi: 10.1371/journal pon.:.000000000 0130140. URL https://:///ai org/10 1371/journal pone. 0130140
[9] J. Bastings and K. Filippova. The elephant in the interpretability room: Why use attention as explanation when we have saliency methods? In A. Alishahi, Y. Belinkov, G. Chrupata, D. Hupkes, Y. Pinter, and H. Sajjad, editors, Proceedings of the Third BlackboxNLP Workshop on Analyzing and Interpreting Neural Networks for NLP, pages 149-155, Online, Nov. 2020. Association for Computational Linguistics. doi: 10. (0.188633v11/ 122.000 https://aclanthol
URL
org/202/ blackboxmlp-1.114
[10] A. Bibal, R. Cardon, D. Alfter, Wilkens, X. Wang, T. François, and Watrin. Is atten- tion explanation? an introduction to the debate. In S. Muresan, P. Nakov, and A. Villav- icencio, editors, Proceedings of the 60th Annual Meeting of the Association for Computa- tional Linguistics (Volume 1: Long Papers). pages 3889-3900 Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: 1/653/v1/002 aocgo... URL https://aclanthology org/ 202//2222.a02 acl-long..........l..lmoaamlll..
llogg. 269/
11

---

[11] S. Borgeaud, A. Mensch, J. Hoffmann, T. Cai, E. Rutherford, K. Millican, G. B. Van Den Driess- che. J.-B. Lespiau, B. Damoc, A. Clark, D. De Las Casas, A. Guy, Menick, R. Ring. T. Hennigan, S. Huang, I.. Maggiore, C. Jones, A. Cassirer, A. Brock, M. Paganini, G. Irving, O. Vinyals, S. Osindero, K. Simonyan, J. Rae, E. Elsen, and L.. Sifre. Improving language mod- els by retrieving from trillions of tokens. In K. Chaudhuri, S. Jegelka, I.. Song, C. Szepesvari, G. Niu, and S. Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 2206-2240 PMLR, URL https::/pprceeeding...e
17-23 Jul 2022.
mlr. prwps/ sacess/vv622/ http///v1//ooooree ht02//ooggaud22:: html
[12] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, Shyam, G. Sastry, A. Askell, Agarwal, A. Herbert- Voss, G. Krueger, Henighan, R. Child, A. Ramesh, D. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei. Lan- guage models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances Advances in Neural Information Processing Systems. volume 33. pages 1877- 1901. Curran Associates, Inc., 2020. URL httpp:::ppoocceeee//ttt/ neurips. sc/paper_ Paper .dd
[13] J. Bulian, C. Buck, W. Gajewski, B. Börschinger, and T. Schuster. Tomayto, tomahto. beyond token-level answer equivalence for question answering evaluation. In Y. Goldberg, 7.. Kozareva, and Y. Zhang, editors. Proceedings of the 2022 Conference on Empirical Methods in Nat- ural Language Processing, pages 291-305, Abu Dhabi, United Arab Emirates, Dec. 2022. Association for Computational Linguistics. doi: 10. 10.18653/v1/200 htt.:/.emmlp.mmmmmeeeea .wm.lpmmainnn000 URL
https://acee http::///aaaatholog.o.o org/ 2022. emnlp-maing 20/
[14] H.-T. Chen, M. Zhang, and E. Choi. Rich knowledge sources bring complex knowledge conflicts: Recalibrating models to reflect confficting evidence. Y. Goldberg, Z. Kozareva, and Y. Zhang, editors, Proceedings of the 2022 Conference Empirical Methods in Natural Language Processing, pages 2292-2307, Abu Dhabi, United Arab Emirates, Dec, 2022. Association for Computational Linguisties. doi: 10.1853//12022.mmmlll
146. URL https://
aclan aclanthology 2022 e11222.mmnlpp. emnlp-main. 146/
[15] K. Clark, U. Khandelwal, O. Levy, and C. D. Manning. What does BERT look at? an analysis of BERT's attention. In T. Linzen, G. Chrupata, Y. Belinkov, and D. Hupkes, editors, Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pages 276-286, Florence, Italy, Aug. 2019. Association for Computational Linguistics. doi: 10.18653/v 1/t1///211999222 URL https:://cclattttt/t
org/ Httg//1--828//////
[16] B. Cohen-Wang. H. Shah, K. Georgiev, and A. Mądry. Contextcite: Attribut- ing model generation to context. In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang, editors, Advances in Neural Infor- mation Processing Systems, volume 37, pages hwtp//ppeeeflee/ppeeeeeeeeeeeeee 2024. URL httpps:/proocediing wtwpeeeeediinnnnnnuunneeeeeeeeee cttpripppccppppee_frreeeee 1abbal1622966d066699444444443433 Poapraaoe---ppape--o-oeoeeeeee
Curran Associates, Inc..
Conf Conference.,r..o .df]
[17] B. Conen-Wang, Y.-S. Chuang, and A. Madry. Learning to attribute with attention, 2025. URL
https://arxiv: www.oogg/////.o..oooooooo. absp/504.//...... 13752
[18) A. Conmy, A. Mavor-Parker, A. Lynch, S. Heimersheim, and A. Garriga-Alonso. To- wards automated circuit discovery for mechanistic interpretability. In A. Oh, T. Nau- mann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neu- ral Information Processing Systems, volume 36. pages 16318-16333 Curran Associates, Inc., 2023. URL http:tttst//prroedinnggiee http:///pproeeddingneeneiieieeep cc httpprrpppppppppppp..ee////eeeep http:ppperr.lleeeppprr///2222222 ptt.ll///ppaprr2////ieeee Pap.aaaaeaeaaaaaaaaaaeaeaeeeeeee Conferee-o pdff]
[19] D. Dai, L. Dong, Y. Hao, Z. Sui, B. Chang, and F. Wei. Knowledge neurons in pretrained transformers. In S. Muresan, P. Nakov, and A. Villavicencio, editors, Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Pupers), pages 8493-8502, Dublin, Ireland, May 2022. Association for Computational Linguistics. doi:
long
hrtp:/antntooooooooo//22222/22 www. aa0022..ccll11 acl-long
581/
12

---

[20] D. Dai, Y. Sun, L. Dong, Y. Hao, S. Ma, Z. Sui, and F. Wei. Why can GPT learn in in-context? language models secretly perform gradient descent as meta-optimizers. In A. Rogers. J. Boyd- Graber, and N. Okazaki, editors, Findings of the Association for Computational Linguistics: ACI. 2023, pages 4005 4019, Toronto, Canada, July 2023. Association for Computational Linguistics. doi: 10. (1883665112000 findings-acl.244 URL 2023. findings- acl 247/
org/
[21] N. Elhage, N. Nanda, C. Olsson, T. Henighan, N. Joseph, B. Mann, A. Askell, Y. Bai, A. Chen, T. Conerly, N. DasSarma, D. Drain, D. Ganguli, Z. Hatfield-Dodde D. Hernandez, A. Jones, J. Kermion, L. Lovitt, K. Ndousse, D. Amodei, T. Brown, Clark, Kaplan, S. McCandlish, and C. Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. 1. https:/ttaastommerriiccisss.pppp pwb/2021
[22] Ferrando and E. Voita. Information flow routes: Automatically interpreting language models at scale. In Y. Al-Onaizan, M. Bansal, and Y.-N. Chen, editors, Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 17432-17445, Miami, Florida, USA, Nov. 2024. Association for Computational Linguistics. doi: 10.18653/ v1/2024.emnlp-minn 965. httpttt://ccaatthhohoea
org org/2024, 2024, emnlp-main. 9657
[23] Ferrando, G. Gállego, and M. R. Costa-jusså. Measuring the mixing of contextual in- formation in the transformer. In Y. Goldberg, Z. Kozareva, and Y. Zhang, editors, Proceed- ings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 8698 (698-8714, 87, Abu Dhabi, United Arab Emirates, Dec. 2022. Association for Computational org/
Linguistics. doi: 2022.
mt1.:22cmmmmmmii.5955 URL htttps://aaaatttoooot
www.aaiin... 5951
[24] A. Fisch, A. Talmor, R. Jia, M. Seo, F. Choi, and D. Chen. MRQA 2019 shared task: Evaluating generalization in reading comprehension. In A. Fisch, A. Talmor, R. Jia, M. Seo, E. Choi, and D. Chen, editors, Proceedings of the 2nd Workshop on Machine Reading for Question Answering, pages 1-13, Hong Kong, China, Nov. 2019. Association for Computational Linguistics. doi: (0.18653/v1DD199 19-5801. URL httppp::/acclaaooooooooooo//////
[25] T. Gao, H. Yen, J. Yu, and D. Chen. Enabling large language models to generate text with citations. In H. Bouamor, J. Pino, and K. Bali, editors, Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 6465-6488, Singapore, Dec. 2023. Association for Computational Linguisties. doi: 10.1 https://aclanthen
URL
org 2rg/2023.... emnlp-main. 398/
[26] A. Geiger, H. Lu, T. Icard, and C. Potts. Causal abstractions of neural networks. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 9574-9586. Curran Associates, Inc., 2021. URL httpp::::ppocceetdeitteeeee neur wiw..rrp...... vc/paper__a f files/ ftt.ee/s///appr/2////iiee pddf
[27] M. Geva, J. Bastings, K. Filippova, and A. Globerson. Dissecting recall of factual associations in ruto-regressive language models. In H. Bouamor, J. Pino, and K. Bali, editors, Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 12216- 12235, Singapore, Dec. 2023. Association for Computational Linguistics. doi: 10.18653/v 1/ 2023.emnlp-main.751.
http:t//:://ac
org/
wrg//023...e.oooooo 2023. www.p-mmin..mm..
[28] A. Grattafiori, A. Dubey, A. Jauhri, A. Pandey, A. Kadian, A. Al-Dahle, A. Letman, A. Mathur, A. Schelten, A. Vaughan, A. Yang, A. Fan, A. Goyal, A. Hartshorn, A. Yang, A. Mitra, A. Sravankumar, A. Korenev, A. Hinsvark, and A. Rao. The Ilama 3 herd of models, 2024. URL https://arxiv org/ abs/2407 .2783
[29] K. Guu, K. Lee, Z. Tung, P. Pasupat, and M. Chang. Retrieval augmented language model pre-training. In H. D. III and A. Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 3929- 3938. PMLR, 13-18 Jul 2020. URL. httpps:/pproceedings mlr. wwwppsevelviie vatsss//v11//uu0aaa guu20a html
13

---

[30] M. Hanna, S. Pezzelle, and Y. Belinkov. Have faith in faithfulness: Going beyond circuit overlap when finding model mechanisms. In First Conference on Language Modeling, 2024. URL wwt.://oooooo...////00co......
[31] R. Hendel, M. Geva, and A. Globerson. In-context learning creates task vectors. H. Bouamor, J. Pino, and K. Bali, editors, Findings of the Association for Computational Linguistics: EMNIPP 2023, pages 9318- 9333, Singapore, Dec. 2023. Association for Computational Lin- guistics. doi: 10.18653/v 1/2023. wte...2....nnnneeeeim.. 2023. fin.ingssnn-ao e.mee.mmlp..... 624/
URL httLhhtts://ee htttps://aclanthologg org/
[32] A. R. Hsu, G. Zhou, Y. Cherapanamjeri. Y. Huang, A. Odisho, P. R. Carroll, and B. Yu. Efficient automated circuit discovery in transformers using contextual decomposition. In The Thirteenth International Conference on Learning Representations, 2025. URL https
7 //openreview.
[33] G. Izacard and E. Grave. Leveraging passage retrieval with generative models for open domain question answering. In P. Merlo, J. Tiedemann, and R. Tsarfaty, editors, Proceedings of the 16th Conference of the European Chapter of the Association for Computational Linguisties: Main Volume, pages 874-880, Online, Apr. 2021. Association for Computational Linguistics. doi: 10. 18653/v1/2021.eam www.:021.eal-mmmn...m. 74. httptttss///acaamt
org/2021.....2 2021 eaclemaimmmmn 74/
[34] A. Jacovi and Y. Goldberg. Towards faithfully interpretable NLP systems: How should we define and evaluate faithfulness? In D. Jurafsky, J. Chai, N. Schluter, and J. Tetreault, editors, Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 4198-4205, Online, July 2020. Association for Computational Linguistics. doi: 18653v1//v1/ 2020.acl- 11:20acl-mmin.333 ma.n.386. URL
org/ 2020 acl wwaimmainn.. 386/
[35] S. Jain and B. C. Wallace. Attention is not Explanation. In J. Burstein. C. Doran, and T. Solorio, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pages 3543-3333333 3556, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 18.18/ 13577
URL httttt:////:ccattioon
www.oogg/////.ooo. -
[36] A. Q. Jiang, A. Sablayrolles, A. Mensch, C. Bamford, S. Chaplot, D. de las Casas, F. Bressand, G. Lengyel, G. Lample, L. Saulnier, L. R. Lavaud, M.-A. Lachaux, P. Stock, T, L. Scao, T. Lavril, T. Wang. T. Lacroix, and W. E. Sayed. Mistral 7b, 2023. URL Stthttt::.. ptattps///arivv org/ /ba/ 2310.06825.
[37] Z. Jin, P. Cao, H. Yuan, Y. Chen, J. Xu, H. Li, X. Jiang, K. Liu, and J. Zhao. Cutting off the head ends the conflict: A mechanism for interpreting and mitigating knowledge conflicts in language models. In L.-W. Ku, A. Martins, and V. Srikumar, editors, Findings of the Association for Computational Linguistics: ACI. 2024, pages 1193-1215, Bangkok, Thailand, Aug. 2024. Association for Computational Linguistics. doi: 10. 1/653/////02..fnmm findings-acl.70.. URL www.ooog.oorrgego 2024. mondings-acl.7000
[38] M. Joshi, F. Choi, D. Weld, and I.. Zettlemoyer. TriviaQA: A large scale distantly supervised challenge dataset for reading comprehension. In R. Barzilay and M.-Y. Kan, editors, Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume Long Papers), pages 1601-1611, Vancouver, Canada, July 2017, Association for Computational URL httpps::/acaaaahhllloooooo/ www.lloggy.rg/111111// org/P17-11477
Linguisties. doi: 10. 1/177-1177
[39] V. Karpukhin, B. Oguz, S. Min, P. Lewis, L. Wu, S. Edunov, D. Chen, and W.-t. Yih. Dense passage retrieval for open-domain question answering. In B. Webber, T. Cohn, He, and Y. Liu, editors, Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 7769-6781, Online, Nov. 2020. Association for Computational Linguistics. doi: 10. 2020. emnlp-main..n......memmmmmmmmmam 550/
URL
org/
[40] U. Khandelwal, O. Levy, D. Jurafsky, I.. Zetrlemoyer, and M. Lewis. Generalization through memorization: Nearest neighbor language models. In International Conference on Learning Representations, 2020. URL http://ppeneeeeeeeeeeeeeeee
CE
14

---

[41] P. Koehn. Europarl: A parallel corpus for statistical machine translation. In Proceedings of Machine Translation Summit X: Papers, pages 79-86. Phuket, Thailand, Sept. 13-15 2005. URL https: https:///allantholloy org/ 2005 www.mmitttt.....mmm....mmmmmm.mt tarpsuumnt-ppprr......om
[42] E. Kortukov, A. Rubinstein, E. Nguyen, and S. Oh. Studying large language model behav- iors under context-memory conflicts with real documents. In First Conference on Language Modeling, 2024. URL htt.itttss//ppenreviiee http:///ppparreeeeeeeeeerrrrrrro htt.fforuummddxxx22ieeeee
[43] O. Kovaleva, A. Romanov, A. Rogers, and A. Rumshisky. Revealing the dark secrets of BERT. In K. Inui. J. Jiang. V. Ng, and X. Wan, editors, Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNEL pages 4365-4374, Hong Kong, China, Nov. 2019. Association for Computational Linguistics. doi: 10. 10.18653/v1/D19.. 1445. URL https://aclantthoogg. Drg/D19-1449
[44] T. Kwiatkowski, J. Palomaki, O. Redfield, M. Collins, A. Parikh, C. Alberti, D. Epstein, I. Polosukhin, J. Devlin, K. Lee, K. Toutanova, L. Jones, M. Kelcey, M.-W. Chang, A. M. Dai, J. Uszkoreit, Q. Le, and S. Petrov. Natural questions: benchmark for question answering research. Transactions of the Association for Computational Linguistics. 7:452-466, 2019. doi: 10.1 10.1162/tacl_a a ___0666.. URL https: hatp:://aallatthllgggoooo ott/
[45] P. Lewis, E. Perez, A. Piktus, F. Petroni, V. Karpukhin, N. Goyal, H. Küttler, M. Lewis, W.-t. Yih, T. Rocktäschel, S. Riedel, and D. Kiela. Retrieval-augmented generation for knowledge- intensive nlp tasks. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 9459-9474. Curran Associates, Inc., 2020. URL http::t/ neur iiteurrrp..... cc/ wwwpppcpp.aeeeeeceeeceeee satper_ffles/ppppppe .
[46] K. Li, O. Patel, F. Viégas, H. Pfister, and M. Wattenberg. Inference-time interven- tion: Eliciting truthful answers from language model. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 41451-41530, Curran Associates, Inc., 2023. URL http: wwtpeeewwwwpppppppee.......
nwwrins. wwwpccppppppppeeeeeee www.flllliiffieeeele httpe/eeieeeeaee/2222////////// /aper/2023/file/ 31883000877022200006666666633363 Confeeee pdf
[47] J. Liu, D. Shen, Y. Zhang, B. Dolan, L. Carin, and W. Chen. What makes good in-context examples for GPT-3? In E. Agirre, M. Apidianaki, and I. Vulić, editors, Proceedings of Deep Learning Inside Out (DeeLIO 2022): The 3rd Workshop on Knowledge Extraction and Integration for Deep Learning Architectures, pages 100-114, Dublin, Ireland and Online, May 2022. Association for Computational Linguistics. doi: 11.16533/122222ddeeiioo .ww.gy..orrr.go. www.//2..........ooo..oooo http:/2222.eeelle. dewlio-1.
URL
[48] S. Longpre, K. Perisetla, A. Chen, N. Ramesh, C. DuBois, and S. Singh. Entity-based knowledge conflicts in question answering. In M.-F. Moens, X. Huang, L. Specia, and S. W.-t. Yih, editors, Proceedings of the 202/ Conference on Empirical Methods in Natural Language Processing, pages 7052-7063, Online and Punta Cana, Dominican Republic, Nov. 2021. Association for Computational Linguistics. doi:
URL https://
ac aaclanthology.orgg org 2021. ww1.:01..emlpmmmmmmmmmmmm.m wwmnp-mmin.... 565/
[49] S. M. Lundberg and S.-I. Lee. A unified approach interpreting model predictions. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL :ttttpp://prrcoeiinggs neurips. cc/ 622766621178622226644446822232
.pff
[50] B. McCann, N. S. Keskar, C. Xiong, and R. Socher. The natural language decathlon: Multitask learning as question answering. 2018. lttps:// /tatpps://ariiv.... wwwrxivvoog//ae http://80000887700
[51] J. Minder, K. Du, N. Stoehr, G. Monea, C. Wendler, R. West, and R. Cotterell. Controllable context sensitivity and the knob behind In The Thirteenth International Conference on Learning Representations, 2025. URL httpp:://pperreiieeeee
15

---

[52] A. Nafar. K. B. Venable, and P. Kordjamshidi, Learning vs retrieval: The role of in- in-context ex- amples in regression with large language models. In Chiruzzo, A. Ritter. and Wang, editors, Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume Long Pupers), pages 2006-8229, Albuquerque, New Mexico, Apr. 2025. Association for Computational Linguistics. ISBN 979-8-89176-188-66 URL BKLLtttss///ccaantaaa
org. 2025 www.naaacnaaaa -loo......... 417/
[53] nostalgebraist. Interpreting gpt: The logit lens, August 2020. URL httph: www lesswrong htap://sssssccckkaaaoaaaaa/ooeen 2025-05-12.
-atenngg--ptttttteeeeeeeeeeeeeee the logit lens
Accessed:
[54] C. Olsson, N. Elhage, N. Nanda, Joseph, N. DasSarma, T. Henighan, B. Mann, A. Askell, Y. Bai, A. Chen, T. Conerly, D. Drain, D. Ganguli, Z. Hatfield-Dodds, D. Hernandez, S. Johnston, A. Jones, J. Kernion, L. Lovitt, K. Ndousse, D. Amodei, Brown, J. Clark, J. Kaplan,
S. McCandlish, and C. Olah. In-context learning and induction heads, 2022. URL https
//arxiv. org//ogg//o
11895
[55] F. Ortu, Z. Jin, D. Doimo, M. Sachan, A. Cazzaniga, and B. Schölkopf. Competition of mechanisms: Tracing how language models handle facts and counterfactuals. In L.-W. Ku, A. Martins, and V. Srikumar, editors, Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 8420-8436, Bangkok, Thailand, Aug. 2024. Association for Computatiooal Linguisties. doi: 1/218863/1112224.ccll htltttts://aet hataa/cccaantoooooyyooo .wwpoog..oo..oooooooo.o...oo.... wttp///4..cllooooooooooe. aalllcllggn.......a. 4588/
[56] L. Ouyang. J. Wu, X. Jiang, D. Almeida, C. Wainwright, P. Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray. Schulman, J. Hilton, F. Kelton, L. Miller, M. Simens, A. Askell, P. Welinder, P. F. Christiano, J. Leike, and R. Lowe. Training language models follow instructions with human feedback. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 27730-27744, Curran Associates, Inc., 2022. htt neur cc/ cttp:epssccpppppffffelllleeeeeee hppp://022221iillleeeess66666661 Paper Conference.
[57] F. Petroni, T. Rocktäschel, S. Riedel, P. Lewis, A. Bakhtin, Y. Wu, and A. Miller. Language models as knowledge bases? In K. Inui, Jiang, V. Ng, and X. Wan, editors, Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP)) pages 2463-2473, Hong Kong, China, Nov. 2019. Association for Computational Linguistics. doi: 10.18653/1//D99/ 1/50--1200
URL httpp:://acaaatttll/////// wwwgggg...rg..ggo...... org/ D190-1991 1250/
[58] F. Petroni, P. Lewis, A. Piktus, T. Rocktäschel, Y. Wu, A. H. Miller, and S. Riedel. How context affects language models' factual predictions. In Automated Knowledge Base Construction. 2020. URL httptttt//peeeeeeeete http://ppeeeeeeeeeeeeeeeeerooaar
net/ wttp:/orrmmmddd0000000P00
[59] M. Post. call for clarity in reporting BLEU scores. In O. Bojar, R. Chatterjee, Federmann, M. Fishel, Y. Graham, B. Haddow. M. Huck, A. J. Yepes, P. Koehn, C. Monz, M. Negri, A. Névéol, M. Neves, M. Post, L. Specia, M. Turchi, and K. Verspoor, editors, Proceedings of the Third Conference on Machine Translation: Research Papers, pages 186-191, Brussels, Belgium, Oct. 2018. Association for Computational Linguistics. doi: 10.18533/1/W118633 URL htt.::ttss:///e http:///alaathhoog..ooooo wwwg/o8-g///100
[60] Qi, G. Sarti, R. Fernández, and A. Bisazza. Model internals- based answer attribution for trustworthy retrieval- augmented generation. In Y. Al-Onaizan, M. Bansal, and Y.-N. Chen, editors, Proceedings of the 2024 Conference on Empirical Methods Natural Language Processing, pages 0037-6053, Miami, Florida, USA, Nov. 2024. Association for Computational Linguistics. doi: 2024. emn wammlp-main.344
URL htt.hattss//aaaaathllooooo org/
[61] C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Zhou, W. Li, and P. J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research, 21 21(140): 2020. URL URL http: /mlrr. .rre arg/papers/v21.
004. html
16

---

[62] O. Ram. Y. Levine, I. Dalmedigos, D. Muhlgay, A. Shashua, K.
and Y. Shoham. In-context retreeeal-augmmettd language models. Transactions of the Associ- ation for Computational Linguisticy, 11:1316-1331, 2023. doi: 10. 10.1162/tacl_ wttp:///22..taal111.
URL
[63] D. Samuel. Berts are generative in-context learners. In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang, editors, Advances in Neu- ral Information Processing Systems, volume 37, pages 2558-2589. Curran Associates, Inc., 2024. URL /ttpa///ot////ooooooeooooooooooa hwtperpps.ccppppeffieeeeeee Conferrrr-oofeeneeeeee pdf
[64] D. Schlangen. Lims as function a Terminoloyy, taxonomy, and questions for evaluation, 2024, URL https:://axiiii http://aaxvv..or/////////.0/....
[65] S. Serrano and N. A. Smith. Is attention interpretable? In A. Korhonen, D. Traum, and L. Märquez, editors, Proceedings of the 57th Annual Meeting of the Association for Computa- tional Linguistics, pages -931-2951, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10. 10.11865/v1/PPP 1282.
1282/
[66] K. Shuster, S. Poff, M. Chen, D. Kiela, and J. Weston. Retrieval augmentation reduces hallucination in conversation. In M.-F. Moens, X. Huang, L. Specia, and S. W.-t. Yih, edi- tors, Findings of the Association for Computational Linguistics: EMNIP 2021, pages 3784- 3803, Punta Cana, Dominican Republic, Nov. 2021. Association for Computational Linguis- tics. doi: hooo:dddo......s////00/1///nnnno findings-emmlp.320. httlhtttt://aclaaattt
.rg.e....rrg/oo.... org/2021
findings emnlp.3222 320/
[67] K. Simonyan, Vedaldi, and A. Zisserman. Deep inside convolutional networks: visualising image classification models and saliency maps. In Proceedings of the International Conference on Learning Representations (ICLR). ICLR, 2014.
[68] M. Sundararajan. A. Taly. and Q. Yan. Axiomatic attribution for deep networks. In D. Precup and Y. W. Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 3319-3328, PMLR, 06-11 Aug 2017. ptLpattss://proeeeddiee mmlr wtt.ggg................. v70/ sttr///annaaaaaaaaaiiaaa .ttll
[69] I. Sutskever, O. Vinyals, and Q V. Le. Sequence to sequence learning with neural net- works. In Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K. Weinberger, edi- tors, Advances in Neural Information Processing Systems, volume 27. Curran Associates, Inc., 2014. URL BtRLLtttpps//prcceddingi neurips...... ppaper_filete httppppra.illeepppppp/////////// pttpp//pppp////iiiii/ pdf
a 71425544227422568844411666336684
[70] A. Syed, C. Rager, and A. Conmy. Attribution patching outperforms automated circuit dis- covery. In Y. Belinkov, N. Kim, J. Jumelet, H. Mohebbi, A. Mueller, and H. Chen, editors, Proceedings of the 7th BlackboxNLP Workshop: Analyzing and Interpreting Neural Networks for NLP, pages 407-416, Miami, Florida, US Nov. 2024. Association for Computational org/
Linguistics. doi: 10.1866534 21255///1/2224.baaccooooo/ 2024. blackboxnlp-1.225
https:///////ttttttt
[71] H. Tan, F. Sun, W. Yang, Y. Wang, Q. Cao, and X. Cheng. Blinded by generated contexts: How language models merge generated and retrieved contexts when knowledge conflicts? In L.-W. Ku, A. Martins, and V. Srikumar, editors, Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume Long Papers), pages 6207-6227, Bangkok, Thailand, Aug. 2024. Association for Computational I Linguistics. doi: 10.18653/v1/ 0024.acl-long.333 URL, ttlt:tttt///ccaatthooo www.hhhooooo.ooooooooooooooooo.e 2024/2224..00 acl long........3311
[72] G. Team, M. Riviere, S. Pathak, P. G. Sessa, C. Hardin, S. Bhupatiraju, I., Hussenot, T. Mesnard, B. Shahriari, A. Ramé, al. Gemma 2: Improving open language models practical size, 2024. URL /estts://aaie http:///axvvvv..o org/ / abs/ / 2210//0088001118..81
[73] J. Tiedemann. Parallel data, tools and interfaces in OPUS. In N. Calzolari, K. Choukri, T. De- clerck, M. U. Dogan, B. Maegaard, J. Mariani, A. Moreno, J. Odijk, and S. Piperidis, editors, Proceedings of the Eighth International Conference on Language Resources and Evaluation
17

---

(LREC) 12), pages 2214-2218, Istanbul, Turkey, May 2012. European Language Resources Association (ELRA). URL http::/www..rre www.- conf org/ /ttp://eennnnn////e//////.//a/t. 463 Pww.pappppppp. pdf
[74] E. Todd, M. Li, A. S. Sharma, A. Mueller, B. C. Wallace, and D. Bau. Function vectors in large language models. In The Twelfth International Conference on Learning Representations, 2024. URL htRlhttts::// tttttp:://openreviee www. httl//orrmmmmmdmieeeeee
[75] A. Vaswani, N. Shazeer, N. Parmar; J. Uszkoreit, L. Jones, A. N. Gomez, L. u. Kaiser, and Polo- sukhin. Attention is all you need. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Syy- tems, volume 30. Curran Associates, Inc., 2017.
mewreiindddinnnnnnn..see.....ess cc/ Paper pdf
paper_files/ http:///aw............//////////
[76] E. Voita, D. Talbot, F. Moiseev, R. Sennrich, and I. Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. In A. Korhonen, D. Traum, and L. Märquez, editors, Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 5797-5808, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10. 1111856//1111911111 1580. UtL.ittp:://aclanttoo P19- P19-15807
org/
[77] J. Von Oswald, E. Niklasson, E. Randazzo, J. Sacramento, A. Mordvintsev, A. Zhmoginov, and M. Vladymyrov. Transformers learn in-context by gradient descent. In A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 35151-351744 PMLR, 23-29 Jul 2023. URL httpps:/prroceedings mlr press/
html
[78] D. Vrandečić and M. Krötzsch. ACM, 57(10):78-85, Sept. 2014,
Wikidata: a free collaborative knowledgebase. Commun. ISSN 0001-0782, doi: 10.1145/2629489, URL https
[79] K. R. Wang, A. Variengien, A. Conmy, B. Shlegeris, and J. Steinhardt. Interpretability in the wild: a circuit for indirect object identification GPT-2 small. In The Eleventh International Conference on Learning Representations, 2023. URL httpps://peerrvieett wwwpe/rrriiie.nn//oorn.n NpVVSSSS
[80] J. Wei, M. Bosma, V. Zhao, K. Guu, A. W. Yu, B. Lester, N. Du, A. M. Dai, and Q. V. Le. Finetuned language models zero-shot learners. In International Conference on Learning Representations, 2022. URL /tctps://openreviee
net/forunt www.orrmmmddddddooeoo 1d=gZZGGoozdRR
[81] S. Wiegreffe and Y. Pinter. Attention is not not explanation. In K. Inui, J. Jiang, V. Ng, and X. Wan, editors, Proceedings of the 2019 Conference on Empirical Methods in Natural Lan- guage Processing and the 9th International Joint Conference on Natural Language P Processing (EMNLP-IJCNLP), pages 11-20, Hong Kong, China, Nov, 2019. Association for Computational Linguistics. doi: 10.18653/v1/DI1
URL https: / acl anthhoooytyo...... org/ 1002/
[82]
T. Wolf, L. Debut, V. Sanh, J. Chaumond, C. Delangue, A. Moi, Cistac, T. Rault, R. Louf, M. Funtowicz, J. Davison, S. Shleifer, P. von Platen, Ma, Y. Jernite, J. Plu, C. Xu, T. Le Scao, S. Gugger, M. Drame, Q. Lhoest, and Rush. Transformers: State-of-the-art natural language processing. In Q. Liu and D. Schlangen, editors, Proceedings of the 2020 Conference on Empir- ical Methods in Natural Language Processing: System Demonstrations, pages 38-45, Online, Oct. 2020. Association for Computational Linguistics. doi: 10. 188533/11/2020emmmmpmmooo URL htt.httttss/aaaaatholllle .rg www./020....0.o.ooo... 2020. emnlp-demss 6/
[83] W. Wu, Y. Wang, G. Xiao, H. Peng, and Y. Fu. Retrieval head mechanistically explains long- context factuality. In The Thirteenth International Conference Learning Representations, 2025. URL http::htttt//ppptt http::/://peervvieeeee
[84] G. Xiao, Y. Tian, B. Chen, S. Han, and M. Lewis. Efficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Representations, 2024.
httpst: http://awneeeeee.eeee/r/rrrrraas
18

---

[85] J. Xie, K. Zhang, J. Chen, R. Lou, and Y. Su. Adaptive chameleon or stubbor sloth: Revealing the behavior of large language models in knowledge conflicts. In The Twelfth International Conference on Learning Representations, 2024. URL httpt::t/pp:speeeeeeeeee http://wwwwwwwtttteeeeeeee net/forum? id=auKAU12M06
[86] J. Xu, X. Sun, Z. Zhang, G. Zhao, and J. Lin. Understanding and improving layer normaliza- tion. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché- Alché-Buc, E. Fox, and R. Garnett editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https: httttpp:://prooceedinggess mew.sedddddnnnnnnneeppee......ee nwwrips. cc/pp/.ccc///co saaper_files/pppeeeeeee hiapper_files/ppppe////1///////1 .ppaperppfff
cc/
2144:037772277777066667778888888
[87] R. Xu, Z. Qi, Z. Guo, C. Wang, H. Wang. Y. Zhang, and W. Xu. Knowledge conflicts for LLMs: A survey. In Y. Al-Onaizan, M. Bansal, and Y.-N. Chen, editors, Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 8541-8565, Miami, Florida, USA, Nov. 2024. Association for Computational Linguistics. doi: 10.18653/ v1/2024.cmnlp- maw..486..666 URL http::ttt////eett pettps://aclanthollle
org/ 2224/204.42o aunlp-main. 4861
[88] X. Xu, Z. Gou, W. Wu, Z.-Y. Niu, H. Wu, H. Wang, and S. Wang- Long time no see! open-domain conversation with long-term persona memory. S. Muresan, P. Nakov, and A. Villavicencio, editors, Findings of the Association for Computational Linguistics: ACI. 2022, pages 2639-26666 2650, Dublin, Ireland, May 2022. Association for Computational Linguis- tics. doi: 11..8653/v120222.nnnlee fww.diiinng--.l.20.. URL https::/aclanthologg .rg. 2022 f findings-acl. ac1. 207/
[89] K. Yin and G. Neubig. Interpreting language models with contrastive explanations. In Y. Gold- berg, Z. Kozareva, and Y. Zhang, editors, Proceedings of the 2022 Conference Empirical Methods Natural Language Processing, pages 1841998,0 Abu Dhabi, United Arab Emirates, Dec. 2022. Association for Computational Linguistics. doi: 10. 1863/112022.emmlpmmmaa URL httpttts//aaaaattoollooooo
org.org///// 2022. http:/22222..mmmllmmmeemmmmm
14/
[90] K. Yin and J. Steinhardt. Which attention heads matter for in in-context learning? In Forty-second International Conference on Machine Learning, 2025. URL https:://penrrevie net/
f forum?id forum?dd-C7xmmByCP7
[91] Q. Yu, J. Merullo, and F. Pavlick. Characterizing mechanisms for factual recall in language models. In H. Bouamor, J. Pino, and K. Bali, editors, Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 9924-9952 Singapore, Dec. 2023. Association for Computational Linguistics. doi: 10. 1/602331/2023.mmmlppp mmn..mmai..615... URL https://aclantholos org 2023. emnlp-main. 615/
[92] Y. Zhangg S. Sun, M. Galley, Y.-C. Chen, C. Brockett, X. Gao, J. Gao, J. Liu, and B. Dolan. DIALOGPT: Large-scale generative pre-training for conversational response generation. In A. Celikyilmaz and T.-H. Wen, editors, Proceedings of the 58th Annual Meeting of the Asso- ciation for Computational Linguistics: System Demonstrations, pages 270-278, Online, July 2020. Association for Computational Linguistics. doi: 10. 11111///1/20200aalleeeeee
URL
https: / ac lanthology org/ 2020. acl demos. 30/
19

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification:
We describe our contributions at the end of Section 1
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We describe the limitation of our work in Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
20

---

Justification:
We do not introduce any theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide descriptions of methods, models, datasets, results, and hyperpa- rameters that use in the main paper and in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
21

---

Answer: [Yes]
Justification:
We provide access to our code and datasets. All models that we use are publicly
available.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results. See the NeurlPS code and data submission guidelines (https
ssicnpolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provide our experimental setup Section 4 and in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Due to computational resources limitation, we only report results for a single run with the best hyperparameter configuration.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
22

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide information on the compute resources in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics http://pw.//.arrc..c.///p//////a
Answer: [Yes]
Justification: Our research to study inner workings of LLMs conforms with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: describe the potential societal impacts of our work Section
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
23

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system Iearns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: We publicly available models and datasets. The biography dataset that we curate consists of publicly available information accessible through Wikidata, and we foresee risks in releasing the data.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We describe the version of models and datasets that use, along with their licenses in the Appendix.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswithcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
24

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer:
[Yes]
Justification:
We provide the details of the biography dataset in the Appendix.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
We do not conduct any human studies.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We do conduct any human studies.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
25

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [Yes]
Justification: We study inner workings of LLMs. We also use LLMs for paraphrasing of the biography dataset.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: //errr////eups.../ /t//oonfeeennee/222222 for what should or should be described.
26

---

A Appendix
released under the Llama 3.1 Community License. Gemma-2- is released under Apache 2.0. NQ-Swap and TriviaQA that we use are derived from MRQA 244, which is released under the MIT license. We construct the biography dataset using Wikidata Query Service @ which is available under CCO. The OPUS Europarl dataset also available under CCO.
Licenses
9B-it is released under the Gemma license agreement. Miitral--B-Instruuct
QA Datasets For NQ-Swap, we use the preprocessed data and split available on HuggingFacd (4,746 examples). The corpus substitution procedure [48] is applied to generate counterfactual contexts. As for TriviaQA, we use the dev split from the MRQA repository (7,785 examples).
To build our biography dataset, start by selecting 100,000 random individuals that have the following attributes in Wikidata: date of birth (P569), place of birth (P19), education (P69), and occupation (P106). Furthermore we filter for individuals that have notable contributions (P800), in an effort to maximize the chance that all LLMs we employ can answer questions regarding them. An entity may have multiple occupations and educations. For simplicity, we only select one of them in random manner. We also only choose entities that have complete labels for all attributes and ensure that they are distinct individuals, resulting the final 4,255 examples.
Table 3: An example of our dataset with the original (ORIG) and paraphrased (PARA) biography entry.
Vladimir Vapnik was educated in V.A. Trapeznikov Institute of Control Sciences. Vladimir
ORIG
Vapnik was born 06 December 1936. Vladimir Vapnik worked as computer scientist. Vladimir Vapnik was born in Tashkent.
Vladimir Vapnik was born on 06 December 1936 in Tashkent, city that would later shape his life's work. As a young man, Vapnik was fascinated by the potential of machines learn and adapt. He went on to study the V.A. Trapeznikov Institute of Control Sciences, where he was exposed to the latest advancements in computer science and artificial intelligence. It was here that Vapnik's passion for machine learning truly took hold. After completing his studies, Vapnik went on to become a renowned computer scientist, making groundbreaking contributions to the field. His work on support vector machines and the
PARA
dimension would have lasting impact on the development of machine lcarning algorithms. Throughout his career, Vapnik has been recognized for his innovative thinking and dedication to advancing the field of computer science. His legacy continues to inspire new generations of researchers and scientists.
In Table 3 we show an example of our dataset. The paraphrased biography entry is obtained with -lama-3.I-8B-Insttte through greedy decoding and we generate until an EOS token is encountered. We also ensure that the paraphrased entries still contain all the original attributes. The safety guidelines applied during the fine-tuning of the Llama model can sometimes prevent it from generating biographies of political figures including birth dates coinciding with sensitive historical events. To address this, we use simple strategy to jump-start the model's 's generation process. Specifically, we initiate the generation with the phrase, "Here is 150-word fictional biography of (name):". We use the following prompt:
prompt = f clsttarttheadrrrrrt.r hhw.aeeeidllssstteeeeeeeeeeeeeee wwtpsiissteeeeeeeeeeeeeeeee Seapandddhhaer_idiie idl> You are a helpful assis- 150 words fictional biography containing the following facts in random order, make sure to include ALL facts VERBATIM as they are: {facts) ) <leot_ idl><Istart header idl=asiittatttlend header_idl>He is 150-word fictional biography of (name):""*
cant.<leot_idd><tte dlstttttttttrt_____. hheder_idlsuseeee httpauueereeeeddeedee_iiiiiiddd
Implementation Details All model checkpoints that we use and their correspond- ing tokenizers are available on HuggingFace: eettallaamaLLane ----I-nntutttt..
httpp://queey.vkkiddaa.ooo
httpp///uuggggggeecccc/sssssssst Swap
http:://gghaall..oo////g/ggghaar 2019
Instruct
27

---

and google/ google/gemmaa. gemma-gmaaa-mm 1et-mma------aa 2--9b-tt I1 All models were ran on mixed precision (bfloat16). We use Pyserini implementation of DPR 39 For BERT matching, we use the checkpoint provided by Kortukov et al. [42] 13 For all experiments, we apply greedy decoding and set maximum limit of generated tokens
Listrrlla/Mistrra-----B-B
20 unless specified otherwise.
Compute Details All the experiments were conducted on 2 24 GB RTX4090 and x 40 GB A100. Computing difference score D to identify in-context and parametric heads takes about hours. Ablating both heads on NQ-Swap and TQA takes about hours for each run. Patching task and parametric FVs takes about 12 hours average for each model, while the NIAH experiment with retrieval heads takes about 8 hours. Our source tracking experiments consume about hours. The machine translation experiment takes about 6 hours.
Table 4: Overview of models' performance NQ-Swap and TQA for reproducibility purposes. Besides recall, we compute traditional exact match accuracy and BERT matching (BEM) [3] to measure semantic match. We also adopt K-Precision
[ to evaluate answers' groundedness.
NQ-Swap
TQA
Recall
EM
BEM K-Prec.
Recall
EM
BEM
K-Prec.
Oracle
Llama-3. 1-8B-Instruct
87.67
63.63 90.75
93.62
88.12
72.77
90.97
97.23
Mistral- Faatall-B-Instruct--d.. Instruct-vO.
87.05
48 06
90.46
93.88
87.13
65.48
90.57
97.54
Gemma-2-9B-i
85.93
66.79
89.68
93.92
87.50
70.26
90.66
96.88
Counterfactual
Llama-3. -.ama-3.188B-Istteee 1-----nsttt
68.73
51.56
70.27
88.61
Mistral-7B-Bh -
67.61
35.08
70.35
89.12
Gemma-2-9B-it
66.67
50.78
68.58
84.86
DPR (39
Llama-3. 3.1 1-8B-Instttct
46.57
34.68
52.97
84.08
66.10
53.44
69.61
82.83
(harral----nstruct--vk.
49.96
26.95
58.81
84.88
69.42
52.15
73.26
83.50
Gemma-2-9B-it
46.12
32.81
54.78
81.03
66.79
54.37
70.38
80.44
Closed-book
Llama-3, 3. --8t-Instt
51.64
32.34
59.52
74.41
61.66
78.01
Mistral-7B -B- 1B-IInstruct-vO.3
46.26
22.10
57.84
73.12
60.06
76.65
Gemma-2 2-9-it
44.61
22.46
54.53
70.97
56.07
75.16
B Details: Localization of In-Context and Parametric Heads
We discuss additional details regarding experiments and results in $
Experiment Details We format our questions with prompt template following Ram et al. [22 For ablations, we set the activation of attention heads to zero after softmax. Besides recall, we also evaluate models' performance with standard exact match (EM) accuracy and BERT matching (BEM) [33, since EM often too strict. ablation results with DPR [39], we only select instances where K-Precision is equal to I in the original run, since we want to focus on cases where models make full use of the contextual information, especially considering that retrieved contexts can be imperfect.
where we observe similar Additional Results We present our additional results in Figure 10
trends to Figure 2 In general, removal of in-context and parametric heads reduces performance in all models across all metrics for open-book and closed-book settings respectively, under various different configurations. The performance drops also holds for TQA, which shows the transfcrability of the identified in-context and parametric heads considering that they are computed only on NQ-Swap. Furthermore, we see that our method yields a more significant performance decrease compared to AWR heads 83), demonstrating its suitability to study heads that contribute to the answer generation.
/ hhtppe///ngggiigddeecooooaaassss 1
28

---

Qualitative Examples In Table 5 we show what happens qualitatively when in in-context and/or parametric heads in L.a-a-3.1-8B-1 ----t-untttuct are ablated. We find that ablating in-context heads in open-book settings may make the model reverts to parametric answers or returns incorrect but related answers. Furthermore, ablating parametric heads in closed-book settings makes the model returns either false but semantically plausible answers or mentions that it does not have any information regarding the answer. Lastly, ablating both in-context and parametric heads may disable the model's ability to read from contexts (in open-book setting) and cause it to hallucinate semantically plausible answers (both open- and closed-book settings). While we have some evidence regarding what happens in the generation space when in in-context and/or parametrie heads are ablated, determining what factors affect which fallback mechanism employed by the model would require further investigation.
Table 5:
Some output examples of Llama-3. 1.ama-3.1881 1-1--B-BB --8t-unttruct on NQ-Swap when in-context and/or parametric heads are ablated. Gold answers are denoted in bold.
(a) In-context heads ablated in open-book setting
Before
After
<P> Louis XIII 's successor, Carric Underwood, had a great interest in Versailles. He settled on the royal hunting lodge at Versailles and over the following decades had it expanded into one of the largest palaces the world. Beginning in 1661, the architect Louis Le Vau landscape architect André Le Nôtre and painter- decorator Charles Lebrun began detailed renovation and expansion of the château This was done to fulfill Carrie Underwood 's desire to establish a new centre for the royal court Following the Treatics of Nijmegen in 1678, he began to gradually move the court to Versailles. The court was officially established there on 6 May 1682. </P>
<P> Louis XIII 's successor, Carric Underwood, had a great interest in Versailles. He settled on the royal hunting lodge at Versailles and over the following decades had it expanded into one of the largest palaces the world. Beginning in 1661. the architect Louis Le Vau landscape architect André T.e Nôtre and painter- decorator Charles Lebrun began detailed renovation and expansion of the château This was done to fulfill Carrie Underwood 's desire to establish a new centre for the royal court Following the Treaties of Nijmegen in 1678, he began to gradually move the court to Versailles The court was officially established there on 6 May 1682. </P>
Based on this text, answer these questions: Q: who expanded the palace of versailles to its present size?
Based on this text, answer these questions:
Q: who expanded the palace of versailles to its present size?
A: Louis XIII (Louis XIV)
A: Louis XIV. (Louis X/V)
<P> Knockin on Heaven 's Door is a song written
<P> Knockin on Heaven 's Door is a song written and sung by Jacques Cousteau. for the soundtrack of the 1973 film Pat Garrett and Billy the Kid. Released as single reached No. 12 on the Billboard Hot 100 singles chart Described by Dylan biographer Clinton Hcylin as " an exereise in splendid simplicity the song in terms of the number of other artists who have covered it, is of Dylan most popular post- compositions. </P>
and sung by Jacques Cousteau, for the soundtrack of the 1973 film Pat Garrett and Billy the Kid. Released as single. . reached No. on the Billboard Hot 100 singles chart Described by Dylan biographer
Clinton Heylin as an exercise in splendid simplicity the song in terms of the number of other artists who have covered it, is one of Dylan most popular post- 19st-1960s compositions. </P>
Based on this text, answer these questions: Q: who wrote knock knock knocking on heavens door?
Based on this text, answer these questions:
Q: who wrote knock knock knocking on heavens door?
A: Jacques Cousteau (Bob Dylan)
A: The song is attributed to the author of the lyrics, which are credited to the composer, but the (Bob Dylon)
(b) Parametric heads ablated in closed-book setting
Before
After
Answer these questions:
Answer these questions:
Q: when was theme from a summer place released?
Q: when was theme from a summer place released?
A: 1959 (1959)
A: 1967 (I959)
Answer these questions:
Answer these questions:
Q: who plays timon in lion king on broadway?
Q: who plays timon in lion king on broadway?
A: Adam Jacobs (Max Casella)
A: I do not have information on who plays Timon in
The Lion King on Broadway. (Max Casella)
29

---

(c) Both in co-tont and parametric heads ablated in open- and blosed-booke settings
Before
After
<P> In 26 August 2016
Olympics 2016 bronze
<P> In 26 August 2016 Olympics 2016 bronze
medallist Kyle Busch was made brand ambassador for
medallist Kyle Busch was made brand ambassador for BBBP </P>
BBBP. </P>
Based on this text, answer these questions:
Based this text, answer these questions:
Q: who has been chosen as the brand ambassador of campaign beti bachao-beti padhao?
Q: who has been chosen as the brand ambassador of the campaign
bachao-beti padhao?
A: Kyle Busch (Kyle Busch)
A: The correct answer is not specified, but the text
mentions that the ambassador of the campaign is unknown. (Kyle Busch)
<P> The fourth season Chicago Fire, an American drama television series with executive producer Dick
<P> The fourth season Chicago Fire, an American drama television series with executive producer Dick Wolf' and producers Derek Haas Michael Brandt
Wolf. and producers Derek Haas Michael Brandt
and Matt Olmstead,
ordered on February 5, 2015 by NBC and premiered on October 13, . 2015 and concluded on May 17, 2016 The season contained 775 episodes </P>
ordered on February 5, 2015 by NBC and premiered on October 2015 and concluded on May 17 2016 The season contained 775 episodes. </P>
and Matt Olmstead
Based on this text, answer these questions: Q: how many episodes are in chicago fire season 4? A: 775 (775)
Based this text, answer these questions:
Q: how many episodes are in chicago season 4? A: 2015, 2016, and the first and last episodes of the season. (775)
Answer these questions:
Answer these questions:
Q: how many episodes are in chicago fire season 4?
Q: how many episodes are in chicago season 4?
A: 23 episodes (23)
A: 22 (23)
C Details: Functional Roles of In-Context and Parametric Heads
Here, we provide additional details regarding experiments and results in $ For prompts, do not use chat template as it yields worse results in our preliminary experiments.
C.1 Functional Maps
Through the bar plot, we show layer-wise relevance distributions by summing the relevance each head assigns to the question or answer tokens. This reveals whether a particular layer oriented towards cnerruction-oolloiinn heads with larger markers.
retrieval. In addition, we highlight the top-scoring task and retrieval
Figure o the functional maps for Mistral-7B-Inntruett- In---nntttuto... and Gemma-2 2-99-- Consistent with 1--.-a...-------------- trated band of in-context heads in the middle layers, flanked by parametric heads in the early and late layers. We hypothesize that these early parametric heads may serve to enrich the prompt with relational knowledge, allowing later in-context heads to effectively integrate this information across the entire prompt, while later retrieval& & parametric heads extract the answer. This intriguing pattern suggests potential general principle governing transformer architectures, raising the question of whether this structure a universal feature of language models. Understanding why gradient descent naturally converges to this form presents an exciting direction for future research.
Figure [ these models exhibit strikingly similar structure: concen-
C.2 Disentangling Functional Roles of In-Context heads
We also compute the sorted in-context scores for results are visible in Figure 7 While some task heads show strong in-context behavior, many fall in the moderate range. Notably, Gemma even exhibits parametric task heads, indicating that task heads not exclusively in-context.
and Gemma-2-9B-it. The
30

---

Doekee Prmmeetr Teads
Hodds Peneeetic Houce
Too de- Hords
Sartinsd Eeleccnc
10
15
Lapper
20
25
30
35
Ifrer
Hoand
(a) Mistral- Instruct-vO.3
(b) (eemma-2-9Ban
Figure 6: Functional map of in-context and parametric heads in Mistral-7B Dre------nntttt..... and Gemma- 2-9B-it. Note that the number of attention heads in Gemma is 672, while Mistral contains 1024 heads. The bar plot shows layer-wise of heads' relevance wrt. question tokens answer tokens located within the context. We highlight the top 40 task and retrieval heads with larger markers.
food
Hooh
00
-14
600
30
600
130
(a) Mistral-7B- -
(b) Gemma-2-999
Figure 7: Sorted in-context scores for all heads of - and Gemma-2-9B-. comparing open-book and closed- book settings via score D. Positive scores indicate in-context behavior, while negative scores reflect parametric behavior. Retrieval heads (green) and task heads (blue) are predominantly high-scoring in-context heads.
C.3 Task Heads
Experiment Details To assess whether the activations of task heads genuinely capture information about the intensional frame - the instruction the model aims to follow we extract the head outputs z3 at the final token position. These outputs are then saved and directly patched into unrelated prompts in zero(one)-ste -shot manner without averaging across heads. Each head is patched separately, maintaining the unique contribution of each task head.
For this evaluation, we use the biography dataset described in Appendix A Each sample is annotated with four distinct attributes: birthdate, birthplace, educational institution, and profession. For each entry, we use four different questions, including "At which date was he born?". "In which place was he born?" "At which institution was he educated?", and "What was his profession?".
append each question (in bold exemplary for one question) to the following biography entry, and extract the FVs in single pass last token position:
31

---

Table 6: Task heads' FVs evaluation. We observe that the random score can occasionally be 20%.
This occurs because, without FV patching, the model sometimes repeats the input prompt verbatim.
Random
(40) usik
7
(40) tnsi..tt
"Ar which date was he born?"
Llama 3.1 Mistral v0.3
19%
97%
77%
6%
97%
86%
Gemma
2%
99%
22%
"AI which institution was educated??
Llama 3.1
6%
98%
96%
Mistral v0.3
2%
89%
43%
Gemma 2
3%
89%
74%
"In which place was he born?"
Llama 3.1
25%
97%
93%
Mistral v0.3
10%
74%
Gemma 2
12%
83%
82%
"What was his profession?"
Llama 3.1
22%
87% %
73%
Mistral v0.3
20%
73%
72%
Gemma 2
13%
81 %
54%
Vladimir Vapnik was born on 06 December 1936 in Tashkent, a city that would later shape his life's ' work. As a young man, Vapnik was fascinated by the potential of machines to learn and adapt. He went on to study at the V.A. Trapeznikov Institute of Control Sciences, where he was exposed to the latest advancements in computer science and artificial intelligence. It was here that Vapnik's passion for machine learning truly took hold. After completing his studies, Vapnik went on become a renowned computer scientist, making groundbreaking contributions to the field. His work support vector machines and the Copnik-Chervonenkis dimension would have a lasting impact on the development of machine learning algorithms. Throughout his career, Vapnik has been recognized for his innovative thinking and dedication to advancing the field of computer science. His legacy continues to inspire new generations of researchers and scientists. Q: At which date was he born? A:
Two key hyperparameters influence this method: 1. the number of task heads selected per model for FV extraction, and 2. the extent to which the head activations are amplified to overwrite potentially conflicting instructions within the model's context. This is defined by
1 Z
(12)
where E IR is scaling factor. We conducted hyperparameter sweep over 5% of the dataset as development set, finding that a scaling factor of a performed well for Llama 3.1 and Mistral vO.3, while 3 was optimal for Gemma 2. Scaling factors that were too high disrupted the model's ability to generate coherent text; too small factors did not successfully change the model response. For consistency, select 40 task heads for Llama 3.1, Mistral v0.3 and Gemma achieving at least 80% recall accuracy across these models.
Table 6 summarizes these results for three different models with roughly similar parameter counts: Llama 3.1 (8B), Mistral v0.3 (7B), and Gemma (9B). To further investigate the role of in-context heads, we compare this performance against two configurations:
Ha0 Selecting the top 40 task heads.
(aD) {hg Hask Hex 1, 1...., :0): Selecting only task heads that are also strong in-context heads.
We observe that the strong in-context heads alone capture a significant portion of the recall score, suggesting they play a critical role in interpreting the intensional frame. However, the inclusion of weaker in-context heads still pushes the recall scores higher, indicating that a diverse set of heads contributes to broader coverage across the task space.
32

---

Qualitative Examples In Table 7 we show a qualitative example for all models. In the forward pass, no question is appended the biography entry. Then, we patch the function vectors and observe the models' response.
Table 7: Qualitative examples of task FVs patching.
Input:
Tim Berners-Lee, a renowned engineer, was born on June 1955 London. Growing up in the bustling city, he developed a passion for innovation and problem-solving. After completing his education, Berners-Leee went to study at The Queen's 's College, where he honed his skills in computer science and engineering. It was during this time that he began to envision the globe. As an engineer, Berners-Lee was well-equipped to bring his vision life. He spent years working tirelessly to develop the World Wide Web, a revolutionary technology that would change the face of communication forever. In 1989, Berners-l.ce le submitted a proposal for the World Wide Weh to his employer, CFRN, and the rest, as they say, is history. Today, Bemerr-L.es is celebrated a pioneer in the field of computer science and visionary.
way of sharing information across
3.ama-3.1-88-Insttn
The Queen's 's College.
The Queen's College where he studied.
The Queen's College
Tim Berners-Loe studied at The Queen's 's College, University of Oxford.
Gemma-2-9B-itt
Hc studied at The Qucen's 's College.
C.4 Parametric Heads
Experiment Details We select random 1,000 examples for each model where their closed-book answer is correct wrt. gold reference (measured by recall with threshold of 0.7). Then we split them randomly with a proportion of 2.5% train, 2.5% dev, and 95% test set in order to find whether it is necessary scale the output of parametric attention heads. We consider the scaling factor E [1, 1.25, 1.5, 1.75, 2, 2.25, 2.5, 2.75, 3). maximizing recall on the dev set. For the number of heads, we also consider Read E [10, 20, 30, 40, 50, 60, 70, 80, 90, 100] taken from the top scoring parametric heads identified in $ In our final results, we extract parametric FVs from the combination of train and dev sets in zero-shot manner, and apply them on the test sets. Table 8 shows recall scores on the development set with their optimal scaling factor and number of heads used.
Table 8: Zero-shot recall scores for parametric FVs along with their optimal scaling factor and number of parametric heads used on the dev set.
Recall
a Bhead
Random FVs
Llama-3, 3.1.....1-BBBB -nssttuttt
7.59
3
50
Mistral- 7B Instruuct---
8.81
50
Gemma-2-9B-it
5.97
2.25
50
Parametric FVs
Llama-3, 3.1-8B-InstBuue
45.69
Mistral- 7B- ---trrruutt--.3
40.02
1.25
50
Gemma-2-9B..
38.00
3
50
Table 9: Cloze-style statements and question prompts used extract and patch parametric FVs.
Attribute
Cloze Statement
Prompt
date of birth
[X] was born on
Answer these questions: Q: what is the birth date of [X]? A:
place of birth
[X] was born in
Answer these questions: Q: where was [X] born? A:
occupation
[X] worked as
Answer these questions: Q: what is the occupation of [X]? A:
education
[X] was educated at
Answer these questions: Q: where was [X] educated? A:
33

---

As illustrated in Figure 4 an attribute of an individual is converted to cloze-style statement, of which the parametric FVs then extracted from the last token position. The attribute is chosen randomly, to demonstrate that parametric FVs indeed contain the entities" information and not just particular attribute. Then, we insert the parametric FVs the final token of a question prompt of another unrelated individual, and also for all subsequent token generations. We show the cloze-style statement and the prompt we used to elicit the answer Table
Qualitative Examples In Table 10 we show several qualitative examples result of parametric FVs patching. We observe that parametrie FVs are able to induce the generation of attributes that belong to the original entity conditioned the question prompts.
Table 10: Qualitative examples of parametric FVs patching for all models.
Llama-3,1-2 --B---trreet
John Backus (computer scientist) born Answer these questions: Q: what is the occupation of
Helena Bonham Carter? A: computer scientist
Julie Gardner ftelevision producer) was educated at - Answer these questions: Q: what is the occupation
of Konrad Zuse? A: Konrad Zuse was British-bom American television producer, writer, and director.
Mistral-7BB-nnstruttt0..333
Santiago Calatrava (Technical University of Valencia) was born on + Answer these questions: Q: where was Hans Zassenhaus educated? A: He was educated at the University of Valencia, Spain, and the University
of Madrid, Spain.
John Steinbeck (Salinas) worked Answer these questions: where was Paul McCartney bom? A:
Salinas, California
Gemma-2-9B-it
Linus Torvalds (University Helsinki) born in Answer these questions: Q: where Chris Carter
educated? Chris Carter was educated at University of Helsinki.
Enissa Amani (comedian) was educated at - Answer these questions: Q: what is the occupation of John von Neumann? A; John von Neumann was comedian.
Performance of Parametric FVs In Table sets. We hypothesize that this is due to how different parametric heads seem to deal with different domains of parametric knowledge, whereas for contextual information this does not seem be the case. To test this, we patch parametric FVs from cloze-style statements into related question prompts where the attribute is shared. In Table TT are used.
that Hparam performs worse compared to other
that the recall score tend to increase more heads
Table 11: Zero-shot recall scores for parametric FVs on the test set, separated by attributes. The scaling factor we use is taken from Table 8
(a) Theaa 50
Date of birth
Place of birth Education Occupation
Parametric FVs
Alama-3.1-88 1..-8---Instree
2
13.33
49,82 55.68
51.10
69.36
Bestral-7--Instrutt-e
1.25
24.13
62.14
51.65
Gemma-2 2-9B-it
9.56
37.20
48.70
66.11
(b) Mhead 100
Date of birth
Place of birth Education Occupation
Parametric FVs Mama--311-BB-nsstie
2
37.41
49,99
52.05
74.59
Sostral-7B-Insttuct--00
1.25
35.21
59.15
62,50
72.96
Gemma-2-9B-it
3
11.87
40.62
42.05
63.84
34

---

C.5 Retrieval Heads
Experiment Details Following the previous analysis of task and parametric heads, we utilize the biography dataset for this experiment. Each entry is provided to the model without accompanying question, and we randomly insert a multi-token needle within the prompt. This process is repeated 10 times, each with different needle. The needles used are famous poem titles from around the world:
1. Al-Burda
2. Auguries of Innocence
3. Der Zauberlehrling
4. Ode Nightingale
5. She Walks Beauty
6. The Raven
7. The Road Not Taken
8. The Second Coming
9. The Waste Land
10. Über die Berge
As hyperparameter, only vary the number of retrieval heads. We use 5% of the dataset development set, where we select the smallest value of (top retrieval heads) that achieves recall score of approximately 90%. Hence, for Llama 3.1 and Mistral v0.3 use 40 heads, while for Gemma 2 we select 30 heads. To activate the copying behavior in the retrieval heads, we modify the attention weights concentrate the tokens of the needle. To allow for some adaptivity of the model, we use the following boosting scheme, such that the model can focus on subtokens inside the needle:
Let Jncedle be the set of token positions corresponding to the multi-token needle. Let Ås., denote the unnormalized attention weights (before applying the softmax function) at last query position. The modification performed in two steps:
Initial Needle Tokens Boost: This step prevents the attention weights from being zero before applying the softmax, ensuring that the model can effectively attend to the needle tokens:
Às. As, +5, Vj E Jneedle
(13)
2. Scaling Adjustment:
This step further amplifies the weights for needle tokens that the model already tends to focus on (having value bigger than zero), reinforcing the copying behavior:
Ás.5 <- 1000 As.
Vj
Jnendle
(14)
This two-step approach encourages the model to selectively copy one token at time, which found to be effective in practice. the initial boost (step 1) omitted, performance tends to degrade, as zero valued attention weights cannot be effectively scaled in the second step.
Qualitative Examples Table 12 illustrates example input with the needle "The Road Not Taken"
highlighted in bold.
D Details: Source Tracking
In the following, we provide details on the training procedure of the probe used for causal source tracking.
Training & Evaluation Let
-= (2((2))(0) h.(m) (((), (()))) /M
35

---

Table 12: Qualitative example of the seedle-in-a-haystace experiment.
Input:
Jimmy Wales, renowned orator, born 07 August 1966 in the vibrant city of The Road Not Taken Huntsville. Growing up. Wales was fascinated by the world of technology and innovation, which would later become the foundation of his illustrious career. After completing his education, Wales went on to pursue a degree in finance from Indiana University Bloomington. However, it was his passion for online communities and collaborative knowledge-sharing that led him to create one of the most influential websites of our time. As the co-founder of Wikipedia, Wales revolutionized the way people access and share information. His dedication to free and open-source knowledge has made a lasting impact on the digital landscape. Today, Wales continues to be a prominent figure in the tech industry. inspiring a new generation of innovators and thinkers. His remarkable journey is testament to the power of vision and perseverance.
-lama-3.1-88-Inttree
The Road Not Taken is reference to the famous poem by Robert Frost
Biarral7B-Instruct--ooe
The Road Not Taken
Geemma-2-9B-it:
*he Road Not Taken=#
be the set of N samples of NQ-Swap dataset with counterfactual contexts, where ((m) E NIVI denotes the target token index in the vocabulary V for sample and
(m) 1, predicted token g(n) is contextual (from external documents),
0, if predicted token ((x) is parametric (from model memory).
All samples include counterfactual entries, filtered to retain only those where: (1) the counterfactual answer object of appears among the top 10 predicted tokens, and (2) the correct closed-book parametric answer object o" also accurately predicted among the top 10 predicted tokens. This approach allows for direct comparison of parametric and contextual retrieval head activations for identical inputs, enhancing the probe's 's training quality.
We then learn weights {wn F over the selected set of retrieval heads Hzet of $ C.5 by solving
Σ
Lt's 11)
argmin wa Eaea
(15)
WE Has
An optimal decision threshold is then chosen via ROC analysis on a held-out development subset of 2, selecting the threshold that maximizes the true positive rate while minimizing the false positive rate.
To test localization, aggregate each head's attention map with the learned terms:
Äs.
Σ wh Ln L('s It) Ab3
(16)
hEH
We also experimented with simple averaging of the attention maps, but this approach resulted in approximately 10% lower scores across all models. We then predict the source token index as
i argmax, As.j-
Since localization is only meaningful for contextual samples, we restrict this evaluation to counterfac- tual samples from 2. Specifically, we compute the top-1 accuracy by checking whether & matches the ground truth token position of the counterfactual entry of.
In Figure g the logit lens scores of the top 40 retrieval heads in the Llama 3.1 model are illustrated. We observe that retrieval heads exhibit heightened activity when the model relies context, as indicated by the elevated logit lens scores in green color. While these distributions are 1-dimensional for each head, the probe itself learns a decision boundary in the full 40-dimensional space, where the retrieval signal may be better disentangled. Interestingly, some heads appear only sporadically highly active, suggesting a high degree of specialization a promising direction for future research.
In Figure g we illustrate the localization capabilities of the aforementioned method on four random samples for all three models. The aggregated attention maps are plotted heatmaps, where red
36

---

colors signify high values and blue colors signify negative colors. Note, that the attention maps are weighted with the probe weights, which can negative, allowing for negative superposition of attention maps. The beginning of sentence token is receiving some attention weight values due to its usage as attention sinks 84
E
Additional Results: Resolving Knowledge Conflicts
Table 13 reports the effect selectively ablating in-context and parametric heads on the NQ-Swap dataset, which contains counterfactuals designed to elicit knowledge conflicts between contextual and parametric information. The ablation performed incrementally by removing an increasing number of heads identified as belonging to either the in-context or parametric category. When in-context heads are ablated, the model is prevented from relying on contextual information the input and forced to depend more heavily its internal (parametric) knowledge. As shown in the upper part of the table, this generally leads to a substantial increase in recall score wrt. original, non-counterfactual gold answers, indicating that the model gets better at resisting the misleading context and grounds its answers in the stored knowledge. However, the improvement is not strictly monotonic across all models suggesting that the influence of an individual head is not uniform.
Conversely, when parametric heads are ablated (lower part), the model is encouraged to ground its predictions more strongly on the provided context while suppressing the reliance on its internal memory. We again observe substantial gains in recall score wrt. counterfactual open-book answers, although the relationship between the number of the removed heads and the performance is non- monotonic for some models.
Table 13: Ablation of in-context vs parametric heads on NQ-Swap. Recall (%) wrt. gold answer (closed-book) for ablating in-context heads and wrt. counterfactual answer (open-book) for ablating parametric heads.
Ablation (Nr. Heads Removed)
20
60
80
100
In-Context Heads
Llama-3.1-8B Instruct
9.5
24.1
24.3
35.1
33.1
Mistral-7B-Insttan 1B-Instruct-vO.3
12.0
35.6
37.3
37.6
45.0
45.9
Gemma-2-99-
13.3
27.7
23.5
20.3
14.6
9.2
Parametric Heads
Llama-3. 3.1-8B 1-8B-Instruet
68.7
72.5
73.2
74.2
75.7
76.4
Mistral-7B-Inssse IB-Isttruct--O.
67.6
72.9
75.0
75.3
72.6
62.0
Gemma-2-9Ba
66.6
69.3
71.9
73.5
72.8
73.9
F
Additional Results: Machine Translation
To further demonstrate the generalizability of our approach, we perform evaluation machine translation using the OPUS Europarl datase (en-fr, en-es) where we align the first 10,000 examples. The attention heads used for patching are retrieved following the procedure described in $ where we select the top 15 heads exhibiting the highest relevance scores on their key positions for the instruction "Translate into German" " in the seed prompt "Trunslate into German: 1 love Al research." We intentionally choose German as neither source nor target language in subsequent evaluations to demonstrate that the extracted feature vector does not encode language-specific information, but rather the general translation instruction.
In the patching setup, we provide the model with inputs of the form "[SENTENCE] - and parse the model output following the arrow as the generated translation. The function vector is applied in zero-shot fashion to induce translation behavior across language pairs. For Llama-3. 3.1-8B....... -nsttru we use scaling factor of and for -nnstutt- scaling factor of
37

---

Table 14: BLEU score for -French and Bagglis-+Frenme translation. implementation is used to compute sentence-level BLEU score.
SacreBLEU [59
Spanish-FFennhh French
English French
Baseline
Random
MT Heads
Baseline
Random
MT Heads
Maana-3.1-SB-nntteee 33.2
1.0
31.5
29.5
0.37
24.9
Mistral-7B- Daaiial-7--nstruet-v0...
29.2
18.2
25.2
0.60
18.3
To assess the effectiveness of this approach, we compare it against two baselines (Table [4] (I) direct prompting with "Translate the sentence into [LANGUAGE]: [SENTENCE] - and (2) randomly patching attention heads. Our method achieves comparable performance to the explicit prompting baseline demonstrating that the patched task heads reliably induce translation behavior across languages. Notably, this effect is achieved using only 15 of the 1024 attention heads in the I-ntsttuttt and Disstall7B-Innsrruee-vee identified translation subspace. In contrast, random patching yields near-zero BLEU score, confirming that translation behavior arises from targeted task heads selection.
model, underscoring the precision of the
Unfortunately, the Gemma-2-9B-it model exhibits instability in this setting: it frequently fails to produce outputs adhering the [SENTENCE] (TRANSLATION] format, making translation parsing unreliable. We leave this issue for future investigation, but we hypothesize that this instability stems from the comparatively small number attention heads Gemma 2, which likely leads to stronger superposition effects and less clearly separable functional subspaces compared to models such Llama 3.1 and Mistral v0.3.
38

---

0:30
6.30
232
248
10
11 Value
310
Loglt Lane Value
Logt Lans Valne
Legit Valto Hend 900
Logt Lare Valoo Ised aED
Lird 500
Hind 780
Hend 845
0.10
Desett
DG
-30
Leght Lere Valten
Leght Lete Valuee
Leght Vilue
Lught Valtee
Lught Voluee
0L3S 0.10
0.0
2.0
0.00
0000
Value
Vilie
Vilue
Hoad 1012
Heal th4
Hoad so6
Head 804
Head 844
0.39
32-
82-
-2
15 Logit Lore Value
12
Lost Leaa: Valuo
Lagth Lome Value
Lust Lema Vilue
Logn Lara Value
8.36 6.13
3.15
3.65
0.00
0:00
0.02
Vilie
Lagit Hood ast
Vilie
Lagit Ville Hard 1015
Hand 454
Head 750
Band 750
0:35
3.9
3.15
3.69
011
15
8
33
Logit LoEs Valae
Loght Loas Viluee
Loght Less Villae
Loght Less Value
Loght Leas Value
Drooty
0.00
Logh:
Loght
Loght
Logil
Logit
Tloott
Tinu
Head deo
Head Wi3
Head
0.1
033
Deratto
4.30
0.33
2.00
3.65
055
a.
Logit Loas Valae
Logil Leas Vilue
Logil Leta Villa
Logil Leto Viluo
Logit Loaa Vidue
Figure 8: Distributions of logit lens scores for the top 40 retrieval heads in Llama-3. -lamm-33.--8B-mme I-8-ttrruut... Shown are the logit lens activations for the ground truth and counterfactual output tokens, comparing cases where the model generates the answer from its parameters (red) versus cases where it retrieves the answer from the context (green) respectively. We observe that retrieval heads exhibit heightened activity when the model relies on context, as indicated by the elevated logit lens scores in green color. While these distributions are 1-dimensional for each head, the probe itself learns decision boundary in the full 30-dimensional space, where the retrieval signal may be better disentangled. Interestingly. some heads appear be only sporadically highly active, suggesting high degree of specialization a promising direction for future research.
39

---

Predicted Loken: Mount
Ulma 3.1 (80)
Mistral voL3 (Ts)
Genma (01)
thatts
watl
the Phelle Hag
Mitool staai
shall actins
actiay, beglin daace
klage, Hiida Salce
deque, uik vokaaa ascuad blngii sugeests aquast ans,
couutatt haluan the boartts of le bed in the oidc
bokice the the kettl bed witk the 100 Soore.
balaee the keatto of berta kool witk The pec serora roll graaiis
Icoaa a Lakeniic Auriic Eae must wtll wol aas is da
la Aurun the andd ander volita ladbe
and
lo kikiumis kumme Tis cask octire rnd mnliaan in kude aeiiii
Predicted token: Jordan
Linnu 1.1 186)
Mateal v0.3 (Ts)
Ocnana
dike calan sanve tae saa,, oik dore, clang
They
34 good soong yrlod dens oe iit
Midjon fook
Vakae idid lodatis lntla
date biikleg souad lene
vil cadki dlaase Bilk tiliiig Monioio, viich has bakoey lathor
Altess
tull cancal by
caudy anstim citeel bline Thes righon
then widh domtl satail Duick taep Midjac Eate Vimeaii Singae pourt
ais Anaee Brnpeel cateppue fes segirs for the
danee giaet Aarunt Yorjoct
dise girs Life will oe des the
Predicted token: 8
Matral v6.
Centr Soatt arrante sase df acleiia Milly
Elmill Sheal Wome claeni lave Btadioe bolbil Mikky tren wal Thee comente ohhnk ment
cotth
statt roed.
to the
ronded pallii too
Yorth
de 12l mceeitac
lnd lid clicke abow Shi507 the tnagg Haakay tooker csein
BlN
aped
of
rilo ta to The
Squater thyy cle such
pr the tea
Eae lick 2002 Airams
Andion wode Red on hat
incil 1802
Predicted token: Philosoph
Tiena 3.1 (80)
Midrral v03 (71)
Gemms R (5h)
coneiquua Noman
aindal Lavoe, Wourd algi Siglink
los dectlutti the old loc the betein tare
well
Tamora to
bely
dlan hand Eae callano
de thr Earth Doyael plpttss dindlged caliiie aoead
bamer work
Figure 9: Heatmaps of the weighted aggregation of retrieval heads' attention maps at the final query position superimposed on the input prompt to pinpoint the retrieved source token. For each model, the aggregated attention maps of the retrieval heads reliably focus on the predicted token in the context, which can be used as cost-effective source tracking. The beginning of sentence token is receiving some attention weight values due its usage as attention sinks (84).
40

---

AttaLRP (In-context)
AttaLRP (Parametric)
AWR
Random
NQ-Swap (Counterfact)
NQ-Swap (Oracle)
NQ-Swap (DPR)
NQ-Swap (Closed-book)
100
Recall
20
20 40
GU
80 100
20 40 80 80 100
20 40 60 80 100
20 40 60 80 100
Num Heads Ablated
Num Heads Ablated
Num Heads Ablated
Num Heads Ablated
TQA (Oracle)
TQA (DPR)
TQA (Closed-book)
100-
80-
Recall
60-
40-
20
0
20
40
60
80
100
0
100 Num Heads Ablated
20 40
60
80
0 20
40
60
90 100
Num Heads Ablated
Num Heads Ablated
Figure 10: Recall analysis for Llama-3.1-88 1--1--8-IIsssreee when either in-context or parametric heads ablated on NQ-Swap and TQA under various configurations. For DPR, we instances where K-Precision 3 is equal to 1 in the non-ablated run.
AttaLRP (In-context)
AttaLRP (Parametric)
AWR
Random
NQ-Swap (Counterfact)
NQ-Swap (Oracle)
NQ-Swap (DPR)
NQ-Swap (Closed-book)
Recall
20
04
20
40
60
so 100
20 40 do 100 Num Heads Ablatod
40 60 so 160
0
40 60 so 100 Num Heads Ablatod
Num Heads Ablated
Num Heads Ablated
TQA (Oracle)
TQA (DPR)
TQA (Closed-book)
100
80-
Rroil
60
40-
20
0
0
20
40
60
80
100
0
20
40
60
80 100
20
40
60
80
100
Num Heads Ablated
Num Heads Ablated
Num Heads Ablated
Figure 11: Recall analysis for Mistral-7B-Inn In when either in-context or parametric heads ablated on NQ-Swap and TQA under various configurations. For DPR, we instances where K-Precision 3 is equal to 1 in the non-ablated run.
41

---

AttaLRP (In-context)
AttaLRP (Parametric)
AWR
Random
NQ-Swap (Counterfict)
NQ-Swap (Oracle)
NQ.Swap (DPR)
NQ-Swap (Closed-book)
100
Recall
20
20 40
60
80 100
20 40 80 80 100
20 40 60 80 100
20 40 60 80 100
Num Heads Ablated
Num Heads Ablated
Num Heads Ablated
Num Heads Ablated
TQA (Oracle)
TQA (DPR)
TQA (Closed-book)
100-
80
60
40-
20
0
20
40
60
80
100
0
20
40
60
80
100
20 40 60 90 100
Num Heads Ablated
Num Heads Ablated
Num Heads Ablated
Figure 12: Recall analysis for Gemma-2-9B-in when either in-context or parametric heads are ablated on NQ-Swap and TQA under various confi configurations. For DPR, use instances where K-Precision is equal to in the non-ablated run.
AttaLRP (In-context)
AttaLRP (Parametric)
AWR
Random
NQ-Swap (Counterfact)
NQ-Swap (Oracle)
NQ-Swap (DPR)
NQ-Swap (Closed-book
EM
20
04
20
40
100
20 40 do 100 Num Heads Ablatod
40 60 so 100
20 40 60 so 100 Num Heads Ablatod
Num Heads Ablated
Num Heads Ablated
TQA (Oracle)
TQA (DPR)
TQA (Closed-book)
100
80-
60
40-
20
0
20
40
60
80
100
0
20
40
60
80
100
0
20 40
60
80 100
Num Heads Abbated
Num Heads Ablated
Num Heads Ablated
Figure 13: EM analysis for Alama-3.1-88 Instruct when either in-context or parametric heads are ablated NQ-Swap and TQA under various configurations. For DPR, we use instances where K-Precision 3 is equal to 1 in the non-ablated run.
42

---

AttaLRP (In-context)
AttaLRP (Parametric)
AWR
Random
NQ-Swap (Counterfict)
NQ-Swap (Oracle)
NQ-Swap (DPR)
NQ-Swap (Closed-book)
100
EM
20
20 40
GU
80
100
20 40 80 80 100
20 40 60 80 100
20 40 60 80 100
Num Heads Ablated
Num Heads Ablated
Num Heads Ablated
Num Heads Ablated
TQA (Oracle)
TQA (DPR)
TQA (Closed-book)
100-
80
60
40
20
0
20
40
60
80
100
0
20
40 60 80 100
0 20 40 60
90 100
Num Heads Ablated
Num Heads Ablated
Num Heads Ablated
Figure 14: EM analysis for Misttal-7B--nnstruct--)) when either in-context or parametric heads ablated on NQ-Swap and TQA under various configurations. For DPR, we instances where K-Precision 3 is equal to 1 in the non-ablated run.
AttaLRP (In-context)
AttaLRP (Parametric)
AWR
Random
NQ-Swap (Counterfact)
NQ-Swap (Oracle)
NQ-Swap (DPR)
NQ-Swap (Closed-book)
EM
20
04
20
40
60
100
20 40 do 100 Num Heads Ablatod
40 60 so 100 20 40 60 so 100 Num Heads Ablatod
Num Heads Ablated
Num Heads Ablated
TQA (Oracle)
TQA (DPR)
TQA (Closed-book)
100
80-
60
40-
20
0
20
40
60
80
100
0
20
40
60
80
100
20 40 60
80
100
Num Heads Abbated
Num Heads Ablated
Num Heads Ablated
Figure 15: EM analysis for Gemma-2-9B-i when either in-context or parametric heads are ablated on NQ-Swap and TQA under various configurations. For DPR. we use instances where K- K-Precision is equal to in the non-ablated run.
43

---

AttaLRP (In-context)
AttaLRP (Parametric)
AWR
Random
NQ.Swap (Counterfact)
NQ-Swap (Oracle)
NQ-Swap (DPR)
NQ-Swap (Closed-bool)
100
BBM
20
20 40
GU
80 100
20 40 80 80 100
20 40 60 80 100
20 40 60 80 100
Num Heads
Ablated
Num Heads Ablated
Num Heads Ablated
Num Heads Ablated
TQA (Oracle)
TQA (DPR)
TQA (Closed-book)
100-
80
BEM
60-
40-
20
0
20
40
60
80
100
0
20
40
60
80 100
0
20
40
60
90 100
Num Heads Ablated
Num Heads Ablated
Num Heads Ablated
Figure 16: BEM score analysis for Llama-3. 1.ama-3.1-8B. --St--nstt when either in-context or parametric heads ablated on NQ-Swap and TQA under various configurations. For DPR, we use instances where K-Precision 3 is equal to 1 in the non-ablated run.
AttaLRP (In-context)
AttaLRP (Parametric)
AWR
Random
NQ-Swap (Counterfact)
NQ-Swap (Oracle)
NQ-Swap (DPR)
NQ-Swap (Closed-book)
100
BEM
20
04
20
40
60
so 100
20 40 do 100
20 40 60 so 160
20 40 60 so 100 Num Heads Ablatod
Num Heads Ablated
Num Heads Ablatod
Num Heads Ablated
TQA (Oracle)
TQA (DPR)
TQA (Closed-book)
100
80-
BEM
60-
40-
20
0
0
20
40
60
80
100
0
20
40 60 80 100
20 40
60
80
100
Num Heads Ablated
Num Heads Ablated
Num Heads Ablated
Figure 17: BEM score analysis for 7B-hnsllttt--nttuttt--- heads are ablated on NQ-Swap and TQA under various configurations. For DPR, instances where K-Precision
when either in-context or parametric
3 is equal to in the non-ablated run.
44

---

AttaLRP (In-contexd)
AttaLRP (Parametric)
AWR
Raadom
NQ-Swap (Counterfact)
NQ-Swap (Oracle)
NQ-Swap (DPR)
NQ-Swap (Closed-book)
100
BBM
20 40
60
80 100
20 40 80 80 100
20 40 60 80 100
20 40 60 80 100
Num Heads Ablated
Num Heads Ablated
Num Heads Ablated
Num Heads Ablated
TQA (Oracle)
TQA (DPR)
TQA (Closed-book)
100-
80-
BEM
60
40-
20
0-
0
20
40
60
80
100
0
100 Num Heads Ablutod
20
40 60
80
0
20
40
60
80
100
Nur Hends Ablated
Num Heads Ablatod
Figure 18: BEM score analysis for Gemma-2 2-9B-it when either in-context or parametric heads are ablated NQ-Swap and TQA under various configurations. For DPR, we use instances where K-Precision g is equal to 1 in the non-ablated run.
45


Reviews:
[{"review_id": "eMeGwLhrPn", "reviewer": "Reviewer_J8iq", "summary": "This work presents probing results aimed at understanding the inner workings of large language models (LLMs) in the context of question answering. Building on the existing AttnLRP method, the authors identify three types of attention heads—in-context (input-dependent), parametric (internal knowledge-centric), and retrieval-based (copying)—and explore their respective roles in solving QA tasks using the NQ-Swap and TriviaQA datasets.\n\nIn Section 5, the authors argue that their method outperforms both a random baseline and Attention Weight Recall (AWR), thereby justifying the use of AttnLRP. In Section 6, they observe that removing in-context heads leads to performance degradation in closed-book settings. Based on this finding, they conjecture that in-context heads not only process contextual information but also play a role in interpreting the so-called “intensional frame”—the semantic structure imposed by the task instruction. Within this framework, the authors suggest that task heads—although it remains unclear whether these are a subset of in-context heads or merely share some overlapping functions—are responsible for encoding the intensional frame (e.g., the subject–relation–object structure), while parametric heads encode information about the subject entity.\nThis hypothesis is empirically supported by experiments involving function vectors derived from the corresponding attention heads. In Section 7, the authors further train a linear probe to examine whether the representations involved in the retrieval mechanism are linearly separable.", "questions": "[Follow‑up questions after rebuttal]\n\n1. About the (dataset and task) coverage of this work\n\nI acknowledge that this work aims to address more complex cases by considering two QA datasets. However, I do not believe this is sufficient to claim that it covers most real‑world scenarios. As the authors themselves note, the evaluation is limited to cases involving short‑form answers. The overall storyline appears to discuss the inner workings of ICL from a general perspective, without restricting the discussion solely to QA datasets. (Although, in the abstract, the authors mention that “we shed light on the mechanism of in‑context retrieval augmentation for question answering,” the introduction, in my reading, presents the discussion in the broader context of general in‑context learning rather than only QA.) If this interpretation is correct, it would be difficult for most readers to agree that evaluation on just two QA datasets is sufficient to support the claim that the findings will generalize to all possible uses of LLMs, including tasks requiring long‑form answer generation.\n\n2. In‑Context Learning vs. In‑Context Retrieval Augmentation\n\nContinuing from the above discussion, what is the exact difference between in‑context learning and in‑context retrieval augmentation?\nAs these terms are used interchangeably in the draft, it would be beneficial to clearly define their commonalities and distinctions. My understanding is that in‑context learning is the broader concept, whereas in‑context retrieval augmentation—the main focus of your work—represents a narrower, more specific instantiation of it.\nIn relation to your experimental settings and target tasks, a clearer and more consistent use of these terms would help contextualize the contribution and scope of your work.\nI am still somewhat unclear on whether this work aims to generalize its findings from QA tasks to broader use cases, or whether it is intended to focus exclusively on QA tasks. \nPlease clarify this point.\n\n3. Insufficient baselines (or [Unsuitable baselines])\n\nAlthough the authors defended their choice of using only one baseline with the following justification: “Regarding baselines, we note that AWR [4] is, to the best of our knowledge, the most recent and relevant method capable of extending to multi‑token settings. Other closely related work [5] is limited to single‑token responses and does not provide code, which makes direct comparisons difficult.”\n\nI appreciate their response, but I remain unconvinced by this justification. \nAs the authors themselves noted, the target datasets—TriviaQA and NQ‑Swap—require relatively short‑form answers. \nFor example, in the TriviaQA dataset (see example instances in https://huggingface.co/datasets/mandarjoshi/trivia_qa/), answers often consist of only two words.\nGiven this, I believe the authors could make a reasonable effort to adapt existing baselines to these settings without substantial modifications. \nAlternatively, they could consider evaluating on subsets of the datasets containing instances with single‑word answers to allow a fairer comparison with existing work.\nIn summary, I believe the contribution of this work could and should be directly compared to prior approaches to more clearly highlight its unique contributions.\n\n4. Discussion about [Unclear novel contributions wrt. prior studies], [Usage of special entity-based dataset], [Knowledge conflict results]\n\nThank you for the clarification.\nIt would help readers better understand the context if the corresponding parts were revised to guide them more clearly in line with your intentions. \nIn other words, the writing could be further refined to improve the clarity and effectiveness of idea communication.\n\n5. Clear definitions of some key words (continued on the question 2)\n\nKeywords used in this work, such as function vectors and in‑context heads, are already widely adopted in the related literature. However, in some cases, the definitions of these terms in this work appear to differ from their usage in prior studies. It would therefore be highly beneficial to clearly define the intended semantics of these keywords in the draft.\nFor example, are the function vectors mentioned in this work equivalent to those described in \"Todd et al., Function Vectors in Large Language Models, ICLR 2024\" [62]? If not, what are the precise differences?\nThis point is also connected to the discussion in the authors’ response under the section [Unclear relations between heads].", "limitations": "[Follow‑up discussion on the limitations of this work after rebuttal]\n\nI believe that the limitation discussed in the draft—namely, that the proposed method is applicable only to attention heads and not to other components of Transformers—is not sufficient to meaningfully guide revisions of this work or inform related future research.\nA more substantive limitation is the narrow coverage in terms of tasks and methodologies, which should be addressed. \nFor instance, investigating whether the findings can be generalized to more complex QA tasks or to entirely different task types would be both meaningful and desirable, potentially as part of the revision of this work itself.\n\nIn this context, I appreciate the authors’ efforts during the rebuttal process to report the [In‑context MT results]. I agree that the inclusion of such additional experiments can strengthen the research.", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths\n- The paper presents a series of analyses grounded in existing and reliable techniques, including AttnLRP.\n- It attempts to verify that specific attention heads in LLMs play distinct roles in processing question answering tasks.\n\nWeaknesses\n- It remains unclear what novel contributions this work offers. The study builds on existing methods and identifies attention head categories—such as in-context and parametric heads—that are already well known. What distinguishes this work from prior studies, beyond reinforcing existing findings?\n- Some methodological details are ambiguous. For instance, why is a special entity-based dataset used specifically in Section 6 but not in other experiments? Additionally, the relationships between key concepts such as (1) in-context heads, (2) task heads, and (3) retrieval heads are not clearly defined. Several terms are introduced without precise explanation, which may confuse readers.\n- Although Section 5 provides empirical evidence that the adopted method outperforms a random baseline and AWR, it is not evident that this is the most suitable technique for the type of analysis conducted in this work. Given the rich body of literature on probing LLMs, further justification or comparisons with additional baselines would strengthen the claim.\n- The attention head types examined in this work (e.g., in-context and parametric heads) are already well-established in related literature. Consequently, the contribution of this study is unclear, especially since the analysis closely resembles previous work. Furthermore, the scope of the analysis is quite limited—it covers only two QA datasets and three relatively small LLMs—raising concerns about the generalizability of the findings and whether the models analyzed are representative of truly large-scale, capable LLMs.", "was_revised": true, "final_justification": "As the discussion phase is still ongoing, the final justification will be determined after its conclusion. \nAt this stage, I acknowledge the authors’ responses; however, I believe that the score I have assigned remains reasonable.", "rebuttal": "We sincerely thank the reviewer for the constructive feedback and for the time spent evaluating our manuscript. We are grateful for the recognition of our rigorous analysis uncovering the distinct roles that attention heads in LLMs fulfill. However, we respectfully offer the following clarifications to address the concerns raised, which we hope will highlight the novelty of our contribution.\n\n**[Unclear novel contributions wrt. prior studies]**: While prior studies have indeed identified heads with behavior similar to in-context and parametric heads, as we acknowledged and discussed in our related work section, our findings substantially advance the understanding of these mechanisms in several important ways:\n- Prior studies mostly focused on **narrow domains where the answer is composed of a single token, far from realistic use cases**. Here, we analyse popular open-domain QA datasets such as NQ-Swap & TriviaQA along with a human biography datasets that we curate, **which have short-form answers and better reflect the complexity of real-world language tasks**.\n- For the first time, we show that **in-context heads are considerably more complex than previously thought and can specialize to task or retrieval heads** which provides a finer-grained categorization than is available in existing literature. We also empirically demonstrate their causal roles in shaping the model’s internal representation (via intensional frames) along with parametric heads.\n- Prior works were restricted to **simple toy tasks due to the computational costs of the employed methods**. Our novel approach of heads filtering, along with the efficiency of AttnLRP allows **better scaling and enables extraction of steerable function vectors (FVs) in a zero-shot manner for complex tasks**, as described in Section 6.2. To substantiate how our method also generalizes beyond QA, we have included additional preliminary experiments on machine translation, detailed below.\n- We also introduce **a novel post-hoc procedure to trace both contextual and parametric knowledge usage** in LLMs, thus enabling the identification of potential hallucinations, knowledge conflicts and citations (Section 7). We have added additional experiments demonstrating that our method can effectively mitigate knowledge conflicts. While we present our findings as a preliminary step, we believe the results are promising and plan to develop this direction further in future work, especially since this line of work has become increasingly important in the community [1,2,3].\n\n**[Usage of special entity-based dataset]**: We would like to clarify that all major experiments (including the extraction of in-context, parametric, retrieval, and task heads)  are conducted on **realistic QA datasets**. The entity-based synthetic dataset is used solely for **controlled experiments** aimed at validating specific functional hypotheses: namely, that parametric heads encode entity attributes, task heads capture intensional frames, and retrieval heads copy answer tokens verbatim. Such fine-grained ground truth supervision is difficult to attain with realistic datasets, making controlled settings essential for isolating these behaviors, especially since entities might not possess the same attributes that we want to investigate across datapoints. We appreciate the reviewer’s feedback and will revise the manuscript to make this distinction clearer.\n\n**[Unclear relations between heads]**: We appreciate the reviewer’s request for clearer distinctions between in-context, task, and retrieval heads. These distinctions are visually and conceptually defined in Figure 1, in the introduction (L36-37), and further explained in Section 6.1. That being said, we agree that their relationship could be emphasized more strongly, and we will revise the manuscript to clearly state that **in-context heads can specialize into either task or retrieval heads, depending on whether they encode intensional frames or retrieve contextual facts**.\n\n**[Unsuitable baselines]**: Regarding baselines, we note that AWR [4] is to the best of our knowledge, the most recent and relevant method capable of extending to multi-token settings. Other closely related work [5] is limited to single-token responses and **does not provide code**, which makes direct comparisons difficult. \n\nWhile we outperform random and AWR baselines, our primary focus here is to enable a fine-grained analysis of attention heads in large-scale, realistic datasets, targeting heads that causally affect the answer generation process during in-context retrieval augmentation. We would like to emphasize, however, that **our method is agnostic to the choice of feature attribution technique and is compatible with a wide range of methods**. As detailed in Section 3 (L82–94), many existing approaches are not well-suited for our setting due to **limitations in scalability or faithfulness**. In particular, our need to trace and filter attention heads across large datasets necessitates a method that is both computationally efficient and faithful in attributing model behavior. For this reason, we employ AttnLRP, as justified in Section 3.1 (L105–113).\n\nTo illustrate the computational challenge with less efficient methods, consider path patching [6], which roughly requires one clean and one corrupted forward pass per head. With 1024 attention heads, this results in 2048 forward passes per sample. On a **smaller test dataset** of NQ-Swap, where a single run takes approximately 15 minutes on NVIDIA GeForce RTX 4090, this would amount to **0.25h × 2048 ≈ 31.3 days, making such methods infeasible at scale**.\n\n**[Limited scope of analysis]**: **Given relevant related work, we can confidently state that the scope of our experiments goes beyond the usual extent of the current state of the art.** Compared to toy settings and single-token answers used in prior works, our analysis on two open-domain QA datasets along with the curated biography dataset represent a substantial step forward, especially wrt. a more realistic use case. We acknowledge that we don’t have access to very large-scale models due to computational constraints. However considering our experimental settings and fine-grained analysis, we would like to point out **our analysis of 7-9B parameter models is much bigger compared to those used in [5] relevant to this work** (the authors mostly experimented on toy data with single-token answers with short contexts), which is made feasible by our careful choice of an efficient and faithful attribution method. As it stands, we believe our contribution is also applicable to industrial-scale LLM analysis.\n\n**[In-context MT results]**: To further highlight the generalizability of our approach, we conducted experiments on machine translation on the OPUS Europarl dataset (on French, Spanish, English). In these experiments, we identified task-specific heads responsible for translation in a zero-shot manner. We then compacted these heads into a feature vector representation and patched them into different prompts (again in a zero-shot fashion) successfully inducing translation behavior into various languages. To evaluate the effectiveness of this approach, we compared it against two baselines: (1) directly prompting the model with “Translate the sentence into [LANGUAGE]:” and (2) randomly patching attention heads. Our method **matched the BLEU score of the explicit prompting baseline**  (implemented with SacreBLEU, 33 points on Es-Fr, 29 points on Fr-En), demonstrating that the patched task heads reliably induce the desired translation behavior. Notably, this was achieved using only 15 of the 1024 attention heads in the LLama 3.1-8B-Instruct model, **highlighting the preciseness of our approach**. Random patching has a BLEU score close to zero. **We believe these results substantially strengthen the contribution of our paper.**\n\n**[Knowledge conflict results]**: We computed the accuracy of LLama 3.1-8B-Instruct on the NQ-Swap with counterfactual contexts. Computing the recall score on the gold answer instead, the model is clearly distracted by the wrong contexts achieving only 9% accuracy. Removing the top 60 in-context heads reduces the reliance on the wrong context pushing the score to 35% which is similar to the closed-book accuracy of the model. Conversely, computing the recall score for counterfactual answers i.e. reliance on the counterfactual inputs, the model achieves 68% accuracy without ablation. Ablating the top 100 parametric heads reduces the reliance on internal knowledge and the model is more likely to accept the counterfactual context increasing the score to 76%.\n\nWe sincerely thank the reviewer for the time spent engaging with our work and hope we could clarify all points raised.\n\nReferences:\n\n[1] Xu, Rongwu, et al. \"Knowledge conflicts for LLMs: A survey.\" arXiv preprint arXiv:2403.08319 (2024).\n\n[2] Gao, Tianyu, et al. \"Enabling large language models to generate text with citations.\" EMNLP (2023).\n\n[3] Qi, Jirui, et al. \"Model internals-based answer attribution for trustworthy retrieval-augmented generation.\" EMNLP (2024).\n\n[4] Wu, Wenhao, et al. \"Retrieval head mechanistically explains long-context factuality.\" ICLR (2025).\n\n[5] Jin, Zhuoran, et al. \"Cutting off the head ends the conflict: A mechanism for interpreting and mitigating knowledge conflicts in language models.\" Findings of ACL (2024).\n\n[6] Goldowsky-Dill, Nicholas, et al. \"Localizing model behavior with path patching.\" arXiv preprint arXiv:2304.05969 (2023)."}, {"review_id": "npDWtyielE", "reviewer": "Reviewer_n2tY", "summary": "It proposes an attribution-based method to identify attention heads that are crucial during in-context retrieval augmentation, showing that these heads operate distinctly on the prompt. This method can categorize attention heads into \"in-context\" and \"parametric\" types. \nIt analyzes how in-context heads specialize in interpreting instructions and retrieving information, mapping their locations across different model layers.  Furthermore, it demonstrates how both in-context and parametric heads influence the answer generation process by compacting them into function vectors or modifying their attention weights. \nThe research presents preliminary findings on enabling causal tracing of source information for retrieval-augmented language models, which indicates promising directions for improving the interpretability of Retrieval-Augmented Generation (RAG) systems.", "questions": "1. Methodological Specificity and Generalizability (of the attribution method):\nThe paper proposes an \"attribution-based method\" to identify specialized attention heads. My question would be: What are the precise mathematical definitions and computational procedures of this method? How robust and generalizable is this approach across different Transformer architectures or even different languages beyond English?\n\n2. Boundaries and Interactions of Attention Head Classification:\nThe paper categorizes attention heads into \"in-context\" and \"parametric.\" Are these classifications absolute, or is there a spectrum or potential for \"hybrid\" heads? Furthermore, how do these identified attention heads interact and cooperate with other core components of the model, such as the Feed-Forward Networks (MLPs), which are also known to store knowledge?\n\n3. Universality of Insights and Limitations in Application:\nThe insights presented are primarily derived from in-context learning in question-answering tasks. To what extent do these findings regarding attention head function and influence universally apply to other types of in-context learning tasks (e.g., summarization, translation, code generation), or remain valid in more complex multi-modal and multi-task learning scenarios? What are the potential limitations of these findings in practical applications, such as model fine-tuning or safety auditing?", "limitations": "1. Limited Scope of Component Interaction Analysis: The study primarily focuses on the role of attention heads. It acknowledges that the full picture of how attention heads interact with other crucial components of LLMs, such as MLP (Multi-Layer Perceptron) modules (e.g., knowledge neurons), to induce specific functions remains an open question and is left for future investigation.\n\n2. Potential for Misuse of Insights: The paper highlights a broader ethical concern regarding the potential for the identified specialized heads to be misused, for instance, to induce or amplify malicious model behaviors. This is a general ethical limitation tied to deep interpretability, rather than a methodological flaw.\n\n3. Untested in Diverse In-Context Learning Scenarios: While the paper demonstrates the mechanism in retrieval-augmented question answering, its findings are not explicitly tested across a wider variety of in-context learning tasks (e.g., in-context summarization, translation, or more complex reasoning tasks) where retrieval might also play a role, or in multi-modal contexts.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. Pioneering Attention Head Classification: The paper introduces a novel attribution-based method that effectively identifies and clearly categorizes attention heads into \"in-context\" and \"parametric\" types. This is a crucial breakthrough for understanding how Large Language Models (LLMs) leverage in-context learning for retrieval augmentation.\n2. Deep Insight into Model Mechanisms: The research provides a detailed analysis of how \"in-context\" attention heads specialize in instruction comprehension and information retrieval. It also demonstrates how these heads influence the answer generation process, significantly enhancing our understanding of LLMs' internal workings.\n3. Foundational for Interpretable RAG Systems: The paper lays the groundwork for more transparent and trustworthy Retrieval-Augmented Generation (RAG) systems by offering preliminary methods for tracing source information in retrieval-augmented LLMs. This is a vital step towards building safer and more controllable AI.\n\nWeaknesses\n1. Limited Scope of Investigation: The study primarily focuses on the role of attention heads. It does not extensively explore the synergistic interactions between attention heads and other model components (such as MLP modules or knowledge neurons) in inducing specific functions, leaving this for future research.\n2. Untested in Cross-Modal or Multi-Task Scenarios: While the paper demonstrates the transferability of attention heads across different datasets, its analysis mainly concentrates on unimodal question-answering tasks. It does not investigate whether these mechanisms apply similarly in more complex cross-modal or multi-task learning environments.\n3. Potential for Malicious Exploitation: The paper acknowledges a potential risk that the identified attention heads could be exploited to induce malicious behavior. While not a flaw in the research methodology itself, it highlights an ethical and safety challenge associated with such deep insights into model control.", "was_revised": true, "final_justification": "Accept ---\nThe authors have provided a clear and thorough rebuttal, directly addressing my concerns and those of the other reviewers. They have committed to adding clarifications, improving presentation, expanding the related work section, and including additional results and analyses that will strengthen the paper. Given the novelty of their mechanistic analysis on in-context retrieval augmentation, the solid empirical evidence, and the promising implications for building more transparent and trustworthy RAG systems, I am satisfied with their responses and recommend acceptance.", "rebuttal": "We sincerely thank the reviewer for the constructive and positive feedback and the time spent on our manuscript. We are encouraged that our approach is seen as a crucial breakthrough in understanding how Large Language Models leverage in-context learning specifically for retrieval augmentation. Furthermore, we are grateful that the foundational nature of our preliminary methods for tracing source information within retrieval-augmented LLMs is appreciated as a vital step toward developing more interpretable, transparent, and ultimately safer AI systems. In the following, we provide detailed responses to the reviewers’ concerns.\n\n**[Methodological specificity and generalizability of attribution method, along with its mathematical definitions and computational procedures]**: We thank the reviewer for raising this important question. We would like to clarify that our study is intentionally designed to be **agnostic to the specific choice of feature attribution technique** and, as such, is compatible with a broad class of attribution methods. As outlined in Section 3 (L82–94), we carefully considered numerous existing attribution techniques but found many to be limited by issues such as **scalability or faithfulness** to the model’s internal mechanisms. Since our goal involves tracing and filtering attention heads across large datasets, our chosen method has to be both computationally efficient and rigorously faithful in attributing the model’s behavior. \n\nRather than developing a novel attribution method, we selected AttnLRP [6] (detailed in Section 3.1, L105–113) due to its well-founded mathematical basis and suitability for our setting. AttnLRP is a (modified) gradient-based method **requiring only a single backward pass** to assign relevance values to all heads, which offers significant computational advantages and takes the entire computational path into account. Its core principle involves approximating the contribution of neurons to the final output logit through Jacobian computations, while specifically correcting gradient distortions introduced by non-linearities common in large language models. For instance, layer normalization notoriously disrupts gradient flow [1], but AttnLRP’s formulation carefully linearizes this operation by treating normalized vectors as constants by detaching them from the computational graph, as rigorously proven in their original work. Relevance values can subsequently be obtained directly from the `.grad` attributes of the corresponding tensors in the PyTorch implementation. We acknowledge that readers may benefit from a concise introduction to AttnLRP. To address this, we are considering adding a brief overview section in the Appendix. For a more detailed treatment, we refer readers to the original AttnLRP paper [6].\n\nWe further illustrate the need for efficiency by comparing with other methods such as path patching [2] that requires one clean and one corrupted forward pass per head (plus an additional backward pass).  With 1024 attention heads, this results in 2048 forward passes per sample. On a **small test dataset** of NQ-Swap, where a single run takes approximately 15 minutes on NVIDIA GeForce RTX 4090, this would amount to **0.25h × 2048 ≈ 31.3 days, making such methods infeasible at scale**.\n\n**[Universality of insights and limitations in application, including robustness and generalizability across different Transformer architectures, different languages than English, and other types of ICL tasks]**: Our study generalizes beyond language modeling and is **applicable to all Transformer-based architectures**, including vision models, and is agnostic to the attribution method as well as. To further highlight the generalizability of our approach, we conducted experiments on **Machine Translation on the OPUS Europarl dataset (on French, Spanish, English)**. In these experiments, we identified task-specific heads responsible for translation in a **zero-shot manner**. We then compacted these heads into a feature vector representation and patched them into different prompts (again in a zero-shot fashion) successfully inducing translation behavior into various languages. To evaluate the effectiveness of this approach, we compared it against two baselines: (1) directly prompting the model with “Translate the sentence into [LANGUAGE]:” and (2) randomly patching attention heads. Our method matched the BLEU score of the explicit prompting baseline (implemented with SacreBLEU, 33 points on Es-Fr, 29 points on Fr-En), demonstrating that the patched task heads reliably induce the desired translation behavior. Notably, this was achieved using only 15 of the 1024 attention heads in the LLama 3.1-8B-Instruct model, highlighting the preciseness of our approach. Random patching, on the other hand, results in a BLEU score of close to zero. **We believe these results substantially strengthen our paper and we thank the reviewer** for suggesting to include machine translation experiments.\n\nFurthermore, we would like to highlight that **QA remains one of the most prevalent and impactful applications of LLMs**, which motivated our decision to focus on this task. While prior studies on attention heads have largely been restricted to simple **toy settings** as we mentioned in our related work section, to the best of our knowledge, we are among the first to scale this line of analysis to real world datasets. **This represents a substantial advancement in the field.** While we are currently unable to provide an analysis on multimodal image-text data, we believe our results are likely to generalize to vision-language models (VLMs). This is because VLMs typically retain the core architecture of language models pretrained on text, with the addition of a vision encoder or projector that maps image inputs into soft text tokens. The underlying LLM remains largely unchanged [3]. We plan to add this as a limitation and an exciting direction for future research.\n\n**[Absoluteness of classifications of attention heads]**: Thank you for the insightful question. As shown in Figure 3, we plotted the in-context scores of all attention heads in sorted order. The results reveal a continuous spectrum ranging from in-context (positive score) to parametric behavior (negative score). Notably, many heads fall in the middle of this distribution, with average scores slightly above zero. This suggests that a substantial number of heads are not specialized exclusively for either in-context or parametric roles and their mixed contributions effectively cancel out in the scoring. We agree that further analysis of these intermediate heads is a compelling direction for future work, and we plan to elaborate on this point in our manuscript.\n\n**[Interaction with MLP modules]**: We agree that investigating the interplay between attention heads and MLP layers is a highly valuable direction. We hypothesize that parametric heads ultimately extract attributes from token positions where MLP layers have previously enriched the representations. Analyzing this interaction is considerably more complex due to the polysemantic nature of concepts encoded in MLPs, hence the growing body of work on sparse autoencoders (SAEs) [4].\n \nDue to the limited scope and space constraints of our current manuscript, we did not explore this aspect in depth. However, we have acknowledged this limitation already explicitly in the current version. It is worth noting that many recent studies have focused **exclusively on analyzing attention heads** [5] and likewise **omitted the study of MLPs**. In this context, we believe our in-depth analysis of attention heads already constitutes a substantial and meaningful contribution to the field.\n\n**[Potential for misuse/exploitation]**: As you rightly mentioned, this is a general limitation tied to deep interpretability, rather than a weakness of our method. We already acknowledged this as a potential risk in the main paper.\n\nWe thank the reviewer again for the time spent engaging with our work and we hope our responses have clarified the points raised.\n\nReferences:\n\n[1] Xu, Jingjing, et al. \"Understanding and improving layer normalization.\" Advances in neural information processing systems 32 (2019).\n\n[2] Goldowsky-Dill, Nicholas, et al. \"Localizing model behavior with path patching.\" arXiv preprint arXiv:2304.05969 (2023). \n\n[3] Liu, Haotian, et al. \"Visual instruction tuning.\" Advances in neural information processing systems 36 (2023)\n\n[4] Gao, Leo, et al. \"Scaling and evaluating sparse autoencoders.\" arXiv preprint arXiv:2406.04093 (2024).\n\n[5] Wu, Wenhao, et al. \"Retrieval Head Mechanistically Explains Long-Context Factuality.\" ICLR (2025).\n\n[6] Achtibat, Reduan, et al. \"AttnLRP: Attention-Aware Layer-Wise Relevance Propagation for Transformers.\" International Conference on Machine Learning. PMLR, 2024."}, {"review_id": "OcwVYbsX77", "reviewer": "Reviewer_FEvZ", "summary": "This paper investigates the mechanism behind in-context retrieval augmentation by identifying attention heads specialized for this job in large language models. Using an attribution-based method, AttnLRP, the authors categorize attention heads into two types: in-context heads that process contextual information from prompts, and parametric heads that store relational knowledge acquired during pre-training. The paper shows that in-context heads further specialize into task heads (for instruction following) and retrieval heads (for pure copy-paste from context). Through controlled experiments on biography datasets and QA tasks, the authors demonstrate that these heads can be compressed into function vectors or have their attention weights modified to causally influence answer generation. The work concludes by presenting a linear probe that can distinguish between parametric and contextual predictions with high accuracy and claim to enable more transparent analysis on context based generation.", "questions": "Q1: Please correct me if my understanding is wrong. I’m wondering if it is surprising that attention tracing for the answer token does localize the token which probably is similar in the context in Table 2. Am I mis-understanding this experiment?", "limitations": "Yes.", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 1, "originality": 2, "strengths_and_weaknesses": "S1: This paper addresses an interesting topic of interpreting what attention heads of a transformer perform in different roles. The experimental design reveals the existence of attention heads delivering parametric and in-context knowledge, providing valuable insights into the mechanistic understanding of retrieval-augmented generation.\n\nS2: The paper presents a relatively thorough investigation with many well-designed control experiments, including ablation studies, function vector extraction, and cross-dataset generalization tests that strengthen the empirical findings.\n\nS3: The attribution-based methodology using AttnLRP provides a principled approach to head identification that goes beyond simple attention weight analysis.\n\n---\n\nW1: The experiments are sound, but it seems like there aren't many actionable items that practitioners could directly use. This would be enhanced if the whole method could be efficiently packaged to apply directly to any transformer-based LM, or if performance gains on some realistic benchmarks could be achieved through the insights gained.\n\nW2: It seems like most experiments still rely on single-token answers, and it is unclear how these findings could help with any interpretability of answers with long chain-of-thought reasoning, which represents the vast majority of LLM use cases in practice.\n\nW3: I think interpretability studies can benefit from showing real-life impact, as pointed out in recent discussions (https://icml.cc/virtual/2025/workshop/39962, this is just showing that this discussion is significant enough to grow into a workshop). However, I cannot directly see a direct impact of this study beyond advancing our theoretical understanding of attention mechanisms. To this end, can this interpretability method be at least compared to simple prompting: Literally asking the model “Why did you answer ###? Is there a reason from the context?”\n\nW4: It seems like some important citations on mechanistic interpretability are missing: https://arxiv.org/pdf/2310.10348, https://arxiv.org/pdf/2407.00886\n\nW5: Correct me if I am wrong, the method seems like a direct application of AttnLRP, or are there improvements delivered?", "was_revised": true, "final_justification": "I confirm my final score of 3.\n\nWhile this paper shows some promising results, the practicality of the results compared to simple baselines such as prompting doesn't seem to be clearly demonstrated.\n\nI believe that the assigned score is adequate.", "rebuttal": "We sincerely thank the reviewer for their encouraging feedback and we are especially grateful for the recognition of our contributions to the mechanistic understanding of retrieval-augmented generation (S1), as well as the acknowledgment of our experimental rigor through comprehensive and well-designed experiments (S2). We also appreciate the reviewer’s appreciation of our attribution-based methodology as a principled and meaningful advancement over prior studies (S3). In the following, we respectfully address the reviewer’s concerns and suggestions in the responses below. We hope our clarifications and revisions further strengthen the contributions and clarity of the work.\n\n**[Sound, but no actionable items (W1) and benefit for real-life impact (W3)]**: We thank the reviewer again for their thoughtful feedback and for acknowledging the soundness of our experimental design. We fully agree that interpretability research should strive to connect with real-world applications that benefit practitioners. Toward this goal, we have aimed to present not only a theoretical advancement in understanding attention head specialization in LLMs, but also practical implications that enhance model transparency and safety.\n\nSpecifically, we **added additional experiments on detecting knowledge conflicts** for Section 7: scenarios in which LLMs must choose between conflicting contextual (external) information and their internal parametric memory [1]. As shown in the paragraph below, our method allows us to guide the model to prefer either contextual knowledge or its memorized answer, depending on the desired use case. As such, a practitioner could provide a high quality knowledge base for a RAG system e.g. legal texts and ground the model on this external information by preventing the LLM from using outdated internal knowledge. This demonstrates how our method can be used to detect and mitigate knowledge conflicts, an important step toward improving LLM reliability.\n\nIn addition, we respectfully draw attention to Section 7, where we leverage our findings to build faithful post-hoc citation mechanisms. We introduce a novel attribution-based procedure for tracing both contextual and parametric knowledge usage within LLMs, enabling the creation of **faithful citations** and hence introducing a control mechanism **against hallucinations**, which is a major problem, particularly given its growing importance in the NLP communities [2, 3].\n\n**[Knowledge conflict results]**: We computed the accuracy of LLama 3.1-8B-Instruct on the NQ-Swap with counterfactual contexts. Computing the recall score on the gold answer instead, the model is clearly distracted by the wrong contexts achieving only 9% accuracy. Removing the top 60 in-context heads reduces the reliance on the wrong context pushing the score to 35% which is similar to the closed-book accuracy of the model. Conversely, computing the recall score for counterfactual answers i.e. reliance on the counterfactual inputs, the model achieves 68% accuracy without ablation. Ablating the top 100 parametric heads reduces the reliance on internal knowledge and the model is more likely to accept the counterfactual context increasing the score to 76%.\n\n**[Direct applications to Transformer LMs (W1)]**: Yes, our method is model-agnostic and can be readily applied to any pretrained LLM without requiring fine-tuning or architectural modifications. Furthermore, it is straightforward to package this approach into a reusable library. We will open source our implementation to facilitate this.\n\n**[Comparability to simple prompting (W3)]**: As noted in L60–61, RAG LLMs cannot guarantee faithful attribution to contextual passages [2], highlighting the need for alternative methods such as ours for post-hoc knowledge tracing. This perspective also aligns with recent works, e.g. by Qi et al. [3], which emphasizes the value of actionable interpretability as it provides actionable insights for both users and developers of LLM systems.\n\nRegarding prompting-based attribution, we also note that LLMs **do not reliably** reflect the true origin of their outputs. For instance, they may exhibit **confirmation bias** when asked leading questions such as *“Is there a reason in the context...”*. Even when the external context conflicts with internal memory, models often confirm answers aligned with parametric knowledge, as shown in [4]. [5] also demonstrates that LLMs strongly prefer generated contexts regardless of correctness, which can still suffer from hallucinations . These findings make the reliability of prompting heuristics questionable and support our argument that interpretability methods are essential for faithful knowledge attribution and conflict resolution.\n\n\n**[Single-token answers (W2)]**: We use multi-token answers for the vast majority of our experiments and highlight how this differs from previous works in the related work section. We will make this clearer in the revised version of our manuscript.\n\n**[Missing citations (W4)]**: Thank you for the references. We will add this to the revised version of our work.\n\n**[Direct applications of AttnLRP (W5)]**: While we opt to use AttnLRP due to our need for a method that is both computationally efficient and faithful in attributing model behavior (L82-94, L105-L103), our framework is general and can be applied to other feature attribution methods.\n\n**[Is it surprising that attention tracing for answer tokens does well for localization? (Q1)]**: Thank you for the thoughtful question. At a first glance, it may indeed seem normal that an answer token attends to contextually similar tokens. However, our work goes beyond this intuition in several important ways.\n\nFirst, we do not simply visualize attention weights. We identify **which specific** attention heads are responsible for tracing the origin of an answer, namely the retrieval heads. Using attribution analysis, we show that only a **small subset** of heads (4% of total heads) meaningfully contribute to source tracing, while others are largely inactive or fulfill orthogonal roles. This is a nontrivial finding: constructing accurate relevance heatmaps requires knowing **which** attention heads to include and how to **weight** their contributions correctly, as described in Section 7 and Appendix D. Without this guidance via attribution, one would **not** obtain the sharply focused heatmaps we present, but instead **noise**. As discussed in Section 3, this selective identification is not achievable through raw attention maps alone. Attribution methods are therefore essential: they reduce the search space and expose the functional roles of heads, rather than assuming all attention is equally explanatory.\n\nSecond, we recognize the broader interpretability debate, particularly the concern that “Attention is not explanation” [6]. Our approach addresses this challenge directly: rather than assuming attention weights are inherently meaningful, we **validate empirically** that certain attention heads produce faithful and reproducible attribution signals, while others do not. \n\nCrucially, our method requires no fine-tuning, supervision, or model modifications, distinguishing it from prior approaches to source attribution that rely on explicitly training models for this purpose [7]. Our **zero-shot, post-hoc** framework is therefore broadly applicable and practical for analyzing existing models without additional overhead for practitioners.\n\nWe thank the reviewer again sincerely for the time spent engaging with our work and we hope our responses have clarified the concerns raised.\n\nReferences:\n\n[1] Xu, Rongwu, et al. \"Knowledge conflicts for LLMs: A survey.\" arXiv preprint arXiv:2403.08319 (2024).\n\n[2] Gao, Tianyu, et al. \"Enabling large language models to generate text with citations.\" EMNLP (2023).\n\n[3] Qi, Jirui, et al. \"Model internals-based answer attribution for trustworthy retrieval-augmented generation.\" EMNLP (2024).\n\n[4] Xie, Jian, et al. \"Adaptive chameleon or stubborn sloth: Revealing the behavior of large language models in knowledge conflicts.\" The Twelfth International Conference on Learning Representations. 2023.\n\n[5] Tan, Hexiang, et al. \"Blinded by generated contexts: How language models merge generated and retrieved contexts when knowledge conflicts?.\" arXiv preprint arXiv:2401.11911 (2024).\n\n[6] Bibal, Adrien, et al. \"Is attention explanation? an introduction to the debate.\" Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2022.\n\n[7] Sun, Zhiqing, et al. \"Recitation-augmented language models.\" arXiv preprint arXiv:2210.01296 (2022)."}, {"review_id": "E55L8zQCXb", "reviewer": "Reviewer_ZnzS", "summary": "This paper provides a mechanistic analysis of how large language models retrieve information in context. The authors employ an attribution-based methodology to identify distinct sets of attention heads that serve different functions: parametric heads that extract knowledge from model parameters, task heads that process task-specific information, and retrieval heads that extract information from the given context. They validate the functionality of these head categories through ablation studies that demonstrate their causal roles in the model's behavior.", "questions": "**Q1.** Regarding the statement \"In addition, we use counterfactual contexts as they contain information that is usually not seen during pre-training and fine-tuning stages, thus forcing models to rely on the text to answer the question correctly\" (p.4, l.149-150), could you clarify where the counterfactual contexts are sourced from?\n\n**Q2.** You mention that \"this leads to a slight performance increase on NQ-Swap with counterfactual contexts, which suggest that the ablation forces the model to rely more on the given context instead of its own parametric knowledge\" (p.6, l.209-210). Depending on your counterfactual data selection method, do you verify that the context and parametric answers differ? This result suggests that removing parametric heads eliminates knowledge conflict and improves performance - it would be valuable to discuss this knowledge conflict aspect more explicitly.\n\n**Q3.** What happens qualitatively when you ablate in-context heads versus parametric heads? When ablating in-context heads, does the model default to parametric knowledge (especially during knowledge conflicts)? What occurs when ablating parametric heads in the CB example? What happens if you ablate both types?\n\n**Q4.** In Figure 1, what exactly do the bar plots represent? Are they showing average task/retrieval relevance for all heads in that layer, or the maximum? Is the marker size representing each head's relevance, thresholded by some minimum value?\n\n**Q5.** Why would task heads be marked as context heads or have particularly high context scores? According to Equation 5, these heads should be relevant for both CB and OB, so they should have values around zero. Do you have intuitions for explaining this pattern?\n\n**Q6.** Since $H_{task}$ and $H_{ret}$ are not mutually exclusive, did you examine whether some heads have high values in both categories?\n\n**Q7.** For the experiments in Section 6.2:\n- Are you using the chat template for the instruct models? Figure 4 appears to use a Q: ... A: ... format, which seems uncommon and potentially out-of-distribution for chat models.\n- Could you clarify the experimental setup for modifying attention weights? The \"Attn Weight\" label is only relevant for $\\cal{H}_{ret}$ right?\n- Did you evaluate model performance without any patching? I expect it to be similar to random, but might still be interesting.\n\n**Q8.** You state that \"Oracle context is often not available; therefore we also use Dense Passage Retriever (DPR)\" (p.4, l.151), and Figure 2 only reports Oracle results. Does Figure 2 only show examples where true Oracle context exists, or does it refer to DPR? The appendix displays DPR separately, which creates some confusion.\n\n**Q9.** While the results for $H_{param}$ are good they are clearly worse than the other two sets. Can you provide any intuitions of why you see poorer performance here? My hypothesis would be that different heads deal with different domains of parametric knowledge, whereas this may be less the case for contextual knowledge. This would mean that, in order to cover many domains, you would just need more parametric heads.\n\n**Expectation Management:** I think this is a great paper with just some minimal issues. If you can resolve W1/Q1 & W2 (without reproducing the other paper) as well as clarify the confusion around W7 & W8, I'll happily raise the score to 5. If you manage to reproduce/explicitely quantitatively compare to https://aclanthology.org/2024.findings-acl.70.pdf, and include discussions about Q2/W13 as well as all minor issues, I will consider a score of 6. Even without addressing Q2/W13, I'd strongly recommend implementing the presentation improvements W4, W5, W6, W9, W10, and W11 for CR as they would enhance the paper's clarity.", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n\n**S1.** Great idea with the synthetic document dataset approach.\n\n**S2.** Good visual presentation with effective use of color and high-quality figures that effectively communicate the results (aside from font size issues).\n\n**S3.** The results in Section 6.2 are impressive and demonstrate strong empirical validation.\n\n**S4.** The paper addresses most questions that arise during reading, providing comprehensive coverage.\n\n## Weaknesses\n\n### Major \n\n**W1.** The method for generating counterfactual data needs clearer explanation.\n\n**W2.** The relationship to the prior work https://aclanthology.org/2024.findings-acl.70.pdf requires better contextualization. While you cite this work, you dismiss it as only looking at single word settings. However, the parallels are substantial - they also compute two sets of heads. Your statement that \"Our method identifies two groups of heads based on their functions: parametric heads that encode relational knowledge and in-context heads responsible for processing information in the prompt\" (p.1, l.33-34) needs better positioning relative to existing work. Consider comparing methods directly. Demonstrating that you find similar heads or that your head sets perform better would be optimal (but not necessarily required, see *Expectation Management* under *Questions*).\n\n### Minor\n\n**W3.** While you discard inhibitory effects, there likely exists interplay where parametric heads inhibit context answers and vice versa. Did you explore this? This might be worth putting into the limitations.\n\n**W4.** Figure fonts are generally too small across multiple figures (Figures 1, 2, axis labels). Consider increasing font sizes where space permits.\n\n**W5.** Figure 1 should be exported as a vectorized image for better quality.\n\n**W6.** Minor notation issue: in my opinion \"for a sequence S\" (p.4, l.127) should be \"sequence of length S\".\n\n**W7.** The statement \"retrieval heads and task heads consistently rank among the highest scoring in-context heads, emphasizing the critical role of in-context heads for retrieval augmented generation\" (p.7, l.262-266) appears overstated. Figure 2 shows that ablating only 20 heads deteriorates performance, while most retrieval heads rank between 100 and 400.\n\n**W8.** The hyperparameter K mentioned as \"will be determined in the next section\" (p.7, l.243) is never formally determined - you simply state that you use top 40 heads (l.260). Explicit definition and justification for K selection would improve clarity. The appendix brings clarity, but then you must refer to that. The same holds for Table 1, where it is completely unclear how many heads you compare. From reading the main text (without looking at the appendix) the conclusion would be that you compare 100 $H_{param}$ heads to 40 of the other heads. From reading the appendix, I assume you're using 40 $H_{param}$ heads as 50 heads already performs better for Llama. You should be more specific about this.\n\n**W9.** Section 6.2 should mention the existence of the FV scaling factor $\\alpha$ in the main text (footnote).\n\n**W10.** Experiment 7 is difficult to understand without consulting Appendix D. Consider expanding the explanation, focusing on your proposed metrics rather than the logit lens function. I claim that most interested readers will be familiar with logit lens. If you get another page for camera ready, I'd recommend spending extra lines explaining this experiment. Further, Figure 5b doesn't visualize the weights, which are important here.\n\n**W11.** Table 2 presentation could be improved. In my understanding, ROC AUC measures parametric vs contextual knowledge, while other metrics measure localization performance. Consider separating these with vertical lines or hierarchical headers for clarity. \nGiven the density of this section, placing these metrics adjacent to each other initially creates confusion as readers may attempt to compare them inappropriately.\n\n**W13.** You are missing an opportunity to report additional metrics on resolving knowledge conflicts in the counterfactual setting. While you demonstrate that ablating context heads destroys recall, you don't discuss what happens to the parametric answer. If you could show that your metric enables outperforming other knowledge conflict interpretation papers (those cited on l.59), this would significantly strengthen your contribution.", "was_revised": true, "final_justification": "The authors addressed all my concerns and questions to my satisfaction, which is why I increased my score as mentioned under 'Expectation Management'. I refrain from giving a score of 6 because, while I believe this paper is very solid and has a significant impact, I don't think it warrants the term 'groundbreaking impact'. Secondly, the authors were unable to provide a direct comparison to the work of Jin et al. [1]. Nevertheless, I recommend that this paper is accepted, as I believe it would make a valuable contribution to NeurIPS.\n\n\n[1] Jin, Zhuoran, et al. \"Cutting off the head ends the conflict: A mechanism for interpreting and mitigating knowledge conflicts in language models.\" Findings of ACL (2024).", "rebuttal": "We are truly grateful to the reviewer for their exceptional feedback. This is genuinely the most insightful review we have ever received in any conference setting we participated in so far. Your comments have substantially improved our presentation, and we deeply appreciate the care you invested in engaging with our work.\n\n**[Counterfactual data W1/Q1]**\nThank you for the opportunity to clarify. We use NQ-Swap, a variant of Natural Questions based on the MRQA 2019 task. The counterfactual data is generated through corpus substitution, where QA instances with named entity answers are identified. We then replace all occurrences of the named entities in the context by other entities sampled from the NQ dataset which have the same type as the original. This is intentional to elicit knowledge conflicts and can be used to evaluate whether LLMs rely more on provided context or on memorized knowledge learned during pre-training (since the substitute entities usually do not appear in the same context during pre-training). We will revise the manuscript to explain this more clearly.\n\n**[Relationship to prior work W2]** We believe the work of Jin et al. [3] is a well-executed contribution to the field. We appreciate their effort in identifying parametric and in-context heads and in exploring head-level ablations to mitigate knowledge conflicts. However, we note that their study is only tangentially related to ours: specifically to our initial experiment in Section 5, where we identify in-context and parametric heads. In our case, this experiment serves primarily as a first step for the more substantive contributions in the subsequent sections. There are several key differences:\n\n- First, Jin et al. focus on narrow-domain tasks with single-token answers, which diverges from realistic QA use cases. In contrast, we conduct our analysis on open-domain QA datasets (NQ-Swap, TriviaQA), as well as a curated human biography dataset that better captures the complexity of real-world settings.\n- While we would have welcomed a direct comparison with Jin et al., this is unfortunately not possible. Their GitHub repository is empty, and no code is available. More importantly, their method—based on path patching [5]—is computationally impractical for QA tasks, requiring 2048 forward passes per sample (1024 heads × 2 passes), which would take 21 days for a single test set run over NQ-Swap on an RTX 4090.\nAs discussed in Section 3 (L82–94), we surveyed various attribution methods and found limitations in scalability or faithfulness. We selected AttnLRP [1] for its competitive attribution quality and significantly lower computational cost. Although Jin et al. briefly explored gradient-based methods (§4.3), they dismissed them due to lack of faithfulness. This criticism applies to standard gradient-based approaches, which are indeed known to suffer from high noise, a limitation well-documented in the literature[1,2]. However, AttnLRP improves on gradient-based methods, hence we selected it.\n- For the first time, we show that in-context heads are considerably more complex and can specialize to task or retrieval heads which provides a finer-grained categorization than is available in literature. We also empirically demonstrate their causal roles in shaping the model’s representation (via intensional frames) along with parametric heads.\nOur method allows us to extract steerable function vectors in a zero-shot manner for complex tasks, as described in (Section 6.2). \n- We also introduce a novel post-hoc procedure to trace both contextual and parametric knowledge usage in LLMs, thus enabling the identification of potential hallucinations by providing citations (Section 7), especially since this line of work has become increasingly important [4].\n\n**[Inhibitory effects W3]** We indeed identify inhibitory heads that appear to suppress the influence of counterfactual inputs, reinforcing the model’s parametric answer. For the rebuttal, we extended our analysis to task heads and found similar inhibitory behavior. Notably, patching such heads fails to induce the intensional frame, likely because they destabilize it. Boosting their activation yielded inconsistent effects: sometimes suppressing the frame, other times increasing verbosity. Overall, the patterns are difficult to interpret without a more systematic analysis and we leave it for future work. We will also put this in the limitation section.\n\n**[Fig. W4+W5, notation issue W6, scaling factor W9, section 7 W10]** We sincerely thank you for your detailed feedback and will gladly incorporate these improvements into the manuscript.\n\n**[Ranking of heads W7]** What we intended to convey is the reverse relationship: as shown in Fig. 3, the highest-scoring in-context heads are primarily composed of retrieval and task heads. We did not mean to imply that all retrieval heads rank among the top in-context heads. We will revise this to more accurately reflect the intended meaning.\n\n**[K-Hyperparams W8, Table 2 W11]**\nTo save space, we moved non-essential details to the App. and briefly noted this in L279. The table is also compact for space reasons. We will expand on this in the camera-ready version and we appreciate the helpful suggestion.\n\n**[Knowledge Conflict Q2 & W13]**\nThank you for this fruitful suggestion. We computed the accuracy of LLama 3.1 on the NQ-Swap with counterfactual contexts. Computing the recall score on the original gold answer instead, the model is clearly distracted by the wrong contexts achieving only 9% accuracy. Removing the top 60 in-context heads reduces the reliance on the wrong context pushing the score to 35% which is similar to the CB accuracy of the model. Conversely, computing the recall score for counterfactual answers instead, the model achieves 68% accuracy without ablation. Ablating the top 100 parametric heads reduces the reliance on internal knowledge and the model is more likely to accept the counterfactual context increasing the score to 76%. In cases where the model still fails, it often asserts that the answer is not present in the context indicating a complete failure to locate it, unrelated to parametric memory.\n\n**[Qualitative analysis for ablation Q3]** \n- The model tends to revert to parametric knowledge when in-context heads are ablated (/w counterfactual contexts), while sometimes returning incorrect but semantically related answers or “The text does not mention this information.”\n- When parametric heads are ablated in CB, the model tends to return false but semantically plausible answers or mention “I do not have information on …”.\n- When both heads are ablated, models seem unable to read from the contexts (/w counterfactual) and hallucinate semantically plausible answers (counterfactual and CB). \nWe will add some qualitative examples.\n\n**[Bar Plots & Marker Size Q4]**\nWe agree this aspect could be clearer and we will revise it accordingly. The bar plot (Section 6.1 L246–248) shows layer-wise relevance distributions by summing, per layer, the relevance each head assigns to the question or answer object tokens. This reveals whether a layer is more task- or retrieval-focused. Top 40 task and retrieval heads are highlighted with larger markers, while marker color indicates in-context vs. parametric behavior (based on the top 100 heads in each category).\n\n**[Task heads in OB and CB Q5]**\nAs shown in Fig. 3 & 6, while some task heads show strong in-context behavior, many fall in the moderate range. Notably, Gemma 2 even exhibits parametric task heads, indicating that task heads are not exclusively in-context. We will clarify this in the revised version. As for why task heads receive more relevance in OB settings: one explanation is that they must handle more contextualization, thereby attracting more relevance. In contrast, in CB settings, parametric heads dominate as the input is more self-contained and absorb more relevance.\n\n**[Exclusivity of heads Q6]**\nTask and retrieval heads are largely mutually exclusive. For example, in LLaMA, the intersection between the top 40 task heads and the top 40 retrieval heads yields only a single shared head. However, we believe there may be a small misunderstanding: in App. (L803–811), we state that in-context heads and task heads tend to show considerable overlap (not retrieval heads). We will clarify this more explicitly.\n\n**[Chat template, attention modification, performance w/o patching Q7]**\nWe did not use chat-specific templates, but followed [6] to exclude possible system prompt interference.\nExactly, only for the evaluation of retrieval heads did we modify the attention weights. This is explained in App. L856-867.\nWe evaluated the performance without any patching in our preliminary experiments and it indeed yielded scores slightly below random. \n\n**[Fig. 2 Q8]** Indeed, it refers to true oracle and not DPR. We will add this to the main text.\n\n**[Performance of H_param Q9]** Your hypothesis appears correct. We patched parametric FVs from cloze-style statements into related question prompts (e.g., “Mike Tyson was born on” → “Q: What is the birth date of Albert Einstein?”). On Llama, we see an average increase of ~7% with 100 heads, which confirms that more parametric heads are needed to cover many domains.\n\nRef:\n\n[1]  Achtibat, Reduan, et al. \"AttnLRP: Attention-Aware Layer-Wise Relevance Propagation for Transformers.\" ICML, 2024.\n\n[2] Smilkov, Daniel, et al. \"Smoothgrad: removing noise by adding noise.\" arXiv preprint  (2017).\n\n[3] Jin, Zhuoran, et al. \"Cutting off the head ends the conflict: A mechanism for interpreting and mitigating knowledge conflicts in language models.\" ACL Findings (2024).\n\n[4] Qi, Jirui, et al. \"Model internals-based answer attribution for trustworthy retrieval-augmented generation.\" EMNLP (2024).\n\n[5] Goldowsky-Dill, Nicholas, et al. \"Localizing model behavior with path patching.\" arXiv preprint (2023). \n\n[6] Ram, Ori, et al. \"In-context retrieval-augmented language models.\" ACL (2023)"}]

Author rebuttal:

Meta-review: The paper mechanistically analyzes retrieval augmentation -- how different heads perform different functions such as retrieving from the context vs. from the memory. 

While the reviewers were on the borderline (two borderline rejects), the paper garnered one strong advocate (`ZnzS `). Based on my further impression of the paper, I recommend acceptance.

More details below:
- `J8iq` raised concerns about novelty -- that this paper combines many ideas, but lacks a significantly novel contribution. `ZnzS ` says that although each individual idea exists prior to this paper, the application is novel and brings out a new finding. I agree with this characterization. 
- `J8iq` also brought up the fact that the study focused on two QA questions, and `FEvZ` says these are only one-token answer questions. `ZnzS` agrees with this concern especially in light of how the authors distinguish themselves from Jin et al., **I agree with the reviewers that this difference against Jin et al., cannot be claimed. Please edit your manuscript accordingly.** 
- Both `ZnzS` and `J8iq ` brought up lack of contextualization against other works, including mechanistic interpretability literature, and a specific paper that is closely related but the authors dismissed it too casually (there was back and forth on this with `FEvz`, `ZnzS`). From my discussion, this did not seem to be debilitating in terms of the contributions of the paper. **Regardless this is very important: I strongly encourage the authors to do  do a thorough literature search, and rewrite how they contextualize these contributions against these other lines of related work. Please also cite and compare with Jin et al., prominently and accurately.**
- `FEvZ` has a main concern about the practicality of the result. I consider this as a subjective concern; there is possible future value in understanding the circuitry behind these behaviors. However, it may not hurt to dwell more on this and provide some suggestions at the end of the paper.
Decision: Accept (poster)