Segment Policy Optimalisatie: Effectieve Toewijzing van Segmentniveau Credits in Reinforcement Learning voor Grote Taalmodellen

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

Samenvatting

Het effectief verbeteren van de redeneervaardigheden van grote taalmodellen met behulp van reinforcement learning (RL) blijft een cruciale uitdaging. Bestaande benaderingen hanteren voornamelijk twee contrasterende granulariteiten voor voordeelschatting: Token-level methoden (bijv. PPO) streven naar fijnmazige voordeelsignalen, maar kampen met onnauwkeurige schattingen door moeilijkheden bij het trainen van een nauwkeurig criticusmodel. Aan het andere uiterste vertrouwen traject-level methoden (bijv. GRPO) uitsluitend op een grofmazig voordeelsignaal van de uiteindelijke beloning, wat leidt tot onnauwkeurige toewijzing van verdienste. Om deze beperkingen aan te pakken, stellen we Segment Policy Optimization (SPO) voor, een nieuw RL-framework dat gebruikmaakt van segment-level voordeelschatting op een tussenliggende granulariteit, waardoor een betere balans wordt bereikt door nauwkeurigere toewijzing van verdienste dan traject-level methoden en minder schattingspunten vereist dan token-level methoden, wat nauwkeurige voordeelschatting mogelijk maakt op basis van Monte Carlo (MC) zonder een criticusmodel. SPO omvat drie componenten met nieuwe strategieën: (1) flexibele segmentpartitie; (2) nauwkeurige segmentvoordeelschatting; en (3) beleidsoptimalisatie met behulp van segmentvoordelen, inclusief een nieuwe waarschijnlijkheidsmaskerstrategie. We concretiseren SPO verder voor twee specifieke scenario's: (1) SPO-chain voor korte chain-of-thought (CoT), met nieuwe partitie op basis van snijpunten en voordeelschatting op basis van ketens, wat 6-12 procentpunt verbeteringen in nauwkeurigheid oplevert ten opzichte van PPO en GRPO op GSM8K. (2) SPO-tree voor lange CoT, met nieuwe voordeelschatting op basis van bomen, wat de kosten van MC-schatting aanzienlijk verlaagt, wat 7-11 procentpunt verbeteringen oplevert ten opzichte van GRPO op MATH500 onder 2K en 4K contextevaluatie. We maken onze code openbaar beschikbaar op 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.