[논문리뷰] When to Memorize and When to Stop: Gated Recurrent Memory for Long-Context Reasoning
ICML 2026. [Paper] [Page]
Leheng Sheng, Yongtao Zhang, Wenchang Ma, Yaorui Shi, Ting Huang, Xiang Wang, An Zhang, Ke Shen, Tat-Seng Chua
Bytedance Seed | National University of Singapore | University of Science and Technology of China
11 Feb 2026
Introduction
매우 긴 컨텍스트에 대한 추론은 책 전체를 읽거나 에이전트 시스템에서 대규모 메모리를 처리하는 것과 같은 실제 응용 분야에서 LLM의 핵심 기능이다. 이 기능은 중요하지만 여전히 어려운 과제이다. LLM은 일반적으로 컨텍스트 길이가 길어질수록 성능이 급격히 저하되며, 최대 context window를 초과하는 컨텍스트를 처리하는 데 어려움을 겪는다.
이러한 한계를 해결하기 위해, 최근 MemAgent에서는 RNN과 유사한 방식으로 긴 컨텍스트 추론을 위한 순환 메모리 패러다임을 탐구했다. 위 그림에서 볼 수 있듯이, MemAgent는 전체 컨텍스트를 한 번의 forward pass로 인코딩하는 대신, 긴 컨텍스트 추론을 순차적이고 순환적인 청크 단위 메모리화 프로세스로 재구성하였다. 구체적으로, 메모리 에이전트가 각 청크에 대하여 순차적으로 메모리를 업데이트하면, 답변 에이전트는 최종 메모리를 기반으로 답변을 예측한다. 두 에이전트는 같은 policy 모델을 공유하며, 프롬프트에 따라 동작이 구분된다.
하지만 MemAgent는 단순한 RNN 스타일의 메모리 업데이트에서 발생하는 몇 가지 단점을 그대로 물려받았다. 본 논문에서는 실제 추론 과정에서 메모리 안정성과 효율성을 저해할 수 있는 두 가지 주요 한계를 파악했다.
- 메모리 폭발 위험: 증거가 없는 청크를 업데이트할 때 메모리 에이전트는 시간이 지남에 따라 관련성이 없거나 노이즈가 있는 콘텐츠를 축적할 수 있다. 이러한 드리프트는 메모리를 점진적으로 증가시켜 메모리 폭발을 유발할 수 있다. 일단 폭발이 발생하면 누적된 노이즈는 후속 업데이트를 더욱 방해하여 나중에 나오는 청크에서 새로운 핵심 증거를 통합하기 어렵게 만든다. 또한, 각 단계에서 이미 폭발한 메모리를 비효율적으로 재생성하는 것은 inference 비용을 증가시킨다.
- 종료 메커니즘 부족: 기본 워크플로우는 모든 청크를 처리하도록 하드코딩되어 있으며 수집된 증거가 충분할 때 조기 종료 메커니즘을 제공하지 않는다. 즉, 충분한 증거가 이미 수집된 후에도 모델은 나머지 청크를 끝까지 처리해야 하므로 불필요한 계산이 발생한다. 증거가 불균등하게 분포되어 있을 때 이러한 비효율성은 더욱 심화된다.
이를 위해 본 논문에서는 RNN에서 긴 시퀀스 모델링 문제를 해결하기 위한 GRU의 gating 효과에서 영감을 받은 GRU-Mem을 제안하였다. 핵심 아이디어는 두 개의 텍스트 제어 게이트, 즉 update gate와 exit gate를 순환 워크플로우에 추가하는 것이다. Update gate는 현재 청크에 대한 메모리를 업데이트해야 하는지 여부를 결정하고, exit gate는 충분한 증거가 수집되면 모델이 조기에 종료될 수 있는지 여부를 결정한다. 기존 설계와 유사하게 답변 에이전트는 최종 메모리를 기반으로 답변을 제공한다.
이러한 기능을 모델에 부여하기 위해, 저자들은 end-to-end RL 내에 두 가지 reward 신호, 즉 update reward와 exit reward를 도입하여 각각 올바른 업데이트 및 종료 동작에 reward를 제공하였다. Update gate는 유용한 정보가 포함된 몇몇 청크에 대해서만 선택적으로 메모리를 업데이트하여 메모리 폭발을 완화하고, Exit gate는 불필요한 계산을 줄이는 유연한 조기 종료 메커니즘을 제공한다. 이 두 가지 게이트를 통해 더욱 안정적이고 효율적인 긴 컨텍스트 추론이 가능해진다.
Method
1. Gated Recurrent Memory Workflow
메모리 폭발 위험과 종료 메커니즘 부족 문제를 해결하기 위해, 순환 워크플로우에 update gate와 exit gate라는 두 가지 gating 메커니즘을 도입하였다.
구체적으로, GRU-Mem은 답변 에이전트 \(\psi_\theta\)는 변경하지 않고 메모리 에이전트 \(\phi_\theta\)만 변경하여, 기존 \(\phi_\theta\)에 gate 제어를 위한 두 가지 추가 동작을 더하였다. GRU-Mem은 후보 메모리 \(\hat{\mathcal{M}}_t\)와 두 개의 binary gating 신호 \(\mathcal{U}_t\), \(\mathcal{E}_t\)를 출력한다. \(\mathcal{U}_t\)는 메모리 업데이트 여부를, \(\mathcal{E}_t\)는 루프 종료 여부를 나타낸다.
\[\begin{equation} \mathcal{U}_t, \hat{\mathcal{M}}_t, \mathcal{E}_t = \phi_\theta (\mathcal{Q}, \mathcal{C}_t, \mathcal{M}_{t-1}) \end{equation}\]
이를 위해 메모리 에이전트 \(\phi_\theta\)는 구조화된 출력 형식을 따른다.
<think>와</think>로 둘러싸인 중간 추론 과정을 생성한다.<check>와</check>사이에 업데이트 결정을 내린다. “yes” (\(\mathcal{U}_t\) == True)는 메모리 업데이트를 트리거하고 “no” (\(\mathcal{U}_t\) == False)는 업데이트를 건너뛴다.<update>와</update>사이에 후보 메모리 \(\hat{\mathcal{M}}_t\)를 출력한다. \(\mathcal{U}_t\) == True이면 메모리는 후보 메모리로 업데이트된다. \(\mathcal{U}_t\) == False이면 이전 메모리 \(\mathcal{M}_{t-1}\)이 채택되고 후보 메모리 \(\hat{\mathcal{M}}_t\)는 폐기된다.<next>와</next>사이에서 증거 수집을 계속할지 여부를 결정한다. “continue”은 순환 루프를 계속하는 것을 의미하고 “end”는 루프를 종료하는 것을 의미한다. 종료되면 최종 메모리 \(\mathcal{M}_t\)는 질문에 답하기 위해 즉시 답변 에이전트 \(\psi_\theta\)로 전송된다.
2. Workflow Optimization with End-to-End RL
메모리 에이전트에게 update gate와 exit gate를 언제 정확하게 활성화해야 하는지 학습시키기 위해, 최종 답변의 정확성뿐만 아니라 gate 상태 생성 동작의 정확성에 대해서도 명시적으로 reward를 제공한다.
Reward Design
Outcome reward. MemAgent와 동일한 reward를 사용한다. 그룹 $g$ 내의 전체 궤적에 있는 모든 대화 \(\{o_{g,t}\}_{t=1}^{T_g}\)에 대해 동일한 outcome reward \(r_g^\textrm{outcome}\)를 할당한다.
\[\begin{equation} r^\textrm{outcome} = \mathbb{I} (\textrm{is_equiv} (\mathcal{A}, \hat{\mathcal{A}})) \end{equation}\]Update reward. Update gate를 학습시키기 위해, 각 step $t$ 대화에서 올바른 update gate 상태 \(\mathcal{U}_t\)를 생성할 때 reward를 제공한다. 구체적으로, 증거가 포함된 청크의 경우, 메모리 에이전트는 <check>yes</check>를 생성할 때 reward를 받고, 증거가 없는 청크의 경우, <check>no</check>를 생성할 때 보상을 받는다.
Exit reward. Exit gate를 학습시키기 위해, 질문에 답하는 데 필요한 마지막 증거가 포함된 턴 \(t_\textrm{last evidence}\)에서 <next>end</next>를 생성할 때 전체 경로에 reward를 제공한다. 그렇지 않은 경우에는 경로 내의 모든 대화가 불이익을 받는다.
(\(t_\textrm{exit}\)은 메모리 에이전트가 워크플로우를 종료하기로 결정한 턴)
올바른 exit gate 상태 생성에 대해서는 페널티가 없으며, 증거 부족으로 인한 조기 종료가 늦은 종료보다 더 큰 페널티를 받는다.
Format reward. 메모리 에이전트의 생성 결과 $o_{g,t}$가 정확하게 파싱될 수 있도록 추가적인 format reward \(r^\textrm{format}\)을 도입한다. 구체적으로, <think></think>, <check></check>, <update></update>, <next></next>의 시퀀스를 만족하는지 확인한다. 또한, <check>와 </check> 사이에는 “yes” 또는 “no”여야 하고, <next>와 </next> 사이에는 “continue” 또는 “end”여야 한다. 모든 생성 결과 \(\{o_{g,t}\}_{t=1}^{T_g}\)가 형식 조건을 만족할 때만 reward 1을 부여하고, 그렇지 않을 때는 0을 부여하는 엄격한 reward를 적용한다. 이는 잘못된 형식이 이전의 파싱 오류로 인한 것인지 추론할 수 없기 때문이다.
Trajectory-level reward. 그룹 $g$에 속한 모든 생성 출력 \(\{o_{g,t}\}_{t=1}^{T_g}\)은 동일한 outcome reward, exit reward, format reward를 공유하므로 다음과 같이 하나의 trajectory-level reward \(r_g^\textrm{traj}\)로 결합할 수 있다.
\[\begin{equation} r_g^\textrm{traj} = r_g^\textrm{outcome} + r_g^\textrm{exit} + r_g^\textrm{format} \end{equation}\]Advantage Calculation
Trajectory-level advantage \(\hat{A}_{g,t,i}^\textrm{traj}\)와 turn-level advantage \(\hat{A}_{g,t,i}^\textrm{turn}\)를 각각 계산한다. 그 후, 이들을 hyperparameter $\alpha$로 결합하여 보다 유연한 제어를 구현한다.
\[\begin{equation} \hat{A}_{g,t,i}^\textrm{traj} = r_g^\textrm{traj} - \frac{1}{G} \sum_{g=1}^G r_g^\textrm{traj}, \quad \hat{A}_{g,t,i}^\textrm{turn} = r_{g,t}^\textrm{update} - \frac{1}{G_t} \sum_{g=1}^{G_t} r_{g,t}^\textrm{update} \\ \hat{A}_{g,t,i} = \alpha \hat{A}_{g,t,i}^\textrm{traj} + (1 - \alpha) \hat{A}_{g,t,i}^\textrm{turn} \end{equation}\]\(\hat{A}_{g,t,i}^\textrm{traj}\)는 서로 다른 그룹의 궤적 간에 계산된다. 또한, \(\hat{A}_{g,t,i}^\textrm{turn}\)는 서로 다른 그룹의 step $t$에서의 턴 간에 계산된다. \(\hat{A}_{g,t,i}^\textrm{turn}\) 계산에서, 워크플로우는 exit gate의 제어 하에 서로 다른 step $T_g$에서 종료될 수 있으므로, step $t$에서의 그룹 크기 $G_t$는 궤적 그룹 크기 $G$와 다를 수 있다. 따라서, 이러한 조기 종료 그룹은 다른 그룹보다 더 작은 $T_g$를 갖는다.
Policy 모델은 계산된 advantage \(\hat{A}_{g,t,i}\)를 기반으로 하는 loss $\mathcal{J}(\theta)$를 사용하여 최적화된다.
3. Inference
본 논문에서는 항상 exit gate의 제어 하에 학습을 진행한다고 가정한다. 즉, 학습 중 \(\mathcal{E}_t\)가 True가 되면 워크플로는 즉시 종료되어 질문에 대한 답변을 생성한다. 그러나 일부 질문의 경우, 전체 맥락을 파악해야 답변이 가능하기 때문에 증거가 충분한지 판단하기 어렵다.
따라서 inference의 유연성을 위해 exit gate를 사용하는 inference 전략(w EG)과 사용하지 않는 전략(w/o EG) 두 가지를 제공한다. w/o EG 모드에서는 \(\mathcal{E}_t\)가 True가 되더라도 생성 워크플로우는 종료되지 않는다.
Experiments
1. Performance and Efficiency Comparison
다음은 다양한 long-context task에 대한 성능 비교 결과이다.
다음은 다양한 컨텍스트 길이에 대한 성능을 비교한 결과이다. (MV task)
2. Study of Gating Mechanisms
다음은 MV task에서 메모리 크기의 변화를 MemAgent와 비교한 그래프이다. (컨텍스트 길이는 512K)
다음은 증거가 상위 20% 위치에서 나타날 때의 성능을 비교한 결과이다. Exit gate 덕분에 inference 시간이 1/4로 줄어들었다.
다음은 컨텍스트 길이에 따른 조기 종료, 정확한 종료, 늦은 종료의 비율을 측정한 결과이다.
3. Ablation Study
다음은 다양한 $\alpha$에 대한 training dynamics를 비교한 결과이다.
다음은 RL 학습에 대한 ablation 결과이다.
