扶摇AI知识笔记AI 前沿知识库
智能体应用

[Agent Memory / 强化学习] MemPO源码学习笔记 --- (3)--- Rollout思路

来源:掘金 行业动态 约 15974 字 agent memorymemory
掘金
转载

本文转载自 掘金,版权归原作者及原发布平台所有。本站仅作知识整理与转载分享,如涉版权问题请联系客服删除。

01核心要点

  • [Agent Memory / 强化学习] MemPO源码学习笔记 --- (3)--- Rollout思路 0x00 概要 0x01 Rollout 主要内容 1

02正文全文

[Agent Memory / 强化学习] MemPO源码学习笔记 --- (3)--- Rollout思路

5.5 问题4:训练8轮 vs 评估64轮的分布漂移

5.6 问题 5:搜索结果 short_text 的非对称设计

现有的基于强化学习的 Memory 管理方法往往缺乏一种有效机制针对 Memory 的更新内容进行引导优化,Memory 的内容难以保证质量。

MemPO(Self-Memory Policy Optimization)使模型对 Memory 进行自管理,并引入了基于有效信息含量的 Memory-level 的优势估计,引导 Memory 保留对解决任务更有效的信息,进而提升记忆有效性。

MemPO的独特切入点:让模型把记忆写在每轮开头(),形式上像“自我对话的草稿纸“,既是记忆又是思考链的一部分。这样,变成可训练的策略变量,用RL信号端到端地教会模型“什么值得记、怎么记"。RL 直接端到端优化这一行为,无需额外的记忆模块。

论文标题:MemPO: Self-Memory Policy Optimization for Long-Horizon Agents

论文地址:arxiv.org/abs/2603.00…

代码地址:github.com/TheNewBeeKi…

模型和数据集地址:huggingface.co/collections…

MemPO 的时间线如下,可以看到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

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

1.训练rollout阶段:模型已经生成了轨迹(包含、、

full_traj(完整轨迹):系统提示+原始问题+第1轮对话+搜索结果 + 第2轮对话 + 搜索结果 + ···第T轮所有上下文

mem_traj(仅记忆上下文):系统提示+第T轮记忆内容

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条序列。之所以这样做而不是分两次调用:

Q -> [mem1][think][search] ──► 结果 ──► [mem2][think][answer Z]

MemPO 额外做的事 (不生成新内容,只算概率):

┌───────────────────────────┐ ┌─┐

│Q + 完整对话历史 (full_traj) │ ──►│Z│ P_full = 0.72

└───────────────────────────┘ └─┘

┌────────────────────────────────┐ ┌─┐

│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是从这条轨迹中提取/构造出来的。

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"不是指"完整的最终轨迹",而是指"那一轮的完整上下文"(相对于mem_traj只有question +的"压缩版")。

最终的轨迹(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

最终轨迹(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的对应位置。

┌─────────┬─────────────────────────┬────────────────────────────────────┐

│ 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 都会有。只有 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的对应位置。不是只用最后一个。

假设第3轮生成了:

"<mem>之前搜索发现HurtLocker获得2010最佳影片</mem><think>现在搜导演...</think><search>.."

mem_traj_ids = mem_sys_prompt_ids + response_mem_ids

↑ ↑

system prompt + 原始question response截取到</mem>为止

具体拼接:

[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>\nKathryn 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无论给什么上下文都预测不好(如人名的中间子词)

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[2N,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

[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字

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字)

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 轮,强制模型做 "记忆蒸馏"

以工具结果为计数单位,第2个工具结果之前全部截断用

short_text(搜索结果前100字,网页前100字)

用 full_text + 前缀 (强制以记忆开头)

每轮模型只能看到:原始问题+上轮记忆+上轮搜索摘要

记忆压缩链如下图所示。类比: 每一站只能带一张 "总结卡片" 上车,车厢里不存过去的东西。

这就是为什么MemPO必须对token单独施加奖励信号: 是唯一的信息传输通道,如果没有额外激励,模型会倾向于写无意义的记忆(反正每轮都重新搜索也能凑对)。

Round1 知识 --> Mem2 --> Mem3 --> ... --> Mem_T

↑每轮记忆必须覆盖所有历史,否则信息丢失

# prepare_llm_query():

short_text = job.get("short_text", job["text"])

# ↑如果没有 short_text 字段,fallback 到完整 text

搜索结果: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篇完整搜索结果 ← 刚放入

]

search_4 → flag=1(还没截够)

llm_4→跳过

search_3 → flag=2截断点在这里!

System: You are a helpful assistant.

User: 原始问题 <- 永远保留

Asst: <mem>round4记忆</mem> <- 仅最后一轮

<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_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> 通过了这个高标准,推理时(还能看到搜索摘要)

表现应该更好

一句话总结:两者确实相互印证—prepare_prompt()是"约束",mem_reward 是"激励",都指向同一个设计目标:让成为一个独立自洽的信息摘要。唯一的细节差异是训练信号比推理约束稍严(不包含 short_text),这实际上是一种"训练比推理更难"的保守设计,通常对泛化有益。

我们接下来进行潜在问题分析,即目前这种设计会导致模型遗忘什么?

原始问题(永久保留) + 上轮的<mem>(上一轮llm输出的完整文本) + 上轮搜索结果(full text或100字)

那么:必须"装下"哪些东西?我们结合几个问题来进行分析。

训练配置:max_response_length = 4096 tokens(run_train.sh 有效配置)

评估配置:模型最大输出(无硬限制,但受模型ctx限制)

一轮 llm 输出 = <mem> +<think>+ <search/answer>

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 已经丢失!

遗忘什么:推理链条的中间节点。如果模型在某轮没有将发现写入,该信息永久消失。

8轮(训练) VS 64轮(评估)→ 评估时遗忘风险远高于训练!

训练:max_assistant_turns =8(run_train.sh 有效配置),在8轮内,需"接力"7次

评估:max_turns=64,在64轮内,需"接力"63次

这类似于:训练时打5关的游戏,测试时要打64关 → 后期技能完全没有训练信号

search_results.short_text = full_text # 搜索:全文 (可能很长)

webpage.short_text = page[:100] # 网页:仅前100字

当前轮:prompt包含网页全文(25000字)

下一轮:prompt中这个网页只有100字(截断99.6%!)

100字约等于3~4句话→极度压缩

网页中的关键段落(位于第3段)→ 完全消失

这要求,模型在上轮必须把关键内容写进,但它不知道下轮只剩100字

MemPO奖励信号:mem_reward = P_mem(正确答案 I 仅靠) - P_full(正确答案 I 完整上下文)

训练压力:让足够好,使得"只看"和"看完整历史"效果接近

奖励以每轮为粒度,无法捕捉"第2轮记忆影响第5轮推理"的长程依赖

03原文直达

本文内容转载自 掘金,如需查看原排版、配图与最新修订,请访问原始出处。

阅读原文(掘金)

正在校验阅读权限…
RELATED

相关阅读

更多 智能体应用