Optimización de Políticas por Segmentos: Asignación Efectiva de Crédito a Nivel de Segmento en Aprendizaje por Refuerzo para Modelos de Lenguaje a Gran Escala

Segment Policy Optimization: Effective Segment-Level Credit Assignment in RL for Large Language Models

May 29, 2025

Autores: Yiran Guo, Lijie Xu, Jie Liu, Dan Ye, Shuang Qiu

cs.AI

Resumen

Mejorar las capacidades de razonamiento de los modelos de lenguaje de gran escala utilizando de manera efectiva el aprendizaje por refuerzo (RL) sigue siendo un desafío crucial. Los enfoques existentes adoptan principalmente dos granularidades de estimación de ventajas contrastantes: los métodos a nivel de token (por ejemplo, PPO) buscan proporcionar señales de ventaja detalladas, pero sufren de estimaciones imprecisas debido a las dificultades en entrenar un modelo crítico preciso. En el otro extremo, los métodos a nivel de trayectoria (por ejemplo, GRPO) dependen únicamente de una señal de ventaja de grano grueso proveniente de la recompensa final, lo que lleva a una asignación de crédito imprecisa. Para abordar estas limitaciones, proponemos Segment Policy Optimization (SPO), un novedoso marco de RL que aprovecha la estimación de ventajas a nivel de segmento en una granularidad intermedia, logrando un mejor equilibrio al ofrecer una asignación de crédito más precisa que los métodos a nivel de trayectoria y requiriendo menos puntos de estimación que los métodos a nivel de token, permitiendo una estimación precisa de la ventaja basada en Monte Carlo (MC) sin un modelo crítico. SPO cuenta con tres componentes con estrategias novedosas: (1) partición flexible de segmentos; (2) estimación precisa de la ventaja del segmento; y (3) optimización de políticas utilizando ventajas de segmento, incluyendo una novedosa estrategia de máscara de probabilidad. Además, instanciamos SPO para dos escenarios específicos: (1) SPO-chain para cadenas de pensamiento (CoT) cortas, que presenta una partición basada en puntos de corte y una estimación de ventaja basada en cadenas, logrando mejoras de 6 a 12 puntos porcentuales en precisión sobre PPO y GRPO en GSM8K. (2) SPO-tree para CoT largas, que presenta una estimación de ventaja basada en árboles, lo que reduce significativamente el costo de la estimación MC, logrando mejoras de 7 a 11 puntos porcentuales sobre GRPO en MATH500 bajo evaluaciones de contexto de 2K y 4K. Hacemos nuestro código públicamente disponible en https://github.com/AIFrameResearch/SPO.

English

Enhancing the reasoning capabilities of large language models effectively using reinforcement learning (RL) remains a crucial challenge. Existing approaches primarily adopt two contrasting advantage estimation granularities: Token-level methods (e.g., PPO) aim to provide the fine-grained advantage signals but suffer from inaccurate estimation due to difficulties in training an accurate critic model. On the other extreme, trajectory-level methods (e.g., GRPO) solely rely on a coarse-grained advantage signal from the final reward, leading to imprecise credit assignment. To address these limitations, we propose Segment Policy Optimization (SPO), a novel RL framework that leverages segment-level advantage estimation at an intermediate granularity, achieving a better balance by offering more precise credit assignment than trajectory-level methods and requiring fewer estimation points than token-level methods, enabling accurate advantage estimation based on Monte Carlo (MC) without a critic model. SPO features three components with novel strategies: (1) flexible segment partition; (2) accurate segment advantage estimation; and (3) policy optimization using segment advantages, including a novel probability-mask strategy. We further instantiate SPO for two specific scenarios: (1) SPO-chain for short chain-of-thought (CoT), featuring novel cutpoint-based partition and chain-based advantage estimation, achieving 6-12 percentage point improvements in accuracy over PPO and GRPO on GSM8K. (2) SPO-tree for long CoT, featuring novel tree-based advantage estimation, which significantly reduces the cost of MC estimation, achieving 7-11 percentage point improvements over GRPO on MATH500 under 2K and 4K context evaluation. We make our code publicly available at https://github.com/AIFrameResearch/SPO.