Understanding Synthetic Context Extension via Retrieval Heads
Authors: Xinyu Zhao
Fangcong Yin
Greg Durrett
Venue: ICML 2025 poster (2025)
Keywords: Large Language Models
Synthetic Data
Long Context
Retrieval Heads
Primary area: deep_learning->large_language_models
TL;DR: We identify when synthetic data falls short on learning long-context abilities and trace an explanation to a specific set of attention heads.

Abstract: Long-context LLMs are increasingly in demand for applications such as retrieval-augmented generation. To defray the cost of pretraining LLMs over long contexts, recent work takes an approach of synthetic context extension: fine-tuning LLMs with synthetically-generated long-context data. However, it remains unclear how and why this synthetic context extension imparts abilities for downstream long-context tasks. In this paper, we investigate fine-tuning on synthetic data for three long-context tasks that require retrieval and reasoning. We vary the realism of "needle'' concepts to be retrieved and diversity of the surrounding "haystack'' context, from using LLMs to construct synthetic documents to using templated relations and creating symbolic datasets. Although models trained on synthetic data underperform models trained on the real data, the impacts of both training settings can be understood via a shared feature of the attention computation, retrieval heads (Wu et al., 2024). The retrieval heads learned from synthetic data have high overlap with retrieval heads learned on real data. Furthermore, there is a strong correlation between the recall of heads learned and the downstream performance of a model, allowing us to interpret and predict the performance of models trained in different settings. Our results shed light on how to interpret synthetic data fine-tuning performance and how to approach creating better data for learning real-world LLM capabilities over long contexts.

Understanding Synthetic Context Extension via Retrieval Heads
Xinyu Zhao'
Fangcong Yin
Greg Durrett
Abstract
disrupts the typical pre-training pipeline: pre-training must be done much data possible, but pre-training long context model would necessarily reduce the number of observed tokens able to fit on the GPU. One solution for this is to rely on synthetic data, now common in post- training settings such as SFT (Xu et al., 2024; Yue et al., 2024; Xu et al., 2025; Chen et al., 2024a) and RLHF/DPO (Yang al., 2024). Recent prior work has proposed using synthetic data to extend the long-context abilities of LLMs after pre-training (Xiong et al., 2025; Zhao et al., 2024; Team, 2024).
Long-context LLMs are increasingly in demand for applications such as sarrival--augmenter gen- eration. To defray the cost of pretraining LLMs over long contexts, recent work takes an approach of synthetic context extension: fine-tuning LLMs with seahhetccall-generatee long-context data. However, remains unclear how and why this synthetic context extension imparts abilities for downstream long-context tasks. In this paper, we investigate fine-tuning on synthetic data for three long-context tasks that require retrieval and rea- soning. We vary the realism of "needle" concepts to be retrieved and diversity of the surrounding "haystack" context, from using LLMs to construct synthetic documents to using templated relations and creating symbolic datasets. Although models trained on synthetic data underperform models trained on the real data, the impacts of both train- ing settings can understood via a shared feature of the attention computation, retrieval heads (Wu et al., 2025). The retrieval heads learned from synthetic data have high overlap with retrieval heads learned on real data. Furthermore, there is a strong correlation between the recall of heads learned and the downstream performance of a model, allowing us to interpret and predict the performance of models trained different set- tings. Our results shed light on how interpret synthetic data fine-tuning performance and how to approach creating better data for learning real- world LLM capabilities over long contexts.
This use of synthetic data is particularly necessary for long context tasks since they are so laborious for humans to man- ually label. Synthetic data is also configurable: it can exhibit different reasoning skills and "teach" models have make certain types of inferences (Du et al., 2017; Yu et al., 2018; Agarwal et al., 2021; Tang et al., 2024; Divekar & Durrett, 2024). One way to do this is using templates to express pieces of information that must be reasoned over and to cre- ate symbolic tasks that are thought to mirror the reasoning required in the real task (Hsieh al., 2024; Prakash et al., 2024; Saparov & He, 2023; Li et al., 2024). However, past work has shown varying results from training this kind of context extension (Fu et al., 2024); we lack gen- eral understanding of what properties are needed for good synthetic data.
data for
In this paper, we explore several methods of creating syn- thetic long context data across three tasks. Our goal is to examine what makes synthetic data effective for this kind of context extension. While more realistic data is often better, it is unreliable: certain types of more synthetic data can exhibit desired long-cou contect patterns even more effectively and with fewer shortcuts than realistic data. However, other types of synthetic data severely underperform on these tasks.
1. Introduction
The quadratic memory scaling of Transformer attention imposes a strong computational constraint ability to train and do inference on long-context models. This
To understand this divergence, we analyze models finetuned different synthetic long-context datasets for the presence of special set of attention heads called retrieval heads (Wu et al., 2025). Figure/ 1 shows two of our key results. First, the retrieval heads learned on poor-performing syn- thetic data tend to be fewer than those learned on realistic or high-quality synthetic data. Second, we find that the cosine similarity of retrieval scores of individual heads learned on synthetic data and realistic data correlates strongly with
'Department of Computer Science, The University of Texas Austin, Texas, USA. Correspondence to: Xinyu Zhao wtwuuzhaa@uutxxsccdiu
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Understanding Synthetic Context Extension via Retrieval Heads
Real Data
Synthetic Data (Realistic Concept) (Mohtashami & Jagai, 2023)
Synthetic Date (Synthetic Concept) (sieh et al. 2024)
Doc 1) Bily Giles was an Ulster Volunteer Force volunteer who later became active in poltics... Saptember 1938 (aged 41)
[Oac The grass is groen. The sky is bive. The sun yellow. United Kingdom is place of death of Billy Giles. The sky is blue. Here we Thore and back again.... [Irelevant distractor sentences)...
[Doc 1) The grase is green. The sky Live. The is yellow: WPFN the place death ABCD. The sky is blue Here we go. There and back egain (Ireeeevnt distractor sentences)....
Billy died
Belfast, Northern Ireland, the United Kingdom...
/Wikppddaa Passagas]
(Doc Ni Pound sterling is the currency of the
(Dac Ai Pound Starling is the currency of the United Kingdomm The grass green. The aky is blue. The sun is yellow. The sky is blue. Here we go...
[Dec Ny WXZY the currency of WPFN. The grass green. The sky The sky blue. Home we go.
ternitonies The pound is the main unit of United Kingdom
its associated
sun is yellow
sterting....
Q: What surrency is used where Bily Giles ded?
Q: What ourrency used where Duly Giles dled?
Q: What currency used where ABCD died?
A: Pound Sterling
A: Pound Sterking
A: WXZY
SFT
SFT
SFT
Retricval Heads
Performance
Retrieval Hoads
Performance
Retrieval Heads
Parforance
0.5
AX 0.4
K35
6.7
03
02
M45 01
465
045
We explore synthetic context extension with different forms of synthetic data across multiple tasks. Examples for two-hop question from MuSiQue (Trivedi et al., 2022) are shown here. A special set attention heads, retrieval heads (Wu et al., 2025), help explain the performance gap between fine-tuning real data and synthetic data.
Figure/ 1.
the downstream performance. Learning the "right" set of retrieval heads seems to be necessary condition for high performance. However, it is not sufficient: even when our synthetic data largely trains the same retrieval heads re- quired for the real task, it does so less effectively. We show that patching heads at the intersection of poor-performing model and a high-performing model can improve perfor- mance of the former: these heads are where important oper- ations are happening, but realistic data teaches them more strongly. While these heads do not tell an entire story of the long- long-context understanding mechanisms, they serve as an effective indicator of the subnetworks affected during fine-tuning.
assume a distribution pt of real-world task instances. We assume that small. limited set of input-label pairs Dy (IT, UT) drawn from the distribution pr is available as seed data. A synthetic dataset Dy is set of input-label pairs sampled from the outputs of a data generator 9 given the seed data or the known properties: Dy ~ p((ā, D) Dr). Benchmarking or training . on such a synthetic DT is expected to evaluate or teach M the capabilities that can be transferred to the real-world distribution PT.
A recent line of work has shown that simple heuristic-based synthetic datasets can be surprisingly effective for context extension, a post-training scenario where LLMs that have been pre-trained on short-context corpora are further trained on long-context tasks to extend the effective context window (Fu et al., 2024; Zhao et al., 2024; Xiong et al., 2025). For example, Xiong et al. (2025) finds that fine-tuning on synthetic simple dictionary key-value retrieval task can even outperform models fine-tuned on realistic in-domain data. Other types of simplified and symbolic data are used in long-contex: benchmarks (Hsieh et al., 2024; Li et al., 2024), despite the fact that the overlap between these and realistic data capabilities has not been thoroughly studied.
Our contributions are: (1) analysis of synthetic data across three synthetic tasks for long-context LLM training to de- termine the effect of varying data realism, including using symbolic data; (2) experimental results showing that learned retrieval head scores correlates with effectiveness of the training data for this setting. Taken together, we believe this work indicates a path forward for how to engineer better synthetic data and how to connect the construction process of synthetic data (a) what it teaches Transformers and (b) how those models perform on downstream tasks.
We call these approaches synthetic context extension: us- ing synthetic data to extend the context window of LLMSs It remains unclear how and why synthetic data, especially when drawn from a very different distribution from the real data, can be effective despite results that support the con- trary (Chen et al., 2024b; Liu et al., 2024b). There is also a lack of general principles for creating synthetic training data beyond dataset-specific constructions in the literature. We
2. Background and Setup
2. Background: Synthetic Data for Training LMs
Formally, consider a supervised learning setting for a pre- trained transformer language model M. Given a task T, we
2

---

Understanding Synthetic Context Extension via Retrieval Heads
start by constructing synthetic datasets varying in systematic ways to unify these variants from the literature.
Training Configuration For each task, we fine-tune two short-context LLMs, Llama-3-SB-Instrucc (Team, 2024) and (Jiang et al., 2023). Prior work indicates that attention heads are largely responsible for implementing algorithms (Olsson et al., 2022) and using information within the context (Stolfo et al., 2023; Lieberum et al., 2023) while MLP layers are responsible for parametric knowledge (Geva et al., 2021). In addition, when adapting to long contexts, attention heads in particular must handle new position embeddings and softmax over more context tokens (Veličković et al., 2024). Therefore, we fine-tune attention heads only."
2.2. Experimental Setup
Following Xiong et al. (2025), we focus on fine-tuning LLMs for long- retrieval and reasoning tasks where training on high-quality symthetic data has been shown to outperform real data. We also extend to multi-hop settings. We experiment on tasks where, given a long context C. and a context- -bseed query 9. a language model .M needs to retrieve one or more "needle concepts" fi...., fm from C (pieces of relevant information), reason over that informa- tion, and then generate response y ~ p(y C, q) where p(y C, q) is the conditional distribution that M places over the vocabulary Σ" given the context and the query. We con- sider extending the context window from 8K to 32K tokens to be representative of synthetic context extension following Chen et al. (2023). We use the following datasets.
To extend models from their original 8K pretrained context length to 32K, we follow (Gradient, 2024) in calculating new RoPE (Su et al., 2024), theta values, using 6315088 for Llama-3-8B IIsttrucctt and 59300 for Mistral-7B-Instrun v0.1. We scale the sliding window accordingly for Mistral- 7B- Instrruuut---.. to 16k context. These are the only adjust- ments we make to the models, following Fu et al. (2024). Our hyperparameters and hardware setup can be found in Appendix C.I.
MDQA (Liu et al., 2024a): MDQA is multi-document question answering (QA) dataset where only one paragraph in C' contains the gold answer to single-hop query; that is, there is a single f which directly addresses q. We extend the original MDQA dataset in 4K context to 32K context by retrieving additional distractor paragraphs from Natural Questions-Open (Kwiatkowski et al., 2019; Lee et al., 2019) with Contriever (Izacard et al., 2022).
3. Synthetic Datasets
3.1. Principles Under Consideration
To create representative range of synthetic data for each task, we partition the input text C into (A) text containing relevant information {Ji,..., Sm) ("needle concepts") and (B) the surrounding context C\ C'ffi...... fm This allows us to categorize any task as having variant of concept ex- pression (how the target information fi is expressed) and context diversity (the naturalness and relevance of the sur- rounding information). In the following paragraphs, we discuss common variants found in synthetic data literature. We single out and emphasize a highly structured variant of concept and context, symbolic tasks, for being devoid of natural language yet noted to transfer to realistic tasks (Xiong et al., 2025).
MuSiQue (Trivedi et al., 2022): MuSiQue is multi-hop QA dataset where the model must identify piece of relevant information from different document for each hop of the question in order to retrieve the final correct answer from the context. We use the linear three-hop subset of MuSiQue and extend the dataset to 32K by adding padding paragraphs to the original context. In this setting, the facts fi, f2, fa are natural language sentences expressing knowledge triples. ?
SummHay Citation (Laban et al., 2024): Summary of a Haystack (SummHay) is long-context retrieval dataset where the model is given a set of documents with controlled "insights," " and asked to produce a list of key points. Addi- tionally, the model must cite the correct documents in sup- port of each key point. We isolate the citation component and construct a task where, given a haystack of 10 docu- ments and a key point ("insight"), the model must correctly identify the two documents that support the point and their associated document IDs. The two facts fi, fa may span multiple sentences and may be substantially paraphrased versions of the insight.
Concept Expression A common procedure for creating synthetic data involves exploiting task asymmetry (Josifoski et al., 2023; Xu et al., 2023; Lu et al., 2024; Chandradevan et al., 2024; Chaudhary et al., 2024; Tang et al., 2024), where asking an LLM to generate natural language data based off of a label (e.g. a sentence based off of a knowledge triple) is easier than predicting the answer from text of the same complexity and domain. In this scenario, the LLM is asked to create diverse "needle" target concept expressions ft. In task-specific cases, it beneficial to make this data less realistic while encouraging generalization.
FFllowing (Mohtashami & Jaggi, 2023). we pad with irrele- vant 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.
NNot that this is different from the two-hop examples in Figure used for demonstrative purposes.
We find similar conclusions when fine-tuning all Llama-3-8B- Instruct modules; see Appendix H.
3

---

Understanding Synthetic Context Extension via Retrieval Heads
Real Data
Synthetic Data
Realistic
Synthetic
iDoc Ij Buy Giee Listrr Volunteor Forna valunteor who became active poltics.... Billy Giles died on 25 September 1996 (eged 41) Belfast, Northern frefand, the United Kingdom.... Giles commemoratedd, along with other prominent Loyatsi paramillariis.
[Doc Bily Giles was an Lister Volunteer volunteer who later became active in pollccs At the age of 41, Billy Giles died in Belfast, Northern Iroland, UK on 25 Soptember 1998.
Dee ABCD was an XL2S who became active in LHSZ.... WFPN the place death ABCD.
Concept Expression
[Dse NJ WXZY the currency of WPFN. The pound is the main unit of WXZY...
(Doc NJ The currency of the UK as wall as nine of its affiliated territories is pound sterling, The pound is the main unit of steriig...
Q: What curranyy is used whem ABCD dlad?
A: WXZY
... /Wiippedia Passages)... ...
[Doc M] Pound sterling the currency the United Kingdom and rine of associsled temtoniis. The pound is the main unit of starking...
Volunteer Forze, Bily Giles lator participated in political activitis.. .. Belfast, Northern Ireland of the UK is the place of dealh of Billy .. Together with other fighters of Loyalst
(Ooc As a volunteer of
[Doc The grass green. sky blue. sun yollow. The sky blue. Hore we go. Belfast, Northern Irefand of the UK is the place of death of Billy Giles. [Dreelevnnt olstractors)....
Context Diversity
Giles nighly respected... ..
[Doc NJ Pound Sterling is the currency of the UK. The graes green. The sky bive. The yellow. sty blue. Here we go...
Q: What currency used where Billy Giles dlea?
[Doc NJ Pound Sterling the currency of the UK.
A: Pound Sterling
Starking primary currency unit of pound...
Figure 2. Examples of elements of synthetic datasets for MuSiQue with varying levels of concept expression and context diversity. The needle sentences fr in the context and the entities in them are bold. High concept expression means more realistic expression of the needle fi. and low expression means more synthetic, including replacing real entities with symbolic entities or transforming. fi into templated sentences. High context diversity means more realistic context surrounding the needles, and low means more synthetic contexts such repeated, irrelevant padding sentences
For example, 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) and prevent overfitting to specific entities. In long context benchmarks (Hsich et al., 2024; Li et al., 2024), it is common to express the needle concepts in short, templated sentences.
3.2. Synthetic Dataset Construction
For each of the long-context tasks, we sample a set of exam- ples Df from the training set and use the principles above to construct various synthetic datasets based on Dy. See Appendix A for the complete set of prompts used to create the synthetic data, and Appendix B for our training prompts.
MDQA Given a training example of MDQA training data gC, E Dr. we combine the query q and answer y into our needle f that will be put into the context and that needs to be retrieved by the model. For f. we use two simplification levels of concept expression by (1) keeping the real entities in the query and answer (high expression), and (2) replacing the real entities with 4-character symbolic entities (low expression). We create the context surrounding the needle claim with two levels of context diversity: (1) prompting GPT-4o-mini to paraphrase the original context from MDQA training data (for the real entities), or generate a Wikipedia-style paragraph that elaborates on the claim (high diversity); (2) padding the context paragraph with repeated sentences (low diversity). The symbolic dataset is the simple dictionary key-value retrieval dataset from (Xiong et al., 2025).
Context Diversity We can also vary the expression of CMf..... }m the "haystack." " This ranges from distrac- tor needles which may have the same form (template) as the target concept to padding with repeated sentences. We use the repeated set of sentences "The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again." as our low-diversity padding to compare with context that is synthetically generated by an LLM, following Hsich et al. (2024) and Mohtashami & Jaggi (2023).
Symbolic Tasks We also experiment with purely symbolic (involving dictionary key-value or list retrieval) versions of our real tasks, since such tasks are believed to recruit similar model abilities as their natural language counterparts. For examplc. prior work has indicated that pre-training on code helps on Entity Tracking (Prakash et al., 2024)and that fine- tuning on a symbolic dictionary key-value retrieval task can provide greater benefits than even real data (Xiong et al., 2025). Additionally, RULER (Hsich et al., 2024) introduced variable assignment task for long-context value tracking. This latter task features expressions like "VAR X/ = 12345 Find all variables that are
MuSiQue The Si here are based on multi-hop knowledge graph relations. Like with MDQA, we create two simplifi- cation levels of concept expression by (1) keeping the real entities in the query and answer (high expression), and (2) replacing the real entities with 4-character symbolic entities (low expression), and constructing f. by prompting GPT- 4o-mini to write sentences or via template. We create two levels of context diversity by (I) prompting GPT-4 to write a paragraph containing the fact (high diversity), and (2)
VAR Y1= = 54321 assigned the value 12345." " that do not contain meaningful natural language. hence why we differentiate this category from natural language synthetic data.
4

---

Understanding Synthetic Context Extension via Retrieval Heads
Table 1. Performance (F1) of fine-tuning LLMSs on different synthetic data for the long-context retrieval and reasoning tasks. A moderate to large gap exists between the most performant synthetic context extension strategy (bold) and fine-tuning on real data. While careful construction synthetic data can help close the gap, there does not exist task agnostic general way of constructing synthetic datasets for extending LLMs' context window on long-context retrieval and reasoning tasks. f indicates that model trained on the bold synthetic dataset in the column outperforms a model trained on the indicated dataset with p < 0.05 according to paired bootstrap test. see that most gains of 2 0.02 accuracy statistically significant.
MDQA
MuSiQue
SummHay Cite Mistral
Concept Exp.
Context Div.
Concept Exp.
Context Div.
LJama3
Mistral
IJama3
Mistral
LJama3
High
High
0.31
0.201
0.371
0.22
High
High
0.70
0.281
High
Low
0.41*
0.23t
0,41
0.23
High
Low
0.61
0,281
Low
High
0.49
0.31
0.29t
0.21
Simplified
High
0.79
0.38
Low
Low
0.47!
0.24'
0.34
0.17
Simplified
Low
0.65
0.281
Symbolic
Symbolic
0.48
0.16!
0.32
0.11
Symbolic
Symbolic
0.54'
0.18!
Real Data (Full)
0.83
0.64
0.45
0.20
Real Data (Full)
0.81
0.40
Real Data (Limited)
0,80
0.59
0.32
0.16
Non-FT
0.45
0.12
0.22
0.03
Non-FT
0.40
0.07
padding each paragraph with repeated text (low diversity). The symbolic task, as demonstrated in Figure 4, consists of a list of dictionaries with 4-character identifier, keys and values. Queries are of the form "What is the PROPERTY_3 of the PROPERTY_2 2 of the PROPERTY_I 11 of DICTIONARY_ The answer is found by multi-hop traversal by accessing subsequent dictionary names associated with the specified properties.
SummHay there is a 2-4% gap between the best synthetic data and real data on Llama 3, and MDQA there is a much larger gap at 33%.
Careful construction of synthetic data can help close a lot of the gap by varying the level of concept expression and con- text diversity beyond the symbolic synthetic dataset. How- ever, the effective way of constructing synthetic data for training is very task-specific and can even be counter- intuitive: there does not exist single construction strat- egy that achieves the best performance across tasks, and sometimes a more "realistic" synthetic dataset can even underperform the more "synthetic" counterparts.
SummHay Citation We derive the f: from the insights in one of two ways. (1) We prompt GPT-4o-mini to rephrase the insights to create the query, and then prompt again to split rephrased insights into multiple sentences to place into the context (yielding multiple fi per insight) (high expression); and (2) We prompt GPT-4o-mini to simplify the insights to create the query, and split each simplified insight into multiple sentences to place into the context (low expression). We create two levels of context diversity by (1) padding each document with distractor insights from the same topic, (high diversity) and (2) padding each document with repeated text (low diversity). The symbolic task, as demonstrated in Figure 4, consists of lists containing 180 random 4-character strings, where the query is 4-character string that appears in two different lists.
These results show a complex picture of fine-tuning LLMs with synthetic data for long-context tasks: the downstream performance cannot be simply "predicted" by how the syn- thetic training dataset is constructed. To interpret the success and failure of synthetic data for training, a more fine-grained explanation is needed beyond some general, task-agnostic data construction desiderata.
4. Retrieval Heads are Necessary for Context Extension
One of the key features of our tasks is the need for retriev- ing needles ft embedded in a long context. Work from the mechanistic interpretability literature has shown that some attention heads in pre-trained (Olsson et al., 2022; Lieberum et al., 2023) fine-tuned (Panigrahi et al., 2023; Yin et al., 2024) transformers specialize in retrieving and synthesizing
3.3. Results
Table shows the performance (F1 scores) of fine-tuning LLMs on different synthetic datasets on the given long- context tasks. We first note that across datasets, fine-tuning on synthetic datasets still falls short compared with fine- tuning on real data, indicating the complexity of the eval- uated long-contex, tasks.* For instance, on MuSiQue and
the results of (Xiong et al., 2025) are obtained on 4K context rather than 32K and the models are fine-tuned with fewer training examples.
'Particularly on MDQA, we note that such observation is very different from the in (Xiong et al., 2025) that finds fine-tuning synthetic data to more effective than real data, We note that
5

---

Understanding Synthetic Context Extension via Retrieval Heads
MDQA p 0.6360
MuSiQue p= 0.7500
SummHay Citation
p 0.7714
0.8
0.450
0.80
Liama Naama--8B-Iin 3-8B-Instruct
0.425
0.75
0.7
0.400
0.70
-
0.5
0.350
0.65
0.5
0.325
0.60
0.300
0.55
Symbolic
0.50
0.85
0.90
0.95
1.00
0.70
0.75
0.80
0.85
0.90
0.95
1.00
0.86
0.89
0.90
0.92
0.94
0.96
0.98
1.00
High
p = 0.5833
p= 0.6071
p
0.9411
Context: High
0.6
0.300
0.40
Contoxt: Low
Mistral-7 Mistral-70-Insti 1o--7-nnsstuct-vo
0.275
0.35
0.5
0,250
0.225
0.30
0.200
0.175
0.150
0.2
0.125
0.30
0.92
0.94
0.95
0.98
1.00
1.0
Retrieval Score Cosine Sim.
Retrieval Score Cosine Sim.
Insight Score Cosine Sim.
Figure 3. Cosine similarity between the retrieval scores on real datasets (R, R) vs. their synthetic versions, and Spearman correlation for each setting. We use multiple umited-relation datasets for MDQA, as described Appendix C.
information from the context in principled ways.² Notably, recent work (Wu et al., 2025) indicates that there exists a special, intrinsic set of attention heads in pre-trained trans- formers that attend to relevant information fi in long context C. given query q and copy it to the output y. Wu et al. (2025) dub them as retrieval heads.
a fine-tuned model M', we evaluate on a dataset D {(C (C",""" y²)} where the answer y" needs to be identified from some needles f* in C* and copied to the model output ÿ. When M generates output token w E ů: we examine whether not an attention head places the most attention probability mass on the same token in the answer span y" in the context. If so, we consider the token w to be retrieved by the attention head. Given an evaluation example (C*, ","""""") """,."" let G, = {wn}} be the set of all tokens w that are retrieved by head h during decoding. We define the retrieval score Sn, for head h on single example
Given the nature of our task, we analyze these attention heads as a proxy for the subnetworks being recruited and learned during fine-tunie tne-tuning with synthetic data. Our core hy- pothesis is that can attribute the performance of synthetic context extension to how well the models learn adapt the attention heads relevant to retrieving and using information from long context, as indicated by the retrieval scores of attention heads. Building on prior work, we extend identifi- cation of retrieval heads to multi-hop settings in MuSiQue and SummHay Citation.
S₁ GB ng" 1y*
(I)
Note that in the SummHay Citation task, the model is prompted to identify the numerical IDs of the documents (e.g. "[3]") that contain the given query insight q. In this case, we find it more useful look at the attention heads that retrieve tokens from the insight needles f* that contain information relevant to q rather than retrieving tokens from the answer y². Note that there are far more tokens in the correct insight needles f than in the answer y here. Thus, the insight score for a single example is is defined as:
4.1. Detecting Retrieval Heads
Following Wu et al. (2025), we detect retrieval heads by computing retrieval scores. To compare across fine-tuned models, we consider any attention head with a positive retrieval score to be retrieval head, and later compute cosine similarity to account for the strength of scores. Given
'For example, (Prakash al., 2024) identifies sparse of heads that are responsible for retrieving and transmitting the positional information of objects from the context in the entity tracking task.
S = 1 [[ n > 0
For each head, we average scores over all evaluation exam- ples from D* to yield the final score.
6

---

Understanding Synthetic Context Extension via Retrieval Heads
Table 2. Cosine similarity of real dataset retrieval scores (+ SummHay insight scores) across tasks.
M
MuSiQue
SummHay Retrieval Mistral
SummHay Insight
Llama3
Mistral
Llama 3
Mistral
Llama3
Llama3
Mistral
MDQA
1.00
1.00
0.84
0.87
0.44
0.74
0.15
0.26
MuSiQue
0.84
0.87
1.00
1.00
0.59
0.69
0.28
0.20
SummHay Retrieval
0.44
0.74
0.59
0.69
1.00
1.00
0.08
0.07
SummHay Insight
0.15
0.26
0.28
0.20
0.08
0.07
1.00
1.00
Given a long-context task T, we detect a set of retrieval heads Heel of the models fine-tuned with real data Dy on an evaluation set of real data. For each model M fine- tuned with synthetic data DT, we detect set of retrieval heads Hyynth on an evaluation set of the corresponding syn- theric data. Hsynth reflects how synthetic context extension enables models to learn modules specialized in retrieving information from synthetic long-context data. Following Wu et al. (2025), we confirm that masking retrieval heads leads to an outsized drop in performance (Figure 9 in Ap- pendix F). The following sections will examine how Haynth explains synthetic data transferability to real long-context data.
task for MuSiQue and SummHay Citation. This holds more strongly for Llama-3-8B-Inse than for Mistral- 7B Instruuctt--O..
To account for score magnitude, we examine the relation- ship between the cosine similarity of vectorized retrieval scores with downstream task performance." Figure 3 shows surprisingly strong relationship here. When synthetic data does not induce retrieval heads matching the real task, per- formance is low.
This, retrieval heads are important indicators for whether synthetic dataset targets the desired real data reasoning ability. We view this result understanding of the mechanisms for target task can allow us to design synthetic tasks that induce the same behavior as required in the real task. However, targeting similar components is not enough; at the same level of similarity, we still observe a wide range of performances, as we will discuss in Section 4.5.
demonstration of how better
4.2. Case Study
We start with a case study of training Llama-3-8B 3--------B----ete on synthetic data for MuSiQue in Figure 1. Highlights show the retrieval score for each head at each layer. The model trained on the real data achieves an FI score of 0.45 on the evaluation set, and has 129 attention heads which receive a positive retrieval score. Notably, the mod- els trained on synthetic data (both realistic and symbolic) achieve lower F1 (0.41 and 0.33 respectively) while ex- hibiting far fewer retrieval-scoring attention heads (112 and 74 heads respectively). The real data retrieval heads have high recall (0.76 and 0.82) against the synthetic data heads indicating when the synthetic data induces fewer retrieval heads, they tend to be subsets of the real attention heads (Appendix E, Table 6), although this relationship is weaker on MDQA and SummHay Citation. 6
4.4. Retrieval Heads Across Tasks
We ask whether all tasks are leveraging the same set of re- trieval heads. Table shows cosine similarity of linearized retrieval scores between tasks. The single-hop and and multi-hop extractive QA tasks, MDQA and MuSiQue, have the highest cosine similarity (Llama-3-8B-Instruss M -B--nsttuuct-OO. 0.87). However, there is much lower similarity between the QA tasks and the SummHay Citation Retrieval Heads, and the least similarity with SummHay Insight Heads. 8 Comparing to Figure 3. we find that our real tasks have relatively high cosine similarity (> 0.66) with their synthetic versions, with the exception of the purely symbolic mhained-dictionary-ooddd citation tasks. This suggests that there are task-specific subsets of retrieval heads, either activated based on rea- soning ability or token diversity; we leave this for future
0.84;
4.3. Real Task Performance and Retrieval Heads
and list-
As noted previously, when the synthetic data induces fewer retrieval heads, they tend to be from the same population of retrieval heads active on the real data. Following this for each synthetic dataset, calculate the recall of non-zero scoring attention heads against the real dataset (first col- umn of Tables 5-10 in Appendix E). As shown in Table 4, we find that this is strongly correlated with F1 on the real
"We find it effective to directly match attention heads by index even when models are fine-tuned different datasets. Visualiza- tion in Appendix D supports this. "SummHay Retrieval Heads attend to the final answer (doc- ument number). whereas SummHay Insight Heads attend to the insight text within the document.
"We present full retrieval head counts and pairwise recall results in Appendix
7

---

Understanding Synthetic Context Extension via Retrieval Heads
in Appendix G.
Table 3. Results on I. Insttrue after patching retrieval heads that comprise the complement and intersection between the real and synthetic data versions, compared to random retrieval heads and original performance. Best patching Fl is underlined, and best F1 in row is bolded. indicates that the patched model outperforms the Orig. performance with p < 0.05 according to paired bootstrap test. Patching the intersection outperforms both random and complement heads.
Our results in Tables 3 and 11 show that patching intersec- tion heads outperforms patching both random and comple- ment heads. The improvement is the greatest for synthetic tasks with the lowest performance on the corresponding real task, and negligible or negative for the best synthetic tasks. The efficacy of patching Hinter indicates that while a synthetic dataset may target the necessary retrieval heads for the real task, they are insufficient in learning how to best utilize the required model components. 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.
Data Variant
Tisk
N
Compl.
Inter.
Rand.
Orig.
Concept
Context
Real
Real
0.82
Limited
Real
68
0.82
0.83
0.82
0.80
MDQA
Low
High
61
0.651
0.661
0.54
0.49
Symb.
Symb.
71
0.43
0.731
0.50
0.48
Low
Low
74
0.707
0.711
0.53
0.47
High
Low
74
0.63!
0.51
0.70
High
High
60
0.521
0.591
0.26
0.37
Real
Real
0.45
5. Related Work
High
Low
71
0.38
0.41
0.33
MuSiQue
High
High
71
0.33
0.29
0.25
0.37
Law
Law
61
0.417
0.33
0.33
0.34
Prior work has shown that benchmarking or training LLMs on synthetic data can reveal or obtain capabilities that can be transferred and generalized to real tasks, especially in settings where numan-annotaaio data is hard to obtain such as long-context tasks. For this purpose, synthetic data commonly used and believed to represent a simple reduction of the kinds of abilities employed in linguistically complex settings. The Needle-In-A- Aeedle-In-AHHaysttckk (NIAH) introduced by (Kamradt, 2023) involves placing a needle statement at a random position within a haystack consisting of unrelated essay text. Subsequent work (Hsieh et al., 2024; Li et al., 2024) has expanded this task to multi-value retrieval and used simple templated needle sentences to include distractor needles in the context. (Hsieh et al., 2024) additionally parameterized its test suite by the diversity of the input context and the target value type. These benchmarks are designed in part to expose shortcomings in a model's 's ability to utilize information throughout its context, known as the lost-in cost-in-the-me the-middle effect (Liu et al., 2024a). To address these shortcomings, prior work has synthesized more realistic long-context data based off of seed corpora to shore up the abilities of models to use their entire pretraining context (An et al., 2024) and further extend up to 1M context (Wang et al., 2024).
Symb.
Symb,
55
0.331
0.36t
0.35*
0.32
Limited
Real
53
0.34
0.34
0.31
0.32
Low
High
58
0.371
0.19
0.27
0.29
Real
Real
0.81
Simpl.
High
27
0.73
0.74
0.73
0.79
Summlly
High
High
19
0.72
0.75
0.79
0.70
Simpl.
Low
26
0.66
0.70
0.62
0.65
High
Low
21
0.57
0.68
0.67
Symb.
Symb.
26
0.53
0.60
0.48
0.54
investigation.
4.5. Synthetic Data Affects Required Model Components Less Effectively
Given different datasets of the same conceptual reasoning and retrieval task, it is peculiar that fine-tuning some datasets results in fewer retrieval heads, and that synthetic datasets can target subsets of the retrieval heads used for real data. Do the retrieval heads common to all datasets better capture the core capability required for the task? For the common attention heads, do models learn a better way of updating them from the real data than the synthetic data? To investigate these, we follow Prakash et al. (2024) to perform cross-model activation patching of retrieval heads in the intersection and complement between the real dataset and the synthetic datasets. Specifically, given the set of retrieval scoring attention heads on the real data, Hreat, and the set of retrieval scoring heads on a synthetic dataset, Hsyath we take the complement compl = Heal \ HHsyuth and the intersection Humer Heeal n Hsymth. For a fair comparison, we sample Wheads = min() Hommpl, (Fieer)) )) without replacement from both sets. Additionally we compare with Theads randomly sampled attention heads. For each set, we patch activations from the model trained on the real data to the model trained on the synthetic data. Implementation details can be found
Leveraging the potential generalizability of synthetic data, a line of work in interpretability literature generates synthetic data perform controlled experiments to probe the inner workings of L.MMs. For example, (Kim & Schuster, 2023) shows that synthetic version of entity tracking can be used to mechanistically understand how fine-tuning enhances cxisting capabilities of pre- pre-trained LLMs via mechanistic intervention techniques, and (Kim al., 2024) shows that the transformer circuit responsible for syllogistic reasoning in LLMs can be identified by evaluating on synthetic logical statements, However, there is a lack of understanding of
8

---

Understanding Synthetic Context Extension via Retrieval Heads
when and how the mechanism discovered from synthetic tasks generalizes to real-world capabiliits
the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Langwage Technologies, pp. 3554-3565, Online, June 2021. Association for Computational Linguisties. doi: 00.1665311/2221..acc/mm.....
Our work bridges these directions by providing mechanis- tic explanations for the transferability of synthetic context extension while motivating the pursuit of better usage of synthetic data to evaluate, enhance, and understand the ca- pabilities of LLMs.
URL. https:// a aclanthology. wtgp nwaclannccaaa. maai.. 278.
An. S., Ma, Z., Lin, Z., Zheng, N., Lou, J.-G., and Chen, W. Make your LLM fully utilize the context. In The Thirty- eighth Annual Conference on Neural Information Pro- cessing Systems, 2024. URL httttt:://opppeeeeee (tt.ffrrmmnn--YTEEEX..
6. Conclusion
In this paper, we investigated the relationship between the nature of synthetic data for synthetic context extension and performance on downstream tasks. Different synthetic datasets give widely varying performance, partially because of the different numbers of retrieval heads they induce in a model. We showed that these heads are causally connected to the performance, and that these heads are necessary (but not sufficient) for strong downstream model. We believe this work paves the way for further mechanistic understand- ing of long context behavior and the ways in which synthetic data induces new capabilities in language models.
Chandradevan, R., Dhole, K., and Agichtein, E. DUQ- Gen: Effective unsupervised domain adaptation of neu- ral rankers by diversifying synthetic query generation. In Duh, K., Gomez, H., and Bethard, S. (eds.), Pro- ceedings of the 2024 Conference of the North Ameri- can Chapter of the Association for Computational Lin- guistics: Human Language Technologies (Volume I: Long Papers), pp. 7437-7451, Mexico City, Mexico, June 2024. Association for Computational Linguistics. doi: 0..18653/v1/20244 www.a-loac.llaaaaae.ee. URL https:// aclanthology org/ / 2024. naacl- 413.
Acknowledgments
Chaudhary, A., Raman, K., and Bendersky, M. It's all relative! a synthetic query generation approach for improving zero-shot relevance prediction. In Duh, K., Gomez, H., and Bethard, S. (eds.), Findings of the Asso- ciation for Computational Linguisties: NAACL 2024, pp. 1645-1664, Mexico City, Mexico, June 2024. Association for Computational Linguistics. doi: 10. 10.18653/v1/2022 findings findings-naacl. 107. URL https: / ////accaantt aclant 22g/2024.2 f wnniinnnnnnnaacccc 107.
We thank the members of the TAUR lab for their helpful suggestions and discussions throughout this project. We also thank Xi Ye for providing feedback on the draft. This work was partially supported by NSF CAREER Award IIS- 2145280, the NSF AI Institute for Foundations of Machine Learning (IFML), the Sloan Foundation via Sloan Re- search Fellowship, and grant from Open Philanthropy. This research has been supported by computing support the Vista GPU Cluster through the Center for Generative AI (CGAI) and the Texas Advanced Computing Center (TACC) at the University of Texas at Austin.
Chen, J., Qadri, R., Wen, Y., Jain, N., Kirchenbauer, J., Zhou, T. and Goldstein, T. GenQA: Gener- ating millions of instructions from a handful of prompts. ArXiv, abs/2406.10323, 2024a. URL https: / //api.ssmaaaaaaee.. www.eeettnnnttotaeeee wtt.:ooooooo.oooooooouoocutt 270560271.
Impact Statement
This paper presents work whose goal is to enable better train- ing of long-con non-contxxt large language models and better under- standing of the role of synthetic data in LLM training. There are many potential impacts of better long-context models and training schemes, but none of which are uniquely en- abled by the methods and analysis we present here.
Chen, J., Zhang, Y., Wang, B., Zhao, W. X., Wen, J.- R., and Chen, W. Unveiling the flaws: Exploring im- perfections in synthetic data and mitigation strategies for large language models. ArXiv, abs/2406. 12397, 2024b. URL https: httpppp.s..o. www.sssemsstttiteeome.. org azg/CorpusID:: :2056622:27778888
References
Chen, S., Wong, S., Chen, L., and Tian, Y. Ex- tending context window of large language models via positional interpolation. ArXiv, abs/2306.15599 2023. URL https: semanticscholar. org/ erg/CorpusID:25522222
Agarwal, O., Ge, H., Shakeri, S., and Al-Rfou, R. Knowledge graph based synthetic corpus generation for knowledge-enhanced language model pre-training- In Toutanova, K. Rumshisky, A., Zettlemoyer, L., Hakkani- Tur, D., Beltagy, I., Bethard, S., Cotterell, R., Chakraborty, T., and Zhou, Y. (eds.), Proceedings of
Divekar, A. and Durrett, G. SymthesizRR: Generating di- verse datasets with retrieval augmentation. In Al-Onaizan,
9

---

Understanding Synthetic Context Extension via Retrieval Heads
Y., Bansal, M., and Chen, Y.-N. (eds.), Proceedings of the 2024 Conference on Empirical Methods in Nat- ural Language Processing, pp. 19200-19227, Miami, Florida, USA, November 2024. Association for Compu- tational Linguisties. doi: 10116655/11/2224.cmmmmmmaaa 1071. URL ht https: http://aaaaattoloooooo/////////a emn fomlp-mas in. 1071/.
Hu, E. J., yelong shen, Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, I.., and Chen, W. LoRA: Low-rank adap- tation of large language models. In International Confer- ence on Learning Representations, 2022. https: / //openrevieeee ew. net/ forum? wtwp:mmmiinnneeeeeeeeeee.
Izacard, G., Caron, M., Hosseini, L., Riedel, S., Bojanowski, P., Joulin, A., and Grave, E. Unsupervised dense in- formation retrieval with contrastive learning. Transac- tions on Machine Learning Research, 2022. ISSN 2835- 8856. URL https::/p//........ http:///pppeeeeeeeeee/eeeee/eeep / id-j px111xxib0.
Du, X., Shao, J., and Cardie, C. Learning to ask: Neural question generation for reading comprehension. In Barzi- lay, R. and Kan, M.-Y. (eds.), Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (V (Volume I: Long Papers), pp. 1342-1352, Vancouver, Canada, July 2017. Association for Compu- tational Linguistics. doi: 10. 121188//1//117-1222 URL ht https: / /aaaaaaaannntloooeoeeeeeo org/ 1123.
wtt.rrrrewweee//mmmm forum?
Jiang, A. Q.. Sablayrolles, A., Mensch, A., Bamford, C., Chaplot, D. S., de Las Casas, D., Bressand, F., Lengyel, G., Lample, G., Saulnier, L., Lavaud, L. R., Lachaux, M.- A., Stock, P., Scao, T. L.,, Lavril, T., Wang. T., Lacroix, T., and Sayed, W. E. Mistral 7B. ArXiv, abs/2310.06825,, 2023. URL https: //ap://apppp api. www.seeeescett.oiee To:
Elsahar, H., Vougiouklis, P., Remaci, A., Gravier, C., Hare, J., Laforest, F., and Simperl, E. T-REx: A large scale alignment of natural language with knowledge base triples. In Calzolari, N., Choukri, K., Cieri, C., Declerck, T., Goggi, S., Hasida, K., Isahara, H., Maegaard, B., Mar- iani, J., Mazo, H., Moreno, A., Odijk, J., Piperidis, S., and Tokunaga, T. (eds.), Proceedings of the Eleventh In- temational Conference on Language Resources and Eval- uation (LREC 2018), Miyazaki, Japan, May 2018. Eu- ropean Language Resources Association (ELRA). URL ht https: / /ccl anthology. org/ 111-1111111 -18-1144.
Josifoski, M., Sakota, M., Peyrard, M., and West, R. Ex- ploiting asymmetry for synthetic training data genera- tion: SynthlE and the case of information extraction. In Bouamor, H., Pino, J., and Bali, K. (eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 1555-1574, Singapore, De- cember 2023. Association for Computational Linguistics. doi: 100.18653/11/2223..oo / //aclanthologe hww.hhhlllllo....22.....e emnlp immiin.... 96.
main.96. URL https:
Fu, Y., Panda, R., Niu, X., Yue, X., Hajishirzi, H., Kim, Y., and Peng, H. Data engineering for scaling language models to 128k context. In Forty-first International Conference on Machine Learning, 2024. URL https: //openreview
Kamradt, G. Needle in a haystack pressure testing L.LMs, 2023. URL https://github. com/ gkamradt/ LLMTest_ Seirest_NeeddeennAHaattt.... ttree/ aarremmiii.. mmain.
ieriidddaaee TaAgeo71Uh.
Kim, G., Valentino, M., and Freitas, A. A Mecha- nistie Interpretation of Syllogistic Reasoning in Auto- Regressive Language Models. ArXiv, hts/2408.08599) 2024. URL /tttps://api... sewanteeeenniiiiii icstholar.....
Geva, M., Schuster, R., Berant, J., and Levy, O. Transformer feed-forward layers are key-value memories. In Moens, M.-F,, Huang. X., Specia, L., and Yih, S. W.-t. (eds.), Pro- ceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 5484-5495, On- line and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10. 10.18653/w1/2021..oo .mmnlp www.emmmmmmmmmm.....mm.m. main.446. URL https:// ac selanthology.. orgg2021111222... 11t11..mnnlpp. momnlp-mainn 1.6
Kim, N. and Schuster, S. Entity tracking in language models. In Rogers, A., Boyd-Graber, J., and Okazaki, N. (eds.), Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 3835-3855, Toronto, Canada, July 2023. Association for Computational Linguistics. URL https:
Gradient. Scaling Rotational Embeddings for Long-Context Language Models, 2024. https:// aalitaiiibaa ai/blog/
doi: 10. 10.16653W112000..dd1llo.g..1 / //aclannthologg..rrg///// http::oooooooo...ooooooo/ee acl-oog 213/.
URL
for- poor-long-ar
www.akowwki.. www.::uww..ooooaaii..eeeeeeeeeed O., Collins, M.. Parikh, A., Alberti, C., Epstein, D., Polosukhin, I., De- vlin, J., Lee, K., Toutanova, K., Jones, L.,, Kelcey, M., Chang, M.-W., Dai. A. M., Uszkoreit, J., Le, Q. and Petrov, S. Natural questions: A benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:452 466, 2019. doi: 10.
Hsieh, C.PP., S., Kriman, S., Acharya, S., Rekesh, D., Jia, F., and Ginsburg. B. RULER: What's 's the real context size of your long-contex language models? In First Conference on Language Modeling, 2024. URL https: / //openreview wwt.riiww.eeeeeeeeeeeeeeeeeeeeee wwt./forufm/eeee d= IoSbe
10

---

Understanding Synthetic Context Extension via Retrieval Heads
1162/aac URL httppss//aalaattholllg//// org/219-10255
Linguisties. URL hwwph:htttteeeeeeeeee bttttp:://accattte htttps://acllnnholooyy.oooo org/ 2024. 1124.ac1-1ong.15
Laban, P., Fabbri, A., Xiong, C., and Wu, C.-S. Summary of haystack: A challenge to long- long-context LLMs and RAG systems. In Al-Onaizan, Y, Bansal, M., and Chen, Y.-N. (eds.). Proceedings of the 2024 Conference on Empiri- cal Methods Natural Language Processing, pp. 9885- 9903, Miami, Florida, USA, November 2024. Association for Computational Linguistics. doi: 10. 10.18653/w1/2024. emnlp-main.552. URL https:/ /aclaattaaoooooyoo..caac org/2024
Mangrulkar, S., Gugger, S., Debut, I.., Belkada, Y., Paul, S., and Bossan, B. PEFT: State-of-the-a Parameter- Efficient Fine-Tuning methods. https: /eet::/ggthttt httpbbuggingffcceeeeeeee 2022.
Mohtashami, A. and Jaggi, M. Random- infi- nite context length for transformers. In Thirty-seventh Conference Neural Information Processing Systems, 2023. URL https:: /ttp://///eeeeeeeeeeeeeee net/ forum? 7
Lee, K., Chang, M.-W., and Toutanova, K. Latent retrieval for weakly supervised open domain question answering. In Korhonen, A., Traum, D., and Mårquez, (eds.), Proceedings of the 57th Annual Meeting of the Asso- ciation for Computational Linguistics, pp. 6086-6096, Florence, Italy, July 2019. Association for Computa- tional Linguistics. doi: 10.18653/v1//19-16611 URL https:/ /acl antho logy. otw.g/orrg/191119990
Olsson, C., Elhage, N., Nanda, N., Joseph, N., Dassarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., Drain, D., Ganguli, D., Hodds, Z., Hernandez, D., Johnston, S., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T. B., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. In-context Learning and Induction Heads. ArXiv, abs/2209. 18895, 2022. URL https: //api. semanti www.ttca.c...r.......o.m....... org / /ctp//oopaD::2555222222
Li, M., Zhang, S., Liu, Y., and Chen, K. NeedleBench: Can LLMs do retrieval and reasoning in million context window? ArXiv, abs/2407.11963 2024. URL https: //arxiv. org/ abs 2aab/22407 11963.
Panigrahi, A., Saunshi, N., Zhao, H., and Arora, S. Task- specific skill localization in fine-tuned language mod- els. In Krause. A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), Proceedings of the 40th International Conference on Machine Learn- ing, volume 202 of Proceedings of Muchine Learn- ing Research, pp. 27011-27033, PMLR. 23-29 Jul 2023. UtL http:///c/ccccee../////./// / v202/ http//aangrrai2233 htmml
Lieberum, T., Rahtz, M., Kram'ar, ar, J., Irving, G., Shah, R., and Mikulik, V. Does circuit analysis inter- pretability scale? evidence from multiple choice capabilities in Chinchilla. ArXiv, abs/2307.09458, 2023. URL https://pp...mmmttt..t. www.ettisccooll........... org/ org/Corpus CorpusI 2599959992999
wwwpeeeedinngmm.r....... mlr.ppress.
Liu, N. F., Lin, K., Hewitt, J., Paranjape, A., Bevilacqua, M., Petroni, F., and Liang, P. Lost in the middle: How language models use long contexts. Transactions of the Association for Computational Linguistics, 12:157-173, 2024a. doi: 10. 10.1162/tacLa..00333 URL https: ac aclanthology. hwrpllloyyooro////// 2024. tacl-1.9.
Prakash, N., Shaham, T. R., Haklay, T., Belinkov, Y., and Bau, D. Fine-tuning enhances existing mechanisms: A case study on entity tracking. In The Twelfth In- ternational Conference on Learning Representations, 2024. URL https: battps:://penreee/// /pppreeieeeeeeeeeeeeeeeeeee net/ forum? id=8sKcAWOF2D.
Liu, R. Wei, J., Liu, F., Si, C., Zhang, Y., Rao, J., Zheng, S., Peng, D., Yang, D., Zhou, D., and Dai, A. M. Best practices and lessons learned on syn- thetic data for language models. ArXiv, abs/2404.07503, 2024b. URL https: http:p/://pp..emmntiissetete www.nticscoolll......... org/CorppsII::20 25904285.
Saparov, A. and He, H. Language Models Are Greedy Reasoners: A Systematic Formal Analysis of Chain- of-Thought. In The Eleventh International Conference on Learning Representations, 2023. URL /ttps::/ openreview net woworrmmidgggFBB-B.
Stolfo, A., Belinkov, Y., and Sachan, M. A Mechanis- tic Interpretation of Arithmetic Reasoning in Language Models using Causal Mediation Analysis. In Bouamor, H., Pino, J., and Bali, K. (eds.), Proceedings of the 2023 Conference Empirical Methods in Natural Lan- guage Processing, pp. 7035-7052, Singapore, December 2023. Association for Computational Linguistics. doi: main.435. URL https:// a aclanthol hogganhhoogy...oooooooooo ogy .rrg http2222..mmmppmmmme emnlp-ma 435.
Lu, Z., Zhou, A., Ren, H., Wang, K., Shi, W., Pan, J., Zhan, M., and Li, H. MathGenie: Generating synthetic data with question back toekssrrnslttoon for enhancing mathe- matical reasoning of LLMs.. In Ku, L.-W., Martins, A., and Srikumar, V. (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computational Linguis- tics (Volume I: Long Papers), pp. 2732-2747. Bangkok, Thailand, August 2024. Association for Computational
11116533/122023.com

---

Understanding Synthetic Context Extension via Retrieval Heads
Su, J.. Ahmed, M., Lu, Y., Pan, S., Bo, W., and Liu, Y. RoFormer: Enhanced transformer with Neurocomput- ing, 568:127063, 2024. ISSN 0925-2312. doi:
Xiong, Z.,, Papageorgion, V., Lee, K., and Papailiopou- los, D. From artificial needles to real haystacks: Im- proving retrieval capabilities in LLMs by finetuning on synthetic data. In The Thirteenth International Confer- ence on Learning Representations, 2025. URL https: / /ooenreview.wco net/ / forum? iidd 8m..dmmmm74k4kkeee
Rotary Position Embedding.
URL httpp:://www.sccieeeedddiie com/ occ.enee/atticllee/i////22222222
Xu, B., Wang. Q., Lyu, Y., Dai, D., Zhang, Y., and Mao, Z. S2ynRE: Two-stage self-training with synthetic data for low-resource relation extraction. In Rogers, A., Boyd- Graber, J., and Okazaki, N. (eds.), Proceedings of the 6 Annual Meeting of the Association for Computa- tional Linguistics (Volume 1: Long Papers), pp. 8186- 8207, Toronto, Canada, July 2023, Association for Com- putational Linguisties. doi: 10. 10..1853//1/2220 1/202/ wwwp:/oooooooogogoooe 2023.
Tang, L., Laban, P., and Durrett, G. MiniCheck: Efficient fact-checking of LLMs on grounding documents. In Al-Onaizan, Y., Bansal, M., and Chen, Y.-N. (eds.). Pro- ceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pp. 2818-8847, Miami, Florida, USA, November 2024. Association for Compu- tational Linguistics. doi: 10.18655/w122224ccmmmmmmmm 499. URL https:/ / /ttp///anntlllo.oorgggo/o/ 499/.
455. URL acl-leng. 4455.
Team, L. The Llama 3 herd of models, 2024. URL https: //arxiv. wwtg.bsg///oo.....oooo... 240/ 21783.
Xu, C., Sun, Q. Zheng, K., Geng, X., Zhao, P., Feng, J., Tao, C., Lin, Q., and Jiang, D. WizardLM: Empower- ing large pre-trained language models to follow com- plcx instructions. In The Twelfth International Confer- ence on Learning Representations, 2024. URL https: / //openeeveww net/ wor.mriddcumdnne 3ND.H.
Trivedi, H., Balasubramanian, N., Khot, T., and Sabharwal, A. MuSiQue: Multihop questions via single-hop ques- tion composition. Transactions of the Association for Computational Linguistics, 10:539-554, 2022, doi: 10. 1162/tacl URL https: htt.::pss://aaaaahooooo..o org/2022. tacll..1ccl.......
Xu, Z., Jiang, F., Niu, L., Deng, Y., Poovendran, R., Choi, Y., and Lin, B. Y. Magpie: Alignment data synthesis from scratch by prompting aligned LLMs with nothing, In The Thirteenth International Conference on Learning Representations, 2025. URL https: /ppnrevien..ee.e net / forum? (icuuummiiii www.dddd--nnnaannnn...
Veličković, P., Perivolaropoulos, C., Barbero, F., and Pas- canu, R. softmax is not enough (for sharp out-of- distribution), 2024. URL attps://arxivv org/ 2ttp/2410.001000
von Werra, L.. Belkada, Y., Tunstall, L., Beeching, E., Thrush, T., Lambert, N., Huang, S., Rasul, K., and Gallouédec, Q. TRL: Transformer reinforcement learn- ing. https://g/ /tiphb. 2020.
Yang, K., Klein, D., Celikyilmaz, A., Peng, N., and Tian, Y. RLCD: Reinforcement learning from con- trastive distillation for LM alignment. In The Twelfth International Conference on Learning Representations, 2024. URL https: http//ppnnreeeeeeeeeeeeee net.reeeieweneee/eeeeeee forum? id=v3XxtxWKi6
thub. wtt. face///t...
Wang, L., Yang, N., Zhang, X., Huang, X., and Wei, F. Bootstrap your context length, 2024.
Yin, F., Ye, X., and Durrett, G.
LoFiT: Localized fine-tuning on LLM representations. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024. URL https: /oppnrevieeeeeeeeeeeeeeeeee net/ forum? id=dfixFbECS2.
Wei, J., Hou, L., Lampinen, A., Chen, X., Huang, D., Tay, Y., Chen, X., Lu, Y., Zhou, D., Ma, T., and Le, Q. Symbol tuning improves in-context learning in language models. In Bouamor, H., Pino, J., and Bali, K. (eds.), Proceedings of the 2023 Conference Empirical Methods in Natural I Language Processing, pp. 968-979, Singapore, Decem- ber 2023, Association for Computational Linguistics. doi: 10.18653/v1/2023 wwwp:///.....m.pmmpmmeeeeeeeeeme www.e-mmin..m... main.61. URL https:// acl 1 antho ogy. , org/ 2023. emn p-ma i n. 61/.
Yu, A. W., Dohan, D., Le, Luong, T., Zhao, R., and Chen, K. Fast and accurate reading comprehension by combining self-attention and convolution. In In- ternational Conference on Learning Representations, 2018. URL https: 1 id=B14T1G-BB
httL:ttt//ppeeeeeeeee http:///pppeeeeeeeeeeeppe net fooum??
Wu, W., Wang. Y., Xiao, G., Peng, H., and Fu, Y. Retrieval head mechanistically explains long-contex, factuality. In The Thirteenth International Conference on Learning Representations, 2025. http:::: oppenreview. n / forum? i RpGGt11
Yue, X., Qu, X., Zhang, G., Fu, Y., Huang, W., Sun, H., Su, Y., and Chen, W. MAmmoTH Building math generalist models through hybrid instruction tuning. In The Tweifth International Conference on Learning Representations,
12

---

Understanding Synthetic Context Extension via Retrieval Heads
2024. URL httpetse//eeeeeereee..eeee// forum? id=yICIGs7701.
Zhao, L., Wei, T., Zeng, L., Cheng, C., Yang, L., Cheng, P. Wang, L.,, Li, C., Wu, X. G., Zhu, B., Gan, Y., Hu, R., Yan, S., Fang, H., and Zhou, Y. LongSky- work: A training recipe for efficiently extending context length in large language models. ArXiv, fbs/2406.006000 2024. URL https: /tpppppp.... http::::/app.aaaaaaaaa.a.. www.stessmstsiiiieeeeeeeeeei
org/CorrpsID::2000 270210030222220000 0363.
13

---

Understanding Synthetic Context Extension via Retrieval Heads
A. Synthetic Dataset Creation Prompts
A.1. MDQA
Given a training example MDQA data (c, q. y) E Dr, we first combine the query and the answer into a sentence and prompt GPT-4o-mini rephrase the sentence with the sentence paraphrasing prompt to make the needle f. Then, for the synthetic dataset with high context diversity, prompt GPT-4o-mini generate Wikipedia-style context paragraph with the context generation prompt.
Prompt A.1: MDQA Sentence Paraphrasing Prompt
System Prompt:
You are helpful Al assistant and you are good at creative writing.
Prompt:
Rewrite the following sentence Wikipedia style with additional details: feent ence]
Make sure that readers can correctly answer the following question by reading your rewritten sentence:
Question: iquestioo)
Answer:
lanswer)
Prompt A.2: MDQA Context Generation Prompt
System Prompt:
You helpful Al assistant and you are good at creative writing.
Prompt:
Please make up 100-word Wikipedia paragraph for the following fake entities: [entsty). Invent details about people, places, and work related each entity, and make sure all details generated paragraph. After making up the paragraph, please generate a hwwp://wwwwwwwwwhhh... (1) is related to the given fake entities:
related any real-world entities. Give short, meaningful title your
question that:
(2) one can use the paragraph to correctly infer the answer within one or two words;
(3) is not direct copy of a sentence from the paragraph. Please also include the gold answer the generated question.
Please give your response in the format:
Title: [title]
Text: [text]
Question: [question]
Answer:[answer]
A.2. MuSiQue
Prompt A.3: MDQA Sentence Paraphrasing Prompt
Prompt:
Please make up single sentence for each of the following fake entities in the style of wikipedia article.
took_entiiiee
Please give your response in the format:
Title: [title]
Text: [text]
Prompt A.4: MuSiQue Context Generation Prompt
Prompt:
Please make up 5-sentence wikipedia paragraph the following entities. Invent details about people, places, work related to each entity.
itake_entities
Please give your response in the format:
Title: [titlc]
Text: [text]
14

---

Understanding Synthetic Context Extension via Retrieval Heads
A.3. SummHay
Prompt A.5: SummHay Query Insight (Concept Expression High) Prompt
Prompt: Please rephrase the sentence: (text]
Prompt A.6: SummHay Query Insight (Concept Expression Simplified) Prompt
Prompt:
Please simplify and shorten the following sentence. Remove details:
centence)
Prompt A.7: SummHay Citation Needle Prompt
Prompt: "Please break up the following sentence into multiple sentences: (text]
B. Training Prompts
Prompt B.I: MDQA and MuSiQue Training Prompt
Prompt:
The following are given passages.
[context]
Answer the question based the given passages. Only give the and do not output any other words.
Question: (question)
Answer:
Prompt B.2: SummHay Citation Training Prompt
Prompt:
following are given documents.
[context]
For the given statement, identify the documents that contain the information by citing the numbers associated with those documents in brackets. For example, if the information in the statement is only found in Document 3, then respond with "[3]". the information contained in both Document 3 and Document then respond with "[3][7]". Only output the answer and do output any other words. (statement)
Statement:
Answer:
C. Additional Data and Training Details
C.1. Data
We use 1400 examples for training MDQA models, 400 examples for MuSiQue models, and 400 examples for SummHay Citation models. Each dataset is partitioned in to a 90/10 train/validation split. We use the validation split to calculate retrieval and insight scores.
15

---

Understanding Synthetic Context Extension via Retrieval Heads
MDQA Example
Context:
Document I: (Title: Don Quixote (Teno)) portion Don Quixote and his horse are visible. The horse appears to be charging forward of the stone with his head raised, mouth open, and hooves kicking. The left foot of the horse formed, intentionally, by Teno. In Don Quixote's hand is a lance of steel. Both figures are loosely modeled and the figures and stone rest oval base measuring which was into three pieces for transport by ship the United States, An inscription on the sculpture reads: King Juan Carlos and Queen Sofía presented the sculpture June 3, 1976, on
Document 10: (Title: Rocinante) Rocinante is Don Quixote's ' horse the novel Don Quixote by Miguel de Cervantes, In many ways, Rocinante not only Don Quixote's is horse, but also his double: like Don Quixote, he is awkward, past his prime, and engaged in task beyond his capacities.
Question:
what don quixote's horse's name
Answer:
Rocinante
SummHay Citation Example
Context:
Document [24]: ... Furthermore, the provision of U.S. dollars by global central banks increased, ensuring adequate liquidity within the international financial system. This measure illustrated the depth of the coordinated efforts among major financial institutions to stave off crises and maintain functional stability. The reverberations of these actions and their impacts on the markets are still unfolding...
Document [27]: ... Turning our gaze the realm global financial oversight, central banks are making coordinated efforts prevent liquidity crunch in the international financial system. Recognizing the importance of maintaining robust liquidity, global central banks have ramped up their provision of U.S. dollars, showcasing united front in ensuring financial stability. Central banks in Canada, Britain, Japan, Switzerland, and the eurozone have initiated daily currency swaps to ensure that banks operating within their jurisdictions have the necessary dollars to function smoothly. This strategy is aimed at providing stability and fostering confidence in the global banking system during uncertain economic times...
Statement:
Global central banks increased their provision of U.S. dollars to ensure adequate liquidity in the international financial system, demonstrating coordinated efforts to prevent liquidity crunch.
Answer:
[24][27]
Real Data (Limited) For MDQA and MuSiQue, we experiment with only training on subset of the relations involved in eval question hops.
On MDQA, create three variants: L, is the subset containing Who, When, and Where questions; L2 is the subset containing When and Where questions; Lg 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 L, results due to space constraints. Fine-tuning Llama-3-8P 3--------t-nntee on these datasets results in the following F1-scores for the target dataset: L1 = 0.80, L2 = 0.63, L3 = 0.65. Fine-tuning Mistral-7B. - on these datasets results in the following F1 -coorr for the target dataset: 0.59, L2 0.52, L3 0.57.
16

---

Understanding Synthetic Context Extension via Retrieval Heads
MuSiQue Symbolic Data
Context
BPUG (..., 'UCCA:: 'QUID', 'TZAM': 'DPWW, 'EJSN': 'TFF',
KVTJ (.... 'UCCA" 'SVV', 'EROG': 'FODR; 'TZAM': "XYTH,
FQDR (..., 'UQCA' 'EHQQ' 'UDPB': 'BPUG' 'EROG': "MMII,, ...)
Q: What is the TZAM of the UDPB of the ERQG of KVTJ?
A: XDPW
SummHay Citation Symbolic Data
Context
Document [16]: [..., 'SIWK' 'NGOW:; 'UXHQ', 'RBZE',
Document (27]: [.., 'DUTT', NGOW', 'LTYM', .P.PP....
Q: NGOW
A: [16][27]
Figure 4. Examples of symbolic data consruction for MuSiQue and SummHay Citation.
MuSiQue, we use the subset of linear 3-hop questions consisting solely of T-REx component questions (Elsahar et al., 2018), identified by ">>". 10.8% of MuSiQue linear 3-hop questions in the training set fit this cri- teria. Additionally, among all component question hops in the training set, 43.0% sourced from T-REX.
C.2. Symbolic Data Construction
See Figure 4 for examples.
C.3. Training
For fine-tuning, we use the Huggingface TRL. (von Werra et al., 2020) and PEFT (Mangrulkar al., 2022) libraries to fine-tune attention heads with LoRA (Hu al., 2022) (rank 8 and alpha 8) using a batch size of and 4 gradient accumulation steps.
We enable Flash Attention and DeepSpeed and use single NVIDIA H100 GPU (96GB) for each training run. We use greedy decoding in all evaluations.
D. Retrieval Score Heatmaps
Attention head retrieval scores for the real tasks are shown in Figure 5.
For each target real task, we present heatmaps comparing the real task retrieval scores to the synthetic dataset retrieval scores: MDQA in Figure 6, MuSiQue in Figure 7, and SummHay Citation in Figure 8.
E. Retrieval Head Recall
In Table 5, Table 6, and Table we examine the overlap between non-zero scoring attention heads on our target tasks and their synthetic versions after fine-tuning Biama----B-Instruue We find that on all tasks, the attention heads with non-zero retrieval scores the real data have high recall (2 0.76) against those identified the synthetic data. On MuSiQue and SummHay Citation, also observe strong relationship (Spearman R=0.75 and R=1.0 respectively) between the non-zero score attention head recall and F1 on the real task.
However, fine-tuning Mistral-7B- ---n----tttttt results in slightly different patterns, as shown in Table 8, Table 9, and Table 10. First, we see more scoring attention heads, which could be caused by the sliding window attention used in the architecture, which only enables subset of heads to any single position. Second, many of the synthetic datasets result in far
17

---

Understanding Synthetic Context Extension via Retrieval Heads
MDQA
MuSiQue
SummHay Citation
Inntruutt
layer [[ick
20
20
layer idx
20
0.50
15-
15-
15
0.05
20
25
15
30
20
25
0.00
idx
head
jdx
MDQA
MuS Que
SummHay Citation
0.75
in..instruct--i
layer [ex
20
layer lex
20
layer lex
20
0.50
15
5-
5
0.05
20
25
10
15
30
0
20
0.00
ldx
head
ldx
Figure 5. Retrieval scores for MDQA. MuSiQue, and Insight scores for SummHay Citation. Top Row: Llama-3-88 3-BB --8--nnstuct.. Bottom The y-axis indicates the layer index and the x-axis indicates the head index within the layer. We note that retrieval heads largely found in the last 2/3 layers of the model, as expected according their involvement in the "final step" of copying the correct answer to the output. By contrast, SummHay Citation insight heads are concentrated in the middle layers, indicative of their intermediate rolc. Within single layer, the specific important attention head indices were likely randomly primed during pretraining to be effectively adapted to the target task.
Row: Misstral7BB-Instruet--0
Figure 6. Retrieval for MDQA and its synthetic dataset versions. Top Row: Llama-3-88 -omma-3-BB-nnstruct.. Bottom Row: Mistral-7B. Instruct-vO.1. The y-axis indicates the layer index and the x-axis indicates the head index within the layer.
more non-zero scoring attention heads, pattern that we see across all tasks. On MuSiQue and SummHay Citation, we observe a slightly weaker relationship (Spearman R=0.40 and R=0.82 respectively) between the non-ze zero score attention head recall and Fl the real task.
18

---

Understanding Synthetic Context Extension via Retrieval Heads
lear
Tea
Figure 7. Retrieval scores for MuSiQue and its synthetic dataset versions. Top Row: Hama-33-B-Insttuu Bottom Row: Mistral-7B- Instruct-vO.1. The y-axis indicates the layer index and the x-axis indicates the head index within the layer.
hdt
200
Made 12
Figure 8. Insight scores for SummHay Citation and its synthetic dataset versions. Top Row: Llama-3-8B- Instruutt Bottom Row: The y-axis indicates the layer index and the x-axis indicates the head index within the layer.
Table Spearman correlation of synthetic data attention head recall with Fl on the real dataset, showing strong relationship.
Model Llama3
Mistral
MDQA
0.22
0.16
MuSiQue
0.75
0.40
SummHay Citation
1.00
0.82
F. Retrieval Head Masking
Figure shows the effect masking attention heads with the top-k retrieval (MDQA, MuSiQUE) or insight (SummHay Citation) scores. Compared to masking the same number of randomly chosen heads (over trials), masking top-k attention heads consistently results in a larger drop in performance than masking random attention heads.
19

---

Understanding Synthetic Context Extension via Retrieval Heads
Table Pairwise recall of Llama-3-8B- -i- attention heads with non-zero retrieval scores for MDQA synthetic datasets. Limited
datasets: L1 Who, When, Where; L2 = When, Where: L3 Who. Retrieval Head recall on the real dataset (first column) is weakly
correlated with Fl on the MDQA data (Spearman 0.22).
R.R
R.R
R.R (La)
R.R
H.HI
H.L
L,
L.L
S,S
# Heads
F1
(L1)
(L2)
0.79
0.84
0.88
0.85
0.78
0.81
0.83
0.87
157
0.82
R,R (L1)
0.76
1.00
0.90
0.86
0.81
0.69
0.81
0.80
0.85
151
0.80
R.R (L2)
0.66
0.74
1.00
0.78
0.71
0,63
0.71
0.70
0.74
124
0,63
R.R (La)
0.63
0.64
0.70
1.00
0.69
0.61
0.66
0.68
0.77
112
0.65
H.H
0.75
0.79
0.86
1.00
0.74
0.78
0.83
0.83
139
0.37
0.73
0.74
0.80
0.78
1.00
0.77
0.80
0.78
147
0.41
LHH
0.80
0.83
0.88
0.90
0.86
0.81
1.00
0,91
0.93
154
0.49
L.L.
0.67
0.68
0.72
0.77
0.76
0.69
0.75
1.00
0.76
127
0.47
S,S
0.64
0.66
0.69
0.79
0.69
0.61
0.70
0.69
1.00
116
0.48
Table Pairwise recall T lamm-33-8B-Instruuc attention heads with non-zero retrieval scores for MuSiQue synthetic datasets. We find that the attention heads identified the real dataset has high recall against all synthetic datasets (>0.76). Retrieval head recall on the real dataset (first column) is also strongly correlated with FI on the real MuSiQue data (Spearman R = 0.75).
R.R
R.R (L.)
H,H
H,I.
I,H
L,L
S,S
Heads
R.R
1.00
0.96
0.81
0.76
0.87
0.82
0.87
129
0.45
R.R (L)
0.41
1.00
0.50
0.41
0.52
0.49
0.42
55
0.32
H.H
0.59
0.85
1.00
0.63
0.70
0.65
0.63
94
0.37
H.L
0.66
0.84
0.76
1.00
0.81
0.78
0.71
112
0.41
L.H
0.45
0.64
0.50
0.48
1.00
0.65
0.53
67
0.29
L.L.
0.47
0.65
0.51
0.52
0.72
1.00
0,55
74
0.34
S.S
0.67
0.76
0.67
0.63
0.79
0.74
1.00
100
0.32
Table Pairwise recall of Llama-3-88 8----nntt attention heads with non-zero insight scores for SummHay Citation synthetic datasets.
Insight head recall on the real dataset (first column) is also strongly correlated with F1 on the real data (Spearman = 1.0)
R.R
H.H
H,L
L.H
L.L
S.S
# Heads
F1
R.R
1.00
0.77
0.94
0.66
0.78
0.87
48
0.81
H.H
0.85
1.00
1.00
0.75
0.82
0.87
53
0.70
H,I.
0.60
0.58
1.00
0.48
0.72
0.70
31
0.61
L.H
0.90
0.92
1.00
1.00
0.88
0.93
65
0.79
L.L
0.65
0.62
0,94
0.54
1.00
0.80
40
0.65
S,S
0.54
0.49
0.68
0.43
0.60
1.00
30
0.54
G. Retrieval Head Patching Details
We implemented retrieval head patching with Baukit." Given an example from the test set and a of attention heads to patch, we run forward pass with the model fine-tuned the real data and extract the attention output from the selected attention heads before being projected and concatenated back to the residual stream. Then, we the same example and run forward pass with the model fine-tuned on synthetic dataset. We replace the attention outputs of the aforementioned selected attention heads with the attention outputs extracted from the model fine-tuned on real data. Using the procedure described above, we patch the attention outputs the selected attention heads into the model fine-tuned on a synthetic dataset for all input tokens. We then use the patched inputs to generate and decode output tokens without patching any activations for the output tokens.
20

---

Understanding Synthetic Context Extension via Retrieval Heads
Table Pairwise recall of -istral-7B-Insstuct--h) attention heads with non-zero retrieval for MDQA synthetic datasets. Limited
datasets: L1 Who, When, Where: L2 = When, Where: L3 Who. Retrieval head recall on the real dataset (first column) is weakly
correlated with Fl on the real MDQA data (Spearman 0.16).
R.R
R.R
R.R (La)
R.R
H.HI
H.L
L,H
L.L
S,S
# Heads
F1
(L1)
(L2)
0.76
0.81
0.74
0.74
0.77
0.69
0.80
178
0.63
R,R (L.)
0.81
1.00
0.85
0.86
0.77
0.80
0.80
0.72
0.82
192
0.59
R.R (L2)
0.81
0.84
1.00
0,86
0.74
0.77
0.80
0.72
0.87
190
0,52
R.R (La)
0.74
0.72
0.73
1.00
0.67
0.71
0.72
0.63
0.74
161
0.57
H.H
0.86
0.84
0.81
0.86
0.83
0.83
0.76
0.84
208
0.20
0.89
0.86
0.93
0.85
1.00
0.86
0.81
0.87
212
0.22
LHH
0.88
0.85
0.86
0.92
0.82
0.83
1.00
0.78
0.85
205
0.31
L.L. 0.93
0.91
0.91
0.94
0.88
0.92
0.91
1.00
0.91
240
0.24
S.S
0.80
0.76
0.81
0.81
0.72
0.73
0.74
0.68
1.00
178
0.15
Table Pairwise recall of Mistral-7BB-nnsruut--o)) attention heads with non-zero retrieval scores for MuSiQue synthetic datasets. Retrieval Head recall the real dataset (first column) also moderately correlated with Fl the real MuSiQue data (Spearman R 0.40)
R.R
R.R (L)
H.H
H.L
L.H
L.L
S,S
Heads
R,R
1.00
0.66
0.56
0.55
0.63
0.66
0.62
111
0.31
R,R (L)
0.63
1.00
0.57
0.57
0.59
0.60
0.53
106
0.14
H.H
0.89
0.95
1.00
0.83
0.88
0.86
0.82
178
0.21
HLL
0.83
0.91
0.78
1.00
0.81
0.81
0.78
167
0.23
L.H
0.84
148
0.21
L.L
0.75
0.72
0.61
0.61
0.70
1.00
0.68
126
0.17
S.S
0.74
0.66
0.61
0.62
0.72
0.72
1.00
133
0.11
Table Pairwise recall
attention with non-zero insight scores SummHay Citation synthetic datasets. Insight Head recall on the real dataset (first column) is also strongly correlated with F1 on the real SammH Citation data (Spearman R -.82)
H.H
H.L
L.H
L.L
S.S
# Heads
F1
R.R
1,00
0.31
0.15
0.16
0.45
0.81
91
0.40
H.H
0.87
1.00
0.34
0.39
0.75
0.88
0.28
H.L
0.79
0.65
1.00
0.57
0.80
0.81
496
0,28
L.H
0.86
0.76
0.57
1.00
0.78
0.88
497
0.38
L.L
0.75
0.44
0.24
0.24
1.00
0.81
150
0.28
S,S
0.14
0.05
0.03
0.03
0.09
1.00
0.17
G.1. Nisttral7E--nssruuct-... Retrieval Head Patching
See Table 11.
G.2. Intersection and Complement Head Retrieval Scores
See Table 12.
H. Full Finetuning
In this section, present results on 1ine-Llama-3-B-IInsttee with fine-tuning of all LoRA modules, and demonstrate that we find similar conclusions.
Phtpps:/ /i thub. com/daavibba//aaiiii
21

---

Understanding Synthetic Context Extension via Retrieval Heads
MDQA
MuSiQue
SummHay Citation
Meta-Llama-3-8B-- 3-8---nstttt
SummHay Citation
Selection Top-k Rand.
Mistral-7B-Instrucct 1
1.0
0.8
0.6
0.4
02
10
20
30
60
50
10
20
30
40
50
10
20
30
40
50
# of Masked Attention Heads
# of Masked Attention Heads
# of Masked Attention Heads
Figure Top row: -lama-3-SB-Insst Instrucctt.. Bottom row: D Effect masking activations from attention heads (following Wu et al. (2025)) with the top-k highest retrieval (MDQA, MuSiQue) or insight (SummHay Citation) scores. We compare with masking the same number of randomly chosen heads, averaged over samples. Masking top-k attention heads consistently results a larger drop in performance than masking random attention heads.
MDQA
MuSiQue
SummHay Citation
0.75
1.88 Instruct
25
25
25
layer [idx
20
byyer ldx
20-
layer ldk
20
0.50
0.30
5-
5-
0.05
0
5
10
15
20
25
10 15
20
25
30
0
5
10
15
20
25
30
head idx
head idx
head idx
Figure 10. Llama-3-88- 8aae--3-----n-suuce (all LoRA modules): Retrieval scores for MDQA, MuSiQue, and Insight scores for SummHay Citation, after fine-tuning on each task. The y-axis indicates the layer index and x-axis indicates head index within the layer.
H.1. Synthetic Data Performance
See Table 13. We find that there are mostly small (<: 0.05) performance differences between fine-tuning only attention heads and all modules. Notable exceptions are found the SummHay Citation task, where the performance of the synthetic datasets increase up to +0.13 (High, High).
H.2. Retrieval Score Heatmaps
See Figure 10.
H.3. Retrieval Head Recall
In Table 14, Table 15, and Table 16, we find that there are generally fewer non-zero scoring attention heads on the synthetic tasks, compared to the real task. On MuSiQue, the non-zero attention heads tend to be subsets the those identified on the real task, as when only finettnnig attention modules.
22

---

Understanding Synthetic Context Extension via Retrieval Heads
Table 11. Results Chatral-B-Instrucct after patching heads that comprise the complement and intersection retrieval heads hetween the real and synthetic data versions, compared to random retrieval heads and original performance. The best patching F1 is underlined. the best each row bolded. indicates that the patched model outperforms Orig. performance with p 0.05 according a paired bootstrap test. Patching the intersection outperforms both random and complement heads.
Data Variant
Task
N
Compl.
Inter:
Rand.
Orig.
Concept
Context
Real
Real
0.63
Limited
Real
80
0.57
0.53
0.49
0.59
MDQA
Low
High
69
0.21
0.341
0.23
Low
Low
86
0.21
0.40
0.26
High
Low
0.31
0.16
High
High
80
0.21
0.261
0.18
0.20
Symb.
Symb.
70
0.01
0.02
0.02
0.15
Real
Real
0.31
High
Low
92
0.23
0.26!
0.20
0.23
MuSiQue
High
High
91
0.16
0.241
0.20
Low
High
73
0.14
0.21
0.17
Low
0.181
Limited
Real
70
0.14
0.20
0.18
0.14
Symb.
Symb.
80
0.14
0.19
0.151
0.11
Real
Real
0.40
Simpl.
High
78
0.34
0.35
0.38
High
Low
72
0.33
0.33
0.35
0.28
SummHay
High
High
70
0.30
0.30
0.29
0.28
Simpl.
Low
0.29
Symh.
Symb.
13
0.14
0.14
0.16
0.18
Table 12. Average retrieval insight scores for attention heads the intersection and complement.
Dataset Variant
Llama-3-888 Bama-3-8B-Instruce
Mistral-7B-Instrueess
Task
Concept
Context
Inter.
Compl.
Inter.
Compl.
Real
Real (Who, When, Where)
0.045
0.011
0.059
0.019
Real
Real (Who)
0.052
0.012
0.062
0.021
Real
Real (When, Where)
0.050
0.012
0.059
0.018
High
High
0.046
0.010
0.057
0.020
MDQA
High
Low
0.045
0,015
0.056
0.018
Low
High
0.044
0,010
0.057
0.013
Low
Low
0.049
0.013
0.054
0.015
Symbolic
Symbolic
0.051
0.013
0.059
0.020
Real
Real (Limited)
0.121
0.049
0.065
0.021
High
High
0.105
0.040
0.053
0.015
High
Low
0.096
0,045
0.055
0.018
MuSiQue
Low
High
0.116
0,048
0.055
0.017
Low
Low
0.106
0.054
0.058
0.021
Symbolic
Symbolic
0.099
0.037
0.057
0.026
High
High
0.071
0.008
0.093
0.036
High
Low
0.092
0.016
0.098
0.039
SummHay
Simplified
Low
0.087
0.017
0.097
0.050
Simplified High
0.068
0.008
0.093
0.041
Symbolic
Symbolic
0.097
0.021
0.213
0.064
H.4. Retrieval Score Cosine Similarity
Across Tasks See Table 17. Similar to fine-tuning only attention-heads, we find the highest similarity between MDQA and MuSiQue retrieval scores, and much lower similarity with SummHay Citation scores, reflecting the different nature of the task (extractive QA vs. citation).
Synthetic Datasets vs. Real Task Performance See Figure 11. Overall, we find that synthetic datasets with lower performance recruit fewer scoring attention heads, although the relationship is weaker than when only fine-tuning attention
23

---

Understanding Synthetic Context Extension via Retrieval Heads
Table 13. Llama-3-8B-Instruct (all LoRA modules): Performance (F1) of fine-tuning different synthetie data the long-context retrieval and reasoning tasks. The results of training on the best synthetic datasets are bolded.
Concept Exp.
Context Div.
MDQA
MuSiQue
Concept Exp.
Context Div.
SummHay
High
High
0.35
0.40
High
High
0.83
High
Low
0.39
0.42
High
Low
0,68
Low
High
0.49
0.30
Simplified
High
0.83
Low
Low
0.47
0.38
Simplified
Low
0,58
Symbolic
Symbolic
0.46
0.37
Symbolic
Symbolic
0.63
Real Data (Full)
0.82
0,45
Real Data (Full)
0.81
Real Data (Limited)
0.84
0.32
Table 14, Llama-3-8B-Instruct (LoRA all modules): Pairwise recall of attention heads with non-zero retrieval scores for MDQA synthetic datasets. Limited datasets: I., Who, When, Where; L2 When, Where; Lg Who. Recall of real data retrieval heads moderately correlated with F1 (Spearman R = 0.60).
R,R
R,R (I.1)
R,R (L.2)
R,R (L.)
H,H
H,I.
I.H
I,I.
S,S
# Heads
F1
R.R
L.00
0.71
0.70
0.76
0.79
0.73
0.71
0.75
0.76
137
0.82
R,R (L1)
0.77
1.00
0.83
0.84
0.79
0.71
0.80
0.83
0.84
148
0.84
R.R (L2)
0.77
0.84
1.00
0.83
0.79
0.71
0.75
0.80
0.84
150
0.73
R.R (L.)
0.72
0.74
0.71
1.00
0.69
0.67
0.70
0.75
0.77
129
0.72
H.H
0.67
0.66
0.67
1.00
0.69
0.70
0.74
0.74
0.35
H.L
0.76
0.69
0.67
0.74
0.79
1.00
0.72
0.78
0.76
143
0.39
L.H
0.82
0.86
0.79
0.86
0.89
0.80
1.00
0.91
0.90
159
0.49
L.L
0.76
0.77
0.74
0.81
0.81
0.76
0.79
1.00
0.82
138
0.47
S,S
0.69
0.70
0.74
0.73
0.66
0.70
0.75
1.00
0.46
Table 15. Llama-3-8B-Instie (LoRA all modules): Pairwise recall attention with non-zero retrieval scores MuSiQue synthetic datasets, Recall of real data retrieval heads is moderately correlated with F1 (Spearman = 0.36).
R.R
R,R (T.)
H,H
H,I.
I,H
L,I.
S,S
Heads
F1
R.R
1.00
0.94
0.82
0.77
0.88
0.88
0.86
135
0.48
R.R (L)
0.44
1.00
0.56
0.41
0.46
0.56
0.39
63
0.41
H.H
0.59
0.87
1.00
0.64
0.75
0.73
0.61
0.40
H.L
0.78
0.89
0.89
1.00
0.90
0.86
0.78
136
0.42
L.H
0.65
0.73
0.77
0.66
1.00
0.83
0.71
100
0.30
L.L.
0.50
0.68
0.57
0.49
0.64
1.00
0.55
77
0.38
S.S
0.75
0.73
0.73
0.68
0.84
0.84
1.00
118
0.37
Table 16. Llama-3-8B -aama----B-Inttteee Instructt (LoRA all modules); Pairwise recall of attention heads with non-zero insight scores for SummHay Citation synthetic datasets. Recall of the real data insight heads moderately correlated with F1 (Spearman R 0.58).
R.R
H.H
H,L
L.H
L.L
S.S
# Heads
Fl
R.R
1.00
0.63
0.77
0.50
0.66
0.71
45
0.81
H.H
0.89
1.00
1.00
0.73
0.81
0.93
63
0.82
H.I.
0.67
0.62
1.00
0.49
0.60
0.76
39
0.68
L.H
0.87
0.90
0.97
1.00
0.84
0.90
78
0.83
L.L
0.84
0.75
0.90
0.63
1.00
0.81
58
0.57
S.S
0.67
0.62
0.82
0.49
0.59
1.00
42
0.62
heads.
24

---

Understanding Synthetic Context Extension via Retrieval Heads
Table 17. Llama-3-8--tnstuctt (all LoRA modules): Cosine similarity of real dataset retrieval scores (+ SummHay insight scores) across
tasks.
MDQA
MuSiQue
SummHay Retrieval
SummHay Insight
MDQA MuSiQue
1.00
0.85
0.35
0.16
0.85
1.00
0.50
0.29
SummHay Retrieval
0.35
0.50
1.00
SummHay Insight
0.16
0.29
0.11
1.00
MDQA
MusiQue R-0.1468
SummHay Citation R-0.5429
R-0. 0.6333
0.475
Liama: 3.88 lei--issttctt
0.450
0.80
Rew Real
Limited
0.7
0.425
0.75
Symboltz
0.600
Concept
0.73
0.350
0.65
Context
0.525
High
0.60
1aw
0.300
0.85
0.90
0.85
0.50
0.85
1.00
0.75
040
Figure 11. Llama-3-8B-Instruucc (all LoRA modules): Cosine similarity between the retrieval scores real datasets (R, R) their synthetic versions, and Spearman correlation for each setting.
Table Llama-3-SB-Insteenn (all LoRA modules): Results after patching heads that comprise the complement and intersection retrieval heads between the real and synthetic data versions, compared to random retrieval heads and original performance. Best patch F1 is bolded, and the improvement over the original FI.
Variant
Task
N
Compl.
Inter.
Rand.
Orig.
A
Concept
Context
Real
Real
0.82
Real
Real (Limited)
75
0.87
0.84
0.85
0.84
0.04
Low
High
70
0.66
0.61
0.56
0.49
0.17
MDQA
Low
Low
67
0.61
0.71
0.44
0.47
0.24
Symbolic
Symbolic
72
0.46
0.33
0.52
0.46
0.06
High
Low
72
0.63
0.27
0.47
0.39
0.24
High
High
63
0.47
0.57
0.64
0.35
0.29
Real
Real
0.48
High
Low
61
0.39
0.35
0.39
0.42
-0.03
Real
Real (Limited)
59
0.40
0.42
0.35
0.41
0.01
MuSiQue
High
High
73
0.41
0.37
0.31
0.40
0.01
Low
Low
68
0.35
0.38
0.02
Symbolic
Symbolic
51
0.43
0.10
0.35
0.37
0.06
Real
0.81
Simplified
High
39
0.77
0.81
0.82
0.83
-0.01
High
High
28
0.76
0.76
0.81
0.82
-0.01
SummHay
High
Simplitied
24
0.60
0.72
0.67
0.68
0.05
Symbolic
Symbolic
27
0.64
0.71
0.66
0.62
0.08
Simplified
Simplified
27
0.64
0.64
0.61
0.57
0.07
H.S. Patching
See Table 18. Notably, we find that patching complement attention head activations is the best in more settings than patching the intersection (7 settings vs. 6 settings). This is despite the results Table 19 showing that the intersection attention heads have higher scores.
25

---

Understanding Synthetic Context Extension via Retrieval Heads
Table 19. complement.
(all LoRA modules): Average retrieval insight scores for attention heads the intersection and the
Dataset Variant
Llama-3 3-Biaa---B-Imnsss I-strruct
Task
Concept
Context
Inter.
Compl.
High
Low
0.047
0.013
Real
Real (Who, When, Where)
0.046
0.013
High
High
0.049
0.010
Low
High
0,045
0.009
MDQA
Low
Low
0.047
0.012
Symbolic
Symbolic
0.049
0.015
Real
Real (Limited)
0.125
0.049
High
High
0.113
0.037
Low
High
0.105
0.039
MuSiQue
High
Low
0.095
0.037
Low
Low
0.119
0.045
Symbolic
Symbolic
0.099
0.031
Simplified
Low
0.067
0.010
High
Low
0.077
0.020
SummHay
Simplified
High
0.065
0.010
High
High
0.064
0.011
Symbolic
Symbolic
0.081
0.013
26


Reviews:
[{"review_id": "VbtzTW3be9", "reviewer": "Reviewer_S7xK", "summary": "This paper examines how synthetic data affects the performance of long-context language models (LLMs) on retrieval-based tasks. The authors find that while models fine-tuned on synthetic data generally underperform compared to those trained on real data, careful construction of synthetic datasets can partially close this performance gap. They identify \"retrieval heads\" as critical attention mechanisms that help models retrieve relevant information from long contexts, and show that synthetic data induces fewer of these heads than real data. However, there's a strong correlation between the presence of these retrieval heads and model performance. The study demonstrates that the cosine similarity between retrieval scores of real and synthetic data is a strong predictor of model effectiveness, providing insights into how to create better synthetic training data for long-context tasks.\n\n\n## update after rebuttal\n\nThe authors have ​​​​ explained the differences between their work and the retrieval head framework proposed by Wu et al. (2024). ​​Additionally​​, the inclusion of p-value experiments ​​further strengthens​​ their arguments. ​​As a result​​, I have ​​increased​​ my score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review!\n\n> Statistical analysis such as confidence intervals, p-values are missing.\n\nTo address this, we will make the following additions to Tables 1, 3 and 11 to show when a performance gain is significant:\n\n**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.\n\n| Concept Exp. | Context Div. | MDQA Llama3 | MDQA Mistral | MuSiQue Llama3 | MuSiQue Mistral | Concept Exp. | Context Div. | SummHay Llama3 | SummHay Mistral |\n|--------------|---------------|--------------|---------------|------------------|-------------------|---------------|----------------|------------------|----------------|\n| High         | High          | 0.31†        | 0.20†         | 0.37†            | 0.22              | High         | High          | 0.70†            | 0.28†          |\n| High         | Low           | 0.41†        | 0.23†         | **0.41**         | **0.23**          | High         | Low           | 0.61†            | 0.28†          |\n| Low          | High          | **0.49**     | **0.31**      | 0.29             | 0.21              | Simplified   | High          | **0.79**         | **0.38**       |\n| Low          | Low           | 0.47†        | 0.24†         | 0.34†            | 0.17†             | Simplified   | Low           | 0.65†            | 0.28†          |\n| Symbolic     | Symbolic      | 0.48         | 0.16†         | 0.32†            | 0.11†             | Symbolic     | Symbolic      | 0.54†            | 0.18†          |\n|              |               |              |               |                  |                   |               |                |                  |                |\n| **Real Data (Full)** ||  **0.83** | **0.64** | **0.45** | **0.20** | **Real Data (Full)** || **0.81** | **0.40** |\n| Real Data (Limited)  || 0.80     | 0.59     | 0.32     | 0.16     | Real Data (Limited) || 0.80     | 0.40     |\n| Non-FT       || 0.45             | 0.12              | 0.22          | 0.03          | Non-FT         || 0.40           | 0.07           |\n\n**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.\n\n> 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.\n\nIn 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.\n\n> 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?\n\nWhile 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).\n\nTo 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.", "claims_and_evidence": "The claims made in the paper are generally well-supported:\n\n- The underperformance of synthetic data fine-tuning compared to real data is demonstrated through performance metrics across multiple tasks and models (Table 1).\n\n- The predictive power of cosine similarity between real and synthetic data retrieval scores is demonstrated through direct comparison of similarity metrics with performance outcomes (Figure 3).\n\n- The task-specific nature of retrieval heads is supported by cosine similarity measurements across different task types (Table 2).", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria in the paper are well-suited to the problem of understanding how synthetic data affects the performance of long-context language models on retrieval-augmented tasks:\n\n- The authors chose three diverse long-context tasks (MDQA, MuSiQue, and SummHay Citation) that represent different aspects of retrieval and reasoning. This provides a comprehensive understanding of how synthetic data impacts various types of long-context processing.\n\n- The systematic variation of concept expression and context diversity in synthetic datasets allows for controlled experimentation on how different aspects of data realism affect model performance. \n\n- The focus on retrieval heads as a specific mechanism for understanding model behavior is appropriate, as these heads have been shown to be critical for information retrieval in long-context settings. This provides a concrete, interpretable measure of how well models are learning to perform the required tasks.\n\n- The paper includes appropriate baselines (fine-tuning on real data) and makes meaningful comparisons between different synthetic data construction methods.", "theoretical_claims": "There are no theoretical claims in the paper. Authors mainly use experiments to support the claims.", "experimental_designs_or_analyses": "The experimental designs and analyses in this paper are generally sound.\n\n- The systematic variation of concept expression and context diversity in synthetic datasets allows for controlled experimentation. This approach effectively isolates variables that might influence fine-tuning outcomes.\n\n- The focus on retrieval heads as a specific mechanism for understanding model behavior is appropriate, as these heads have been shown to be critical for information retrieval in long-context settings. This provides a concrete, interpretable measure of how well models are learning to perform the required tasks.\n\nHowever, While the paper shows numerical differences in performance metrics, it doesn't provide detailed statistical analysis (confidence intervals, p-values) to confirm whether these differences are statistically significant.", "relation_to_broader_scientific_literature": "The contributions of this paper may shed light on **Retrieval Mechanisms in Transformers** and **Synthetic Data for Language Model Training** research domains. Specifically, The paper builds upon work by Wu et al. (2024), which identified retrieval heads as critical mechanisms for long-context factuality in LLMs. This paper extends that research by examining how synthetic data affects the development and effectiveness of these retrieval heads across different tasks and model architectures.", "essential_references_not_discussed": "The paper did a good work for literature review.", "other_strengths_and_weaknesses": "Strengths:\n- The paper creatively combines existing ideas about synthetic data and mechanistic interpretability, focusing specifically on how synthetic data affects the development of retrieval heads in long-context LLMs.\n- The findings have practical implications for developing more effective synthetic training data for long-context LLMs, which is important given computational constraints of training on real long-context data.\n- The writing is generally clear and accessible, with appropriate technical detail for the intended audience.\n\nWeaknesses:\n\n- Statistical analysis such as confidence intervals, p-values are missing.\n\n- 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.\n\n- 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?", "other_comments_or_suggestions": "See above.", "questions_for_authors": "See **Other Strengths And Weaknesses**", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Oql5xy6nac", "reviewer": "Reviewer_Dczr", "summary": "The paper explores the impact of fine-tuning large language models (LLMs) with synthetic data for long-context tasks, particularly in retrieval and reasoning. The study evaluates different methods of synthetic data generation, varying both the realism of the \"needle\" (key concept) and the diversity of the \"haystack\" (context). The core finding is that the effectiveness of synthetic data can be interpreted through retrieval heads—specialized attention heads responsible for extracting relevant information. The study demonstrates that retrieval heads learned from synthetic data correlate well with those learned from real data, but synthetic training remains less effective. The authors also introduce a method of patching retrieval heads to improve model performance. The findings contribute to understanding synthetic data's role in LLM training and provide insights into designing better synthetic datasets.\n\n====== update after rebuttal ======\n\nDuring the rebuttal process, overall, I think the authors did not provide sufficiently empirical or insightful direct responses to some of the points reviewers raised. For instance, phrases such as \"leave for future work,\" \"we do not explore it here,\" and \"out of the scope of this work\" were largely used, which may indicate a lack of deeper engagement. However, considering the strengths and merits of the work as it currently stands, my final opinion is to maintain my original score of weak accept.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review!\n\n> 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].\n\nThank 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.\n\n> 1. Retrieval heads seem necessary but not sufficient for strong task performance (MoreHopQA: More than multi-hop reasoning).\n\nWe 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.\n\n> 2. Would your results generalize to generative synthetic data tasks (e.g., reasoning, math, or code)?\n\nOur 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.\n\n> 3. Does retrieval head behavior vary across different transformer architectures (e.g., mixture-of-experts models)?\n\nWe 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.\n\n> 4. How does synthetic data performance change as model size scales up? Does a larger model mitigate synthetic data limitations?\n\nDue 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.)\nNevertheless, 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.", "claims_and_evidence": "Most claims are clearly supported.\n1. Claim: Synthetic data fine-tuning can extend the effective context of LLMs, but it underperforms real data training.\n- Evidence: The authors compare performance on three long-context tasks (MDQA, MuSiQue, SummHay) and show that even the best synthetic datasets have a significant performance gap compared to real data.\n\n2. Claim: Retrieval heads play a key role in model performance on long-context tasks.\n- Evidence: Through a mechanistic interpretability analysis, the results demonstrate that models trained on synthetic data have a subset of retrieval heads found in models trained on realistic data.\n\n3. Claim: Patching retrieval heads from models trained on realistic data into models trained on synthetic data can improve performance.\n- Evidence: The authors conduct intervention experiments, showing that patching retrieval heads from real-data models improves performance on long-context tasks.", "methods_and_evaluation_criteria": "The experimental design is sound and appropriate for the problem domain. The study fine-tunes two well-known LLMs (Llama-3-8B-Instruct and Mistral-7B-Instruct) and evaluates them on three long-context tasks, covering single-hop retrieval (MDQA), multi-hop retrieval (MuSiQue), and citation retrieval (SummHay).", "theoretical_claims": "The paper does not primarily present new theoretical results, but it builds upon and extends existing work on retrieval heads (*Retrieval Head Mechanistically Explains Long-Context Factuality*). The mechanistic explanations are empirically validated through experiments rather than formal proofs.", "experimental_designs_or_analyses": "The main results are mostly sound, I also checked the following aspects:\n1.\tModels are fine-tuned with consistent training procedures.\n2.\tHyperparameters are documented (Appendix C).\n3.\tRetrieval heads are measured consistently across different models and datasets.", "relation_to_broader_scientific_literature": "The paper is well-positioned within existing literature:\n- It builds on work on synthetic data for LLM fine-tuning (From Artificial Needles to Real Haystacks: Improving Retrieval Capabilities in LLMs by Finetuning on Synthetic Data)\n- It extends retrieval head analysis from single-hop settings to multi-hop and long-context tasks.\n- It connects with mechanistic interpretability work on transformer circuits (In-context Learning and Induction Heads.)\n\nIt 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].\n\n---\n[1] Img-diff: Contrastive data synthesis for multimodal large language models\n[2] Data-juicer 2.0: cloud-scale adaptive data processing for foundation models\n[3] Make your llm fully utilize the context\n[4] Long context alignment with short instructions and synthesized positions\n[5] Understanding and Mitigating the Bias Inheritance in LLM-based Data Augmentation on Downstream Tasks\n[6] LLM-Based Synthetic Datasets: Applications and Limitations in Toxicity Detection", "essential_references_not_discussed": "None.", "other_strengths_and_weaknesses": "Strengths:\n- Clear mechanistic insights into why synthetic data works.\n- Systematic comparison of synthetic data types.\n\nWeaknesses:\nMy major concern lies in the generalization beyond task and model specifications.\n- It is suggested to add more discussion on other LLM capabilities like reasoning.\n- The different transformer architectures and model size scaling up should also be considered.", "other_comments_or_suggestions": "None.", "questions_for_authors": "1. Retrieval heads seem necessary but not sufficient for strong task performance (Morehopqa: More than multi-hop reasoning). \n2. Would your results generalize to generative synthetic data tasks (e.g., reasoning, math or code)?\n3. Does retrieval head behavior vary across different transformer architectures (e.g., mixture-of-experts models)?\n4. How does synthetic data performance change as model size scales up? Does a larger model mitigate synthetic data limitations?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DXehrXepJe", "reviewer": "Reviewer_wNwy", "summary": "This work aims to answer an important research question in the field of long-context modeling: how could the training on synthetic long-context data improve LLMs.\nThe authors present a novel investigation into the fine-tuning of LLMs using synthetically-generated long-context data.\nOne of the key contributions of this paper is the exploration of varying the realism in the \"needle\" concepts (the information to be retrieved) and the diversity of the \"haystack\" context (the broader dataset).\nThe paper's findings reveal that while models trained on synthetic data do not perform as well as those trained on real data, the underlying effectiveness of synthetic data can be reflected by the patterns in retrieval heads.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review!\n\n> 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.\n\nRegarding 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!\n\n> 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.\n\nWu 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.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "This work provides an empirical study. There is no theoretical proof.", "experimental_designs_or_analyses": "Yes. The design of various synthetic datasets makes sense to me.", "relation_to_broader_scientific_literature": "This work contributes to prior studies about (1) interpretability of attention mechanism and (2) the effectiveness of synthetic long-context data.", "essential_references_not_discussed": "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].\nAlthough 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.\n\n[1] Make your llm fully utilize the context.\n\n[2] Bootstrap Your Own Context Length.", "other_strengths_and_weaknesses": "No.", "other_comments_or_suggestions": "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).\nIf there are some observations in line with the conclusions in this work, it could greatly contributes to the soundness of this work.", "questions_for_authors": "No.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1ck3xYrP2s", "reviewer": "Reviewer_6T1q", "summary": "The paper demonstrates that synthetic context extension can partially emulate real data’s effects on LLMs but falls short due to less effective training of retrieval heads. It provides a framework for understanding this gap through retrieval heads, offering both a diagnostic tool and a path toward improving synthetic data generation for long-context tasks. These findings could inform strategies to create synthetic datasets that better target the necessary model components, reducing reliance on costly real long-context data.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review!\n\n> 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\n\nParametric 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.\n\n> 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.\n\nReasoning complexity: for our experiments, we explore tasks at 3 different levels of complexity: single-hop (MDQA), two-hop (SummHay Citation), and three-hop (MuSiQue).\n\nDistractor 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.\n\n> 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.\n\nDifferent 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.\n\nBroader 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.", "claims_and_evidence": "1. The paper presents experimental results across three tasks: Multi-Document Question Answering (MDQA), Multi-Hop Situated Question Answering (MuSiQue), and SummHay Citation. For each task, models fine-tuned on synthetic data consistently show lower F1 scores compared to those fine-tuned on real data. For example, on MDQA, the best synthetic data yields an F1 score of 0.49, while real data achieves 0.83 using Llama-3-8B-Instruct. Similar gaps are observed for MuSiQue and SummHay Citation, though the differences vary in magnitude. This claim is strongly supported by the evidence. The consistent performance gap across multiple tasks and models provides clear and convincing support for the claim that synthetic data underperforms compared to real data in long-context tasks.\n2. The paper identifies retrieval heads as attention heads specialized in retrieving key information from the context. It shows that models trained on synthetic data have fewer retrieval heads with positive retrieval scores (e.g., 112 and 74 for synthetic data vs. 129 for real data on MuSiQue). Additionally, there is a strong correlation between the recall of retrieval heads (the overlap of synthetic data heads with real data heads) and downstream task performance, with a Spearman correlation of 0.75 for Llama-3-8B-Instruct on MuSiQue. This claim is well-supported by the evidence. The reduction in retrieval heads and the high correlation with performance provide a convincing explanation for the performance gap. However, the paper’s focus on retrieval heads as the primary mechanism could be slightly overstated, as other components (e.g., multi-layer perceptrons, or MLPs) might also contribute, though this is partially addressed by a footnote stating similar conclusions hold when fine-tuning all modules.\n\nWeekness: 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; 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.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "The proofs for theoretical claims in the paper are correct.", "experimental_designs_or_analyses": "The experimental designs or analyses in the paper are sound.", "relation_to_broader_scientific_literature": "The paper investigates the use of synthetic data for training LLMs on long-context tasks, specifically Multi-Document Question Answering (MDQA), Multi-Hop Situated Question Answering (MuSiQue), and SummHay Citation. It explores how varying the realism of \"needle\" concepts (key information to retrieve) and the diversity of the \"haystack\" context (surrounding information) impacts model performance. This systematic approach sheds light on the properties of synthetic data that influence real-world long-context capabilities. Previous research has used synthetic data for a variety of NLP tasks. In contrast, this paper extends the application of synthetic data to long-context tasks, systematically analyzing how realism and diversity affect performance across multiple domains.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "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.", "other_comments_or_suggestions": "No.", "questions_for_authors": "No.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper presents a timely and well-executed empirical study on the effectiveness of synthetic context extension for training long-context LLMs, offering mechanistic insights via the analysis of retrieval heads. The experimental design is solid, spanning multiple datasets, models, and task types, and the authors' identification of retrieval heads as key to performance provides a concrete diagnostic tool for evaluating synthetic dataset quality. The inclusion of patching experiments further strengthens the causal claims about the role of these heads.

While some reviewers raised concerns about limited novelty beyond prior work (particularly Wu et al. 2024), the authors make a compelling case that their contributions lie in extending retrieval head analysis to multi-hop and complex reasoning tasks, as well as offering principled methods for evaluating synthetic datasets in these settings. Additionally, reviewers pointed out the initial lack of statistical significance testing, which the authors addressed with a robust rebuttal including bootstrap analyses. Broader generalization—across models, tasks, or LLM capabilities like reasoning—remains an open question, but the empirical evidence presented here is strong and the methodology is reproducible and impactful.

Overall, while the work may not be a clear top-tier paper in terms of conceptual innovation, its solid experimental foundation, practical relevance for synthetic data construction, and interpretability contributions justify acceptance.
Decision: Accept (poster)