Leveraging KV Similarity for Online Structured Pruning in LLMs

TL;DR

The paper tackles the latency and inefficiency of large language model inference by introducing Token Filtering, an online structured pruning method that skips redundant attention computations based on joint key–value similarity. It leverages per-head KV similarity with an anchor-based running statistic and a variance-aware fusion to robustly identify uninformative tokens, combined with a tail-focused, layer-wise thresholding strategy and a warm-up phase to maintain stability. Extensive zero-shot evaluations across LLaMA-2, LLaMA-3, Mistral, and Phi-4 show Token Filtering preserves accuracy and perplexity close to dense models, while delivering substantial latency and memory reductions, particularly at higher pruning ratios and longer contexts. The approach outperforms prior structured pruning methods on challenging benchmarks such as MMLU and maintains strong long-context performance, all without calibration data or fine-tuning. This makes online pruning a practical and scalable path to speed up inference in diverse LLM deployments.

Abstract

Pruning has emerged as a promising direction for accelerating large language model (LLM) inference, yet existing approaches often suffer from instability because they rely on offline calibration data that may not generalize across inputs. In this work, we introduce Token Filtering, a lightweight online structured pruning technique that makes pruning decisions directly during inference without any calibration data. The key idea is to measure token redundancy via joint key-value similarity and skip redundant attention computations, thereby reducing inference cost while preserving critical information. To further enhance stability, we design a variance-aware fusion strategy that adaptively weights key and value similarity across heads, ensuring that informative tokens are retained even under high pruning ratios. This design introduces no additional memory overhead and provides a more reliable criterion for token importance. Extensive experiments on LLaMA-2 (7B/13B), LLaMA-3 (8B), and Mistral (7B) demonstrate that Token Filtering consistently outperforms prior structured pruning methods, preserving accuracy on commonsense reasoning benchmarks and maintaining strong performance on challenging tasks such as MMLU, even with 50% pruning.

Figures (3)

• Figure 1: Technique of Token Filtering. ① Tokens first pass through the Token Filtering layer before entering the attention layer. ② Cosine similarity is computed between the key/value and the anchor key/value, where the anchor key/value represents the average of previous keys and values. ③ The key similarity and value similarity are added to obtain the similarity score. ④ If the similarity score is high, the attention layer is skipped.
• Figure 2: Visualization of inverse KV similarity and attention weights for a representative head (Layer 37, Head 5). For clarity, we plot the inverse cosine similarity so that higher values correspond to lower similarity. Peaks in inverse similarity generally coincide with peaks in attention weights, illustrating that tokens with lower similarity tend to receive higher attention.
• Figure 3: Latency and memory reduction (%) of Token Filtering under a 50% pruning ratio across different batch sizes on LLaMA-2-13B