通过自预测机制跳过冗余分词,让长文本推理快又准。
Probe and Skip: Self-Predictive Token Skipping for Efficient Long-Context LLM Inference
- 先探测关键层的分词影响,再决定跳过哪些
- 预填充和生成阶段分别提速2.46倍和2.29倍
- 无需训练,适合追求高效长文本推理的人
长上下文推理虽提升大模型推理能力,但带来显著计算开销。现有基于分词的方法虽有效降低延迟,但仍存在结构优化不足、选择标准陈旧和冗余干扰等问题,导致速度与精度平衡不佳。为此,我们提出无需训练的自预测分词跳过(SPTS)框架。核心思想是先探测目标层在跳过前的影响:针对多头注意力设计部分注意力探测(PAP),通过部分前向计算筛选重要分词;针对前馈网络设计低秩变换探测(LTP),用低秩代理网络预测分词变化。此外,采用多阶段延迟剪枝(MSDP)策略,在各层动态分配跳过预算,逐步移除冗余分词。大量实验表明,该方法在预填充阶段实现最高2.46倍加速,端到端生成阶段达2.29倍加速,同时保持最先进的准确性。代码将在录用后公开。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。