Efficient Prompt Compression with Evaluator Heads for Long-Context Transformer Inference

TL;DR

This work tackles the high computational and memory demands of long-context processing in transformer-based LLMs by introducing EHPC, a training-free prompt compression method that leverages specialized evaluator heads. Evaluator heads identify salient tokens during the pre-filling stage, enabling a compact prompt with minimal loss in inference quality. EHPC achieves state-of-the-art results on prompt compression benchmarks and substantially accelerates long-context inference, offering both Extended Model Inference (EMI) and Native Model Inference (NMI) deployment modes and competitive performance relative to KV-cache-based acceleration. The approach reduces API costs for commercial models and lowers local memory requirements, with robust performance across QA, code, and multi-hop tasks, illustrating practical impact for real-world long-context applications.

Abstract

Although applications involving long-context inputs are crucial for the effective utilization of large language models (LLMs), they also result in increased computational costs and reduced performance. To address this challenge, we propose an efficient, training-free prompt compression method that retains key information within compressed prompts. We identify specific attention heads in transformer-based LLMs, which we designate as evaluator heads, that are capable of selecting tokens in long inputs that are most significant for inference. Building on this discovery, we develop EHPC, an Evaluator Head-based Prompt Compression method, which enables LLMs to rapidly "skim through" input prompts by leveraging only the first few layers with evaluator heads during the pre-filling stage, subsequently passing only the important tokens to the model for inference. EHPC achieves state-of-the-art results across two mainstream benchmarks: prompt compression and long-context inference acceleration. Consequently, it effectively reduces the complexity and costs associated with commercial API calls. We further demonstrate that EHPC attains competitive results compared to key-value cache-based acceleration methods, thereby highlighting its potential to enhance the efficiency of LLMs for long-context tasks.

Figures (3)

• Figure 1: Visualization of attention scores from a single attention head during inference on the "Needle-in-a-Haystack" long-text benchmark. This benchmark requires the LLM to follow instructions and retrieve "needles" -- specific pieces of information randomly inserted into a long text -- to answer a given question. The evaluator heads are identified as those that accurately locate the relevant facts, thereby achieving high scores.
• Figure 2: Heatmap of evidence scores for three different LLMs in the pilot experiment, illustrating scores across layers and heads, with heads re-ranked in descending order for clarity.
• Figure 3: Illustration of the proposed method. (a) LLM inference comprises two stages: the pre-filling stage and the decoding stage. (b) The proposed prompt compression approach leverages the efficiency of the pre-filling stage, thereby reducing inference latency for both stages during inference with compressed context.