The Atlas of In-Context Learning: How Attention Heads Shape In-Context Retrieval Augmentation

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.

[Counterfactual data W1/Q1] Thank 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.

[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:

• 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.
• 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. As 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.
• 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. Our method allows us to extract steerable function vectors in a zero-shot manner for complex tasks, as described in (Section 6.2).
• 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].

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

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

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

[K-Hyperparams W8, Table 2 W11] To 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.

[Knowledge Conflict Q2 & W13] Thank 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.

[Qualitative analysis for ablation Q3]

• 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.”
• 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 …”.
• When both heads are ablated, models seem unable to read from the contexts (/w counterfactual) and hallucinate semantically plausible answers (counterfactual and CB). We will add some qualitative examples.

[Bar Plots & Marker Size Q4] We 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).

[Task heads in OB and CB Q5] As 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.

[Exclusivity of heads Q6] Task 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.

[Chat template, attention modification, performance w/o patching Q7] We did not use chat-specific templates, but followed [6] to exclude possible system prompt interference. Exactly, only for the evaluation of retrieval heads did we modify the attention weights. This is explained in App. L856-867. We evaluated the performance without any patching in our preliminary experiments and it indeed yielded scores slightly below random.

[Fig. 2 Q8] Indeed, it refers to true oracle and not DPR. We will add this to the main text.

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

Ref:

[1] Achtibat, Reduan, et al. "AttnLRP: Attention-Aware Layer-Wise Relevance Propagation for Transformers." ICML, 2024.

[2] Smilkov, Daniel, et al. "Smoothgrad: removing noise by adding noise." arXiv preprint (2017).

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

[4] Qi, Jirui, et al. "Model internals-based answer attribution for trustworthy retrieval-augmented generation." EMNLP (2024).

[5] Goldowsky-Dill, Nicholas, et al. "Localizing model behavior with path patching." arXiv preprint (2023).

[6] Ram, Ori, et al. "In-context retrieval-augmented language models." ACL (2023)

Thanks — this addresses all my concerns. I believe this is now a solid paper that should be accepted. I will update my score to 5 and defend this stance during the reviewer discussion. I've also increased the clarity score, considering you correctly implement all the discussed improvements. I won't consider a score of 6 because of the missing direct comparison, even though I agree that missing code is a problem. I also believe that the paper's impact doesn't classify as 'groundbreaking' (a score of 6 is described as 'Technically flawless paper with groundbreaking impact on one or more areas of AI [...]').

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.

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

Specifically, 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.

In 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].

[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%.

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

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

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

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

[Missing citations (W4)]: Thank you for the references. We will add this to the revised version of our work.

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

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

First, 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.

Second, 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.

Crucially, 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.

We thank the reviewer again sincerely for the time spent engaging with our work and we hope our responses have clarified the concerns raised.

References:

[1] Xu, Rongwu, et al. "Knowledge conflicts for LLMs: A survey." arXiv preprint arXiv:2403.08319 (2024).

[2] Gao, Tianyu, et al. "Enabling large language models to generate text with citations." EMNLP (2023).

[3] Qi, Jirui, et al. "Model internals-based answer attribution for trustworthy retrieval-augmented generation." EMNLP (2024).

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

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

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

[7] Sun, Zhiqing, et al. "Recitation-augmented language models." arXiv preprint arXiv:2210.01296 (2022).

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.

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

Rather 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].

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

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

Furthermore, 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.

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

[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].

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

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

We thank the reviewer again for the time spent engaging with our work and we hope our responses have clarified the points raised.

References:

[1] Xu, Jingjing, et al. "Understanding and improving layer normalization." Advances in neural information processing systems 32 (2019).

[2] Goldowsky-Dill, Nicholas, et al. "Localizing model behavior with path patching." arXiv preprint arXiv:2304.05969 (2023).

[3] Liu, Haotian, et al. "Visual instruction tuning." Advances in neural information processing systems 36 (2023)

[4] Gao, Leo, et al. "Scaling and evaluating sparse autoencoders." arXiv preprint arXiv:2406.04093 (2024).

[5] Wu, Wenhao, et al. "Retrieval Head Mechanistically Explains Long-Context Factuality." ICLR (2025).

[6] Achtibat, Reduan, et al. "AttnLRP: Attention-Aware Layer-Wise Relevance Propagation for Transformers." International Conference on Machine Learning. PMLR, 2024.

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.

[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:

• 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.
• 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.
• 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.
• 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].

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

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

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

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

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

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

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

[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%.

We sincerely thank the reviewer for the time spent engaging with our work and hope we could clarify all points raised.

References:

[1] Xu, Rongwu, et al. "Knowledge conflicts for LLMs: A survey." arXiv preprint arXiv:2403.08319 (2024).

[2] Gao, Tianyu, et al. "Enabling large language models to generate text with citations." EMNLP (2023).

[3] Qi, Jirui, et al. "Model internals-based answer attribution for trustworthy retrieval-augmented generation." EMNLP (2024).

[4] Wu, Wenhao, et al. "Retrieval head mechanistically explains long-context factuality." ICLR (2025).

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

[6] Goldowsky-Dill, Nicholas, et al. "Localizing model behavior with path patching." arXiv preprint arXiv:2304.05969 (2023).