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

要約

強化学習（RL）を使用して効果的に大規模な言語モデルの推論能力を強化することは、依然として重要な課題です。
既存のアプローチは、主に2つの対照的な利点推定の粒度を採用しています。トークンレベルの方法（例：PPO）は、微調整されたアドバンテージシグナルを提供することを目的としていますが、正確な批評家モデルのトレーニングが難しいために不正確な推定に悩まされています。
もう1つの極端に、軌道レベルのメソッド（GRPOなど）は、最終報酬からの粗粒のアドバンテージ信号にのみ依存しており、不正確なクレジット割り当てにつながります。
これらの制限に対処するために、セグメントポリシー最適化（SPO）を提案します。これは、中間の粒度でセグメントレベルのアドバンテージ推定を活用する新しいRLフレームワークであり、軌跡レベルの方法よりも正確なクレジット割り当てを提供することにより、より正確なクレジット割り当てを提供することにより、トークンレベルの方法よりも少ない推定ポイントを必要とすることにより、MC. Carlo（MC）の場合に基づく推定値を強化します。
SPOは、新しい戦略を備えた3つのコンポーネントを備えています。（1）柔軟なセグメントパーティション。
（2）正確なセグメントアドバンテージの推定。
（3）新しい確率マスク戦略を含むセグメントの利点を使用したポリシーの最適化。
さらに、2つの特定のシナリオのSPOをインスタンス化します。（1）新しいカットポイントベースのパーティションとチェーンベースのアドバンテージ推定を特徴とする短いチェーンオブサート（COT）のSPOチェーン。
（2）長いCOTのSPO-Treeは、新しいツリーベースのアドバンテージ推定を特徴としており、MCの推定コストを大幅に削減し、2Kおよび4Kコンテキスト評価の下でMath500のGRPOよりも7ドル$ 11 $ 11ドルのポイント改善を達成します。
コードをhttps://github.com/aiframeresearch/spoで公開しています。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google