PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2602.08382v1cs.CLcs.AI

TL;DR

Long-context processing is limited by computational cost, forgetting, and retrieval fragmentation. LycheeMemory compresses chunked documents into memory, selectively recalls blocks for iterative reasoning, and jointly trains compression and reasoning with reinforcement learning. It reaches ultra-long context scales while retaining competitive accuracy and improving efficiency relative to strong baselines.

  • Problem

    Long-context LLM processing faces quadratic costs, information forgetting, and context fragmentation in retrieval-augmented generation.

  • Method

    LycheeMemory compresses document chunks into KV-cache representations, uses a Gate for selective recall, and jointly optimizes the Compressor and Reasoner through end-to-end RL.

  • Results

    2× lower peak GPU memory and 6× faster inference than MemAgent are reported, while the method scales context length from 7K to 1.75M tokens with competitive accuracy.

  • Takeaways & Limitations

    Compressed memory with selective, state-dependent retrieval provides a favorable accuracy–efficiency trade-off for ultra-long-context reasoning.

Abstract

from arXiv · show

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.

1. Introduction

LycheeMemory addresses long-context inefficiency and fragmentation by compressing documents into memory and dynamically reasoning over selected representations. It jointly trains compression and reasoning while maintaining competitive accuracy and scaling to 1.75M tokens.

  • Long-context methods face computational, forgetting, and context-fragmentation challenges, especially in extremely long sequences and multi-hop retrieval.
  • LycheeMemory divides labor between a compressed memory bank and dynamic working memory for efficient long-context reasoning.Documents are chunked into compressed KV-cache representations, then processed through dynamic recall and reasoning.
  • The framework combines a Compressor, Gate, and Reasoner, with end-to-end RL for the Compressor and Reasoner and separate classifier training for the Gate.
  • 1.75M tokens is the reported context-length scale, while accuracy remains competitive on multi-hop reasoning benchmarks.

2 Related Work

LycheeMemory bridges explicit and implicit memory approaches by storing compressed KV-cache representations while performing state-dependent retrieval and reasoning through a plaintext working memory.

  • Explicit memory methods improve memory management but standard RAG can suffer from context fragmentation and limited multi-hop precision.
  • Implicit memory methods compress or route internal representations to reduce the cost of processing long contexts.
  • LycheeMemory combines compressed KV-cache storage with state-dependent retrieval and reasoning through a plaintext working memory.This design preserves compression scalability while providing an interpretable trace over selected evidence chunks.

3 Methodology

LycheeMemory compresses chunked documents into a memory bank, selectively recalls blocks with a state-aware Gate, and iteratively reasons with working memory. End-to-end RL jointly optimizes compression and reasoning for downstream inference.

  • Architecture: The architecture assigns compression, relevance filtering, and complex reasoning to separate Compressor, Gate, and Reasoner roles.
  • Memory Construction: Documents are segmented into K sequential chunks and encoded into compact latent representations stored in the compressed memory bank Θ.The experiments use chunk size sz = 4096.
  • Dynamic Recall and Reasoning: The dynamic workflow scans memory blocks while iteratively updating plaintext working memory from an initially empty state.
  • Memory Construction: Memory tokens are interleaved with chunk tokens so each token captures preceding content, producing a compact KV-cache-style representation.The Compressor maps each chunk Ci to a latent representation θi.
  • Compression Training: The compression module is pretrained with reconstruction, QA, and creative-generation objectives before end-to-end optimization.
  • Dynamic Recall and Reasoning: The Gate filters memory blocks using relevance probabilities conditioned on the query and evolving working memory, updating memory only when P > τ.
  • End-to-End RL Optimization: End-to-end RL treats the Compressor and Reasoner as one joint policy so final-answer rewards optimize compressed memory for downstream reasoning.
  • Complexity and Efficiency: The Gate lowers linear-scan cost by using scalar classification to filter irrelevant blocks before expensive reasoning updates.Its per-block cost is far below reasoning cost, denoted Cgate ≪ Creason.

4. Experiments

LycheeMemory is evaluated for long-context QA, length extrapolation, zero-shot generalization, inference efficiency, retrieval, and compression choices. Across these experiments, it maintains competitive accuracy while improving scalability and efficiency through compressed memory and dynamic gating.

  • Main Results: 7K–1.75M tokens are evaluated for length extrapolation, with LycheeMemory showing only a mild performance drop as context length increases.The comparison spans standard contexts through an ultra-long 1.75M-token setting.
  • Main Results: LycheeMemory with Gate trades a small accuracy drop for substantially improved inference efficiency relative to MemAgent and the w/o Gate ablation.The w/o Gate variant achieves higher accuracy across most evaluated context lengths.
  • Inference Efficiency Analysis: The full-context baseline has quadratic O(N^2) latency growth and fails with OOM at 128K, whereas compressed-memory methods avoid this scaling behavior.MemAgent and LycheeMemory without Gate show linear O(N) complexity; the latter is faster because it processes a shorter effective sequence.
  • Inference Efficiency Analysis: 6× speedup over MemAgent and 3.5× over the w/o Gate baseline are achieved at 128K, while the closest reported 112K accuracy drop is 6%.The Gate keeps inference time nearly flat as context grows from 8K to 128K.
  • Different Compression Ratios: 4× compression preserves over 80% accuracy at 128K, matches 2× performance within a gap below 1%, and avoids the sharp degradation of 16× compression.The study adopts α = 4 as the default because it halves the memory footprint of α = 2 without statistically significant reasoning loss.
  • Ablation on Gate: 86.3% recall at 112K and 84.1% at 224K show that the query-and-working-memory-conditioned Gate outperforms static and query-only retrieval strategies.The Gate models retrieval conditioned on both the query and evolving working memory, enabling adaptive evidence discovery across reasoning steps.

5. Conclusion

LycheeMemory combines compressed memory, dynamic working memory, and end-to-end reinforcement learning to support efficient long-context reasoning. It scales to 1.75M tokens while maintaining competitive accuracy and improving efficiency over MemAgent.

  • End-to-End RL achieves the best reported performance, with 68.75 Avg. sub-EM from joint optimization that reaches the compressor.
  • LycheeMemory integrates a Compressor, Gate, and Reasoner for iterative reasoning over a compressed memory bank.
  • LycheeMemory scales context length to 1.75M tokens while maintaining competitive multi-hop reasoning accuracy.
  • Compared to MemAgent, LycheeMemory reduces peak GPU memory by 2× and provides a 6× inference speedup.

A.4 Evaluation and Baselines

The evaluation uses normalized sub-EM on three long-context question-answering benchmarks and compares LycheeMemory with strong retrieval and memory baselines. The section also examines GPU storage constraints and mitigation strategies for very long contexts.

  • Evaluation: Evaluation covers RULER-HQA, 2WikiMultihopQA, and StreamingQA using normalized sub-EM.
  • Evaluation: RULER-HQA test contexts range from 7k to 1.75M tokens with randomized evidence positions and distractor documents.
  • Baselines: The baselines include a standard RAG Agent and Search-R1, both retrieving the top-8 relevant chunks with text-embedding-3-large.
  • Storage and Computation Trade-off: At 1.75M tokens, compressed KV-cache storage still consumes substantial VRAM, limiting activation space and larger batch sizes.
  • Storage and Computation Trade-off: For contexts beyond 1M tokens, optional JIT compression reduces storage overhead by regenerating latent representations when needed.

B Training Convergence Analysis

Joint optimization produces a more capable compressor despite higher early reward variance, while error analysis identifies streaming and reasoning-order limitations. Later evidence can be lost when the system cannot revisit earlier chunks.

  • B Training Convergence Analysis: End-to-End RL initially has higher reward volatility because updates propagate through reasoning into the compression module.
  • B Training Convergence Analysis: The Frozen Compressor converges rapidly near 0.67 but plateaus because the reasoner remains constrained by a static memory bank.
  • B Training Convergence Analysis: Despite RL variance, joint policy optimization converges with an upward reward trend as the compressor learns task-critical features.
  • C.1 Unidirectional Dependency Mismatch (35%): Approximately 35% of errors arise from single-pass streaming when early evidence becomes relevant only after later information appears.
  • C.1 Unidirectional Dependency Mismatch (35%): The Gate may discard an initially irrelevant clue before a later bridge entity makes that clue relevant, and the system cannot backtrack.
  • C.2 Premature Inference Anchoring (21%): Approximately 21% of errors involve premature anchoring on partial evidence, which creates confirmation bias in working memory.

C.4 Other Error Types (27%)

Additional error modes include working-memory overflow and instruction misalignment, while the computational analysis contrasts full-context processing, per-chunk memory updates, and sparse gated retrieval.

  • C.4 Other Error Types (27%): Context overflow can fill the working-memory limit and flush out earlier correct evidence.
  • C.4 Other Error Types (27%): Instruction misalignment occurs when the model retrieves the right evidence but produces an answer in the wrong format.
  • Computational Complexity: The complexity comparison distinguishes full-context processing, MemAgent’s autoregressive update for every chunk, and LycheeMemory’s sparse gated retrieval.

D.1 FLOPs Formulation

The analysis compares full-context, MemAgent, and LycheeMemory scaling, showing how compressed processing and selective reasoning reduce the cost of long-context inference. LycheeMemory retains linear scaling while lowering practical FLOPs through compression and sparse generation.

  • D.1 FLOPs Formulation: Full-context processing scales quadratically with document length because global self-attention operates over the entire sequence.As context reaches millions of tokens, this quadratic term becomes prohibitive.
  • D.1 FLOPs Formulation: MemAgent is asymptotically linear but performs full generation for every chunk, producing a steep computational-cost increase.Its per-step generation includes KV-cache reading and autoregressive writing.
  • D.1 FLOPs Formulation: LycheeMemory uses 4× compression and parallel chunk processing to encode KV-cache representations without autoregressive decoding.Compression cost is proportional to input tokens, while the compressed sequence reduces the effective length seen by the Gate.
  • D.1 FLOPs Formulation: The Gate scores all K compressed blocks in one forward pass per block, and the Reasoner processes only the top-T relevant chunks where T ≪ K.The total cost is decomposed as COurs = Ccomp + Cgate + Creason.
  • D.1 FLOPs Formulation: MemAgent has cost N/sz × Cgen, whereas LycheeMemory has O(N) + T × Cgen through sparse generation.The comparison separates linear compression and gating from expensive generation calls.
  • D.1 FLOPs Formulation: When T ≪ N/sz, LycheeMemory reduces expensive generation calls and produces a much flatter scaling curve.The analysis gives retrieving only 10% of chunks as an example condition.
  • D.2 Quantitative Comparison: Across increasing context lengths, LycheeMemory and MemAgent scale linearly, while LycheeMemory consistently incurs lower FLOPs than MemAgent.The advantage is attributed to 4× compression and adaptive gating that activates the Reasoner for only a small fraction of chunks.
  • D.2 Quantitative Comparison: The supplied comparison tables describe LycheeMemory’s competitive LongBench summarization performance and show that a 1024-token working-memory capacity performs best.These results are listed as contextual comparisons alongside the FLOPs analysis.

E Out-of-Distribution (OOD) Generalization Analysis

The paper tests whether LycheeMemory’s compressed memory transfers from multi-hop QA to long-document summarization. It reports strong zero-shot generalization across both targeted extraction and global aggregation settings.

  • E. Out-of-Distribution (OOD) Generalization Analysis: LycheeMemory is evaluated zero-shot on GovReport and MultiNews summarization tasks despite being optimized primarily for multi-hop QA.The evaluation reports ROUGE-1, ROUGE-2, and ROUGE-L scores.
  • E. Out-of-Distribution (OOD) Generalization Analysis: The results show competitive summarization performance and suggest that the compressed memory bank captures transferable semantic features.The reported transfer spans QA-style targeted extraction and summarization-style global aggregation.

F Ablation on Working Memory Capacity

The ablation examines working-memory capacity and compression-related failure modes. Smaller active memory can improve reasoning precision, whereas aggressive compression can merge fine-grained attributes.

  • F Ablation on Working Memory Capacity: Working-memory capacity is varied from 1024 to 4096 tokens on RULER-HQA across 7K–56K-token contexts.Documents use 4096-token chunks, and the Gate is excluded to isolate capacity effects.
  • F Ablation on Working Memory Capacity: 79.49% average accuracy is achieved with LWM = 1024, the highest result among the tested capacities.Increasing the memory budget does not improve performance in this ablation.
  • F Ablation on Working Memory Capacity: Larger working-memory buffers can introduce noise tokens that distract attention during final answer synthesis.The reported effect is degradation in reasoning precision rather than improvement.
  • G Dynamic Evolution of Working Memory: Working-memory length initially grows as evidence is collected but stabilizes near 500 tokens around Turn 23.The figure describes bounded memory use rather than unbounded growth.
  • G Dynamic Evolution of Working Memory: Across 32 reasoning steps and 128K-context validation samples, memory stabilizes at approximately 503 tokens below the 1024-token limit.The analysis attributes this to retaining task-critical content rather than passively appending retrieved information.
  • Compression Failure Analysis: A synthetic narrative compares 4× and 16× compression to examine fine-grained retrieval under confounding entities.The example uses similar names and dates to stress attribute binding.
  • Compression Failure Analysis: The 16× model merges distinct birth-year attributes, indicating that aggressive compression can cause feature collapse and accumulated hallucinations.The 4× model retains precision in the associated table summary.

H.2 Associative Reasoning and Self-Correction

The case study illustrates associative reasoning with iterative self-correction in working memory. The model first follows a misleading association, then updates its answer when a later chunk provides the city’s exact nickname.

  • H.2 Associative Reasoning and Self-Correction: The model temporarily infers “Sparkling Pearl Embedded in a Highland” as the city’s nickname from proximal Dian Lake descriptions.This is an explicitly tentative inference during the intermediate reasoning step.
  • H.2 Associative Reasoning and Self-Correction: A later chunk states that Kunming is known as the “Spring City,” supplying the exact nickname needed to resolve the ambiguity.The updated working memory includes both the lake nickname and Kunming’s city nickname.
  • H.2 Associative Reasoning and Self-Correction: An initial association links the Yunnan Provincial Museum to Kunming, while nearby lake information introduces a competing nickname hypothesis.The working-memory states preserve both the museum location and lake-related distractors.
  • H.2 Associative Reasoning and Self-Correction: The reasoning process differentiates the lake’s nickname from the city’s alias and corrects the earlier working-memory hypothesis.The case demonstrates overriding a tentative inference with a verified fact.
  • H.2 Associative Reasoning and Self-Correction: The case connects state-dependent retrieval with correction of late-hop evidence that query-only retrieval may miss.The broader comparison frames LycheeMemory as an alternative paradigm rather than a replacement for RAG.
Loading 2602.08382v1…