QSpec実行時の2つの量子化による推論の一致率を調べる

はじめに

以前、Speculative Decoding の一手法である QSpec の論文を読みました。

QSpec は、ドラフト生成に高速な W4A4、検証に高精度な W4A16 を用いる Speculative Decoding の手法です。
W4A4 は W4A16 よりも演算精度が低いものの、両者の予測結果が十分に類似していれば、W4A4 モードが提案したトークンの多くを W4A16 モードで受容でき、推論を高速化できます（W4A4 モードと W4A16 モードについては後述します）。

この前提を検証するため、論文では W4A4 と W4A16 がそれぞれ算出した最有力トークンの確率や両者の一致率、さらに W4A4 の提案したトークンが最終的に受容されるか棄却されるかを調べています。

そこで今回は、同様な傾向を確認するため、QSpec において W4A4 と W4A16 が予想した最有力トークンを比較する実験を行いました。

実験の前提知識

Speculative Decoding

Speculative Decoding （投機的デコーディング）は、LLM の出力品質を維持しながら、推論を高速化するための手法です。
通常の自己回帰型 LLM では、入力されたトークン列をもとに、後続のトークンを1つずつ順番に生成します。
一方、Speculative Decoding では、推論を2段階に分けて処理します。

1. Draft フェーズ
   小型または低精度で高速なモデル (Draft モデル)が、複数の候補トークンを生成します。
2. Verify フェーズ
   高精度なモデル (Target モデル) が候補トークンをまとめて検証し、条件を満たすトークンを受容します。条件を満たさないトークンが見つかった場合は、その位置以降の候補を破棄し、Target モデルの予測結果に基づいて生成を続けます。

上の図は Speculative Decoding の概要を表しています。

図の最上段のように、Draft フェーズで I saw a dog ride in the bus というトークン列が生成されたとします。図の中段は、そのトークン列の in the bus の部分を検証している様子です。
具体的には、次の各トークン列を Target モデルへ入力し、それぞれの次に来るトークンを予測します。

• I saw a dog ride
• I saw a dog ride in
• I saw a dog ride in the

greedy decoding の場合、Draft モデルと Target モデルの予測した Top-1 トークンが一致すれば、そのトークンは受容 (Accept) されます。

一方、予測が一致しない場合、図中の bus と car の例では、Draft モデルが提案した bus は棄却され、Target モデルが予測した car が採用されます。その後、car までのトークン列をもとに、Draft モデルが再び後続の候補列を生成します。

Speculative Decoding による高速化のポイントは、Verify フェーズで、候補列の各位置に対する予測を並列で処理できる点にあります。

逐次的な生成が必要な Draft フェーズには高速なモデルを使用し、候補列をまとめて処理できる Verify フェーズには高精度なモデルを使用します。これにより、Target モデルと同等の出力品質を維持しながら、生成速度の向上を図ります。
両モデルの予測が一致するほど、多くのトークンを一度に生成でき、高い高速化効果が得られます。

W4A4 と W4A16 量子化

LLM の推論において、量子化とはモデル内部の数値を FP16 などの高精度形式から INT4 などの低精度形式へ変換し、推論に必要なメモリ容量や計算量を削減する技術です。
量子化方式は、どの値を何 bit で表現するかによって区別されます。本記事で扱う W4A4 と W4A16 は以下の量子化を意味しています。

W4A16: 重みのみを 4-bit に量子化

W4A16 は、モデルの重みを 4-bit、活性値を 16-bit で扱う方式です。
モデルの重みを 4-bit で保存するため、FP16 モデルと比べて重みのメモリ使用量を大きく削減できます。一方、行列演算では 4-bit の重みを FP16 などに変換しながら、高精度な活性値との積を計算します。
活性値を高精度のまま維持するため、量子化による推論品質の低下を比較的小さく抑えられることが特徴です。

W4A4 : 重みと活性値の両方を 4-bit に量子化

W4A4 は、モデルの重みだけでなく、演算時の活性値も 4-bit に量子化する方式です。
重みと活性値の両方を低精度にすることで、専用の W4A4 カーネルを使用でき、メモリ転送料と行列演算のコストをさらに削減できます。
QSpec の論文では、W4A4 は一般的なタスクでは比較的良い結果を維持できる一方で、GSM8K などの多段階推論を必要とするタスクでは、W4A16 よりも精度低下することが報告されています。

QSpec が利用する性質

W4A4 では多段階推論の最終精度が低下するものの、すべての位置で W4A16 と異なるトークン予測するわけではありません。QSpec の論文では、Top-1 トークン、つまり各位置で最も確率が高いと予測されたトークンは、その大部分が W4A4 と W4A16 の間で一致することが示されています。

QSpec はこの性質を利用し、次の2段階でトークンを生成します。

1. 高速な W4A4 モードで、複数トークンからなる候補列を生成する
2. 高精度な W4A16 モードで、候補列をまとめて検証する

greedy decoding では、各位置における W4A4 モードと W4A16 モードの Top-1 トークンが同じ場合、W4A4 モードが提案したトークンを受容します。Top-1 トークンが異なる場合はその候補を棄却し、W4A16 モードが予測したトークンを採用します。

QSpec の W4A4 モードと W4A16 モード

QSpec では、W4A4 と W4A16 の量子化方式を使った2つのモデルを個別に保持するのではなく同一の QSpec モデルを、W4A4 モードと W4A16 モードで切り替えて実行します。

同じ 4-bit 重み
├─ activation を 4-bit で処理 → W4A4 モード
└─ activation を 16-bit で処理 → W4A16 モード

同じ 4-bit の重みを共有することで、モデルサイズ（メモリ容量）を抑えることができます。

実験：W4A4 と W4A16 の Top-1 一致率を調べる

QSpec 論文の Figure 2 では、GSM8K（算数の文章問題からなるデータセット）を対象に、Atom ベースの W4A4 と W4A16 を適用した Llama-3-8B-Instruct のトークン予測結果を比較しています。
この図には、W4A4 と W4A16 が各トークンに割り当てた確率と、両者の予想が一致したかどうかが示されています。論文では、この結果から次の2点が述べられています。

1. W4A4 と W4A16 のどちらにおいても、大部分のトークンの予測確率が 80% を超えており、高い確率で予測されたトークンの多くが受容されている
2. 受容されたトークンと比べて棄却されたトークンはごく僅かであり、W4A4 と W4A16 の予測結果は高い類似性を示している

そこで今回は、同様の傾向を確認するため、QSpec モデルのデコーディング中に W4A4 モードと W4A16 モードの Top-1 トークンおよびその確率を記録し、両者の一致率を調べます。

なお、本記事の実験は Figure 2 の厳密な再現ではありません。論文の Figure 2 では、W4A16 が生成した回答列を固定し、同一の prefix に対する W4A4 と W4A16 の予測を比較しています。一方、今回測定するのは、QSpec の実際の draft–verify 処理で評価されたトークン位置です。

実験環境

実験には、AWS EC2のg6e.2xlargeインスタンスを使用しました。

項目	内容
インスタンスタイプ	g6e.2xlarge
GPU	NVIDIA L40S
OS	Ubuntu 22.04 ベースの Docker コンテナ
CUDA	12.5.1
Python	3.10
推論基盤	QSpec 版 vLLM

実験条件

主な実験条件は次のとおりです。

項目	内容
データセット	GSM8K test split
サンプル数	1,319
few-shot数	8
最大生成token数	512
decoding	greedy
temperature	0.0
top_p	1.0
QSpec draft長	3
最大同時sequence数	1
乱数seed	0
高confidence判定	Top-1確率 > 0.8

今回は推論性能のベンチマークではなく、内部処理の記録を目的としています。そのため、--max-num-seqs=1 と --enforce-eager を指定し、trace hookによる測定のしやすさを優先します。

測定スクリプト

測定スクリプト (trace_qspec_tokens.py on Github) では、QSpec 版 vLLM の SpecDecodeWorker._verify_tokens に実行時の hook を設定します。QSpec のソースコード自体は変更せず、W4A4 モードが提案したdraft token と、W4A16 モードによる検証結果を取得します。

処理の流れは次のとおりです。

• GSM8K の 8-shot prompt を構築する
• W4A4 で 3 token を生成する
• W4A16 で生成列をまとめて検証する
• 両モードの確率分布と Top-1 token を記録する
• 実際に token が受理されたかを記録する
• 全 request の終了後に集計結果を出力する

結果

測定結果をプロットしたものを次に示します。

散布図の横軸は W4A16 モードの Top-1 確率、縦軸は W4A4 モードの Top-1 確率を示しています。Top-1 確率とは、それぞれのモードが最有力と予測したトークンに割り当てた確率です。

W4A4 モードと W4A16 モードでは、異なるトークンが Top-1 として選ばれる場合があります。そのため、横軸と縦軸の値は、必ずしも同じトークンに対する確率ではありません。

散布図の各点は、QSpec の draft–verify 処理で評価された1つの draft 位置を表します。W4A4 モードと W4A16 モードの Top-1 トークンが一致した点を Accepted、一致しなかった点を Rejected として色分けしています。
また、グラフ上部と右側の Density は、それぞれ W4A16 モードと W4A4 モードにおける Top-1 確率の周辺分布を表しています。

Figure 2 とは測定条件が異なるため、結果を直接比較することはできませんが、同様の傾向が確認できました。得られたデータを集計した結果は、次のとおりです。

項目	結果
リクエスト数	1319
draft 位置数	275,049
Top-1 一致数	263,984
Top-1 不一致数	11,065
Top-1 一致率	95.98%
Top-1 不一致率	4.02%

275,049か所の draft 位置において、W4A4 モードと W4A16 モードの Top-1 一致率は 95.98% でした。両モードの Top-1 トークンは、大部分の位置で一致していることが分かります。

Top-1 確率を集計した結果は、次のとおりです。

指標	W4A4	W4A16
Top-1確率の平均	92.51%	92.99%
中央値	99.966%	99.986%

Top-1 トークンの多くには非常に高い確率が割り当てられていることが分かります。
さらに、両モードの Top-1 確率がともに 0.8 を超える領域は、全体の 82.58% を占めていました。この領域における Top-1 一致率は 99.96% でした。
一方、両モードの Top-1 確率がともに 0.8 以下の領域では、一致率が 70.56% まで低下しました。

以上の結果から、今回の実験条件では、次の傾向が確認できました。

• W4A4 モードと W4A16 モードのどちらにおいても、多くの位置で Top-1 確率が高い
• 両モードの Top-1 トークンは 95.98% の位置で一致する
• 両モードの Top-1 確率がともに 0.8 を超える場合、Top-1 トークンはほぼ常に一致する

これらは、 W4A4 と W4A16 のトークン予測が大部分の位置で類似しているという、QSpec 論文の主張と整合する結果です。

おわりに

今回は QSpec の有効性の前提となっている以下の性質について、測定により確認することができました。

1. W4A4 と W4A16 のどちらにおいても、大部分のトークンの予測確率が 80% を超えており、高い確率で予測されたトークンの多くが受容されている
2. 受容されたトークンと比べて棄却されたトークンはごく僅かであり、W4A4 と W4A16 の予測結果は高い類似性を示している

測定環境

測定環境の詳細：https://github.com/jun-shibata/ai-paper-notes/blob/main/qspec/README.md

コミットハッシュ

QSpec : 1d8124fbb4e4950d69dc35b6fe963d4e559c1104
モデル(AnonymousResearcher2025/Llama3_8B_Instruct_QSpec): 2a1cc64f6abe757e4755393339ce7cbcbebe41cc
データ(openai/gsm8k): 740312add88f781978c0658806c59bc2815b9866