Source-linked AI summary
FlowRL: Matching Reward Distributions for LLM Reasoning
Xuekai Zhu, Daixuan Cheng, Dinghuai Zhang, Hengli Li, Kaiyan Zhang, Che Jiang, Youbang Sun, Ermo Hua, Yuxin Zuo, Xingtai Lv, Qizheng Zhang, Lin Chen, Fanghao Shao, Bo Xue, Yunchong Song, Zhenjie Yang, Ganqu Cui, Ning Ding, Jianfeng Gao, Xiaodong Liu, Bowen Zhou, Hongyuan Mei, Zhouhan Lin
TL;DR
Reward-maximizing RL can collapse LLM reasoning onto dominant solution modes, limiting diversity. FlowRL matches a normalized reward distribution through reverse-KL and flow balancing, achieving consistent gains over PPO and GRPO on math and code tasks while generating more varied reasoning trajectories.
Problem
Reward-maximizing RL tends to overfit dominant reward modes, limiting reasoning-path diversity and generalization to less frequent valid solutions.
Method
FlowRL uses a learnable partition function to normalize scalar rewards into a target distribution and minimizes reverse KL through a trajectory-balance objective, with length normalization and importance sampling for long chain-of-thought training.
Results
FlowRL improves over GRPO by 10.0% and PPO by 5.1% on math benchmarks, surpasses both on code reasoning tasks, and produces more diverse solution approaches.
Takeaways & Limitations
Reward distribution matching provides an effective approach for improving exploration and diversity in LLM reasoning RL across math and code tasks.
Takeaways & Limitations
Reverse KL is used because the policy can be sampled but the target reward distribution cannot be sampled directly.
Abstract
from arXiv · showhide
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 $10.0\%$ over GRPO and $5.1\%$ 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. Introduction
FlowRL addresses mode collapse in reward-maximizing RL for LLM reasoning by matching the full reward distribution and encouraging diverse exploration. Across math and code tasks, it improves over PPO and GRPO while producing more diverse reasoning paths.
- Reward-maximizing RL can overfit dominant reward modes, limiting reasoning-path diversity and generalization to less frequent valid solutions.The problem is especially pronounced for long chain-of-thought reasoning, where diverse plausible solutions matter.
- FlowRL matches the full reward distribution instead of maximizing rewards, using a learnable partition function and reverse KL divergence to encourage mode coverage.The method reframes exploration as distribution matching and derives its objective from GFlowNet trajectory balance.
- FlowRL combines length normalization and importance sampling to address gradient explosion and sampling mismatch in variable-length chain-of-thought training.
- FlowRL outperforms GRPO by 10.0% and PPO by 5.1% across six challenging math benchmarks, and surpasses both on three coding benchmarks.The evaluation uses mainstream RL baselines, base and distilled models, and both math and code domains.
- Diversity analysis finds that FlowRL generates substantially more diverse rollouts than baseline methods, supporting exploration of multiple solution strategies.
2. Preliminaries
The preliminaries formulate LLM reasoning as conditional policy optimization and contrast REINFORCE, PPO, and GRPO with GFlowNets. GFlowNets transport flow through states so terminal objects are sampled in proportion to reward.
- LLM reasoning is modeled as conditional generation, with a policy producing answers under task-specific rewards.
- GRPO samples answer groups from an old policy and updates the model using normalized within-group reward advantages.
- REINFORCE uses direct policy gradients, while PPO adds a critic and importance sampling; GRPO removes value functions and relies on group comparisons.
- GFlowNets balance forward and backward flows so stochastic policies sample discrete objects in proportion to their rewards.The initial flow enters at the initial state, is transported through intermediate states, and the terminal flow equals the outcome reward.
3. Methodology
FlowRL replaces reward maximization with reward-distribution matching: a learnable partition function converts scalar rewards into a target distribution, optimized through a trajectory-balance objective. Length normalization and importance sampling make this approach practical for long chain-of-thought training.
- 3.1. From Reward Maximization to Distribution Matching: FlowRL aligns policy outputs with a reward-induced distribution, addressing dominant-mode collapse while preserving diverse high-reward trajectories.The target is optimized with reverse KL because the policy can be sampled but the target reward distribution cannot.
- 3.1. From Reward Maximization to Distribution Matching: A learnable partition function Z_ϕ(x) normalizes exp(βr(x,y)) into a valid target distribution for reverse-KL optimization.
- 3.1. From Reward Maximization to Distribution Matching: Minimizing the reverse-KL objective is gradient-equivalent to GFlowNet trajectory balance, enabling a stable squared-loss surrogate with a learned partition function.
- 3.2. FlowRL: Long chain-of-thought sequences can cause exploding gradients because sequence-level log probabilities decompose into token-wise sums whose gradient norms scale with length.
- 3.2. FlowRL: Importance sampling reweights stale trajectories by the policy ratio to address mismatch between off-policy rollouts and the current policy.The gradient through the current-policy ratio is detached to prevent excessive policy drift.
- 3.2. FlowRL: FlowRL incorporates a reference-model prior and group-normalized outcome rewards into its practical objective.
- 3.2. FlowRL: Length normalization rescales sequence log probabilities, balancing long and short responses and stabilizing the learning signal.
- 3.2. FlowRL: The resulting objective updates policy parameters using clipped importance weights and normalized rewards, with implementation details provided separately.
4. Experiment Settings
The experiments use matched RL baselines and evaluate FlowRL across math and code reasoning with multiple model scales and benchmarks.
- Baselines: FlowRL is compared with REINFORCE++, PPO, and GRPO using the same learning rate, batch size, training steps, and convergence criteria.
- Training Configuration: The study trains on DAPO math data and DeepCoder code data, using Qwen-2.5 models for math and DeepSeek-R1-Distill-Qwen-7B for code.
- Evaluation Configuration: Evaluation covers six math benchmarks and three code benchmarks, reporting average accuracy over 16 rollouts as Avg@16.
- Evaluation Configuration: Tables 1 and 2 report Avg@16 math results across 7B and 32B scales and metrics for all three code benchmarks.
5. Results
FlowRL consistently outperforms reward-maximization baselines across math and code reasoning, with gains reported across model scales and benchmarks.
- Main Results: FlowRL consistently outperforms all reward-maximization baselines across math and code reasoning domains.The comparison spans both model scales and multiple benchmark families.
- Main Results: 35.6% and 48.4% are FlowRL’s highest average math accuracies with 7B and 32B models, respectively.
- Main Results: 5.1% and 10.1% are FlowRL’s reported improvements over PPO and GRPO on the 32B math model, respectively.
- Main Results: 37.43% Avg@16 on LiveCodeBench, 1549.47 on Codeforces, and 83.28% on HumanEval+ are FlowRL’s reported code results.The passage states that FlowRL outperforms all baselines on these benchmarks.
- Ablations: Importance sampling substantially improves FlowRL performance across all reported math reasoning benchmarks.
6. Analysis
The analysis links FlowRL’s performance to diversity: its tuned β setting and flow-balanced optimization produce more varied reasoning paths than reward-maximizing baselines.
- Diversity Analysis: β=15 achieves the best performance in FlowRL’s hyperparameter ablation.
- Diversity Analysis: FlowRL’s GPT-judged diversity score nearly doubles that of PPO on AIME 24/25 rollouts.The evaluation uses GPT-4o-mini and compares FlowRL with baseline methods.
- Diversity Analysis: In a representative AIME case, GRPO repeats AM-GM and identity-loop patterns, while FlowRL derives a cubic equation and reaches the correct answer.
7. Related Work
Related work covers RL for LLM reasoning, diversity-oriented flow methods, and approaches that address mode collapse through entropy or reward-distribution matching.
- Reinforcement Learning for Reasoning: Prior LLM reasoning methods include REINFORCE, PPO, and GRPO, while entropy regularization promotes diversity but can be difficult to influence long CoT trajectories.
- Reinforcement Learning for Reasoning: FlowRL differs from entropy-based approaches by shifting the RL objective from reward maximization to reward distribution matching.
- GFlowNets: GFlowNets balance forward and backward probability flows across states and sample discrete objects in proportion to reward.
- Flow-Matching Policies: Flow-matching policy work uses learned transport or advantage-weighted ratios to reformulate policy optimization without expensive likelihood computations.
8. Conclusion
FlowRL matches the policy to a normalized reward distribution through a learnable partition function and reverse-KL minimization. Across math and code benchmarks, it improves reasoning performance and produces more varied solution approaches than reward-maximizing baselines.
- FlowRL transforms scalar rewards into normalized target distributions and minimizes reverse KL divergence between the policy and target distribution.
- FlowRL is theoretically equivalent to trajectory-balance objectives and implicitly maximizes both reward and entropy.
- FlowRL consistently improves performance over GRPO and PPO across math and code reasoning benchmarks.
- Diversity analyses and case studies show that FlowRL generates more varied solution approaches while avoiding repetitive patterns.
A. Proof of Proposition 1
The proof establishes that the reverse-KL objective used by FlowRL has the same expected gradients as the GFlowNet trajectory-balance objective. This equivalence provides the theoretical bridge between distribution matching and trajectory-balance optimization.
- The proof takes gradients of the KL divergence and the trajectory-balance objective to compare their optimization directions.
- Minimizing the KL divergence is equivalent up to a constant to minimizing the trajectory-balance loss.
B. Theoretical Analysis
FlowRL connects reverse-KL reward distribution matching to GFlowNet trajectory balance and interprets the resulting objective as joint reward and entropy maximization. The method adds length normalization, importance sampling, and a learnable partition function to support long-CoT training and practical optimization.
- Theoretical Analysis: FlowRL can be interpreted as jointly maximizing expected reward and policy entropy, encouraging broader exploration of high-quality reasoning solutions.
- Theoretical Analysis: The objective aligns the policy with the full reward-weighted distribution instead of collapsing onto a few high-reward modes.
- Theoretical Analysis: FlowRL builds on GFlowNets, which use flow-consistency objectives at detailed, trajectory, and sub-trajectory levels to support diversity-seeking generation.
- Theoretical Analysis: FlowRL extends trajectory balance with length normalization and importance sampling to address gradient explosion and sampling mismatch in long-CoT reasoning.
- Theoretical Analysis: The learnable partition function and reverse-KL formulation distinguish FlowRL from reward-maximizing and related KL-regularized methods.
E. Implementation of Partition Function 𝑍𝜙
The partition function estimates the initial-state probability flow and normalizes rewards across possible paths. In implementation, it receives a prompt representation formed by averaging final-layer hidden states across prompt tokens.
- The partition function estimates the denominator formed by rewards across possible paths, converting rewards into a probability distribution.
- FlowRL computes the partition function from the mean of final-layer hidden states over all prompt tokens.
- The section includes a diversity-evaluation prompt that rates variation in solution approaches from 1 to 5.