Segment Policy Optimization: Effektive Zuweisung von Gutschriften auf Segmentebene im Reinforcement Learning für große Sprachmodelle

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

May 29, 2025

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

cs.AI

Zusammenfassung

Die Verbesserung der Fähigkeiten großer Sprachmodelle zum logischen Schlussfolgern mithilfe von Reinforcement Learning (RL) bleibt eine entscheidende Herausforderung. Bestehende Ansätze verfolgen hauptsächlich zwei kontrastierende Granularitäten der Vorteilsschätzung: Token-Level-Methoden (z. B. PPO) zielen darauf ab, feingranulare Vorteilssignale bereitzustellen, leiden jedoch unter ungenauen Schätzungen aufgrund von Schwierigkeiten beim Training eines präzisen Kritiker-Modells. Auf der anderen Seite verlassen sich Trajektorien-Level-Methoden (z. B. GRPO) ausschließlich auf ein grobgranulares Vorteilssignal aus der Endbelohnung, was zu ungenauen Kreditzuweisungen führt. Um diese Einschränkungen zu überwinden, schlagen wir Segment Policy Optimization (SPO) vor, ein neuartiges RL-Framework, das eine Segment-Level-Vorteilsschätzung mit mittlerer Granularität nutzt und so eine bessere Balance erreicht, indem es präzisere Kreditzuweisungen als Trajektorien-Level-Methoden bietet und weniger Schätzpunkte als Token-Level-Methoden benötigt. Dies ermöglicht eine genaue Vorteilsschätzung basierend auf Monte Carlo (MC) ohne ein Kritiker-Modell. SPO umfasst drei Komponenten mit innovativen Strategien: (1) flexible Segmentpartitionierung, (2) präzise Segmentvorteilsschätzung und (3) Policy-Optimierung mithilfe von Segmentvorteilen, einschließlich einer neuartigen Wahrscheinlichkeitsmasken-Strategie. Wir konkretisieren SPO weiter für zwei spezifische Szenarien: (1) SPO-Chain für kurze Chain-of-Thought (CoT), das eine neuartige Schnittpunkt-basierte Partitionierung und kettenbasierte Vorteilsschätzung beinhaltet und eine Verbesserung der Genauigkeit um 6–12 Prozentpunkte gegenüber PPO und GRPO auf GSM8K erzielt. (2) SPO-Tree für lange CoT, das eine neuartige baumbasierte Vorteilsschätzung bietet, die die Kosten der MC-Schätzung erheblich reduziert und eine Verbesserung der Genauigkeit um 7–11 Prozentpunkte gegenüber GRPO auf MATH500 unter 2K- und 4K-Kontextbewertungen erreicht. Unser Code ist öffentlich verfügbar unter 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.