DistillSpec: Improving Speculative Decoding via Knowledge Distillation

要約

投機的デコード (SD) は、複数のトークンを生成するためにより高速なドラフト モデルを採用することで大規模な言語モデルの推論を加速します。その後、これらのトークンはより大きなターゲット モデルによって並行して検証され、ターゲット モデルの分布に従ってテキストが生成されます。
ただし、ターゲット モデルとよく一致するコンパクトなドラフト モデルを特定するのは困難です。
この問題に取り組むために、SD を適用する前に知識の蒸留を使用してドラフト モデルをターゲット モデルとより適切に調整する DistillSpec を提案します。
DistillSpec は 2 つの重要な設計上の選択を行います。これらは、ドラフトとターゲットの整合性を改善するために重要であることが体系的な研究を通じて実証されています。それは、ドラフト モデルからのポリシーに基づくデータ生成の利用、もう 1 つはタスクとデコード戦略に合わせて発散関数を調整することです。
特に、DistillSpec は、貪欲なサンプリングと非貪欲なサンプリングの両方を使用して、さまざまな標準ベンチマークで標準 SD よりも 10 ～ 45% の大幅な高速化を実現します。
さらに、DistillSpec と非可逆 SD を組み合わせて、レイテンシーとタスクのパフォーマンスのトレードオフに対するきめ細かい制御を実現します。
最後に、さまざまなサイズのモデルを使用する実際のシナリオでは、最初に蒸留を使用してターゲット モデルのパフォーマンスを向上させ、次に DistillSpec を適用して適切に調整されたドラフト モデルをトレーニングすると、標準と比較してパフォーマンスの低下を最小限に抑えながらデコード レイテンシを 6 ～ 10 倍削減できます。
蒸留せずにデコードします。

要約(オリジナル)

Speculative decoding (SD) accelerates large language model inference by employing a faster draft model for generating multiple tokens, which are then verified in parallel by the larger target model, resulting in the text generated according to the target model distribution. However, identifying a compact draft model that is well-aligned with the target model is challenging. To tackle this issue, we propose DistillSpec that uses knowledge distillation to better align the draft model with the target model, before applying SD. DistillSpec makes two key design choices, which we demonstrate via systematic study to be crucial to improving the draft and target alignment: utilizing on-policy data generation from the draft model, and tailoring the divergence function to the task and decoding strategy. Notably, DistillSpec yields impressive 10 – 45% speedups over standard SD on a range of standard benchmarks, using both greedy and non-greedy sampling. Furthermore, we combine DistillSpec with lossy SD to achieve fine-grained control over the latency vs. task performance trade-off. Finally, in practical scenarios with models of varying sizes, first using distillation to boost the performance of the target model and then applying DistillSpec to train a well-aligned draft model can reduce decoding latency by 6-10x with minimal performance drop, compared to standard decoding without distillation.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google