TAPS: 推論的サンプリングのためのタスク対応提案分布

要旨

speculative decodingは、軽量なドラフトモデルが将来のトークンを提案し、より大規模なターゲットモデルがそれを並列に検証する仕組みにより、自己回帰的生成を高速化します。しかし実際には、ドラフトモデルは通常、広範な汎用コーパスで学習されており、 speculative decodingの品質がドラフトモデルの学習分布にどの程度依存するかは明らかではありません。本研究では、MathInstruct、ShareGPT、および混合データで学習した軽量なHASSおよびEAGLE-2ドラフトモデルを用いてこの問題を検証し、MT-Bench、GSM8K、MATH-500、SVAMPで評価しました。受け入れ長で測定すると、タスク特化型学習による明確な専門化が確認されます：MathInstructで学習したドラフトは推論ベンチマークで最も強く、ShareGPTで学習したドラフトはMT-Benchで最も強力です。混合データ学習は頑健性を向上させますが、大きな混合データセットがすべてのデコード温度で優位に立つわけではありません。また、推論時に専門化されたドラフトモデルを組み合わせる方法も検討しました。単純なチェックポイント平均化は性能が低く、信頼度ベースのルーティングは単一ドメインのドラフトを上回り、マージドツリー検証は両バックボーンにおいて全体で最高の受け入れ長を達成しました。最後に、信頼度はエントロピーよりも有用なルーティング信号です：拒否されたトークンはエントロピーが高い傾向がありますが、信頼度はベンチマークレベルではるかに明確なルーティング判断を生み出します。これらの結果は、 speculative decodingの品質がドラフトモデルのアーキテクチャだけでなく、ドラフトモデルの学習データと下流ワークロードの一致にも依存し、専門化されたドラフトモデルは重み空間で結合するよりも推論時に組み合わせる方が優れていることを示しています。

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.

TAPS: 推論的サンプリングのためのタスク対応提案分布

TAPS: Task Aware Proposal Distributions for Speculative Sampling

要旨

Support