Adaptive Layer Selection for Layer-Wise Token Pruning in LLM Inference

Chuan Xiao; Makoto Onizuka; Rei Taniguchi; Yuyang Dong

arxiv: 2601.07667 · v2 · submitted 2026-01-12 · 💻 cs.CL · cs.AI· cs.LG

Adaptive Layer Selection for Layer-Wise Token Pruning in LLM Inference

Rei Taniguchi , Yuyang Dong , Makoto Onizuka , Chuan Xiao This is my paper

Pith reviewed 2026-05-16 14:58 UTC · model grok-4.3

classification 💻 cs.CL cs.AIcs.LG

keywords KV cache reductionlayer-wise token pruningadaptive layer selectionLLM inferenceattention score varianceprefilling stagetoken selectionone-shot pruning

0 comments

The pith

ASL selects the KV pruning layer on the fly by measuring variance in attention-based token ranks, improving accuracy on hard tasks without training.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper introduces ASL, a training-free method that adaptively picks the layer for layer-wise token pruning during LLM prefilling. It calculates the variance across tokens in how they are ranked by attention scores and uses that spread to decide where to apply one-shot selection so the KV cache stays within a user budget. Fixed pruning layers often lose critical information on difficult tasks such as KV retrieval, but the variance signal lets ASL shift the pruning point to preserve accuracy while still reducing cache size. The approach works in the prefilling stage and combines directly with existing decoding-stage methods like SnapKV. Tests on InfiniteBench, RULER, and NIAH show it outperforms prior layer-wise pruning techniques specifically where those methods degrade.

Core claim

By exploiting the variance of token ranks ordered by attention score, ASL adaptively selects the layer at which to perform one-shot token selection for KV cache reduction, allowing it to trade inference speed for accuracy in a task-dependent manner without any training or calibration.

What carries the argument

Variance of token ranks ordered by attention score, used as the signal to choose the pruning layer adaptively during prefilling.

If this is right

ASL meets any user-specified KV budget while balancing accuracy across easy and hard tasks.
It outperforms fixed-layer methods on difficult benchmarks such as KV retrieval.
It can be combined with existing decoding optimizations like SnapKV.
The method requires no training or per-task calibration and runs only in prefilling.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

Variance signals from attention ranks could guide adaptive compression in other stages or model components.
LLM serving systems might reduce per-task tuning by relying on this internal attention statistic.
The approach suggests attention dynamics during prefilling contain usable information about optimal compression points.

Load-bearing premise

Variance in how tokens rank by attention score is a reliable signal for picking the right pruning layer regardless of the task.

What would settle it

A task where the layer chosen by attention-rank variance produces noticeably lower accuracy than the best fixed-layer choice under the same KV budget.

read the original abstract

Due to the prevalence of large language models (LLMs), key-value (KV) cache reduction for LLM inference has received remarkable attention. Among numerous works that have been proposed in recent years, layer-wise token pruning approaches, which select a subset of tokens at particular layers to retain in KV cache and prune others, are one of the most popular schemes. They primarily adopt a set of pre-defined layers, at which tokens are selected. Such design is inflexible in the sense that the accuracy significantly varies across tasks and deteriorates in harder tasks such as KV retrieval. In this paper, we propose ASL, a training-free method that adaptively chooses the selection layer for KV cache reduction, exploiting the variance of token ranks ordered by attention score. The proposed method balances the performance across different tasks while meeting the user-specified KV budget requirement. ASL operates during the prefilling stage and can be jointly used with existing KV cache reduction methods such as SnapKV to optimize the decoding stage. By evaluations on the InfiniteBench, RULER, and NIAH benchmarks, we show that ASL, equipped with one-shot token selection, adaptively trades inference speed for accuracy, outperforming state-of-the-art layer-wise token pruning methods in difficult tasks.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Desk Editor's Note private letter to a colleague

ASL adds a training-free adaptive layer picker for token pruning using attention rank variance, but the supporting evidence and justification for the heuristic look thin.

read the letter

The main takeaway is that this paper offers ASL, a simple training-free way to pick which layer to apply token pruning at during prefilling. It uses the variance of how tokens rank by attention scores as the signal, instead of locking in fixed layers ahead of time. The goal is to keep accuracy steadier across tasks while staying inside a user-set KV budget, and it can layer on top of methods like SnapKV for the decode phase. Evaluations are reported on InfiniteBench, RULER, and NIAH, with the claim that it handles harder retrieval-style tasks better than prior fixed-layer pruning work. That adaptive angle is the clearest new piece; most earlier layer-wise pruning papers stick to static choices that the authors say degrade unevenly. The approach stays lightweight and requires no extra training or per-task tuning, which is a practical plus for deployment. The soft spots sit in the justification and the results. The paper treats variance of attention ranks as a reliable, task-agnostic cue for the right pruning layer, yet supplies no analysis showing why this statistic tracks downstream accuracy better than mean attention, entropy, or rank stability. If that correlation is weak or flips on other long-context distributions, the adaptivity reduces to an untested rule. The abstract mentions benchmark wins, but without tables, ablations, error bars, or statistical checks visible here, it is hard to judge how large or consistent the gains really are. This is the kind of work that interests people building efficient long-context inference systems who are looking for quick, training-free tweaks to try. A reader already working on KV cache reduction would get a usable idea to test, though the paper would need tighter validation to move beyond a heuristic. I would send it to peer review so the full experiments and the choice of variance can be examined directly.

Referee Report

3 major / 2 minor

Summary. The paper introduces ASL, a training-free adaptive layer selection method for layer-wise token pruning during LLM inference. It computes the variance of per-token ranks (ordered by attention scores) in the prefilling stage to dynamically choose the pruning layer, aiming to respect a user-specified KV cache budget while preserving accuracy better than fixed-layer baselines. Evaluations on InfiniteBench, RULER, and NIAH benchmarks claim that ASL with one-shot selection outperforms prior layer-wise pruning methods, especially on difficult long-context tasks, and is compatible with methods like SnapKV.

Significance. If the empirical results hold under rigorous verification, ASL would provide a practical, zero-training overhead improvement to KV cache reduction for long-context inference. The adaptive heuristic addresses a clear limitation of static layer choices and could be broadly useful if the variance signal proves robust across models and distributions.

major comments (3)

[§3.2] §3.2: The core assumption that variance of attention-based token ranks is a reliable, task-agnostic indicator for selecting the pruning layer is presented without derivation or correlation analysis. No evidence is given showing why variance (rather than mean attention, rank stability, or entropy) predicts accuracy preservation on hard retrieval tasks, leaving the adaptive rule as an unmotivated heuristic.
[Evaluation section (Tables 1-3)] Evaluation section (Tables 1-3): The reported outperformance on RULER and NIAH lacks error bars, statistical significance tests, or multiple random seeds. Given that attention scores are stochastic and task-dependent, this undermines the claim that ASL reliably trades speed for accuracy versus fixed-layer SOTA methods.
[§4.2] §4.2: The manuscript states that ASL meets the user-specified KV budget, but provides no explicit mechanism or ablation showing how the chosen layer is adjusted to enforce the exact budget when variance indicates a suboptimal layer; this is load-bearing for the flexibility claim.

minor comments (2)

[Abstract] Abstract and §1: The claim of 'outperforming state-of-the-art' should be qualified with the specific baselines and KV budgets used, as the abstract currently lacks quantitative details.
[§3.1] Notation in §3.1: The definition of token rank variance is introduced without a clear equation number or pseudocode, making reproduction difficult.

Simulated Author's Rebuttal

3 responses · 0 unresolved

We thank the referee for the constructive and detailed comments. We address each major point below and will incorporate revisions to improve clarity, motivation, and rigor in the manuscript.

read point-by-point responses

Referee: [§3.2] The core assumption that variance of attention-based token ranks is a reliable, task-agnostic indicator for selecting the pruning layer is presented without derivation or correlation analysis. No evidence is given showing why variance (rather than mean attention, rank stability, or entropy) predicts accuracy preservation on hard retrieval tasks, leaving the adaptive rule as an unmotivated heuristic.

Authors: We agree that the motivation for using variance can be strengthened. In the revised manuscript, we will expand §3.2 with a new correlation analysis subsection. This will include empirical plots correlating variance of per-token attention ranks with accuracy on hard retrieval tasks from RULER and NIAH, along with a brief comparison to alternatives such as entropy and mean attention to justify the choice of variance as the selection signal. revision: yes
Referee: [Evaluation section (Tables 1-3)] The reported outperformance on RULER and NIAH lacks error bars, statistical significance tests, or multiple random seeds. Given that attention scores are stochastic and task-dependent, this undermines the claim that ASL reliably trades speed for accuracy versus fixed-layer SOTA methods.

Authors: We acknowledge the need for greater statistical rigor. In the revision, we will rerun the relevant experiments using multiple random seeds (accounting for any stochastic elements in attention computation or sampling), add error bars to Tables 1-3, and include statistical significance tests (e.g., paired t-tests) for the key comparisons against fixed-layer baselines to substantiate the reliability of the reported improvements. revision: yes
Referee: [§4.2] The manuscript states that ASL meets the user-specified KV budget, but provides no explicit mechanism or ablation showing how the chosen layer is adjusted to enforce the exact budget when variance indicates a suboptimal layer; this is load-bearing for the flexibility claim.

Authors: This is a fair observation. The current text assumes budget compliance but does not detail the adjustment procedure. We will revise §4.2 to explicitly describe the mechanism (e.g., iterative layer adjustment or pruning ratio scaling) used to enforce the exact KV cache budget. We will also add a targeted ablation study quantifying the performance impact of this adjustment step. revision: yes

Circularity Check

0 steps flagged

No circularity: ASL selection rule defined directly from observable attention variance

full rationale

The paper defines ASL as a training-free heuristic that computes variance of per-token ranks (ordered by attention scores) during prefilling and uses that scalar to pick the pruning layer under a fixed KV budget. No equations are presented that reduce this choice to a fitted parameter, a self-cited uniqueness theorem, or any input by construction. The rule is stated as an explicit, observable statistic rather than a derived quantity that tautologically reproduces its own inputs. No load-bearing self-citations or ansatz smuggling appear in the derivation chain. The method therefore remains self-contained against external benchmarks and receives a score of 0.

Axiom & Free-Parameter Ledger

0 free parameters · 2 axioms · 0 invented entities

The method rests on the domain assumption that attention scores reflect token importance and on the paper-specific heuristic that rank variance signals a good pruning layer; no free parameters or new entities are introduced.

axioms (2)

domain assumption Attention scores provide a valid ranking of token importance for pruning decisions
Standard assumption in attention-based KV cache pruning literature.
ad hoc to paper High variance in token rank order across layers indicates an effective point for token selection
Core design choice of ASL not derived from prior equations.

pith-pipeline@v0.9.0 · 5524 in / 1188 out tokens · 23075 ms · 2026-05-16T14:58:05.483650+00:00 · methodology

discussion (0)

Lean theorems connected to this paper

Citations machine-checked in the Pith Canon. Every link opens the source theorem in the public Lean library.

IndisputableMonolith/Cost/FunctionalEquation.lean washburn_uniqueness_aczel unclear

?

unclear
Relation between the paper passage and the cited Recognition theorem.

we propose ASL, a training-free method that adaptively chooses the selection layer for KV cache reduction, exploiting the variance of token ranks ordered by attention score

What do these tags mean?

matches: The paper's claim is directly supported by a theorem in the formal canon.
supports: The theorem supports part of the paper's argument, but the paper may add assumptions or extra steps.
extends: The paper goes beyond the formal theorem; the theorem is a base layer rather than the whole result.
uses: The paper appears to rely on the theorem as machinery.
contradicts: The paper's claim conflicts with a theorem or certificate in the canon.
unclear: Pith found a possible connection, but the passage is too broad, indirect, or ambiguous to say the theorem truly supports the claim.

Forward citations

Cited by 2 Pith papers

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score.

When Does Value-Aware KV Eviction Help? A Fixed-Contract Diagnostic for Non-Monotone Cache Compression
cs.LG 2026-05 unverdicted novelty 6.0

A fixed-contract probe shows value-aware KV eviction recovers needed evidence in 72.6% of accuracy-improving cases on LongBench but only 32.4% otherwise, suggesting an order of recover evidence, rank value, then prese...
When Less Latent Leads to Better Relay: Information-Preserving Compression for Latent Multi-Agent LLM Collaboration
cs.LG 2026-04 unverdicted novelty 6.0

Orthogonal Backfill compression for latent KV caches in multi-agent LLMs reduces communication by 79.8-89.4% while achieving comparable or superior performance to full relay on 7 of 9 benchmarks.