Understanding Synthetic Context Extension via Retrieval Heads

Thank you for your review!

The paper heavily attributes the performance gap to retrieval heads, potentially underplaying the role of other transformer components like MLPs, which are known to handle parametric knowledge. Although the authors note in a footnote that similar conclusions hold when fine-tuning all modules, this is not fully explored in the main text

Parametric knowledge should not be required for strong performance on the tasks we examine, since all the relevant information is contained within the input context. In fact, the closed-book F1 on MuSiQue is 0.086 for Llama3 and 0.042 for Mistral when prompted with the question directly and none of the relevant context documents, an indication that the model cannot “cheat” by drawing on parametric knowledge to answer the question. MDQA’s closed-book accuracy is higher (0.333 for Llama3 and 0.209 for Mistral). As a result, we think our focus on attention heads is consistent with the experimental evidence we’ve found as well as intuition from prior work. However, we can add more thorough discussion of this point in the main text in any future version.

The paper varies synthetic data along concept expression and context diversity, but other factors (e.g., reasoning complexity or distractor presence) might also affect performance and are not explored.

Reasoning complexity: for our experiments, we explore tasks at 3 different levels of complexity: single-hop (MDQA), two-hop (SummHay Citation), and three-hop (MuSiQue).

Distractor presence: In our framework, distractor presence is a variation of the context diversity rather than its own axis. In future work, it would be interesting to clearly define what makes a document (or sentence) “distracting”, such as if it has high token overlap with the question or needle sentence but is not actually useful for answering the question. Formalizing these notions would require tailoring to each “conception expression” variant, so we do not explore it here.

The experiments are conducted on two models (Llama-3-8B-Instruct and Mistral-7B-Instruct-v0.1) and three tasks. While the results are consistent, they may not fully generalize to other models with different architectures or pretraining, or to a broader range of tasks.

Different architectures or pretraining: Retrieval heads have been shown to be present across model families with different architectures (attention variants and mixture of experts) by Wu et al. (2024). We study how they are affected by fine-tuning in this work, particularly in two architectures with different attention variants, which we believe to have the most impact on the behavior of retrieval heads.

Broader range of tasks: While our results are most applicable to tasks that have a long input context, and where the final answer must be retrieved from that context, the concept of identifying attention heads which attend to the input context is also broadly applicable: as we demonstrate with SummHay “Insight Heads”, looking at attention heads that correctly identify relevant intermediate information in a multi-step reasoning task is also a strong indicator of synthetic data performance.

Thank you for your review!

For constructing synthetic long-context data, there has been some existing work that proposed general principles for creating synthetic training data beyond dataset-specific constructions, such as [1,2]. Although these work did not focus on the principles proposed in this work, it would be nice to have some discussion on these general methods for synthetic long-context data.

Regarding the construction of general synthetic context extension datasets, our results indicate the effectiveness of a diverse source of documents and tasks as done in [1] and [2]. We will add these citations to the paper, thanks!

It could be better if there are some analysis on the patterns of retrieval heads in existing long-context LLMs (such as Mistral-v0.2 and Llama-3.1). If there are some observations in line with the conclusions in this work, it could greatly contributes to the soundness of this work.

Wu et al. (2024) showed that retrieval heads are clearly identifiable in existing long-context LMs such as Mistral-v0.2, which are trained on a wide variety of data mixtures. Our research question pertains to how specific types of context extension data influence the training of retrieval heads. While examining the retrieval heads used by existing long-context LMs on MDQA, MuSiQue, and SummHay after training on large data mixtures is an interesting extension, we consider it to be out of the scope of this work.

Thank you for the review!

It is suggested to discuss its position in the context of general-purpose data synthesis for LLMs [1,2]; context compression for long-context LLMs [3,4]; and analysis of synthetic data biases and their effects on downstream performance [5,6].

Thank you for the suggestion, we’ll add the discussion to the paper. We note for the other reviewers and the AC that these papers do not directly impact the conclusions and comparisons in our paper.

1. Retrieval heads seem necessary but not sufficient for strong task performance (MoreHopQA: More than multi-hop reasoning).

We agree with this view: retrieval heads are a useful indicator for comparing the model components targeted by different datasets, but do not comprise the whole set of model circuitry required for a given task. The advantage is that the concept of identifying attention heads which attend to the input context is very broadly applicable: as we demonstrate with SummHay “Insight Heads”, looking at attention heads that correctly identify relevant intermediate information in a multi-step reasoning task is also a strong indicator of synthetic data performance. This can be extrapolated to tasks like MoreHopQA that ultimately have generative answers–as long as the input context has required relevant information, attention heads that correctly attend to the relevant information can be studied.

2. Would your results generalize to generative synthetic data tasks (e.g., reasoning, math, or code)?

Our results are most applicable to tasks that require a long input context, and generative coding tasks that include an existing codebase or libraries as input would fit this description. As for tasks with shorter input context but long generative outputs, like most math problems in existing datasets, looking at the behavior of attention heads that might attend to the relevant parts of the intermediate generated output (“chain of thought”) is an interesting question that we leave for future work.

3. Does retrieval head behavior vary across different transformer architectures (e.g., mixture-of-experts models)?

We experimented with two attention variants in our work–full attention in Llama-3-8B-Instruct and sliding window attention in Mistral-7B-Instruct-v0.1, since these attention differences are significant for long-context retrieval behavior. We show consistent observations across these variants. Since mixture-of-experts is an architecture variant that changes the MLP modules and does not affect the attention mechanism, we do not think this would make a significant difference–in fact, Wu et al, 2024 showed that the basic properties of retrieval heads are similar for both Mistral-7B-v0.2 and a MoE variant in the same model family, Mixtral-8x7B-v0.1.

4. How does synthetic data performance change as model size scales up? Does a larger model mitigate synthetic data limitations?

Due to computational limitations, we restricted our investigation to 7B-scale models. (Note that long-context models have more extreme memory demands than standard-context models.) Nevertheless, dynamically activated retrieval heads have been shown to comprise a similar fraction of the total attention heads across model sizes (Wu et al., 2024). Since retrieval heads are unevenly activated based on context, we can speculate that synthetic data might face similar limitations for larger models.

Thank you for the review!

Statistical analysis such as confidence intervals, p-values are missing.

To address this, we will make the following additions to Tables 1, 3 and 11 to show when a performance gain is significant:

Table 1: $\dagger$ indicates that a model trained on the bold synthetic dataset in the column outperforms a model trained on the indicated dataset with $p < 0.05$ according to a paired bootstrap test. We see that most gains $\geq 0.02$ are statistically significant.

Tables 3 and 11: Patching Results (summarized due to character limits): We perform a paired bootstrap test to test whether any patch model outperforms the original model. In Table 11, all performance improvements for patched synthetic data models on MDQA and MuSiQue are significant with $p < 0.05$, and gains on SummHay are not significant. In Table 11, all performance improvements for patched synthetic data models on MDQA and MuSiQue are significant with $p < 0.05$, and gains on SummHay are not significant.

While the analysis employs a well-established methodology—specifically adopting the retrieval head framework defined by Wu et al. (2024)—the study primarily extends prior work by applying this technique to examine LLM behavior on synthetic datasets. This approach, while methodologically sound, results in limited novelty, as it does not introduce significant conceptual or technical innovations beyond the foundational framework.

In addition to showing a strong relationship between the retrieval heads recruited by synthetic datasets and downstream performance, our work takes a novel step in demonstrating that mechanistic interpretability can give insight into realistic data tasks involving complex reasoning. In contrast, Wu et al. (2024) only examined a single-step retrieval task. We also show that attention heads that attend to intermediate information within the context can be clearly identified and used to understand dataset performance, which extends the breadth of tasks that can be examined beyond purely extractive tasks.

You note that different tasks leverage different sets of retrieval heads. What do you believe accounts for these differences, and how should this influence the design of synthetic data for specific types of long-context tasks?

While our work does not do a full characterization of the circuits that are relevant to solve each task, each task requires different upstream capabilities, which we think leads to different attention heads recruited for the final retrieval step. (And different attention heads are recruited for different intermediate reasoning capabilities, as indicated by the SummHay Insight Head analysis).

To design synthetic data for different long-context tasks, our work shows that only a small (e.g. ~40 MuSiQue examples) number of real examples are needed to identify relevant retrieval heads, and then this can be used to assess the performance of promising synthetic datasets (using a few hundred examples) before scaling up synthetic dataset generation.