PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Weizhi Fei, Xueyan Niu, Guoqing Xie, Yingqing Liu, Bo Bai, Wei Han

arXiv:2501.12959v2cs.CL

TL;DR

Long-context LLM inference can incur high computational and memory demands. EHPC uses evaluator heads to select important prompt tokens, then passes the compressed prompt to the model. It reports state-of-the-art prompt-compression results, competitive long-context acceleration, and up to 40% improvement over direct inference on QA datasets.

  • Problem

    Long-context tasks increase computational and memory demands for transformer-based LLM inference, contributing to latency on lengthy prompts.

    Source: §1 · p. 1

  • Method

    EHPC identifies evaluator heads that score important prompt tokens, uses the first few layers to skim inputs, and forwards retained tokens for inference without training a separate model.

    Source: §1 · pp. 2–3

  • Results

    EHPC achieves state-of-the-art prompt-compression performance and competitive long-context acceleration versus KV-cache methods, improving direct inference by up to 40% on question-answering datasets.

    Source: §1 · p. 3, §5.2 · p. 10

  • Takeaways & Limitations

    EHPC is reported to reduce complexity, API costs, and memory use for commercial LLMs while accelerating long-context inference on local models.

    Source: §1 · pp. 2–3

    Prompt compression declines in performance on code-related and few-shot learning tasks.

    Source: §5.2 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

Long-context inference raises computational and memory demands, while existing prompt compressors can require repeated LLM calls or a separately trained model. EHPC selects important tokens using evaluator heads and reports strong prompt-compression and long-context acceleration results.

1 Introduction

  • Long-context tasks are increasingly needed, but transformer inference on lengthy prompts faces high computational and memory demands and increased latency.The introduction highlights long-document summarization, reasoning, and autonomous agents as relevant tasks.Source: §1 · p. 1
  • Prior prompt-compression methods may require repeated LLM calls for chunked processing or training a secondary model, whereas EHPC is training-free.The paper also attributes EHPC's efficiency to sparse attention scores and highly parallelizable prefill attention computations.Source: §1 · pp. 1–2
  • Evaluator heads are attention heads identified as capable of selecting long-input tokens significant for inference.The paper identifies them through synthetic-data experiments and investigates their properties using synthetic and real-world data.Source: §1 · pp. 2–3
  • EHPC uses evaluator-head scores in the first few layers to skim prompts, then passes only important tokens to the model for inference.The method uses the LLM's intrinsic attention mechanisms and is described as training-free and computationally efficient.Source: §1 · pp. 2–3
  • EHPC reaches state-of-the-art prompt-compression results on LongBench and ZeroScroll and improves direct inference by up to 40% on question-answering datasets.Table 1 compares average performance and latency on LongBench with a 2048-token compressed prompt.Source: Table 1 · p. 1, §1 · p. 3

2 Related Work

The related-work section distinguishes implicit KV-cache reduction from explicit prompt compression, then reviews representative methods in each area.

2 Related Work

  • Implicit approaches reduce the KV cache, while explicit approaches reduce the number of input tokens; explicit methods can also apply to black-box LLMs.The KV cache stores prior-token key and value matrices to reduce redundant calculations and improve decoding efficiency, but its memory needs grow with input length.Source: §2 · p. 3, §2.1 · p. 3

2.1 Implicit Token Reduction

  • KV-cache methods retain selected entries using attention patterns, including heavy hitters, adaptive policies, pooled prompt-attention clusters, or retrieval-head-specific caches.H2O, FastGen, SnapKV, and head-wise methods are described as examples of these strategies.Source: §2.1 · p. 3

2.2 Explicit Prompt Compression

3 Methodology

EHPC replaces a prompt with a shorter one while limiting changes to the model’s conditional distribution, then uses evaluator-head attention to select tokens for deletion. Pilot experiments identify these heads, and in Native Model Inference the two pre-filling stages can have lower complexity when both compression factors are at least two.

3.1 Prompt Compression Strategy

  • EHPC replaces the prompt with a shorter sequence while minimizing divergence between the model’s conditional distributions, using token deletion rather than generated summaries.Source: §3.1 · p. 4

3.2 Evaluator Heads

  • Evaluator heads are selected specialized heads whose attention scores estimate token utility; unlike retrieval heads, they assess prompt tokens rather than preserve the KV cache.Source: §3.2 · p. 4

3.3 Pilot Experiments for Detecting Evaluator Heads

  • On known-evidence long prompts, EHPC aggregates last-row attention over the evidence, averages scores by head, then selects top-k heads from the highest-scoring layer.Source: §3.3 · p. 5

3.4 Prompt Compression Using Evaluator Heads

  • Evaluator-head attention yields token utility scores; EHPC deletes non-essential tokens in original order and pools neighboring similarly scored tokens to improve readability.Source: §3.4 · pp. 5–6

3.5 Complexity

  • For NMI, compression pre-fills L/κ1 layers and then processes N/κ2 retained tokens; combined pre-filling complexity is lower when κ1, κ2 ≥2.Source: §3.5 · p. 6

4 Properties of the Evaluator Heads

The section tests whether evaluator heads can be identified using synthetic QA data and applied across real-world tasks. The identified heads appear sparse, practical on LongBench, and robust across task types.

Existence

  • Several heads in the middle layers identify key information in long-context QA, and evaluator heads are sparse and concentrated mainly in those layers.The pilot used Needle-in-a-Haystack data, inserting a sentence into a variable-length context and asking the model to retrieve it.Source: §4 · p. 7

Practicality

  • Across three LongBench datasets, four highest-scoring heads from the top-scoring layer outperform four comparison heads from that same layer.Table 2 compares the two head sets on real-world downstream tasks.Source: §4 · p. 7, Table 2 · p. 7

Robustness

  • Task-customized heads do not significantly improve their corresponding reasoning and coding tasks over the QA-identified heads, supporting task-agnostic robustness.The comparison covers multi-hop reasoning and code completion on LongBench.Source: §4 · p. 7, Table 3 · p. 7

5 Experimental Results

EHPC is evaluated for prompt compression and long-context inference acceleration. Across the reported benchmarks, it performs strongly on average, with particularly favorable QA results and a trade-off on code and few-shot tasks.

5.1 Prompt Compression Benchmark

  • EHPC achieves superior average performance across LongBench and ZeroSCROLLS under compressed-length constraints, and is competitive on summarization and few-shot tasks.The reported results also say it surpasses original prompts on average and improves QA-task accuracy; Table 4 covers target lengths of 2,000 and 3,000 tokens.Source: Table 4 · p. 8, §5.1 · p. 9
  • Each method reduces the average prompt from 14,354 to 2,048 tokens, and EHPC is faster than LongLLMLingua and LLMLingua-2 in the reported latency comparison.The comparison uses a 10-example RepoBench-P subset on a 40GB-VRAM GPU; the authors attribute EHPC’s lower latency to its use of the efficient pre-filling stage.Source: §5.1 · p. 9, Table 5 · p. 9

5 Experimental Results

  • EHPC compresses prompts before inference to reduce KV-cache memory and computation; in native-model inference, its acceleration is comparable to KV cache compression at the same compression ratio.The evaluation compares prompt and KV cache compression against H2O, SnapKV, and GemFilter using Llama-3.1-8B and Phi-3.8B.Source: §5.2 · pp. 9–10

5.2 Acceleration of LLM Inference

  • EHPC achieves the best average performance across baselines at different compression rates, with especially strong QA results, including up to 40% average improvement on Phi-3.8B QA tasks.The QA tasks include Single-Document QA and Multi-Document QA; prompt compression performs less well on code-related and few-shot learning tasks, where KV cache compression remains more advantageous.Source: §5.2 · p. 10, Table 6 · p. 10

6 Conclusion

The conclusion presents EHPC as a training-free prompt compression method that uses evaluator heads and pre-filling attention scores to select important tokens. Experiments across two benchmarks support its potential to reduce commercial API costs and accelerate long-text inference.

6 Conclusion

  • EHPC uses evaluator heads and pre-filling attention scores to select important tokens, with experiments across two benchmarks supporting its potential to reduce API costs and accelerate long-text inference.The method is evaluated in native-model and extended-model inference settings.Source: §6 · p. 10

A Background

This background introduces transformer attention and the two stages of LLM inference. It also describes KV-cache compression and pooling as ways to manage long-context processing and preserve coherent token selections.

Multi-head Attention

  • Multi-head attention aggregates information from input prompts, with attention weights determining which value tokens contribute to each attention block.The model computes attention from queries and keys, then applies the weights to values.Source: A Background · p. 14

LLM Inference and KV Cache Compression

  • LLM inference has a parallelized pre-filling stage that computes keys and values, followed by slower autoregressive decoding using the KV cache.The decoding stage loads the cached states and generates each output token through a forward pass.Source: A Background · p. 14
  • KV-cache compression methods use attention weights to score input tokens and decide which cache entries to retain as long-context cache size grows.The passage describes averaging attention-matrix rows to represent token scores.Source: A Background · p. 14
  • Pooling token scores can make selected tokens continuous rather than isolated, yielding more coherent semantics.The described pooling options include max and average pooling.Source: A Background · p. 14

B Hyper-parameters

The section lists evaluator heads selected for three LLMs and specifies the pooling settings used to evaluate them.

B Hyper-parameters

  • For Llama-3.1-8B-Instruct, which has 32 layers and 32 heads, the reported selection uses layer 13 and heads [18, 13, 21, 8, 11, 1, 4, 3].Source: B Hyper-parameters · p. 15
  • For CodeLlama-7B, which has 32 layers and 32 heads, the reported selection is layer 14 with heads [24, 3, 18, 7, 29, 2, 9, 1].Source: B Hyper-parameters · p. 15
  • For Phi-3.5-mini-3.8B-Instruct, which has 32 layers and 32 heads, the selected layer is 17, with chosen heads [7, 17, 30, 2, 6, 16, 25, 18].Source: B Hyper-parameters · p. 15
  • All experiments use average pooling because its observed difference from maximum pooling was negligible.Source: B Hyper-parameters · p. 15
  • The observed window and pooling kernel are 16 and 32 for Llama-3.1-8B-Instruct, and 4 and 32 for Phi-3.5-mini-3.8B-Instruct.A larger kernel typically produces a more continuous compressed context, motivating the general preference for larger kernels.Source: B Hyper-parameters · p. 15

C Complexity

The section describes pre-filling and decoding complexity for transformer inference, then notes that EMI has similar complexity when its inference model remains transformer-based.

C Complexity

  • Pre-filling complexity is O(LHdkN^2), for L layers, H heads, head dimension dk, and an N-token prompt.Source: C Complexity · p. 15
  • Decoding t tokens has overall complexity O(LHdk(Nt + t^2)).At generation step i+1, each head's softmax operation has complexity O(dk(N + i)).Source: C Complexity · p. 15
  • EMI complexity remains similar when the inference model is still transformer-based.Source: C Complexity · p. 15

D Datasets

The section introduces synthetic and real-world benchmarks for assessing LLMs' ability to handle long contexts.

D Datasets

  • The section presents synthetic and real-world benchmarks to assess LLM long-context capability.Source: D Datasets · p. 15

D.1 Synthetic Benchmarks

The synthetic benchmarks include retrieval from Needle-in-a-Haystack contexts and more varied, complex tasks in RULER.

D.1 Synthetic Benchmarks

  • Needle-in-a-Haystack tests long-context retrieval by inserting a sentence into a variable-length context and querying the LLM to retrieve it.Source: §D.1 · p. 15
  • RULER extends Needle-in-a-Haystack with varied needle types and quantities, plus tasks such as multi-hop tracing and aggregation.These tasks evaluate behaviors beyond direct searching within context.Source: §D.1 · p. 15

D.2 Real-world Benchmarks

This material describes English real-world benchmark coverage, prompt-compression baselines, synthetic data for evaluator-head identification, and an EHPC compression example.

D.2 Real-world Benchmarks

  • LongBench evaluation covers six task types across 16 English datasets, while ZeroSCROLLS covers summarization, question answering, and aggregation across ten datasets.ZeroSCROLLS evaluation uses the validation set because it is the only set with ground-truth data.Source: §D.2 · pp. 16–17
  • MultiNews has 200 entries evaluated with Rouge-L; NarrativeQA and Qasper each have 200 entries evaluated with F1.MultiNews clusters 2–10 same-topic news articles; NarrativeQA draws on books and movie screenplays, while Qasper covers NLP publications.Source: §D.2 · p. 16

E Baselines of prompt compression

  • Selective-Context ranks tokens by self-information, LLMLingua allocates segment budgets by perplexity, and LongLLMLingua adds question-aware compression.LLMLingua-2 uses fine-tuned smaller BERT models and an extractive compression dataset annotated by ChatGPT.Source: E Baselines of prompt compression · p. 18

F Additional Details for Evaluator Heads

  • Task-aware evaluator-head identification uses synthetic multi-hop tracing data and a separate long-code dataset with manually inserted completion evidence.The code data is unrelated to LongBench to prevent data leakage.Source: F Additional Details for Evaluator Heads · p. 18

G Example Prompt Compression Results Using EHPC

  • EHPC’s example shows a long story prompt alongside compressed excerpts that retain selected story details.The full prompt includes Kay and Gerda’s return journey; the excerpts retain details about warm milk and the robber girl.Source: G Example Prompt Compression Results… · p. 19

Related papers

Loading 2501.12959v2…