SEAL: Scaling to Emphasize Attention for Long-Context Retrieval
This work addresses long-context retrieval issues for LLM users, representing an incremental advancement by building on existing attention mechanisms.
The paper tackles the problem of quality degradation in large language models (LLMs) when handling long-context retrieval, and introduces SEAL, a learning-based method that adjusts attention heads to achieve significant improvements in retrieval performance across various tasks and models.
While many advanced LLMs are designed to handle long sequence data, we can still observe notable quality degradation even within the sequence limit. In this work, we introduce a novel approach called Scaling to Emphasize Attention for Long-context retrieval (SEAL), which enhances the retrieval performance of large language models (LLMs) over long contexts. We observe that specific attention heads are closely tied to long-context retrieval, showing positive or negative correlation with retrieval scores, and adjusting the strength of these heads boosts the quality of LLMs in long context by a large margin. Built on this insight, we propose a learning-based mechanism that leverages generated data to emphasize these heads. By applying SEAL, we achieve significant improvements in long-context retrieval performance across various tasks and models. Additionally, when combined with existing training-free context extension techniques, SEAL extends the contextual limits of LLMs while maintaining highly reliable outputs.