FlowRL: Alineación de Distribuciones de Recompensa para el Razonamiento en Modelos de Lenguaje

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

Resumen

Proponemos FlowRL: igualar la distribución completa de recompensas mediante el equilibrio de flujos en lugar de maximizar las recompensas en el aprendizaje por refuerzo (RL) de modelos de lenguaje grandes (LLM). Los modelos avanzados de razonamiento recientes adoptan métodos de maximización de recompensas (por ejemplo, PPO y GRPO), que tienden a sobreoptimizar las señales de recompensa dominantes mientras descuidan rutas de razonamiento menos frecuentes pero válidas, reduciendo así la diversidad. En contraste, transformamos las recompensas escalares en una distribución objetivo normalizada utilizando una función de partición aprendible, y luego minimizamos la divergencia KL inversa entre la política y la distribución objetivo. Implementamos esta idea como un método de optimización equilibrado por flujos que promueve una exploración diversa y trayectorias de razonamiento generalizables. Realizamos experimentos en tareas de razonamiento matemático y de código: FlowRL logra una mejora promedio significativa del 10.0% sobre GRPO y del 5.1% sobre PPO en benchmarks matemáticos, y se desempeña consistentemente mejor en tareas de razonamiento de código. Estos resultados destacan la igualación de la distribución de recompensas como un paso clave hacia una exploración eficiente y un razonamiento diverso en el aprendizaje por refuerzo de LLM.

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.