投機的デコード

投機的デコード

投機的デコードは、現代のAIにおける最も賢い最適化手法の一つです。小さくて高速なモデルを使って、大きなモデルがより速くテキストを生成できるようにします。基本的な考え方はシンプルです。小さなモデルに大きなモデルが何を言うかを推測させ、その推測を大きなモデルが一度に検証するというものです。

通常、大きなモデルは一度に1つのトークンを生成し、各トークンはメモリからモデルのすべての重みをロードする必要があります。これは、以前に説明したメモリ帯域幅のボトルネックのために遅くなります。しかし、大きなモデルが一度に複数のトークンを生成できれば、はるかに高速になります。問題は、大きなモデルは中間トークンを見ずに複数のトークンを先に予測できないことです。

投機的デコードは、メインモデルの10分の1程度の小さなドラフトモデルを使用して、複数の候補トークンを生成することでこの問題を解決します。ドラフトモデルは小さいため高速です。その後、大きなモデルがすべての候補トークンを一度に並列処理し、どのトークンに同意するかをチェックします。ドラフトモデルはかなり優れているため、その推測のほとんどは正しく、大きなモデルはそれらを受け入れます。

重要な洞察は、トークンのバッチを検証するのは、1つのトークンを生成するのとほぼ同じ速さであるということです。なぜなら、重みはバッチごとに一度だけロードすればよいからです。したがって、ドラフトモデルが5つのトークンを生成し、大きなモデルがそれらすべてを受け入れた場合、通常は1つのトークンを生成するのにかかる時間で5つのトークンを生成したことになります。これは、ドラフトモデルが完全に正確である場合、5倍の高速化を実現します。

実際には、投機的デコードは実際のワークロードで2倍から3倍の高速化を実現します。高速化の度合いは、ドラフトモデルが大きなモデルにどれだけ適合しているかに依存します。ドラフトモデルが大きく異なる場合、その推測のほとんどが拒否され、高速化は最小限になります。しかし、適切に適合したドラフトモデルを使えば、投機的デコードはフリーランチです：品質を損なうことなく、より高速な生成を実現します。