arXiv論文:KVキャッシュのランダム削除で推論最大43%高速化

a black and white photo of a wall AI

Salesforce AI Researchの研究チームが、LLM(Large Language Model、大規模言語モデル)の推論速度を高める新手法「Random Attention」を発表しました。長い思考過程を経て答えを出すAIモデルでは、途中経過を保持するメモリ(KVキャッシュ)が肥大化し処理が重くなる問題があります。研究チームは高度な選別ロジックを使わず、トークンをランダムに間引くだけで既存の最良手法と同等の精度を保ちながら、vLLM(AIモデルを効率良く動かすための代表的な配信基盤)上で32〜43%のスループット向上を確認しました。

a black and white photo of a wall

背景と文脈

近年のAIモデルは、答えを出す前に「chain-of-thought(思考の連鎖、途中の推論過程を段階的に文章化する手法)」で長い推論文を生成してから最終的な回答にたどり着く方式が主流になっています。数学の問題を解く前に途中式を書き出すようなイメージです。この方式は正答率を大きく引き上げる一方、生成する文章が長くなるほど計算コストも膨らむという副作用を抱えています。

LLMは一語ずつ単語を生成する際、それまでの文脈をすべて計算し直すのではなく、KVキャッシュと呼ばれるメモリに過去の計算結果を保存して使い回しています。推論が長引くほどこのキャッシュは膨張し、GPUメモリを圧迫してサーバーの処理能力を頭打ちにします。これを防ぐため、重要度の低いトークンをキャッシュから追い出す「エビクション(eviction、退避・削除の意)」という技術が研究されてきました。従来手法の多くは、どのトークンを残すべきかをアテンションスコア(注目度を示す数値)などから複雑に計算して判定しており、その計算自体にもコストがかかっていました。

技術/ビジネス面

green and black circuit board
Photo by Akshar Dave🌻 on Unsplash

Random Attention: Rethinking KV Cache Eviction for Efficient Reasoningと題されたこの論文の要点は、「複雑な選別は不要」という逆転の発想です。プロンプト(最初にモデルへ渡す指示文)の部分だけは削除せずに守り、それ以降の推論トレースについては各アテンションヘッド(注意機構の処理単位)ごとにランダムにトークンを間引きます。研究チームによると、推論トレースには2種類の冗長性があるため、この乱暴に見える手法でも情報が失われにくいといいます。ひとつはモデル自身が必要な情報を文章中で繰り返し書き直す「テキスト面の冗長性」、もうひとつは複数のアテンションヘッドがそれぞれ独立に同じ情報を保持する「構造面の冗長性」です。

実験では、プロンプト部分さえ保護しておけば、それ以降はランダムに間引いても既存の最も強力な選別手法と同等の精度を維持できると確認されました。その上でvLLM環境での実測では、スコア計算が不要になった分だけ処理が軽くなり、スループットが32〜43%向上しています。コードはGitHub上で公開されており、既存のvLLM環境に組み込んで検証しやすい形になっています。

これからどうなるか

この結果は、LLMを自前でホスティングしている開発者やスタートアップにとって直接的なコスト削減につながる可能性があります。高度なキャッシュ管理ロジックを自前で実装・チューニングする手間が減り、既存の推論基盤に軽い変更を加えるだけで処理能力を底上げできるからです。特に、長い思考過程を生成するタイプのモデルをAPI経由で大量に呼び出すプロダクトでは、GPUの台数を増やさずにリクエストのさばける量を増やせる可能性があります。

今後は、この手法が商用の推論サービスにどこまで採用されるかが焦点になりそうです。ランダム性に依存する手法だけに、精度が不安定になるケースがないか、より幅広いタスクでの追試が期待されます。

まとめ

Salesforce AI Researchは、KVキャッシュのトークンをランダムに間引くだけで既存の複雑な選別手法と同等の精度を保ちつつ、vLLM上で処理速度を32〜43%高められる新手法「Random Attention」を発表しました。実装のシンプルさから、LLMを自前で運用する開発者にとって実用性の高い選択肢になりそうです。

参考リンク

アイキャッチ画像: Photo by Google DeepMind on Unsplash

タイトルとURLをコピーしました