原題:LoopVQ: Sub 1-bit KV Cache Quantization for Looped Transformers
ループ型Transformerは同じ層群を繰り返し呼び出すことで、パラメータ数を抑えつつ数倍大きい通常のモデルに匹敵する性能を発揮しますが、各ループステップごとに独自のKVキャッシュを保存する必要があり、キャッシュサイズがループステップ数に比例して線形に増加するため、これが主要なメモリボトルネックとなっています。著者らは、KVがループ軸上で高い冗長性を持つ一方で、ループステップによって量子化に対する感度が大きく異なることを観察し、LoopVQを提案しました。これは、連続する複数のループステップのチャネルを1つのサブベクトルとして統合し、単一のコードブックインデックスを共有する(ループ軸方向のベクトル量子化)手法であり、さらにFisher重み付きk-meansによるコードブックの微調整を行うことで、量子化後のモデル出力が全精度モデルに近づくようにします(学習時には各ループステップで量子化されたキャッシュを読み込みます)。この手法は通常のLLM向けの方法と直交しており、TaSQのキー変換と併用可能です。
いいね 8 · コードスター数 3
理由
- 著者によると、要素あたり0.5ビットの場合、OuroおよびHuginnモデルにおける汎用タスク、長い思考チェーン推論、長文脈検索タスクにおいて、既存の通常のLLM向けKVベクトル量子化手法を上回ります。TaSQを併用した場合でも、0.25ビットでBF16精度の大部分を維持します。
- システム上の利点が実証されています:著者がSGLangに実装した結果、32Kコンテキストにおいてデコード遅延がほぼ半減し、Ouro 1.4B / 2.6Bのピークサービシングスループットは、BF16ベースラインに対してそれぞれ4.3倍および6.1倍向上しました。
- コードはApache-2.0ライセンスです:キャリブレーションフロー、パッキングされたインデックスストレージと融合デコードカーネルを含むSGLangのフォーク、および論文内のすべての表を再現するスクリプトがリポジトリに含まれています。
使い方
- ./setup_envs.shを実行して3つの環境(loopvq-serveエンジン、loopvq-eval評価用、loopvq-buildキャリブレーション用)を作成し、export MODELS_DIRでOuro-1.4B、Ouro-2.6B、Ouro-*-Thinking、およびhuginn-0125を指します。汎用タスクの表を再現するには、MODEL=ouro_1p4b GPUS="0 1 2 3" bash experiments/loopvq_general/build.shを実行し、続いてMODEL=ouro_1p4b bash experiments/loopvq_general/run_all.shを実行、最後にpython3 experiments/loopvq_general/make_table.pyを実行します。
- サービス起動:loopvq-serve環境内でPREFIX_TOKENS=8 RECENT_TOKENS=64 PREFILL_BACKEND=triton ./scripts/serve_method.sh $MODELS_DIR/Ouro-1.4B vq <コードブックディレクトリ> --port 30000 --disable-piecewise-cuda-graphを実行すると、SGLangのOpenAI互換インターフェースが得られます。先頭および直近のトークンはbf16のまま保持されます。
リスク
- ループ型Transformerにのみ適用可能です(論文ではOuro 1.4B / 2.6BおよびHuginn-3.5Bを使用)。通常のLLMには直接適用できません。キャリブレーションにはGPQAの16件のシーケンスのみが使用されました。
- 遅延およびスループットの数値は著者の環境によるものです。いいね数8、コードスター数3(注目度を示すに過ぎません)。