MemPO 소스 학습 노트 3편, Rollout 사고방식 정리

0x00 개요

0x01 Rollout 주요 내용

- 1.1 타임라인

- 1.2 순수 Rollout

- 1.3 batch 계산

- 1.4 순전파

- 1.5 Rollout 궤적 생성

0x02 full_traj 구성 방식

- 2.1 스냅샷

- 2.2 사용

0x03 mem_traj 구성 방식

0x04 기억 압축 메커니즘

- 4.1 핵심 로직

- 4.2 구체적 사례

- 4.3 설계 의도 분석

- 4.4 설계 일관성 검증

0x05 잠재적 문제 분석

- 5.1 핵심 정보 흐름 모델

- 5.2 문제 1: 공간이 제한적이라 모든 이력을 다 쓸 수 없음

- 5.3 문제 2: 쓰는 것은 지난 라운드 전문이지 요약이 아님

- 5.4 문제 3: 멀티홉 추론 체인 단절

- 5.5 문제 4: 학습 8라운드 vs 평가 64라운드의 분포 드리프트

- 5.6 문제 5: 검색 결과 short_text의 비대칭 설계

- 5.7 총정리: 정보 망각 위험 매트릭스

- 5.8 MemPO가 완화하는 방법(단, 완전히 해결하지는 못함)

0xFF 참고

기존의 강화학습 기반 Memory 관리 방법은 흔히 Memory의 업데이트 내용을 겨냥해 유도·최적화하는 효과적인 메커니즘을 결여하고 있어, Memory의 내용 품질을 보장하기 어렵다.

MemPO(Self-Memory Policy Optimization)는 모델이 Memory를 스스로 관리하도록 하며, 유효 정보 함량에 기반한 Memory 수준(Memory-level)의 어드밴티지 추정을 도입하여, Memory가 과제 해결에 더 효과적인 정보를 보존하도록 유도함으로써 기억의 유효성을 높인다.

MemPO의 독특한 진입점: 모델이 매 라운드 시작 부분에 기억을 쓰게 하며(), 형식상 '자기 대화의 초안지'처럼 보여, 기억이자 사고 연쇄의 일부가 되게 한다. 이렇게 함으로써 가 훈련 가능한 정책 변수가 되고, RL 신호로 엔드투엔드로 모델에게 '무엇을 기록할 가치가 있고, 어떻게 기록할 것인가'를 가르친다. RL이 이 행위를 직접 엔드투엔드로 최적화하며, 별도의 기억 모듈이 필요 없다.

MemPO의 정보는 다음과 같다:

- 논문 제목: MemPO: Self-Memory Policy Optimization for Long-Horizon Agents

- 논문 주소: arxiv.org/abs/2603.00…

- 코드 주소: github.com/TheNewBeeKi…

- 모델 및 데이터셋 주소: huggingface.co/collections…

이번 글에서는 몇 가지 rollout 설계 아이디어를 살펴본다.

MemPO의 타임라인은 다음과 같으며, Rollout 단계를 볼 수 있다:

3-타임라인

Rollout 단계

순수 Rollout = 궤적 생성 + 데이터 수집 과정, 즉 AgentLoopManager.generate_sequences() 함수의 주된 부분:

순수 Rollout 단계(token 생성 + 도구 상호작용): B6 AgentLoopManager.generate_sequences ← 스케줄러, 16개 동시 실행 시작 A4 _handle_generating_state ← 매 라운드 생성 후 < mem > 위치와 내용 수집 C3 ToolParser.parse ← < search > / < access > 태그 파싱 C4 AsearcherSearchTool.execute ← RAG 검색 호출 B5 RewardManagerWorker.compute_score ← 궤적 완료 후 비동기 트리거(rollout과 병렬) └→ B1 NaiveRewardManager.__call__ ← 디코딩+스코어링 호출 └→ B2 compute_score ← 주 스코어링 진입점 ├→ C2 extract_solution ← < answer > 추출 ├→ B3 validate_format ← 형식 검증 └→ B4 em_check ← 정확 매칭

Rollout 후처리(여전히 같은 함수 내, 단 모든 궤적이 완료됨):

A1 _postprocess ← 추가 순전파, P_mem-P_full 계산

Rollout 단계에 속하지 않음

MemPO는 일부 내용이 뒤섞여 있다. 예를 들어: ans_mask와 threshold는 rollout 단계의 꼬리에 속한다——AgentLoopManager.generate_sequences() 함수 내에서, 16개 궤적 rollout이 모두 완료된 후 실행되는 추가 순전파다.

하지만 실제로 ans_mask와 threshold는 '순수 rollout' 단계(token 생성)에도, '순수 reward' 단계(B 시리즈 em_check)에도 속하지 않고, rollout 완료 후의 Memory Reward 계산 단계(A1)에 있다——여전히 rollout 함수 내부에 있지만, 논리적으로는 memory reward 계산에 속한다.

B4-algo compute_grpo_outcome_advantage ← PPO 업데이트 전 advantage 정규화 A2 compute_grpo_memory_advantage ← PPO 업데이트 전 advantage 정규화 A3 compute_advantage ← 중첩 A5 AgentMemory.prepare_prompt ← 평가 전용, 학습에서는 호출 안 함

따라서 우리는 이번 글에서 순수 rollout 단계뿐 아니라 ans_mask와 threshold 같은 '경계를 넘나드는' 내용도 소개할 것이다.

advantage는 batch 내의 각 궤적마다 계산된다.

핵심 포인트: 정규화는 question별로 그룹화되어 이루어진다(같은 그룹 16개 궤적끼리 서로 비교), 전 batch 통일 정규화가 아니다. 이는 난이도가 다른 question끼리 서로 간섭하지 않도록 보장한다 ✅

하나의 batch(가정: 3개 question × 16개 궤적/question = 48개): Question Q1: traj_1, traj_2, ..., traj_16 Question Q2: traj_17, traj_18, ..., traj_32 Question Q3: traj_33, traj_34, ..., traj_48 Outcome Advantage: Q1 그룹: scores = [ 1 , 0 , 1 , 1 , 0 ,...] → mean= 0.6 , std= 0.5 adv_1 = ( 1 - 0.6 )/ 0.5 = + 0.8 adv_2 = ( 0 - 0.6 )/ 0.5 = - 1.2 ... Q2 그룹: 독립적으로 mean/std 계산 Q3 그룹: 독립적으로 mean/std 계산 → outcome_adv [48, seq_len] 각 궤적마다 하나의 값, 그 모든 token에 브로드캐스트 Memory Advantage: Q1 그룹: 16개 궤적의 모든 라운드 mem_reward 풀링(~48개 값) → 통일 mean/std → 각 궤적의 각 라운드 <mem> 구간에 각각 할당 → mem_adv [48, seq_len] 각 궤적의 <mem> 구간마다 서로 다른 값 최종: final_adv [48, seq_len] = outcome_adv + mem_adv → 각 궤적, 각 token 위치마다 하나의 확정된 advantage 값 → 전부 PPO loss에 넣어 함께 업데이트

순전파(생성 단계의 통상적 샘플링) 외에, MemPO에는 순전파가 한 번 더 있다(extra forward pass), 바로 Teacher Scoring이다. Teacher Scoring의 특별한 점은:

이 compute_log_prob 호출은 detached되어 있다(PPO의 역전파에 참여하지 않음). 오직 mem_reward 수치를 계산하는 데 사용되며, 상수 계수로 advantage에 들어간다.

extra forward pass가 하는 일은: '이미 생성된 답안'에 대해 확률을 다시 계산하는 것이지, 다시 생성하는 것이 아니다. P_mem과 P_full의 계산 과정은 '완전한 추론'이 아니라 '주어진 텍스트의 확률을 계산'하는 것이다.

- 추론(Inference) = 모델이 스스로 새 token을 생성(자기회귀 디코딩)

- 순전파(Forward Pass) = 주어진 기존 텍스트에 대해 각 token의 확률을 계산

실제 방법

compute_log_prob은 A1에서 한 번의 호출로 2N개 입력을 동시에 처리한다——하지만 실제로는 한 번의 순전파(batch 추론)이지, 두 번의 독립적 순전파가 아니다.

- 병합 방식(실제 사용): 1회 순전파, batch=2N→높은 GPU 활용률

- 분할 방식(채택 안 함): 2회 순전파, 각각 batch=N→두 배의 호출 오버헤드, 게다가 KVcache 재사용 불가

A1의 방법은 다음과 같다:

1. 학습 rollout 단계: 모델이 이미 궤적을 생성함(,,

, 답안 Z 포함)

2. 두 가지 '입력 컨텍스트' 준비:

- full_traj(완전 궤적): 시스템 프롬프트+원본 문제+제1라운드 대화+검색 결과 + 제2라운드 대화 + 검색 결과 + ···제T라운드 모든 컨텍스트

- mem_traj(기억 컨텍스트만): 시스템 프롬프트+제T라운드 기억 내용

3. '정답 Z'를 뒤에 이어 붙여, 모델이 확률을 계산하게 함:

P_full = P(Z| full_traj) P_mem = P(Z| mem_traj)

구체적으로는 다음과 같다

full_traj와 mem_traj를 하나의 큰 batch로拼接 traj_input.input_ids = pad([ full_1, full_2, ..., full_N, ← N개 완전 컨텍스트 mem_1，mem_2,..., mem_N, ← N개 <mem> 요약만 포함 J) #answer 두 번 반복 traj_input.responses =[ans_1,...,ans_N, ans_1,...,ans_N] #한 번의 순전파(batch size=2N), 두 가지 컨텍스트의 logp 동시 계산 log_probs = model.compute_log_prob(traj_input) #결과 분리 full_logp = log_probs[:N] # ←앞부분: 완전 컨텍스트→답안의 log 확률 mem_logp = log_probs[N:2N] # ←뒷부분: mem 요약→답안의 log 확률 # 핵심: 이것은 한 번의 batch 순전파로, GPU가 2N개 시퀀스를 동시에 처리한다. 두 번 호출로 나누지 않은 이유:

도식

rollout이 생성한 궤적(완료됨):

Q -> [mem1] [think] [search] ──► 결과 ──► [mem2] [think] [answer Z]

MemPO가 추가로 하는 일(새 내용을 생성하지 않고 확률만 계산):

상황 A: 완전 컨텍스트를 주면, 답안 Z의 확률은 얼마인가?

┌───────────────────────────┐ ┌─┐ │Q + 완전 대화 이력 (full_traj) │ ──►│Z│ P_full = 0.72 └───────────────────────────┘ └─┘

상황 B: 내용만 주면, 답안 Z의 확률은 얼마인가?

┌────────────────────────────────┐ ┌─┐ │sys_prompt + <mem>기억 내용</mem> │──► │Z│ P_mem = 0.68 └────────────────────────────────┘ └─┘

mem_reward = P_mem - P_full = 0.68 - 0.72 = -0.04 ──►기억을 충분히 잘 쓰지 못함, 기억만으로는 완전 컨텍스트를 보는 것보다 4% 나쁨 ──►token에 부정적 보상을 부여하여 모델이 기억을 개선하도록 촉진

Rollout은 오직 한 가지 궤적——완전한 다중 라운드 대화 궤적——만을 생성한다. full_traj와 mem_traj는 이 궤적에서 추출/구성된 것이다.

Rollout 단계에서 실제로 일어나는 일:

B6: 동일한 question을 SGLang에 16번 보냄 → 16개의 독립적인 동시 궤적(각각은 상태 기계 GENERATING>TOOL_CALLING>...를 거침) → LLM 샘플링에 무작위성(temperature)이 있기 때문에 16개 궤적의 내용은 서로 다름

단일 궤적의 생성 과정(다중 라운드 예시): Round 1: LLM 생성 → < mem > .:. < think > ... < search > query </ search > → 도구가 결과 반환 Round 2: LLM 생성 → < mem > ... < think > ... < search > query </ search > → 도구가 결과 반환 Round 3: LLM 생성 → < mem > :.. < think > ... < answer > xxx </ answer > → 종료

이것이 유일하게 생성된 "궤적" → response_ids

A4가 각 라운드 생성 후 "겸사겸사" 추출: Round 2 생성 후: full_traj[0]=deepcopy(raw_input_ids) ← 당시의 완전한 컨텍스트 스냅샷 mem_traj[0]=mem_sys_prompt_ids + " < mem > R2 요약 </ mem > " ← 수작업拼接 Round 3 생성 후: full_traj[1]=deepcopy(raw_input_ids) ← 더 길어진 완전한 컨텍스트 mem_traj[1]=mem_sys_prompt_ids + " < mem > R3 요약 </ mem > " ← 수작업拼接

핵심 구분:

- 생성된 궤적: 오직 한 종류뿐(16개의 완전한 대화 궤적, LLM이 실제로 생성한 token)

- 추출/구성된 데이터: full_traj는 컨텍스트 스냅샷이고, mem_traj는 인공적으로拼接한 짧은 시퀀스

- full_traj와 mem_traj는 LLM을 통해 생성된 것이 아니라, 기존 데이터에서 잘라내거나拼接한 것

스냅샷

full_traj는 궤적 전체가 아니라, 특정 라운드가 생성을 시작하기 전의 "컨텍스트 스냅샷"이다.

5라운드 궤적을 가정하면: Round 1 : [system + question] → R1 response Round 2 : [system + question + R1 + tool_result_1] → R2 response Round 3 : [system + question + R1 + tool_result_1 + R2+ tool_2] → R3 response Round 4 : [...더 김] → R4 response Round 5 : [...가장 김] → R5 (answer) full_traj_list가 수집하는 것(Round 2 + 이면서 <mem>이 있는 라운드만): full_traj [0] = [system + question + R1+ tool_result_1] → R2 전의 컨텍스트 full_traj [1] = [system + question + R1 + tool_1+R2+ tool_2] → R3 전의 컨텍스트 full_traj [2] = [system + question + R1 + tool_1+R2+tool_2+R3+tool_3] → R4 전의 컨텍스트 ↑ 각각은 해당 라운드 생성 전의 완전한 prompt(deepcopy(raw_input_ids)) ↑ 해당 라운드에서 생성된 response는 포함하지 않음

따라서 full_traj="현재 라운드까지의 완전한 대화 이력(해당 라운드 출력 제외)"이며, "모델이 지금까지의 모든 대화를 보았다면 정답을 예측할 확률은 얼마인가?"라는 질문에 답하기 위해 사용된다. 이와 대조적으로 mem_traj="system+question + 해당 라운드 요약만 봄"이다.

비교

실제로 full_traj라는 이름은 오해를 부르기 쉽다 — full_traj의 "full"은 "완전한 최종 궤적"을 뜻하는 것이 아니라 "그 라운드의 완전한 컨텍스트"(question만 있는 mem_traj의 "압축판"에 상대하여)를 뜻한다.

최종 궤적(rollout이 산출한 진짜 궤적): = prompt_ids + response_ids = [system + question] + [R1 + tool_1 + R2+ tool_2+ R3 +..+ R5_answer] ← 모든 라운드의 완전한 대화를 포함하며, PPO Update에 사용되는 데이터 full_traj [k] (A4가 수집한 것): = k+2번째 라운드 생성 전의 컨텍스트 스냅샷 ← 최종 궤적이 아니라 중간의 "단면" ← 오직 A1 Memory Reward 계산에만 사용 ← PPO Update에는 사용하지 않음

최종 궤적과 full_traj의 관계는 다음과 같다:

최종 궤적(response_ids): [R1_tokens | R2_tokens | R3_tokens | R4_tokens |.R5_tokens] ← 이것이 PPO 훈련의 대상 full_traj [] : full_traj [0] = R2 시작 전까지의 컨텍스트 ← 최종 궤적의 "접두사 잘라내기" full_traj [1] = R3 시작 전까지의 컨텍스트 ← 더 긴 접두사 full_traj [2] = R4 시작 전까지의 컨텍스트 ← 더 긴 접두사 ← 오직 MemoryReward의 확률 비교에만 사용

계산

만약 최종 궤적이 5라운드라면, 우리는 5라운드 전부를 계산하지 않고 "일부 라운드"만 계산하며, 최종 정규화 후 final_adv의 해당 위치에 기록한다.

5라운드 궤적의 전형적인 경우:

┌─────────┬─────────────────────────┬────────────────────────────────────┐ │ Round 1 │ 첫 번째 라운드, 이력 없음 │ → < mem >을 생성하지 않음 → ✗ 계산 안 함 │ │ Round 2 │ < mem > 있음 │ → ☑ mem_reward_R2 계산 │ │ Round 3 │ < mem > 있음 │ → ☑ mem_reward_R3 계산 │ │ Round 4 │ < mem > 있음 │ → ☑ mem_reward_R4 계산 │ │ Round 5 │ 마지막 라운드에서 answer 제공 │ → < mem >이 있을 수도 있고 없을 수도 있음 │ │ │ │ │ 있으면 → ☑ 계산; 없으면 → ✗ 계산 안 함 │ └─────────┴────────────────────────┴─────────────────────────────────────┘

흐름은 다음과 같다:

A1: < mem >이 있는 3-4개 라운드에 각각 mem_reward 계산 A2: 모든 mem_reward를 궤적 간·라운드 간 풀링 → 정규화 A3: 정규화된 mem_adv를 final_adv의 해당 < mem > ... </ mem > 구간에 기록 final_adv [seq_len]: [R1 tokens | < mem > R2 </ mem > | other R2 | < mem > R3 </ mem > | other R3 | ...] [outcome_adv | outcome+adv_2 | outcome | outcome+adv_3 | outcome | ...] ↑ 기록 기록 안 함 ↑ 기록 기록 안 함

핵심 규칙:

validate_format의 규칙 8은 매 라운드마다 기록하도록 강제하므로, 정상적인 경우 Round 2-5는 모두 <mem>이 있다. Round 1만 이력이 없기 때문에 생성하지 않는다. 만약 특정 라운드가 형식 오류로 누락되면 그 라운드도 건너뛴다.

모든 full_traj는 A1(Memory Reward 계산)에서 전부 사용된다. 주의: full_traj는 Memory Reward 계산(A1)에만 사용되며, Outcome Reward(B 계열)에도, PPO Update의 actor 전방 계산에도 사용되지 않는다. PPO Update는 rollout이 산출한 원본 response_ids를 사용한다.

batch=3개의 question × 16개 궤적/question=48개 궤적이고, 각 궤적이 3라운드를 생성한다고 가정 → 3개의 full_traj

A1의 입력: concat_fu1l = 48개 궤적 × 3라운드 = 144개 full_traj(전부 평탄화) concat_mem = 48개 궤적 × 3라운드 = 144개 mem_traj(전부 평탄화) → 2× 144 = 288개로拼接 → 1회 compute_log_prob 전방 → 144개 P_full과 144개 P_mem 획득 → 144개 mem_reward = P_mem - P_full 궤적으로 재조립: 궤적1의 mem_rewards = [r_R2, r_R3, r_R4] 궤적2의 mem_rewards = [r_R2, r_R3] ... → 전부 A2의 advantage 정규화에 사용 → 최종적으로 PPO loss에서 <mem> token의 기울기 방향에 영향

A1 계산 시, 각 full_traj[k]는 각각 하나의 mem_reward를 산출한다:

full_traj [0] + answer → P_full_R2 ─┐ mem_traj [0] + answer → P_mem_R2 ─┤→ mem_reward_R2 = P_mem_R2 - P_full_R2 full_traj [1] + answer → P_full_R3 ─┐ mem_traj [1] + answer → P_mem_R3 ─┤→ mem_reward_R3 = P_mem_R3 - P_full_R3 full_traj [2] + answer → P_full_R4 ─┐ mem_traj [2] + answer → P_mem_R4 ─┤→ mem_reward_R4 = P_mem_R4 - P_full_R4 → mem_rewards = [mem_reward_R2, mem_reward_R3, mem_reward_R4] → 3개 값이 전부 A2 정규화로 들어감 A2 정규화 후 final_adv의 서로 다른 구간에 기록: final_adv: [R1_tokens | <mem>R2</mem> tokens | ... | <mem>R4</mem> tokens | R5] 0 mem_adv_R2 ... mem_adv_R4 0

각 full_traj[k]는 한 라운드의 구간에 대응하며, 각각 독립적으로 mem_reward를 계산하고 각각 final_adv의 해당 위치에 기록한다. 마지막 것만 사용하는 것이 아니다.

구성 과정(A4 단계에서 수집)은 다음과 같다:

3라운드에서 다음과 같이 생성되었다고 가정: " < mem > 이전 검색에서 HurtLocker가 2010년 최고 작품상을 받은 것을 발견 </ mem > < think > 이제 감독을 검색... </ think > < search > .."

mem_traj 구성

mem_traj_ids = mem_sys_prompt_ids + response_mem_ids ↑ ↑ system prompt + 원본 question </ mem >까지 잘라낸 response 구체적拼接: [ system : "You are a helpful assistant."] [ user : "Who directed the 2010 Best Picture winner?"] [assistant："<mem>이전 검색에서 HurtLocker가 2010년 최고 작품상을 받은 것을 발견</mem>"] ↑ 잘라냄, 뒤의 think / search는 전부 버림 full_traj와 비교: [ system + question + Round1 완전한 대화 + Round2 완전한 대화 + Round3 시작...] ↑ 완전한 다중 라운드 이력 컨텍스트(매우 김)

mask가 있는가? ans_mask가 있지만, mem_traj 자체에 mask를 하는 것이 아니라 답변 token에 mask를 한다:

answer_ids = tokenize( " \n <think>... \n </think> \n <answer> \n Kathryn Bigelow \n </answer>" ) ans_mask: [ 0 ， 0 ， 0 ， 0 , 1 ， 1 , 1 ， 1 ， 1 , 1 ， 1 ， 1 ， 0 , 0 , 0 , 0 ] ↑ Kathryn Bigelow 의 token ↑ \n </ answer > 의 4 개 token 의미: "핵심 답변 내용"의 log_prob만 보고, < think >/< answer > 태그 token은 무시한다

또 한 층의 threshold 필터 mask가 있는데, 그 의미는 모델이 "전혀 확신하지 못하는" token을 무시한다는 것이다

- 어떤 token은 어떤 맥락을 주든 예측을 잘 못한다(예: 인명의 중간 서브워드)

- 그것들을 걸러내고, 모델이 확신하는 token만 본다

- 이렇게 하면 P_mem-P_full 비교가 더 안정적이다

full_ans_mask = ans_mask AND (full_logp > log( 0.5 )) mem_ans_mask = ans_mask AND (mem_logp > log( 0.5 ))

전체 도식은 다음과 같다:

mem_traj (입력): [sys_prompt| question|<mem>요약</mem>] ← mask 없음 full_traj (입력):· [sys_prompt|question|전체 멀티턴 히스토리...] ← mask 없음 answer (목표): [\n<think>...\n<answer>\n{gt}\n</answer>] model .compute_log_prob (input, answer): → log_prob[ 2 N，ans_len]←각 답변 token의 조건부 확률 필터링: ans_mask x (logp > threshold) → 핵심 답변 token 중 신뢰도 높은 부분만 유지 → P = exp ( mean (filtered_logp))

이제 단계별로 살펴보자.

이 소절은 주로 AgentMemory·prepare_prompt() 메모리 압축 메커니즘에 대한 심층 분석이다.

def prepare_prompt ( self ): # 1. 고정 보존: 시스템 프롬프트+원본 질문 prompt = [{ "role" : "system" , "content" : "You are a helpful assistant." }] # 원본 질문은 절대 잃지 않음 prompt.append({ "role" : "user" , "content" : self.memory[ 0 ].text}) # initial prompt # 2. 끝에서 앞으로 스캔하며 "절단점 i"를 찾음 flag = 0 flag = 0 for i in range (- 1 , - len (self.memory)- 1 , - 1 ): r = self.memory[i] if r. type == "prompt" : #到头了 flag = 2 elif r. type in [ "search_results" , "webpage" ]: # 도구 결과를 만날 때마다 +1 flag += 1 elif r. type == "llm_gen" : #llm 생성은 건너뜀 continue else : raise RuntimeError( f"Unknown record type: {r. type } " ) if flag == 2 : #두 번째 도구 결과→절단 break for j in range (i + 1 , 0 ): # 3. 절단점 이후의 내용만 유지 r = self.memory[j] if r. type in [ "search_results" , "webpage" ]: # 도구 결과는 요약 사용 prompt.append({ "role" : "user" , "content" : r.short_text}) elif r. type == "llm_gen" : #llm 생성은 전체 텍스트 사용 prompt.append({ "role" : "assistant" , "content" : "<mem>" + r.text}) else : raise RuntimeError( f"Unknown record type: {r. type } " ) return prompt

4 라운드 대화 memory 배열 내용:

[0] prompt = "원본 질문" [1] llm_gen_1 = " < mem > 기억1 </ mem > < think > ... </ think > < search > ... </ search > "" [2] search_1 = "information>검색결과1(완전, 매우 길 수 있음) </ information > " [3] llm_gen_2 = " < mem > 기억2 </ mem > .." [4] search_2 = " < information > 검색결과2(완전) </ information > " [5] llm_gen_3 = " < mem > 기억3 </ mem > .." [6] search_3 = " < information > 검색결과3(완전) </ information > " ◄─── short_text = 앞 100자 [7] llm_gen_4 = " < mem > 기억4 </ mem > ..." [8] search_4 = " < information > 검색결과4(완전) </ information > " ◄─── short_text = 앞 100자

backward 스캔(5번째 라운드를 준비할 때):

i =- 1 :search_4 → flag= 1 (멈추지 않음) i =- 2 :llm_gen_4 → skip i =- 3 :search_3 → flag= 2 ← BREAK! 절단점 =i=- 3

forward 조립 range(-2,0)=[-2，-1]:

memory [-2] = llm_gen_4 → assistant: "<mem>기억4</mem>. " memory[- 1 ] = search_4 → user:short_text (앞 100 자)

최종적으로 LLM에 전달되는 prompt:

System: You are a helpful assistant. User: 원본 질문 Asst: < mem > 기억4 </ mem > < think > ... </ think > < search > ... </ search > <- 가장 최근 1 라운드만 User: 검색결과4 (앞100자) <- short_text 압축 Asst: < mem > <- 모델이 계속 생성

전혀 보이지 않음: round1、round2、round3의 내용!

이것이 "달팽이 등껍질" 식 메모리 압축 설계다.

"달팽이 등껍질" 식 메모리 압축 설계 ──────────────────────────────────────────────────────── Round 1: [Q] -> [Mem1] -> [search] -> [결과1(전체)] -> 생성 Round 2: [Q] + [결과1(단축)] -> [Mem2] ↑ 지난 라운드 결과 요약만 보이고, 나머지는 전부 Mem1에 의존 Round 3: [Q] + [결과2(단축)] -> [Mem3] ↑ 위와 같음, Mem2가 Mem1 + round1의 정보를 운반함 Round T: [Q] + [llm_gen_{T-1}] + [결과_{T-1}(단축)] -> 생성 ↑최근 1 라운드만 노출되어 모델이 "기억 증류"를 강제로 수행함

메커니즘은 다음과 같다:

메커니즘 세부 사항 절단 입도 도구 결과를 카운트 단위로 하여, 두 번째 도구 결과 이전까지 전부 절단 도구 결과 사용 short_text(검색 결과 앞 100자, 웹페이지 앞 100자) 모델 생성 full_text + 접두사 사용 (기억으로 시작하도록 강제) 원본 질문 영구 유지(memory[0], 절단하지 않음) 강제 효과 매 라운드 모델은 다음만 볼 수 있음: 원본 질문+지난 라운드 기억+지난 라운드 검색 요약

기억 압축 체인

기억 압축 체인은 아래 그림과 같다. 비유: 각 정거장마다 "요약 카드" 한 장만 들고 타며, 차량 안에는 과거의 것을 저장하지 않는다.

이것이 바로 MemPO가 token에 별도로 보상 신호를 부여해야 하는 이유다: 이것이 유일한 정보 전송 통로이며, 추가적인 인센티브가 없으면 모델은 무의미한 기억을 쓰는 쪽으로 기울 것이다(어차피 매 라운드 다시 검색해도 맞출 수 있으니까).

Round1 지식 - -> Mem2 - -> Mem3 - -> ... - -> Mem_T ↑매 라운드 기억이 모든 히스토리를 커버해야 하며, 그렇지 않으면 정보가 소실됨

short_text의 진짜 의미

short_text의 두 가지 경우

prepare_llm_query(): short_text = job.get( "short_text" , job[ "text" ]) # ↑short_text 필드가 없으면 전체 text로 fallback

도구 결과 유형 text short_text search_results 완전 검색 결과(5편×각 편 앞 5000자) =text(압축 없음!) webpage 완전 웹페이지 분할(각 단락 25000자) 앞 100자(진짜 압축)

결론:

- 검색 결과: short_text는 완전 내용과 동일하며, 압축 없음

- 웹페이지 내용: short_text는 앞 10o자만 유지하며, 대폭 압축

절단점의 진짜 의미

절단점은 "텍스트 내용"에 대한 절단이 아니라, "히스토리 라운드"에 대한 절단이다: 4 라운드가 이미 진행되었고 5번째 라운드를 준비한다고 가정:

memory =[ [ 0 ] prompt [ 1 ] llm_gen_1 "<mem>round1기억</mem><think>...</think><search>...</search>" [ 2 ] search_1 topk= 5 편 완전 검색 결과 [ 3 ] llm_gen_2 "<mem>round2기억</mem>..." [ 4 ] search_2 topk= 5 편 완전 검색 결과 [ 5 ] llm_gen_3 "<mem>round3기억</mem>..." [ 6 ] search_3 topk= 5 편 완전 검색 결과 [ 7 ] llm_gen_4 "<mem>round4기억</mem>..." [ 8 ] search_4 topk= 5 편 완전 검색 결과 ← 방금 넣음 ]

backward 스캔, 도구 결과 세기:

search_4 → flag = 1 (아직 충분히 절단 안 됨) llm_4→건너뜀 search_3 → flag = 2 절단점은 여기!

prompt 조립, 절단점 이후의 내용만 취함:

llm_gen_4 + search_4

최종적으로 LLM이 보는 prompt:

System: You are a helpful assistant. User: 원본 질문 <- 영구 유지 Asst: < mem > round4기억 </ mem > <- 마지막 1 라운드만 < think > ... </ think > < search > ... </ search > User: [round4의 완전 검색 결과] <- short_text(검색=전문) Asst: < mem > <- 모델이 계속 생성하기를 기다림

round1、round2、round3→완전히 사라짐!

직관적 이해: "텍스트를 압축"하는 것이 아니라 "히스토리를 버리는" 것이다:

Round 1 Round 2 Round 3 Round4 Round5 [llm] [검색] [llm] [검색] [llm] [검색] [llm] [검색] ？ xxx xxx xxx √√√ x = 완전히 버림(prompt에 넣지 않음) √ = prompt에 넣음(검색 결과는 short_text 사용, llm 출력은 접두사 <mem>)

모델은 round4에서 round1-3의 모든 중요한 정보를 읽어내야 한다. 이것이 바로 잘 쓰느냐 못 쓰느냐가 모델이 여러 라운드 후에 정답을 줄 수 있느냐를 직접 좌우하는 이유다-----훈련의 품질이 바로 MemPO의 핵심 목표다.

mem_reward와 prepare_prompt()가 서로를 뒷받침하는지 살펴보자.

가설 확인

먼저 둘 각자의 "가설"을 확인한다.

prepare_prompt()가 추론 시에 모델에게 보여주는 것:

- sys_prompt - 원래 문제(memory[0]) - 이전 라운드 llm 출력(<mem>기억</mem><think>...</think> - 이전 라운드 검색 결과(short_text)

mem_traj(훈련 시 P_mem을 계산하는 데 사용)

mem_sys_prompt_ids = deepcopy(prompt_ids) + response_mem_ids(현재 라운드 <mem>...</mem> 내용) # prompt_ids = 완전한 초기 prompt, 시스템 프롬프트 + 원래 문제 포함(이것이 rollout 시작 시의 input)

둘의 비교

prepare_prompt()(추론) mem_traj(훈련 신호) ───────────────────────────────────────────────────────────────────────────── 시스템 프롬프트 √포함 포함(prompt_ids의 일부) 원래 문제 √포함 포함(prompt_ids의 일부) 현재 라운드 <mem> √포함(이전 라운드의 기억) 포함(현재 라운드의 기억) 이전 라운드 검색 결과 √포함(short_text) X불포함 이전 대화 X불포함(윈도우 잘림) X불포함

결론: 대체로 일치하며, 하나의 핵심 차이가 있다:

- 일치하는 부분(핵심 로직 정렬): 둘 다 "+원래 문제만으로" 답할 수 있어야 한다고 강조 → mem_reward 훈련 목표 = 추론 시의 실제 제약 √

- 차이점: prepare_prompt()는 "이전 라운드 검색 결과의 short_text"(100자)도 보여주지만, mem_traj는 이 short_text를 포함하지 않는다

영향:

P_mem의 계산 조건이 추론 시의 실제 조건보다 더 엄격하다(추론 시에는 100자 검색 요약을 볼 수 있지만, 훈련 보상은 기억만 볼 수 있다고 가정) -> 이는 훈련 신호가 실제로 "높은 기준 버전"이라는 뜻 -> 모델의 <mem>이 이 높은 기준을 통과하면, 추론 시(검색 요약도 볼 수 있음)에는 성능이 더 좋아야 한다

설계 일관성 검증도

3-설계 일관성 검증도

한 문장 요약: 둘은 실제로 서로를 뒷받침한다—prepare_prompt()는 "제약"이고, mem_reward는 "동기 부여"이며, 둘 다 동일한 설계 목표를 가리킨다: (모델이) 독립적이고 자기완결적인 정보 요약이 되도록 하는 것. 유일한 세부 차이는 훈련 신호가 추론 제약보다 약간 더 엄격하다는 것(short_text 미포함)이며, 이는 실제로 "훈련이 추론보다 더 어렵다"는 보수적 설계로, 일반적으로 일반화에 유익하다.

다음으로 잠재적 문제 분석, 즉 현재 이러한 설계가 모델로 하여금 무엇을 망각하게 만드는지 진행한다.

매 라운드마다 모델이 보는 prompt:

원래 문제(영구 보존) + 이전 라운드의 <mem>(이전 라운드 llm 출력의 전체 텍스트) + 이전 라운드 검색 결과(full text 또는 100자)

그렇다면: 무엇을 "담아내야" 하는가? 몇 가지 문제와 결합해 분석해 보자.

- 훈련 설정: max_response_length = 4096 tokens(run_train.sh 유효 설정)

- 평가 설정: 모델 최대 출력(하드 제한 없음, 단 모델 ctx 제한을 받음)

한 라운드 llm 출력 = <mem> + <think> + <search/answer>

만약 (에) 500 tokens를 썼다면, +

524 tokens만 남음

5 라운드 대화에서 축적된 지식이 하나의 <mem>에 압축되어야 함: Round1 발견: A는 모회사 창업자(기록 필요) Round2 발견: A 회사는 1998년 설립(기록 필요) Round3 발견: 1998년 모 산업 정책(관련 배경, 기록할까?) Round4 발견: 경쟁사 정보(기록할까?) → 5 라운드 지식 → 500 token<mem>, 선택적 망각이 불가피

무엇을 망각하는가: 모델은 "중요해 보이지 않는" 중간 사실을 버리는 경향이 있는데, 이런 사실들은 이후 추론에서 핵심적일 수 있다.

prepare_prompt() prompt.append({"role":"assistant","content":" < mem > " + r.text}) # ↑전문! # r.text = 이전 라운드의 완전한 llm 출력, <think>, <search> 등 태그 포함 이전 라운드 llm 출력(r.text) 구조: < mem > 제4라운드 기억 요약 </ mem > < think > 이번에 xxx를 검색했고, 이전의 YYY와 결합하면... </ think > < search > 검색어 </ search > 이번 라운드 prompt에서 보이는 것: < mem > + 위의 전문 = < mem > < mem > 제4라운드 기억 요약 </ mem > < think > ... </ think > . 모델은 중첩된 < mem > < mem > ... 구조에서 정보를 추출해야 함

잠재적 망각: 중첩 구조는 정보 추출 난이도를 높이며, 모델은 최외곽 접두사 내용만 주목하고 안쪽의 추론 과정을 무시할 수 있다.

전형적 멀티홉 문제:

Q: "X의 창업자는 어느 대학에서 박사 학위를 받았는가?" Round 1: "X회사" 검색 → 창업자가 John임을 발견 Round 2: "John 학력 배경" 검색 → John이 MIT에서 박사 과정을 밟았음을 발견 Round 3: 답변 준비, <mem>에 포함되어야 할 것: (a)X의 창업자=John(round1 발견) (b)John이 MIT에 있음(round2 발견) 하지만 round 1의 검색 결과는 round 3에서 이미 사라졌다!

무엇을 망각하는가: 추론 사슬의 중간 노드. 모델이 어느 라운드에서 발견한 것을 (mem에) 쓰지 않으면, 그 정보는 영구히 사라진다.

확률적 망각 분석:

- 매 라운드마다 역사를 압축해야 함 → 누적 망각률이 라운드 수에 따라 지수적으로 증가

- 8라운드(훈련) VS 64라운드(평가)→ 평가 시 망각 위험이 훈련보다 훨씬 높다!

- 훈련: max_assistant_turns =8(run_train.sh 유효 설정), 8라운드 내에서 7번 "이어받기"해야 함

- 평가: max_turns=64, 64라운드 내에서 63번 "이어받기"해야 함

모델은 63번 이어받기한 샘플을 본 적이 없다 → 심층 기억 실패

이는 다음과 유사하다: 훈련 시에는 5관짜리 게임을 하지만, 테스트 시에는 64관을 해야 함 → 후반부 기술에는 전혀 훈련 신호가 없다

search_results.short_text = full_text # 검색: 전문 (매우 길 수 있음) webpage.short_text = page[: 100 ] # 웹페이지: 처음 100자만

webpage 유형에 대해 다음과 같이 분석한다:

현재 라운드: prompt에 웹페이지 전문(25000자) 포함 다음 라운드: prompt에서 이 웹페이지는 100자만 남음(99.6% 잘림!) 100자는 약 3~4문장에 해당 → 극도로 압축됨 웹페이지의 핵심 단락(제3단락에 위치) → 완전히 사라짐

이것은 모델이 이전 라운드에서 핵심 내용을 (mem에) 써 넣어야 함을 요구하지만, 모델은 다음 라운드에서 100자만 남는다는 것을 모른다

정보 유형 망각 확률 원인 초기 라운드 검색 결과의 세부 사항 높음 prompt에서 완전히 삭제됨 웹페이지 내용의 비(非)도입부 부분 매우 높음 short_text가 처음 100자만 보존 중간 추론 사슬의 중간 노드 중-높음 모델의 선택적 압축 원래 문제의 원문 영(零) memory[0]에 항상 보존됨 이전 라운드의 전문 영(零) prompt에 직접 들어감 수치 정밀도(인명, 날짜 등) 중간 공간 제한으로 인한 생략

MemPO 보상 신호: mem_reward = P_mem(정답 I (mem)만으로) - P_full(정답 I 완전한 컨텍스트)

훈련 압력: (mem을) 충분히 좋게 만들어서 "(mem)만 보는 것"과 "완전한 역사를 보는 것"의 효과가 비슷해지도록

한계:

- 보상은 최종 답변만 보며, 중간 추론 사슬의 완전한 전달을 장려하지 않음

- 보상은 매 라운드를 단위로 하여, "제2라운드 기억이 제5라운드 추론에 영향을 준다"는 장거리 의존성을 포착할 수 없음

- 훈련은 8라운드뿐이라, 모델이 심층 망각 시나리오를 본 적이 없음

TransFormer-커버

본문은 markdown.com.cn으로 조판됨

뤄시의 생각

422

글

417k

읽음

436

팬