Kati8/ICLR_2023_2025 huggingface dataset:
authors: Wenhao Wu
authors: Yizhong Wang
authors: Guangxuan Xiao
authors: Hao Peng
authors: Yao Fu
reviews: 
Comments: The paper investigates how transformer-based models extract relevant information from long context. It identifies a specific type of attention head, named retrieval head, which plays a significant role in the information retrieval process. The authors demonstrate how they detect retrieval heads and describe their characteristics through various experimental settings. Additionally, they conducted experiments that involved pruning the retrieval heads to show that these heads are essential for recalling specific information amidst vast amounts of data. Their findings indicate that retrieval heads are crucial for extractive question-answering and chain-of-thought reasoning.
Strengths: - The paper shows originality by exploring the retrieval capabilities of transformer language models, a topic that has not been extensively studied.
- The paper is well-organized and easy to follow. The paper first defines the special type of attention head that plays a significant role in recalling relevant information during generation. Then, it successfully proves the existence of the retrieval heads within Needle-in-a-Haystack (NIAH) task and further illustrates several properties of the retrival heads which are quite interesting.
- Figures are clear and intuitive to comprehend  
- The paper is impactful as it proposes prospective research directions involving retrieval heads.
Weaknesses: - The work is limited to the Needle-in-a-Haystack (NIAH) task. Although NIAH is a good task to prove the existence of the retrieval heads, we do not know if the similar findings and significance would transfer to other tasks where the LM needs to paraphrase or utilize the previous context (not just copy-and-paste), which are more complex and closely related to real-world applications.
Recommendation: 8
Reviewer confidence: 4
Soundness/correctness: 4
Clarity: 4
reviews: 
Comments: This paper identifies and analyzes the properties of retrieval heads, a specialized type of attention head primarily responsible for retrieving information. Key properties include: 

1. **Universality** – retrieval heads are present across all explored models.
2. **Sparsity** – only a small subset of attention heads serve this retrieval function.
3. **Intrinsic nature** – these heads exist in pretrained models, even those trained on short contexts.
4. **Dynamic activation** – their activation varies depending on specific tokens and contexts.
5. **Causality** – pruning these heads leads to significant performance drops.

By examining the influence of retrieval heads across various tasks, these findings shed light on which internal model components actively seek information from input tokens.
Strengths: 1. Provides in-depth analysis and extensive experiments on various properties of retrieval heads.

2. Well-written, with clear graphics, easy-to-follow explanations, and a well-organized structure.

3. Numerous examples and case studies effectively illustrate the properties, making the concepts easy to understand.
Weaknesses: Nothing major needs to be addressed. Please address some discussion questions in question sections.
Recommendation: 8
Reviewer confidence: 2
Soundness/correctness: 3
Clarity: 3
reviews: 
Comments: The paper provides a systematic examination of a specific type of attention head, termed "retrieval heads," which primarily handle information retrieval from input data. It introduces an approach based on the Needle-In-a-Haystack (NIAH) setup to empirically identify retrieval attention heads across various transformer-based architectures. The findings demonstrate that: 
1) retrieval heads are present across a diverse set of models,
2) only a small subset of attention heads function as retrieval heads,
3) these retrieval heads exist even in models pretrained with limited context, suggesting they are an intrinsic artifact of pretraining,
4) they are dynamically activated rather than continuously active, and
5) there is a causal link between retrieval heads and the model’s capability to retrieve relevant information.
Strengths: Overall, the paper is well-structured, the assumptions are clear, and, with a few exceptions listed in the "weaknesses", the methodology is clear. The results and experiments robustly support the authors’ claims.
Weaknesses: The main concern with the paper is the lack of details on the benchmark used to generate the needle-and-haystack pairs. The paper does not clarify how these pairs are created, the diversity of pairs (e.g., in topic, token variety), or the validation methods used. Including these details would provide reviewers with valuable insights into the experimental design, helping them better assess the generalizability and universality of the findings.


Additional minor points:
1) The acronym "KV" is not defined anywhere in the paper. Based on the context, it likely stands for "Key-Value," but this should be explicitly stated.
2) The caption for Figure 3 could benefit from significant revision. It’s challenging to interpret without detailed reference to the discussion, so adding clarifying information in the caption itself would help.
3) Figure 5 is difficult to interpret; it’s unclear what is being visualized. Overlaying the heatmaps for comparison could enhance clarity, or, if this visualization is redundant given the results in Figure 6, consider omitting it.
4) The discussion in Section 4, in particular, would benefit from qualitative examples (perhaps as an appendix) to illustrate and substantiate claims related to downstream tasks.
Recommendation: 8
Reviewer confidence: 3
Soundness/correctness: 4
Clarity: 4
reviews: 
Comments: This paper investigates the mechanism with which transformer-based language models "retrieve" information in the long context. It experimented with four model families, six model scales, and three types of post-training variants, and reveals that a special type of attention heads are largely responsible for retrieving information (either copy-paste or paraphrase) from long contexts. Such attention heads are named “retrieval heads”. The authors find that these retrieval heads
- (1) exist in all the explored models,
- (2) are only 5% of the attention heads,
- (3) exist in models large-scale-pretrained with long or only short contexts and remain the same when the models are continually pretrained on longer contexts,
- (4) are dynamically activated given different contexts, and
- (5) will cause degradation in retrieval abilities or chain-of-thought abilities if pruned.
Strengths: - The paper is well-written, with a clear overarching research question (”How do transformer-based language models acquire long-context capabilities?”), the substantial findings of the existence of “retrieval heads” and their properties, and experiments to support each finding.
- The experiments are extensively conducted on LLaMA, Yi, Qwen, and Mistral model families, at various scales from 6B to 8x7B, on base and chat models, and model dense and Mixture-of-Experts models.
- To identify which attention head is contributing to the retrieval from contexts, the authors proposed a novel retrieval score to measure the frequency of a head’s copy-paste behavior during autoregressive decoding. This retrieval score is both analyzed in different models and used to guide the experiments that prune or mask retrieval heads to understand the causal importance of retrieval heads.
- The authors also considerately report empirical results that these identified retrieval heads are activated during paraphrasing, QA, CoT reasoning tasks, and not just in copy-paste tasks.
Weaknesses: - It’d facilitate reading to clarify that "$k$ is a sentence that is irrelevant to $x$" in L146, instead of, for example, a short phrase or a single word. Can add a reference to Figure 2 so that readers see an example.
- The paper misses dataset details (L195, L355, L427). Are NIAH samples created manually or by prompting large language models? What datasets are used to begin with in Sec 2-4? What additional evaluation tests did you create in Sec 4.1-4.3?
- The paper misses experimental details, such as prompts used, links to existing assets used, etc.
- The questions below need to be addressed.
Recommendation: 8
Reviewer confidence: 4
Soundness/correctness: 3
Clarity: 3