エンドツーエンド強化学習による圧縮メモリを活用した動的長文脈推論

要旨

大規模言語モデル（LLMs）は、二次的な計算コスト、情報の忘却、および検索拡張生成（RAG）に内在する文脈断片化など、長文脈処理において重大な課題に直面している。本論文では、すべての生トークンを処理するのではなく、チャンク単位の圧縮と選択的記憶想起に基づく、認知科学にヒントを得た効率的な長文脈推論フレームワークを提案する。このフレームワークは、長い入力をチャンクに分割し、学習済みの圧縮器を用いて各チャンクを圧縮された記憶表現に符号化する。ゲーティングモジュールは関連する記憶ブロックを動的に選択し、それらは下流タスクを解決するために、進化する作業記憶を持つ推論モジュールによって反復処理される。圧縮器と推論器はエンドツーエンドの強化学習によって共同最適化され、ゲーティングモジュールは分類器として別個に訓練される。実験結果から、提案手法はRULER-HQAなどのマルチホップ推論ベンチマークで競争力のある精度を達成し、文脈長を7Kトークンから1.75Mトークンへ外挿可能であり、強力な長文脈ベースラインと比較して有利な精度と効率のトレードオフを提供することを示す。特に、MemAgentと比較して、ピークGPUメモリ使用量を最大2倍削減し、推論速度を6倍向上させる。

English

Large Language Models (LLMs) face significant challenges in long-context processing, including quadratic computational costs, information forgetting, and the context fragmentation inherent in retrieval-augmented generation (RAG). We propose a cognitively inspired framework for efficient long-context inference based on chunk-wise compression and selective memory recall, rather than processing all raw tokens. The framework segments long inputs into chunks and encodes each chunk into compressed memory representations using a learned compressor. A gating module dynamically selects relevant memory blocks, which are then iteratively processed by a reasoning module with an evolving working memory to solve downstream tasks. The compressor and reasoner are jointly optimized via end-to-end reinforcement learning, while the gating module is trained separately as a classifier. Experimental results show that the proposed method achieves competitive accuracy on multi-hop reasoning benchmarks such as RULER-HQA, extrapolates context length from 7K to 1.75M tokens, and offers a favorable accuracy-efficiency trade-off compared to strong long-context baselines. In particular, it achieves up to a 2 times reduction in peak GPU memory usage and a 6 times inference speedup over MemAgent.

エンドツーエンド強化学習による圧縮メモリを活用した動的長文脈推論

Dynamic Long Context Reasoning over Compressed Memory via End-to-End Reinforcement Learning

要旨

Support