KeyDiff: Key Similarity-Based KV Cache Eviction for Long-Context LLM Inference in Resource-Constrained Environments

Thank you for your insightful feedback. We have responded to each of your comments below:

Regarding the incremental overhead of anchor updating: a naive approach would simply recompute the anchor from the current KV cache each pass through the model. As you suggest, an incremental approach is also possible. If we chose the key mean as the anchor, we recieve a single new key $K_c$ and we evict a single key $K_e$, we can rescale the new and evicted keys then add and subtract them appropriately: $K_A^{'} = K_A + \frac{K_c}{N} - \frac{K_e}{N}$A similar relation holds during block prompt processing. This means the incremental anchor update complexity is proportional to the number of new keys plus the number of evicted keys.

Regarding KeyDiff performance with quantized keys: preliminary internal analysis does show that cosine similarity remains reliable in downstream benchmarks. However, we do not have a comprehensive analysis of quantized eviction performance at this time.

Regarding the sliding window size appropriate for chat mode: our goal was to address the lack of recency bias in KeyDiff that is present in attention-based methods and show that the sliding window corrects for this. We believe the optimal window size would be application-dependent and likely needs domain-specific experimentation and benchmarking to determine the correct window size.

Regarding RAG integration: yes, KeyDiff can work with RAG. Prompt and RAG tokens could be naively combined and processed as a single long input prompt. KeyDiff could potentially be amended to retain prompt tokens and only evict RAG tokens to improve instruction-following.

Regarding on-device performance: we were unable to properly benchmark a full LLM inference pipeline comparable to the existing NVIDIA GPU TTFT measurements during the rebuttal window. However, we anticipate that performance will behave similarly, assuming the each KV cache eviction scheme does not introduce additional overhead on-device. To measure key scoring overhead on an edge device, we present relative latency on a recent Samsung Android device to compute the eviction scores of KeyDiff, TOVA, SnapKV and H2O, relative to the KeyDiff scoring latency with a cache budget of 512 KVs. Slight runtime variations can be attributed to kernel optimizations and proprietary hardware specifications.

We can see that KeyDiff matches competing methods at small cache sizes and outperforms them at larger cache sizes by a wide margin.

Regarding per-layer vs. per-head eviction: KeyDiff does in fact rank and evict keys on a per-head basis. This allows KeyDiff to attain excellent benchmark performance.

Thank you for your helpful comments. We have replied to each of them below.

Regarding a more comprehensive analysis of the hybrid attention + sliding window eviction strategy: we do not have a comprehensive analysis of each strategy on LongBench at this time, but we can include such results in the final version of the paper. Attention-based eviction methods tend to have an implicit recency bias [2], which gives them an advantage compared to KeyDiff, which strictly looks at key similarity and ignores temporality. We can also state this motivation for adding the sliding window more clearly in the main text of the final version.

Regarding replacing attention-based eviction with cosine similarity-based eviction: Attention-based eviction is limited by the number of tokens present in current block when evicting: it cannot know future attention scores based solely on the current block. By using cosine similarity instead, we have an eviction criteria that is independent of future tokens (i.e. key similarity) while remaining strongly correlated with future attention scores. We describe in Section 3.1 more precisely how we address the challenge presented in Section 2.4.

Regarding the paper organization, we will do our best to reorganize the paper layout for the final submission to reduce the number of forward references and disreuptions to the reading flow.

[2] Oren, Matanel, et al. "Transformers are multi-state rnns." arXiv preprint arXiv:2401.06104 (2024).

Thank you for your helpful feedback. We have responded to each of your comments below.

Regarding NIAH performance of KeyDiff: Thank you for your helpful comment; it prompted us to revisit the NIAH results. We re-ran our tests over more trials (increasing from the average of five random NIAH trials to fifty) to smooth out any noise in the tests. KeyDiff performs significantly better, and we plan to update the final paper with these results.

To summarize our findings here: below, we show the difference in NIAH scores between the eviction-free baseline and KeyDiff with a KV cache of 6k (averaged over fifty samples), i.e., $NIAH_{\mathrm{KeyDiff}} - NIAH_{\mathrm{Full}}.$ KeyDiff performs similarly to the eviction-free baseline for context lengths 1000 and 4222, where no eviction occurs, and performs roughly on par with the baseline at longer contexts. A key takeaway from this table is that the eviction-free baseline also cannot solve NIAH problems with 100% accuracy, so matching the model's NIAH performance is the best any eviction method can hope to achieve.

Regarding evaluation on the RULER benchmark: while we don't have our own results comparing KeyDiff with competing baselines on the RULER benchmark, we would like to highlight that NVIDIA's KVPress package on GitHub (and the corresponding kvpress-leaderboard on Huggingface) has recently adopted KeyDiff. KVPress demonstrates KeyDiff's competitive performance on RULER and thus provides independent, third-party validation of our method. Note that we are not affiliated with NVIDIA or KVPress in any way.

According to their evaluation -- including RULER as a subset -- KeyDiff achieves higher leaderboard scores than other eviction methods such as TOVA, SnapKV, QFilter, and Knorm, with a clear margin for Llama-3.1-8B-instruct and Qwen3-8B. Although KVPress benchmarking is done without applying the proposed block prompt processing (BPP), KeyDiff still outperforms competing methods. According to our results, KeyDiff performs best under the memory-constrained BPP setup; therefore, we expect a greater performance gap in KeyDiff's favor when the methods are evaluated with BPP enabled.

We also want to note that KeyDiff appears ranked fourth in the plot of KVPress. However, the methods ranked above KeyDiff are meta-heuristic approaches (e.g., DuoAttention, AdaKV, or SnapKV with Compressed Question) built upon an underlying score-based eviction scheme. KeyDiff is fully compatible with these meta-heuristics. Consequently, as a base scoring method, KeyDiff remains a strong KV cache eviction method.

Regarding the extension of KeyDiff to MLA: we believe that, with the exception of the first decoder layer, the geometric observations underpinning KeyDiff are intrinsic to each layer's token embeddings, since $W_K$ is just a linear projection. The challenge is how to analyze the RoPE operator, since it is clear that RoPE introduces necessary some temporal bias while allowing keys to cluster together on a per-head basis. Evicting token embeddings per-layer without performance degradation is the main challenge of future work.

Regarding adversarial evaluation: we would like to emphasize that KeyDiff operates on keys directly rather than tokens, i.e., after we have applied the RoPE and the learned projection $W_k$ to the input to a decoder block. This, of course, makes it tricky to construct an input prompt that will cause KeyDiff to fail: as one might infer from Figure 4, the keys and queries tend to cluster together regardless of the input prompt. Lemma 3.1 shows that attention scores are bounded by a function of the corresponding cosine similarity values, assuming a bounded key norm. If one can construct an input prompt that forces certain keys to have a massive norm relative to those already in the cache, KeyDiff will likely evict incorrect keys. However, we are not aware of a reliable method to construct such a prompt across all heads and layers.

Thank you for your thoughtful and considered feedback. We have replied to each of your comments below.

Regarding the correlation between negative key cosine similarity and attention scores: below we report the Spearman rank correlation for each layer of the Llama 3.2 3B Instruct model, averaged over each head. We averaged the correlations over several randomly selected samples from the musique task of LongBench. We observe high correlation across layers, validating our primary claim. We can also infer from this correlation that a large percentage of attention scores will be retained if we evict KVs according to key cosine similarity. We can include this table in the final version of the paper, along with similar data for other models.

Regarding retrieval-critical benchmarks: while we were unable to implement the Counting Star benchmark during the rebuttal window, we have initial results on the Phonebook Lookup task, which requires exact recall of a phone number corresponding to a target name [1] from list of names and numbers. We report results averaged over five random phonebook queries with KeyDiff and TOVA using a 6k cache budget. We see that KeyDiff generally outperforms TOVA, but still lags behind the baseline once it begins evicting KVs from the cache.

Regarding evaluation on the RULER benchmark: while we don't have our own results comparing KeyDiff with competing baselines on the RULER benchmark, we would like to highlight that NVIDIA's KVPress package on GitHub (and the corresponding kvpress-leaderboard on Huggingface) has recently adopted KeyDiff. KVPress demonstrates KeyDiff's competitive performance on RULER and thus provides independent, third-party validation of our method. Note that we are not affiliated with NVIDIA or KVPress in any way.

According to their evaluation -- including RULER as a subset -- KeyDiff achieves higher leaderboard scores than other eviction methods such as TOVA, SnapKV, QFilter, and Knorm, with a clear margin for Llama-3.1-8B-instruct and Qwen3-8B. Although KVPress benchmarking is done without applying the proposed block prompt processing (BPP), KeyDiff still outperforms competing methods. According to our results, KeyDiff performs best under the memory-constrained BPP setup; therefore, we expect a greater performance gap in KeyDiff's favor when the methods are evaluated with BPP enabled.

We also want to note that KeyDiff appears ranked fourth in the plot of KVPress. However, the methods ranked above KeyDiff are meta-heuristic approaches (e.g., DuoAttention, AdaKV, or SnapKV with Compressed Question) built upon an underlying score-based eviction scheme. KeyDiff is fully compatible with these meta-heuristics. Consequently, as a base scoring method, KeyDiff remains a strong KV cache eviction method.

Regarding the NIAH recall saturation: we have noted this behavior as well and also attribute it to intrinsic model limitations. We will add a baseline NIAH heatmap without cache eviction to validate this in the final version. We have not yet experimented with larger models, but anticipate the NIAH scores to improve. Below, we show the difference in NIAH scores between the eviction-free baseline and KeyDiff with a KV cache of 6k (averaged over fifty samples), i.e., $NIAH_{\mathrm{KeyDiff}} - NIAH_{\mathrm{Full}}.$ We can see that KeyDiff performs similarly to the eviction-free baseline for context lengths 1000 and 4222 and is roughly on par with the baseline at longer contexts.

Regarding PCA analysis interpretation: The full KeyDiff scoring criteria (without an anchor, labeled "Pairwise" in the ablation study) is the sum of negative cosine similarity among keys. In order for a key to score highly, it must have a very low cosine similarity with all other keys in the KV cache. This occurs when keys are more uniformly distributed in key space rather than tightly clustered in a single region. This means that a set of more uniformly distributed keys will lead to higher KeyDiff scores and ultimately higher attention, as validated by the correlation analysis above. The PCA plots are an effort to demonstrate that the keys retained by KeyDiff are more uniformly distributed than competing approaches, while evicted keys remain similarly uniform (i.e., not introducing bias).

Regarding the relation between key and beginning-of-sequence (BOS) token norms: This is a great point with a slightly nuanced answer. The BOS or sink token indeed has a smaller norm than other keys, due to its role in computing an approximate no-op in each attention head to avoid overmixing. We discuss this phenomenon in more detail in Section C.3. The smaller norm affects the inequality of Lemma 3.1, tightening the gap between cosine similarity with attention. However, the BOS token is typically not aligned with the majority of keys. Figure 16 shows that the key that is least similar to the key mean (which is typically the sink/BOS token) tends to have a high cosine similarity with the query mean. KeyDiff will always select this token due to its geometric irregularity, despite it having a lower norm than the remaining keys. This can also be seen in Figure 4, where the key with the highest KeyDiff score (the BOS token) can be seen as an outlier.

[1] Jelassi, Samy, et al. "Repeat after me: Transformers are better than state space models at copying." arXiv preprint arXiv:2402.01032 (2024).

Dear Reviewers, The author-rebuttal phase for NeurIPS 2025 ends in 48 h (Aug 8, 23:59 AoE). To safeguard the integrity of the conference—and to give every author the fair hearing you would want for your own work—please complete the following steps by the deadline: read the author rebuttal, engage in discussions, fill in "Final Justification" text box and update “Rating” accordingly. After these, one important step you should take is to submit “Mandatory Acknowledgement” button. Otherwise, NeurIPS may take some penalties to those reviewers who skip the acknowledgement step, such as dest rejection of your submissions. I know your schedules are demanding, yet your partication in the discussion makes a world of difference to the authors who have spent months on their research.

Thank you for your time

Thank you for your response.

• It's encouraging to see that the correlation between cosine similarity and attention score is close to 1.
• However, I would still appreciate seeing the cumulative attention score achieved by your methods, as this was not included in the rebuttal.
• Additionally, I remain somewhat concerned about the feasibility of KeyDiff when applied to more complex tasks. The gap towards dense attention for Phonebook is still significant and the RULER benchmark is only showing the results for 4K context length.