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

We thank the reviewer for their insightful and positive feedback. We are encouraged that they found our paper well written and our method effective and adequately evaluated. We answer the questions below.

W1: the empirical evidence is strong, the paper lacks deeper theoretical analysis on the evaluator heads.?

We provide the properties of the evaluator heads in Section 4, which provides an empirical explanation of the found heads. We hope these discussions are useful to understand why our methods work.

Q1:EHPC selects heads from the highest scoring layer, would combining heads from multiple layers helps?

The ablation study in Table 4 below shows that selecting the top heads from multiple layers does not improve performance. Moreover, it forces us to compute attention for every chosen layer, which sharply increases latency. Thus, we simply select heads from the highest scoring layer.

Table 4; The ablation study on single/multi-layer. A single layer represents selecting the heads from the layer with the highest scores. Multi-layer selected the top heads from any layers.

We thank the reviewer for their constructive and positive feedback. We are encouraged that they found our method efficient, adequately evaluated, and original. We answer specific questions in the following and will incorporate all feedback in the final version.

W1.1: I would implore the authors to evaluate on more models covering different sizes and architectures, like 30B or 70B models, if feasible, and architectures like MoE models

Thank you for the valuable feedback. We include properties of evaluator heads on Llama2 13B 64k below. If we have more computation, we will include the models with larger sizes and MoE models.

• Existence

The score ratio of the selected layer is 70%, and the score ratio of selected heads at selected layers is 55%. This shows that our selected evaluator heads have significantly high scores

• Generalizability

We present the results of evaluator heads on real-world tasks in Table 2 below. The findings indicate that NMI (Llama2 13B 64K) effectively understands the compressed data. However, due to limited memory, we had to truncate the context length to 8K, which resulted in relatively poor performance.

Table2; The results of real-world tasks in LongBench. For Llama2-13b-64k, we truncate the context length to 8k.

• Robustness

llama 2 13b 64k: 
Layer of evaluator heads: 13 layer
Needle QA :[23, 12, 13, 14, 24, 21, 26, 17],
Variable Tracking (reasoning): [26, 21, 17, 23, 24, 14, 36, 11],
Common heads: [14, 17,21, 23,4,26]

Due to limited memory, we just conduct the results of variable tracking and Needle QA. The results show that these heads found by different tasks have 6/8 common heads.

W1.2: The generalizability study in Section 4 does not provide details on how the downstream performance is calculated on the given set of tasks with a subset of the heads. Additionally, all the subsets of tasks shown are from LongBench. It would be nice to see some tasks from RULER as well.

• Sure. Below we spell out the exact pipeline we used to evaluate downstream performance.

1. Given a prompt and a set of heads, we first compute the attention scores of these heads.

2. Following Eq. 1 in the paper, we compute the $s$ score for every token in the prompt. Then we select the top-k tokens to form the compressed prompt.

3. The compressed prompt is fed to the LLM for inference. Finally, task-specific metrics are computed on the model’s output. . These baseline tasks are discussed in Appendix H.

• For the generalizability, we aim to evaluate real-world tasks while the tasks in Ruler are synthesized. We also include the tasks from Infinity Bench [2] in Table 4 (Section 4, main paper).

W1.3: The robustness study does not tell me whether the actual heads identified using different tasks are the same or different. Even if downstream performance does not change, the actual heads identified might be different. This is a crucial distinction between saying the evaluator heads are task agnostic versus evaluator heads identified using task X have a good overlap with evaluator heads of task Y.

Llama 3.1 8B Instsrut :
Layer of evaluator heads: 13 layer
Needle QA : [1, 4, 8, 11, 13, 17, 18, 21]
Variable Tracking (reasoning):  [1, 2, 8, 9, 11, 13, 18, 21],
Code completion:  [3,4,5,6,13,18,21,27]
Common heads: [8,13,18,21] 

mistral-nemo-instruct-2407: 14 layer
Needle QA :  [10, 11, 25,  8, 29, 24, 17,  9]
VR : [17, 25, 11, 10, 16,  0, 30, 29]
Code: [29, 25, 17, 30, 10, 11,  8, 28]
Common heads: [10, 11, 17,25, 29]

Thank you for the question. The above results show that:

For Llama 3.1 8B Instruct, there are 4 heads in common on three tasks in the selected 8 heads.

For mistral-nemo-instruct-2407, there are 5 heads in common on three tasks in the selected 8 heads.

W2: A small discussion on whether and how the NMI performance can be improved would be helpful.

Thank you for the valuable suggestion. We provide the discussion below and have included this as future work in the revised manuscript.

Table 6 (main paper) demonstrates that NMI struggles with few-shot learning and code tasks, which demand a high level of context integrity. Existing compression methods prioritize semantic salience at the expense of fluency and structure, leading to downstream accuracy drops. Preserving both structure and fluency during compression is therefore a critical direction for future work.

W3: EHPC does not apply to decode heavy workloads

We provide the results of Infinity Bench in the table above, which covers 32 k-input summarization and long-form generation tasks. For tasks involving heavier decode workloads, we consider for future work.

Q1: I understand that the purpose of evaluator and retrieval heads is different, but is there an overlap between them - i.e. are most evaluator heads the same as retrieval heads?

Our evaluator heads are designed in the same layer, whereas the retrieval heads lie in different layers. Consequently, the two sets are almost disjoint; on Llama-2-7B-80K and llama2-13b-64k we observe zero overlap between the heads reported in [1] and the ones we identified as evaluators.

Q2: When identifying evaluator heads, what is the rationale behind extracting only the last row of the attention matrix?

We follow common practice and use the last row of the attention matrix as the scoring vector for each token. This row contains the attention scores that the final (query) token assigns to all earlier positions, which is exactly the distribution the model itself uses for next-token prediction. Treating these scores as the token’s “importance” therefore aligns our evaluator with the model’s own generative signal.

Q3; In the robustness study, can you confirm if the actual evaluator head indices identified over different tasks are the same or have significant overlap?

Sure, see response to W1.3.

Q4: Can you provide details on how the downstream performance was measured in the generalizability study with a limited number of attention heads?

See the response to W1.2.

3. The compressed prompt is fed to the LLM for inference. Finally, task-specific metrics are computed on the model’s output. . These baseline tasks are discussed in Appendix H.

Q5: I think the authors should use LongBench v2, which is an updated version of the LongBench dataset with much larger context lengths to show the effectiveness of their method further.

Thank you for the valuable suggestion. We are still working on evaluating EHPC LongBench v2 and hope to include more results in the final version. We provide the results of 32k on the Infinity Bench for validating our performance on a larger context length. The results in Table 3 show that our NMI setting achieves consistently strong results compared with SnapKV.

Table 3; The results of Infinity Bench. We evaluate SnapKV and our NMI setting on the 32k token length.

Q6: A small discussion on whether and how the NMI performance can be improved would be helpful.

See response to W2.

Q7: In Table 5, the inference time (second column) is different for different prompt compression methods. Why is that the case?

In the second column, the 2048 tokens represent the target token length. However, the compression methods in these baselines often do not strictly equal 2048, resulting in varying inference times.

[1] Retrieval Head Mechanistically Explains Long-Context Factuality, ICLR, 2025. [2] \infty${B}ench: Extending Long Context Evaluation Beyond 100{K} Tokens, ACL, 2024

We thank the reviewer for their insightful and positive feedback. We are encouraged that they found our method compelling, easily applicable, and effective. We answer specific questions below and will incorporate all feedback in the final version.

Q1: Does NMI perform better than EMI when the computation budget is controlled the same?

Thank you for the insightful question. We conduct the experiments to infer the compressed prompt in Table 1 below. The diagonal results are NMI, while the others are EMI. For averaged results, we find that NMI is better than EMI because the same model shares a similar pattern.

Table 1: The averaged results with 2048 compressed prompts on EMI on different settings.

W1: The methodology part is somewhat confusing. For instance, when detecting evaluator heads in Section 3.4, the authors use only the last row of the attention matrix, but all the rows are included in the calculation of utility scores (Equation 1). Moreover, some denotations have mistakes. For instance, in line 157, $A\in \mathbb{R}^{N\times N}$, not $A\in \mathbb{R}^{H\times L}.$

• Thank you for catching this; the typo in Line 157 has been fixed in the revised manuscript.

• In Equation 1, we select the last few rows from N_r to N, where N is the prompt length and N_r is a parameter adjusting the number of selected rows. Selecting the last few rows is generalized from selecting the last row, where previous research [1] has shown that this can capture more information.

• When detecting evaluator heads in plot experiments, we just select the last row in Section 3.4 for convenience. While computing scores by the evaluator heads, we select the last few rows in Equation 1 for strong performance.

W2: The differences between EHPC and existing attention-based compression methods (e.g., SnapKV) are not clearly discussed. Therefore, the novelty of the method is not very clear.

Thank you for highlighting the need for a clearer distinction. We provide two key differences below and have included them in the revised manuscript.

• 1. SnapKV and other implicit methods retain the original prompt length and only sparsify the KV cache after computing the full attention map. EHPC, in contrast, removes tokens to reduce the number of tokens in the prompt. The compressed prompts can be applied to other language models, which SnapkV and other implicit methods fail to do.
• 2. We design pilot experiments to find the evaluator heads for improving the compression quality and accelerating the compression time. In Table 6, we have also empirically compared our method with acceleration methods under the same KV cache memory usage.

---

[1] SnapKV: LLM Knows What You Are Looking for Before Generation, NeurIPS, 2024

We thank the reviewer for their overall positive feedback. After carefully reading the comments, we noticed that some of the raised points (e.g., Seer’s effect on end-task performance) do not seem to apply to our submission. We suspect a possible mix-up with another paper’s review. Nevertheless, we have addressed every question below to the best of our understanding. We would be grateful if the reviewer could share any clarifying remarks during the discussion period, and we are happy to provide further details as needed.

1. Adaptation to vision-language models or audio-augmented transformers where KV size and reuse patterns.

Reducing the token load for multi-modal models is an important and pressing issue. While the current paper focuses on LLMs, we have not yet had the opportunity to conduct more experiments on models such as vision-language and audio-augmented transformers. We are optimistic that the proposed method has the potential in such multi-modal models, and we hope to have more definite results in future research.

2. Adaptive $\lambda$ weights in the score function.

Thank you for the question; however, we are not sure if the reviewer is referring to the utility scores in Eq. (1). The score is computed with a pooling operation, where the set of evaluator heads $\mathcal{C}_f$ is chosen according to the accumulated score of the evidence in Eq. (2). We selected the top-$k$ heads according to line 228, where we control the cardinality of $\Lambda$ to be less than $k$.

3. End-task performance (e.g., summarization ROUGE scores, QA accuracy) beyond token-level match.

Thank you for the question. In line with previous research, we have conducted experiments on the LongBench benchmark, which involves six tasks: single-document QA, multi-document QA, summarization, few-shot learning, code completion, and synthetic tasks for retrieval and counting. The results with respect to task performance can be found in Tables 4 and 5 (main paper).

4. Incorporation of existing inference accelerations like speculative decoding, FlashAttention, and prefill skipping.

Thank you for the insightful suggestion. Prompt compression is orthogonal to existing inference-acceleration techniques, so the two can be readily combined: after our method shortens the prompt, any subsequent speed-up (e.g., speculative decoding, FlashAttention, prefill skipping) can be applied whether the serving model is the same or different, and whether it runs in the cloud or on the edge. Table 6 already shows that our approach yields the best average performance under KV cache constraints; coupling it with these orthogonal optimizations will therefore amplify the gains at inference time. We will incorporate this discussion in the final version.

5. Performance across model sizes and the relationship between cache threshold and model size.

Thank you for the valuable suggestions. However, supporting models with a larger size involves substantial work and timehave lots of work to do; we are still working on this. If we have additional results, we will include them intoin our revised manuscripts.