Raisonnement Dynamique sur Contexte Long via une Mémoire Comprimée et un Apprentissage par Renforcement End-to-End

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

February 9, 2026

papers.authors: Zhuoen Chen, Dongfang Li, Meishan Zhang, Baotian Hu, Min Zhang

cs.AI

papers.abstract

Les modèles de langage de grande taille (LLM) rencontrent des défis significatifs dans le traitement de contextes longs, incluant des coûts computationnels quadratiques, l'oubli d'informations et la fragmentation contextuelle inhérente à la génération augmentée par récupération (RAG). Nous proposons un cadre cognitivement inspiré pour l'inférence efficace en contexte long, basé sur la compression par segments et la récupération sélective en mémoire, plutôt que sur le traitement de tous les tokens bruts. Le framework segmente les entrées longues en segments et encode chaque segment en représentations mémoire compressées à l'aide d'un compresseur appris. Un module de gating sélectionne dynamiquement des blocs de mémoire pertinents, qui sont ensuite traités itérativement par un module de raisonnement doté d'une mémoire de travail évolutive pour résoudre des tâches en aval. Le compresseur et le raisonneur sont optimisés conjointement via un apprentissage par renforcement de bout en bout, tandis que le module de gating est entraîné séparément comme un classifieur. Les résultats expérimentaux montrent que la méthode proposée atteint une précision compétitive sur des benchmarks de raisonnement à sauts multiples tels que RULER-HQA, extrapole la longueur de contexte de 7K à 1,75M de tokens, et offre un compromis précision-efficacité favorable par rapport à des bases de référence solides en contexte long. En particulier, elle permet jusqu'à une réduction par 2 de l'utilisation maximale de mémoire GPU et une accélération de l'inférence par 6 par rapport à MemAgent.

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.