Leveraging KV Similarity for Online Structured Pruning in LLMs

TL;DR

This paper presents Token Filtering, an online structured pruning method for LLMs that reduces inference costs by skipping redundant attention computations based on token similarity, without requiring calibration data.

Contribution

It introduces a novel online pruning technique using token similarity that improves stability and performance over prior offline methods, with no extra memory overhead.

Findings

Outperforms prior pruning methods on LLaMA-2, LLaMA-3, and Mistral models.

Maintains high accuracy on reasoning benchmarks with 50% pruning.

No additional memory overhead introduced by the method.

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.