出力層射影を HNSW で近似する LLM 推論高速化手法──Gemma 3 270M でデコードスループットを 82% 向上

論文概要

大規模言語モデル (LLM) の自己回帰的デコード処理において、大規模な語彙を持つ小型モデルの出力埋め込み行列が深刻なメモリ帯域ボトルネックを引き起こしている。本研究では、出力層の射影および top-k トークン選択の処理を、トークン埋め込みに対する最大内積探索 (Maximum Inner Product Search) として再定式化し、密な語彙射影を HNSW (Hierarchical Navigable Small World) ベースのベクトルインデックスに置き換える手法を提案する。

Gemma 3、Llama 3.2、Qwen 3 モデルを用いた CPU 推論実験において、本手法は出力射影を大幅に高速化し、Gemma 3 270M のバッチサイズ 1 におけるエンドツーエンドのデコードスループットを最大 82% 向上させた。同時に、AlpacaEval 評価において生成品質が維持されることを実証している。

関連研究

従来の LLM 推論最適化においては、KV キャッシュの圧縮や量子化、あるいはモデルのプルーニングや蒸留が中心的なアプローチであった。しかし、特にパラメータ数が少ない小型モデルや多言語対応で語彙サイズが大きいモデルにおいて、最後の出力層における高次元な全語彙ロジット計算とメモリ帯域の圧迫は十分に解決されていなかった。本研究は、近傍探索アルゴリズムを推論の最終段に適用することで、従来のエグザクトな行列積演算の代替を試みる点が異なる。

新規性と貢献

本研究の主要な貢献は、出力層の全語彙射影を近似最近傍探索に置き換えることで、メモリ帯域ボトルネックを根本から迂回した点にある。密な行列乗算を HNSW インデックスによる検索に置き換え、取得した少数の高スコアトークンのロジットをスパースなフル語彙テンソルに散布 (scatter) する仕組みを導入した。これにより、遅延に敏感な小バッチ・CPU 推論環境において、精度劣化を抑えながら実用的な高速化を実現した。

提案手法の詳細

提案手法の中核は、出力トークン選択を近似的な最大内積探索として解くことにある。従来のデコードでは全ての語彙に対する内積を計算してソフトマックス関数に入力するため、埋め込み行列のメモリ読み込みがボトルネックとなる。本手法では HNSW インデックスを用いることで、すべての語彙を総当たりで計算する代わりに、上位の候補集合のみを効率的に取得する。取得された少数の候補はスパースなテンソルに展開され、既存のデコードパイプラインにそのまま統合可能となっている。

評価・考察

論文では Gemma 3、Llama 3.2、Qwen 3 モデルを用いた CPU 推論による検証が行われている。特に Gemma 3 270M モデルのバッチサイズ 1 の条件では、エンドツーエンドのデコードスループットが最大 82% 向上した。また、AlpacaEval による評価では、近似探索を用いた場合でも生成品質の劣化が無視できる範囲に収まることが示されている。ef = 200 の設定におけるプロファイリング結果からも、出力射影処理のレイテンシが大幅に削減されていることが確認できる。

応用例と今後の展望

本手法は、エッジデバイスや CPU 環境などメモリ帯域が限られた環境で稼働するレイテンシ敏感な小型 LLM アプリケーションにおいて直接的な実務インパクトを持つ。例えば、組み込み機器やオンデバイス AI を開発するテック企業やコンシューマー向けソフトウェア開発において、スループットの改善とコスト削減に寄与する。今後は、GPU 環境や大規模モデルへのスケーラビリティ、およびインデックス構築オーバーヘッドの最適化が課題となる。

結論

本研究は、LLM の出力層における全語彙射影を HNSW ベースのベクトルインデックスに置き換えることで、メモリ帯域ボトルネックを解消し推論を加速する手法を提案した。Gemma 3 270M で最大 82% のスループット向上を達成し、低レイテンシな小バッチ推論における実用的な代替手段であることを示した。

注釈

• HNSW (Hierarchical Navigable Small World): グラフ構造を用いて高速な近似近傍探索を実現するアルゴリズム。
• AlpacaEval: LLM の指示追従性能や対話品質を比較・評価するための自動評価ベンチマーク。