{"id":"8d5f0181-d5b4-4cde-be16-d56c967a5a66","arxiv_id":"2602.10238","paper_version":2,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":7.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":4,"one_line_summary":"Per-head RL agents that rank KV-cache tokens by predicted future utility outperform heuristic eviction baselines and generalize zero-shot to new tasks.","lead":"This paper trains small per-head reinforcement-learning agents to rank which tokens to keep in an LLM's key-value cache, aiming to evict the least useful ones. If it works, it could cut memory use during long-context inference without changing the model.","discovery_kind":"new_application","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Future-attention reward (Eq. 3.5) may be invalid under eviction; central claim lacks stability check.","rationale":"The reader's weakest assumption is that the future-attention reward remains a valid proxy after eviction. I agree this is the most load-bearing concern because the central claim is not just that KVP works empirically, but that it learns to predict future token utility. If the reward is not stable under eviction, the learned ranking is not actually predicting post-eviction utility, and the conceptual contribution is weakened. The paper's empirical curves could still be explained by a heuristic that correlates with utility in the full cache but not necessarily after eviction. The proposed concrete test directly evaluates this stability by comparing importance scores before and after eviction. I do not see a more fundamental issue: the proof of Proposition 1 is correct, the RL setup is reasonable, and the zero-shot results partially address overfitting. The abstract's 128K overclaim is a separate reporting issue, not a threat to the method's validity. Therefore, the reader's CONDITIONAL verdict remains appropriate, pending the stability check and other minor revisions.","tokens_in":48693,"tokens_out":9789,"duration_ms":104012,"concrete_test":"On a held-out set (e.g., OASST2 test), take a KVP-trained agent and a set of sequences. For each sequence and head: (1) compute full-cache importance scores s_i (Eq. 3.5); (2) for budgets b in {100, 500, 1000}, evict the lowest-scored tokens according to the learned ranking and rerun the LLM with the compressed cache to obtain new attention patterns; (3) recompute importance scores s_i' for the remaining tokens from this compressed run; (4) measure the Spearman rank correlation (and Kendall's tau) between s and s' restricted to retained tokens. If the correlation falls below, say, 0.5 at small budgets, the reward is not a stable proxy. A second check: compare the total future attention actually received by the retained set under compression with the predicted value from the full-cache ranking; a large gap indicates the objective is miscalibrated.","verdict_should_be":"UNCHANGED","load_bearing_attack":"KVP's reward (Eq. 3.5) is the cumulative future attention in the uncompressed model. The policy is trained to rank tokens by this static score. However, the deployment scenario evicts tokens, changing the token set available to future attention. The paper never checks whether the importance scores remain valid after eviction. If attention redistributes (e.g., to remaining tokens or attention sinks), the learned ranking may be systematically misaligned with true post-eviction utility. This is a structural premise: the method's conceptual claim — that it learns to predict future token utility — depends on the scores being stable under the very operation they are used for. The paper's downstream results are encouraging, but they do not isolate this mechanism; a heuristic that happens to work would produce similar curves. No ablation or analysis in Section 4.2 or the appendix addresses the stability of the reward under eviction.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes KVP, a framework that reframes KV cache eviction as a ranking problem. For each KV head, a lightweight MLP policy scores tokens using only keys, values, and positions; eviction keeps the top-b tokens. The policies are trained offline with REINFORCE/RLOO on pre-computed generation traces, using a reward that sums the (normalized) cumulative future attention of evicted tokens over all cache budgets, computed from the uncompressed model. Proposition 1 states that under uniqueness and nestedness, eviction is equivalent to ranking; the proof is given in Appendix A.1. Experiments on RULER-4k and OASST2-4k with Qwen2.5-7B-Chat and Phi-4-14B report accuracy/perplexity improvements over StreamingLLM, SnapKV, TOVA, KeyDiff, LagKV, K-Norm, and Random, plus zero-shot transfer to BoolQ, ARC, MMLU, HellaSwag, and GovReport. An ablation in Section 4.2 argues that RL is preferable to a differentiable sorting surrogate.","tokens_in":48978,"tokens_out":8185,"duration_ms":86764,"significance":"The paper is a serious attempt to replace heuristic scoring with learned, query-agnostic, per-head rankings for KV cache management. Its strengths include a clean ranking formulation with a correct and elementary proof, an offline training scheme that avoids additional LLM inference during policy learning, a budget-agnostic holistic reward, low inference overhead (about 1% prefill FLOPs, 0.71 ms per 10k-token compression), and extensive per-head diagnostics and wall-clock measurements. If the evaluation limitations below are addressed, the work would be a solid contribution to efficient LLM inference. The empirical results, while encouraging, are currently presented with overbroad claims and without statistical support.","major_comments":[{"comment":"The abstract states that KVP is evaluated on RULER 'up to 128K tokens', but Section 4 specifies RULER-4k (approximately 4500-token sequences) as the only RULER configuration; there is no 128K experiment in the main text or appendix. The abstract's zero-shot list also includes 'LongBench passage retrieval', which is not among the downstream tasks evaluated in Section 4.1 (BoolQ, ARC, MMLU, HellaSwag, GovReport). Align the claims with the actual experiments.","section":"Abstract vs. Section 4"},{"comment":"The reward R_b sums future attention computed with the uncompressed model. At deployment the cache is evicted, so future attention is redistributed among the surviving keys; the paper never checks whether the static importance ranking remains valid for the actual post-eviction attention. The ablation in Section 4.2 (Fig. 5) evaluates −R_b against the same uncompressed future-attention target, which does not simulate the deployment condition. Add a stability test: re-compute future attention after masking evicted tokens and compare the ranking, or measure the correlation between the reward under eviction and under no eviction. Without this, the paper's central mechanistic claim is not directly supported.","section":"Eq. (3.5), Section 3.1.1 and 4.2"},{"comment":"All evaluations compress the cache exactly once after prefill and then generate. This is a one-shot compression protocol, not a general streaming KV-cache eviction policy where decisions must be made repeatedly as new tokens are generated. This protocol favors a full ranking over the entire context and limits the scope of the 'eviction' claim. Please either evaluate periodic or streaming eviction (e.g., re-ranking or incremental eviction during decoding) or explicitly restrict the contribution to single-pass post-prefill compression and adjust the title/abstract accordingly.","section":"Section 4, Budgeting and Compression Schedule"},{"comment":"No error bars, confidence intervals, or number of seeds are reported for any downstream accuracy/perplexity curve. The abstract and Section 4.1 use 'significantly outperforms', but with single runs there is no statistical basis for this. Report means and variances over at least three seeds (or bootstraps) for the main comparisons, and state evaluation randomness (e.g., Gumbel noise) across runs.","section":"Figures 2–5, Section 4.1"}],"minor_comments":[{"comment":"Typos: 'surragates' should be 'surrogates'; 'Blondelet al.' should be 'Blondel et al.'.","section":"Section 4.2 and Fig. 5 caption"},{"comment":"Define the attention notation A(x_i, x_j) explicitly (which term is the query and which is the key), and specify the GQA max rule. Also give the future-token window f used in the experiments; the paper leaves f unspecified.","section":"Section 3.1.1, Eq. (3.5)"},{"comment":"Report the RLOO batch size K and state whether the always-retained first 4 / last 16 tokens are applied to all baselines. These implementation choices affect reproducibility and comparability.","section":"Appendix A.2"},{"comment":"The abbreviations KVPS and KVPR are used in the legends but defined only in Section 4.1. Define them in the captions or in the main text before first use.","section":"Figures 3 and 4 legends"},{"comment":"The claim of 'consistently outperforming all baselines' is too strong; the per-subtask breakdown shows KVP behind SnapKV/TOVA on several RULER subtasks (e.g., CWE and QA2). Soften the claim or add a discussion of per-task variability.","section":"Section 4.1, Figure 13"}],"recommendation":"major_revision","confidential_remarks":"The paper is within scope and the core idea is timely and potentially publishable. The main risks are overclaiming relative to the experimental protocol and the lack of a stability check for the future-attention reward under actual eviction. I would ask the authors to correct the abstract, add the stability analysis, and provide statistical support for the main comparisons."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Plainly: this is a solid paper with a genuinely new angle. They reframe KV eviction as a per-head RL ranking problem, train tiny MLP agents on pre-computed traces using only keys, values, and positions, and optimize a budget-agnostic reward based on cumulative future attention. The per-head specialization and Gumbel-Sort training are clean, the wall-clock overhead is negligible, and the empirical sweep is broad: RULER, OASST2, zero-shot on BoolQ/ARC/GovReport/MMLU/HellaSwag, plus a Phi-4 transfer. The RL-vs-surrogate ablation is convincing, and they beat or match attention-aware baselines despite not using queries. Proposition 1 is a simple observation, but it gives the ranking framing a solid footing.\n\nThe main soft spot is the reward. It is computed from future attention in the uncompressed model—that is, on traces where no tokens have been evicted. The policy is trained to predict exactly that static quantity. But at deployment you actually evict tokens, and no section or ablation checks whether those importance scores remain valid after removal. Attention can redistribute to the remaining tokens or to sinks, so the learned ranking may be systematically misaligned with true post-eviction utility. The downstream results show the policy is doing something right, but the paper's central claim—that it learns to predict future utility—is not directly tested. This is a structural gap, not a nit, and it should be addressed: measure how much the future-attention distribution shifts under eviction, or compare against a reward computed on evicted traces.\n\nTwo more concerns. The abstract claims RULER up to 128K, but the body only reports RULER-4k; I see no 128K experiment anywhere. That is an overclaim and should be corrected. And the protocol compresses once after prefill, which is a favorable setup for ranking methods and says little about sustained eviction during decoding. Add error bars and release code, and the paper would be much easier to trust.\n\nBottom line: worth refereeing, and if the authors close the stability gap, it's a strong contribution. Cite this if you work on cache management; I'd bring it to a reading group. Recommend a serious referee.","headline":"Learned per-head RL eviction ranks are a real step forward, but the abstract's 128K claim and the unverified stability of the future-attention reward keep it from being a clean win.","tokens_in":49373,"tokens_out":3639,"would_cite":true,"duration_ms":40593,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"This paper claims that evicting tokens from an LLM's KV cache is best done by learning to rank tokens by their predicted future attention, and that lightweight per-head reinforcement-learned rankers outperform heuristic baselines and transf","keywords":["KV cache eviction","reinforcement learning","token ranking","future attention","long-context inference","per-head policy","zero-shot generalization","cache management"],"falsifier":"Compare, on a held-out set, the policy's predicted token-utility scores against the actual drop in generation quality (e.g., downstream accuracy or next-token log-likelihood) when each candidate token is individually removed from a cache that has already been compressed to a given budget. If the two orderings diverge substantially, the reward—future attention in the uncompressed model—does not survive eviction.","tokens_in":48653,"feed_emoji":"🧠","tokens_out":6805,"duration_ms":74516,"temperature":0.7,"pith_summary":"Large language models reuse previously processed tokens through a key–value (KV) cache, which grows with sequence length and can dominate memory during generation. This paper claims that deciding which cache entries to evict is best framed as a ranking problem: learn, for each attention head, a lightweight scoring function that orders tokens by how much attention they will receive from future tokens, so that dropping the lowest-ranked tokens under any memory budget preserves the most useful information. The authors prove that when the optimal retained sets for increasing budgets are nested, eviction and ranking are equivalent, and they train per-head stochastic ranking policies with a policy-gradient method on pre-computed generation traces, using only keys, values, and positions as inputs. Across long-context and dialogue benchmarks, the learned policies outperform recency-, norm-, and attention-based heuristics, and they generalize zero-shot to other tasks and to longer sequences with no modification of the underlying LLM. If correct, the work turns cache eviction from a hand-tuned heuristic into an offline-learned, budget-agnostic signal for memory-efficient inference.","feed_headline":"Learned rankers beat heuristics at evicting LLM cache tokens","feed_subtitle":"Per-head policies, trained on keys and values alone, predict future token usefulness and transfer to new tasks and longer contexts.","key_machinery":"The argument rests on four pieces: (1) Proposition 1, which reduces budget-constrained eviction to total-order ranking under uniqueness and nestedness of optimal subsets; (2) a per-head scoring MLP that takes only the cached key vector, value vector, and position as its input, so the policy is attention-free and query-independent; (3) a stochastic ranking policy that samples a full permutation by adding extreme-value noise to scores and sorting, enabling parallel sampling and policy-gradient training; and (4) a holistic reward equal to the negative cumulative future attention of evicted tokens, summed over all cache budgets and normalized against the optimal ranking, which removes the need f","core_discovery":"KV cache eviction can be exactly reformulated as learning a single budget-agnostic ranking of tokens. The paper shows that if the best b-token subset is unique and nested inside the best (b+1)-token subset for every b, then there is a total order whose top-b entries realize the optimal eviction for every budget (Proposition 1). The authors then learn such an order with a softmax-sequential stochastic ranking policy, parameterized by a small MLP per attention head that maps the token's key, value, and position to a score. The reward is the negative cumulative future attention of the evicted tokens, summed over all possible cache sizes and normalized by the optimal ranking's cost; this makes t","pith_inferences":["The paper stops short of testing whether the future-attention reward remains valid after tokens are actually removed; a natural follow-up is to re-estimate utilities on the compressed cache and retrain or fine-tune with a reward measured under eviction.","If future attention is a stable individual utility, the same per-head rankings could be combined across heads to allocate non-uniform budgets adaptively, which the paper mentions but does not implement.","The offline training on traces means the framework transfers to new model families by recollecting traces and re-training small MLPs; the one-time data-generation forward pass is the dominant cost, so the method scales with the model rather than with the training loop.","The ranking could be composed with compression techniques: rank tokens first, then merge or quantize only the retained ones, yielding a two-stage memory reduction the paper notes but does not evaluate."],"forward_implications":["A cache can be compressed to any target size by taking the top-b of a single learned ranking, so one policy serves every memory budget without re-training or re-scoring.","Because the ranking uses only keys, values, and positions, it adds no overhead during autoregressive decoding; all compression cost is confined to the prefill stage.","Per-head specialization matters: the learned policy stays effective across all layers and heads, while fixed heuristics are good for some heads and harmful for others.","The learned utility scores transfer zero-shot to unseen tasks and to sequence lengths well above the training distribution, suggesting the ranking captures a general property of the model's attention.","The same ranking can act as a principled signal for hierarchical memory management, with lowest-ranked tokens being natural candidates for offloading rather than hard eviction."],"fun_headline_variants":["RL trains per-head cache eviction, not heuristics","Cache eviction as a learned ranking problem","Per-head agents learn to rank KV tokens","From heuristic to learned: KV cache eviction","Learn a rank: smarter LLM cache eviction"],"cache_read_input_tokens":2304,"weakest_assumption_plain":"The load-bearing premise is that cumulative future attention in the uncompressed model remains a correct measure of a token's usefulness after other tokens have been evicted; if eviction changes future attention patterns, the learned ranking is optimizing a mis-specified objective.","fun_headline_variants_meta":{"raw":{"variants":["RL trains per-head cache eviction, not heuristics","Cache eviction as a learned ranking problem","Per-head agents learn to rank KV tokens","From heuristic to learned: KV cache eviction","Learn a rank: smarter LLM cache eviction"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000142,"raw_usage":{"total_tokens":1029,"prompt_tokens":790,"completion_tokens":239,"prompt_tokens_details":{"cached_tokens":256},"prompt_cache_hit_tokens":256,"prompt_cache_miss_tokens":534,"completion_tokens_details":{"reasoning_tokens":167}},"tokens_in":534,"tokens_out":239,"duration_ms":3442,"temperature":1.0,"reasoning_tokens":167,"cache_read_input_tokens":256,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-03T01:11:38.059715+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Compare, on a held-out set, the policy's predicted token-utility scores against the actual drop in generation quality (e.g., downstream accuracy or next-token log-likelihood) when each candidate token is individually removed from a cache that has already been compressed to a given budget. If the two orderings diverge substantially, the reward—future attention in the uncompressed model—does not survive eviction.","supporting_citations":[],"review_version":1}