はじめに
Speculative Decoding と量子化を組み合わせた手法である QSPEC について、個人的な勉強でその要点となる内容をまとめました。
QSPEC は、低精度で高速な W4A4 をドラフト生成に、高精度な W4A16 を検証に利用する Speculative Decoding の手法です。両者で量子化済み Weight と KV キャッシュを共有することで、メモリ使用量を抑えつつ、W4A16 相当の精度と、推論の効率化を両立したと報告されています。
Juntao Zhao, Wenhao Lu, Sheng Wang, Lingpeng Kong, and Chuan Wu, "QSPEC: Speculative Decoding with Complementary Quantization Schemes", 2025.
https://arxiv.org/abs/2410.11305
基本技術
論文の背景知識となる技術用語
Weight-only 量子化
モデルの重みだけを低精度化する方式。例えば、W4A16 では、モデルの重みを4ビットに量子化して保存し、活性化 (Activation) には、FP16 などの16ビット形式を使用する。
Weight-Activation 量子化
重みと活性化の両方を同時に量子化する方式。高精度への逆量子化を行うことなく、低精度カーネルで演算する。低精度演算に対応したハードウェアを活用して高速に推論することができる。
W4A16 と W4A4 の違い
W4A16 は、モデルの重みを INT4 などで保存し、活性化は FP16 や BF16 で扱う。実際の実装では、4ビットの重みを GPU から読み出し、計算時には FP16 などへ逆量子化して行列積を行うことがある。
W4A4 は、重みと活性化の両方を4ビットに量子化する。
行列積の入力は4ビットだが、積和演算の結果は値が大きくなるため、通常は INT32 などの高い精度で累積する。その後、スケールを適用して FP16 や INT4 などへ変換する。
重みの逆量子化
重みの逆量子化 (dequantization) とは、INT4 などで小さく保存された重みを、行列積で扱える FP16/BF16 などの値へ戻す処理である。
例: W4A16 での逆量子化
W4A16 では、重みは4ビット、活性化は FP16/BF16 となる。
Y = X_{FP16} \hat W_{FP16}
ここで、
\hat W_{FP16} = s \big(W_{INT4} - z\big)
として、4ビット重みを FP16 などへ逆量子化してから行列積を行う。
低精度カーネルによる高速化
INT4, INT8, FP8 などのデータ型に合わせて実行の効率化を図ったカーネル。具体的には以下のような理由で高速化する。
- メモリから読み取るデータを減らす
FP16 の重みを W4 にすると、理論上の重みサイズは 1/4 となり、メモリから転送するデータ量を削減することができる - キャッシュやレジスタを有効活用する
データを小さくすることで、キャッシュや共有メモリへ、より多くの要素を置くことができる - 専用演算器のスループットを高められる
GPU や NPU には、低精度の行列積を高速に処理する専用演算器がある。一般に低精度になるほど1命令で扱うことのできる要素数は増やすことができる
Decode がメモリ帯域律速になりやすい理由
最大の理由は、1トークンを生成するために、巨大な重みと KV キャッシュを読み込む必要がある一方、小バッチではそれらを十分に再利用できないため。
Prefill フェーズでは、プロンプト中の多数のトークンを並列計算しトークンを並列計算することができるが、 Decode フェーズでは、原則として1ステップずつ逐次生成する必要がある。
| 段階 | 入力 | 特徴 |
|---|---|---|
| Prefill | プロンプト中の多数のトークン | 多数のトークンを並列計算できる |
| Decode | 新しい1トークン | 原則として1ステップずつ逐次生成する |
Draft モデルと Target モデル
Speculative Decoding は、小さく高速な Draft モデルが複数の候補トークンを先に生成し、大きな Target モデルがそれらをまとめて検証することで、出力品質を維持しながら Decode を高速化する手法である。
Speculative Decoding の Accept/Reject の流れ
それまでの文脈 $ h $ に対して、Draft モデルの確率分布を $ q(x|h) $ 、Target モデルの確率分布を $ p(x|h) $ とした場合、Draft が候補 $ x $ を生成したとき、基本的な Speculative Sampling では次の確率で受理する。
a(x) = min \big( 1, \frac{p(x|h)}{q(x|h)}\big )
常に最大確率のトークンを選ぶ Greedy Decoding では、Target モデルの最大確率トークンが Draft 候補と一致すれば受理される。
受容率と高速化率の関係
受容率が高いほど、Target モデル1回の検証で確定できるトークン数が増えるため、基本的に高速化率も高くなる。
ただし、Speculative Decoding では、最初の棄却が起きた時点で、それより後ろの候補を採用できないため、先頭からのトークン受理率が重要である。
1 ラウンドで受理される Draft トークン数 A の期待値は以下で表される。
E[A] = \alpha + \alpha^2 + ... + \alpha^K = \frac{\alpha(1-\alpha^K)}{1-\alpha}
標準的な Speculative Sampling では、Draft 候補が途中で棄却された場合、Target モデルの分布から補正トークンを1個生成する。また、すべての候補が受理された場合にも、Target が計算した追加の1トークンを使用できる。したがって、1ラウンドで確定するトークン数 $ L $ は、
L = A + 1
と考えられ、その期待値は以下で表される。
E[L] = 1 + \alpha + \alpha^2 + ... + \alpha^K = \frac{1-\alpha^{K+1}}{1-\alpha}
通常の Target モデルが1トークンを生成する時間を $ T_{target} $ とする。
Speculative Decoding の1ラウンドに必要な時間を以下とする。
T_{round} = T_{draft} + T_{verify} + T_{overhead}
この場合の高速化率は以下で表される。
S \approx \frac{E[L]T_{target}}{T_{draft}+T_{verify}+T_{overhead}}
Motivation
ここでは、論文に記載されている次の2つの観察について記載します。
- W4A4 は、多段階推論タスクにおいて、最終的な性能が大きく低下する
- 一方で、トークン単位で見ると、W4A4 と W4A16 の予想はよく似ている
W4A4 の活性化量子化は多段階推論に影響しやすい
Atom や QuaRot などの Weight-Activation 量子化手法は、一般的なベンチマークでは、Weight-only 量子化に近い性能を保ちながら、推論を高速化できると報告されています。
しかし、QSpec の論文著者らは、従来の評価で使用されてきたタスクだけでは、活性化量子化による性能低下を十分に捉えることはできないと述べています。また、それを示すベンチマークとして、Llama-3-8B-instruct に Atom ベースの量子化を適用し、次のタスクで W16A16, W4A16, W4A4 を比較しています。
- WikiText-2 : 文章に対する予測問題
- PIQA : 物理的な常識推論
- GSM8K : 複数のステップを必要とする算数問題
- MBPP : Python プログラムの生成問題
結果、WikiText-2 や PIQA では、W4A4 の性能低下は比較的小さいですが、GSM8K や MBPP では、W4A4 の性能が平均で約 30% 低下するのに対し、W4A16 は約 4% の低下にとどまっています。
この結果から、W4A4 の活性化量子化は、GSM8K や MBPP など多段階推論タスクに影響するとされています。
W4A4 と W4A16 のトークン予測はよく似ている
論文の著者らによると、トークン単位の予想では、W4A4 と W4A16 の Top-1 予想には類似性があります。この性質を確認するために、GSM8K テストセットを用いて、以下の実験をしています。
- Atom ベースの A4W16 を用いた貪欲サンプリングを行い、各問題に対する回答トークンを生成する
- Atom ベースの W4A16 が選択した Top-1 トークンと、その予想確率を記録する
- 問題文と W4A16 が生成した回答を連結し、W4A4 で順伝播を1回実行する
- 各トークン位置における W4A4 と W4A16 の予想を比較する
論文の図2から、次の2点が確認できます。
- W4A4 と W4A16 のどちらも、トークン予測確率の多くが 80% を超えている
- 棄却されるトークンは、受容されるトークンと比べてわずかである
つまり、W4A4 は最終的なタスク性能こそ大きく低下するものの、生成するトークンの大部分は W4A16 と一致しています。
ほとんどのトークンが一致しているにも関わらず、W4A4 の性能は大きく低下する理由について、著者らは、少数の重要なトークンに生じる予測の違いであると述べています。自己回帰型の LLM では、それまでに生成したすべてのトークンから、次のトークンを予測するため、ある推論ステップで誤ったトークンを生成すると、その誤りを含む文脈に基づいて、後続のトークンが生成され、その後の推論全体が誤った方向へ進む可能性があるからです。
QSPEC
先の実験結果から、大部分のトークンを高速な W4A4 で生成し、少数の重要な誤りだけを W4A16 で検出・修正するという QSPEC の着想が得られます。
ここでは、QSPEC の構成について述べます。
Draft フェーズ
QSPEC では、次の $ \gamma $ 個のトークン $ \hat T_{i+1:i+\gamma} $ と関連する分布 $\hat {p}_{i+1:i+\gamma}(t) $ を予想するために、W4A4 による Weight-Activation 量子化と重みを共有する量子化スキームを採用しています。活性化の精度を低減しているため、高速なトークン生成が可能となります。
Verify フェーズ
W4A4 の性能低下を補うため、Verify フェーズでは、高精度な重みのみの量子化を用います。
具体的には、高精度量子化モデル $M_h$ が
\begin{align}
&T_{\le i} \\
&\hat{T}_{i+1:i + \gamma}
\end{align}
を連結したものを入力として受け取り、高品質な予測確率 $p_{i+1:i+\gamma+1}(t)$ を出力します。
重みの共有
QSPEC では、W4A4 と W4A16 の計算において別々の重みを保存せず、同じ4ビットの量子化済み重みを用います。つまり、W4A4 と W4A16 の違いは重みではなく、主に活性化の精度と計算カーネルです。
KV キャッシュの上書き
QSPEC では、受理されたトークンに対して、W4A4 による低精度な KV キャッシュを W4A16 によるキャッシュで置き換えます。これにより、後続のデコーディング処理において、高品質なコンテキスト情報を利用できるようにします。
単一モデル内で重みを共有し、KV キャッシュを再利用することで、QSPEC は二重のキャッシュ管理を不要にし、精度を犠牲にすることなくメモリ使用量を削減できます。
論文中の実験結果
Fidelity
論文の Table 3 では、Atom と QuaRot の2種類の量子化手法において、W16A16, W4A16, QSPEC, W4A4 の精度を比較しています。結果として、W4A4 では GSM8K, MATH, MBPP など複雑な推論タスクで精度が大きく低下していますが、QSPEC では W4A16 と同程度の精度を保つことが示されています。
QSPEC は W4A4 の結果を W4A16 の予想と比較して検証しているため、原理的に W4A16 と同等の精度が得られます。しかし、PyTorch の非決定的な演算や、数値誤差による影響で、完全に一致するわけではありません。
Throughput
論文の Table 4 では、3B, 7B, 13B の Llama モデルでバッチ数を8, 16, 32と変化させて W4A16, W4A4, QSPEC の推論速度を比較している。結果として、W4A16 に比べ、平均して1.4倍、最大1.64倍の性能向上をしています。
これは、以下の理由によるものと考えられます。
- W4A4 の低精度カーネルによる高速なドラフト
- W4A16 による複数ドラフトトークンの並列検証
Memory
論文の Table 5 では、Llama-2-7b-chat-hf における QSPEC と EAGLE の比較を示しています。EAGLEは、単一シーケンスの入力（バッチサイズ1）において最適な性能を発揮しますが、バッチサイズが増加する（8および16）につれて、効率低下します。また、EAGLE のドラフトモデルは KV ストレージが増えるため、16バッチで OOM となっています。それに対し、QSPEC は優れたスケーラビリティとメモリ効率を示しているとしています。
終わりに
本記事では、Speculative Decoding に量子化を組み合わせた QSPEC の論文を読み、その要点となる内容をまとめました。次は、論文で示されている事柄について、実際に再現検証を行ってみたいと考えています。