You Only Use Reactive Attention Slice For Long Context Retrieval

TL;DR

This paper introduces YOURA, an attention-based retrieval method for long context LLMs that improves inference efficiency by up to 30% without sacrificing quality, using a novel reaction score heuristic.

Contribution

The paper presents a new attention-based retrieval technique, YOURA, with a unique reaction score heuristic and embedding-agnostic sentence yield, addressing limitations of embedding-based retrieval for long contexts.

Findings

Achieves up to 30% inference throughput improvement.

Maintains nearly identical quality to truncate-middle baseline.

Effective across multiple LLM models and datasets.

Abstract

Supporting longer context for Large Language Models (LLM) is a promising direction to advance LLMs. As training a model for a longer context window is computationally expensive, many alternative solutions, such as Retrieval Augmented Generation (RAG), have been used. However, most existing RAG methods adopt embedding-based retrieval that falls short on long contexts. To address such challenges, we propose an attention-based retrieval technique, You Only Use Reactive Attention slice (YOURA). YOURA leverages a novel retrieval heuristic called reaction score to rank the relevance of each sentence in the input context with the query sentence. Intuitively, we measure how the per-token attention score "reacts" to the query and greedily retrieves the most reactive sentences. Internally, YOURA generates a token-indexed vector (called reaction vector) for the whole input context. To map each sentence to the token-indexed vector, we propose an Embedding-Agnostic Sentence Yield (EASY), a best-effort token wiggling algorithm. We evaluate our retrieval technique on three open-source pre-trained LLM models across six LongBench QA datasets. Our technique achieves up to 30% vLLM inference throughput improvement for serving long-context queries with a nearly identical quality score to the simple yet effective truncate-middle approach.