FlowRL: Correspondência de Distribuições de Recompensa para Raciocínio em LLMs

FlowRL: Matching Reward Distributions for LLM Reasoning

September 18, 2025

Autores: 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

cs.AI

Resumo

Propomos o FlowRL: correspondência da distribuição completa de recompensas por meio de balanceamento de fluxo, em vez de maximizar recompensas no aprendizado por reforço (RL) de modelos de linguagem de grande escala (LLM). Modelos avançados de raciocínio recentes adotam métodos de maximização de recompensas (\eg, PPO e GRPO), que tendem a otimizar excessivamente sinais de recompensa dominantes enquanto negligenciam caminhos de raciocínio menos frequentes, porém válidos, reduzindo assim a diversidade. Em contraste, transformamos recompensas escalares em uma distribuição alvo normalizada usando uma função de partição aprendível e, em seguida, minimizamos a divergência KL reversa entre a política e a distribuição alvo. Implementamos essa ideia como um método de otimização balanceada por fluxo que promove exploração diversificada e trajetórias de raciocínio generalizáveis. Realizamos experimentos em tarefas de raciocínio matemático e de código: o FlowRL alcança uma melhoria média significativa de 10,0% em relação ao GRPO e 5,1% em relação ao PPO em benchmarks matemáticos, e apresenta desempenho consistentemente melhor em tarefas de raciocínio de código. Esses resultados destacam a correspondência da distribuição de recompensas como um passo crucial para exploração eficiente e raciocínio diversificado no aprendizado por reforço de LLMs.

English

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.