MemReward: Graph-Based Experience Memory for LLM Reward Prediction with Limited Labels

本文提出了 MemReward，一种基于图的经验记忆框架，通过构建包含查询、思考过程和答案的异构图并利用图神经网络在少量标签下传播奖励信号，显著提升了大语言模型在复杂推理任务中的强化学习微调效率，在仅使用 20% 标签的情况下即可达到接近甚至超越全监督 Oracle 的性能。

要点

这篇论文介绍了一种名为 MemReward 的新方法，旨在解决大语言模型（LLM）在“强化学习”训练过程中面临的一个巨大难题：如何在不花大价钱请人打标签的情况下，让模型学会更好地思考？

为了让你轻松理解，我们可以把整个过程想象成**“一个天才学生（LLM）在备考，但老师（人类专家）太忙了，没时间批改每一道题”**的故事。

1. 背景：为什么我们需要 MemReward？

想象一下，你想训练一个 AI 像数学家或程序员一样思考。最好的方法是让它做很多题，然后告诉它：“这道题做对了，给你奖励；做错了，扣分。”

• 理想情况（Oracle）： 每一道题做完，都有专家老师立刻批改，给出正确答案。但这太贵、太慢了！比如证明一个复杂的数学定理，需要顶尖专家花几个小时；回答开放性问题，甚至没有标准答案。
• 现实困境： 我们只有很少的“专家批改”（比如只有 20% 的题有答案），剩下的 80% 题做完了却没人知道对错。如果只练那 20% 的题，AI 进步很慢；如果不管那 80%，又太浪费。

MemReward 就是为了解决这个“老师不够用”的问题而诞生的。

2. MemReward 的核心创意：建立“错题本”与“相似题库”

MemReward 不像传统方法那样孤立地看每一道题，而是把 AI 做过的所有题目（包括思考过程和最终答案）都存进一个**“超级经验图书馆”**里。

核心比喻：聪明的“学习小组”

想象 AI 是一个正在备考的学生，MemReward 就是它组建的一个**“超级学习小组”**。

1. 建立图书馆（构建异构图）：

   • 学生把做过的每一道题都记下来，包含：题目是什么（Query）、是怎么思考的（Thinking）、最后答案是什么（Answer）。
   • MemReward 把这些记录整理成一张巨大的关系网（图）。
   • 连线规则： 如果两道题长得像（比如都是关于“二次方程”的数学题），它们之间就拉一根线；如果一道题的思考过程和答案有关联，它们之间也拉线。这就形成了一个**“异构图”**（Heterogeneous Graph），就像把题目、思路、答案都编织在了一起。

2. 学习小组长（GNN 图神经网络）：

   • 在这个图书馆里，只有 20% 的题目有“老师批改过的红叉或红勾”（标签）。
   • MemReward 训练了一个**“小组长”（GNN）**。这个小组长的任务是：通过观察那些有红勾红叉的题目，去推测那些没批改的题目是对是错。
   • 怎么推测？ 如果一道新题和旁边几道“有红勾”的旧题长得特别像，小组长就会推断：“这道新题大概率也是对的！”

3. 在线考试（强化学习优化）：

   • 在正式训练时，AI 做新题。
   • 如果是那 20% 有老师批改的题，直接用老师的评分。
   • 如果是那 80% 没老师批改的题，“小组长”立刻跳出来，在图书馆里找最相似的几道旧题，根据旧题的评分，给新题一个**“预测分”**。
   • AI 拿着这个“预测分”继续学习，就像有了老师的指导一样。

3. 为什么这个方法这么厉害？

论文通过实验发现，MemReward 的效果惊人：

• 四两拨千斤： 只用 20% 的真实老师评分，MemReward 就能达到 97.3% 的“全知全能老师（Oracle）”的效果。也就是说，省了 80% 的精力，却几乎没损失成绩。
• 举一反三（泛化能力）： 最有趣的是，当 AI 遇到从未见过的新领域（比如训练时只学了数学和代码，考试时考物理常识）时，MemReward 甚至比全知全能老师表现得更好！
  • 比喻： 就像这个学习小组不仅记住了题目，还学会了“解题的底层逻辑”。遇到新题时，它能通过类比，比死记硬背标准答案的老师猜得更准。
• 结构很重要： 论文还发现，把“题目”、“思考过程”和“答案”分开记录（异构图）非常重要。
  • 如果只记答案（像死记硬背），效果就差。
  • 如果连上“思考过程”，小组长就能看懂 AI 是怎么一步步推导的，从而更准确地判断对错。

4. 总结：MemReward 到底做了什么？

简单来说，MemReward 就是给大语言模型装了一个**“智能记忆库”**。

• 以前： 模型做 100 道题，只有 20 道有人改，剩下 80 道只能瞎猜，或者干脆不做。
• 现在： 模型做 100 道题，20 道有人改。剩下的 80 道，模型会去“记忆库”里找相似的 20 道旧题，看看旧题是怎么被评价的，然后**“照猫画虎”**，自己给自己打分。

一句话总结：
MemReward 利用**“物以类聚”的原理，让大模型通过“参考相似题目的历史经验”**，在只有少量老师指导的情况下，也能学会像专家一样思考，大大降低了训练 AI 的成本。

技术摘要

MemReward 论文技术总结

1. 研究背景与问题 (Problem)

在大语言模型（LLM）的强化学习（RL）微调中，尤其是针对复杂推理任务（如数学证明、代码生成、开放式问答），奖励标签（Reward Labels）的获取成本极高。

• 挑战：获取大规模奖励标签通常需要昂贵的人工标注或耗时的验证程序（例如数学证明需要专家审核，开放式问答缺乏标准答案）。
• 现状：当奖励标签稀缺时（例如仅有少量标注数据），传统的 RL 微调效果受到严重限制，因为模型无法从大量未标注的生成轨迹（Rollouts）中有效学习。
• 核心问题：如何利用结构化的标注推理经验记忆，在标签稀缺的情况下，有效预测未标注轨迹的奖励，从而实现高效的 RL 微调？

2. 方法论 (Methodology)

论文提出了 MemReward，一种基于图的经验记忆框架（Graph-Based Experience Memory Framework），用于在标签稀缺场景下预测 LLM 的奖励。

核心架构

MemReward 将推理过程建模为一个异构图（Heterogeneous Graph），包含三种节点和三种边：

1. 节点类型：
   • 查询节点 (Query Nodes)：代表输入问题。
   • 思维节点 (Thinking Nodes)：代表模型的推理过程（Thinking Process）。
   • 答案节点 (Answer Nodes)：代表最终生成的答案。
2. 边类型：
   • Query-Query 边：基于查询嵌入的余弦相似度连接语义相似的查询（捕获跨问题的奖励模式）。
   • Query-Thinking 边：连接查询与其对应的思维过程。
   • Thinking-Answer 边：连接思维过程与其对应的答案。

工作流程

MemReward 分为两个阶段：

阶段一：暖启动训练 (Warmup Phase)

• 数据准备：使用初始策略 $\pi_0$ 为训练集中的标注查询生成推理轨迹（包含思维和答案）。
• 图构建：将标注的查询、思维、答案构建为异构图。
• GNN 训练：训练一个异构图神经网络（Heterogeneous GNN）。
  • 利用注意力机制聚合不同类型的邻居信息（如从相似查询聚合奖励信号）。
  • 通过二元交叉熵损失函数（Binary Cross-Entropy Loss）训练 GNN，使其能够根据节点特征预测轨迹的奖励分数（正确为 1，错误为 0）。

阶段二：在线策略优化 (Online Policy Optimization)

• 混合奖励获取策略：在 GRPO（Group Relative Policy Optimization）训练过程中：
  • 标注查询：直接使用真实奖励（Ground Truth）。
  • 未标注查询：
    1. 生成新的推理轨迹。
    2. 将未标注查询嵌入与暖启动图中的标注查询进行相似度匹配，连接 Top-K 个最近邻。
    3. 利用训练好的 GNN 传播信息，预测未标注轨迹的奖励。
• 奖励整合：将 GNN 预测的奖励与真实奖励混合，用于计算优势函数（Advantage），指导策略更新。

3. 关键贡献 (Key Contributions)

1. 基于图的奖励传播机制：首次将异构图神经网络引入 LLM 的奖励预测，利用语义相似性和结构化依赖（查询 - 思维 - 答案）在标注数据稀缺时传播奖励信号。
2. 跨域泛化能力：设计了一个共享的 GNN 架构，能够在不同领域（数学、问答、代码）之间迁移奖励模式。实验证明，仅用域内数据训练的 GNN 在域外任务上也能超越全监督基线。
3. 细粒度的推理结构建模：明确区分并建模了“思维过程”节点。消融实验表明，保留思维节点对于捕捉多步推理任务（如数学和代码）中的中间逻辑至关重要，这是仅基于查询或答案的模型无法做到的。
4. 高效的标签利用率：证明了在仅有 20% 标注数据的情况下，该方法能达到接近全监督（Oracle）的性能，显著降低了 RL 微调的标注成本。

4. 实验结果 (Results)

实验在 Qwen2.5-1.5B 和 Qwen2.5-3B 模型上进行了评估，涵盖数学、问答和代码生成三大领域共 13 个基准测试。

• 接近全监督性能：
  • 在 3B 模型上，仅使用 20% 标签，MemReward 达到了 97.3% 的全监督（Oracle）性能。
  • 在 1.5B 模型上，达到了 96.6% 的全监督性能。
  • 相比仅使用 20% 标签丢弃其余数据的基线（R1-p），MemReward 在 3B 模型上提升了 1.35 分，在 1.5B 模型上提升了 5.38 分。
• 超越全监督的泛化能力：
  • 在域外（Out-of-Domain）任务上，MemReward 甚至超越了全监督 Oracle。
  • 例如在 3B 模型上，MemReward 平均得分为 66.96，而 Oracle 为 66.07。这表明 GNN 预测的奖励信号有助于模型学习更通用的奖励模式，而非过拟合特定标注数据。
• 不同任务的收益：
  • 数学推理受益最大（如 GSM8K 提升 11.56 分），因为数学问题具有高度结构化的相似性，利于图传播。
  • 随着标注比例从 20% 增加到 70%，性能平滑提升，在 70% 时达到 Oracle 的 99.4%。
• 消融实验：
  • 移除图结构（使用 MLP）导致问答任务性能大幅下降。
  • 移除“思维节点”导致数学和代码任务性能显著下降，证实了中间推理步骤对奖励预测的重要性。

5. 意义与影响 (Significance)

• 降低 RL 门槛：MemReward 为解决 LLM 强化学习中“奖励标注昂贵”这一核心瓶颈提供了有效方案，使得在标注预算有限（如减少 80% 标注）的情况下，仍能进行高质量的 RL 微调。
• 提升泛化性：通过图结构挖掘数据间的潜在关联，模型不仅能利用标注数据，还能从未标注数据中学习到更鲁棒的奖励信号，从而在未见过的任务上表现更好。
• 方法论创新：将半监督学习中的标签传播思想与图神经网络结合，并具体化为“查询 - 思维 - 答案”的异构图结构，为 LLM 的推理能力增强提供了新的技术路径。

总结：MemReward 通过构建基于经验的异构图记忆，成功利用少量标注数据预测大量未标注轨迹的奖励，不仅大幅降低了 RL 微调的标注成本，还意外地提升了模型在域外任务上的泛化能力，是 LLM 推理能力训练领域的一项重要进展。