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

Paper: 2602.08382 Authors: Zhuoen Chen, Dongfang Li, Meishan Zhang, Baotian Hu, Min Zhang Categories: cs.CL, cs.LG

Abstract

This paper addresses the computational challenges of reasoning over extended contexts in large language models. The proposed framework segments inputs into chunks, converts each chunk to compressed memory representations via a learned compressor, and employs a gating mechanism to selectively retrieve relevant memory blocks. A reasoning module processes these iteratively with an evolving working memory. Both the compressor and reasoner are jointly trained via end-to-end reinforcement learning. The system achieves competitive performance on multi-hop reasoning tasks, demonstrates context extrapolation from 7K to 1.75M tokens, provides a 2x reduction in peak GPU memory, and delivers 6x inference speedup over MemAgent.

Key Contributions

• Learned memory compression: A trainable compressor that distills long chunks into compact memory representations
• Gated selective retrieval: A gating mechanism that retrieves only the most relevant memory blocks for each reasoning step
• Iterative working memory: A reasoning module that maintains and updates working memory across multiple processing steps
• End-to-end RL training: Joint optimization of compressor and reasoner through reinforcement learning
• Extreme context extrapolation: Generalizes from 7K training context to 1.75M tokens at inference

Background: Long Context Challenges

Processing long contexts in transformers faces two fundamental bottlenecks:

1. Quadratic attention cost: Self-attention scales as $O(n^2)$ with sequence length $n$
2. Linear memory cost: KV cache grows linearly, consuming GPU memory

Existing approaches include:

• Sparse attention: Reduces computation but may miss critical long-range dependencies
• Retrieval-augmented generation (RAG): External retrieval but disconnected from the model’s reasoning
• Memory-augmented models: Fixed compression schemes that don’t adapt to the reasoning task

Architecture

Memory Compressor

The input sequence is divided into chunks $\{c_1, c_2, \ldots, c_N\}$. Each chunk is processed by the compressor to produce a compact memory representation:

$m_i = \text{Compress}_\phi(c_i) \in \mathbb{R}^{k \times d}$

where $k \ll |c_i|$ is the compressed memory size and $d$ is the hidden dimension. The compressor learns to retain the most task-relevant informach chunk.

Gated Memory Retrieval

Not all memory blocks are relevant to every reasoning step. A gating mechanism computes relevance scores:

$g_i = \sigma(W_g [q; m_i])$

where $q$ is the current reasoning query and $m_i$ is the $i$-th memory block. Only memory blocks with $g_i > \tau$ are retrieved, dramatically reducing the effective context size.

Iterative Reasoning Module

The reasoner maintains a working memory $w$ that evolves across reasoning steps:

$w^{(t+1)} = \text{Reason}_\psi(w^{(t)}, \{m_i : g_i^{(t)} > \tau\})$

At each step, the reasoner:

1. Queries the compressed memories via the gating mechanism
2. Integrates retrieved information into working memory
3. Decides whether to continue reasoning or produce an output

End-to-End RL Training

The entire pipeline — compressor, gate, and reasoner — is trained jointly using reinforcement learning. The reward signal comes from task performance (e.g., correctness on multi-hop QA):

$\mathcal{J}(\phi, \psi) = \mathbb{E}\left[R(\text{answer}, \text{ground\_truth})\right]$

RL training is essential because the discrete gating decisions and iterative reasoning steps are not differentiable through standard backpropagation.

Context Extrapolation

A remarkable property of this framework is its ability to extrapolate fond training context lengths. Trained on sequences of only 7K tokens, the model successfully processes contexts up to 1.75M tokens — a 250x extrapolation factor. This is possible because:

• The compressor operates on fixed-size chunks regardless of total length
• The gating mechanism scales sublinearly with the number of memory blocks
• The iterative reasoner can perform as many steps as needed

Results

Metric	Value
Context Extrapolation	7K → 1.75M tokens (250x)
Peak GPU Memory	2x reduction
Inference Speed	6x faster than MemAgent
Multi-hop Reasoning	Compe with full-context baselines

Significance

1. Practical long-context processing: Enables reasoning over million-token contexts on standard hardware
2. RL for memory management: Demonstrates that RL can effectively train memory compression and retrieval jointly
3. Extreme generalization: 250x context extrapolation suggests the learned compression captures fundamental information structure
4. Efficiency gains: 2x memory and 6x speed improvements make long-context reasoning practical for deployment
5. Modular design: The compressor, gate, and reasoner can potentially be improved independently