ベクトル検索でLLMの出力層を高速化する新しい推論手法

導入

大規模言語モデルは、トークンを1つ生成するたびに、隠れ状態を語彙全体のスコアへ変換する。通常は出力埋め込み行列との密な行列積を実行し、その後でTop-k選択やサンプリングを行う。語彙が大きい多言語モデルでは、この処理は計算量だけでなく、巨大な行列を繰り返し読み出すメモリ帯域の面でも負担になる。特にCPU上の小型モデルやバッチサイズ1の推論では、演算よりデータ移動がレイテンシーを左右する場合がある。

出力射影をベクトル検索として扱う

この研究は、出力射影とTop-k選択の組み合わせを最大内積検索として定式化する。隠れ状態と各トークン埋め込みの内積を語彙全体について計算する代わりに、トークン埋め込みをベクトル集合とみなし、HNSW（Hierarchical Navigable Small World）インデックスから内積の大きい候補を近似的に探す。

仕組みの要点は次の通りだ。

• トークン埋め込みを索引化：出力埋め込みをHNSW構造に登録し、近似ベクトル検索を可能にする。
• 候補を絞り込む：各生成ステップで高スコアが見込まれる少数のトークンだけを取得し、全語彙へのアクセスを避ける。
• 既存のデコーダーと接続：検索した候補のlogitを疎な全語彙テンソルへ散在させることで、既存のサンプリングやTop-k処理に組み込みやすくする。
• 近似と速度を交換する：全候補の厳密な順位を保証するのではなく、生成に必要な有力候補を残すことを目指す。

評価結果と注意点

著者らはGemma 3、Llama 3.2、Qwen 3を使い、CPU推論で評価した。ベクトルインデックスを用いる出力ヘッドは、出力射影の処理を大きく短縮した。エンドツーエンドのバッチサイズ1デコードでは、Gemma 3 270Mで最大82%のスループット向上が報告されている。またAlpacaEvalによる評価では、生成品質を維持できたとしている。

ただし、この結果から密な出力層があらゆる環境で不要になるとは言えない。効果が期待されるのは、語彙が大きく、バッチが小さく、出力射影がレイテンシーの大きな部分を占める条件だ。大バッチでは密な行列演算の並列化が有利になる可能性があるほか、近似検索にはインデックス構築、メモリ配置、候補の再現率、疎なlogitの散在処理といったコストもある。提示された情報だけでは、これらの条件の完全な比較までは分からないため、実運用には個別の検証が必要だ。

意義

この研究は、LLMの出力ヘッドを固定的な行列積だけでなく、トークンベクトルの検索コンポーネントとして設計できることを示している。今後、インデックスのオーバーヘッドや品質・遅延の関係がさらに検証されれば、CPU推論、エッジ環境、低同時実行数のサービスで有用な選択肢になり得る。

メモリアクセスがデコードの制約になる場合、全語彙行列の演算を細かく最適化するより、そもそも読み出す語彙の範囲を減らす方が直接的な改善につながる可能性がある。現時点では、万能な置き換えではなく、レイテンシー重視の推論向け最適化として捉えるのが適切だ。

出典：arXiv