[論文レビュー] DistillSpec: Improving Speculative Decoding via Knowledge Distillation

Yongchao Zhou, Kaifeng Lyu+6 more|arXiv (Cornell University)|Oct 12, 2023

Topic Modeling被引用数 4

ひとこと要約

DistillSpecは、オンポリシーで生成されたデータとタスク固有の分散関数を用いて、コンパクトなドラフトモデルを大きなターゲットモデルと一致させる知識蒸留を提案する。これにより、標準的な特許的推論と比較して10–45%の高速化が達成され、さらにターゲットモデルを事前に蒸留することで最大10倍の遅延低減が実現される。

ABSTRACT

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.

研究の動機と目的

• 特許的推論における高い受容率を達成するため、ドラフトモデルとターゲット言語モデルの間の整合性を向上させること。
• 標準的な知識蒸留の限界を克服し、特許的推論に最適化された分布の一致を目的とすること。
• ドラフトモデルの品質とターゲットモデルとの互換性に影響を与える要因としての設計選択（データ生成戦略と分散関数の選定）を体系的かつ包括的に調査すること。
• 損失を伴う特許的推論と組み合わせることで、遅延と性能のトレードオフを細かく制御可能にすること。
• まず大規模モデルをより小型で高品質なターゲットモデルに蒸留し、その後にDistillSpecを適用して高相性のドラフトモデルを訓練する実用的な推論パイプラインの構築。

提案手法

• 学生モデルが自らの妥当な完成形から学べるように、ドラフトモデルからのオンポリシーでのデータ生成を用いて蒸留目的を訓練する。
• 推論戦略（グリーディ vs. 非グリーディ）とタスクに応じて、分散関数（例：KLダイバージェンス、交差エントロピー）を最適化し、ターゲットモデルの出力分布との整合性を向上させる。
• コンパクトなドラフトモデルを知識蒸留によって訓練し、ターゲットモデルのトークンレベルおよびシーケンスレベルの分布を、たとえ独自のタスクパフォーマンスが低下しても、可能な限りよく近似させる。
• 損失を伴う特許的推論と統合することで、ターゲットモデルの分布からの制御された逸脱を可能にし、遅延と性能のトレードオフをより柔軟に制御できる。
• 二段階のパイプラインを提案：まず大規模モデルをより小型で高品質なターゲットモデルに蒸留し、その後にDistillSpecを用いて高相性のドラフトモデルを訓練する。

実験結果

リサーチクエスチョン

• RQ1 ドラフトモデルからのオンポリシーでのデータ生成は、特許的推論におけるドラフトモデルとターゲットモデルの整合性にどのように影響するか？
• RQ2 特許的推論における知識蒸留の最適な分散関数は何か？また、推論戦略（グリーディ vs. 非グリーディ）に応じてどのように変化するか？
• RQ3 DistillSpecは、多様な自然言語処理タスクやモデルサイズにおいて、特許的推論の受容率とブロック効率を向上させることができるか？
• RQ4 DistillSpecを損失を伴う特許的推論と組み合わせた場合、標準的な特許的推論と比較して遅延と性能のトレードオフにどのような影響を与えるか？
• RQ5 複数の異なるサイズのモデルが利用可能な状況下で、最も効果的な推論パイプラインは何か？蒸留処理は特許的推論の前に行うべきか、それとも後に実施すべきか？

主な発見

• DistillSpecは、GSM8K、XSum、CNNDMなど複数のベンチマークで、グリーディおよび非グリーディ推論の両方において、標準的な特許的推論と比較して10–45%の高速化を達成した。
• 蒸留の過程でオンポリシーで生成されたデータを用いることで、ドラフトモデルとターゲットモデルの整合性が顕著に向上し、受容率とブロック効率が向上した。
• 蒸留の最適な分散関数は推論戦略に強く依存する：グリーディ推論ではKLダイバージェンスが最も優れた性能を示し、非グリーディサンプリングでは交差エントロピーがより効果的であった。
• DistillSpecと損失を伴う特許的推論を組み合わせることで、より平坦で高品質な遅延-性能トレードオフ曲線が得られ、性能の低下を最小限に抑えながらより大きな遅延の削減が可能になった。
• ターゲットモデルをまず蒸留し、その後にDistillSpecを適用する順序で実施した場合、標準的な推論と比較して、推論遅延を6–10倍まで低減したが、性能の劣化は最小限に抑えられた。
• GSM8Kでは、非グリーディサンプリングと最適なハイパーパrameterを用いた場合、DistillSpecは標準的な特許的推論と比較してRouge2スコアを15%向上させ、ブロック効率を25%向上させた。