ニュース

AI“メモリ爆食い”の救世主なるか、Googleの極限圧縮技術「TurboQuant」

 米グーグル(Google)の研究部門であるGoogle Researchは3月24日、大規模言語モデル(LLM)やベクトル検索エンジンのメモリ消費を大幅に削減する新たな量子化アルゴリズム「TurboQuant」を発表した。AIの応答速度を低下させるボトルネックが解消され、AIや検索基盤の劇的な効率化を実現するとうたう。

開発の背景とAIメモリの課題

 AIが情報を理解する基本単位であるベクトルは、画像の特徴や言葉の意味など複雑な情報を捉える高次元なものになるほど膨大なメモリが消費される。

 これにより、頻繁に使う情報を一時保存するキーバリュー(KV)キャッシュが逼迫し、システム全体の処理遅延が引き起こされていた。

 従来からベクトル量子化によるデータ圧縮は利用されていたが、データブロックごとに量子化定数が計算・保存されることでメモリオーバーヘッドが1ビットから2ビット発生し、圧縮効果が部分的に相殺される課題が存在していた。

メモリの無駄をゼロにする2つの新圧縮技術

 TurboQuantでは、「PolarQuant」と「QJL」という2つの仕組みをあわせ、メモリの無駄(オーバーヘッド)が排除される。

 PolarQuantでは、データの位置を示す方法が大きく変更される。例えば「東に3ブロック、北に4ブロック」という縦横の指示が、「斜め37度の方向に5ブロック」という角度と距離の指示に変換されるような処理が実行される。

 データが「強さ」と「方向」に整理され、規則的で予測しやすい「円形」のマス目に当てはめられるため、負荷の高いデータ正規化(調整計算)のステップが省かれ、余分なメモリ消費が抑えられる。

 さらにQJL技術により、最初の圧縮時に生じたわずかなズレが補正される。この残存した誤差データが「+1」か「-1」という最小単位のシンプルな符号に変換される。

 データ同士の重要な関係性は保たれたまま、メモリを消費しない「速記」のように機能し、AIがどの情報を重視すべきか正確に判断するためのエラーチェッカーとして活用されるという。

最大8倍の高速化、Geminiへの応用も

 各種ベンチマークテストにおいて、AIモデルの精度が損なわれることなく、KVメモリのサイズが6分の1以下へと縮小された。

LongBenchベンチマークにおける、各手法のKVキャッシュ圧縮時の精度(スコア)を示したグラフ。カッコ内は消費するメモリのビット幅。TurboQuantは、わずか2.5~3.5ビットの少ないメモリ消費量でありながら、16.0ビットのフルキャッシュ(非圧縮)とほぼ同等の高いスコアを維持している

 事前のトレーニングや微調整が実行されることなく、KVキャッシュが3ビットまで量子化される。H100 GPU環境での検証では、32ビットの非量子化キーと比較して、4ビットのTurboQuantにより最大8倍の処理高速化が確認されている。

アテンションロジット計算における、最適化されたベースラインに対する処理の高速化(Speedup)を示したグラフ。入力されるデータ量(シーケンス長)が増加しても大幅なパフォーマンスの向上が確認できる

 本技術は、「Gemini」のような大規模モデルにおけるKVキャッシュのボトルネック解決の大きな鍵になるとされる。

 さらに、意味ベースのベクトル検索においても最小限のメモリとほぼゼロの処理時間で大規模なインデックス構築が可能となり、Googleのセマンティック検索のさらなる高速化と効率化が期待される。

高次元ベクトル検索において、TurboQuantが他の最先端手法(PQやRabitQ)と比較して優れた再現率(Recall)を達成している