Retrieval Head Mechanistically Explains Long-Context Factuality

We thank the reviewer for their support and detailed comments! The reviewer is mainly concerned how the model behavior would transfer to other tasks like paraphrasing. Here we note:

Experiments on paraphrasing task

Below we show an example about how masking out retrieval heads breaks the model’s understanding of semantic dependency while masking out random non-retrieval heads does not:

• Input: The Whispering Forest of Lunthera, home to bioluminescent insects, is famed for its murmuring trees and an ancient legend about lost travelers finding their way home.
• Mask 30 retrieval: Lunthira’s Whispering Woods, home to luminescent creatures, is renowned for its lulling trees and a mythical tale about lost adventurers who happened upon the forest by chance.
• Mask 30 random: Lunthera’s Whispering Forest, abounding in luminescent insects, is renowned for its conversing trees and a fable about lost travelers finding their way home.

In this example: if one mask out the retrieval head, the model outputs “luminescent creatures” without specifying they are insects; it outputs lulling trees while the input says murmuring trees. The input also says lost travelers “find their way home”, but this information is missed. In contrast, masking out random heads do not have these problems

Below is another example showing how masking out retrieval heads make the model hallucinate about information that does not exist in the input:

• Input: The glowing sands of the Duskveil Desert, enriched with rare minerals, shimmer in the dark and are said to hold heat for days after the sun sets.
• Mask 30 retrieval: In the Dusk Veil Desert, void of life and composing largely of glass, the sand is imbued with solar energy and preserves heat under the twilight.
• Mask 30 random: The verdant sands of the Duskveil Desert, composed of rare minerals, linger with a glowing glow at sunset and are said to retain heat for days.

In this example, masking out retrieval heads makes the model saying “composed of rare minerals” while the input says “rare minerals”. Masking out random heads do not have these problems.

In the updated paper’s Appendix Figure 15, we give more examples about how masking out retrieval heads influence model’s behavior on paraphrasing and question answering, while masking out random non-retrieval heads do not significantly change model’s behavior.

Other important comments

• We fixed the repetition sentence in section 3.
• We fixed the typo “needel” in Figure 9.

We thank the reviewer for the detailed comments and the support! Regarding questions:

1. Are retrieval heads consistent across different architectures? Does the role of retrieval heads vary with different transformer architectures (e.g., decoder-only vs. encoder-decoder models), or are these properties universally applicable?

We tend to believe that retrieval heads are consistent across architectures and the current experimental results on a variety of transformer variants supports this hypothesis. Specifically, we have conducted experiments with:

• Grouped-query attention (e.g., Qwen1.5),
• Mixture of Experts models (e.g., Mixtral), and
• Hybrid models incorporating state-space layers (e.g., Jamba, Please refer to the comments above titled [Update on Experiments] We also found Jamba for more details).

As today's mainstream models are decoder-only, we currently did not include encoder-decoder architecture. Our belief is that retrieval heads is a property emerge with attention layers, so as long as the architecture has at least one layer of global attention, we tend to believe there will be retrieval heads within it, either decoder-only or encoder-decoder.

2. How does the model determine when retrieval heads should be dynamically activated? Is there a mechanism or threshold within the model that dictates when these retrieval heads become active, especially in different contexts?

We are also fascinated by how exactly the firing / triggering of retrieval happens within the architecture. Currently we are unable to ping point an triggering mechanism -- or maybe they are intrinsically hard to identify because attention is basically dot product between two vectors, and whether retrieval happens basically trace back to the exact numbers of the key-value vectors, which are outputs of previous layers, thus a little bit hard to determine.

We envision that future research could adopt approaches similar to the ongoing exploration of the "physics of LLMs", starting with smaller models and synthetic datasets to isolate individual model behavior. For example, keeping all other factors the same, control only one aspect of the data, such that one model can do retrieval and another model cannot do. We would be excited to see how this would be studied in future work.

We appreciate the reviewer’s detailed comments and support. Below are our responses:

Diversity of needle-haystack

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)

In our initial submitted version before the rebuttal, we use three sets of needle manually written by our authors. These needles are:

1. Needle: A new report from the WMO shows that records were once again broken, and in some cases smashed, for greenhouse gas levels, surface temperatures, ocean heat and acidification. Question: What does a new report from WMO shows?
2. Needle: The best thing to do in Beijing is to take a walk in Chaoyang Park and have a cup of Espresso in the evening. Question: What is the best thing to do in Beijing?
3. Needle: Mr Green is disliked by everyone because he is a mean person and also he can't ride a horse or dive a car. Question: Why does everyone dislikes Mr Green?

The haystacks are randomly sampled documents from Slimpajama.

Since the reviewer is concerned about the diversity of the needle, we conduct a follow-up experiment to demonstrate that the detection results of retrieval head do not change when increasing the diversity of needles. Specifically, we use the above three cases as in-context examples to a language model, and ask the language model to generate 100 more examples. By doing this we get a synthetic data of 1.2K unique tokens spanning 10 topics (Technology, Transportation, Education, Festivals, Health .etc). Below are three examples of the synthetic data:

1. Needle: The newly constructed SkyBridge connects three cities—Everdale, Pinehurst, and Riverpoint—allowing citizens to travel with ease, admire scenic views, and save significant commute time. Question: What does the newly constructed SkyBridge connect and offer?
2. Needle: The upcoming GalaxyFest will feature over 100 sci-fi authors, 50 exclusive book signings, and a virtual reality experience of Mars colonization. Question: What features will the upcoming GalaxyFest have?
3. Needle: The Aurora Conservatory is renowned for its collection of rare Arctic flora, cutting-edge climate research, and eco-friendly glass dome architecture. Question: Why is the Aurora Conservatory renowned?

Repeating the retrieval head detection algorithm on Mistral 7B with our newly generated 100 needles, we get the same set of retrieval heads as before. This is to say, we confirm that out conclusion holds when scaling the number of needles from 3 to 100. We have added the details in appendix Figure 14.

Other important comments

1. We added the clarification noting “KV” means “key-value” in the last line of abstract.
2. We elaborated the meaning of Figure 3 in the caption by explaining the color and the sparsity of retrieval heads.
3. We added the elaboration of Figure 5. Note that Figure 5 is complimentary with Figure 6 because it not only shows chat model and base model share the same set of retrieval heads, but also shows that these heads are mostly within middle layers.
4. We have added examples of masking out retrieval heads v.s. random heads for paraphrasing and question answering in the updated appendix Figure 15. These examples consistently demonstrate the influence of retrieval heads in downstream tasks.

We thank the reviewer for the detailed comments. Below we note:

About missing details:

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

We appreciate the comment on missing details. Please see our [Update on Experimental Details], where we have included comprehensive descriptions of the datasets, methods for creating NIAH samples, and evaluation tests. We would also be happy to provide further follow-up explanations. These updates will also be incorporated into the revised version of our paper. Our full reproducible code and data is open-sourced, and we will reveal the link after the anonymity period.

About $R$:

L156: By $R$, do you mean real numbers? If so, perhaps use $\mathbb{R}$ instead and clarify that refers to unnormalized attention scores.

Yes, $R$ refers to real numbers, and $\mathbb{R}$ is indeed more precise. We have updated the notation and clarified its meaning in the revised paper.

Question on Retrieval Head detection:

L194: Does it happen that the model generates a word that is not the next token in the needle that should be generated? If this happens, do you skip the example? Or consider that as a case when all attention heads do not perform a copy-paste operation, even if an attention head actually pays the most attention to the token that should be generated next?

• Our retrieval head detection is specifically targeting the copy-paste behavior. So yes, if the model generates a tokens that is not in the needle, we do not count it, no matter whether the attention head is attending to the target token or not.
• That being said, since the needle sentence contain multiple tokens, to identify retrieval heads the model does not necessarily need to generate all of the tokens in the needle. Say if a needle "best place to visit in SF ..." contain 20 tokens, as long as the model copy-paste a fair portion of the tokens, say 14 out of 20, this level of copy-paste is enough for us the identify how strong/ frequently an attention head is doing retrieval.

L203: What do you mean by “stabilize” and “converge”? Please either provide definitions or plots to illustrate.

• Since heads detected with strong retrieval scores in one context may not exhibit strong retrieval scores in another context, so we conduct retrieval head detection over multiple (question, answer pairs) and multiple long-contexts (haystacks).
• As we increase the number of trials for detecting retrieval heads, “stabilize” and "converge” refers to the ranking of attention heads (based on retrieval scores) becoming consistent across repeated trials of retrieval head detection.
• We have added definitions and included supporting plots in Section 7.2, Figure 13, in the updated Appendix.

Writings

It’d facilitate reading to clarify that " is a sentence that is irrelevant to " 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.

Figure 3: Seems that in fact less that 1% of attention heads are activated more than 50% times. The 5% in the caption could probably be changed to 1%.

Figure 7: Could be nice to include the dataset name in the caption.

Thanks for the suggestions and we have modified the paper in in our new revision.