PaperlayerView exact revision on arXiv

Source-linked AI summary

Information-Aware KV Cache Compression for Long Reasoning

Jushi Kai, Zhuiri Xiao, Alexandra Birch, Zhouhan Lin

arXiv:2606.26875v1cs.CLcs.AI

TL;DR

Long-context reasoning makes KV-cache compression important, but attention-based token selection can miss information needed for future reasoning. InfoKV adds forward influence and entropy-aware signals to attention-based compression, consistently outperforming existing attention-based methods across long prefilling and decoding benchmarks.

  • Problem

    Attention-based KV compression captures short-range relevance but may miss tokens that remain important for evolving future reasoning trajectories.

  • Method

    InfoKV combines Forward Influence, predictive entropy, layer-wise representation evolution, and attention scores to select tokens for KV-cache compression.

  • Results

    InfoKV consistently outperforms existing attention-based compression methods across long prefilling and decoding benchmarks, models, and reasoning tasks.

  • Takeaways & Limitations

    Forward-looking information signals complement attention and improve KV-cache compression for long-context reasoning.

  • Takeaways & Limitations

    Entropy remains an indirect approximation of future utility rather than an explicit optimization objective.

Abstract

from arXiv · show

Reasoning capability has advanced rapidly in large language models (LLMs), leading to an increasing size of key-value (KV) cache in both prefilling and decoding stages. Existing KV cache compression methods mainly rely on attention weights to estimate token importance. While attention effectively captures contextual relevance, it overlooks complementary information-theoretic signals related to predictive uncertainty and token informativeness. In this paper, we revisit token importance from a forward-looking perspective and introduce \textit{Forward Influence}, a metric that measures how compressed tokens affect future contexts. Our analysis reveals that tokens selected by attention scores mainly influence nearby contexts, whereas tokens associated with high predictive uncertainty exhibit substantially stronger influence on distant future contexts. Based on the observation, we propose \textbf{InfoKV}, an entropy-aware KV cache compression framework that incorporates information-theoretic signals. It combines token-level predictive uncertainty with layer-wise representation evolution and integrates the resulting entropy scores with attention scores during reasoning. Experiments on long-context reasoning benchmarks with Llama-3.1, Llama-3.2, and DeepSeek-R1 demonstrate that InfoKV consistently outperforms existing attention-based KV compression methods in both long prefilling and decoding scenarios.

1. Introduction

The introduction identifies KV-cache growth as a major bottleneck in long-context reasoning and argues that attention-based compression misses tokens with important future effects. It presents InfoKV, which combines predictive entropy, layer-wise representation evolution, and attention to preserve more informative tokens.

  • Motivation: Long-sequence reasoning is costly because attention computation grows quadratically while KV-cache memory grows linearly with sequence length.The bottleneck is especially pronounced when reasoning handles thousands of tokens.
  • Existing Approaches: Existing KV-cache compression methods typically retain tokens with large attention weights from a fixed recent-token observation window.These methods selectively preserve a subset of past tokens and discard the remainder.
  • Limitations: Attention-based compression is short-term and backward-looking, primarily capturing local dependencies rather than tokens crucial to evolving future reasoning trajectories.This mismatch is especially problematic during long decoding, where reasoning paths change dynamically across generation steps.
  • Forward Influence: Forward Influence measures how removing a token from the KV cache changes future predictive distributions, providing a forward-looking measure of token utility.The introduction contrasts attention’s emphasis on recent-context relevance with tokens’ effects on future generation steps.
  • InfoKV: InfoKV combines token-level predictive entropy, layer-wise representational evolution, and attention scores for entropy-aware KV-cache token selection.Experiments report consistent gains over attention-based compression on long prefilling and substantial improvements in long decoding across multiple reasoning benchmarks.

2. Related Work

Prior KV cache compression methods predominantly use attention-based token pruning or merging, but these backward-looking heuristics are limited for long decoding. Recent work therefore extends compression online and incorporates information-theoretic signals to measure intrinsic token importance.

  • KV Cache Compression: SnapKV, PyramidKV, and FastKV estimate token importance from attention weights and discard less-attended tokens through heuristic pruning.Other approaches merge tokens to approximate full-cache attention.
  • KV Cache Compression: Attention-based compression reduces memory usage but is backward-looking and mainly benefits long prefilling tasks with short answers.This limitation motivates methods designed for longer reasoning and decoding contexts.
  • Compression for Long-decoding: Recent studies extend KV cache compression from prefilling to decoding because long reasoning paths increase decoding latency and KV cache growth.RPC generalizes SnapKV to online decoding by periodically operating during generation.
  • Information Signals for Token Importance: Information-theoretic methods characterize token importance intrinsically, with Selective Context using self-information to quantify informativeness and prune redundant input content.Unlike attention weights, a token’s carried information is described as its native importance.

3. Methodology

The methodology reframes KV-cache token importance around predictive uncertainty and future influence rather than attention alone. It defines Forward Influence and develops InfoKV by combining entropy-based informativeness, layer-wise representation evolution, and attention for per-layer token selection.

  • Attention-based compression preserves tokens important to recent contexts but assumes they remain important for future generation.
  • Higher predictive entropy indicates lower confidence and richer contextual information, motivating entropy as a signal for long-term token utility.
  • Forward Influence measures a token’s effect on a future context as the average KL divergence between predictions with and without that token.Influence estimation uses Llama-3.1-8B-Instruct, with all layers sharing the same token choice during compression.
  • Attention-selected tokens mainly affect nearby contexts, whereas high-entropy tokens exert substantially stronger influence on distant future contexts.Combining attention and entropy balances short-range and long-range influence across short and long future horizons.
  • InfoKV computes layer-specific entropy from predictive informativeness and semantic representation evolution, then combines it with attention to retain top-ranked KV-cache tokens.Larger hidden-state shifts indicate unresolved semantic information, while Top-k Restricted Entropy limits uncertainty estimation to the most probable tokens.

4. Experiments

InfoKV is evaluated for long prefilling and decoding across long-context reasoning, instruction-following, mathematical-reasoning, and coding benchmarks. It achieves state-of-the-art or highly competitive prefilling performance and consistently improves over attention-based baselines, including RPC, in long decoding.

  • Long prefilling: Long prefilling evaluates Llama-3.1-8B-Instruct and Llama-3.2-3B-Instruct on LongReason across 16k–64k contexts and compressed prompt KV caches.LongReason contains extensive supporting evidence and distractor information, stressing preservation of reasoning-critical information under limited KV budgets.
  • Long decoding: Long decoding evaluates DeepSeek-R1-Distill-Qwen-7B and DeepSeek-R1-Distill-Llama-8B on IFEval, AIME 2024, and LiveCodeBench while periodically compressing generated-token KV caches.Models generate reasoning steps and final answers with a maximum output length of 32768 tokens.
  • Long prefilling: InfoKV achieves SOTA or highly competitive performance across most LongReason settings, demonstrating the effectiveness of entropy-aware information signals.Results are reported across multiple context lengths in Table 1.
  • Long prefilling: InfoKV consistently outperforms attention-based baselines at both 40% and 20% cache budgets on Llama-3.1-8B-Instruct, with larger gains at longer sequence lengths.The results indicate that entropy-aware selection better retains globally informative tokens for future reasoning than recent-attention heuristics.
  • Long decoding: InfoKV outperforms RPC across all three long-decoding task categories, and at 25% and 12.5% retaining ratios on IFEval surpasses the full cache of R1-Distill-Llama-8B.The result suggests that long reasoning trajectories contain redundancy and that retaining all historical tokens can introduce distracting or less informative contexts.

5. Analysis

The analysis finds that τ=1 and top-256 restricted entropy perform best, while a moderate entropy–attention balance with α=0.9 gives the strongest compression results. Layer-adaptive budgets can help some models but are less consistent than uniform allocation.

  • Choice of τ: τ=1 achieves the best overall performance and more stable behavior across settings at a 40% retaining ratio.Increasing τ reduces the contribution of layer-wise representation distance and makes entropy scores rely more on final-layer predictive uncertainty.
  • Top-k Restricted Entropy: k=256 achieves the best performance across different τ values, while top-k restricted entropy consistently outperforms full-vocabulary entropy computation.Restricting entropy to the most probable tokens avoids contributions from low-probability tokens that provide limited useful uncertainty information.
  • Balance between Entropy and Attention: α=0.9 achieves the best results for both full-vocabulary and top-256 restricted entropy, whereas pure attention selection performs worse.Further reducing α degrades performance, indicating that attention-based short-range dependency remains important alongside entropy-based uncertainty.
  • Layer-wise Adaptive Budgets: Adaptive allocation assigns larger KV budgets to higher-entropy layers and compresses more redundant layers more aggressively.This strategy is motivated by richer uncertainty and broader contextual information in early and middle layers, compared with increasing confidence and redundancy in higher layers.
  • Layer-wise Adaptive Budgets: At a 25% retaining ratio, adaptive allocation improves R1-Distill-Llama-8B over uniform allocation but causes less consistent gains on R1-Distill-Qwen-7B.The authors attribute degradation to excessively imbalanced budgets over-compressing certain layers and therefore adopt the uniform strategy.

6. Conclusion

The paper introduces Forward Influence to measure how compressed tokens affect future predictive distributions, motivating InfoKV’s entropy-aware token selection. InfoKV combines predictive entropy, layer-wise representation evolution, and attention scores for long-context reasoning.

  • Forward Influence measures the effect of compressed tokens on future predictive distributions.
  • Attention weights mainly capture short-range dependencies, while high-predictive-uncertainty tokens exert substantially stronger influence on distant future contexts.
  • InfoKV combines predictive entropy, layer-wise representation evolution, and attention scores to select tokens during long-context reasoning.
  • Extensive experiments evaluate InfoKV on long prefilling and long decoding benchmarks.

Limitations

The paper notes that entropy is only an indirect proxy for future utility, and adaptive layer-wise budget allocation can improve some models while destabilizing others.

  • Limitations: Entropy captures stronger forward influence than attention metrics but remains an indirect approximation of future utility rather than an explicit optimization objective.Attention weights primarily reflect the relevance of nearby historical contexts to the current query.
  • Limitations: Adaptive layer-wise budget allocation improves performance for some models but can destabilize reasoning performance for others, indicating architecture-specific behavior.The observation suggests that different architectures may exhibit distinct information-processing characteristics.

A. Visualizations of Token Scoring

Figures 7 and 8 visualize entropy- and attention-based token scores on two reasoning-task examples. Attention emphasizes tokens closely relevant to the current query, whereas entropy captures query-independent content words carrying important information.

  • Visualizations: Figures 7 and 8 visualize entropy and attention token scores on two reasoning-task examples.The visualizations cover two examples of reasoning tasks.
  • Attention scoring: Attention assigns high weights to “Option” and the following “A”, “B”, “C” and “D” when deriving the final answer.The attention scores are obtained from “The answer is” in the last sequence.
  • Attention versus entropy: Attention tends to retrieve tokens closely relevant to the current query, while entropy measures token informativeness independently of the query.Entropy is observed to capture more content words that carry important information.

B. Experiment Details · B.1. Long Prefill

The long-prefill experiments standardize implementation and pooling configurations across methods, while fixing entropy-attention weighting and prompting templates to isolate KV cache compression effects.

  • B.1. Long Prefill: All compared methods use official implementations and consistent public hyperparameter settings.These choices support fair comparison across KV cache compression methods.
  • B.1. Long Prefill: The observation window size is fixed at 64 tokens.This is part of the common experimental configuration.
  • B.1. Long Prefill: Pooling uses average pooling with a kernel size of 9.The pooling function and kernel are shared across common configurations.
  • B.1. Long Prefill: Top-k restricted entropy is computed from the top 256 predicted tokens.The entropy computation uses the top-256 prediction set.
  • B.1. Long Prefill: The entropy-attention balancing coefficient α is set to 0.9.This coefficient controls the relative weighting of entropy and attention scores in Eq. (8).
  • B.1. Long Prefill: Fixed prompting templates are used for both direct-answer and Chain-of-Thought reasoning settings in LongReason.The same prompts are applied to every compared method.
  • B.1. Long Prefill: Using identical prompts ensures performance differences mainly reflect KV cache compression strategies rather than prompting variations.This controls prompting as a potential source of experimental variation.

B.2. Long Decoding

Long decoding uses task-specific attention–entropy weighting: α=0.95 for AIME 2024 and LiveCodeBench, and α=0.9 for IFEval. Token-score visualizations illustrate entropy and attention during Llama-3.1-8B-Instruct decoding on two cases.

  • Hyper-parameter settings: α=0.95 is used for AIME 2024 and LiveCodeBench, where mathematical notation and symbolic reasoning make attention reliable for preserving locally important contexts.These settings are part of the hyper-parameter configuration for the three long decoding benchmarks.
  • Hyper-parameter settings: α=0.9 is used for IFEval to provide stronger entropy guidance and better capture token importance.The smaller attention weight increases the relative contribution of entropy scores.
  • Token-score visualizations: Figure 7 visualizes entropy and attention token scores for Llama-3.1-8B-Instruct decoding on a step-by-step reasoning task, while Figure 8 shows the corresponding scores for a second case.The figures cover two decoding-stage examples, identified as case 1 and case 2.
Loading 2606.26875v1…