L ong H eads: Multi-Head Attention is Secretly a Long Context Processor (2024)

Abstract: Large language models (LLMs) have achieved impressive performance in numerous domains but often struggle to process lengthy inputs effectively and efficiently due to limited length generalization and attention’s quadratic computational demands. Many sought to mitigate this by restricting the attention window within the pre-trained length. However, these methods introduce new issues such as ignoring the middle context and requiring additional training. To address these problems, we propose LongHeads, a training-free framework that enhances LLM’s long context ability by unlocking multi-head attention’s untapped potential. Instead of allowing each head to attend to the full sentence, which struggles with generalizing to longer sequences due to out-of-distribution (OOD) issues, we allow each head to process in-distribution length by selecting and attending to important context chunks. To this end, we propose a chunk selection strategy that relies on the inherent correlation between the query and the key representations, efficiently distributing context chunks to different heads. In this way, each head ensures it can effectively process attended tokens within the trained length, while different heads in different layers can collectively process longer contexts. LongHeads works efficiently and fits seamlessly with many LLMs that use relative positional encoding. LongHeads achieves 100% accuracy at the 128k length on passkey retrieval task, verifying LongHeads’ efficacy in extending the usable context window for existing models.

## introduction llms are usually required to handle tasks with long contexts, such as in-context learning @xcite , tool learning @xcite , and * equal contribution. ## method in this section, we describe how the longheads utilizes the inherent ability of multi-head attention to encode and generate long sequences without additional training. ## experiment we evaluate the proposed longheads primarily using the llama-2 @xcite considering its wide adoption and popularity. the effectiveness of longheads is evaluated on three kinds of tasks: language modeling, synthetic retrieval task and long context benchmark. ## conclusion we present longheads, a novel, training-free framework for efficiently processing long contexts in pre-trained llms. utilizing the intrinsic capabilities of attention heads, longheads smartly segments and assigns long text to relevant heads, streamlining the handling of extended sequences without extra computational load. experiment results validate longheads's superiority in re-stricted attention setups and its competitive edge against full attention methods when applied to the longbench suite. our approach paves the way for performance breakthroughs in long context llm operations, leveraging existing model structures to unlock new potential without further training.