Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Zimeng Wu, Donghao Wang, Chaozhe Jin, Jiaxin Chen, Yunhong Wang

TL;DR

Long-context LLM inference incurs peak compute costs due to dense attention and FFN processing. SPTS introduces a training-free token-skipping framework that leverages self-predictive signals: Partial Attention Probing for MHA and Low-rank Transformation Probing for FFN, complemented by Multi-Stage Delayed Pruning to allocate skipping budgets across layers. By reducing work on less informative tokens while preserving the core information flow, SPTS achieves notable speedups (up to 2.46× TTFT and 2.29× E2E) with minimal accuracy loss across multiple models and tasks. The approach demonstrates strong generalization and practical applicability for long-context inference in real systems.

Abstract

Long-context inference enhances the reasoning capability of Large Language Models (LLMs) while incurring significant computational overhead. Token-oriented methods, such as pruning and skipping, have shown promise in reducing inference latency, but still suffer from inherently limited acceleration potential, outdated proxy signals, and redundancy interference, thus yielding suboptimal speed-accuracy trade-offs. To address these challenges, we propose SPTS (Self-Predictive Token Skipping), a training-free framework for efficient long-context LLM inference. Specifically, motivated by the thought of probing the influence of targeted skipping layers, we design two component-specific strategies for selective token skipping: Partial Attention Probing (PAP) for multi-head attention, which selects informative tokens by performing partial forward attention computation, and Low-rank Transformation Probing (LTP) for feed forward network, which constructs a low-rank proxy network to predict token transformations. Furthermore, a Multi-Stage Delayed Pruning (MSDP) strategy reallocates the skipping budget and progressively prunes redundant tokens across layers. Extensive experiments demonstrate 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 model performance. The source code will be publicly available upon paper acceptance.

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

TL;DR

Long-context LLM inference incurs peak compute costs due to dense attention and FFN processing. SPTS introduces a training-free token-skipping framework that leverages self-predictive signals: Partial Attention Probing for MHA and Low-rank Transformation Probing for FFN, complemented by Multi-Stage Delayed Pruning to allocate skipping budgets across layers. By reducing work on less informative tokens while preserving the core information flow, SPTS achieves notable speedups (up to 2.46× TTFT and 2.29× E2E) with minimal accuracy loss across multiple models and tasks. The approach demonstrates strong generalization and practical applicability for long-context inference in real systems.

Abstract

Long-context inference enhances the reasoning capability of Large Language Models (LLMs) while incurring significant computational overhead. Token-oriented methods, such as pruning and skipping, have shown promise in reducing inference latency, but still suffer from inherently limited acceleration potential, outdated proxy signals, and redundancy interference, thus yielding suboptimal speed-accuracy trade-offs. To address these challenges, we propose SPTS (Self-Predictive Token Skipping), a training-free framework for efficient long-context LLM inference. Specifically, motivated by the thought of probing the influence of targeted skipping layers, we design two component-specific strategies for selective token skipping: Partial Attention Probing (PAP) for multi-head attention, which selects informative tokens by performing partial forward attention computation, and Low-rank Transformation Probing (LTP) for feed forward network, which constructs a low-rank proxy network to predict token transformations. Furthermore, a Multi-Stage Delayed Pruning (MSDP) strategy reallocates the skipping budget and progressively prunes redundant tokens across layers. Extensive experiments demonstrate the effectiveness of our method, achieving up to 2.46 and 2.29 speedups for prefilling and end-to-end generation, respectively, while maintaining state-of-the-art model performance. The source code will be publicly available upon paper acceptance.
Paper Structure (24 sections, 16 equations, 3 figures, 3 tables)

This paper contains 24 sections, 16 equations, 3 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. LLM Inference and Efficiency
  4. Token Pruning
  5. Token Skipping
  6. Methodology
  7. Framework Overview
  8. Partial Attention Probing for MHA
  9. Low-rank Transformation Probing for FFN
  10. Low-rank Proxy FFN
  11. Conditioned Transformation for Token Selection
  12. Multi-Stage Delayed Token Pruning
  13. Experimental Results and Analysis
  14. Experimental Settings
  15. Models
  16. ...and 9 more sections

Figures (3)

  • Figure 1: Comparison of efficiency and accuracy across different methods for LLaMA-3.1-8B-Instruct llama8b using LongBench longbench dataset. (a) Our proposed SPTS accelerates both prefilling and decoding without sacrificing either phase, achieving consistent speedups in time to first token (TTFT) and end-to-end generation latency (E2E). (b) SPTS maintains higher accuracy at a more aggressive TTFT acceleration ratio.
  • Figure 2: Framework overview of our proposed method. (a) Partial Attention Probing (PAP) performs token skipping in MHA by pre-computing a small fraction of attention to identify informative tokens, which then participate in reduced attention computation. (b) Low-rank Transformation Probing (LTP) enables token skipping in FFN by leveraging an offline-derived low-rank proxy network to probe token transformations, which are combined with attention-based information to guide token selection for reduced feed-forward. (c) Multi-Stage Delayed Pruning (MSDP) partitions the network into multiple stages, where the set of candidate tokens remains fixed within each stage and is reduced through pruning at stage boundaries.
  • Figure 3: Comparison of (a) TTFT acceleration ratios and (b) E2E acceleration ratios with different methods across distinct context lengths on the LLaMA3.1-8B-Instruct model.