Raciocínio Dinâmico em Contextos Longos sobre Memória Comprimida via Aprendizado por Reforço de Ponta a Ponta

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

February 9, 2026

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

cs.AI

Resumo

Os Modelos de Linguagem de Grande Porte (LLMs) enfrentam desafios significativos no processamento de contexto longo, incluindo custos computacionais quadráticos, esquecimento de informações e a fragmentação de contexto inerente à geração aumentada por recuperação (RAG). Propomos uma estrutura cognitivamente inspirada para inferência eficiente em contexto longo, baseada na compressão por blocos e na recuperação seletiva de memória, em vez de processar todos os tokens brutos. A estrutura segmenta entradas longas em blocos e codifica cada bloco em representações de memória compactadas usando um compressor aprendido. Um módulo de portão seleciona dinamicamente blocos de memória relevantes, que são então processados iterativamente por um módulo de raciocínio com uma memória de trabalho em evolução para resolver tarefas subsequentes. O compressor e o módulo de raciocínio são otimizados conjuntamente via aprendizado por reforço de ponta a ponta, enquanto o módulo de portão é treinado separadamente como um classificador. Resultados experimentais mostram que o método proposto alcança precisão competitiva em benchmarks de raciocínio multi-etapa, como o RULER-HQA, extrapola o comprimento de contexto de 7K para 1,75M de tokens e oferece um equilíbrio favorável entre precisão e eficiência em comparação com fortes baselines de contexto longo. Em particular, alcança uma redução de até 2 vezes no uso de memória GPU de pico e uma aceleração de inferência 6 vezes maior em relação ao 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.