Ragionamento Dinamico su Contesti Lunghi tramite Memoria Compressa mediante Apprendimento per Rinforzo End-to-End

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

February 9, 2026

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

cs.AI

Abstract

I modelli linguistici di grandi dimensioni (LLM) affrontano sfide significative nell'elaborazione di contesti lunghi, inclusi i costi computazionali quadratici, la dimenticanza delle informazioni e la frammentazione del contesto intrinseca nella generazione aumentata dal recupero (RAG). Proponiamo un framework ispirato alla cognizione umana per l'inferenza efficiente su contesti lunghi, basato sulla compressione per segmenti e sul richiamo selettivo dalla memoria, anziché sull'elaborazione di tutti i token grezzi. Il framework suddivide gli input lunghi in segmenti (chunk) e codifica ciascun segmento in rappresentazioni compresse di memoria utilizzando un compressore appreso. Un modulo di gating seleziona dinamicamente i blocchi di memoria rilevanti, che vengono poi elaborati iterativamente da un modulo di ragionamento con una memoria di lavoro in evoluzione per risolvere compiti downstream. Il compressore e il motore di ragionamento sono ottimizzati congiuntamente tramite apprendimento per rinforzo end-to-end, mentre il modulo di gating viene addestrato separatamente come classificatore. I risultati sperimentali mostrano che il metodo proposto raggiunge un'accuratezza competitiva su benchmark di ragionamento multi-hop come RULER-HQA, estrapola la lunghezza del contesto da 7K a 1,75 milioni di token e offre un favorevole compromesso accuratezza-efficienza rispetto a solidi baseline per contesti lunghi. In particolare, raggiunge una riduzione fino a 2 volte nell'utilizzo di picco della memoria GPU e un'accelerazione dell'inferenza di 6 volte rispetto a 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.