Probe and Skip: Self-Predictive Token Skipping for Efficient Long-Context LLM Inference

TL;DR

This paper introduces Self-Predictive Token Skipping (SPTS), a training-free method that improves long-context LLM inference efficiency by selectively skipping tokens using probing strategies, achieving significant speedups with maintained accuracy.

Contribution

The paper proposes a novel, training-free token skipping framework with probing strategies and multi-stage pruning, addressing structure optimization and redundancy issues in long-context LLM inference.

Findings

Achieves up to 2.46× speedup in prefilling

Achieves up to 2.29× speedup in end-to-end generation

Maintains state-of-the-art accuracy

Abstract

Long-context inference enhances the reasoning capability of Large Language Models (LLMs), but incurs significant computational overhead. Token-oriented methods, such as pruning and skipping, have shown great promise in reducing inference latency, yet still suffer from inherently insufficient structure optimization, outdated selection criteria, and redundancy interference, resulting in suboptimal speed-accuracy trade-off. To address these issues, we propose a novel training-free framework dubbed Self-Predictive Token Skipping (SPTS), for efficient long-context LLM inference. Specifically, motivated by probing the influence of target layers prior to skipping, we design two selective token skipping strategies for typical structures, including Partial Attention Probing (PAP) for multi-head attention and Low-rank Transformation Probing (LTP) for feed forward network. The former selects informative tokens via partial forward attention computation, while the latter constructs a low-rank proxy network to predict token transformations. In addition, a Multi-Stage Delayed Pruning (MSDP) strategy reallocates skipping budgets and progressively removes redundant tokens across layers. Extensive experiments display the effectiveness of our method, achieving up to 2.46$\times$ and 2.29$\times$ speedups for prefilling and end-to-end generation, respectively, while maintaining state-of-the-art accuracy. We will release the source code upon acceptance.