PaperlayerView exact revision on arXiv

Source-linked AI summary

When to Memorize and When to Stop: Gated Recurrent Memory for Long-Context Reasoning

Leheng Sheng, Yongtao Zhang, Wenchang Ma, Yaorui Shi, Ting Huang, Xiang Wang, An Zhang, Ke Shen, Tat-Seng Chua

arXiv:2602.10560v1cs.CLcs.AI

TL;DR

Long-context reasoning is important but difficult because LLMs degrade with increasing context length and struggle with sparse evidence. GRU-Mem adds text-controlled update and exit gates, trained with dedicated rewards, and generally outperforms vanilla MemAgent with up to 400% times inference speed acceleration. Its evaluation is limited to QA, while the extra rewards reduce training stability and lengthen convergence.

  • Problem

    Long-context reasoning remains challenging because LLM performance degrades as context grows, especially when relevant evidence is sparsely distributed.

  • Method

    GRU-Mem adds text-controlled update and exit gates to recurrent memory, training their behaviors end-to-end with separate update and exit rewards.

  • Results

    GRU-Mem generally outperforms vanilla MemAgent across diverse long-context reasoning tasks and achieves up to 400% times inference speed acceleration.

  • Takeaways & Limitations

    Selective memory updates and early termination provide a stable and efficient paradigm for long-context reasoning within the evaluated tasks.

  • Takeaways & Limitations

    The study is limited to QA, while its extra rewards reduce training stability and require a smaller off-policy degree and longer convergence time.

Abstract

from arXiv · show

While reasoning over long context is crucial for various real-world applications, it remains challenging for large language models (LLMs) as they suffer from performance degradation as the context length grows. Recent work MemAgent has tried to tackle this by processing context chunk-by-chunk in an RNN-like loop and updating a textual memory for final answering. However, this naive recurrent memory update faces two crucial drawbacks: (i) memory can quickly explode because it can update indiscriminately, even on evidence-free chunks; and (ii) the loop lacks an exit mechanism, leading to unnecessary computation after even sufficient evidence is collected. To address these issues, we propose GRU-Mem, which incorporates two text-controlled gates for more stable and efficient long-context reasoning. Specifically, in GRU-Mem, the memory only updates when the update gate is open and the recurrent loop will exit immediately once the exit gate is open. To endow the model with such capabilities, we introduce two reward signals $r^{\text{update}}$ and $r^{\text{exit}}$ within end-to-end RL, rewarding the correct updating and exiting behaviors respectively. Experiments on various long-context reasoning tasks demonstrate the effectiveness and efficiency of GRU-Mem, which generally outperforms the vanilla MemAgent with up to 400\% times inference speed acceleration.

1 Introduction

Long-context reasoning is important but difficult because LLM performance degrades with growing context, while recurrent chunk-by-chunk memory methods risk memory explosion and unnecessary computation. GRU-Mem addresses these issues with update and exit gates trained by dedicated rewards, improving stability, efficiency, and speed.

  • LLMs experience dramatic performance degradation as context length grows and struggle with corpora exceeding their maximum context window.
  • MemAgent processes long contexts chunk-by-chunk, recurrently updating textual memory before answering from the final memory.
  • Indiscriminate updates on evidence-free chunks can accumulate irrelevant content, inflate memory beyond its budget, and increase inference cost.
  • Without early exit, MemAgent processes every remaining chunk after sufficient evidence has been collected, causing avoidable computation.
  • GRU-Mem uses text-controlled update and exit gates to selectively update memory and terminate once sufficient evidence is collected.
  • GRU-Mem generally outperforms vanilla MemAgent across diverse tasks and model sizes, with up to 400% times inference speed acceleration.

2 Preliminary

Long-context reasoning splits a large context into chunks and recurrently compresses them into memory for answering, but the vanilla workflow remains limited by indiscriminate updates and no early exit.

  • Task Formulation: Long-context QA requires locating and reasoning over a few evidence pieces sparsely distributed across very large contexts.This sparse-evidence setting is commonly associated with the needle-in-a-haystack problem.
  • Recurrent Memory Paradigm: The recurrent memory paradigm divides context C into fixed-size chunks, then processes them sequentially with a memory agent.Most chunks contain no information needed to answer the question, avoiding direct ingestion of the entire context at once.
  • Recurrent Memory Paradigm: At each step, the memory agent reads Q, C_t, and M_t−1 to generate M_t; after all chunks, the answer agent predicts  from Q and M_T.The memory and answer agents share the same parameterized policy model, differentiated by prompt assignment.
  • Workflow Optimization with End-to-End RL: End-to-end reinforcement learning optimizes each memory or answer conversation turn using its corresponding advantage within a multi-turn workflow.The loss uses policy and reference models, clipping factors, and importance-sampling weights.
  • Limitations: The vanilla workflow can enhance long-context reasoning, but indiscriminate updates and the absence of early exit limit memory stability and inference efficiency.These limitations motivate selective memory updating and termination mechanisms in GRU-Mem.

3 Methodology

GRU-Mem augments recurrent memory reasoning with update and exit gates controlled by the memory agent, while training explicitly rewards correct gate behavior. Its workflow selectively updates memory, can terminate scanning early, and then answers from the terminal memory.

  • Gated Recurrent Memory Workflow: GRU-Mem equips the recurrent memory workflow with update and exit gates to improve long-context reasoning stability and efficiency.The gates address memory explosion and the lack of an exit mechanism.
  • Gated Recurrent Memory Workflow: The memory agent outputs a candidate memory M̂_t and binary signals U_t and E_t indicating whether to update memory and terminate the loop.The answer agent remains unchanged while the memory agent receives two additional gate-control actions.
  • Gated Recurrent Memory Workflow: At each step, a positive update decision adopts M̂_t, while a negative decision retains M_t−1 and discards the candidate memory.The agent then decides whether to continue collecting evidence or end the recurrent loop.
  • Inference Procedure: When the exit gate signals termination, the final memory is sent immediately to the answer agent for answering.The algorithm optionally enables exit-gate termination after processing each chunk.
  • Training: GRU-Mem training explicitly rewards correct update-gate and exit-gate statuses in addition to final-answer correctness.The two gate-specific rewards are combined within end-to-end reinforcement learning.

GRU-Mem

GRU-Mem trains gated recurrent memory with rewards for correct updates, exits, and output formatting, then combines trajectory-level and turn-level advantages for policy optimization. Its exit reward targets the last required evidence, while inference may also require reading the full context for questions needing exhaustive evidence.

  • Reward Design: Update reward favors opening the gate on evidence-containing chunks and keeping it closed on evidence-free chunks.This trains selective memory updates rather than indiscriminate accumulation.
  • Reward Design: Exit reward is neutral at the last required evidence, penalizes early exit more than late exit, and penalizes all trajectory conversations for incorrect timing.The specified rewards are −0.75 for texit < tlast evidence, 0 for equality, and −0.5 for texit > tlast evidence.
  • Reward Design: Format reward grants a reward of 1 only when every generated turn follows the required output format; otherwise, all outputs receive 0.The checked structure includes <think>, <check>, <update>, and <next> sequences, with strict correctness across the trajectory.
  • Advantage Calculation: Trajectory-level and turn-level advantages are combined with hyperparameter α into the total advantage used to optimize the policy loss J(θ).Trajectory-level advantages compare trajectories across groups, whereas turn-level advantages compare turns at step t across groups.
  • Advantage Calculation: Because the exit gate can terminate workflows at different steps, early-exited groups may have smaller trajectory lengths and different turn-level group sizes.The turn-level group size Gt can differ from the trajectory group size G when workflows exit at different Tg.
  • Inference: Some questions require processing the whole context, so the workflow cannot always determine that evidence is sufficient for an early exit.The paper gives questions asking for all special magic numbers as an example.

4 Experiments

Experiments evaluate GRU-Mem across tasks, model sizes, context lengths, gating behaviors, and reinforcement-learning ablations. GRU-Mem generally improves performance and inference efficiency over MemAgent while reducing memory growth and enabling earlier exits.

  • 4.1 (RQ1) Performance and Efficiency Comparison: Experiments compare GRU-Mem and MemAgent across QA, NIAH, and multi-values benchmarks using Qwen2.5-3B-Instruct and Qwen2.5-7B-Instruct.Evaluation averages performance and inference time across context sizes and includes multiple inference modes.
  • 4.1 (RQ1) Performance and Efficiency Comparison: GRU-Mem generally outperforms MemAgent across diverse datasets, with especially strong results on out-of-distribution NIAH tasks and the MK task series.The advantage is particularly pronounced for the 3B backbone, where MemAgent shows a sharp MK-series performance drop.
  • 4.1 (RQ1) Performance and Efficiency Comparison: 400% faster inference is achieved in several cases with early exit, while GRU-Mem generally achieves around 200% acceleration without the exit gate.The reported acceleration occurs without harming performance in the cited early-exit cases.
  • 4.2 (RQ2) Study of Gating Mechanisms: GRU-Mem’s memory grows more slowly because its update gate changes memory only on a few critical evidence-containing chunks, whereas MemAgent reaches the 1024-token maximum.Indiscriminate updates can make memory growth hurt performance and increase inference overhead.
  • 4.2 (RQ2) Study of Gating Mechanisms: When the last evidence appears in the top 20% of documents, GRU-Mem with the exit gate reduces inference time to 1/4 of MemAgent and usually identifies the exit position.This demonstrates flexibility under unbalanced evidence occurrence.
  • 4.3 (RQ3) Ablation Study: An α of 0.9 provides a more stable validation trend than α values of 1.0 and 0.5, while exit accuracy exceeds 0.8 across all tested settings.Higher α improves evidence-present updates but increases unnecessary updates on evidence-free chunks.
  • 4.3 (RQ3) Ablation Study: Reinforcement-learning training generally improves performance across diverse tasks, with larger benefits on harder tasks such as HQA, SQuAD, and the MK series.The comparison uses the 7B model with and without RL training.

5 Limitations

The paper’s scope and training procedure impose two stated limitations: evaluation is concentrated in QA, and the extra rewards reduce training stability.

  • The study is limited to the QA domain, leaving tasks such as summarization largely underexplored.
  • Extra rewards in GRU-Mem reduce training stability, requiring a smaller off-policy degree and longer convergence time.

6 Conclusion

The conclusion frames GRU-Mem as a gated extension of recurrent long-context memory that addresses indiscriminate updates and unnecessary computation. Across diverse tasks, it outperforms vanilla MemAgent and reaches up to 400% inference speed acceleration.

  • GRU-Mem adds text-controlled update and exit gates to recurrent chunk-by-chunk memory reasoning.The update gate changes memory only when necessary, while the exit gate terminates once sufficient evidence is collected.
  • The two gates are trained end-to-end with separate rewards for correct updating and exiting behaviors.
  • Across diverse long-context reasoning tasks, GRU-Mem outperforms vanilla MemAgent while achieving up to 400% inference speed acceleration.

A.3 Reinforcement Learning with Multiple Tasks

This section motivates role-specific rewards for training one policy model to learn multiple behaviors in long-context reasoning.

  • Role-specific rewards can train a single policy model to learn both memory updating and early exiting behaviors.The motivation builds on work showing that LLMs can learn multiple objectives by rewarding different learning targets.

B Implementation Details

The implementation uses specified RL-training hyperparameters and structured prompts for GRU-Mem’s memory and answer agents. Experiments evaluate performance and inference efficiency across varying context lengths on multiple tasks.

  • Training configuration: RL training uses a 5,000-token chunk size, 8,192-token maximum prompt, and 2,048-token maximum response.The setup also uses clip ratio 0.2 and learning rate 1 × 10−6.
  • Prompting: The memory-agent and answer-agent prompts are presented separately in Figure 10a and Figure 10b.The supplied implementation passages identify these prompt locations but do not reproduce their full contents.
  • Evaluation scope: The study reports performance and efficiency across varying context lengths on HQA, SQuAD, SK-1/2/3, MK-1/2/3, MQ, and MV.Figures 11–20 provide these context-length evaluations.
  • Evaluation scope: Figure color shading encodes acceleration ratio, with deeper hues indicating higher acceleration ratios.Numerical values represent performance metrics.
  • Evaluation scope: GRU-Mem consistently achieves higher inference speed, with acceleration becoming more apparent as context length increases.This comparison is reported across the context-length figures.

D.2 Training Dynamics

Training dynamics show rapid acquisition of correct formatting, progressively improved exit behavior, and reduced response length from selective memory updates. Under top-10% evidence placement, GRU-Mem maintains performance while achieving approximately 80% correct exits and faster inference.

  • Training dynamics: Format correctness quickly reaches around 100% accuracy under all tested α settings.This occurs after only a few training steps.
  • Training dynamics: The update gate lowers average response length by restricting memory updates to evidence-present chunks.Higher α produces a sharper response-length increase, while the gate reduces unnecessary updates on evidence-free chunks.
  • Training dynamics: Exit deviation converges near zero as exact exits increase, while early and late exits decrease across training.The absolute exit deviation also decreases stably across α settings.
  • Unbalanced evidence distribution: Around 80% of exits are correct when the last evidence appears within the top 10% of documents.Under this unbalanced evidence distribution, GRU-Mem maintains the same performance while achieving faster inference.

E Case Study

Case studies illustrate how GRU-Mem updates memory when a chunk contains useful evidence and retains queried facts across subsequent reasoning. Examples include identifying an Animorphs-related answer and recalling an exact UUID.

  • Memory updating: The memory agent reasons about whether each new section contains useful information before updating its memory.The case-study section presents examples of memory-agent behavior and updating dynamics.
  • Case 1: In Case 1, the memory update connects The Hork-Bajir Chronicles with the Animorphs series and enslaved Hork-Bajir.The reasoning identifies the companion-book relationship as relevant to the question.
  • Case 2: Case 2 reports a ground-truth answer of 276,170 inhabitants for the queried city’s inhabitant count.The supplied case-study passage provides the numerical ground truth without additional reasoning details.
  • Case 3: Case 3 demonstrates exact recall of the UUID ac5f891e-ac5d-4d3b-b830-a7efe7cf807a.The prediction matches the listed ground truth.
Loading 2602.10560v1…