[Paper Review] Attention Sorting Combats Recency Bias In Long Context Language Models

Alexander Peysakhovich, Adam Lerer|arXiv (Cornell University)|Sep 28, 2023

Topic Modeling4 citations

TL;DR

This paper proposes attention sorting, a lightweight inference-time technique that improves long-context language models by reordering context documents based on attention weights to reduce recency bias. By iteratively sorting documents so relevant ones move toward the end of context, the method boosts extractive QA performance—especially for open-source models—doubling accuracy at 30k context length and enabling them to match proprietary models like Claude-2.

ABSTRACT

Current language models often fail to incorporate long contexts efficiently during generation. We show that a major contributor to this issue are attention priors that are likely learned during pre-training: relevant information located earlier in context is attended to less on average. Yet even when models fail to use the information from a relevant document in their response, they still pay preferential attention to that document compared to an irrelevant document at the same position. We leverage this fact to introduce ``attention sorting'': perform one step of decoding, sort documents by the attention they receive (highest attention going last), repeat the process, generate the answer with the newly sorted context. We find that attention sorting improves performance of long context models. Our findings highlight some challenges in using off-the-shelf language models for retrieval augmented generation.

Motivation & Objective

• To investigate why long-context language models underutilize relevant information placed early in context, despite attending to it more than distractors.
• To address the performance degradation in retrieval-augmented generation (RAG) when relevant documents are distant from the response.
• To develop a practical, inference-time method that improves model performance without retraining or architectural changes.
• To evaluate whether attention weights can serve as a reliable signal for reordering context in long-context tasks.

Proposed method

• The method performs one decoding step and computes average attention weights across all layers and heads for each document in context.
• Documents are sorted in descending order of attention score, with the highest-attentive document moved to the end of the context.
• The process is repeated iteratively—typically two passes—before final generation with the re-sorted context.
• The approach is applied to a synthetic, controlled QA benchmark (SynthWiki) with one relevant document and multiple distractors.
• It leverages the fact that even when models fail to use early information, they still attend more to relevant documents than distractors at the same position.
• The technique is evaluated on both open-source and proprietary long-context models, with ablation on number of passes and model variants.

Figure 1: Performance of all long-context models that we study on question answering degrades when the relevant information is embedded in a long context of irrelevant distractor text.

Experimental results

Research questions

• RQ1 Does attention weight distribution reflect document relevance even when models fail to use early information in generation?
• RQ2 Can iterative reordering of context based on attention scores improve long-context QA performance?
• RQ3 Why do open-source models degrade more than proprietary models as distractor count increases in RAG tasks?
• RQ4 How effective is attention sorting in mitigating recency bias in models pre-trained with a recency-inductive bias?
• RQ5 Can attention-based reordering outperform baseline models without fine-tuning or architectural changes?

Key findings

• Attention sorting significantly improves QA accuracy on long-context models, with open-source models seeing performance boosts of over 100% at 30k context length.
• After two passes, attention sorting moves relevant documents into the 'useable' part of the context window, where attention is most effective.
• TogetherLlama-7B-32k, fine-tuned on long-context QA, performs best and matches or exceeds proprietary models like Claude-2 and GPT-3.5.
• The method improves performance across all tested open-source models, including WizardCoder and YaRN, especially with repeated passes.
• Models with stronger recency bias (e.g., Llama-2 base) suffer greater degradation with distractors, but attention sorting effectively counteracts this.
• The approach works best when the model has been fine-tuned for long-context QA, suggesting synergy between model inductive bias and the method.

Figure 2: QA accuracy on SynthWiki as a function of the position of the relevant document in the context. We see a replication of the ‘lost in the middle’ [ 15 ] effect on this dataset in which accuracy is lower when the relevant information is in the middle of a long context. The recency (informati