Ottimizzazione delle Politiche a Livello di Segmento: Assegnazione Efficace del Credito a Livello di Segmento nell'Apprendimento per Rinforzo per Modelli Linguistici di Grande Dimensione

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

May 29, 2025

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

cs.AI

Abstract

Migliorare efficacemente le capacità di ragionamento dei modelli linguistici di grandi dimensioni utilizzando l'apprendimento per rinforzo (RL) rimane una sfida cruciale. Gli approcci esistenti adottano principalmente due granularità di stima del vantaggio contrastanti: i metodi a livello di token (ad esempio, PPO) mirano a fornire segnali di vantaggio granulari, ma soffrono di stime imprecise a causa delle difficoltà nell'addestrare un modello critico accurato. All'estremo opposto, i metodi a livello di traiettoria (ad esempio, GRPO) si basano esclusivamente su un segnale di vantaggio grossolano proveniente dalla ricompensa finale, portando a un'assegnazione del credito imprecisa. Per affrontare queste limitazioni, proponiamo Segment Policy Optimization (SPO), un nuovo framework RL che sfrutta la stima del vantaggio a livello di segmento con una granularità intermedia, raggiungendo un migliore equilibrio offrendo un'assegnazione del credito più precisa rispetto ai metodi a livello di traiettoria e richiedendo meno punti di stima rispetto ai metodi a livello di token, consentendo una stima accurata del vantaggio basata su Monte Carlo (MC) senza un modello critico. SPO presenta tre componenti con strategie innovative: (1) partizione flessibile dei segmenti; (2) stima accurata del vantaggio dei segmenti; e (3) ottimizzazione della politica utilizzando i vantaggi dei segmenti, inclusa una nuova strategia di mascheramento delle probabilità. Istanziamo ulteriormente SPO per due scenari specifici: (1) SPO-chain per brevi catene di ragionamento (CoT), caratterizzato da una partizione basata su punti di taglio e una stima del vantaggio basata su catena, ottenendo miglioramenti di 6-12 punti percentuali in accuratezza rispetto a PPO e GRPO su GSM8K. (2) SPO-tree per CoT lunghe, caratterizzato da una stima del vantaggio basata su alberi, che riduce significativamente il costo della stima MC, ottenendo miglioramenti di 7-11 punti percentuali rispetto a GRPO su MATH500 con valutazioni di contesto di 2K e 4K. Rendiamo il nostro codice disponibile pubblicamente all'indirizzo 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.