Focus Directions Make Your Language Models Pay More Attention to Relevant Contexts

ICLR 2026 Conference Submission15755 Authors

Keywords: Long context language models, Mechanistic interpretability

Abstract:

Long-context large language models (LLMs) are prone to being distracted by irrelevant contexts. The reason for distraction remains poorly understood. In this paper, we first identify the contextual heads, a special group of attention heads that control the overall attention of the LLM to the contexts. Then, we demonstrate that distraction arises when contextual heads fail to allocate sufficient attention to relevant contexts and can be mitigated by increasing attention to these contexts.
We further identify focus directions, located at the key and query activations of these heads, which control the amount of attention activated from the attention sink to the contexts.
With a proper amount of attention activation, the contextual heads could allocate more attention to relevant contexts.
Motivated by this, we introduce an automated magnitude control method that keeps attention activation within a proper range, enabling practical use of focus directions.
We comprehensively evaluate the effect of focus direction on various long-context tasks and find that focus directions can help mitigate the poor task alignment of long-context LLMs.
We believe our findings could promote further research on long-context LLM alignment.

Primary Area: foundation or frontier models, including LLMs

Submission Number: 15755