Re Reference AI

技術

投機的デコーディングとは

Speculative Decoding

小規模な下書きモデルに複数トークンを先に生成させ、本体モデルで一括検証し生成を高速化する手法

推論最適化デコーディング

ひとことで言うと

小さいAIに先に答えの候補を作らせ、大きいAIがまとめて確認することで生成を速くする方法。

概要

投機的デコーディング(Speculative Decoding)とは、軽量で高速な「下書きモデル(draft model)」に複数トークンをあらかじめ生成させ、それを大規模な「本体モデル(target model)」で一括検証する手法。 一致した部分をまとめて確定させることで、LLMの生成を高速化する。 本体モデルは1トークンずつ逐次生成する代わりに、下書きモデルが提案した複数トークンの妥当性を並列に検証するだけで済むため、出力の分布を変えずに生成速度を向上できる点が特徴。 下書きモデルの予測が本体モデルの判断と一致する頻度に応じて、高速化の効果は大きくなる。下書きモデルには、本体モデルと似た傾向を持つ小型モデルが選ばれやすい。

背景

自己回帰型のLLMは1トークンずつ逐次生成する必要があり、この逐次性がGPUの並列計算能力を十分に活かせないボトルネックになっていた。 投機的デコーディングは、生成の逐次性そのものは残しつつ複数トークン分の検証を1回のモデル呼び出しにまとめることで、GPUの並列処理能力を有効活用し体感速度を高めるために考案された。

歴史

2023年: LeviathanらとChenらが、それぞれ独立に投機的デコーディングの手法を発表し、大規模モデルの生成速度を保証された品質のまま高速化できることを示す。

アーキテクチャ

下書きモデルが複数トークン(例: 4〜8個)を先に生成し、本体モデルはこれらのトークン全てを含む系列を1回のフォワードパスで処理して各トークンの確率を計算する。 本体モデルの確率と下書きモデルの提案が十分に一致するトークンまでをそのまま採用し、最初に食い違ったトークン以降は本体モデルの分布から生成し直すことで、最終的な出力分布を本体モデル単体の場合と統計的に一致させる。

ワークフロー

下書きモデルが複数トークンを高速に生成 → 本体モデルがそれら全てを1回のフォワードパスで検証 → 一致した部分までを採用し確定 → 食い違った箇所以降を本体モデルの分布から生成し直す。

コード例

下書きモデルと本体モデルによる投機的デコーディングの骨格

def speculative_generate(prompt, draft_model, target_model, k=5, max_tokens=100):
    tokens = tokenize(prompt)
    while len(tokens) < max_tokens:
        draft_tokens = draft_model.generate(tokens, num_tokens=k)
        accepted = target_model.verify(tokens, draft_tokens)
        tokens += accepted
        if len(accepted) < k:
            tokens += target_model.generate(tokens, num_tokens=1)
    return detokenize(tokens)

利点

  • 出力の品質・分布を損なうことなく、生成速度を向上できる
  • 本体モデルへの追加学習を必要とせず、推論時の仕組みだけで導入できる
  • 本体モデル単体で生成する場合と比べ、GPUの並列計算資源をより有効に活用できる

欠点

  • 下書きモデルと本体モデルの予測が大きく食い違うタスクでは、高速化の効果が小さくなる
  • 下書きモデルの用意・運用等、システム構成が複雑になる
  • 下書きモデルと本体モデルの両方をメモリ上に保持する必要があり、メモリ使用量が増える

比較

  • デコーディング投機的デコーディングは、デコーディングの速度を高速化する手法
  • KVキャッシュ投機的デコーディングは、KVキャッシュ等と組み合わせてさらに推論を最適化できる
  • SLM投機的デコーディングの下書きモデルには、多くの場合SLMのような軽量なモデルが使われる

関連用語

デコーディングKVキャッシュSLM推論

よくある質問

投機的デコーディングで出力の品質は落ちない?

理論上、最終的に採用されるトークンは本体モデル単体で生成した場合と統計的に同じ分布に従うよう設計されており、出力品質を保ったまま速度だけを改善できるとされる。

下書きモデルは本体モデルと同じアーキテクチャである必要がある?

必須ではないが、トークナイザーを共有し出力の傾向が近いモデルほど提案が一致しやすく、高速化の効果を得やすいとされる。

参考文献

関連Zenn記事