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-the-middle 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 found-in-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 point. These findings open up future directions in understanding LLM attention bias and its potential consequences.

## introduction effective prompting of large language models (llms) @xcite @xcite has enabled a variety of userfacing applications, including conversational interfaces (chatbots) @xcite , search and summarization @xcite , open-domain question answering @xcite , tool usage @xcite , fact checking @xcite , and collaborative writing @xcite . some of these applications, such as search and summarization @xcite @xcite , 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 @xcite and api documentation @xcite ) from external knowledge sources and makes them available in the llms' input prompt @xcite @xcite . despite the widespread utility of rag @xcite @xcite , recent experiments highlight a striking deficiency: llms struggle to locate relevant documents when they are placed in the middle of their input prompts @xcite . they call this the lost-in-the-middle phenomenon. to overcome this phenomenon, a few mechanistic strategies have been proposed @xcite . 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 @xcite @xcite . 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 @xcite @xcite . 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 @xcite . 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 @xcite . 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 @xcite , 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 @xcite . we hope the work opens up future directions in understanding llm's attention biases and their effect on downstream tasks. ## positional attention bias overpowers mid-sequence context recent work has produced language models capable of handling increasingly long input contexts @xcite . however, many of these models struggle to locate relevant information placed in the middle of the input sequence @xcite , 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 @xcite , utilizing an opendomain question answering task @xcite for our exploratory study. in the lostin-the-middle setup (liu et al., 2023), a model is tasked to answer a user query x q using a set of k related documents retrieved from an external data source @xmath0, x distract 1 , . . . , x distract k-1 }, where only the gold document x gold contains the correct answer. the question and documents are typically serialized as an input sequence write a high-quality answer for the given question using only the provided search results (some of which might be irrelevant). question: what is mercy mercy me by marvin gaye about search results: document [1] ...the follow-up single, "mercy mercy me (the ecology)", peaked at number-four on the hot 100... document [10] mercy mercy me (the ecology)... became regarded as one of popular music's most poignant anthems of sorrow regarding the environment document [20]... question: what is mercy mercy me by marvin gaye about answer: the song was written solely by gaye and became one of popular music's most poignant anthems of sorrow regarding the environment write a high-quality answer for the given question using only the provided search results (some of which might be irrelevant). question: what is mercy mercy me by marvin gaye about search results: document [1] ...the follow-up single, "mercy mercy me (the ecology)", peaked at number-four on the hot 100... document [10] mercy mercy me (the ecology)... became regarded as one of popular music's most poignant anthems of sorrow regarding the environment document [20]... question: what is mercy mercy me by marvin gaye about answer: …the song was released as the follow-up single to the album's title track and peaked at number-four on the hot 100… original prompt / model's response write a high-quality answer for the given question using only the provided search results (some of which might be irrelevant). question: what is mercy mercy me by marvin gaye about search results: document [1] …the album became gaye's first million-selling album launching two more top ten singles document [13] mercy mercy me (the ecology)... became regarded as one of popular music's most poignant anthems of sorrow regarding the environment document [20]... question: what is mercy mercy me by marvin gaye about answer: …the song was released as a single in august 1971 and became gaye's first million-selling single… shuffled prompt / model's response original prompt / our model's response this dependence persists even after randomly shuffling the document order, irrespective of its relevance to the query. we measure this dependence by computing the tf-idf similarity score between the response and each document (gold document originally at position 10). ## found-in-the-middle: modeling and isolating positional attention bias 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. ## improving long-context utilization with found-in-the-middle 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. ## related work retrieval augmented generation. while llms exhibit strong capabilities @xcite @xcite , their knowledge is inherently limited in its pretraining data, and they are observed to struggle in handling knowledge intensive tasks @xcite . 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 @xcite @xcite @xcite . although rag has powered many recent language model applications from questionanswering @xcite to automatic task completion @xcite , recent work show that llms tend to lost-in-the-middle, significantly hindering the full potential of rag @xcite . 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. ## discussion 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.