推論速度を向上させる Speculative Decoding(投機的デコーディング)とは

こんにちは Qualiteg 研究部です。

投機的デコーディングとは何か？

投機的デコーディングは、大規模言語モデル（LLM）の推論速度を向上させる技術です。

たいていのモデルを1.4～2.0倍程度、高速化できることが報告されています。

このアプローチでは、小さなモデル（ドラフトモデル）を使用して初期の予測を行い、その結果を大きなモデル（ターゲットモデル）が検証することで、全体の推論プロセスを高速化します。

ざっくりいうと、

大きなモデルは計算負荷も高く計算速度も遅いので、まず、小さなモデルで高速に計算したあとで、その計算結果をうまくつかって大きなモデルでの計算負荷をさげ、スピードを向上させようというアイデアです。

基本的に大きなモデルと、小さなモデルはサイズ以外は基本的にまったく同じネットワーク構造をしていることが前提となります。

たとえば 70Bの Llama3 と 8B の Llama3 を組み合わせてつかうイメージです。

当然70B の Llama3 の推論計算のほうが 8B よりも重たくなりますので、小さい8BのLlama3 で先回りして推論計算することで高速化を行うというテクニックとなります。

投機的デコーディングのメカニズム

投機的デコーディングでは、小さなモデル（ドラフトモデル）の予測結果を大きなモデル（ターゲットモデル）で使用するかどうかを判断する際、主に以下の手順と考慮点があります

1. ドラフトモデルの生成: ドラフトモデルは、予測の初期段階で複数の候補トークンを高速に生成します。このモデルはターゲットモデルよりもはるかに小さいため、予測を迅速に行うことができます。
   
2. ターゲットモデルによる検証: ターゲットモデルは、ドラフトモデルが生成したトークンを検証し、それらが妥当であるかどうかを判断します。このプロセスでは、ドラフトモデルの出力とターゲットモデルの予測を比較し、一致するトークンのみが最終的な出力として採用されます。
   
3. TAR（Token Acceptance Rate）の計算: TARは、ドラフトモデルが生成したトークンのうち、ターゲットモデルが受け入れたトークンの割合を示します。この割合が高いほど、ドラフトモデルの予測がターゲットモデルの基準に適合していることを意味し、スループットの向上に貢献します。
   
4. スループットとレイテンシーのトレードオフ: ドラフトモデルを使用する主な目的は、推論プロセスのスループットを向上させることです。ドラフトモデルのレイテンシーが十分に低く、かつTARが高い場合、このアプローチは全体の推論時間を短縮し、効率を向上させることができます。
   
5. パフォーマンスのベンチマーク: 実際にドラフトモデルとターゲットモデルを使用する際には、異なるドラフトモデルの構成とサイズで複数の実験を行い、最適な設定を見つける必要があります。これにより、どのドラフトモデルが最も効果的であるかを科学的に判断することが可能です。

以上の手順と考慮点によって、小さなモデルの予測結果が大きなモデルで実用的に使用できるかどうかを判断することができます。

ドラフトモデルでの計算結果をターゲットモデルが評価するときに、結局ターゲットモデルでの推論計算が走るから、計算量削減にはならないのではないか？

そんな疑問が浮かびませんか？

ターゲットモデルで計算を行うとなると、なぜ小さなモデルを使うのか疑問に思うのは理解できます。

投機的デコーディングの利点（というか、コアとなるアイデア）は、ターゲットモデルの計算負荷を効率的に管理する点にあります。ここでは、計算が削減されるメカニズムを具体的に説明します。

投機的デコーディングの基本プロセス

1. ドラフトモデルの利用:
   ドラフトモデルは、低レイテンシーで多数の候補トークンを生成します。これはターゲットモデルよりもはるかに迅速に行われます。
   
2. バッチ処理
   ターゲットモデルでは、ドラフトモデルが生成した複数のトークンを一度に検証します。これは通常のオートリグレッシブ生成（トークンを1つずつ生成）と比べて、モデルが一度に多くのデータを処理できるため、GPUなどの計算リソースを効率的に利用できます。
   
3. プリフィル手法:
   ターゲットモデルは、ドラフトモデルが生成した複数のトークンに基づいて予測を行い、これを一種のプリフィル（事前充填）として使用します。ターゲットモデルがすべての候補を1つずつ独立に生成する代わりに、有効なトークンのセットを確認し、受け入れることで、計算を省略します。

実際の計算削減のポイント

• 並列処理
  ターゲットモデルがドラフトモデルから提供されたトークン群をバッチで処理することにより、トークンごとの生成ではなく、効率的な並列処理が可能になります。
  
• 選択的検証
  ターゲットモデルは有効と判断したトークンのみを受け入れます。これにより、全体的な生成プロセスのステップ数が減少し、無駄な計算が省かれます。
  
• 効率的なデータ処理
  ドラフトモデルからの入力を使用することで、ターゲットモデルは入力の一部としてこれを活用し、全体の計算負荷を削減します。

まとめ

いかがでしたでしょうか、今回はなるべく数式を用いずに、投機的デコーディングについて解説してみました。

投機的デコーディングでは、確かにターゲットモデルで最終的な計算が行われますが、ドラフトモデルの出力を利用して効率的に処理を行うことで、全体の計算コストとレイテンシーを削減できます。この方法により、ターゲットモデルの計算負担が軽減され、より迅速かつ効率的なデータ処理が可能になります。

参考文献

https://arxiv.org/pdf/2211.17192
https://arxiv.org/pdf/2302.01318

論文「2402.01528v2」と「2211.17192v2」によりますと、投機的デコーディングの有効性はドラフトモデルの選定に大きく依存しているようです。

これらの研究では、異なるドラフトモデルがどのようにターゲットモデルの性能に影響を与えるかを検証していますが、とくにトークン受容率（TAR）＝ドラフトモデルが生成したトークンのうち、ターゲットモデルがどれだけ受け入れるかが、スループット向上の鍵を握るようです。当然といえば当然で、ドラフトモデルがイケてるトークン(logits)をどれだけ出せるか、ですね。