STEPQuant：誤差の寿命と位置から設計するDelta則の再帰状態量子化

通常のAttentionでは、過去のキーと値を保持するKVキャッシュが系列とともに増える。線形Attentionは、過去の情報を固定サイズの再帰状態にまとめる。再帰状態とは、新しいトークンを処理するたびに書き換え、次の処理へ引き継ぐ記憶だ。STEPQuantが圧縮するのは、この状態である。

系列が長くなっても状態の大きさが増えないことと、サービス全体のメモリ負担が小さいことは別の話だ。状態はリクエストごとに必要になる。さらに、処理に備えて予約する状態スロットもメモリを使う。著者らは、こうした状態の集まりが高い並行度で負担になると位置付けている。

そこで考えられるのが量子化だ。数値を少ないビットで表せる値へ丸めれば、保存量を減らせる。ただし、状態は一度保存して終わるデータではない。丸めた状態を次の更新に使い、更新後にまた丸める。各回の小さな差が、その後の生成へ持ち越される点が問題になる。

対象はDelta則で状態を更新する線形Attentionだ。更新では、以前の状態を保持ゲートで減衰させ、キーに応じた変換を行い、新しいキーと値の情報を書き込む。出力は、クエリを使って状態から読み出す。保持ゲートは、過去の情報をどれだけ残すかを制御する値である。

論文が扱うGDNとKDAでは、ゲートの単位が違う。GDNは、Attentionの処理単位であるヘッド全体に同じ保持ゲートを使う。KDAはキーの成分ごとにゲートを持つ。そのため、誤差の残りやすさを扱う単位も、前者ではヘッド、後者では状態行列のキー側の行になる。

STEPQuantの出発点は、量子化誤差の大きさだけでは影響を判断できないという観察だ。同じ大きさの誤差でも、すぐ減衰する部分と長く残る部分では、後の更新への影響が違う。また、キー行ごとに出力への響き方が異なり、数値の大きさも行と列の両方で変わる。著者らは、この時間方向と空間方向の違いを、精度配分と丸め方に反映する。

著者らは、量子化した経路とFP32の参照経路に同じキー、値、クエリ、ゲートを与える条件で、誤差の動きを整理する。

ここで、$t$は更新の時点、$E_t$はその時点の状態誤差、$A_t$は以前の状態に作用する変換、$\varepsilon_t$は今回の量子化で新しく加わる誤差だ。つまり、過去の誤差を引き継ぐ項に、新しい誤差が足される。

キーの大きさ、書き込み係数、保持ゲートが所定の条件を満たすと、変換自体は誤差を拡大しない。それでも、誤差が繰り返し加われば蓄積し得る。誤差を増幅しないことだけでは、低ビット化の安全性は保証できない。

この式は、二つの経路が同じ入力を受ける場合の解析である。自己回帰生成では、出力の違いが後の入力の違いにつながるため、生成全体の誤差を保証する式としては読めない。STEPQuantの寿命推定も、ゲートによる減衰を使った近似である。

STEPQuantはSpatial-Temporal Quantizationの略称で、学習済みモデルを校正して量子化するPTQ、つまり学習後量子化の手法だ。時間方向では「寿命を考慮したビット配分」を行う。

まず校正データから、各単位の平均対数保持率を求める。それを使い、有限の期間に誤差がどれほど残るかを寿命の重みとして見積もる。次に、候補となる精度で量子化した場合の歪みを測る。歪みは元の状態との差であり、Qwenでは行の感度を反映した平均二乗誤差、KDAでは行ごとの平均二乗誤差を使う。

配分は、寿命で重み付けした歪みの合計を、平均ビット予算の範囲内で小さくするように決める。長く残り、量子化の差も大きい単位へ精度を回す考え方だ。さらに少数の高リスク単位を、FP16という浮動小数点形式で保護する。精度と保護領域はオフラインで決め、リクエスト間で固定する。

校正には全評価構成でWikiText-2の訓練部分から32区間、各2,048トークンを使う。名目4ビット予算の整数候補は2、4、6、8ビット、6ビット予算では4、6、8ビットだ。「4ビット設定」は、すべての場所を4ビットにするという意味ではない。

空間方向の工夫は「キー行の感度を考慮した二軸の当てはめ」だ。量子化で使うスケールとは、整数の値を元の数値の大きさへ戻すための係数である。基本比較の一様量子化は、各キー行の絶対値の最大値に基づいてスケールを決める。STEPQuantは行と列の二つの係数を使う。

まず、状態更新の変換とクエリから、各キー行の誤差が読み出しにどれだけ響くかを校正で推定する。状態内の値が大きい行と、出力に敏感な行は、同じとは限らない。行の係数には平均絶対値と感度を組み込み、高感度の行では量子化の刻みを細かくする。

列の係数は、行の感度で重み付けした再構成誤差が小さくなるように合わせる。状態を、行係数、列係数、整数コードの積で表すことで、両方向の数値の偏りを扱う。寿命に応じて保存精度を選ぶ部分と、同じ精度でも誤差の置き方を変える部分が組み合わされる。

生成中は、前の状態を復元し、浮動小数点で更新と現在の出力の読み出しを行う。その後、更新した状態を量子化して保存する。したがって、その回に新しく注入した誤差が作用するのは次回以降である。著者らはSGLangに統合し、状態の圧縮を含む処理として評価している。

精度評価はQwen3.8-27BとKimi-Linear-48B-A3B-Instructを対象に、BF16重みと4ビットAWQ量子化重みで実施した。比較の基準はFP32状態と、キー行ごとのスケールを使う対称一様INT4、INT6、INT8だ。計算にはA800を4枚使っている。

長文生成の7課題はLiveCodeBench v6、EvalPlus、AIME 2026、MATH-500、HMMT February 2026、GPQA Diamond、IFBench。問題当たりのサンプル数はこの順に5、5、64、4、64、8、4で、温度1.0、top-k 20、top-p 0.95、最大65,536生成トークンという条件を使う。短文はMMLU、ARC-C、OpenBookQA、HellaSwag、WinoGrande、LAMBADAの6課題だ。毎回最も確率の高いトークンを選ぶgreedy生成で回答を採点し、抽出できない回答は誤答とする。短文の点数には回答形式への対応も含まれる。

機構の検証では、C4の4系列でFP32と入力を共有し、2,046回更新した。一様INT6の累積二乗状態誤差と、ゲート由来の半減期とのSpearman相関は、Qwenで0.8004、KDAで0.8017だった。これは寿命が誤差に関係するという観察を支えるが、自由な生成での保証ではない。

性能測定はBF16重み、A800 4枚のテンソル並列、固定バッチ32〜512で行った。入力128トークン、生成1,024トークンのgreedy生成を使い、warmup後の3実行の中央値を報告する。デコード中の当てはめや書き戻しは含むが、入力をまとめて処理するprefillや最終出力配信は含まない。

以下はすべて著者らの報告値だ。BF16重みでの長文7課題の平均正解率は、名目6ビットSTEPQuantでQwenが80.59%、Kimiが61.47%。FP32状態の80.60%、61.52%に近く、一様INT6の45.04%、45.70%との差は大きい。

名目4ビットでもQwenは80.51%、Kimiは58.52%で、一様INT8の71.86%、56.02%を上回る。ただし、FP32との差はQwenで0.09ポイント、Kimiで3.00ポイントの低下だ。少ないビットで一様INT8を上回ったことと、FP32と同等であることは分けて評価する必要がある。

短文6課題では、名目4ビットがQwenで87.63%、Kimiで68.11%。FP32からの低下は0.15、0.25ポイントだった。4ビットAWQ重みとの組み合わせでも、名目6ビットの長文平均は79.27%、58.62%で、同じ重みのFP32状態から0.05、0.33ポイントの低下にとどまった。

構成要素を外す試験も、組み合わせの意義を支える。QwenのBF16重み、名目4ビットでAIME、GPQA、LiveCodeBenchの平均を見ると、STEPQuantは84.72%、空間成分のみは73.95%、時間成分のみは12.87%。移植したQ-Mambaは7.64%、FP32は84.61%だった。この限定された比較では、寿命による配分だけでも、二軸の当てはめだけでも、全体の精度には届かない。

名目ビット予算と実保存量は区別したい。スケールとFP16保護領域を含むcompact表現は、6ビット設定でQwenが6.361、KDAが6.300 bits/value、4ビット設定で4.621、4.300 bits/valueとなる。KDAの値は解析的な会計で、並列処理に伴うpaddingやメモリ割り当ての付加分を除く。状態の表現量、予約プール、総サービングメモリは別の範囲を測っている。

A800とSGLangを使うバッチ512、名目6ビットの指定構成では、Qwenの4ビット重みで総メモリが419.73から131.18 GiBへ68.7%減少した。状態メモリは5.03倍の圧縮だ。対応するKimiの評価では総メモリが149.70から69.36 GiBへ53.7%減り、状態メモリは5.08倍に圧縮された。いずれも同じモデルのFP32状態構成との比較である。

一方、モデル全体のデコード性能の改善幅は、状態更新単体より小さい。BF16重み、バッチ512では、Qwenのスループットが6,040から7,280 tokens/sへ20.53%向上し、Kimiは21,241から23,748 tokens/sへ11.80%向上した。Qwenの状態更新単体の2.91倍高速化を、そのまま全モデルの速度として扱うことはできない。これらの測定から、リクエスト全体の待ち時間や運用費用の削減率までは分からない。

正解率が近くても、生成に必要な量が変わる場合がある。Qwenの4ビットAWQ重みで長文7課題を見ると、名目4ビットSTEPQuantの平均生成長は9.50Kトークンで、FP32状態の7.22Kより約31.6%長かった。平均正解率は78.98%で、FP32の79.32%から0.34ポイント低下した。31.6%は生成長の相対的な増加、0.34ポイントは正解率の差である。

この生成長には、思考部分だけでなく誤答や生成上限に達した出力も含まれる。長く生成したことを、推論の改善とは解釈できない。名目4ビットを検討するなら、Kimiの長文精度低下とともに、Qwenの生成量増加も負担として確認する必要がある。

理論面にも境界がある。寿命の重みはゲートによる減衰の近似で、時間とともに変わるキー依存の変換や、注入された誤差同士の交差項を完全には扱わない。実験はその近似が評価対象で役立つことを示すが、一般的な誤差保証にはならない。

適用の根拠があるのは、評価されたGDN/KDAモデルに近い並行デコードの条件だ。他の再帰アーキテクチャ、別のハードウェア、動的なリクエスト構成で同じ効果が出るかは未検証である。また、並行研究DAMPとの比較は実装を再現した直接比較ではなく、生成設定も異なる。手法間の優劣を確定する材料にはできない。