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

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Sheng, Leheng, Zhang, Yongtao, Ma, Wenchang, Shi, Yaorui, Huang, Ting, Wang, Xiang, Zhang, An, Shen, Ke, Chua, Tat-Seng
Format: Preprint
Publié: 2026
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866917267017564160
author Sheng, Leheng
Zhang, Yongtao
Ma, Wenchang
Shi, Yaorui
Huang, Ting
Wang, Xiang
Zhang, An
Shen, Ke
Chua, Tat-Seng
author_facet Sheng, Leheng
Zhang, Yongtao
Ma, Wenchang
Shi, Yaorui
Huang, Ting
Wang, Xiang
Zhang, An
Shen, Ke
Chua, Tat-Seng
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2602_10560
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle When to Memorize and When to Stop: Gated Recurrent Memory for Long-Context Reasoning
Sheng, Leheng
Zhang, Yongtao
Ma, Wenchang
Shi, Yaorui
Huang, Ting
Wang, Xiang
Zhang, An
Shen, Ke
Chua, Tat-Seng
Computation and Language
Artificial Intelligence
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.
title When to Memorize and When to Stop: Gated Recurrent Memory for Long-Context Reasoning
topic Computation and Language
Artificial Intelligence
url https://arxiv.org/abs/2602.10560