Diffusion은 신이고 GPT는 무적이다

ICLR 2026 Oral. [Paper] [Github]
Hongli Yu, Tinghong Chen, Jiangtao Feng, Jiangjie Chen, Weinan Dai, Qiying Yu, Ya-Qin Zhang, Wei-Ying Ma, Jingjing Liu, Mingxuan Wang, Hao Zhou
Bytedance Seed | Tsinghua University
3 Jul 2025

Introduction

LLM 시스템은 여전히 긴 컨텍스트를 효과적으로 처리하는 데 어려움을 겪는다. 책 한 권을 통째로 처리하거나, 여러 단계를 거쳐 복잡한 추론 과정을 실행하거나, 에이전트 시스템의 장기 기억을 관리하는 것과 같은 복잡한 작업은 모두 방대한 양의 텍스트를 생성하여 현재 LLM 시스템의 일반적인 컨텍스트 처리 용량을 빠르게 초과할 수 있다.

따라서, 강력한 긴 컨텍스트 처리 능력을 갖춘 성공적인 LLM은 무한한 길이의 텍스트 처리, 성능 저하 없이 scaling, 선형 복잡도로 효율적인 디코딩이라는 세 가지 요소를 모두 충족해야 한다. 인간은 긴 컨텍스트 정보를 처리할 때, 핵심적인 내용을 추상화하여 텍스트 전체의 본질을 파악하는 경향이 있다. 이때 중요한 디테일들을 메모하거나 요점을 기록하고, 중복되거나 관련 없는 데이터는 버린다. 모든 사실이나 사소한 정보를 암기하려고 애쓰지 않고, 당면한 문제의 더 중요한 측면에 지적 에너지를 집중한다. 이러한 선택적 집중은 과정을 단순화할 뿐만 아니라 복잡한 문제를 더욱 효율적으로 해결하는 데 도움이 된다.

이러한 인간 중심적 직관에 따라, 본 논문은 LLM에 동적으로 업데이트되는 고정 길이 메모리를 탑재하기 위해 강화 학습(RL)을 새롭게 활용하는 MemAgent를 제안하였다. Inference 시에 LLM은 입력 텍스트를 세그먼트별로 처리한다. 각 세그먼트를 읽으면서 모델은 메모리를 능동적이고 선택적으로 업데이트하고, 모든 관련 메시지가 메모리에 통합되고 시너지 효과를 내도록 조정하여 최종 출력을 생성한다. 이 메커니즘을 통해 LLM은 메모리 길이가 고정되어 있어 모델의 context window 크기가 고정되므로 선형 시간 복잡도를 유지하면서 임의의 텍스트 길이를 유연하게 처리할 수 있다.

이러한 세그먼트 기반 접근 방식은 하나의 긴 텍스트 입력에서 여러 출력을 생성하며, 이를 위해 여러 번의 메모리 업데이트와 최종 응답 생성을 위한 마지막 업데이트가 필요하다. 여러 독립적인 컨텍스트에 걸쳐 대화를 가능하게 하는 이러한 유형의 에이전트 워크플로우를 학습하는 것은 현재 LLM 연구에서 아직 탐구되지 않은 영역이다. MemAgent는 각 컨텍스트 독립적인 대화를 최적화 대상으로 간주한다. 저자들은 DAPO 알고리즘을 기반으로, 검증 가능한 reward를 통해 임의의 에이전트 워크플로우를 최적화하는 Multi-Conv DAPO를 구현했다.

8K의 적당한 context window를 사용하고 3.2만 개의 문서로 학습된 모델은 최대 400만 토큰 규모의 문서에 대한 성능 저하 없이 선형적인 계산 비용으로 일관되게 뛰어난 성능을 보였다. 이는 본 논문의 긴 컨텍스트 메모리 접근 방식의 효율성과 scalability를 입증한다.

Method

1. The MemAgent Workflow: RL-shaped Memory for Unbounded Contexts

MemAgent는 임의로 긴 문서를 하나의 덩어리가 아니라 제어된 증거 흐름으로 간주한다. 매 step마다 모델은 정확히 두 가지를 인식한다.

  1. 다음 텍스트 청크
  2. 지금까지 중요하다고 판단된 모든 내용을 요약한 압축된 고정 길이 메모리

중요한 것은 이 메모리가 context window 내의 일반 토큰들의 시퀀스에 불과하므로 기본 LLM의 핵심 생성 프로세스는 변경되지 않는다는 점이다.

새로운 청크를 읽은 후, 모델은 이전 메모리를 업데이트된 메모리로 덮어쓴다. 이러한 덮어쓰기 전략은 매우 단순해 보이지만, 바로 이 전략 덕분에 시스템이 확장될 수 있다. 메모리 길이가 절대 늘어나지 않기 때문에 청크당 총 연산량은 $O(1)$로 유지되고, end-to-end 복잡도는 청크 수에 대해 선형적이다.

저자들은 덮어쓰기 결정을 RL 문제로 구성하였다. 에이전트는 나중에 유용할 정보를 유지하고 귀중한 토큰을 낭비하게 할 방해 요소를 제거할 때 reward를 받는다. 모델은 답변에 중요한 사실을 보존하면서 적극적으로 압축하는 방법을 학습한다.


워크플로우는 inference를 두 가지 모듈로 자연스럽게 분해한다. 컨텍스트 처리 모듈은 청크를 순회하며 프롬프트 템플릿으로 메모리를 업데이트한다. 스트림이 모두 처리되면 답변 생성 모듈이 호출되어 모델은 문제 설명과 메모리만을 참조하여 답변을 생성한다. 동일한 tokenization과 attention 레이아웃이 두 ​​모듈 모두에 적용되어 아키텍처 수정 없이 모델의 내재된 length extrapolation 기능을 활용할 수 있다.

따라서 MemAgent는 이 설계로부터 세 가지 이점을 얻는다.

  1. 무제한 길이: 문서가 스트림으로 처리되므로 수백만 개의 토큰을 가질 수 있다.
  2. 성능 저하 없음: RL은 메모리가 필요한 정보만 정확하게 유지하도록 하여 거의 손실 없는 extrapolation을 제공한다.
  3. 선형적 비용: 일정한 window 크기는 디코딩 시간과 메모리 소비가 입력 길이에 따라 선형적으로 증가함을 의미한다. 이는 적당한 컨텍스트 크기의 LLM을 최소한의 엔지니어링 오버헤드로 효율적인 긴 컨텍스트 추론 모델로 변환하는 실용적인 방법을 제공한다.

2. Training MemAgent with Multi-conv RL

본 논문에서는 RL을 통해 최적화될 policy의 일부로서, 컨텍스트 처리 과정에서 메모리 업데이트를 고려하였다. 이를 위해 RLVR 방식을 채택하여 MemAgent를 학습시킨다.

저자들은 기본 알고리즘으로 단순성과 RLVR에서의 효율성을 고려하여 GRPO를 채택하였다. GRPO의 rollout 단계에서 policy 모델은 입력 $x$에 대해 $G$개의 개별 응답 \(\{o_i\}_{i=1}^G\)을 샘플링한다. \(\{R_i\}_{i=1}^G\)를 sequence-level reward라고 할 때, $i$번째 응답의 advantage는 다음과 같이 계산된다.

\[\begin{equation} \hat{A}_{i,t} = \frac{r_i - \textrm{mean}(\{R_i\}_{i=1}^G)}{\textrm{std}(\{R_i\}_{i=1}^G)} \end{equation}\]

GRPO objective는 다음과 같다.

\[\begin{equation} \mathcal{J}_\textrm{GRPO} (\theta) = \mathbb{E}_{(q,a) \sim \mathcal{D}, \{o_i\}_{i=1}^G \sim \pi_{\theta_\textrm{old}} (\cdot \vert q)} \left[ \frac{1}{G} \sum_{i=1}^G \frac{1}{\vert o_i \vert} \sum_{t=1}^{\vert o_i \vert} \min \left( r_{i,t} (\theta) \hat{A}_{i,t}, \textrm{clip} \left( r_{i,t} (\theta), 1 - \epsilon, 1 + \epsilon \right) \hat{A}_{i,t} \right) - \beta D_\textrm{KL} (\pi_\theta \| \pi_\textrm{ref}) \right] \\ \textrm{where} \quad r_{i,t} (\theta) = \frac{\pi_\theta (o_{i,t} \mid q, o_{i,< t})}{\pi_{\theta_\textrm{old}} (o_{i,t} \mid q, o_{i,< t})} \end{equation}\]

하지만 MemAgent 접근 방식의 특성상 하나의 쿼리에 대해 여러 개의 컨텍스트에 독립적인 대화가 생성된다. 따라서 단순히 attention mask를 적용하는 것만으로는 policy 최적화를 구현할 수 없다.


이 문제를 해결하기 위해, 각 대화를 독립적인 최적화 대상으로 취급한다. 그룹 내 주어진 샘플 $(q_i, a_i)$에 대해 $n_i$개의 대화 \(\{o_{i,j}\}_{j=1}^{n_i}$가 생성된다고 하자. $o_{i,j}$는 다시 토큰 수준 출력\){o_{i,j,t}}{t=1}^{\vert o{i,j} \vert}$$로 분해된다. 최종 답변이 포함된 최종 대화를 기준으로 샘플당 reward $R_i$를 계산하고, advantage를 모든 관련 대화에 분배한다.

컨텍스트에 독립적인 다중 대화 rollout을 위한 MemAgent 알고리즘 내에서 advantage와 loss는 다음과 같이 계산된다.

\[\begin{equation} \hat{A}_{i,j,t} = r_i - \textrm{mean} (\{ R_i \}_{i=1}^G) \\ \mathcal{C}_{i,j,t} = \min \left( r_{i,j,t} (\theta) \hat{A}_{i,j,t}, \textrm{clip} \left( r_{i,j,t} (\theta), 1 - \epsilon_\textrm{low}, 1 + \epsilon_\textrm{high} \right) \hat{A}_{i,j,t} \right) \\ \mathcal{J}_\textrm{DAPO} (\theta) = \mathbb{E}_{(q,a) \sim \mathcal{D}, \{o_{i,j}\}_{i=1}^G \sim \pi_{\theta_\textrm{old}} (\cdot \vert q, o_{i,j-1})} \left[ \frac{1}{\sum_{i=1}^G \sum_{j=1}^{n_i} \vert o_{i,j} \vert} \sum_{i=1}^G \sum_{j=1}^{n_i} \sum_{t=1}^{\vert o_{i,j} \vert} \left( \mathcal{C}_{i,j,t} - \beta D_\textrm{KL} (\pi_\theta \| \pi_\textrm{ref}) \right) \right] \\ \end{equation}\]

Advantage 값은 최종 답변이 포함된 대화에서 도출된 후, 동일한 샘플에서 시작된 모든 대화에 균일하게 적용된다. Loss는 DAPO의 token-level averaging loss와 유사하다. 또한, loss 계산의 차원을 기존의 (그룹, 토큰) 구조에서 (그룹, 대화, 토큰) 구조로 확장했다. DrGRPO를 따라, advantage 값을 reward 표준 편차로 정규화하지 않는다.

3. Reward Modeling

RLVR 방식들을 따라, 규칙 기반 검증기가 계산한 최종 reward를 사용하여 모델을 학습시킨다. RULER와 같이 하나의 질문에 여러 개의 정답이 있는 경우, 이러한 정답들이 동일한 것으로 간주된다. 여러 개의 정답 집합 \(Y = \{y_1, \ldots, y_n\}\)이 주어졌을 때, reward는 다음과 같이 정의된다.

\[\begin{equation} R(\hat{y}, Y) = \max_{y \in Y} (\mathbb{I}(\textrm{is_equiv}(y, \hat{y}))) \end{equation}\]

($\hat{y}$는 예측 정답, $\mathbb{I}(\cdot)$는 indicator function)

다른 task들의 경우, 모든 정답이 최종 출력에 포함되어야 한다 (ex. Multi-Value Needle in a Haystack). 이러한 경우 reward는 다음과 같이 정의된다.

\[\begin{equation} R(\hat{y}, Y) = \frac{\vert y \in Y \mid \mathbb{I}(y \in \hat{y}) \vert}{\vert Y \vert} \end{equation}\]

4. Rethinking MemAgent from Autoregressive Modeling Perspectives

MemAgent 디자인을 autoregressive 언어 모델링 관점에서 생각해 볼 수 있다. LLM은 시퀀스 \(\textbf{x}_{1:N}\)의 joint likelihood를 다음과 같이 인수분해한다.

\[\begin{equation} p(\textbf{x}_{1:N}) = \prod_{n=1}^N p (x_n \mid \textbf{x}_{1:n-1}) \end{equation}\]

이는 모든 이전 토큰이 활성 컨텍스트에 남아 있어야 한다는 가정을 암묵적으로 전제로 한다. 바로 이 때문에 attention이 긴 컨텍스트 병목 현상을 일으키는 것이며, 전체적인 복잡도는 $O(N^2)$이다.


MemAgent는 무한한 히스토리를 고정 길이 메모리 \(\textbf{m} \in \mathbb{V}^M\)으로 대체한다. 입력 텍스트는 $K$개의 연속된 청크 \(\textbf{c}^1, \ldots, \textbf{c}^K\)로 스트리밍된다 (각 청크의 길이는 최대 $C$). 청크 $k$를 읽은 후, 모델은 지금까지 관찰된 모든 증거를 요약한 새로운 벡터 \(\textbf{m}^k\)로 패널을 덮어쓴다. $\vert \textbf{m}^k \vert = M$이 상수이므로 step별 계산 및 메모리 사용량은 모두 $O(C + M)$이며, 전체적인 복잡도는 $O(N)$이다.

시퀀스 \(\textbf{m}^{1:K-1}\)을 도입하면 원래 likelihood는 다음과 같이 분해된다.

\[\begin{equation} p(\textbf{x}_{1:N}) = \sum_{\textbf{m}^{1:K-1}} \prod_{k=1}^K \underbrace{p(\textbf{c}^k \mid \textbf{m}^{k-1})}_{\textrm{read}} \underbrace{p(\textbf{m}^k \mid \textbf{c}^k, \textbf{m}^{k-1})}_{\textrm{write}} \quad \textrm{where} \; \textbf{m}^0 = \varnothing \end{equation}\]

각 청크 내부에서는 일반적인 transformer 디코더를 실행하지만, 상수 context window \((\textbf{c}^k, \textbf{m}^{k-1})\)로 컨디셔닝된다. 읽기 경로와 쓰기 경로 모두 토큰 단위로 인수분해된다. 개념적으로, 위의 likelihood 식은 transformer를 사용자가 state 크기를 제어할 수 있는 recurrent network로 변환한다.

로컬-글로벌 모델이나 선형 어텐션 모델은 긴 컨텍스트를 feature space에서 압축하는 하기 때문에 이러한 모델의 요약은 암묵적이고 불투명하다. 반면, MemAgent는 컨텍스트를 토큰 수준으로 압축하기 때문에, MemAgent의 요약은 토큰 공간에 존재하므로 모든 중간 메모리는 사람이 읽을 수 있고 검사하거나 편집할 수도 있다. 이는 RL reward를 설계할 때 활용된 속성이다.

RL이 필수적인 이유

메모리 토큰은 discrete한 덮어쓰기 규칙을 통해 업데이트되므로, backpropagation만으로는 모델이 무엇을 유지하고 무엇을 버려야 하는지 학습할 수 없다. 본 논문의 RL 알고리즘은 각 읽기-쓰기-읽기 루프를 RL transition으로 처리하여, 최종 답변에 대한 정답으로 이어지는 메모리에 직접적인 reward를 제공한다.

Experiments

  • backbone: Qwen2.5-Instruct

1. Main Results

다음은 다양한 컨텍스트 길이에 대한 모델 성능을 비교한 결과이다.

2. Ablation Study

다음은 RL 유무에 대한 ablation 결과이다. (RULER-HotpotQA)

3. Case Study

다음은 MemAgent-14B로 생성된 궤적의 예시이다. 입력 질문은 The director of the romantic comedy 'Big Stone Gap' is based in what New York city?이다.


첫 번째 라운드에서 모델은 뉴욕에 기반을 둔 제작팀을 가리키는 ‘Ghost’라는 정보를 입력받는다. 모델은 이 잠재적으로 유용한 정보를 향후 사용을 위해 저장한다. 두 번째 라운드에서는 관련 컨텍스트 정보가 제공되지 않지만, 모델은 에이전트 상태를 유지한다. 세 번째 라운드에서는 관련 입력값이 모두 제시된다. 모델은 핵심 정보를 정확하게 식별하고 그에 따라 메모리를 업데이트하여 정답을 도출하였다. 이 시점에서 추론 과정이 완료된다. 나머지 라운드에서 모델의 메모리는 변경되지 않고 최종 응답을 생성하는 데 사용된다.

You may also enjoy