In plain words: It picks sentences by how much the model's own attention to each token shifts when the query is added, instead of matching sentence embeddings to it. It sped up answering long texts by up to 30% over cutting the middle, with nearly identical quality.
Abstract · You Only Use Reactive Attention Slice For Long Context Retrieval
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.
Yun Joon Soh, Hanxian Huang, Yuandong Tian, Jishen Zhao
arXiv:2409.13695 · cs.CL, cs.AI, cs.IR · submitted Sep 3, 2024
abstract · pdf · html