动态选择层以优化大模型推理中的注意力缓存,提升复杂任务性能。
Adaptive Layer Selection for Layer-Wise Token Pruning in LLM Inference
- 根据注意力得分的令牌排序方差,自适应选择保留缓存的层
- 在难任务上显著优于现有方法,且满足用户指定的缓存预算
- 无需训练,可与现有缓存压缩方法协同使用
由于大型语言模型(LLMs)的普及,减少推理过程中的键值(KV)缓存受到广泛关注。近年来提出的层间令牌剪枝方法通过在特定层选择部分令牌保留于缓存中,剪枝其余令牌,成为主流方案之一。然而,这些方法通常采用预定义的层进行选择,缺乏灵活性,导致不同任务下性能差异大,尤其在硬任务如KV检索中表现下降明显。本文提出ASL,一种无需训练的自适应层选择方法,利用按注意力得分排序的令牌方差来动态选择适合的层以实现缓存缩减。该方法在填充阶段运行,可与现有缓存压缩技术(如SnapKV)联合使用,优化解码阶段。在InfiniteBench、RULER和NIAH基准上的评估表明,结合一次性令牌选择的ASL能自适应地在推理速度与精度之间权衡,在困难任务中超越当前最优的层间剪枝方法。
原文摘要 · Abstract (English)
Due to the prevalence of large language models (LLMs), key-value (KV) cache reduction for LLM inference has received remarkable attention. Among numerous works that have been proposed in recent years, layer-wise token pruning approaches, which select a subset of tokens at particular layers to retain in KV cache and prune others, are one of the most popular schemes. They primarily adopt a set of pre-defined layers, at which tokens are selected. Such design is inflexible in the sense that the accuracy significantly varies across tasks and deteriorates in harder tasks such as KV retrieval. In this paper, we propose ASL, a training-free method that adaptively chooses the selection layer for KV cache reduction, exploiting the variance of token ranks ordered by attention score. The proposed method balances the performance across different tasks while meeting the user-specified KV budget requirement. ASL operates during the prefilling stage and can be jointly used with existing KV cache reduction methods such as SnapKV to optimize the decoding stage. By evaluations on the InfiniteBench, RULER, and NIAH benchmarks, we show that ASL, equipped with one-shot token selection, adaptively trades inference speed for accuracy, outperforming state-of-the-art layer-wise token pruning methods in difficult tasks.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。