facebookresearch/LayerSkip

解決する課題

LayerSkipは、大規模言語モデル（LLM）の推論における計算コストと遅延を軽減します。生成のドラフト段階でモデルが自身の深いレイヤーのスタックから「早期に終了」できるようにすることで、各トークン生成に必要な計算量を削減し、より高速なトークン生成を実現します。

動作方法

LayerSkipは、単一のモデルがドラフトモデルと検証モデルの両方として機能できるようにするトレーニング手法を採用しています。自己推測的デコードを実装しており、指定されたレイヤー（exit_layer）でモデルが終了し、迅速にドラフトトークンを生成します。その後、これらのトークンはフルモデルで一回のパスで検証され、従来の自己回帰的デコードと比べて大幅に高速化されます。

対象ユーザー

Llamaベースのモデルの推論速度を正確性を損なわず高速化したいAI研究者や開発者向けです。特に要約やコード生成などの生成タスクに適しています。

主な特徴

• 自己推測的デコード: ドラフトと検証に同じモデルを使用するため、別途小さなドラフトモデルを用意する必要がありません。
• 早期終了推論: 速度と精度のバランスを取るために、終了するレイヤーをユーザーが指定できます。
• 広範なモデル対応: Llama 2、Llama 3、Llama 3.2、CodeLlamaのさまざまなモデルに対応する事前学習済みチェックポイントを提供しています。
• 統合: Hugging Face transformers および PyTorch torchtune に統合されています。
• 評価ツール: ベンチマーク、正しさの検証、ハイパーパラメータスイープ用のスクリプトを提供し、最適な終了レイヤーや推測回数を特定できます。