🤖 machine learning

NSNQuant: A Double Normalization Approach for Calibration-Free Low-Bit Vector Quantization of KV Cache

NSNQuantは、独自の「正規化・シフト・正規化」変換とアダマール変換を組み合わせることで、トークンの分布を標準正規分布に適合させ、キャリブレーション用データセットに依存することなく、ロバストな低ビット圧縮と最大3倍のスループット向上を実現する、LLMのKVキャッシュのためのキャリブレーションフリーなベクトル量子化手法である。

原著者： Donghyun Son, Euntae Choi, Sungjoo Yoo

公開日 2026-07-16

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Donghyun Son, Euntae Choi, Sungjoo Yoo

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたは、非常に長い道のりを歩きながら、バックパックの中に膨大な量の思い出のライブラリを運ぼうとしているところだと想像してみてください。これは、物語を書き、数学を解き、あなたとチャットをする超スマートなコンピューターの脳である大規模言語モデル（LLM）が、長い会話を処理しようとする時に起こることです。単語を一つ読むたびに、文脈を理解するために、それより前に来たすべてのことを覚えておかなければなりません。この「記憶」は**KVキャッシュ（KV Cache）**と呼ばれます。問題は、会話が長くなるにつれて、このバックパックがどんどん重くなり、最終的にはスペースがいっぱいになって、コンピューターの動きが極端に遅くなってしまうことです。

これを解決するために、科学者たちは、衣類を真空パックに入れるように、記憶を圧縮することでバックパックを小さくしようと試みてきました。一つの人気のある手法が、**ベクトル量子化（Vector Quantization: VQ）**です。これは、似たもの同士をグループ化し、それらを「辞書」やコードブックからの単一のラベルに置き換えることだと考えてください。すべての靴下の正確な青色の濃淡を覚える代わりに、「青グループ4」とだけ覚えるのです。しかし、落とし穴があります。既存の手法のほとんどは、その辞書を作るために、特定の衣類のセット（キャリブレーション・データセット）を事前に学習する必要があります。もし、全く異なる種類の衣類（新しいタイプの会話）をパッキングしようとした場合、その辞書は適合せず、圧縮は失敗してしまいます。この論文は、まさにその問題、つまり、事前に衣類を学習することなく、いかにして記憶を圧縮するかという課題に取り組んでいます。

問題点：一つのワードローブにしか適合しない辞書

ソウル大学の研究者である著者たちは、LLMのメモリ圧縮における現在の最先端手法である**結合量子化（Coupled Quantization: CQ）**にある、もどかしい不具合に気づきました。CQを、一人の人物のサイズに基づいてカスタムスーツを作る仕立て屋だと想像してみてください。もしその人物が、異なる体型の人々で溢れる部屋に入ってきたら、そのスーツは最初の一人には完璧にフィットしますが、他の全員にはひどい見た目になってしまいます。

AIの世界では、この「人物」とは、モデルがキャリブレーションされたデータ（WikiText-2のような特定のテキストデータセット）のことです。モデルが異なる種類のテキスト（ウェブページの大規模な集合であるC4データセットなど）を処理しようとすると、「スーツ」が適合しなくなります。著者らは、このミスマッチが、特に句読点の扱いにおいてモデルに愚かなミスを引き起こすことを発見しました。例えば、モデルはカンマについて混乱するかもしれません。なぜなら、学習データから学んだ「辞書」には、新しいテキストにおけるカンマの現れ方に適したラベルが存在しなかったからです。これは大きな問題です。なぜなら、モデルが自分のコンフォートゾーンの外に出たとき、信頼性が損なわれることを意味するからです。

解決策：NSNQuant – ユニバーサルなパッキング・キューブ

これを解決するために、チームは、特定のデータを事前に学習する必要のない、メモリ圧縮の巧妙な新手法であるNSNQuantを導入しました。新しいワードローブごとにカスタム辞書を作ろうとする代わりに、NSNQuantは、すべての衣類を標準的な既製品のパッキング・キューブに収まるように強制します。

彼らがどのように行っているかを、**Normalize-Shift-Normalize（NSN）**と呼ばれる3ステップの「マジック・トリック」を使って説明します。

Normalize（第1ステップ）： 小さすぎるものもあれば巨大なものもある、靴下の山を想像してください。最初のステップは、すべての靴下のサイズを同じ大きさに引き伸ばしたり縮めたりすることです。これにより、巨大な靴下がスペースを独占してパッキングを台無しにするのを防ぎます。
Shift（中間ステップ）： 次に、靴下はすべて同じサイズだが、すべて左に傾いている状態を想像してください。「Shift」のステップは、それらをすべて中央に戻し、完璧にバランスが取れた状態にします。
Normalize（最終ステップ）： 念のため、すべてが依然として均一であることを確認するために、最後にもう一度サイズをチェックします。

この3ステップのダンスの後、著者らは最後に**アダマール変換（Hadamard Transform）**というひねりを加えます。これは、靴下の山全体を特定の数学的な方法で回転させることだと考えてください。魔法のような点は、この回転を行った後、元のテキストが何であったかにかかわらず、靴下の乱雑で予測不可能な形が、突然、完璧で滑らかなベルカーブ（標準正規分布）のように見えるようになることです。

データがこの予測可能なベルカーブのようになるため、著者らは、そのベルカーブに合わせて特別に設計された単一の既製「辞書（コードブック）」を使用することができます。データを事前に見る必要はありません。データをそのような形に強制したため、データが辞書に適合することを分かっているのです。

彼らが発見したもの：ユニバーサルな鍵

チームは、LLaMAやMistralファミリーを含む、いくつかの有名なAIモデルでこのアイデアをテストしました。彼らは、単純な物語から複雑なコードや数学の問題まで、さまざまな種類のテキストを用いて、NSNQuantを古い手法（CQやKIVIなど）と比較しました。

結果は素晴らしいものでした。

優れた汎用性： 古い手法（CQ）は、データセットを切り替えるとつまずきましたが、NSNQuantは強力なパフォーマンスを維持しました。それは、あらゆるドアを開けることができるユニバーサルな鍵を持っているようなものであり、古い鍵は、それが作られたドアにしか機能しないものでした。
低ビットでの成功： チームは、メモリをわずか1ビットおよび2ビットに圧縮するテストを行いました。1ビットの設定（メモリを絶対的な最小サイズまで縮小する場合）において、NSNQuantは競合を圧倒しました。例えば、GSM8Kと呼ばれる数学的推論タスクにおいて、古い1ビットの手法は約24のスコアでしたが、NSNQuant-1bは53.45を記録しました。これは2倍以上のパフォーマンスです！
速度とスペース： メモリが非常に小さいため、コンピュータはより多くの会話を同時に処理できます。著者らは、彼らの手法が、標準的な非圧縮バージョンよりも3倍多いデータスループットを処理でき、かつ大幅に少ないメモリを使用できることを示しました。

細部についての注釈

著者らは、この手法は大きな改善ではあるものの、完璧な魔法ではないことも注意深く述べています。AIモデルの非常に最初のレイヤーにおいて、靴下がベルカーブに完璧に適合しない「外れ値」が時折残ることがあると彼らは発見しました。しかし、これらの小さな不具合があっても、全体的なパフォーマンスは非常に高い水準を維持していました。

また、彼らはこの手法が「キャリブレーション・フリー（校正不要）」であることを強調しています。辞書を構築するために特定のデータを数時間学習する必要があった古い手法とは異なり、NSNQuantの辞書は、単一のグラフィックスカード上で5分未満で構築でき、その後あらゆるモデルに再利用可能です。これは、実世界での使用において非常に実用的です。

要約すると、NSNQuantは、あらゆる乱雑な記憶の山を、整然とした予測可能な形へと強制するユニバーサルなパッキング・システムであり、これによりAIモデルが、思考能力を失うことなく、より小さなバックパックに長期的な記憶を運べるようにするものです。データを圧縮する「前」に標準化することで、AIが全く新しい、あるいは未知のトピックを扱っている場合でも、より高速で、安価で、信頼性の高いものにできることを示唆しています。

技術要約: NSNQuant

問題提起

大規模言語モデル（LLM）の推論は、特に長いシーケンスや大きなバッチサイズを扱う際に、ますますメモリ制限（memory-bound）の影響を受けるようになっています。主要なボトルネックは、シーケンス長に対して線形に増加するKey-Value（KV）キャッシュです。ベクトル量子化（VQ）は有望な圧縮技術として浮上していますが、Coupled Quantization（CQ）のような既存の手法は、セントロイド（コードブック）を学習するためにキャリブレーション・データセットに依存しています。著者らは、これらのキャリブレーションに基づく手法が、分布外（OOD）シナリオにおいて深刻な性能低下を引き起こすことを指摘しています。具体的には、入力分布がキャリブレーション・セットと異なる場合（例：WikiText-2でキャリブレーションを行い、C4でテストする場合）、学習されたセントロイドが句読点などの重要なトークンを正確に表現できず、アテンション重みの重大な誤差やモデル性能の低下（例：パープレキシティの上昇）を招きます。

手法: NSNQuant

キャリブレーション・データへの依存とそれに伴う分布シフトに対処するため、著者らはキャリブレーション・フリーのVQ技術であるNSNQuantを提案しています。その核心となる革新は、外部データを使用せずに再利用可能な単一のコードブックの使用を可能にするため、KVキャッシュの分布を標準正規分布に整列させる変換パイプラインです。

1. Normalize-Shift-Normalize (NSN)

この手法は、KVテンソル（ $v \in \mathbb{R}^{l \times d}$ ）に対して3段階の変換を適用します：

Normalize (Token-wise): 各トークンのノルムを $\sqrt{d}$ になるよう正規化します。これにより、アウトライヤー（外れ値）トークンが後続のステップを支配することを防ぎ、小さな大きさを持つトークンが相対的に大きな誤差を被らないようにします。
Shift (Channel-wise): チャネルごとの平均を計算して減算し、分布をゼロを中心に中心化します。
Normalize (Token-wise): トークンのノルムを $\sqrt{d}$ に戻すために、2度目の正規化を適用します。

このプロセスにより、元のテンソルを再構成するために使用される副産物（ $s_1, o, s_2$ ）が生成されます。理論的には、アダマール変換と組み合わせることで、このプロセスはチャネル分布を標準正規分布に整列させます。これは、Lemma 1で導出された分散境界および中心極限定理によって裏付けられています。

2. Hadamard Transform

NSNプロセスは、アダマール変換（具体的にはRandomized Hadamard TransformまたはRHT）と組み合わされます。最終的な実装では効率性のためにナイーブなアダマール変換が使用されますが、RHTは結果としてのチャネル分散が1に近くなるという理論的保証を提供します。

3. Vector Quantization and Codebook

変換後、データは標準正規分布に特化した単一のグローバル・コードブックを用いて量子化されます。

Codebook Construction: キャリブレーション・セットからセントロイドを学習する代わりに、コードブックは合成された標準正規データ（torch.randn）を用いたK-Meansによって初期化され、元のベクトルと量子化されたベクトルの間のコサイン距離を最小化するように微調整されます。このプロセスは高速（数分）であり、モデルの重みに対するバックプロパゲーションを必要としません。
Scale Adjustment: 量子化されたベクトルが元のベクトルに平行な成分を保持するように、適応的なスケーリング戦略が採用されており、直交誤差を最小限に抑えます。
Double Quantization (DQ): オーバーヘッドをさらに削減するため、NSNプロセスで使用されるパラメータ（平均 $o$ およびスケール $s_1$ ）およびコードブックのエントリ自体が量子化（4ビット）され、無視できる程度のメモリ・オーバーヘッド（平均約0.23ビット）しか追加しません。

4. 実装の詳細

Residuals: KIVIに従い、KVキャッシュは量子化された部分とフル精度の残差（residual）に分割されます。残差が満たされると、それはフラッシュ（排出）され、量子化されて追加されます。
Attention Computation: Keyについては、NSNはRotary Position Embedding (RoPE) の前に適用されます。著者らは、この順序がより高い品質をもたらすことを見出しています。アテンション・スコアの計算は、NSNの副産物（ $s_1, s_2, o$ ）をカスタムCUDAカーネル内で直接組み込むように修正されています。
Efficiency: アダマール変換、ベクトル量子化、およびデクオンタイゼーションのために、DRAMアクセスを最小限に抑える共有メモリを活用したカスタムCUDAカーネルが実装されています。

主な貢献

経験的観察: 著者らは、KV分布の入力データへの強い依存性が、既存のVQ手法（CQなど）をOODシナリオで失敗させる原因であることを実証しました。これは特にLLaMA3における句読点トークンに影響を与えます。
Calibration-Free VQ: NSNQuantは、NSNプロセスとHadamard変換を通じて、KV分布を既知の事前分布（標準正規分布）に整列させることで、キャリブレーション・データセットを不要にする手法を導入しました。
性能と汎用性: 広範な実験により、NSNQuantは1ビットおよび2ビットの設定の両方において、多様なデータセット（WikiText-2, C4, LongBench, GSM8K, MMLU）にわたってベースライン（KIVI, KVQuant, CQ）を一貫して上回ることが示されました。また、キャリブレーション手法が劣化する場面でも優れた汎用性を示します。
効率性: 実装は、フル精度ベースラインに対して最大3倍のスループット向上を実現しており、メモリ使用量の減少により、より大きなバッチサイズ（最大4倍）をサポートしています。

結果

Perplexity (PPL): C4データセットにおいて、NSNQuantはCQを大幅に上回ります。例えば、LLaMA3-8Bにおいて、CQ-4c9bのPPLは13.97ですが、NSNQuant-2bは9.08です。1ビット設定では、NSNQuant-1b（PPL 16.69）はCQ-8c10b（PPL 43.78）を圧倒します。
LongBench: NSNQuant-1bは、他の1ビット・ベースラインを大幅に上回ります。2ビット設定では、他の手法と同等か、わずかに上回る結果となりますが、著者らは特定のタスクにおけるメトリックのノイズについても言及しています。
Reasoning and Generation: GSM8K（数学的推論）およびMMLU（マルチタスク理解）において、NSNQuantは特に厳格な推論パスが要求されるfew-shot Chain-of-Thought (CoT) 設定で優れています。
Latency: プリフィル（prefill）段階ではNSNとHadamardのオーバーヘッドによりレイテンシが増加しますが、デコーディング段階は、メモリ制限のあるアテンション計算が緩和されるため、フル精度ベースラインよりも高速になります。

重要性と主張

本論文は、NSNQuantがデータ依存型のキャリブレーションから分布整列型の量子化への転換を象徴していると主張しています。Hadamard変換とNSNを組み合わせることで、KV分布が標準正規分布を生成するという特性を利用することで、モデルに依存しないユニバーサルなコードブックの使用を可能にします。これにより、この手法は分布シフトに対して堅牢になり、キャリブレーションの計算コストとデータ要件を排除します。著者らは、NSNQuantを、高圧縮（1ビット/2ビット）、強力な汎用性、および効率的な実行性能のバランスを提供し、メモリ制約のあるハードウェア上で長文脈LLM推論を可能にするための実用的なソリューションとして位置付けています。

著者らは、標準化が分散の高いアウトライヤー・チャネルが存在する初期レイヤーでは完全ではないこと、および本手法はCQとは異なりチャネル間の相関を明示的にモデル化していないことを認め、謙虚に限界を述べています。しかし、彼らは、特定のキャリブレーション分布への過学習から得られる潜在的な利点よりも、堅牢性と汎用性のトレードオフの方が価値があると主張しています。