PaperlayerView exact revision on arXiv

Source-linked AI summary

Jacap: Robust KV Cache Eviction via Jacobian-Based Nonlinear Information Capacity Preservation

Jiaming Yang, Chenwei Tang, Liangli Zhen, Chenyang Zhang, Jiancheng Lv

arXiv:2609.08131v1cs.CL

TL;DR

Long-context inference makes KV-cache eviction necessary, but existing heuristics do not rigorously characterize token utility under nonlinear softmax attention. The paper models attention as a nonlinear information channel, derives Jacobian Information Capacity, and introduces JACAP for capacity-aware selection. Experiments report superior performance in most scenarios, particularly under high compression, while the method remains limited by its first-order local approximation and simplified softmax competition model.

  • Problem

    Existing KV-cache eviction methods rely largely on heuristics and do not rigorously characterize token utility or interactions under nonlinear softmax attention.

  • Method

    The paper models retained attention as a nonlinear information channel, approximates its local capacity with a Jacobian, and uses softmax-sensitive weighting for capacity-aware token selection.

  • Results

    JACAP delivers superior performance in most scenarios and shows consistent improvements on LongBench, NIAH, and AIME25, especially at high compression ratios.

  • Takeaways & Limitations

    Jacobian-capacity-aware eviction prioritizes informative and non-redundant tokens under tight cache constraints.

  • Takeaways & Limitations

    JACAP relies on a first-order local approximation and token-wise sensitivity weights that simplify pairwise softmax competition, potentially limiting expressiveness in highly dynamic decoding.

Abstract

from arXiv · show

Key-value (KV) cache eviction is essential for scaling long-context inference in Large Language Models. However, existing policies predominantly rely on empirical heuristics, lacking a rigorous characterization of token utility under the inherently nonlinear softmax attention mechanism. In this work, we rethink KV cache eviction through the lens of local information geometry, modeling the attention process as a nonlinear Gaussian communication channel. By performing a first-order Taylor expansion of the attention mapping, we derive the Jacobian Information Capacity, a novel objective that explicitly captures query relevance, softmax sensitivity, and structural diversity. Guided by this theory, we introduce Jacap, a capacity-aware eviction method that utilizes softmax-aware importance weighting and statistical leverage scores for subset selection. Extensive experiments across diverse architectures and benchmarks demonstrate that \textsc{Jacap} delivers superior performance in most scenarios, particularly in high-compression regimes.

1 Introduction

Long-context inference makes KV-cache memory and latency a critical bottleneck, while heuristic eviction methods lack a principled account of token utility and interactions. The paper introduces a nonlinear information-capacity framework and JACAP, which aims to preserve relevant, diverse cache information under compression.

  • Motivation: KV-cache memory overhead and latency grow with context length, constraining long-context LLM inference.The cache stores past attention states for autoregressive decoding.
  • Problem: Existing eviction methods rank tokens using attention, recency, or value norms, but lack a principled characterization of information worth preserving.These methods can treat tokens independently and miss redundancy or complementarity.
  • Problem: Prior information-theoretic eviction work uses a linear-Gaussian attention surrogate, limiting its description of nonlinear softmax dynamics.The earlier framework measures mutual information between future queries and attention outputs.
  • Method: The proposed framework models retained attention as a nonlinear information channel and derives local capacity from the Jacobian of the attention mapping.Its objective incorporates query-key alignment, softmax competition, and value-space diversity.
  • Method: JACAP combines softmax-sensitive weighting with capacity-based token selection to approximate nonlinear information capacity.The method is designed for practical KV-cache eviction under a memory budget.
  • Results: JACAP shows consistent improvements over prior methods on LongBench, NIAH, and AIME25, especially at high compression ratios.The contribution summary reports the cross-benchmark pattern rather than a single-task result.

2 Related Works

KV-cache compression methods include quantization, token reduction, and structural eviction, with structural methods selectively retaining tokens during decoding. Existing eviction approaches often use independent token scores, whereas information-theoretic alternatives improve principled selection but may simplify nonlinear softmax behavior.

  • Compression approaches: Quantization reduces KV-cache memory through low-precision storage, while token merging shortens context by consolidating redundant tokens.These approaches operate at the representation and sequence levels, respectively.
  • Compression approaches: Structural methods compress the KV cache by selectively retaining tokens during decoding and are complementary to quantization and system-level optimizations.JACAP belongs to this structural category.
  • Heuristic eviction: Heuristic eviction methods use attention behavior, key dissimilarity, norms, or expected future attention to estimate token importance.Examples include H2O, SNAPKV, KEYDIFF, KNORM, and Expected Attention.
  • Heuristic eviction: Independent token scoring fails to capture higher-order interactions such as redundancy and complementarity among retained tokens.This limits how well independently ranked subsets represent cache structure.
  • Information-theoretic methods: CAPKV formulates selection through mutual information between future queries and attention outputs using a linear-Gaussian approximation.Under that approximation, the objective becomes a tractable log-determinant form.
  • Information-theoretic methods: JACAP instead models attention as a nonlinear information channel and derives a Jacobian-based local capacity objective that accounts for softmax normalization and competition.This provides a more faithful characterization of information flow in the cache than a fixed-channel approximation.

3 Methodology

JACAP formulates KV cache eviction as preserving query-dependent information through a nonlinear softmax attention channel. It locally linearizes this channel with a Jacobian, then approximates capacity-aware subset selection using softmax-aware weights and output-space diversity.

  • 3.3 Jacobian Information Capacity for KV Cache Selection: JACAP models retained attention as a nonlinear query-to-output mapping rather than a fixed linear channel.The retained cache induces a softmax-normalized channel in which query–key accessibility and token competition jointly affect outputs.
  • 3.3 Jacobian Information Capacity for KV Cache Selection: The method approximates this nonlinear channel around the future-query distribution center with a first-order Jacobian tangent channel.The local model assumes queries concentrate around µQ with covariance ΛQ, and uses the Jacobian as the effective channel gain.
  • 3.3 Jacobian Information Capacity for KV Cache Selection: The Jacobian-capacity objective captures key accessibility, output directions, softmax competition, query statistics, and model uncertainty.The softmax Jacobian represents both token-specific sensitivity and normalization-induced competition, while the capacity is evaluated under ΛQ and Σnoise.
  • 3.4 JACAP: A Jacobian-Capacity Aware Eviction Policy: Because exact subset-dependent log-determinant optimization is computationally prohibitive, JACAP preserves local relevance, sensitivity, and output-space diversity in a lightweight score.The practical algorithm computes subset-independent local attention weights over the candidate pool and selects the top-B marginal Jacobian-capacity contributions.
  • 3.4 JACAP: A Jacobian-Capacity Aware Eviction Policy: The nonlinear importance weight combines query relevance, softmax sensitivity, and query-space variability, while selection accounts for weighted output-space capacity.Moderately responsive tokens can receive more local capacity than saturated high-attention tokens, and the resulting matrix supports marginal scoring.

4 Experiment

Experiments compare JACAP with heuristic and capacity-aware eviction baselines across LongBench, NIAH, AIME25, and a temperature ablation. JACAP is generally strongest under severe cache constraints, while remaining competitive across models and decoding settings.

  • LongBench: JACAP is evaluated against CAPKV, SNAPKV, KEYDIFF, EA, and KNORM across LongBench compression ratios using three language models.The study emphasizes compression ratios 0.75 and 0.9, where memory constraints are most severe.
  • LongBench: At 0.75 compression, JACAP generally matches or slightly exceeds CAPKV and significantly outperforms heuristic baselines.The comparison includes SNAPKV, KEYDIFF, EA, and KNORM.
  • LongBench: At 0.9 compression, JACAP leads the average score by approximately 5 points over the next-best baseline.This result is reported as evidence of stronger retention under extreme cache constraints.
  • LongBench: JACAP maintains stronger SD-QA and MD-QA performance, more stable degradation on summarization and few-shot tasks, and consistent trends across all three models.These results indicate robustness across task types and architectures.
  • NIAH: On NIAH at 0.75 compression, JACAP provides more stable retrieval than baselines across 10k–120k-token contexts and Needle Depths of 0.15–0.95.Its advantage is especially pronounced for deep needles and long contexts, while CAPKV remains competitive at moderate depths.
  • Decoding Eviction: On AIME25 online decoding, JACAP leads at 2048 and 16384 reserved tokens, matches CAPKV at 4096, and trails CAPKV at 8192.The results show competitiveness across changing decoding cache budgets.
  • Temperature Ablation: LongBench ablations identify τ around 10 as the best-performing temperature, balancing query-dependent relevance against diversity.Small τ values overfocus on highly aligned tokens, whereas large τ values oversmooth the relevance prior.

5 Conclusion

The paper concludes by framing KV-cache eviction as a local nonlinear information-capacity problem and deriving a Jacobian-based approximation through first-order linearization. The appendix formalizes this approximation as a Gaussian channel and derives its mutual-information objective.

  • Capacity Objective: The resulting Jacobian Information Capacity is obtained by treating J_C as the effective local channel gain.The theorem defines retained-cache capacity under the local Gaussian and first-order assumptions.
  • Local Channel Model: The framework models retained attention as a noisy nonlinear channel with future queries distributed around a center µ_Q with covariance Λ_Q.The target quantity is the mutual information between queries and retained-cache outputs.
  • Local Linearization: A first-order Taylor expansion around µ_Q converts the nonlinear attention mapping into a local linear channel governed by the Jacobian J_C.The Jacobian is evaluated at the query center.
  • Jacobian Derivation: The attention Jacobian is derived by differentiating logits and softmax weights through the chain rule.The derivation evaluates the attention weights at q = µ_Q and uses the softmax derivative.
  • Capacity Derivation: The appendix treats the perturbed channel as linear Gaussian, computes Gaussian covariance matrices, and expresses mutual information through log-determinants.The determinant identity yields the final closed-form objective and completes the theorem derivation.

A.2 Analysis of the Jacobian Information Capacity

The Jacobian Information Capacity extends linear information-theoretic objectives by modeling softmax attention sensitivity and competition locally. It values tokens that combine query alignment, local sensitivity, and distinct output directions.

  • Softmax competition: The Jacobian capacity introduces the Softmax Competition Matrix SC to capture nonlinear attention dynamics.Its off-diagonal terms model how increasing one token’s logit decreases the attention weights of others.
  • Local sensitivity and saturation: Softmax sensitivity is highest near attention weight 0.5, where tokens offer the greatest local information gain.The diagonal sensitivity term is maximized when α⋆_i ≈ 0.5.
  • Local sensitivity and saturation: Sensitivity vanishes as attention weights approach 0 or 1, so ignored and fully dominant tokens have low marginal information value.This saturation effect contrasts with heuristics that treat larger attention weights as uniformly more important.
  • Capacity interpretation: The objective rewards subsets that align with future queries while also retaining high local sensitivity and distinct output directions.These factors are integrated through the Jacobian structure rather than treated as independent heuristics.
  • Practical reduction: The appendix connects the full Jacobian Information Capacity to the efficient one-shot JACAP scoring rule through principled approximations.The reduction supports practical computation while preserving the central nonlinear sensitivity effect.

B.1 Derivation of the JACAP Algorithm

JACAP derives an efficient eviction score by approximating coupled softmax and query-response structures with diagonal terms, then selecting tokens using leverage scores. The resulting criterion combines nonlinear importance weighting with structural diversity.

  • Approximation: JACAP simplifies the coupled Jacobian objective through diagonal approximations of softmax competition and query-response covariance.These approximations decouple token-wise scoring while retaining the sensitivity-versus-saturation effect.
  • Softmax approximation: The diagonal softmax term preserves the effect that saturated tokens have low local influence while omitting negative pairwise competition.The omission is presented as a practical simplification for efficient scoring.
  • Query-response covariance: The query-response covariance describes how token responses vary and correlate under the future-query distribution.The scalar κ_i measures token-logit variance, while diagonalization assumes independent key responses for efficient scoring.
  • Approximated objective: Under isotropic noise, the approximated objective becomes maximizing the log-determinant of a sum of weighted rank-one matrices.The combined token weight absorbs constants and the noise scale.
  • Leverage score selection: JACAP constructs a global capacity matrix and selects tokens greedily using statistical leverage scores as approximations to marginal log-determinant gains.The leverage term is evaluated as w_i u_i^⊤A^-1u_i using the matrix determinant lemma.

B.2 Computational Complexity Analysis

JACAP incurs additional head-dimension-dependent matrix costs beyond independent heuristic scoring. The paper reports that this overhead is moderate in practice and is justified by improved retention quality under high compression.

  • Complexity stages: JACAP’s per-step computation includes query-key scoring, capacity-matrix construction, matrix inversion, and leverage-score evaluation.The capacity matrix construction and inversion dominate the overall complexity.
  • Complexity stages: With diagonal query covariance, computing token importance costs O(Nd_h), while capacity construction and leverage evaluation introduce quadratic head-dimension dependence.The latter stages involve rank-one updates and operations on d_h × d_h matrices.
  • Comparison with heuristics: Compared with SNAPKV, H2O, KEYDIFF, and KNORM, JACAP adds quadratic dependence on head dimension through inter-token matrix interactions.The cited heuristic methods typically use O(Nd_h) or O(N) independent per-token scoring.
  • Practical trade-off: The additional overhead is reported as moderate because attention head dimensions are usually small and matrix operations parallelize on modern GPUs.The paper states that improved retention quality in high-compression regimes justifies the added cost.

C More Experiment Results

The experiments were conducted on four NVIDIA RTX Pro 6000 GPUs with 1024 GB of system memory under Ubuntu 22.04.

  • Experimental setup: All experiments used 4× NVIDIA RTX Pro 6000 GPUs and 1024 GB system memory under Ubuntu 22.04.These are the reported experimental hardware and operating-system settings.

C.1 Runtime Efficiency Analysis

JACAP maintains runtime comparable to existing eviction methods across context lengths and compression ratios, while its practical overhead remains negligible compared with overall inference cost.

  • Runtime comparison: JACAP’s runtime is highly comparable to existing methods across all tested context lengths and compression settings.The evaluation generates 100 tokens with Qwen3-8B on a single GPU across 8k–64k-token contexts at compression ratios 0.6 and 0.8.
  • Computational overhead: Theoretical complexity is O(Nd_h^2), but typical head dimensions such as d_h = 128 allow GPUs to handle the matrix operations efficiently.
  • Computational overhead: JACAP achieves significant high-compression performance gains without incurring prohibitive computational cost.

C.2 More Experiment Result on NIAH

At compression ratio 0.5 on Qwen3-8B, JACAP achieves the best overall NIAH performance and remains robust across context lengths and needle depths.

  • NIAH performance: JACAP achieves the best overall performance on NIAH at compression ratio 0.5 with Qwen3-8B.
  • Robustness: JACAP consistently retrieves the needle across the full spectrum of context lengths and depths.
  • Baseline behavior: KEYDIFF and KNORM fail significantly at long contexts, while SNAPKV degrades at extreme lengths above 100k tokens and deeper positions.
  • Mechanism: The results indicate that nonlinear sensitivity helps identify and preserve crucial information sparsely located deep within very long sequences.
Loading 2609.08131v1…