PaperlayerView exact revision on arXiv

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

arXiv:2509.15207v3cs.LGcs.AIcs.CL

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.

    Source: §1 · p. 2

  • 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.

    Source: §1 · p. 2, §8 · pp. 10–11

  • 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.

    Source: §1 · pp. 2–3, §8 · p. 11

  • Takeaways & Limitations

    Reward distribution matching provides an effective approach for improving exploration and diversity in LLM reasoning RL across math and code tasks.

    Source: §1 · p. 2, §8 · p. 11

    Reverse KL is used because the policy can be sampled but the target reward distribution cannot be sampled directly.

    Source: §3.1 · p. 4

Abstract

from arXiv · show

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.Source: §1 · p. 2
  • 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.Source: §1 · p. 2
  • FlowRL combines length normalization and importance sampling to address gradient explosion and sampling mismatch in variable-length chain-of-thought training.Source: §1 · pp. 2–3
  • 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.Source: §1 · pp. 2–3
  • Diversity analysis finds that FlowRL generates substantially more diverse rollouts than baseline methods, supporting exploration of multiple solution strategies.Source: §1 · pp. 2–3

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.Source: §2 · p. 3
  • GRPO samples answer groups from an old policy and updates the model using normalized within-group reward advantages.Source: §2 · p. 3
  • REINFORCE uses direct policy gradients, while PPO adds a critic and importance sampling; GRPO removes value functions and relies on group comparisons.Source: §2 · p. 3
  • 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.Source: §2 · p. 3, Figure 2 · p. 3

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.Source: §3.1 · p. 4
  • A learnable partition function Z_ϕ(x) normalizes exp(βr(x,y)) into a valid target distribution for reverse-KL optimization.Source: §3.1 · p. 4
  • Minimizing the reverse-KL objective is gradient-equivalent to GFlowNet trajectory balance, enabling a stable squared-loss surrogate with a learned partition function.Source: §3.1 · p. 4

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.Source: §3.2 · p. 5
  • 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.Source: §3.2 · p. 5
  • FlowRL incorporates a reference-model prior and group-normalized outcome rewards into its practical objective.Source: §3.2 · p. 5
  • Length normalization rescales sequence log probabilities, balancing long and short responses and stabilizing the learning signal.Source: §3.2 · p. 5
  • The resulting objective updates policy parameters using clipped importance weights and normalized rewards, with implementation details provided separately.Source: §3.2 · p. 6

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.Source: §4 · p. 6

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.Source: §4 · p. 6

Evaluation Configuration

  • Evaluation covers six math benchmarks and three code benchmarks, reporting average accuracy over 16 rollouts as Avg@16.Source: §4 · p. 6
  • Tables 1 and 2 report Avg@16 math results across 7B and 32B scales and metrics for all three code benchmarks.Source: Table 1 · p. 7, Table 2 · p. 7

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.Source: §5.1 · pp. 7–8
  • 35.6% and 48.4% are FlowRL’s highest average math accuracies with 7B and 32B models, respectively.Source: §5.1 · p. 7
  • 5.1% and 10.1% are FlowRL’s reported improvements over PPO and GRPO on the 32B math model, respectively.Source: §5.1 · p. 7
  • 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.Source: §5.1 · p. 7

Ablations

  • Importance sampling substantially improves FlowRL performance across all reported math reasoning benchmarks.Source: §5.2 · p. 8, Table 3 · p. 8

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.Source: Figure 3 · p. 8
  • 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.Source: §6.1 · p. 8
  • In a representative AIME case, GRPO repeats AM-GM and identity-loop patterns, while FlowRL derives a cubic equation and reaches the correct answer.Source: §6.2 · p. 9, Table 4 · p. 9

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.Source: §7 · pp. 9–10
  • FlowRL differs from entropy-based approaches by shifting the RL objective from reward maximization to reward distribution matching.Source: §7 · p. 10

GFlowNets

  • GFlowNets balance forward and backward probability flows across states and sample discrete objects in proportion to reward.Source: §7 · p. 10

Flow-Matching Policies

  • Flow-matching policy work uses learned transport or advantage-weighted ratios to reformulate policy optimization without expensive likelihood computations.Source: §7 · p. 10

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.Source: §8 · p. 10
  • FlowRL is theoretically equivalent to trajectory-balance objectives and implicitly maximizes both reward and entropy.Source: §8 · p. 11
  • FlowRL consistently improves performance over GRPO and PPO across math and code reasoning benchmarks.Source: §8 · p. 11
  • Diversity analyses and case studies show that FlowRL generates more varied solution approaches while avoiding repetitive patterns.Source: §8 · p. 11

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.

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

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.

Related papers

Loading 2509.15207v3…