自己スペキュラティブデコーディングによる高速テキスト生成

LayerSkipメソッドによって実装された自己スペキュラティブデコーディングは、単一の大規模言語モデル(LLM)が自身の早い層を使ってトークンをドラフトし、より深い層でそれらを検証することにより、テキスト生成を高速化します。このアプローチは、2つの別々のモデルを必要とする従来のスペキュラティブデコーディングに関連するメモリフットプリントと計算レイテンシを削減します。

自己スペキュラティブ対従来のスペキュラティブデコーディング

従来のスペキュラティブデコーディングは、トークンのシーケンスを生成する小さなドラフトモデルと、それらを検証する大きな検証モデルに依存しています。これに対し、自己スペキュラティブデコーディングは両方の役割に同じモデルを使用します。中間層でモデルを終了してドラフトトークンを生成することにより、セカンダリモデルの重みの必要性がなくなり、これによりメモリが節約され、大規模推論における全体のハードウェアフットプリントが削減されます。

技術的実装: アーリーイグジットとアンエンベディング

自己スペキュラティブデコーディングを有効にするためには、モデルは「アーリーイグジット」が可能である必要があります。ここで、生成プロセスは事前に指定された中間層で停止します。

ロジットのアンエンベディング

言語モデル(LM)ヘッドは通常、最終層のみで訓練されるため、中間層からのロジットを「アンエンベディング」するための特定の訓練が必要です。このプロセスは、中間層の出力をLMヘッドに投影して次のトークンを予測します。

トレーニングの変更点

プリトレーニングまたはファインチューニング中に、アーリーイグジットを可能にするために2つの主要な変更が使用されます：

レイヤードロップアウト: モデルは特定のレイヤーをスキップするように訓練され、ドロップアウト率は深いレイヤーほど増加します。これにより、後続のレイヤーへの依存が減少し、一般化が向上します。
アーリーイグジット損失: 各エグジット（中間層）に正規化された損失が適用され、LMヘッドがモデル全体のさまざまな深さからの出力をアンエンベディングする方法を学ぶように強制されます。

推論プロセス: ドラフトと検証

推論は2つの異なる段階で行われます：

セルフドラフト: 中間層でモデルを終了することによりトークンが生成されます。スペキュラティブトークンの数と特定のエグジットレイヤーは、速度と精度のバランスを取るハイパーパラメータです。
セルフベリフィケーション: フルモデルがドラフトトークンを検証します。これを最適化するために、システムはドラフトフェーズ中にキャッシュされた早い層の結果を再利用し、検証のために残りの層のみを計算します。

ハードウェアとパフォーマンスの最適化

自己スペキュラティブデコーディングは、冗長な計算を最小限に抑えるためにKVQキャッシュ（KVキャッシュとエグジットクエリキャッシュの組み合わせ）を使用します。これにより、2モデルシステムに対して3つの主要な利点があります：

重みの共有: 最初の$E$層はドラフトと検証の両方で再利用されます。
KVキャッシュの共有: 最初の$E$層のキー・バリューペアが再利用されます。
計算の共有: エグジットクエリキャッシュはエグジットレイヤー$E-1$のクエリベクトルを格納し、これにより検証プロセスはレイヤー$0$から$E-1$の計算をスキップできます。

ベンチマークとパフォーマンスのトレードオフ

A100 GPUで行われたベンチマークによると、Llama 3.2 1B、Llama 3 8B、Llama 2 13B、Llama 2 7Bを含むほとんどのモデルサイズにおいて、早期エグジット自己スペキュラティブデコーディングは従来の2モデルスペキュラティブデコーディングよりも高速です。

アーリーイグジットの『スイートスポット』

エグジットレイヤーとパフォーマンスの間には非線形の関係があります。エグジットが早すぎると精度が低く、ドラフトトークンの受容率も低くなります。逆にエグジットが遅すぎるとドラフト段階の計算オーバーヘッドが増加します。ほとんどのモデルでは、予測精度と生成オーバーヘッドのトレードオフが最適化され、毎秒のトークン数を最大化する『スイートスポット』が存在します。

モデル固有の結果

Llama 2 70B: 自己回帰デコーディングよりも大幅に高速ですが、他のモデルと比較して速度向上が限定的であることが示されました。これは、Llama 2 7B バリアント（52Bトークン）と比較してトークン数が少ない（328Mトークン）状況で継続的にプリトレーニングされたためと考えられます。
ベースラインモデル: LayerSkipレシピで訓練されていないモデルは、早い層が出力を予測するように訓練されていないため、トークンの受容率が非常に低く、速度向上が見られません。

Transformers との統合

自己スペキュラティブデコーディングは、Hugging Faceのtransformersライブラリに統合されています。モデルチェックポイントがLayerSkipレシピで訓練されている場合、generate()関数にassistant_early_exit引数を追加することで有効にできます。