通过精简词表平衡生成覆盖率与延迟,提升大模型推测解码速度。
Balancing Coverage and Draft Latency in Vocabulary Trimming for Faster Speculative Decoding
- 基于训练数据中的助手回复,优化词表选择以平衡覆盖率与延迟。
- 词表缩小97%仍保持高覆盖率,领域任务延迟降低16%,吞吐提升20%。
- 适合追求高效推理的部署场景,尤其在领域特定任务中表现突出。
推测解码通过轻量级草稿模型生成候选词元,并由大型目标模型并行验证来加速大语言模型推理。以往研究发现,草稿模型常成为性能瓶颈,因其需串行生成词元且语言建模头开销随词表规模增大而升高。这暴露了草稿模型设计的根本权衡:更大的词表可提升词元覆盖率和与目标模型的一致性,但会增加延迟;更小的词表虽降低延迟,却可能遗漏关键词元。本文提出词表剪枝方法,基于领域任务仅使用全词表一小部分的观察,将词表选择建模为受覆盖率约束的优化问题。覆盖率基于训练数据中的助手回复计算,延迟则采用考虑架构特征的浮点运算量(FLOPs)估算,以捕捉语言建模头成本随词表大小的变化。利用树状帕岑估计器(Tree-structured Parzen Estimator)高效探索覆盖-延迟帕累托前沿,在最小覆盖率约束下实现最优平衡。实验表明,该方法在保持高覆盖率的同时,将草稿词表缩减高达97%,显著提升推测解码吞吐。在领域特定任务中,延迟最高降低16%,吞吐提升20%;在多样化分布外任务上,吞吐最高提升6.7%。
原文摘要 · Abstract (English)
Speculative decoding accelerates inference for Large Language Models by using a lightweight draft model to propose candidate tokens that are verified in parallel by a larger target model. Prior work shows that the draft model often dominates speculative decoding latency, since it generates tokens sequentially and incurs high cost from its language modeling head as vocabulary size grows. This exposes a fundamental trade-off in draft model design: larger vocabularies improve token coverage and agreement with the target model, but incur higher draft latency, while smaller vocabularies reduce latency at the risk of missing tokens required for accurate draft generation. We address this trade-off through vocabulary trimming for draft models, motivated by the observation that domain-specific workloads use only a small fraction of the full vocabulary. We cast draft vocabulary selection as a constrained optimization problem that balances token coverage and draft latency. Coverage is computed over assistant responses in the training data, while latency is estimated using architecture-aware FLOPs that capture the cost of the language modeling head as a function of vocabulary size. We optimize a utility function with a Tree-structured Parzen Estimator to efficiently explore the coverage-latency Pareto frontier under a minimum coverage constraint. Experiments show improved speculative decoding throughput while reducing draft vocabularies by up to 97% with high coverage. On domain-specific tasks, we achieve up to 16% latency reduction and 20% throughput improvement, and up to 6.7% throughput gains on diverse out-of-distribution tasks.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。