Balancing Coverage and Draft Latency in Vocabulary Trimming for Faster Speculative Decoding
作者: Ofir Ben Shoham
分类: cs.CL, cs.AI, cs.LG
发布日期: 2026-03-05
💡 一句话要点
提出词汇表裁剪方法，平衡覆盖率与延迟，加速推测解码。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 推测解码 词汇表裁剪 语言模型加速 约束优化 树状Parzen估计器
📋 核心要点
- 现有推测解码中，draft模型的延迟成为瓶颈，而增大词汇量虽能提高覆盖率，但会显著增加延迟。
- 论文提出词汇表裁剪方法，通过约束优化平衡token覆盖率和draft延迟，选择领域相关的精简词汇表。
- 实验表明，该方法在大幅减少词汇量的同时，提高了推测解码的吞吐量，尤其在特定领域任务上效果显著。
📝 摘要（中文）
推测解码通过使用轻量级的draft模型生成候选token，并由大型target模型并行验证，从而加速大型语言模型的推理。以往研究表明，draft模型通常主导推测解码的延迟，因为它顺序生成token，并且其语言模型头的成本随着词汇量的增长而增加。这暴露了draft模型设计中的一个根本权衡：更大的词汇量提高了token覆盖率和与target模型的一致性，但导致更高的draft延迟，而更小的词汇量降低了延迟，但存在错过准确draft生成所需token的风险。我们通过draft模型的词汇表裁剪来解决这种权衡，其动机是特定领域的任务只使用完整词汇表的一小部分。我们将draft词汇表选择视为一个约束优化问题，该问题平衡了token覆盖率和draft延迟。覆盖率是在训练数据中的助手响应上计算的，而延迟是使用架构感知的FLOPs估计的，该FLOPs捕获了语言模型头的成本作为词汇量的函数。我们使用树状Parzen估计器优化效用函数，以在最小覆盖率约束下有效地探索覆盖率-延迟Pareto前沿。实验表明，在保持高覆盖率的同时，推测解码吞吐量得到提高，并且draft词汇量最多减少了97%。在特定领域的任务上，我们实现了高达16%的延迟降低和20%的吞吐量提高，并且在各种分布外任务上实现了高达6.7%的吞吐量增益。
🔬 方法详解
问题定义：推测解码中，draft模型的延迟是影响整体性能的关键因素。传统的做法是使用完整的词汇表，但这会增加语言模型头的计算成本，尤其是在词汇量很大时。现有方法的痛点在于，无法有效地平衡词汇量大小与token覆盖率之间的关系，导致要么延迟过高，要么覆盖率不足。
核心思路：论文的核心思路是针对特定任务，裁剪draft模型的词汇表，只保留最常用的token。通过减少词汇量，降低语言模型头的计算成本，从而减少draft延迟。同时，通过优化选择策略，保证裁剪后的词汇表仍然具有足够的token覆盖率，以确保draft模型的准确性。
技术框架：整体框架包含以下几个主要步骤：1) 数据收集：收集目标任务的训练数据，特别是助手响应数据。2) 覆盖率计算：在训练数据上计算每个token的出现频率，评估其覆盖率。3) 延迟估计：使用架构感知的FLOPs模型，估计不同词汇量大小下的draft延迟。4) 优化问题构建：将词汇表选择问题建模为一个约束优化问题，目标是最大化效用函数，该函数平衡了token覆盖率和draft延迟，同时满足最小覆盖率约束。5) 优化求解：使用树状Parzen估计器（TPE）来高效地探索覆盖率-延迟Pareto前沿，找到最优的词汇表。
关键创新：最重要的技术创新点在于将词汇表选择问题建模为一个约束优化问题，并使用TPE算法来求解。与简单的频率截断方法相比，该方法能够更有效地平衡覆盖率和延迟，找到更优的词汇表。此外，使用架构感知的FLOPs模型来估计延迟，使得优化过程更加准确。
关键设计：关键设计包括：1) 效用函数：效用函数的设计需要平衡覆盖率和延迟，可以采用加权平均的方式。2) 最小覆盖率约束：设置最小覆盖率约束，以保证裁剪后的词汇表仍然具有足够的表达能力。3) TPE算法：使用TPE算法来高效地探索搜索空间，找到最优的词汇表。4) 架构感知的FLOPs模型：使用与目标硬件架构相关的FLOPs模型来估计延迟，以提高优化结果的准确性。
🖼️ 关键图片
📊 实验亮点
实验结果表明，该方法在保持高覆盖率的同时，可以将draft模型的词汇量最多减少97%。在特定领域的任务上，实现了高达16%的延迟降低和20%的吞吐量提高。在各种分布外任务上，也实现了高达6.7%的吞吐量增益。这些结果表明，该方法能够有效地平衡覆盖率和延迟，提高推测解码的性能。
🎯 应用场景
该研究成果可广泛应用于各种需要加速大型语言模型推理的场景，例如智能客服、对话机器人、代码生成等。通过针对特定领域进行词汇表裁剪，可以显著降低推理延迟，提高用户体验，并降低计算成本。未来，该方法可以进一步扩展到多语言场景和更复杂的模型架构。
📄 摘要（原文）
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.