arXiv:2603.27027cs.CLcs.AI2026-03

让草稿模型按任务定制,能显著提升生成速度和质量。

TAPS: Task Aware Proposal Distributions for Speculative Sampling

  • 用特定任务数据训练草稿模型,使其更匹配下游任务。
  • 数学任务专用模型在推理题上接受长度最高,达1.8倍提升。
  • 基于置信度路由比平均权重更有效,适合实际部署。

推测解码通过轻量级草稿模型预先提出未来词元,由大型目标模型并行验证以加速生成。然而,草稿模型通常在通用语料上训练,其性能与训练分布的关系尚不明确。本文研究了轻量级HASS和EAGLE-2草稿模型在MathInstruct、ShareGPT及混合数据上的表现,评估基准包括MT-Bench、GSM8K、MATH-500和SVAMP。结果显示,任务特定训练带来明显专化:MathInstruct训练的草稿在推理类任务上最优,而ShareGPT训练的在MT-Bench上表现最好;混合数据提升鲁棒性,但更大混合比例未在不同温度下持续占优。进一步发现,简单检查点平均效果差,而基于置信度的路由优于单一领域模型,合并树验证在两类骨干网络上均达到最高接受长度。此外,置信度是比熵更有效的路由信号——被拒绝词元熵更高,但置信度能产生更清晰的基准级决策。结果表明,推测解码质量不仅取决于草稿结构,更受训练数据与下游任务匹配度影响,且专用草稿模型应在推理时组合而非合并权重。

原文摘要 · Abstract (English)

Speculative decoding accelerates autoregressive generation by letting a lightweight draft model propose future tokens that a larger target model then verifies in parallel. In practice, however, draft models are usually trained on broad generic corpora, which leaves it unclear how much speculative decoding quality depends on the draft training distribution. We study this question with lightweight HASS and EAGLE-2 drafters trained on MathInstruct, ShareGPT, and mixed-data variants, evaluated on MT-Bench, GSM8K, MATH-500, and SVAMP. Measured by acceptance length, task-specific training yields clear specialization: MathInstruct-trained drafts are strongest on reasoning benchmarks, while ShareGPT-trained drafts are strongest on MT-Bench. Mixed-data training improves robustness, but larger mixtures do not dominate across decoding temperatures. We also study how to combine specialized drafters at inference time. Naive checkpoint averaging performs poorly, whereas confidence-based routing improves over single-domain drafts and merged-tree verification yields the highest acceptance length overall for both backbones. Finally, confidence is a more useful routing signal than entropy: rejected tokens tend to have higher entropy, but confidence produces much clearer benchmark-level routing decisions. These results show that speculative decoding quality depends not only on draft architecture, but also on the match between draft training data and downstream workload, and that specialized drafters are better combined at inference time than in weight space.

推测解码草稿模型任务适配路由策略

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。