FlowRL: Matching Reward Distributions for LLM Reasoning
TL;DR 精炼摘要
本文提出了一种新方法FlowRL,通过流平衡匹配完整的奖励分布,以替代传统的奖励最大化策略。与传统方法不同,FlowRL促进多样化探索和更具泛化性的推理轨迹。在数学和代码推理任务的实验中,FlowRL分别比GRPO和PPO提高了10%和5.1%的性能,显示了奖励分布匹配在LLM强化学习中的关键作用。
摘要
We propose FlowRL: matching the full reward distribution via flow balancing instead of maximizing rewards in large language model (LLM) reinforcement learning (RL). Recent advanced reasoning models adopt reward-maximizing methods (\eg, PPO and GRPO), which tend to over-optimize dominant reward signals while neglecting less frequent but valid reasoning paths, thus reducing diversity. In contrast, we transform scalar rewards into a normalized target distribution using a learnable partition function, and then minimize the reverse KL divergence between the policy and the target distribution. We implement this idea as a flow-balanced optimization method that promotes diverse exploration and generalizable reasoning trajectories. We conduct experiments on math and code reasoning tasks: FlowRL achieves a significant average improvement of over GRPO and over PPO on math benchmarks, and performs consistently better on code reasoning tasks. These results highlight reward distribution-matching as a key step toward efficient exploration and diverse reasoning in LLM reinforcement learning.
思维导图
论文精读
中文精读
1. 论文基本信息
1.1. 标题
FlowRL: Matching Reward Distributions for LLM Reasoning (FlowRL:通过匹配奖励分布实现大语言模型推理)
1.2. 作者
Xuekai Zhu, Daixuan Cheng, Dinghuai Zhang 等。 主要作者来自 上海交通大学 (Shanghai Jiao Tong University)、上海人工智能实验室 (Shanghai AI Laboratory)、微软研究院 (Microsoft Research)、清华大学 (Tsinghua University) 等机构。 通讯作者为 Hongyuan Mei 和 Zhouhan Lin。
1.3. 发表期刊/会议
arXiv (预印本) 发布日期:2025年9月18日 (根据元数据)。
1.4. 摘要
本文提出了 FlowRL,一种用于大语言模型(LLM)的新型强化学习(RL)方法。传统的推理模型通常采用奖励最大化方法(如 PPO 和 GRPO),这往往会导致模型过度优化单一的高奖励路径(Mode Collapse,模式坍塌),从而忽略了其他虽不常见但依然正确的推理路径,降低了生成的多样性。FlowRL 的核心思想是通过一个可学习的配分函数(Partition Function),将标量奖励转化为归一化的目标分布,并最小化策略与目标分布之间的反向 KL 散度。作者通过流平衡(Flow Balance)实现了这一思想,并引入了长度归一化和重要性采样来解决长链推理中的训练不稳定性。实验表明,FlowRL 在数学和代码推理任务上显著优于 GRPO 和 PPO,能够生成更多样化且高质量的推理轨迹。
1.5. 原文链接
2. 整体概括
2.1. 研究背景与动机
在提升大语言模型(LLM)推理能力的后训练阶段,强化学习(RL)扮演着至关重要的角色。目前主流的方法如 PPO(Proximal Policy Optimization)和 GRPO(Group Relative Policy Optimization)均基于 奖励最大化(Reward Maximization) 的目标。
然而,这种目标存在一个根本性的缺陷:模式坍塌(Mode Collapse)。模型倾向于“走捷径”,只收敛到某一种特定的、能获得高奖励的解题模式,而忽略了其他同样有效但路径不同的解法。这种单一性降低了模型的泛化能力,特别是在复杂的长思维链(Chain-of-Thought, CoT)推理中,探索多样化的解决方案对于解决难题至关重要。
下图(原文 Figure 1)形象地展示了这一问题:右侧的 GRPO 仅关注单一的高峰(High-reward peak),而左侧的 FlowRL 能够覆盖多个有效的奖励模式。
该图像是图表,展示了FlowRL与传统奖励最大化方法(GRPO)的对比。上部分为奖励分布的匹配与最大化方法的KL散度比较,FlowRL的KL为0.11,而GRPO的KL为8.68,显示出FlowRL在多模态上保持了更高的多样性。下部分为两种方法在数学和代码评分上的表现,FlowRL的得分优于GRPO。
2.2. 核心贡献
- 范式转变: 提出 FlowRL,将 RL 的目标从“最大化标量奖励”转变为“匹配奖励分布”。利用流平衡(Flow Balance)原理,鼓励模型按照奖励的大小比例来生成多样化的推理路径。
- 理论与实践结合: 证明了最小化反向 KL 散度与 GFlowNets 中的 轨迹平衡(Trajectory Balance) 目标在梯度上的等价性,并针对 LLM 长文本生成的特点,提出了 长度归一化(Length Normalization) 和 重要性采样(Importance Sampling) 两项关键技术改进。
- 实验效果: 在数学任务(如 MATH-500, AIME)上,FlowRL 相比 GRPO 平均提升 10.0%,相比 PPO 提升 5.1%;在代码任务上也取得了一致的优胜,且多样性分析证实其生成的解法更加丰富。
3. 预备知识与相关工作
3.1. 基础概念
为了理解 FlowRL,初学者需要掌握以下几个核心概念:
- 强化学习中的策略 (Policy, ): 即 LLM 本身。它根据给定的问题 ,生成回答 的概率分布。
- 思维链 (Chain-of-Thought, CoT): LLM 在给出最终答案前,生成的一系列中间推理步骤。本文关注的是长 CoT 推理。
- 模式坍塌 (Mode Collapse): 在生成模型训练中常见的问题。比如有 10 种解数学题的方法,模型发现其中一种最简单,之后就只生成这一种,完全“遗忘”了其他 9 种。这会导致模型面对稍有变化的题目时无法举一反三。
- KL 散度 (Kullback-Leibler Divergence): 衡量两个概率分布之间差异的指标。本文使用的是 反向 KL 散度 (Reverse KL),用于让策略分布 去逼近目标分布。
3.2. GFlowNets (生成流网络)
这是理解本文方法论的理论基石。
-
核心思想: 传统的 RL 试图找到奖励最高的那条路;而 GFlowNet 试图让生成某条路径的概率,与该路径获得的奖励成正比()。
-
流平衡 (Flow Balance): 想象水流流过一个管道网络。GFlowNet 要求进入某个状态的“流量”总和必须等于流出该状态的“流量”总和。
-
轨迹平衡 (Trajectory Balance): GFlowNet 的一种训练目标,要求整条生成轨迹上的前向概率流与后向概率流保持平衡。
下图(原文 Figure 2)展示了 GFlowNets 的流平衡视角:初始流 注入环境,通过策略 分配到不同的路径,最终在终态积累成与奖励成比例的分布。
该图像是示意图,展示了GFlowNets在强化学习中的流平衡视角。图中展示了初始流注入概率质量到环境,并通过策略在中间状态间传输,最终以与标量奖励成比例的方式在终态累积。状态及其对应的流向显示了在强化学习中的流动路径和奖励积累过程。
3.3. 现有方法的局限
- PPO: 经典的 RL 算法,稳定但计算复杂,且作为奖励最大化方法,容易陷入局部最优。
- GRPO: DeepSeek 提出的方法,去掉了 PPO 中的价值函数(Critic),通过组内相对奖励来更新策略。虽然节省了资源,但依然是奖励最大化导向,容易导致多样性丧失。
4. 方法论
4.1. 方法原理:从最大化到匹配
FlowRL 的根本目标是不再盲目追求最高分,而是构建一个目标分布,使得回答 出现的概率与它获得的奖励 的指数成正比。
为了实现这一点,作者引入了一个可学习的 配分函数 (Partition Function) 。这是一个归一化常数,确保所有可能回答的概率之和为 1。
我们的目标是最小化策略 与这个目标分布之间的反向 KL 散度: 其中:
- : 输入的问题。
- : 生成的回答(包含 CoT)。
- : 奖励函数。
- : 一个控制分布平滑程度的超参数(温度系数的倒数)。
- : 一个由神经网络预测的标量,代表给定问题 下所有可能路径的奖励总和的估计。
4.2. 核心方法详解:FlowRL 的推导与实现
作者通过以下步骤将上述理论转化为可训练的 FlowRL 算法:
4.2.1. 步骤一:转化为轨迹平衡 (Trajectory Balance)
直接优化 KL 散度很难,因为我们无法遍历所有可能的 来计算 。作者利用 Proposition 1 证明了,在梯度期望意义下,最小化上述 KL 散度等价于最小化 GFlowNet 中的 轨迹平衡损失:
直观理解: 这个公式在说,模型预测的“总奖励池” () 加上“选择这条路径的概率” (),应该等于“这条路径实际获得的奖励” ()。如果这个等式成立,就达到了流平衡。
4.2.2. 步骤二:引入参考模型与奖励重塑
为了防止模型在训练中遗忘原有的语言能力,作者在奖励函数中加入了一个 参考模型 (Reference Model, ) 的约束。这类似于 PPO 中的 KL 惩罚,但在 FlowRL 中,它是直接乘在奖励目标上的: 将这一项代入轨迹平衡公式,原来的奖励项 变成了 。此外,作者对奖励 进行了组内归一化(Group Normalization),记为 。
此时的损失函数变为:
4.2.3. 步骤三:解决长文本梯度爆炸 (长度归一化)
问题: 在长思维链推理中, 可能包含数千个 Token。 是所有 Token 对数概率的总和,这会导致该项数值极大,引发梯度爆炸。 解决方案: 作者引入了 长度归一化,即除以序列长度 。这使得长短序列的贡献相对平衡。
4.2.4. 步骤四:解决采样不匹配 (重要性采样)
问题: 轨迹平衡损失理论上要求数据必须是从当前策略 采样的(On-policy)。但为了提高效率,RL 通常使用经验回放(Replay Buffer)或微批次更新,这导致数据可能来自稍旧的策略 。 解决方案: 借鉴 PPO,引入 重要性采样 (Importance Sampling, IS) 权重 来修正这种分布偏差,并对其进行截断(Clip)以保证稳定性。
4.2.5. 最终 FlowRL 目标函数
综合以上所有改进,FlowRL 的最终损失函数如下:
其中, 和 的定义为:
符号解释:
- : 重要性采样权重,带有梯度截断(detach),仅用于加权损失,不反向传播梯度到权重本身。
- : 截断范围超参数(如 0.1),防止权重过大。
- : 生成序列的长度(Token 数)。
- : 采样时使用的旧策略。
- : 对一个采样组(Group)内的奖励进行标准化后的值。
5. 实验设置
5.1. 数据集
实验覆盖了 数学 和 代码 两个领域的推理任务:
- 数学领域:
- 数据集来自 DAPO 收集的训练集。
- 评估基准 (Benchmarks): 包含 6 个高难度数据集:AIME 2024/2025(美国数学邀请赛)、AMC 2023、MATH-500、Minerva、OlympiadBench。
- 代码领域:
- 训练集来自 DeepCoder。
- 评估基准: LiveCodeBench, CodeForces, HumanEval+。
5.2. 评估指标
- Avg@16 (Average Accuracy @ 16 samples):
- 定义: 对每个问题进行 16 次采样(推演),计算这 16 次回答的平均准确率。这不仅反映了模型解决问题的能力,也侧面反映了模型输出的稳定性。
- Pass@16: (主要用于代码) 16 次采样中至少有一次通过测试的概率。
- CodeForces Rating/Percentile: 模拟在 CodeForces 竞赛中的评分和排名百分位。
- 多样性评分 (Diversity Score):
- 使用 GPT-4o-mini 作为裁判,对模型针对同一问题生成的 16 个不同解法进行 1-5 分的多样性打分。1 分代表高度重复,5 分代表探索了完全不同的解题策略。
5.3. 对比基线
- REINFORCE++ (R++): 基础的策略梯度算法。
- PPO (Proximal Policy Optimization): 工业界最常用的 RL 微调算法,引入了 Critic 网络。
- GRPO (Group Relative Policy Optimization): DeepSeek 推出的高效算法,去除了 Critic,利用组内相对优势进行更新。
6. 实验结果与分析
6.1. 核心结果分析
实验结果表明,FlowRL 在不同规模(7B, 32B)的模型上,无论是数学还是代码任务,均一致优于所有基线模型。
以下是原文 Table 1 的数学推理基准测试结果:
| Models | AIME24 | AIME25 | AMC23 | MATH500 | Minerva | Olympiad | Avg |
|---|---|---|---|---|---|---|---|
| Qwen2.5-32B-Base, Max Response Len = 8K tokens | |||||||
| Backbone | 4.58 | 2.08 | 28.59 | 52.48 | 26.99 | 21.37 | 22.68 |
| R++ | 14.79+10.21 | 9.17+7.08 | 52.65+24.06 | 44.35-8.13 | 17.37-9.62 | 24.52+3.15 | 27.14 |
| PPO | 26.87+22.29 | 20.41+18.33 | 76.40+47.81 | 69.17+16.69 | 28.79+1.80 | 37.90+16.53 | 43.25 |
| GRPO | 23.12+18.54 | 14.58+12.50 | 76.87+48.28 | 61.60+9.12 | 18.95-8.04 | 34.94+13.57 | 38.34 |
| FlowRL | 23.95+19.37 | 21.87+19.79 | 73.75+45.16 | 80.75+28.27 | 38.21+11.22 | 51.83+30.46 | 48.39 |
| Qwen2.5-7B-Base, Max Response Len = 8K tokens | |||||||
| Backbone | 4.38 | 2.08 | 30.78 | 54.47 | 22.38 | 24.03 | 23.02 |
| R++ | 11.04+6.66 | 5.41+3.33 | 66.71+35.93 | 54.25-0.22 | 24.37+1.99 | 27.33+3.30 | 31.52 |
| PPO | 9.38+5.00 | 7.29+5.21 | 63.43+32.65 | 57.98+3.51 | 26.53+4.15 | 27.25+3.22 | 31.98 |
| GRPO | 13.54+9.16 | 9.79+7.71 | 64.53+33.75 | 57.05+2.58 | 23.06+0.68 | 26.88+2.85 | 32.48 |
| FlowRL | 15.41+11.03 | 10.83+8.75 | 54.53+23.75 | 66.96+12.49 | 31.41+9.03 | 34.61+10.58 | 35.63 |
以下是原文 Table 2 的代码任务基准测试结果:
| Models | LiveCodeBench | CodeForces | HumanEval+ | ||
|---|---|---|---|---|---|
| Avg@16 | Pass@16 | Rating | Percentile | Avg@16 | |
| DeepSeek-R1-Distill-Qwen-7B, Max Response Len = 8K tokens | |||||
| Backbone | 30.68 | 49.46 | 886.68 | 19.4% | 80.90 |
| R++ | 30.46-0.22 | 52.68+3.22 | 1208.03+321.35 | 56.8%+37.4% | 76.61-4.29 |
| PPO | 35.10+4.42 | 54.48+5.02 | 1403.07+516.39 | 73.7%+54.3% | 82.32+1.42 |
| GRPO | 32.75+2.07 | 52.32+2.86 | 1313.82+427.14 | 67.1%+47.7% | 80.13-0.77 |
| FlowRL | 37.43+6.75 | 56.27+6.81 | 1549.47+662.79 | 83.3%+63.9% | 83.28+2.38 |
分析:
- 数学任务: FlowRL 在 32B 模型上的平均得分为 48.39,显著高于 PPO (43.25) 和 GRPO (38.34)。特别是在 MATH-500 和 Olympiad 这类难题上,优势尤为明显。
- 代码任务: FlowRL 在 CodeForces 上达到了 1549.47 的 Rating(超越基线 600 多分),证明了其强大的泛化能力。
6.2. 多样性与探索
FlowRL 的核心优势在于“多样性”。通过 GPT-4o-mini 的打分,作者展示了各方法生成解法的多样性。
下图(原文 Figure 4)清晰地表明,FlowRL 的多样性得分(2.28)几乎是 GRPO(1.23)和 R++(1.11)的两倍。这验证了 FlowRL 成功避免了模式坍塌,探索了更多不同的解题路径。
该图像是一个柱状图,展示了不同方法生成的多样性得分。FlowRL的多样性得分为2.28,显著超过了R++(1.11)、GRPO(1.23)和PPO(1.31),表明FlowRL在生成多样化解决方案方面表现更佳。
6.3. 消融实验
作者验证了两个关键点:
-
重要性采样 (Importance Sampling, IS) 的必要性: 如下表 3 所示,如果去掉 IS(w/o IS),性能从 35.63 骤降至 26.71。这说明修正采样分布偏差对于 FlowRL 在 On-policy/Off-policy 混合训练中至关重要。
-
超参数 : 下图(原文 Figure 3)显示, 是最佳选择。过小会导致奖励信号太弱,过大则可能退化回奖励最大化(Overfitting)。
该图像是一个柱状图,展示了不同 eta值(5, 10, 15, 30)对应的平均得分(%)。其中,eta = 15 的得分最高,达到 35.63,明显优于其他eta值。
7. 总结与思考
7.1. 结论总结
FlowRL 通过引入 GFlowNet 的流平衡思想,成功将 LLM 的强化学习目标从“寻找最高分”转变为“匹配奖励分布”。配合长度归一化和重要性采样等技术改进,FlowRL 不仅解决了长思维链训练中的梯度和采样难题,还在数学和代码推理任务上实现了比 GRPO 和 PPO 更优越的性能和更高的解题多样性。
7.2. 局限性与未来工作
- 计算开销: 虽然文章未详细展开,但引入额外的配分函数网络 会增加一定的训练显存和计算开销。
- 理论与工程的差距: 长度归一化虽然在工程上解决了梯度爆炸,但在理论上是对轨迹平衡方程的一种启发式修改(Heuristic),如何从理论上更严格地统一这两者仍值得探索。
7.3. 个人启发与批判
- GFlowNet 的潜力: 这篇论文是 GFlowNet 思想在 LLM 大规模推理任务上的成功应用案例。它证明了在生成任务中,“多样性”本身就是一种极具价值的“性能”,尤其是在只有稀疏标量奖励(Scalar Reward)的场景下。
- 对 RLHF 的反思: 现有的 RLHF/RLAIF 往往导致模型输出趋同(Mode Collapse)。FlowRL 提供了一种新的视角:如果我们希望 AI 像人类一样具有创造性思维,我们的训练目标就不应该仅仅是逼近那个唯一的“标准答案”,而是逼近整个可行解的分布。
相似论文推荐
基于向量语义检索推荐的相关论文。