Pause-Tuning for Long-Context Comprehension: A Lightweight Approach to LLM Attention Recalibration

TL;DR

Pause-tuning is a lightweight method that improves long-context comprehension in large language models by redistributing attention through fine-tuning with artificially inserted pause tokens, significantly enhancing performance on lengthy inputs.

Contribution

This paper introduces pause-tuning, a novel fine-tuning technique that addresses the Lost-in-the-Middle problem by redistributing attention in LLMs for better long-context understanding.

Findings

Significant performance improvements on the Needle-in-a-Haystack benchmark.

LLaMA 3.2 3B Instruct model improves by 10.61%.

LLaMA 3.1 8B Instruct model improves by 3.57%.

Abstract

LLMs have demonstrated remarkable proficiency in understanding tasks but continue to struggle with long-context comprehension, particularly with content located in the middle of extensive inputs. This limitation, known as the Lost-in-the-Middle (LITM) problem, hinders models from fully processing and utilizing information across lengthy contexts. To address this issue, we introduce pause-tuning, a technique that redistributes attention to enhance comprehension of long-context inputs. Our approach involves fine-tuning language models on datasets with artificially inserted pause tokens, which serve to segment the input into smaller, more manageable parts. We evaluate pause-tuning against alternative approaches using the Needle-in-a-Haystack benchmark, where models must retrieve information embedded within contexts of up to 128K tokens. Experimental results demonstrate significant performance gains, with the LLaMA 3.2 3B Instruct model and the LLaMA 3.1 8B Instruct model improving by 10.61% and 3.57% respectively on average, suggesting that pause-tuning successfully enhances attention redistribution and improves long-context retention. The code and data are available at https://anonymous.4open.science/r/LITM-PauseTokens-7357.