TAPS：推論サンプリングのためのタスク認識プロポーザル分布

要旨: スペキュラティブデコーディングは、軽量なドラフトモデルにより将来のトークンを提案させ、その提案をより大きなターゲットモデルが並列に検証することで、自己回帰的生成を高速化する。しかし実際には、ドラフトモデルは通常、幅広い汎用コーパスで訓練されるため、スペキュラティブデコーディングの品質がドラフトの訓練分布にどの程度依存するのかは不明である。私たちは、この問題を軽量なHASSおよびEAGLE-2のドラフタで検討する。これらのドラフタはMathInstruct、ShareGPT、ならびに混合データの変種で訓練され、MT-Bench、GSM8K、MATH-500、SVAMPで評価する。受理長（acceptance length）による測定では、タスク固有の訓練により明確な専門化が観測される。すなわち、MathInstructで訓練したドラフトは推論ベンチマークで最も強く、ShareGPTで訓練したドラフトはMT-Benchで最も強い。混合データの訓練は頑健性を高めるが、混合のサイズを大きくしても、デコーディング温度の全域で優位になるわけではない。さらに、推論時に専門化したドラフタをどのように組み合わせるかも検討する。素朴なチェックポイント平均は不十分であるのに対し、信頼度に基づくルーティングは単一ドメインのドラフトや統合ツリー検証よりも改善し、そして2つのバックボーンの両方において、最も受理長が長くなるのは、マージされたツリー検証である。最後に、信頼度はエントロピーよりも有用なルーティング信号である。拒否されるトークンはエントロピーが高い傾向があるが、信頼度はベンチマーク単位でのより明確なルーティング判断をもたらす。これらの結果は、スペキュラティブデコーディングの品質がドラフトのアーキテクチャだけでなく、ドラフトの訓練データと下流のワークロードの適合度にも依存すること、そして専門化したドラフタは重み空間で組み合わせるよりも推論時に組み合わせる方が良いことを示している。