Apple Machine Learning Researchは2026年5月(現地時間)、Transformer言語モデルのKey-Values (KV) キャッシュのメモリ要件を削減する新手法「Stochastic KV Routing (ストキャスティック KV ルーティング)」を発表した。この研究は、オートレグレッシブ生成におけるKVキャッシュの大きなメモリフットプリントとサービングコストへの対処を目指す。従来のKVキャッシュ削減手法が時間軸での最適化に焦点を当てていたのに対し、本手法は深さの次元での最適化を提案する点で特徴を持つ。これにより、メモリ効率の向上と計算コストの削減が期待される。
本研究で提案されたStochastic KV Routingは、レイヤーごとのキャッシュを削除することが情報損失なしに効率的な最適化を可能にすることを示している。提案される訓練アプローチはrandom cross-layer attention (ランダムクロスレイヤーアテンション)と称され、訓練中に各レイヤーが自身のKey-Values (KV) 状態、または先行するレイヤーのKey-Values (KV) 状態のいずれかにランダムにアテンションを向ける。この確率的なプロセスにより、モデルは様々な深さ方向のキャッシュ共有戦略に対して適応性を持ち、展開時の未知のハードウェア制約への柔軟性を確保できる。
評価結果によると、このスキームを事前訓練 (pre-training) またはファインチューニング (fine-tuning) 中に適用することで、様々なモデルファミリー (model families) において深さ方向のキャッシュ共有が可能になることが示されている。特に、データ制約のある設定における大規模モデル (larger models in data-constrained settings) の場合、このアプローチは正則化 (regularization) に似た効果を示唆している。これにより、キャッシュのメモリフットプリントを大幅に削減しつつ、性能を維持または向上させる効果が見られる。
関連研究として、Apple Machine Learning Researchは、長期会話型質問応答のためのKey-Value (KV) キャッシュ管理に関する「EpiCache」を2025年9月23日に、また並列Key-Value (KV) キャッシュ生成によるスケーラブルなLLM推論に関する「KV-Runahead」を2024年5月14日にそれぞれ発表している。
参考: Apple ML Research — 2026年5月5日 09:00 (JST)
原文ハイライト"Stochastic KV Routing: Enabling Adaptive Depth-Wise Cache Sharing"