Found in the middle: Calibrating Positional Attention Bias Improves Long Context Utilization (2024)

Abstract: Large language models (LLMs), even when specifically trained to process long input contexts, struggle to capture relevant information located in the middle of their input. This phenomenon has been known as the lost-in-themiddle problem. In this work, we make three contributions. First, we set out to understand the factors that cause this phenomenon. In doing so, we establish a connection between lost-in-the-middle to LLMs’ intrinsic attention bias: LLMs exhibit an U-shaped attention bias where the tokens at the beginning and at the end of its input receive higher attention, regardless of their relevance. Second, we mitigate this positional bias through a calibration mechanism, found-in-the-middle, that allows the model to attend to contexts faithfully according to their relevance, even though when they are in the middle. Third, we show foundin-the-middle not only achieves better performance in locating relevant information within a long context, but also eventually leads to improved retrieval-augmented generation (RAG) performance across various tasks, outperforming existing methods by up to 10 percentage points. These findings open up future directions in understanding LLM attention bias and its potential consequences.

heading: 1 Introduction
text: Effective prompting of large language models (LLMs) (Brown et al., 2020; Anil et al., 2023; Touvron et al., 2023) has enabled a variety of userfacing applications, including conversational interfaces (chatbots) (Thoppilan et al., 2022), search and summarization (Min et al., 2024), open-domain question answering (Izacard and Grave, 2021), tool
*Work done while the author was a student researcher at Google Cloud AI Research. Correspondence to: ChengYu Hsieh <cydhsieh@cs.washington.edu>, Chen-Yu Lee <chenyulee@google.com>
†The authors contributed equally to this work.
usage (Hsieh et al., 2023), fact checking (Asai et al., 2023), and collaborative writing (Lee et al., 2019). Some of these applications, such as search and summarization (Ji et al., 2023; Min et al., 2023; Asai et al., 2023), require the ability to retrieve information from external knowledge sources. As a result, retrieval-augmented generation (RAG) has become a powerful solution. RAG fetches relevant documents (e.g. structured tables (Wang et al., 2024) and API documentation (Karpukhin et al., 2020)) from external knowledge sources and makes them available in the LLMs’ input prompt (Khandelwal et al., 2020; Borgeaud et al., 2022; Izacard et al., 2022b; Xu et al., 2023a). Despite the widespread utility of RAG (Li et al., 2023a; Xiong et al., 2023;
14982
OpenAI, 2022; Gemini Team, 2023), recent experiments highlight a striking deficiency: LLMs struggle to locate relevant documents when they are placed in the middle of their input prompts (Liu et al., 2023; Li et al., 2023a). They call this the lost-in-the-middle phenomenon.
To overcome this phenomenon, a few mechanistic strategies have been proposed (Jiang et al., 2023; Peysakhovich and Lerer, 2023). These methods re-rank the relevance of different documents and re-order the most relevant ones to either the beginning or end of the input context. Unfortunately, re-ranking usually requires additional supervision or dedicated finetuning for performant RAG performance (Karpukhin et al., 2020; Shi et al., 2023c; Sun et al., 2023). Worse, re-ranking methods do not fundamentally improve LLMs’ ability to utilize and capture relevant information from the provided input contexts. The underlying causes of this behavior remains unclear, even though it has been observed across multiple decoder-only LLMs (Touvron et al., 2023; Li et al., 2023a; OpenAI, 2022).
In this work, we make three contributions: First, we set out to understand the potential factors leading to the lost-in-the-middle problem. We establish a connection between lost-in-the-middle to LLMs’ intrinsic attention bias (see Figure 1). Specifically, we find that models often demonstrate a U-shaped attention distributions, with higher attention values assigned to the beginning and end of the input prompt. This correlates well with the U-shaped RAG performance observed in prior literature (Liu et al., 2023). Interestingly, this focus on the beginning and end also extends to content utilization: models preferentially use information from the beginning and end of their prompts (Ravaut et al., 2023; Peysakhovich and Lerer, 2023). This leads us to hypothesize that the positional attention bias may contribute to the phenomenon, wherein the bias could lead to overreliance on content at the beginning/end of the input, regardless of its true relevance.
Second, we verify our hypothesis by intervening on this attention bias to determine its impact on performance. We propose a mechanism to disentangle positional bias from model’s attention. We first esitmate this bias through measuring the change in attention as we vary the relative position of a fixed context in the LLM’s prompt. By quantifying and then removing this bias from the attention scores for a given query, we can obtain the calibrated attention scores across the retrieved
documents. This calibrated attention proves to be better correlated to the ground truth relevance of the document to a user query. In open-domain question answering tasks (Kwiatkowski et al., 2019), our proposed calibrated attention outperforms popular existing approaches for ranking the relevance of retrieved documents (up to 48 Recall@3 points). This finding challenges the recent belief that LLMs struggle to capture relevant context embedded in the middle of inputs, suggesting they may indeed be capable of doing so, but are only hindered by the overwhelming positional bias.
Third, we operationalize our calibration mechanism as a solution for this phenomenon, naming our attention intervention found-in-the-middle (FITM). We show that calibrating the attention leads to improvements across two popular LLMs with different context window lengths on two RAG tasks. Our experiments demonstrate improvements over standard model generation by up to 10 percentage point on NaturalQuestion dataset (Kwiatkowski et al., 2019). We hope the work opens up future directions in understanding LLM’s attention biases and their effect on downstream tasks.

heading: 2 Positional attention bias overpowers mid-sequence context
text: Recent work has produced language models capable of handling increasingly long input contexts (Xiong et al., 2023; Li et al., 2023a). However, many of these models struggle to locate relevant information placed in the middle of the input sequence (Liu et al., 2023), a phenomenon known as the “lost-in-the-middle” problem. While this problem is widely recognized, the potential factors contributing to this behavior remain poorly understood. In this work, we seek to deepen our understanding of the problem through a suite of exploratory qualitative and quantitative studies.
Setup. We adhere to the original experimental setup outlined in Liu et al. (2023), utilizing an opendomain question answering task (Kwiatkowski et al., 2019) for our exploratory study. In the lostin-the-middle setup (Liu et al., 2023), a model is tasked to answer a user query xq using a set of k related documents retrieved from an external data source D = {xgold, xdistract1 , . . . , xdistractk−1 }, where only the gold document xgold contains the correct answer. The question and documents are typically serialized as an input sequence
xprompt = [xq, xdoc1 , ..., x doc k , x q], prompting a language model to generate the final answer1. Observations indicate that model performance significantly decreases when xgold is placed within the middle of the input prompt (i.e., xdocbk/2c), compared to scenarios where xgold is placed at the beginning or end. Here, we reproduce lost-in-the-middle phenomenon with a Vicuna-7b-v1.5-16k (Vicuna) model (Li et al., 2023a) to gain deeper insights into the characteristics of the model’s errors. We focus our error analysis on the setting where we have a
1We repeat the question before and after the documents so that the model can better attend to relevant contexts (Liu et al., 2023; Xu et al., 2023b).
total of 20 documents (K = 20). We specifically look at the examples where the model makes incorrect predictions when the gold document is placed at the middle (10-th) position.

heading: 2.1 U-shaped attention bias
text: We first examine responses generated when gold documents are placed in the middle of input prompts. Qualitatively, the model’s response exhibits a strong bias towards the document at the first position, regardless of the gold document’s location (Figure 2). This bias persists whether the input documents retain their original order or are randomly shuffled.
The strong correlation between the model’s output and the first document could suggest that they
are highly relevant, distracting the model (Shi et al., 2023a). However, quantitatively, the model’s response strongly depends on the document at the first position (Figure 3). This dependence persists even after randomly shuffling the document order, irrespective of its relevance to the query. We measure the dependence by computing the TF-IDF similarity between the response and each document (gold document originally at position 10).
To investigate the potential origins of positional bias, we visualize the model’s self-attention weights, as the weights has been shown to correlate with models’ generations, although not necessarily causal (Dong et al., 2021; Zhang et al., 2023). More formally, given an input prompt consisting of K documents xprompt = [xdoc1 , ..., x doc K ], where each document xdock = {xdock,i } Nk i=1 contains Nk tokens, let Attn : X × N→ R denote a function that computes the average attention weights assigned to document xdock as Attn(x
prompt, k) =∑Nk i=1 attn(x doc k,i )/Nk, where attn(x doc k,i ) is the attention weight value allocated to token xdock,i when predicting the next |xprompt|+ 1 token.
Specifically, we visualize the average selfattention weights assigned to each document across all tokens, decoder layers, and heads. We investigate how these weights vary based on document position within the input prompt. Interestingly, Figure 4 (blue curve) reveals a U-shaped attention pattern. Documents near the beginning and end of the input receive higher weights, while those in the middle receive lower weights. Crucially, the U-shaped pattern persists even after randomly shuffling document order (Figure 4, orange curve), suggesting that this bias does not depend on the documents’ actual content.

heading: 2.2 Does attention favor relevant context?
text: Observation 1: Model prioritizes relevant contexts from the same position. In Figure 4, we observe a significant difference in attention values at xdoc10 when comparing examples with original document order (blue) and randomly shuffled order (orange). Specifically, the attention value is notably higher when when xdoc10 is controlled to be xgold. This contrasts with instances where xdoc10 is uncontrolled, suggesting that apart from U-shaped positional bias, the model exhibits an ability to prioritize relevant context.
Observation 2: Model prioritizes highlyweighted documents for generation. Based on
these observations, we hypothesize that positional attention bias significantly influence the model’s tendency to rely heavily on the first documents during output generation. Specifically, the models are more likely to incorporate the document receiving the highest attention (often the first) into its output. To validate this, for each of the examples of interest (where the model makes incorrect predictions), we divide their documents into first half receiving higher model attention and second half receiving lower attention. We then count the number of examples in which the first or second half contains the document that is most likely used in the model’s generation (i.e., having the highest TF-IDF score with model’s response). In Table 1, we show that documents receiving higher attention positively correlates with them being used in the model’s generation.
From the above studies, we see that not only the model exhibits a U-shape positional attention bias, but this bias also correlates strongly with the model’s biased tendency in using documents placed at certain positions in forming its response. We thus conjecture that lost-in-the-middle happens because of the dominating force of positional bias.

heading: 3 Found-in-the-middle: modeling and isolating positional attention bias
text: Ideally, a model should leverage contexts in the input prompts—faithfully according to their relevance—for generating the response, instead of biasing towards contexts placed at certain positions within the input. Towards this goal, we are interested in modeling the positional attention bias and mitigating it such that model attention can reflect the true relevance of the input context and ultimately improve models’ effective utilization of the full context window.

heading: 3.1 Two main factors in model attention
text: In Sec. 2, we find that there are two main forces driving the model attention assigned to different documents of an input prompt: (a) where the document locates within the entire input, and (b) the relevance of the document.
Our hypothesis. We thus consider modeling the observable attention weights allocated to the k-th document of an input xprompt as:
Attn(xprompt, k) = f(rel(xdock ),bias(k)), (1)
where rel(·) measures the relevance of an input document, bias(·) characterizes the positional attention bias, f(·) is some unknown monotonically increasing function w.r.t. to both rel(xdock ) and bias(k). For ease of exposition, in the remainder of the paper, we overload Attn(xdoc, k) to denote the attention value assigned to document xdock placed at the k-th position within an input prompt containing K documents.
Corroborating our assumed model. Here, we conduct a suite of controlled experiments using NaturalQuestion with K = 20 and a Vicuna-7b-v1.5-16k model to corroborate our assumed model. Specifically, for Eq. 1 to hold, it implies that: Condition 1: When the relevance term is fixed, model attention increases as positional bias increases. That is, given two documents xdoc1 and xdoc2: if Attn(xdoc1, k) > Attn(xdoc1, l), then Attn(xdoc2, k) > Attn(xdoc2, l). Condition 2: Similarly, when the document position k is fixed, model attention increases as the relevance of the document increase: if Attn(xdoc1, k) > Attn(xdoc2, k), then Attn(xdoc1, l) > Attn(xdoc2, l).
We validate Condition 1 and 2 on 100 randomly sampled examples from NaturalQuestion dataset, each with K = 20 documents. For validating Condition 1, given a pair of documents (xdoc1, xdoc2) and positions (k, l), we can compute whether the relationship holds across all possible pairs. We can similarly test for Condition 2. In Table 2, we see that the percentage of valid example pairs are decently high, 83% and 72% respectively, for both conditions, providing supports to our hypothesis.
Recall that our goal is to disentangle positional attention bias from model attention such that the model can faithfully attend to relevant contexts, independent from their positions. So far, while we
have established the monotonic increasing nature of f in Eq. 1, we have yet characterize the actual form of f to remove the positional bias term from model attention.
To approximate f , we consider simple linear models by following machine learning principles (a.k.a. Occam’s razor), for robust estimation:
Attn(xdoc, k) = rel(xdoc) + bias(k) + , (2)
where is a noise. To test how the model captures the underlying relationship, we compute Spearman’s rank correlation between Attn(xdoc1, k)−Attn(xdoc2, k) and Attn(xdoc1, l)−Attn(xdoc2, l)) over quadruplets of (xdoc1, xdoc2, k, l) collected from NaturalQuestion. A high correlation indicates small discrepancy between Attn(xdoc1, k)−Attn(xdoc2, k) and Attn(xdoc1, l)−Attn(xdoc2, l)). From our study, the linear model results in decently high correlation, 0.76, suggesting its effectiveness despite the simplicity. We therefore adopt Eq. 2 as our model and leave other alternatives with more degree of freedoms as future work 2.

heading: 3.2 Disentangling positional attention bias
text: Most notably, having a simple form of f allows us to isolate the effect of positional bias from model attention. Specifically, following from Eq. 2, we can first obtain a reference model attention value with a dummy document xdum by:
Attn(xdum, k) = rel(xdum) + bias(k) + . (3)
By subtracting Eq. 2 and Eq. 3, we can offset the bias term and obtain:
rel(xdoc) (4)
= Attn(xdoc, k)−Attn(xdum, k) + rel(xdum)
Consider using a consistent dummy document xdum which has a constant rel(xdum), we are then
2In Appendix C, we also explore log-linear models, which results in competitive 0.75 rank correlation.
able to obtain the true relevance of different documents xdoc, free from the positional bias. We refer to Attn(xdoc, k)−Attn(xdum, k) as calibrated attention as it removes the baseline attention, and call the overall calibration mechanism found-inthe-middle.
Calibrated attention finds relevant contexts in the middle. Eq. 4 allows us to leverage calibrated attention to estimate and rank the relevance of different documents within an input prompt. To validate the effectiveness of our model, we evaluate using calibrated attention to re-rank documents in an input prompt w.r.t. a given query. We evaluate on NaturalQuestion with the Vicuna model where we focus on the most challenging setting when the gold document in placed in the middle of the input prompt. We compare our model to:
• Vanilla attention: Using uncalibrated attention Attn(xprompt, k) to rank the documents.
• Query generation (Sun et al., 2023): Using likelihood of the model in generating the query based on the document.
• Relevance generation (Sun et al., 2023): Prompting the model to answer whether a document is relevant to a query.
In Table 3, we compare Recall@3 of different methods where we vary the total number of documents retrieved. We see that the proposed calibrated attention consistently outperforms vanilla attention by a large margin, and also shows superior performances when compared to the other two re-ranking metrics. The results validate that our proposed modeling approach is effective, and that if calibrated appropriately, language models can locate relevant information even when they are hidden in the middle of the input.

heading: 4 Improving long-context utilization with found-in-the-middle
text: Having validated that calibrated attention through found-in-the-middle is effective in locating relevant information within a long input context, we are ultimately interested in leveraging it to tackle lostin-the-middle problem and practically improve a model’s RAG performance.

heading: 4.1 Attention calibration
text: To allow the model to attend to contexts without being dictated by positional bias, we propose to intervene the model’s attention based on the proposed calibrated attention. Specifically, given an input xprompt, instead of allocating rel(xdock ) + bias(k) attention to the k-th document, our ideal model attention Attncalibrated(xdock ) would reflect only the relevance of the context rel(xdock ).
To achieve this, we propose to redistribute the attention values assigned to {xdock }Kk=1 according to rel(xdock ). Specifically, for each document x doc k , we propose to rescale the attention values on the tokens within the document, {xdock,i } Nk i=1, by:
attncalibrated(x doc k,i ) = (5)
αk
Attnoriginal(x doc k ) · attnoriginal(xdock,i ) · C,
where αk = Softmax(rel(xdock ), t), t is the temperature hyperparamter, and C is a normalization constant to ensure the total attention ∑ k,i x doc k,i remains unchanged. With the rescaling, we effectively make the final attention on xdock :
Attncalibrated(x doc k ) ∝ Softmax(rel(xdock ), t),
(6) where higher attention is allocated to more relevant context, and t controls the disparity level.

heading: 4.2 Calibrated v.s. uncalibrated attention
text: We evaluate the performance of the proposed attention calibration method. We conduct experiments on two multi-document question answering tasks (more details in Appendix A), NaturalQuestion (Kwiatkowski et al., 2019) and SynthWiki (Peysakhovich and Lerer, 2023), with two models supporting different context window length: Vicuna-7b-v1.5-16k (Vicuna) (Li et al., 2023a) and tulu-2-7b (Tulu) (Wang et al., 2023) with 16k and 8k context window respectively. For each dataset, we consider two settings with different number of retrieved documents, K = {10, 20}.
We leave further implementation details in Appendix B.
Found-in-the-middle improves long-context utilization across various datasets and models. In Figure 5, we see that found-in-the-middle attention calibration consistently outperforms the uncalibrated baseline by a large margin (up to 10 percentage point (pp) improvement) across different tasks and models. On the most challenging scenario when the gold document is placed midsequence, attention calibration consistently offers improvements from 6-10 pp. Notably, we see that attention calibration’s performance curve lies entirely above the vanilla baseline curve, validating the effectiveness of our method in improving models’ long context utilization.

heading: 4.3 Attention calibration in practice
text: In practice, to avoid the lost-in-the-middle effect, one commonly adopted workaround is to reorder the document positions, where documents considered more relevant are placed towards the beginning (or end) of the input. While these methods have led to performance improvements over the baseline without reordering, without handling the model’s intrinsic bias, reordering-based methods’ performance relies heavily on the correct ranking of the documents. We are thus interested in validating whether attention calibration can be applied on top of re-ordering methods to provide another layer of improvements.
Attention calibration improves existing RAG pipelines. We continue using NaturalQuestion and SynthWiki for evaluation. We compare to existing reordering methods including:
• Prompt reordering (Sun et al., 2023; Liang et al., 2023): Reorder documents based on relevance score generated through prompting.
• LongLLMLingua-rk (Jiang et al., 2023): Reorder documents using query generation as the reranking metric.
• Attention sorting (Peysakhovich and Lerer, 2023): Reorder documents using vanilla model attention assigned to the documents.
In Figure 6, we note that LongLLMLingua-rk and prompt reodering are invariant to the gold document’s position since they compute the relevance of each document independently. First, we see that
reordering methods do alleviate lost-in-the-middle problem where models’ performances increase when gold documents is placed mid-sequence. Furthermore, by applying attention calibration on top of a reodring mechanism (LongLLMLingua-rkin this case), LongLLMLingua-rk with calibration consistently achieve the highest performance across datasets and models, suggesting a way to further improve current RAG pipeline.

heading: 5 Related work
text: Retrieval augmented generation. While LLMs exhibit strong capabilities (Gemini Team, 2023; OpenAI, 2022; Touvron et al., 2023), their knowledge is inherently limited in its pretraining data, and they are observed to struggle in handling knowledge intensive tasks (Petroni et al., 2020). To tackle this, retrieval augmented generation (RAG) is an effective framework that retrieves relevant information from external knowledge sources to aid and ground language models’ generation (Lewis et al., 2020; Khandelwal et al., 2020; Borgeaud et al., 2021; Izacard and Grave, 2021; Izacard et al., 2022b).
Although RAG has powered many recent language model applications from questionanswering (Izacard and Grave, 2021) to automatic task completion (Shen et al., 2023), recent work show that LLMs tend to lost-in-the-middle, significantly hindering the full potential of RAG (Liu et al., 2023). In this work, we take a step further to understand the lost-in-the-middle problem from the viewpoint of attention bias. Moreover, we propose a remedy through attention calibration, which improves upon existing RAG frameworks.
Long-context utilization in language models. There is a rich literature on enabling LLMs to handle longer input contexts, including designing efficient training and finetuning schemes (Dao et al., 2022; Li et al., 2023b,a; Shi et al., 2023b) and inference-time methods that extend an LLM’s context length (Press et al., 2021; Ratner et al., 2023; Xiao et al., 2023; Bertsch et al., 2023). Nonetheless, even models specifically trained for long-context suffer lost-in-the-middle problem (Liu et al., 2023; Li et al., 2023a).
To improve LLMs’ performance on handling long contexts, recent methods design better prompting techniques and pipelines that mechanically work around the lost-in-the-middle problem (Chen et al., 2023; Jiang et al., 2023; Peysakhovich and
Lerer, 2023; Junqing et al., 2023). For instance, to avoid having the models process long input contexts, (Chen et al., 2023; Junqing et al., 2023) proposes to split long inputs into shorter contexts for models to better understand. To avoid relevant context being missed by the model, (Jiang et al., 2023; Peysakhovich and Lerer, 2023) proposes to rank the relevance of different parts of the input and re-order the most important parts to either the beginning or end of the entire input, where the models tend to focus more.
While these existing solutions lead to improved model performances by manipulating the input contexts, they do not fundamentally improve LLMs’ underlying long-context utilization capability. In contrast, we set out to directly improve LLMs’ long-context utilization capability to mitigate lostin-the-middle problem.
Self-attention and attention bias. The attention mechanism is initially introduced in RNN-based encoder-decoder architectures (Bahdanau et al., 2015; Luong et al., 2015). Building upon the selfattention mechanism, transformers (Vaswani et al., 2017) have achieved state-of-the-art performance in various domains (Devlin et al., 2018; Dosovitskiy et al., 2020). Self-attention has also been widely used as a proxy to understand and explain model behaviors (Clark et al., 2019; Hao et al., 2021; Vashishth et al., 2019).
However, the relationship between the lost-inthe-middle problem and LLM’s self-attention has been under-explored. As an initial trial, “attention sorting” (Peysakhovich and Lerer, 2023) sorts documents multiple times by the attention they receive to counter lost-in-the-middle. Recently, He et al. (2023) construct a dataset for training LLMs to focus on the most relevant documents among long
contexts. Unlike the method, which necessitate significant investment in data collection and LLM tuning, our method offers an efficient solution by mitigating lost-in-the-middle problem with off-theshelf LLMs.

heading: 6 Discussion
text: In this work, we understand and address the lostin-the-middle phenomenon, by establishing a connection between the phenomenon and models’ positional attention bias. We mitigate the bias by attention calibration which directly modifies the model’s attention mechanism, enabling LLMs to more faithfully attend to contexts based on their relevance, rather than their position. Experiments show that attention calibration improves the performance compared to its uncalibrated counterpart especially when relevant context occurs in the middle of the input. We additionally show attention calibration can be applied on top of existing reordering pipelines to further improve models’ performance.
Limitations
While our study presents significant advances in addressing the "lost-in-the-middle" problem and improving RAG performance in LLMs, several limitations are noteworthy:
Simplification of the mechanism behind positional attention bias. We proposed a simple hypothesis to model the positional attention bias, as shown in Eq. 1. However, the intrinsic mechanisms that drive this bias could be more intricate and dynamic than our current model accounts for. It is possible that some aspects of attention bias are learnable or adaptive, responding to subtle aspects of the data or training process that our current approach does not consider.
Computational overhead. Our method of calibrating positional attention bias, while effective, introduces additional computational overhead. Specifically, we require extraO(K) model forward passes to calibrate attention at each position, compared to vanilla model generation. However, in this study we aim to discover and calibrate the positional attention bias from a scientific perspective. We expect that our discovery can enable future research into developing more calibration methods with lower computational overhead.
Positional attention bias may be beneficial. Our method aims to completely remove positional
attention bias. However, it is important to note that this positional bias might actually be beneficial in certain contexts. In some specific tasks or scenarios, the natural tendency of models to focus more on the beginning and end of inputs could align well with the structure of the task or the nature of the data. Therefore, understanding the tasks and the applications is required before adopting our proposed calibration method.
The root cause of attention bias is unclear. In this work, we aim to discover and understand the connection between the lost-in-the-middle problem and LLMs’ intrinsic attention bias. However, our work does not definitively pinpoint the root cause of attention bias in LLMs. The cause of such a bias could be attributed to the distribution of pretraining corpora, the transformer model architecture, and the optimization process. Future research needs to delve deeper into the origins of this phenomenon.
Ethical Statement
In our research, we focus on enhancing the performance of large language models using existing public datasets, ensuring that no personal or sensitive data was collected or utilized. Our attention calibration method is aimed at improving the efficiency and accuracy of retrieval-augmented generation, with potential benefits across various domains including search engines, question-answering systems, and other text-based applications. It is important to acknowledge that as our technique builds upon pre-trained language models, it may inadvertently inherit and propagate existing biases inherent in these models. Apart from this significant concern, we do not identify any other immediate risks arising from the methodologies or findings presented in our paper.

heading: A Multi-doc QA datasets
text: We use NaturalQuestions (Kwiatkowski et al., 2019)3 (released in Apache-2.0 license) and SynthWiki (Peysakhovich and Lerer, 2023)4 to conduct the experiments. Both datasets contains questionanswer pairs, a gold document contains the answer, and K − 1 distractor documents, where K = 10 and 20.
The NaturalQuestions dataset is the subset with 2655 queries selected by Liu et al. (2023)5 where the annotated long answer is a paragraph. The k−1 distractor passages are Wikipedia chunks retrieved by Contriever (Izacard et al., 2022a) that are most relevant to the query but do not contain any of the annotated answers in NaturalQuestions. The distractor documents are presented in the context in order of decreasing relevance.
The SynthWiki dataset (Peysakhovich and Lerer, 2023) is a synthetic multi-doc QA dataset with 990 entries. All the documents in SynthWiki are GPT-4 generated Wikipedia paragraphs for fictional people, thus it can minimize the knowledge contamination issue from pre-training and ensure the LLMs can only use information from the provided context. The distractor documents are randomly sampled and randomly ordered in SynthWiki.
NaturalQuestions is collected from public English Wikipedia articles and SynthWiki is collected by GPT-4 automatic generation of English fake Wikipedia articles. These two dataset should not contain any information that names or uniquely identifies individual people or offensive content. We ensure that the use of these two datasets was consistent with their intended purpose for academic research and in accordance with their specified licensing agreements.
B Implementation details
In our experiments, we utilize tulu-2-7b and Vicuna-7b-v1.5-16k as the base models. Both models consist of 32 decoder layers, each with 32 attention heads. In applying attention calibration method to intervene model attention, we apply only to the last 16 decoder layers (and all of their attention heads). We find that intervening early layers may lead to unstable generation. We leave finding
3https://github.com/google-research-datasets/ natural-questions
4https://github.com/adamlerer/synthwiki 5https://github.com/nelson-liu/
lost-in-the-middle
the best set of attention heads to intervene as future directions (Zhang et al., 2023).
In the experiments, we find attention calibration to be robust to the temperature term t in Eq. 5. We set t = 5e−5 for all experiments.

heading: C Additional experiment results
text: Different model formulations. To approximate (1), in addition to linear models as shown in (2), we also investigate log-linear models, which is defined as
log Attn(xdoc, k) = rel(xdoc)+bias(k)+ , (7)
where is a noise. We compute rank correlation as described in Sec. 3. The result is shown in Table 4. The log-linear model and linear are competitive to each other, which all result in rank correlation above 0.76.
Experiment tables. Table 5 shows the exact numbers in our experiments.

heading: D Compute and inference details
text: In the experiments, we use the Huggingface Transformer package6 with the two models: Tulu-2-7B7 and Vicuna-7B-v1.5-16k8 both contains 7B parameters. We run the experiments with two NVIDIA A100 GPUs. The inference time is roughly 1 to 3 hours on both datasets. We run our experiments with all greedy decoding without any nondeterministic factor, so we only need to run the experiments for once. Our method is a pure inference method, so there is no need to do training or hyperparameter searching.
6https://github.com/huggingface/transformers 7https://huggingface.co/allenai/tulu-2-7b 8https://huggingface.co/lmsys/vicuna-7b-v1.
5-16k

heading: Future Work
text: com/nelson-liu/
lost-in-the-middle
the best set of attention heads to intervene as future directions (Zhang et al. , 2023). In the experiments, we find attention calibration to be robust to the temperature term t in Eq. 5. We set t = 5e−5 for all experiments