LLM推論を高速化する2つのアプローチ：CDLMとAttention Matching KV圧縮

LLMを安く、速く動かしたい――その要求に応える研究が2本、ほぼ同時期に出てきた。一つはTogether AIが発表した「Consistency Diffusion Language Models（CDLM）」、もう一つはMIT・HarvardチームによるKVキャッシュ圧縮手法「Attention Matching」だ。

拡散言語モデルの速度問題を解く：CDLM

LLMの主流はAutoregressive（AR）モデルだが、最近「拡散言語モデル（Diffusion Language Model, DLM）」が出てきている。DLMはマスクされたトークンをノイズから同時に復元するアプローチで、双方向コンテキストを活用できる。DreamやLLaDAが代表格だ。ただし推論が遅すぎる。CDLMはその速度問題を潰しにいった研究で、arXiv:2511.19269として公開されている。

DLMが遅い2つの理由

DLMの推論ボトルネックは2つある。

1つ目は、リファインメントステップ数の多さ。DLMはマスクされたトークン位置を反復的に更新していくが、生成長と同じだけのステップが必要になる。例えばDream-7Bで256トークンを生成する場合、256ステップ分の全トークン再計算が走る。

2つ目は、KVキャッシュとの相性の悪さ。ARモデルは因果マスクのおかげで過去トークンの計算結果をキャッシュして再利用できる。DLMは双方向アテンションを使うため、あるステップでトークンが更新されると、それより前のトークンのKey/Valueも無効化される。キャッシュが使えないから毎ステップ全コンテキストを再計算する必要があり、これが致命的に遅い。

CDLMの設計：ブロック因果マスクと一貫性学習

CDLMはこの両方に手を入れる。教師-生徒フレームワークで、教師モデル（通常のDLM、双方向アテンション）から生徒モデル（ブロック因果マスク付き）へ知識を蒸留する。

ブロック因果マスクの仕組み

生成対象のテキストを固定長ブロック（サイズB）に分割する。ブロック間には因果マスクを適用し、各ブロックは自分より前のブロックとプロンプトにしかアテンションを向けない。ブロック内では双方向アテンションを維持する。

この設計により、完了済みブロックのKVキャッシュを再利用できる。新しいブロックのデコード中は、そのブロック内のトークンだけが変化するため、過去ブロックのKey/Valueは有効なまま残る。ARモデルのKVキャッシュと同じ原理だが、単位がトークンではなくブロックになっている。

3つの学習目標

CDLMの学習では以下の3つの損失関数を同時に最適化する。

蒸留損失: 教師DLM（双方向）の隠れ状態から復元したロジットを使い、新しくアンマスクされた位置で生徒の予測分布を教師に近づける。ホワイトボックス蒸留で、教師の最終隠れ状態をlm_headに通して再構成した分布とのKLダイバージェンスを最小化する
一貫性損失: ブロック内の「途中の状態y」と「ブロック完了状態y*」で、まだマスクされている位置の予測分布が一致するように制約する。つまり、ブロック内のどのステップで予測しても、そのブロック終了時と同じ予測を出すべきだという制約。これにより少ないステップでもブロック完了時と同等の品質が得られるようになる
DLM損失: 通常のマスク言語モデリング損失。正解テキストに対するクロスエントロピーで、全体的な言語モデリング性能を維持する

最終損失は L = w_distill * L_Distillation + w_cons * L_Consistency + w_dlm * L_DLM の重み付き和になる。

トラジェクトリ収集

学習データは教師DLMの推論軌跡から作る。各プロンプトに対して教師がブロック単位でデコードし、各ステップの入力状態と最終隠れ状態を記録する。温度パラメータを変えて複数の軌跡を生成し、データの多様性を確保する。DLMでは温度がトークンの確定順序にも影響するため、ARモデル以上にデータ拡張の効果がある。

推論時の信頼度閾値付き並列デコード

推論時は、各ブロック内で信頼度の高いトークンを複数同時に確定させる。一貫性損失のおかげで、ブロック内の途中状態でもブロック完了時と近い予測が出るため、1トークンずつ確定させる必要がない。これがステップ数削減の直接的な原因になる。

ベンチマーク結果

Dream-7B-Instructをベースとした評価結果を論文のTable 1から引用する。

ベンチマーク	手法	レイテンシ	ステップ数	スコア
GSM8K-CoT	Dream-7B（ベースライン）	23.5s (x1.0)	256.0 (x1.0)	79.1
GSM8K-CoT	CDLM-Dream	2.1s (x11.2)	44.1 (x5.8)	78.8
HumanEval-Instruct	Dream-7B（ベースライン）	13.4s (x1.0)	256.0 (x1.0)	48.2
HumanEval-Instruct	CDLM-Dream	2.2s (x6.1)	49.6 (x5.2)	50.0
MBPP-Instruct	Dream-7B（ベースライン）	21.7s (x1.0)	256.0 (x1.0)	51.8
MBPP-Instruct	CDLM-Dream	1.5s (x14.5)	33.2 (x7.7)	53.0

GSM8K-CoTではスコア79.1 -> 78.8とほぼ劣化なしでレイテンシ11.2倍改善。MBPP-Instructではスコアがむしろ向上（51.8 -> 53.0）しつつ14.5倍の高速化を達成している。コーディングタスクで特に効果が大きいのは、コードのような構造化テキストではブロック内の並列確定が効きやすいためだろう。

既存のDLM高速化手法（dLLM-Cache、Fast-dLLM）と比較しても、CDLMはレイテンシ・スループット両面で上回っている。特にFast-dLLM（並列デコード+デュアルキャッシュ）との差が顕著で、ブロック因果マスクによるKVキャッシュ完全対応の効果が出ている。

制約と限界

論文のDiscussionセクションでいくつかの制約が述べられている。

CDLMはpost-training手法であり、教師モデルの推論軌跡を事前に収集する必要がある。この収集コストがそれなりにかかる
ブロック因果マスクによって双方向コンテキストの範囲が制限されるため、教師モデル（完全双方向）より理論的な表現力は下がる
MATHベンチマークではスコアが38.0 -> 32.4と低下しており、高度な数学推論ではステップ削減の影響が出やすい
LLaDA-8Bにも適用されているが、Dream-7Bほどの改善幅は出ておらず、ベースモデルとの相性がある

参照: Consistency Diffusion Language Models / arXiv:2511.19269

KVキャッシュをもっと直接的に圧縮する：Attention Matching

CDLMとは全く別の方向で推論コストを削るのが、MIT・HarvardチームのAttention Matching（arXiv:2602.16284）だ。こちらはARモデルを前提に、長文処理でボトルネックになるKVキャッシュのメモリ消費を圧縮する。コードは https://github.com/adamzweiger/compaction で公開されている。

KVキャッシュ圧縮の問題設定

Transformerでは各トークンのKey/Valueペアをキャッシュとして保持し、後続のアテンション計算に再利用する。コンテキストが長くなるほどキャッシュサイズは線形に増加し、100Kトークンを超えるような長文処理ではGPUメモリを圧迫する。

従来の圧縮手法は大きく3種類ある。

トークン選択: H2O、SnapKV、PyramidKV、KVzipなど。アテンションスコアに基づいて重要なトークンだけ残す。高圧縮率では急速に劣化する
トークンマージ: CaMなど。近いトークンのKVを統合する
潜在空間での最適化: Cartridgesなど。コンパクトなKVキャッシュをエンドツーエンドで学習する。品質は高いが、1つのコンテキストの圧縮にGPU数時間かかる

「summarization can be highly lossy, substantially harming downstream performance」――テキスト要約によるKV削減は、アテンション分布を変えてしまい下流タスクの精度が落ちる。Cartridgesは50倍圧縮でも品質を維持できるが、勾配ベースの最適化が遅すぎて実用的でない。

なぜKeyとValueを別々に圧縮できるのか

KVキャッシュの圧縮を考えるとき、KeyとValueをセットで扱う必要はない。これはTransformerのアテンション計算の構造から来ている。

アテンションの計算は2段階に分かれる。

スコア計算: クエリQとキーKの内積でアテンションスコアを求め、softmaxで正規化する。ここではValueは一切関与しない
出力計算: 正規化されたアテンション重みでValueの加重和を取り、最終出力を得る。ここではKeyは直接関与しない

つまりKeyは「どのトークンに注目するか」を決めるインデックス的な役割、Valueは「注目先から何の情報を取り出すか」を決めるコンテンツ的な役割を担う。この2つの役割が計算グラフ上で直列に接続されつつも、変数としては独立しているため、別々の圧縮戦略を適用できる。

Attention Matchingの定式化でC_kとC_vを別々のサブ問題として解いているのも、この構造的分離が根拠になっている。C_k（とバイアスβ）を先に決めてアテンション重みを確定させ、その後でC_vを最小二乗で求める――この逐次的な解法が成立するのは、KとVが計算上独立だからだ。

KとVの圧縮耐性は対称ではない

KeyとValueを別々に扱えるだけでなく、両者の圧縮耐性には明確な非対称性がある。2024年に出た複数の研究がこの現象を報告している。

Keyは低精度量子化に強い

KIVI（arXiv:2402.02750）はKeyをINT2、ValueをINT2にそれぞれ量子化する手法だが、実験の過程でKeyの量子化がValueより品質劣化に与える影響が小さいことを示している。KVQuant（arXiv:2401.18079）も同様の観察を報告しており、Keyに対してはper-channel量子化が有効で、INT4やINT2まで落としてもperplexityの上昇が比較的小さい。

この理由はKeyの分布特性にある。RoPE（回転位置埋め込み）が適用されたKeyベクトルは、チャネルごとに値域が大きく偏る傾向がある。一見量子化に不利に見えるが、per-channel量子化（チャネルごとにスケールを設定）で対処すると、チャネル内の値域が狭いため少ないビットでも精度よく表現できる。加えて、アテンションスコアはsoftmax前の内積値であり、ランキング（どのトークンが上位に来るか）さえ保存されれば、スコアの絶対値が多少ずれてもsoftmaxが吸収する。

Valueは圧縮に敏感

一方でValueの量子化や圧縮は、出力品質に直接響く。Valueはアテンション重みとの加重和で最終出力になるため、Valueの誤差がそのままレイヤー出力の誤差として伝搬する。softmaxのような正規化で吸収される余地がない。

KIVI、KVQuant、GEAR（arXiv:2403.05527）のいずれも、Valueの量子化ビット数を下げたときの劣化がKeyより顕著であることを報告している。GEARはこの問題に対して、低ランク近似＋残差の均一量子化＋外れ値の個別保持というハイブリッド戦略を取っている。QServe（arXiv:2405.04532）もKVキャッシュの量子化でKeyをINT4に圧縮しつつ、Valueにはより慎重なスキームを適用する設計になっている。

Attention Matchingとの接続

この非対称性は、Attention Matchingの3段階分解と自然に対応する。C_kの選択（ステップ1）とバイアスβのフィッティング（ステップ2）はアテンション分布の再現を担い、ここでは「ランキング保存」が本質的に求められる。C_vのフィッティング（ステップ3）は出力値の再現を担い、最小二乗でレイヤー出力を直接近似する。Valueの方が品質に直結するからこそ、閉形式の最小二乗解で精密にフィットする設計は理にかなっている。

Attention Matchingの定式化

Attention Matchingは圧縮問題を「アテンション出力の再現」と「アテンション質量の保存」の2条件に落とし込む。

元のKVキャッシュ (K, V) を、より小さい (C_k, C_v) に置き換えたい。T個のトークンをt個に圧縮する（t << T）。任意のクエリqに対して、圧縮後のアテンション出力が元のアテンション出力を再現すればよい。

ここで2つの条件を課す。

アテンション出力マッチング: 圧縮KVから計算したローカルアテンション出力（softmax重み付き和）が、元のKVから計算した出力と一致する
アテンション質量マッチング: 圧縮KVブロック全体のexpスコア合計（Mass）が、元のKVブロックのそれと一致する

なぜ質量保存が必要かというと、圧縮KVは他のKVブロック（直近の会話ターンや将来の生成トークンなど）と連結されて使われるため。連結時のアテンション重みは各ブロックの質量の比で決まるので、質量がずれると圧縮ブロックへの注意配分が狂う。

スカラーバイアスの導入

t < Tの場合、キーだけでは質量マッチングが不可能になる。例えばq=0のとき、Mass(0; K) = TだがMass(0; C_k) = tで、どんなC_kを選んでもTにできない。

そこで各圧縮キーにスカラーバイアスβを導入する。アテンションスコアがq * C_k^T + βになるように修正し、β_j = log(w_j)として、w_jが「この圧縮キーが元のキー何個分の質量を代表するか」を表す。メモリオーバーヘッドは(2d+1)/(2d)倍とごくわずかで、PyTorch SDPAやFlexAttentionでサポートされている。

閉形式解による高速圧縮

圧縮処理にエンドツーエンドの勾配降下法がいらない。最適化問題を3つのサブ問題に分解し、それぞれ順番に解く。

圧縮キーC_kの選択: 元のキーから重要なものを選ぶか、Orthogonal Matching Pursuit（OMP）でグリーディに選択する
バイアスβのフィッティング: C_kが決まったら、質量マッチング条件は非負最小二乗問題（NNLS）になる。行列Aの要素がA_ij = exp(q_i * C_k_j^T)、目標ベクトルmが元の質量。w_j >= 0の制約付き最小二乗を解いてβ_j = log(w_j)とする。これは閉形式解を持つ
圧縮バリューC_vのフィッティング: アテンション出力マッチング条件は通常の最小二乗問題になる。圧縮されたアテンション重みが決まっている状態で、重み付き和が元の出力を再現するC_vを求める。これも閉形式で解ける

参照クエリQ_refの取得方法も工夫されている。

Repeat-prefill: コンテキストを「繰り返し」指示と共に再度入力し、復元中のクエリベクトルを抽出する
Self-study: コンテキストに対して合成的な質問応答を生成し、そのクエリベクトルを使う
On-policy: レイヤーを順番に圧縮し、各レイヤーで前のレイヤーが既に圧縮された状態からクエリを抽出する。分布シフトを軽減する

バリアント一覧

論文では4つの手法バリアントを評価している。速度と品質のトレードオフが異なる。

AM-OMP: On-policyクエリ（self-study + repeat-prefill）、OMPでキー選択、NNLSでバイアス、最小二乗でバリュー。最高品質だが最も遅い
AM-OMP-fast: AM-OMPの高速版。1回のOMP反復で4キーを選択、バイアスの再フィットは2反復に1回
AM-HighestAttnKeys: On-policyクエリ、アテンションスコア上位でキー選択、NNLSとLS。高速
AM-HighestAttnKeys-fast: repeat-prefillのみでクエリ取得。最速

実験結果

Qwen3-4B、Llama3.1-8B、Gemma3-12Bの3モデルで、QuALITY（長文読解QA、5-8Kトークン）とLongHealth（医療記録QA、60Kトークン）を評価している。

50倍圧縮・Qwen3-4BでのQuALITYの結果（Figure 1から）:

AM-OMPはCartridgesと同等の精度を達成しつつ、圧縮時間は2桁速い（数分 vs 数時間）
H2O、SnapKV、PyramidKVなどのトークン選択ベースラインは50倍圧縮では大幅に劣化
テキスト要約ベースラインも同様に劣化が大きい
AM手法群がパレートフロンティアを形成している

20倍から100倍までの圧縮率を変えた実験（Figure 3）では、特に高圧縮域（20倍以上）でAttention Matchingの優位性が際立つ。LongHealthのような情報密度の高いデータセットでは全手法が早く劣化するが、相対的な順位は変わらない。

圧縮時間の内訳（Gemma3-12B、60Kトークン、H200 GPU、Table 1から）:

context-prefill: 7秒
repeat-prefill: 8秒
self-study: 139秒
OMP圧縮: 分単位

self-studyが最も時間を食うが、省いても（AM-HighestAttnKeys-fast）実用的な品質は得られる。

Gemma3-12Bはsliding-window attention（局所5層に対して大域1層）を使うハイブリッドアーキテクチャだが、大域アテンション層だけを圧縮対象にすることで同様の改善が得られている。

制約と限界

圧縮対象のコンテキスト全体をまず通常通りprefillする必要がある。圧縮自体は速いが、prefillのコストは残る
圧縮キーにRoPE位置埋め込みが適用済みの前提。圧縮後のキャッシュはlogical lengthとしてTを保持し、後続トークンのRoPE位相を正しく維持する設計
FP32で圧縮キー/バリューを計算してBF16にキャスト。量子化との組み合わせは未検証
現時点ではオフライン（バッチ）圧縮が前提。ストリーミングでのオンライン圧縮は付録で簡単に触れられているが、主要な評価対象ではない

論文情報: Fast KV Compaction via Attention Matching / Adam Zweiger, Xinghong Fu, Han Guo, Yoon Kim / 2026年2月18日 / https://arxiv.org/abs/2602.16284