NVIDIAが開発、8倍のKVキャッシュ圧縮を実現するDynamic Memory Sparsification (DMS) 技術が大規模言語モデルの推論効率を大幅に向上させる

NVIDIAが発表した「動的メモリ疎化(DMS)」 近日、NVIDIAとエディンバラ大学の研究チームが、「動的メモリ疎化（Dynamic Memory Sparsification: DMS）」という新しい手法を発表しました。DMSは、推論時の大規模言語モデル（Large Language Models: LLMs）のパフォーマンスを向上させるために、キーバリュー（KV）キャッシュの圧縮に焦点を当てています。この手法は、データ効率が高く、既存のモデルにも容易に実装できる特徴が注目されています。 推論時のボトルネック：KVキャッシュ Transformerベースのモデル（GPT、LLaMA、Qwenなど）は、オートレグレッシブ生成のために過去のトークン表現をKVキャッシュに保存します。ただし、このキャッシュはシーケンス長や並列性（複数の思考チェーン）とも比例して増大し、大量のGPUメモリを消費するために推論速度が低下する問題があります。 既存の最適化技術の限界 KVキャッシュの最適化には、訓練不要のヒューリスティック（注意重みに基づくトークン除去）と、訓練後の重いリフィット（Dynamic Memory Compression: DMC）などが存在しています。前者は精度を低下させる傾向があり、後者は計算的に高価であり、どちらも完全な解決策とはなりきっていません。 前例のない圧縮：DMSの導入 DMSは、これらの欠点を克服するために、ハイブリッドアプローチを取り入れました。伝統的な剪定法のようにKVキャッシュを疎化し、訓練オーバーヘッドを最小限に抑え（約1,000ステップ）、一時的なトークン保持機能（遅延排出）を導入しています。こうすることで、重要コンテキスト情報を維持しながら突然の精度低下を避けられるようになっています。 DMS的核心は、訓練時における eviction （除去）決定を不同分可能にするGumbel-sigmoidベースのサンプリングメカニズムです。未来の除去が予測されるトークンは、スライ딩ウィンドウ期間中は引き続き使用可能で、その後排出されます。これにより、モデルは情報の価値をさらに効果的に吸収できます。 高効率のリフィット DMCと異なり、DMSは追加のパラメータを導入せず、単一のニューロンを利用してevictionを予測します。これにより、既存のモデルを大幅なアーキテクチャ変更なしに容易にアップグレードできます。1,000回程度の訓練ステップで、DMSは最大8倍のKVキャッシュ圧縮を実現し、推理タスクでのモデル性能を維持または向上させました。 ベンチマーク結果：パフォーマンス向上とコスト削減 Qwen-1.5B、7B、32Bといった異なるモデルサイズでテストされたDMSは、AIME（+9.1ポイント）、GPQA（+7.6ポイント）、LiveCodeBench（+9.6ポイント）といった推論重視のベンチマークで、正確性を保ちつつパフォーマンスを向上させることが確認されました。 QuestやTOVAなどのトップクラスのベースラインと比較しても、DMSはKVキャッシュの読み出し効率（実行時間の代理指標）とピークメモリ使用量において一貫して優れ、より効果的なParetoフロンティアを達成しました。 多様なタスクへの有用性 DMSは推論以外のタスクでも有効性を示しました。短文コンテキストのMMLU、GSM8K、HellaSwagでは、最大4倍の圧缩比でわずかな精度低下（約3.5ポイント）に留まりました。また、長いコンテキストを持つNeedle-in-a-HaystackやVariable Trackingなどのタスクでは、純粋なモデルよりも優れた性能を発揮し、情報の過剰押しなどの問題を軽減する可能性があることが示唆されました。 結論 動的メモリ疎化（DMS）は、Transformerベースの言語モデルの推論時効率を強化するための実践的かつ拡張性の高いソリューションを提示しています。DMSは、少量の再訓練でKVキャッシュを効率的に圧縮し、より長いシーケンスや並列思考チェーンにモデルが対応できるようにします。一連の推論や一般用途のタスクでの一致する改良によって、DMSの多様性と効果が証明されており、限られたリソース環境でのLLMデプロイメントに有望な道を提供します。 業界関係者からは、「DMSは、推論時のメモリmanagementを効率化する新しいアプローチであり、既存のモデルへの適用性が高い」との評価が出ています。 NVIDIAは、この研究を通じてLLMの限界をさらに押し広げ、実際の推論ワークロードに適応する新しい技術を開発することで知られています。