Dynamisch Redeneren over Lange Contexten via Gecomprimeerd Geheugen met End-to-End Versterkingsleren

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

February 9, 2026

Auteurs: Zhuoen Chen, Dongfang Li, Meishan Zhang, Baotian Hu, Min Zhang

cs.AI

Samenvatting

Grote Taalmodellen (LLM's) kampen met aanzienlijke uitdagingen bij de verwerking van lange contexten, waaronder kwadratische rekenkosten, het vergeten van informatie en de contextfragmentatie die inherent is aan retrieval-augmented generation (RAG). Wij stellen een cognitief geïnspireerd raamwerk voor voor efficiënte inferentie bij lange contexten, gebaseerd op chunkgewijze compressie en selectief geheugenrecall, in plaats van het verwerken van alle ruwe tokens. Het raamwerk segmenteert lange invoer in chunks en codeert elke chunk in gecomprimeerde geheugenrepresentaties met behulp van een aangeleerde compressor. Een gatingmodule selecteert dynamisch relevante geheugenblokken, die vervolgens iteratief worden verwerkt door een redeneermodule met een evoluerend werkgeheugen om downstreamtaken op te lossen. De compressor en de redeneermodule worden gezamenlijk geoptimaliseerd via end-to-end reinforcement learning, terwijl de gatingmodule apart wordt getraind als een classifier. Experimentele resultaten tonen aan dat de voorgestelde methode een competitieve nauwkeurigheid bereikt op multi-hop redeneerbenchmarks zoals RULER-HQA, de contextlengte extrapoleert van 7K naar 1,75M tokens, en een gunstige nauwkeurigheid-efficiëntie-afweging biedt in vergelijking met sterke lange-context-baselines. In het bijzonder bereikt het een reductie van tot 2 keer in het piek-GPU-geheugengebruik en een 6 keer snellere inferentie dan 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.