End-to-End Context Compression at Scale
Abstractの概要
本論文は、増加するKVキャッシュによるメモリのボトルネックに対処するため、長文コンテキスト言語モデル推論のためのエンコーダ・デコーダ型コンテキスト圧縮を再検討しています。著者らは、長いトークンシーケンスをより短い潜在埋め込みに圧縮し、デコーダが元のコンテキストの代わりにそれを消費するLatent Context Language Models (LCLMs) を導入しました。彼らはアーキテクチャ探索と段階的な学習手法を組み合わせ、その後、4倍、8倍、16倍の圧縮率で、それぞれ3500億トークン以上を用いて0.6Bエンコーダと4Bデコーダのモデルを継続的に事前学習しました。長文コンテキストのベンチマークやシステム測定を通じ、エンドツーエンドの潜在圧縮が標準的な推論フレームワークとの互換性を保ちながら、品質、圧縮速度、メモリ使用量のトレードオフを改善できると論じています。
新規性
これまでKVキャッシュ手法に遅れをとっていたエンコーダ・デコーダのソフトトークン圧縮が、大規模なエンドツーエンド学習と慎重に最適化されたアーキテクチャによって競争力を持ち得ることを示した点が主な新規性です。さらに本論文は、モデルが圧縮されたコンテキストを全体的に調べ、必要に応じて選択したセグメントを展開できるエージェント設定へと潜在圧縮を拡張しています。
成果
報告された結果によると、LCLMは高い精度を維持しつつ、長いコンテキスト長において最初のトークン生成までの時間を短縮し、ピークGPUメモリを削減することで、RULER、LongBench、LongHealthなどのベースラインを超える新たなパレートフロンティアを形成しています。さらに、プロンプト全体を圧縮した場合、GSM8Kにおいてすべての圧縮率でLCLMが最高の精度を達成し、展開可能なエージェントが「干し草の山から針を探す(needle-in-a-haystack)」タスクでの検索性能を大幅に向上させ、非圧縮コンテキストの性能に匹敵する場合があることも報告しています。
論文の注目点
- LCLMはデコーダのプレフィル前にコンテキストを圧縮することでデコーダ側のシーケンス長を直接短縮し、報告された実験において、KVキャッシュ退避アプローチよりも優れたレイテンシとメモリスケーリングをもたらしている。
- 著者らはアーキテクチャ探索と段階的な学習を通じて効果的な設計を特定しており、0.6Bエンコーダと4Bデコーダ、因果マスク、MLPアダプタ、および圧縮率に依存したプーリングの選択が有効であるとしている。
- エージェントベースの設定において、モデルは圧縮されたコンテキストを広く読み取り、512トークンのセグメントを選択的に展開できるため、固定の16倍圧縮コンテキスト単独の場合よりも「干し草の山から針を探す」タスクでの厳密な検索精度が向上する。
参考リンク
- arXiv: https://arxiv.org/abs/2606.09659v1
- Fugu-MT: https://fugumt.com/fugumt/paper_check/2606.09659v1
- Hugging Face Papers: https://huggingface.co/papers/2606.09659