用小模型预判重要性,高效压缩长文本推理
Draft-based Approximate Inference for LLMs
- 用小模型提前预测关键令牌和缓存重要性
- 在保持效率的同时,准确率优于现有方法
- 适合需要高速长文本处理的场景
由于Transformer架构存在二次计算复杂度和线性内存开销,优化长上下文大语言模型的推理愈发重要。现有近似推理方法如键值缓存丢弃、稀疏注意力和提示压缩,通常依赖对令牌或键值对重要性的粗略预测。本文提出一种统一框架,利用小型草稿模型更准确地预测令牌与键值对的重要性。我们提供了新颖的理论与实证分析,证明前瞻式重要性估计的有效性。在此框架下,提出:(i) SpecKV,首个结合草稿模型前瞻机制实现精准键值缓存丢弃的方法;(ii) SpecPC,利用草稿模型注意力激活识别并剔除低重要性提示令牌;(iii) SpecKV-PC,一种级联压缩策略,融合两种技术。在多个长上下文基准上的实验表明,该方法在保持相同内存、延迟与吞吐量优势的同时,持续获得更高准确率。
原文摘要 · Abstract (English)
Optimizing inference for long-context large language models (LLMs) is increasingly important due to the quadratic compute and linear memory cost of Transformers. Existing approximate inference methods, including key-value (KV) cache dropping, sparse attention, and prompt compression, typically rely on coarse predictions of token or KV pair importance. We unify and extend recent work by introducing a framework for approximate LLM inference that leverages small draft models to more accurately predict token and KV pair importance. We provide novel theoretical and empirical analyses justifying lookahead-based importance estimation techniques. Within this framework, we present: (i) SpecKV, the first method to use lookahead with a small draft model to enable precise KV cache dropping; (ii) SpecPC, which leverages draft model attention activations to identify and discard less important prompt tokens; and (iii) SpecKV-PC, a cascaded compression strategy combining both techniques. Extensive experiments on long-context benchmarks demonstrate that our methods consistently achieve higher accuracy than existing baselines while retaining the same efficiency gains in memory usage, latency, and throughput.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。