FlowRL: 匹配LLM推理的奖励分布

发表
Daixuan ChengDaixuan Cheng 提交
作者: zhuXuekai Zhu, Daixuan ChengDaixuan Cheng, Dinghuai ZhangDinghuai Zhang, Hengli Li, Kaiyan ZhangKaiyan Zhang, Che Jiang, Youbang Sun, Ermo Hua, Yuxin ZuoYuxin Zuo, Xingtai LvXingtai Lv, Qizheng ZhangQizheng Zhang, LChenLin Chen, Fanghao Shao, Bo Xue, Yunchong Song, Zhenjie YangZhenjie Yang, Ganqu CuiGanqu Cui, Ning Ding, Jianfeng GaoJianfeng Gao, Xiaodong Liu, Bowen Zhou, Hongyuan Mei, Zhouhan LinZhouhan Lin

摘要

AI 生成总结
FlowRL 通过流量平衡匹配完整的奖励分布,增强了大型语言模型的强化学习,相较于奖励最大化方法,提高了多样性和性能。
我们提出FlowRL:通过流平衡匹配完整的奖励分布,而不是在大型语言模型(LLM)强化学习(RL)中最大化奖励。最近的先进推理模型采用了奖励最大化方法(例如PPO和GRPO),这些方法倾向于过度优化主导奖励信号,而忽略不频繁但有效的推理路径,从而降低了多样性。相比之下,我们使用一个可学习的划分函数将标量奖励转换为归一化的目标分布,然后最小化策略和目标分布之间的反向KL散度。我们将这个想法实现为一个流平衡优化方法,该方法促进了多样化的探索和可泛化的推理轨迹。我们在数学和代码推理任务上进行了实验:FlowRL在数学基准上比GRPO平均提高了10.0%,比PPO提高了5.1%,并在代码推理任务上表现始终优于PPO。这些结果凸显了奖励分布匹配是LLM强化学习中高效探索和多样化推理的关键一步。
查看 arXiv 页面查看 PDF
FlowRL: 匹配LLM推理的奖励分布

评论

Daixuan ChengDaixuan Cheng
论文作者
论文提交者
  • 我们提出了 FlowRL,一种策略优化算法,它将奖励最大化转移到通过流平衡进行奖励分布匹配,从而鼓励探索多样化的推理路径,同时解决了现有 RL 方法固有的模式崩溃限制。
  • 我们引入了长度归一化和重要性采样,以便在可变长度的 CoT 推理上进行有效训练,解决了梯度爆炸和采样不匹配问题。
  • FlowRL 在数学基准测试中分别比 GRPO 和 PPO 高出 10.0% 和 5.1%,并在代码推理任务上展现出强大的泛化能力,多样性分析证实了更广泛的解决方案探索。
AndrewAndrew
在仅解码器架构中,Z_φ(x) 的 MLP 应输入什么?

分区函数 Z_φ(x) 被实现为一个 3 层 MLP,以提示表示 x 作为输入。对于仅解码器模型,MLP 到底应该输入什么?

选项:
  1. 最后一个提示 token - 生成开始前最后一个 token 的隐藏状态
  2. 提示池化 - 对所有提示 token 隐藏状态进行平均/最大池化
  3. 分隔符 token - 在提示和响应之间添加特殊 token

哪种方法对于此用例最常见?

zhuzhu
论文作者

对于 Log_Z 的困惑部分,我深表歉意!我将为您详细解释,并尽快更新我们的论文。

从流程的角度来看:LogZ 衡量从初始状态 S0 的概率流。直观地说,它估计一个分母——所有可能路径的总奖励,因此我们可以通过 reward/Z 将其转换为一个分布。

从实现的视角来看:由于这是初始状态,我们使用 LM 最后一层隐藏状态编码的提示。为了将可变长度的提示转换为标量,我们经验性地取了平均值。这里肯定还有其他我们尚未探索的方法。

zhuzhu
论文作者

这里有一篇讨论此实现的帖子,链接如下:https://x.com/zhu_xuekai/with_replies。
我们收到了关于使用提示的最后一个 token 来节省计算量的建议——我们将在规模化实验中尝试这一点。

欢迎进一步讨论!

AndrewAndrew

感谢您对 Z_φ(x) 实现的澄清!