Optimisation de Politique par Segments : Attribution Efficace de Crédit au Niveau des Segments pour l'Apprentissage par Renforcement dans les Grands Modèles de Langage

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

May 29, 2025

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

cs.AI

Résumé

Améliorer les capacités de raisonnement des grands modèles de langage de manière efficace en utilisant l'apprentissage par renforcement (RL) reste un défi crucial. Les approches existantes adoptent principalement deux granularités contrastées pour l'estimation de l'avantage : Les méthodes au niveau des tokens (par exemple, PPO) visent à fournir des signaux d'avantage fins, mais souffrent d'une estimation imprécise en raison des difficultés à entraîner un modèle critique précis. À l'autre extrême, les méthodes au niveau des trajectoires (par exemple, GRPO) reposent uniquement sur un signal d'avantage grossier provenant de la récompense finale, conduisant à une attribution de crédit imprécise. Pour pallier ces limitations, nous proposons l'Optimisation de Politique par Segments (SPO), un nouveau cadre RL qui exploite l'estimation de l'avantage au niveau des segments avec une granularité intermédiaire, atteignant un meilleur équilibre en offrant une attribution de crédit plus précise que les méthodes au niveau des trajectoires et en nécessitant moins de points d'estimation que les méthodes au niveau des tokens, permettant une estimation précise de l'avantage basée sur Monte Carlo (MC) sans modèle critique. SPO se compose de trois éléments avec des stratégies novatrices : (1) partition flexible des segments ; (2) estimation précise de l'avantage des segments ; et (3) optimisation de la politique utilisant les avantages des segments, incluant une nouvelle stratégie de masque de probabilité. Nous instancions en outre SPO pour deux scénarios spécifiques : (1) SPO-chain pour les courtes chaînes de raisonnement (CoT), avec une partition basée sur des points de coupure et une estimation d'avantage basée sur la chaîne, obtenant des améliorations de 6 à 12 points de pourcentage en précision par rapport à PPO et GRPO sur GSM8K. (2) SPO-tree pour les longues CoT, avec une estimation d'avantage basée sur un arbre, qui réduit significativement le coût de l'estimation MC, obtenant des améliorations de 7 à 11 points de pourcentage par rapport à GRPO sur MATH500 sous des évaluations de contexte de 2K et 4K. Nous rendons notre code public à l'adresse 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.