Source-linked AI summary
Found in the Middle: Calibrating Positional Attention Bias Improves Long Context Utilization
Cheng-Yu Hsieh, Yung-Sung Chuang, Chun-Liang Li, Zifeng Wang, Long T. Le, Abhishek Kumar, James Glass, Alexander Ratner, Chen-Yu Lee, Ranjay Krishna, Tomas Pfister
TL;DR
LLMs often miss relevant information in the middle of long inputs, while the causes of this lost-in-the-middle problem remain poorly understood. The paper connects it to U-shaped positional attention bias and calibrates attention to remove that bias. Found-in-the-middle improves relevance-based retrieval and RAG performance, including up to 15 percentage points over standard generation on NaturalQuestion.
Problem
LLMs struggle to locate relevant documents placed in the middle of long input prompts, and the factors underlying this widely observed problem remain poorly understood.
Method
The paper estimates positional attention bias by varying a fixed context’s prompt position and removes that bias from query attention scores to obtain calibrated attention.
Results
Found-in-the-middle improves document relevance ranking by up to 48 Recall@3 points and improves RAG performance by up to 15 percentage points over standard generation.
Takeaways & Limitations
Attention calibration enables LLMs to attend more faithfully to context relevance, including when relevant context occurs in the middle of the input.
Takeaways & Limitations
The proposed positional-bias model is a simplification, and intrinsic attention mechanisms may be more intricate, dynamic, learnable, or adaptive than it captures.
Abstract
from arXiv · showhide
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 a 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 15 percentage points. These findings open up future directions in understanding LLM attention bias and its potential consequences.
1 Introduction
The paper links lost-in-the-middle behavior to a U-shaped positional attention bias and introduces attention calibration to improve relevance-based context use. Experiments show calibrated attention improves document ranking and RAG performance, including when relevant information appears mid-context.
- Lost-in-the-middle describes LLMs’ difficulty locating relevant documents when they occur in the middle of input prompts, despite RAG’s broad use for external knowledge retrieval.
- LLMs assign higher attention to the beginning and end of prompts, and this positional bias correlates with U-shaped RAG performance and preferential use of boundary content.
- The proposed calibration estimates positional bias by varying a fixed context’s relative position, then removes that bias from query attention scores.
- Calibrated attention better tracks document relevance and improves ranking by up to 48 Recall@3 points, suggesting positional bias rather than inability to capture mid-context information is the main obstacle.
- Up to 15 percentage points improvement over standard generation on NaturalQuestion demonstrates the practical benefit of found-in-the-middle across two LLMs and two RAG tasks.
2 Positional attention bias overpowers mid-sequence context
In an open-domain question-answering setup, the model often relies on document position rather than relevance, producing a U-shaped attention pattern that disadvantages mid-sequence evidence. The exploratory studies connect this positional bias to document use during generation and show that relevant context can still receive elevated attention.
- 2 Positional attention bias overpowers mid-sequence context: The reproduced setup places one gold document among 20 retrieved documents and evaluates cases where the gold document is in the 10th, middle position.The model answers a user query using one relevant document and distractors, with the question repeated around the document sequence.
- 2.1 U-shaped attention bias: The model’s responses remain biased toward the first document when the gold document is positioned in the middle, even after document order is randomly shuffled.Figure 2 illustrates this qualitative first-position bias for gold documents at the 10th and 13th positions.
- 2.1 U-shaped attention bias: The response’s dependence on the first document persists irrespective of that document’s relevance to the query, measured using TF-IDF similarity.This quantitative dependence is reported with the gold document originally at position 10.
- 2.1 U-shaped attention bias: Average self-attention follows a U-shaped positional pattern: documents near the beginning and end receive higher weights than documents in the middle, even after shuffling.Weights are averaged across document tokens, decoder layers, and attention heads.
- 2.2 Does attention favor relevant context?: Attention is not purely positional: the model assigns higher attention to the gold document when it occupies the controlled position than when that position contains an uncontrolled document.This indicates an ability to prioritize relevant context in addition to the positional bias.
- 2.2 Does attention favor relevant context?: Documents receiving higher attention are more likely to be used in the model’s generation, supporting positional bias as a driver of its tendency to rely on certain document positions.The study compares documents in higher- versus lower-attention halves and reports a positive correlation with the document most likely used in the response.
3 Found-in-the-middle: modeling and isolating positional attention bias
This section models attention as the combination of document relevance and positional bias, then removes the positional component to rank contexts more faithfully. The resulting calibrated attention improves retrieval of relevant documents when they appear in the middle of long inputs.
- 3.1 Two main factors in model attention: The model’s attention is hypothesized to reflect both document relevance and document position within the input.The formulation treats relevance and positional bias as separate factors contributing to observed attention.
- 3.1 Two main factors in model attention: The hypothesized factors receive empirical support: 83% of tested pairs satisfy the positional condition and 72% satisfy the relevance condition.These tests use 100 randomly sampled NaturalQuestion examples containing 20 documents each.
- 3.1 Two main factors in model attention: A simple linear attention model achieves a Spearman rank correlation of 0.76, motivating its use for estimating positional bias.The authors note that a log-linear alternative reaches a competitive 0.75 correlation.
- 3.2 Disentangling positional attention bias: Subtracting attention to a consistent dummy document removes the positional bias and yields calibrated attention that estimates document relevance.The method compares each document’s attention with the baseline attention assigned to a dummy document at the same position.
- 3.2 Disentangling positional attention bias: The evaluation ranks documents using calibrated attention against vanilla attention, query generation, and relevance generation baselines.The comparison targets the challenging setting where the gold document is placed in the middle of the prompt.
- 3.2 Disentangling positional attention bias: Calibrated attention consistently outperforms vanilla attention and two re-ranking metrics on Recall@3 when retrieved documents vary in number.The evaluation uses NaturalQuestion with gold documents placed in the middle of the input context.
4 Improving long-context utilization with found-in-the-middle
Found-in-the-middle calibrates attention to reduce positional bias and improve long-context utilization. Across datasets, models, and retrieval settings, calibration improves performance over vanilla attention and complements document reordering in RAG pipelines.
- 4.1 Attention calibration: Calibration redistributes document-token attention so final attention reflects context relevance rather than positional bias, while temperature t controls disparity and normalization preserves total attention.The method rescales attention values within each document and allocates more attention to more relevant contexts.
- 4.2 Calibrated v.s. uncalibrated attention: In NaturalQuestion and SynthWiki, experiments use Vicuna-7b-v1.5-16k and tulu-2-7b with 10- or 20-document contexts.The models provide 16k- and 8k-token context windows, respectively.
- 4.2 Calibrated v.s. uncalibrated attention: Up to 15 percentage points: attention calibration outperforms vanilla attention across tasks and models, including 6–15-point gains when gold documents appear mid-sequence.Its performance curve exceeds the vanilla baseline in 22 of 24 cases.
- 4.3 Attention calibration in practice: Attention calibration complements reordering because reordering methods depend on correctly ranking documents and do not directly remove the model’s intrinsic positional bias.The paper evaluates prompt reordering, LongLLMLingua-rk, and attention sorting as existing alternatives.
- 4.3 Attention calibration in practice: Applying calibration to LongLLMLingua-rk achieves the highest performance across datasets and models, suggesting a complementary improvement for existing RAG pipelines.Reordering methods alleviate lost-in-the-middle effects, while calibration provides an additional improvement.
5 Related work
Related work addresses retrieval and long-context failures through RAG, context extension, prompting, splitting, and reordering. This paper instead targets the underlying long-context utilization capability through attention calibration.
- RAG retrieves external information to support knowledge-intensive tasks, but LLMs can still lose relevant documents placed in the middle of the input.
- Existing long-context methods include efficient training, finetuning, inference-time context extension, input splitting, and relevance-based reordering.Reordering places important content at the beginning or end, where models tend to focus more.
- These solutions improve performance by manipulating inputs but do not fundamentally improve LLMs’ underlying long-context utilization capability.
- Prior attention-focused work includes attention sorting and datasets designed to train models to focus on relevant documents, whereas this method mitigates the problem with off-the-shelf LLMs.
6 Discussion
The paper connects lost-in-the-middle to positional attention bias and addresses it by calibrating attention toward relevance rather than position. Experiments show gains on middle-position retrieval and additional gains when calibration is added to reordering pipelines.
- Attention calibration directly modifies the attention mechanism so LLMs attend to contexts according to relevance rather than position.
- The method improves performance over uncalibrated attention especially when relevant context occurs in the middle, and it further improves existing reordering pipelines.
Limitations
The study identifies limitations in its simplified bias model, computational cost, uncertain task-dependent benefits, and unresolved causes of positional attention bias.
- The positional-bias model may oversimplify mechanisms that are more intricate, dynamic, learnable, or adaptive to data and training.
- Calibration requires extra O(K) model forward passes at each position compared with vanilla generation, creating computational overhead.
- Removing positional bias may be counterproductive in tasks where attention to input beginnings and ends aligns with task or data structure.
- The work does not definitively identify whether pretraining data, transformer architecture, optimization, or another factor causes attention bias.
Ethical Statement
The study uses existing public datasets without collecting personal or sensitive data, but its pretrained-model foundation may propagate inherited biases.
- The method uses existing public datasets and collects no personal or sensitive data, while potentially inheriting and propagating biases from pretrained models.
A Multi-doc QA datasets
Experiments use NaturalQuestions and SynthWiki multi-document question-answering datasets, with controlled distractor documents and specified implementation choices for attention calibration.
- A Multi-doc QA datasets: The experiments use NaturalQuestions and SynthWiki, each pairing questions and answers with one gold document plus K−1 distractors, where K equals 10 or 20.
- A Multi-doc QA datasets: NaturalQuestions contains 2,655 queries with Wikipedia distractors ranked by decreasing query relevance, whereas SynthWiki contains 990 synthetic entries with randomly sampled and ordered distractors.
- A Multi-doc QA datasets: The datasets use public English Wikipedia or GPT-4-generated fictional Wikipedia paragraphs and are described as containing no uniquely identifying or offensive information.
- B Implementation details: Calibration is applied to all heads in the last 16 of 32 decoder layers because intervening in early layers may destabilize generation.
- B Implementation details: The calibration temperature is fixed at t = 5e−5 after experiments found performance robust to this term.
C Additional experiment results
Additional experiments compare linear and log-linear formulations for modeling attention, finding both competitive with rank correlations above 0.75.
- Rank correlations exceed 0.75 for both the linear and log-linear formulations, which perform competitively in the additional experiment.
- The additional experiment reports exact values in Table 5, while Table 4 presents rank correlations for the two formulations.
D Compute and inference details
Experiments use two 7B-parameter models with deterministic greedy decoding and require no training or hyperparameter search. The proposed calibrated-attention intervention stably improves RAG performance over reordering-based baselines.
- Experiments use Tulu-2-7B and Vicuna-7B-v1.5-16k, each containing 7B parameters.
- The experiments use greedy decoding without nondeterministic factors, so each experiment runs once.
- The method is inference-only and requires neither training nor hyperparameter searching.
- Experiments run on two NVIDIA A100 GPUs, with inference taking roughly 1 to 3 hours on both datasets.
- Calibrated attention stably improves RAG performance compared with existing reordering-based baselines.