用目标模型的KV缓存提升长程推测解码准确率
When Hidden States Drift: Can KV Caches Rescue Long-Range Speculative Decoding?
- 让草稿模型复用目标模型的KV缓存,获得更完整的上下文信息
- 实验显示长程接受率显著提升,但整体加速效果仍有限
- 揭示了浅层草稿模型与梯度稀疏两大关键瓶颈
推测解码可加速大模型推理,但现有基于隐藏状态的草稿模型存在长程衰减问题:推测步数越多,准确性越低。尽管已有工作归因于训练-推理不匹配并提出测试时训练(TTT)作为解决方案,但我们发现即使在TTT训练后,长程衰减依然存在。本文从上下文信息保留角度重新审视该问题:隐藏状态重用将目标隐藏状态视为有偏的上下文压缩——它根据当前位置的注意力查询聚合历史标记信息,得到一个仅优化即时下一个词预测的紧凑表示,可能抑制对后续推测步骤重要的信息。而目标模型的KV缓存则显式保留所有标记的键值对表示。因此我们提出KV重用假设:允许草稿模型复用目标模型的KV缓存,可提供更丰富的长程信号。为此,我们设计了KVShot诊断框架,对比三种重用范式:仅隐藏状态、仅KV缓存、混合模式。在Qwen3-8B上的大量实验表明,KV重用能有效提升长程接受率,但当前训练流程下端到端加速仍不明显。分析揭示两个关键结构瓶颈:浅层草稿模型难以准确估计目标查询,且草稿侧的KV投影接收稀疏梯度信号。这表明要实现KV感知解码的全部潜力,需超越TTT,转向块级训练范式。通过暴露这些瓶颈,KVShot为下一代推理架构设计提供了基础诊断平台和清晰路线图。
原文摘要 · Abstract (English)
Speculative decoding accelerates LLM inference, but SOTA hidden-state-based drafters suffer from long-range decay: draft accuracy degrades as the speculative step increases. Existing work attributes this decay to train-inference mismatch and proposes test-time training (TTT) as a remedy, yet we observe that long-range decay persists even in TTT-trained drafters. We revisit long-range decay from the perspective of context information preservation. In hidden-state reuse, we argue the target hidden state acts as a biased context compression: it aggregates historical token information according to the attention query at the current position, yielding a compact representation optimized for immediate next-token prediction. This compression can suppress information less relevant to the current query but important for later speculative steps. In contrast, the target model's KV cache serves as an explicit context, retaining the complete set of token-wise KV representations. We therefore posit the KV-Reuse Hypothesis: allowing the draft model to reuse the target KV cache can provide richer signals for long-horizon drafting. To test this hypothesis, we introduce KVShot, a diagnostic framework that compares three reuse paradigms: hidden-only, KV-only, and hybrid. Extensive evaluations on Qwen3-8B show that KV-Reuse improves long-range acceptance, although end-to-end speedups remain marginal under current training pipelines. Our analysis identifies two key structural bottlenecks: shallow drafters struggle to estimate target queries accurately, and draft-side KV projections receive sparse gradient signals. These findings suggest that realizing the full potential of KV-aware decoding requires moving beyond TTT toward block-wise training paradigms. By exposing these bottlenecks, KVShot provides a foundational diagnostic testbed and a clear roadmap for designing next-generation inference architectures.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。