確率的KVルーティング：適応的な深さ方向キャッシュ共有の実現

Apple Machine Learning Researchチームは、大規模言語モデル推論時のKVキャッシュメモリ使用量を削減する新しい手法「確率的KVルーティング」を発表しました。本論文はAnastasiia Filippova、David Grangier、Marco Cuturi、João Monteiroによって執筆され、2026年5月に公開されました。KVキャッシュは自己回帰生成における冗長計算を避けるために不可欠ですが、そのメモリフットプリントはサービスコストに大きな影響を与えます。

既存の研究は主に時間軸に沿ったKVキャッシュの圧縮や削除に焦点を当てていますが、Appleのチームは深さ次元が直交的で堅牢な最適化の余地を提供すると主張します。過去の研究ではすべての層に完全なキャッシュが必要ないことが示唆されていますが、層間キャッシュ共有の実装は実用的な課題であり、既存の手法はスループット低下や最初のトークンまでの時間増加を招くことが多いです。本論文の核心は、トレーニング中にランダムなクロスレイヤーアテンションを導入することです。各層は自身のKV状態か前の層のKV状態に注意を向けるかをランダムに選択します。この確率的プロセスにより、モデルは様々な深さ方向キャッシュ共有戦略に対してロバストになり、デプロイ時の未知のハードウェア制約に対する柔軟性が確保されます。

評価実験では、このスキームを事前学習またはファインチューニングに適用することで、様々なモデルファミリーで深さ方向キャッシュ共有が可能になることが示されました。さらに、データ制約のある大規模モデルでは、このアプローチは正則化効果を示し、キャッシュメモリを大幅に削減しながら性能を維持または向上させることが頻繁に見られました。この研究は効率的な大規模言語モデル推論の新たな方向性を示し、展開コストの削減とより広範な応用を促進する可能性があります。詳細はApple Machine Learning Researchのウェブサイトでご覧いただけます。