FlowRL: Matching Reward Distributions for LLM Reasoning (2509.15207v1)

Abstract: We propose FlowRL: matching the full reward distribution via flow balancing instead of maximizing rewards in 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.

• The paper introduces FlowRL, demonstrating that matching reward distributions overcomes mode collapse in LLM reasoning.
• FlowRL employs flow-balanced optimization with a learnable partition function, length normalization, and importance sampling for stable training.
• Empirical results show up to a 10% improvement in math and code tasks and nearly double the solution diversity compared to traditional methods.

FlowRL: Reward Distribution Matching for LLM Reasoning

Motivation and Limitations of Reward-Maximizing RL

The paper introduces FlowRL, a reinforcement learning algorithm for LLM reasoning that fundamentally shifts the objective from reward maximization to reward distribution matching. Traditional RL methods for LLM post-training—such as PPO and GRPO—optimize expected reward, which leads to overfitting on dominant solution modes and mode collapse, thereby reducing diversity in generated reasoning paths. This limitation is particularly acute in complex chain-of-thought (CoT) tasks, where diverse solution strategies are essential for generalization.

Figure 1: FlowRL matches the full reward distribution, maintaining diversity across multiple modes, while reward-maximizing methods like GRPO collapse to a single high-reward peak and exhibit higher KL divergence. FlowRL consistently outperforms GRPO across math and code domains.

FlowRL: Distribution Matching via Flow-Balanced Optimization

FlowRL reframes RL for LLMs as a distribution matching problem. Instead of maximizing scalar rewards, FlowRL transforms rewards into a normalized target distribution using a learnable partition function $Z_\phi(\mathbf{x})$, and minimizes the reverse KL divergence between the policy $\pi_\theta(\mathbf{y}|\mathbf{x})$ and the reward-induced distribution:

$$\min_\theta \mathcal{D}_{\mathrm{KL}}\left( \pi_\theta(\mathbf{y}|\mathbf{x}) \,\|\, \frac{\exp(\beta r(\mathbf{x}, \mathbf{y}))}{Z_\phi(\mathbf{x})} \right)$$

where $\beta$ is a temperature hyperparameter. This objective encourages the policy to sample trajectories in proportion to their rewards, promoting coverage of multiple solution modes.

The KL objective is shown to be gradient-equivalent to the trajectory balance loss from GFlowNets:

$$\left( \log Z_\phi(\mathbf{x}) + \log \pi_\theta(\mathbf{y}|\mathbf{x}) - \beta r(\mathbf{x}, \mathbf{y}) \right)^2$$

This equivalence enables practical optimization using a stable squared loss and a learnable partition function, bridging generative modeling and policy optimization.

Technical Innovations for Long CoT Reasoning

Applying trajectory balance to long CoT tasks introduces two challenges: (1) exploding gradients due to sequence-level objectives over long trajectories, and (2) sampling mismatch when using off-policy rollouts. FlowRL addresses these with:

• Length normalization: The log-probability term is rescaled by sequence length, $$\frac{1}{|\mathbf{y}|} \log \pi_\theta(\mathbf{y}|\mathbf{x})$$, stabilizing gradients for variable-length outputs.
• Importance sampling: Off-policy rollouts are reweighted using clipped importance ratios, $$w = \text{clip}\left(\frac{\pi_\theta(\mathbf{y}|\mathbf{x})}{\pi_{\text{old}}(\mathbf{y}|\mathbf{x})}, 1-\epsilon, 1+\epsilon\right)$$, with gradients detached to prevent instability.

The final FlowRL objective is:

$$\mathcal{L}_{\text{FlowRL}} = w \cdot \left( \log Z_\phi(\mathbf{x}) + \frac{1}{|\mathbf{y}|} \log \pi_\theta(\mathbf{y}|\mathbf{x}) - \beta \hat{r}(\mathbf{x}, \mathbf{y}) - \frac{1}{|\mathbf{y}|}\log \pi_{\mathrm{ref}}(\mathbf{y}|\mathbf{x}) \right)^2$$

where $\hat{r}$ is group-normalized reward and $\pi_{\mathrm{ref}}$ is a fixed reference policy.

Empirical Results: Math and Code Reasoning

FlowRL is evaluated on six math and three code reasoning benchmarks using Qwen-2.5-7B/32B and DeepSeek-R1-Distill-Qwen-7B backbones. FlowRL achieves a 10.0% improvement over GRPO and 5.1% over PPO on math tasks (32B), and consistently outperforms all baselines on code tasks (e.g., 37.4% Avg@16 on LiveCodeBench, 1549.5 Codeforces rating, 83.3% HumanEval+ accuracy). These results are robust across model scales and temperature settings.

Diversity Analysis and Case Studies

Diversity of generated solutions is quantitatively assessed using GPT-4o-mini, showing that FlowRL nearly doubles diversity scores compared to PPO and GRPO. Case studies on AIME problems reveal that reward-maximizing baselines exhibit repetitive solution patterns (e.g., repeated AM-GM applications), while FlowRL explores alternative strategies (e.g., symmetry assumptions, polynomial factorization), leading to correct answers and broader exploration.

Theoretical Interpretation

Minimizing the FlowRL objective is shown to be equivalent to jointly maximizing expected reward and policy entropy:

$$\max_\theta \; \mathbb{E}_{\mathbf{y} \sim \pi_\theta} \left[ \beta r(\mathbf{x}, \mathbf{y}) - \log Z_\phi(\mathbf{x}) + \log \pi_{\mathrm{ref}}(\mathbf{y}|\mathbf{x}) \right] + \mathcal{H}(\pi_\theta)$$

This interpretation aligns FlowRL with maximum entropy RL, ensuring both high performance and diverse solution coverage.

Implementation Considerations

• Partition function $Z_\phi$: Implemented as a 3-layer MLP, trained jointly with the policy.
• Rollout generation: Group size of 8, batch size 512 (math) or 64 (code), max response length 8K tokens.
• Importance sampling: Essential for stability and data efficiency; ablation studies show substantial performance drops without it.
• Hyperparameter $\beta$: Optimal value found to be 15; ablation confirms sensitivity.

Implications and Future Directions

FlowRL demonstrates that reward distribution matching is a key step for efficient exploration and generalizable reasoning in LLM RL. The approach is theoretically grounded, empirically validated, and practically scalable to long-sequence tasks. The integration of GFlowNet principles into LLM RL opens avenues for further research in diversity-driven policy optimization, structured exploration, and robust generalization. Future work may extend FlowRL to multimodal reasoning, hierarchical planning, and adaptive curriculum learning, leveraging its entropy-maximizing and mode-covering properties.

Conclusion

FlowRL introduces a principled framework for matching reward distributions in LLM reasoning, overcoming the mode-collapse limitations of reward-maximizing RL. By leveraging flow-balanced optimization, length normalization, and importance sampling, FlowRL achieves superior performance and diversity on challenging math and code tasks. Theoretical analysis confirms its joint reward-entropy maximization, and empirical results highlight its effectiveness in promoting diverse, generalizable reasoning. FlowRL sets a new standard for diversity-driven RL in LLMs and provides a foundation for future advances in reasoning model post-training.