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

1. Problem Statement

Large Language Models (LLMs) face significant memory bottlenecks during inference due to the Key-Value (KV) cache, which stores past token vectors to enable efficient generation. While numerous KV cache reduction techniques exist, layer-wise token pruning methods (e.g., FastKV, GemFilter) have emerged as a popular solution. These methods select a subset of important tokens at specific layers to retain in the KV cache, pruning the rest.

However, existing methods suffer from a critical limitation: inflexibility. They rely on pre-defined, fixed selection layers determined independently of the specific task.

• The Issue: Different tasks exhibit different attention patterns. Simple tasks (e.g., standard QA) allow for early token identification, while complex tasks (e.g., KV retrieval, Needle-in-a-Haystack) require deeper layers to resolve semantic similarities before tokens can be reliably selected.
• The Consequence: A fixed layer that works well for simple tasks causes severe accuracy degradation in difficult tasks. Conversely, to maintain accuracy on hard tasks, these methods must either delay selection (reducing memory savings) or increase the KV budget, negating their efficiency benefits.

2. Methodology: Adaptive Selection Layer (ASL)

The authors propose ASL, a training-free, task-aware method that dynamically determines the optimal layer for token selection during the prefilling stage.

Core Mechanism

ASL exploits the observation that attention scores evolve across layers:

1. Early Layers: Attention scores are uniformly distributed across context tokens.
2. Deep Layers: Attention scores localize to a stable, small subset of tokens.

ASL monitors the variance of token ranks ordered by attention scores to detect when this localization occurs.

Algorithm Steps:

1. Observation Window: Starting from a minimum layer ($L_{min}$), ASL observes the next $L_{obs}$ layers.
2. Score Pooling: It aggregates attention scores across attention heads and applies 1D average pooling to smooth noise and capture contiguous context regions.
3. Rank Calculation: For the most recent $L_{obs}$ layers, it identifies the top-$k$ tokens and calculates the variance of their ranks.
4. Relative Variance: To make the metric task-agnostic, the current variance is normalized by the initial variance at $L_{min}$:
   $$\text{relative\_variance}(L) = \frac{\text{vartop}(L)}{\text{vartop}(L_{min})}$$
5. Adaptive Decision: A user-specified threshold ($\tau$) is set. As the model processes layers, if the relative variance drops below $\tau$, it indicates that token ranks have stabilized. The current layer is immediately designated as the selection layer.
6. Execution:
   • One-Shot Selection: Once the layer is selected, tokens are pruned, and only the selected tokens are propagated to all subsequent layers.
   • Integration: ASL can be combined with existing methods (e.g., SnapKV for pre-selection compression, GemFilter for two-pass strategies) to meet specific KV budget constraints.

3. Key Contributions

1. Identification of Inflexibility: The paper demonstrates that fixed selection layers in state-of-the-art (SOTA) methods lead to significant performance variance across tasks, particularly failing on high-difficulty retrieval tasks.
2. ASL Proposal: Introduction of a novel, training-free adaptive method that determines the selection layer based on the variance of token ranks, enabling the model to adapt to task difficulty without manual tuning.
3. Integration Capability: ASL is designed to work seamlessly with existing KV cache reduction techniques (SnapKV, GemFilter) to optimize both the prefilling and decoding stages while adhering to user-specified memory budgets.
4. Comprehensive Evaluation: Extensive experiments across three major benchmarks (InfiniteBench, RULER, NIAH) with context lengths up to 256k, showing superior accuracy-efficiency trade-offs compared to SOTA methods.

4. Experimental Results

The authors evaluated ASL on Llama-3.1-8B-UltraLong and Qwen2.5-7B models.

• Accuracy:

  • Difficult Tasks: ASL significantly outperforms fixed-layer methods (FastKV, GemFilter) on hard tasks like KV Retrieval and Multi-key NIAH. For example, on InfiniteBench's KV retrieval task with a 2048 KV budget, ASL achieved near-full-KV accuracy where FastKV dropped significantly.
  • Overall Performance: On InfiniteBench, ASL (2-pass version) achieved the highest average score for Llama-3.1-8B-UL. On RULER (128k context), ASL consistently outperformed FastKV, especially as context length increased.
  • NIAH: ASL achieved full scores across all context lengths (1k–256k) on the Needle-in-a-Haystack test, matching full KV performance, whereas other methods struggled at specific depths.

• Efficiency (TTFT & TPOT):

  • Time to First Token (TTFT): ASL is slightly slower than FastKV because it often selects tokens at deeper layers (higher $L_{select}$) to ensure accuracy on hard tasks. However, it remains significantly faster than full KV.
  • Time Per Output Token (TPOT): ASL achieves decoding speeds comparable to other pruning methods, as the decoding stage relies on the reduced token set regardless of when selection occurred.
  • Throughput: While ASL trades some throughput for accuracy, the gap narrows in real-world scenarios with longer output sequences where decoding time dominates.

• Memory Usage: ASL introduces negligible memory overhead (only storing pooled scores and rank caches for a few layers) and achieves memory reduction comparable to other pruning methods (e.g., reducing 18GB usage to ~0.3GB for Llama-3.1-8B).

5. Significance

This work addresses a fundamental rigidity in current LLM inference optimization. By shifting from static, pre-defined heuristics to dynamic, data-driven decision-making, ASL allows LLMs to:

• Balance Accuracy and Speed: Automatically trade off inference speed for accuracy based on the specific difficulty of the input task.
• Handle Long Contexts Robustly: Maintain high performance on complex retrieval tasks (like finding specific keys in a massive document) where previous methods fail.
• Generalize Without Training: Operate as a plug-and-play module without requiring fine-tuning, making it immediately applicable to existing LLM deployments.

The paper concludes that adaptive layer selection is a crucial step forward for deploying LLMs in real-world scenarios where task difficulty varies dynamically, ensuring reliable performance without sacrificing the memory efficiency gains of token pruning.