提出可预测关键令牌的轻量级动态压缩方法,显著降低大模型推理开销。
TokenButler: Token Importance is Predictable
- 基于查询感知的细粒度预测器,动态识别对生成关键的令牌
- 在长文本任务上实现约1.6倍的GPU加速和7.6倍延迟降低
- 适合需要高效长上下文推理的部署场景,如客服、文档理解
大型语言模型依赖键值缓存(KV-Cache)存储历史标记以实现高效解码。随着缓存增长,其成为主要的内存与计算瓶颈。已有研究发现,每个解码步骤中仅有少量标记具有实质性贡献。但关键标记具有动态性且高度依赖输入查询。现有方法或因永久剔除标记而影响质量,或依赖不准确的代理指标。为此,我们提出TokenButler,一种高粒度、查询感知的预测器,能学习识别关键标记。它在固定深度步长下预测低维重要性查询,并结合缓存键的可学习投影,低成本评分标记,实现在固定预算下的动态逐标记选择,同时保留完整缓存。通过蒸馏模型的掩码因果注意力分布训练,仅引入极小参数开销。我们在新设计的小上下文共指检索任务上评估,表现接近理想基准,现有方法在此失败。此外,在RULER和LongBench等长上下文基准上取得竞争性或更优性能,采用提出的‘预测区间+邻近获取’策略,实现约1.6倍的GPU加速,误差控制在≈1.1%以内,相比密集注意力加CPU卸载,延迟降低高达7.6倍。代码已开源。
原文摘要 · Abstract (English)
Large Language Models (LLMs) rely on the Key-Value (KV) Cache to store token history, enabling efficient decoding of tokens. As the KV-Cache grows, it becomes a major memory and computation bottleneck. However, there is an opportunity to alleviate this bottleneck, prior research has shown that only a small subset of tokens contribute meaningfully to each decoding step. A key challenge in finding these critical tokens is that they are dynamic, and heavily input query-dependent. Existing methods either risk quality by evicting tokens permanently, or retain the full KV-Cache but rely on retrieving chunks of tokens and many existing KV-Cache sparsity methods rely on inaccurate proxies for token importance. To address these limitations, we introduce TokenButler, a high-granularity, query-aware predictor that learns to identify these critical tokens. TokenButler predicts low-dimensional importance queries at a fixed depth stride, and combines them with a learned projection of the real KV-cache keys to score tokens cheaply, enabling dynamic per-token selection under a fixed budget while preserving the full KV cache. We train TokenButler by distilling the model's masked causal attention distributions, optimizing a lightweight predictor with minimal parameter overhead. We evaluate TokenButler on a novel synthetic small-context co-referential retrieval task, demonstrating near-oracle accuracy where existing methods fail. Furthermore, TokenButler achieves competitive or superior performance on long-context benchmarks (RULER, LongBench), up to $\approx1.6\times$ on-GPU speedup using our proposed *prediction interval with neighbor fetching* that amortizes predictor cost while maintaining accuracy within $\approx$1.1\%, and up to 7.6$\times$ reduction in latency compared to Dense Attention with CPU offloading. Code is available: https://github.com/abdelfattah-lab/TokenButler
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。