Mask-Guided KV Cache Eviction in Block Diffusion Language Models
Abstractの概要
本論文では、デノイジング中に長いプレフィックスがメモリに保存され繰り返し読み出されるブロック拡散言語モデルを対象に、推論時のKey-Value（KV）キャッシュ管理を研究しています。著者らは、プレフィックスに対して、現在のブロックに対する一時的な読み出し選択と、将来のブロックに対する永続的な保持決定という2つの異なる判断が必要であると主張しています。そこで、現在のブロックのマスククエリを用いてスパースな読み出しを選択し、将来ブロックのマスクプローブを用いてどのKVエントリを永続キャッシュに残すかをランク付けする、学習不要の手法「MaskAhead」を提案しています。本手法はValueを考慮したスコアリング規則を用いて固定予算下でキャッシュを物理的に圧縮し、論文内では4ビットストレージ派生版であるQ-MaskAheadも評価しています。
新規性
本研究の独自性は、読み出し選択と不可逆な保持決定を分離し、即座の読み出しには現在のマスク、エビクションには将来マスクプローブというように、2つの決定に対して異なるクエリの時間範囲を使用する点にあります。また、この再帰的エビクション機構をValueを考慮したランキングと組み合わせ、Q-MaskAheadにおいてはbf16シャドウキャッシュを用いずにコンパクトな4ビット格納KV表現からランキングとアテンションの双方を実行します。
成果
主要な6つの長プロンプトQA設定において、MaskAheadはエビクション手法の中で（1件の同点を含め）最も高いF1を達成し、高密度推論と比較して平均1.2ポイントのF1低下に抑えつつKVメモリを平均9.5倍削減しました。Q-MaskAheadはメモリ使用量をさらに削減し、高密度推論に対して平均2.3ポイントのF1低下で平均20.1倍のKV削減を達成しています。バッチプロファイル試験では、バッチサイズ32においてMaskAheadとQ-MaskAheadがKVメモリをそれぞれ85.1%および89.9%削減し、高密度bf16と比較して記録されたデコード段階の時間を改善しました。
論文の注目点
- MaskAheadはスパース読み出しとキャッシュ保持を別個の推論決定として扱い、それぞれに異なるマスククエリの時間範囲を使用する。
- 単なる読み出し削減や完全なキャッシュオフロードではなく、物理的なキャッシュ圧縮を伴う学習不要でValueを考慮したKVエビクションを実行する。
- 3つのブロック拡散言語モデルによる実験で比較的小さな品質低下で大幅なメモリ削減が示され、4ビット版のQ-MaskAheadは追加の品質コストと引き換えにメモリをさらに削減する。