Otimização de Política por Segmentos: Atribuição Eficaz de Crédito em Nível de Segmento em RL para Modelos de Linguagem de Grande 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

Resumo

Melhorar as capacidades de raciocínio de modelos de linguagem de grande escala de forma eficaz usando aprendizado por reforço (RL) continua sendo um desafio crucial. As abordagens existentes adotam principalmente duas granularidades contrastantes de estimativa de vantagem: Métodos de nível de token (por exemplo, PPO) visam fornecer sinais de vantagem refinados, mas sofrem com estimativas imprecisas devido às dificuldades em treinar um modelo crítico preciso. No outro extremo, métodos de nível de trajetória (por exemplo, GRPO) dependem exclusivamente de um sinal de vantagem de granularidade grossa proveniente da recompensa final, levando a uma atribuição de crédito imprecisa. Para abordar essas limitações, propomos o Segment Policy Optimization (SPO), um novo framework de RL que aproveita a estimativa de vantagem em nível de segmento com uma granularidade intermediária, alcançando um melhor equilíbrio ao oferecer uma atribuição de crédito mais precisa do que os métodos de nível de trajetória e exigindo menos pontos de estimativa do que os métodos de nível de token, permitindo uma estimativa de vantagem precisa baseada em Monte Carlo (MC) sem um modelo crítico. O SPO apresenta três componentes com estratégias inovadoras: (1) partição de segmento flexível; (2) estimativa de vantagem de segmento precisa; e (3) otimização de política usando vantagens de segmento, incluindo uma nova estratégia de máscara de probabilidade. Além disso, instanciamos o SPO para dois cenários específicos: (1) SPO-chain para cadeias de pensamento (CoT) curtas, apresentando partição baseada em pontos de corte e estimativa de vantagem baseada em cadeia, alcançando melhorias de 6 a 12 pontos percentuais em precisão sobre PPO e GRPO no GSM8K. (2) SPO-tree para CoT longas, apresentando estimativa de vantagem baseada em árvore, que reduz significativamente o custo da estimativa MC, alcançando melhorias de 7 a 11 pontos percentuais sobre GRPO no MATH500 sob avaliação de contexto de 2K e 4K. Disponibilizamos nosso código publicamente em 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.