Understanding Synthetic Context Extension via Retrieval Heads

Thank you for the feedback! We address your comments below:

W1: We believe that our results exhibit clear patterns, namely: Retrieval heads for real and synthetic versions of the same task have high overlap, as we noted in the paper: There is higher cosine similarity between retrieval heads for the same conceptual tasks (i.e. real and synthetic data for the same task) than across tasks (MuSiQue vs. SummHay Citation) The synthetic data retrieval heads are largely subsets of the real data retrieval heads. To make this more apparent, we have added recall tables and discussion in Appendix E: “We find that on all 3 tasks, the attention heads with non-zero retrieval scores on the real data have high recall ($\geq$ 0.76) against those identified on the synthetic data, while the reverse is more often lower.” We see greater retrieval head presence in the middle-to-last layers, as visualized in Figure 6 (formerly Figure 5) in Appendix D. SummHay insight heads follow this pattern with one clear difference: they are less prevalent in the last 2-3 layers of the model, reflecting their “intermediate” role of helping determine where downstream attention heads should look for the final answer tokens. Within a single layer, the relevant attention heads were randomly selected to be “primed” during pretraining for performing the target fine-tuning task. We have added this observation to the figure caption.

Additionally, retrieval heads can be identified for all tasks where the answer (or information relevant to the answer) are within the context, and therefore our analysis can be extended broadly.

W2: For the experiments in our work, the random attention head dropout results are averaged over three samples; since prior work has found that a small subset of the model’s components can match the full model performance on the Entity Tracking task (Prakash et al., 2024), we think that some of the random heads may not only be outside the relevant components, but also actively interfering with them. This idea is very compelling–but we leave the identification of attention heads with destructive interference to future work, since all of the attention heads selected to be randomly dropped have the same retrieval score of 0, and we need additional metrics to determine destructive interference.

W3: We evaluate across three tasks, which have been commonly used in this space, with objective automatic evaluation metrics. If this does not fully address your concerns, we’d appreciate additional clarification of this identified weakness!

Q1: We experimented with short-context fine-tuning of Llama-3-8B-Instruct on MuSiQue dataset variants. The resulting F1 scores from fine-tuning are: R,R = 0.59; R,R (L) = 0.46; H,H = 0.47; H,L = 0.48; L,H = 0.47; L,L = 0.50; S,S = 0.44. The F1 of the non-FT model is 0.44. In short context, the synthetic datasets result in marginal improvements over the non-FT model, and we find negligible correlation (Spearman R = 0.07) between F1 and retrieval score cosine similarity, despite similar ranges of cosine similarity (0.72-0.79) between synthetic data and the real data. We hypothesize that synthetic data is able to transfer better to the real data at longer context because it teaches the relevant attention heads to handle the new positional embeddings and produce sharp softmaxes over greater numbers of input tokens–abilities that should have been acquired at the shorter context during pretraining. This results in substantial gains over the non-FT model when fine-tuning on synthetic data at long context, which we have added to Table 1 for comparison. (For example, fine-tuning on the S,S dataset at short context results in +0.00 improvement, but leads to large (+0.10) improvement at 32K context.)

Thank you for the review!

W1: We’ve increased the label text in Fig. 1, Fig. 4, and Fig. 5. Additionally, we’ve added a heatmap legend to Fig. 1 and fixed the leftmost labels.

W2: Apologies for the confusion. We clarified our findings in the general response, in that this is not a strict subset relationship, but instead a “soft” one: the real data has high retrieval head recall (>=0.76 for Llama-3-8B-Instruct on all tasks) against the synthetic data, but not the other way around. We’ve added Appendix E showing pair-wise retrieval head recall, where the first rows of Tables 5, 6, & 7 show that the real data retrieval/insight heads for Llama-3-8B-Instruct have high recall against all synthetic dataset retrieval heads. We’ve also updated the text throughout the paper accordingly.

This “soft subset” result is significant because retrieval head recall is highly correlated with real task performance, which we have made more clear by adding Table 4 to Appendix E. This shows that, when synthetic datasets have higher retrieval head recall against the real task, performance is also higher, with strong Spearman correlation for both Llama3 and Mistral.

Q1.1 In short, the limited subset is defined as a subset of the real dataset that satisfies certain constraints of the distribution. We’ve added the details to Appendix C, which says:

On MDQA, we create three variants: $\text{L}_1$ is the subset containing Who, When, and Where questions; $\text{L}_2$ is the subset containing When and Where questions; $\text{L}_3$ is the subset containing only Who questions. These comprise 65.8%, 31.0%, and 34.8% of all questions in the MDQA training set respectively. In Table 1, Table 3, and Table 11, we only report $\text{L}_1$ results due to space constraints.

On MuSiQue, we use the subset of linear 3-hop questions consisting solely of T-REx component questions (Elsahar et al., 2018), as identified by ``$>>$". 10.8% of MuSiQue linear 3-hop questions in the training set fit this criteria. Additionally, among all component question hops in the training set, 43.0% are sourced from T-REX.

Q1.2: To experiment with scaling up the synthetic datasets, we double the size (to 800 examples) of the synthetic datasets for MuSiQue and fine-tune Llama-3-8B-Instruct. This results in the following F1 improvements: H,H = +0.13; H,L = +0.07; L,H = -0.03; L,L = +0.11; S,S = +0.08. Notably, both H,H and H,L variants outperform fine-tuning on 400 real dataset examples, and three datasets–(H,H), (H,L), and (L,L)--outperform fine-tuning on the 400 limited relation dataset. On these scaled synthetic datasets, retrieval head recall against the real data (400 examples) model is still correlated with performance, with Spearman R = 0.50.

Q1.3: We conducted an experiment on Llama-3-8B-Instruct to explore hybrid limited+synthetic data variants for MuSiQue. Namely, for each synthetic dataset $D_\text{synth}$, we trained a new model on the dataset $D_\text{synth} \cup D_\text{limited}$, resulting in 792 total examples. We find that performance generally increases over using only the synthetic dataset: H,H = +0.05; H,L = -0.02; L,H = +0.08; L,L = +0.09; S,S = +0.02. However, retrieval head recall and retrieval score cosine similarity against the real data change minimally (-0.04 to +0.04).

Q2: Added a footnote to the bottom of page 3 (highlighted): “We pad with irrelevant repeated text ‘The grass is green. The sky is blue…’ to ensure that the added paragraphs do not interfere with the answer to the original question.”

Q3: Thanks for pointing this out, the Mistral patching results table in Appendix F (now Table 11) has been fixed.

Q4: Full-rank finetuning on all parameters for 32K context is prohibitively expensive due to quadratic GPU memory requirements. In view of the limited computational resources we have, we ran LoRA fine-tuning on all parameters to approximate the full-rank fine-tuning setting and the results are presented in Appendix G, showing that we find similar conclusions, namely: there are mostly small (<0.05) performance differences between fine-tuning only attention heads and all modules. the exceptions occur on the SummHay Citation task, where fine-tuning all modules outperforms. when the synthetic data induces fewer retrieval heads (i.e. on MuSiQue), the synthetic data retrieval heads tend to be a “core” subset of real data retrieval heads–with “core” indicating higher retrieval scores. retrieval score recall and cosine similarity are moderately correlated with real task performance. We expect full-rank fine-tuning to yield similar results to LoRA.

Thank you for your review!

W1: You are correct that in this work, we presume that the real data is high quality and construct lower quality synthetic data, and focus on a quantity controlled comparison. In other situations where the real data is of low quality, then high quality synthetic data may outperform. We view the retrieval heads on high quality real data as indicative of the subnetworks required to perform the tasks adequately. When the synthetic data is higher quality and results in greater performance than fine-tuning on real data, a different analysis might be useful, as implied by our work: identify retrieval heads of the synthetic-data-model on the real data, and examine what additional components are contributing to performance. This is an interesting future direction!

W2: In our setup, we consider tasks where the answer to a query is contained in an input text $\mathcal{C}$, which contains spans of text that are relevant to answering the query, {$f_1, ...., f_m$}, which we deem our “needle concepts”. The remaining text $\mathcal{C} \setminus${$f_1, ...., f_m$} is considered the context. By partitioning the input text into these two sets–needle concepts and surrounding context, we can categorize any data as having some variant of “concept expression” or “context diversity”. Both concept and context text can range from being extremely diverse to extremely simplistic or structured. Within this framework, the symbolic data is a highly structured variant of concept expression and concept diversity that we felt was important to emphasize separately, since prior work [1][2] has observed that training on highly structured code data improves natural language entity tracking abilities. We’ve made this more clear in Section 3.1 and highlighted the change.

Among concept and context variants, we aim to explore ones that are most commonly found in prior synthetic data work, as we describe in Section 3.1. There are certainly far more variants that could be enumerated, and we leave this to future work.

W3: We use the sentences from [3][4] which are two popular long-context retrieval benchmarks. The use of padding data in the literature is to isolate the pure “context utilization” (ability to retrieve information regardless of position) performance from any “distractor distinguishing” performance. “Context utilization” by itself requires training when extending models to long context, since the attention heads will be exposed to new position embeddings (even when interpolating) and required to softmax over more context tokens. Table 1 shows that training using this “meaningless” padding does indeed result in greater performance (+0.12 to +0.19 F1 for Llama3 MuSiQue) over the Non-FT model.

As shown in Figure 2, top right, we also have variants that contain distractor context sentences that are similar to the “needle” concepts. We agree that there is a measure of “whether the context is distracting or not”--as measured by whether the context has high lexical, structural, or semantic overlap with the needle text–that may not be accounted for by a straightforward diversity metric such as token distribution entropy. We believe that our variants are effective at representing major prior works, and leave finer-grained taxonomy for future work.

Q1: Please see our response to W2. Namely, our framework is extensible to contextual reasoning tasks, by classifying all text within an input as either “needle concepts” {$f_1, ...., f_m$} or the surrounding “context” $\mathcal{C} \setminus${$f_1, ...., f_m$}. Symbolic tasks are a special variant of this that we chose to highlight. While there are far more possible variants than can be enumerated in a single paper, we aim to explore those most commonly found in prior synthetic data work.

Q2: As discussed in our response to W1, there is interesting future work suggested here. We think about retrieval heads as indicative of the subnetworks trained to accomplish a task, and for any given pretrained model, there are some components primed for further recruitment during fine-tuning. This causes high overlap of the “existing mechanisms” being enhanced when trained on variants of the same conceptual task. For some tasks, more components may be required that are not effectively targeted during fine-tuning due to the low quality or amount of data. We look forward to exploring this further!

[1] Kim et al. (2024). Code Pretraining Improves Entity Tracking Abilities of Language Models. ArXiv, abs/2405.21068

[2] Prakash et al. Fine-Tuning Enhances Existing Mechanisms: A Case Study on Entity Tracking. In Proceedings of the 2024 International Conference on Learning Representations, 2024

[3] Hsieh et. al. RULER: What’s the Real Context Size of Your Long-Context Language Models? In First Conference on Language Modeling, 2024

[4] Mohtashami & Jaggi. Random-Access Infinite Context Length for Transformers. In Thirty-seventh Conference on Neural Information Processing Systems, 2023

Thank you for the detailed suggestions! We would like to point out a significant misunderstanding regarding the “Lack of contributions” weakness, in that paper [1] is the present paper, uploaded during the anonymous review period, which is in compliance with https://iclr.cc/Conferences/2025/CallForPapers . It therefore does not detract from the contributions of the present work.

We address the other comments below:

Comparison to [2]: We believe that our findings in realistic, multi-hop reasoning settings constitute a significant contribution over the task used in [2], in which the model must retrieve the box containing an object which is stated explicitly in the context (no move operations, essentially a single-hop task, synthetically constructed). In addition, [2] focused on improvement from fine-tuning broadly on code data and does not show how different constructions of structured fine-tuning data might recruit the required entity tracking circuits in different or similar ways, while we do in our controlled synthetic data settings. Specifically, by identifying retrieval heads for the synthetic and symbolic datasets, we show that mostly the same components are performing the equivalent functionality after fine-tuning, as in the real task.

Missing Explanations:

Retrieval head definition: Added to Section 4.1 “Detecting Retrieval Heads”: “To compare across fine-tuned models, we consider any attention head with a positive retrieval score to be a retrieval head, and later compute cosine similarity to account for the strength of scores.”

Detecting common retrieval heads: No matching algorithm is used other than matching by head indices. Our results support the idea that specific attention heads are primed during pre-training to be more easily recruited and trained for certain functions. To visualize the efficacy of matching by head index, we updated the paper with heatmaps showing retrieval scores after fine-tuning on different datasets in Appendix D (these are currently a bit small but we’ll make them clearer in a later version). We note this in a (highlighted) footnote in Section 4.3.

Patching explanation: We’ve added a more detailed explanation in Appendix F.

Experiment request

This has been added to Table 1, showing that the non-FT model has extremely poor performance on the 32K context tasks. One major factor is that, during context length extension, attention heads are exposed to new positional embeddings (even if they are kept within pretraining range by appropriately scaling the RoPE theta or using position interpolation) and required to softmax over a greater number of context tokens [3], which requires model adaptation.

[3] Veličković, P., Perivolaropoulos, C., Barbero, F., & Pascanu, R. softmax is not enough (for sharp out-of-distribution). ArXiv, abs/2410.01104.

Questions

Q1: Reworded: “The quadratic memory requirement of Transformer attention imposes a strong computational constraint on our ability to train and do inference on long-context models. This disrupts the typical pre-training pipeline: pre-training must be done at as large a scale as possible, but pre-training a long context model would necessarily reduce the number of observed tokens able to fit on the GPU.”

Q2: We experiment with multiple subsets of relations for MDQA, resulting in 3 variants, as explained in Appendix C: L1 = Who, When, Where (65.8% of the full relation set), L2 = When, Where (31.0%), L3 = Who (34.8%) .We added a (highlighted) note to the Figure 4 caption.

Q3: Clarification added to a (highlighted) footnote in 4.3: “SummHay Retrieval Heads attend to the final answer (document number), whereas SummHay Insight Heads attend to the insight text within the document.”

Q4: Rephrased (and highlighted) for clarity: “One explanation is that fine-tuning induces upstream changes so that a different representation distribution is passed to the retrieval heads when learning on synthetic data. This allows retrieval heads to learn to be effective for the synthetic task while failing on out-of-distribution real data representations.”

Q5: As stated in Section 3.1: “Prior synthetic datasets have made use of fictional entities (Saparov & He, 2023) or nonsense phrases (Wei et al., 2023) in place of real entities and properties, or swapped out nouns to augment the dataset (Lu et al., 2024).” These effectively encourage generalization by preventing the model from overfitting to specific entities. Added a clarification!

Q6: We’ve replaced this text with “corresponding real task”.

Misc.

Where do numbers from paragraph in line 356 come from? Figure 1 does not have 0.35 EM or 0.32 and 0.20. Where do number of heads 129, 112 and 39 heads come from?

Figure 1 shows the F1 score, not the exact match accuracy; we’ve updated the text to reflect this.

We’ve also fixed the remaining instances of unclear writing and typos in the updated version! Dataset examples have been added to Appendix C.1.

Thank you for the score increase and for considering our response!

We appreciate the prior contributions of [1] & [2] but believe that there is need for work like ours that demonstrates that mechanistic interpretability can give insight into realistic data with complex reasoning. Additionally, while [1] & [2] look at how the target toy task circuit changes after fine-tuning on various data, they compare base models with versions fine-tuned on large mixtures of data and do not (i) identify the specific data that increases target task performance, or (ii) explain how fine-tuning on data with very little token overlap (e.g. realistic vs. symbolic data) manages to improve the target task performance. Transformers can conceivably create different intermediate representations of these different data types, thereby either (A) inducing high activations on disparate sets of model components, or (B) inducing high activations on the same sets of model components in ways that do not affect data with out-of-distribution representations. Our work shows that such dissimilar data can still induce changes in the desired realistic task components during long-context fine-tuning.

Ultimately, bridging toy data and realistic data is an important piece to demonstrate broader applicability of mechanistic interpretability, provide a concrete foundation for using toy tasks as long-context benchmarks, and shed light into why fine-tuning on structured data such as code can improve performance on realistic tasks.

Again, thank you for your consideration!