KeysAndValues、固定32K KV cacheで4Bモデルを128Kコンテキスト向けに微調整

AWS、アムステルダム大学、ベルリン工科大学の研究者らは、長コンテキスト向けのsparse attention微調整手法を提案し、[KeysAndValues](https://github.com/awslabs/keys_values) 0.2.0をApache 2.0ライセンスで公開した。中心となる問題は、モデルをfull attentionで学習し、デプロイ時になって初めて容量固定のKV cacheを適用すると、どの情報がevictionされるかをモデルが一度も学習していないため、出力品質が劣化し、正常に生成を停止できなくなる可能性さえあることだ。

新手法は、forward passでcache policyによる上書き判断を記録し、backpropagation時にそれを再生する。さらに、二層のactivation checkpointing、CPU offload、KV bufferの差分表現によってメモリ使用量を制御する。離散的なeviction戦略を微分する必要がなく、理論上は任意の容量固定戦略と組み合わせられる。チームはH2Oパスも書き直し、FlashInfer SDPAとTriton kernelを組み合わせ、attentionの計算中に各KV slotの重みを直接累積するようにした。

[論文](https://arxiv.org/abs/2608.19920)では、Qwen3-4B-Instruct-2507、LoRA rank 16、32K cacheを使用し、64Kおよび128KのHELMETで評価した。128KのTREC coarseタスクでは、cache policyと共同で微調整したモデルが96.0%〜96.4%を達成したのに対し、full attentionで微調整したモデルをsparse inferenceで実行した場合は23.2%〜77.6%にとどまった。CLINC150でも、61.6%〜68.0%から94.0%〜97.4%へ向上した。エンジニアリング上の重要な点は、長コンテキスト学習でシーケンスを複数GPUへ分割することが必須ではなくなったことだ。サンプルは40GBのA100 1基で実行でき、追加のGPUはbatch throughputの向上に利用できる。

一方、制約も明確だ。4つの自然言語質問応答データセットでは結果にばらつきがあり、NQとHotpotQAの多くの設定ではsequence parallelismのほうが依然として優れている。また、論文で主に評価されたのは1種類の4Bモデルのみだ。FlashInfer拡張にはNVIDIA compute capability 8.0以上が必要で、プロジェクトも本番環境へのデプロイにはなお多大なエンジニアリングが必要だと明記している。今後は、この手法をvLLMやSGLangのpaged batchingへ統合できるか、また、より大規模なモデル、量子化cache、マルチテナント負荷の下でも品質とレイテンシの優位性を維持できるかが注目される。