动态剪枝冗余词元,加速大模型长文本推理。
Saliency-driven Dynamic Token Pruning for Large Language Models
- 根据词元重要性逐层动态剪枝,降低计算负担。
- 剪掉65%词元,节省33%~47%算力,提速最高1.75倍。
- 适配多种模型与数据集,可与缓存压缩协同使用。
尽管大语言模型(LLMs)取得成功,但在长序列推理中仍面临注意力机制导致的二次计算复杂度挑战。受神经网络特征归因可解释性启发,我们发现并非所有词元贡献相同。为此提出新颖的词元剪枝框架——显著性驱动动态词元剪枝(SDTP),基于输入上下文逐步动态剪除冗余词元。具体地,设计轻量级显著性预测模块,利用隐藏状态估计每个词元的重要性得分,并嵌入不同层实现分层剪枝。此外,提出基于排名的优化策略,最小化显著性得分与预测重要性得分间的排名偏差。大量实验表明,该框架在多种模型和数据集上具有泛化性:通过分层剪除65%输入词元,显著减少33%~47% FLOPs,推理速度提升最高达1.75倍,且性能保持相当。进一步证明,SDTP可与键值缓存压缩方法结合以实现更高效压缩。
原文摘要 · Abstract (English)
Despite the recent success of large language models (LLMs), LLMs are particularly challenging in long-sequence inference scenarios due to the quadratic computational complexity of the attention mechanism. Inspired by the interpretability theory of feature attribution in neural network models, we observe that not all tokens have the same contribution. Based on this observation, we propose a novel token pruning framework, namely Saliency-driven Dynamic Token Pruning (SDTP), to gradually and dynamically prune redundant tokens based on the input context. Specifically, a lightweight saliency-driven prediction module is designed to estimate the importance score of each token with its hidden state, which is added to different layers of the LLM to hierarchically prune redundant tokens. Furthermore, a ranking-based optimization strategy is proposed to minimize the ranking divergence of the saliency score and the predicted importance score. Extensive experiments have shown that our framework is generalizable to various models and datasets. By hierarchically pruning 65\% of the input tokens, our method greatly reduces 33\% $\sim$ 47\% FLOPs and achieves speedup up to 1.75$\times$ during inference, while maintaining comparable performance. We further demonstrate that SDTP can be combined with KV cache compression method for further compression.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。