自适应层选择用于LLM推理中的层级token修剪
计算与语言
2026-04-17 v2 人工智能
机器学习
摘要
由于大型语言模型（LLM）的流行，针对LLM推理的键值（KV）缓存降低问题受到广泛关注。在近年来提出的众多方法中，层级token修剪方法——在特定层选择保留KV缓存中一部分token并修剪其他token——是最受欢迎的方案之一。它们主要采用预定义的层进行token选择。这种设计方式不够灵活，因为在不同任务上的准确率差异很大，在诸如KV检索等较难任务中表现会恶化。本文提出了ASL（Adaptive Layer Selection），这是一种无需训练的方法，通过利用按注意力得分排序后的token秩（rank）的方差，自适应选择KV缓存降低的选择层。该方法在不同任务之间实现性能平衡，同时满足用户指定的KV预算要求。ASL在填充阶段运行，可与诸如SnapKV等现有KV缓存降低方法联合使用，以优化解码阶段。通过在InfiniteBench、RULER和NIAH基准测试上的评估，我们展示，具备一次性token选择的ASL能够自适应地在推理速度与准确率之间进行权衡，在困难任务中超越了最先进的层级token修剪方法。
引用
@article{arxiv.2601.07667,
title = {Adaptive Layer Selection for Layer-Wise Token Pruning in LLM Inference},
author = {Rei Taniguchi and Yuyang Dong and Makoto Onizuka and Chuan Xiao},
journal= {arXiv preprint arXiv:2601.07667},
year = {2026}
}
备注
ACL 2026 Findings. Source code available at https://github.com/TANIGUCHIREI/ASL