Оптимизация политики на уровне сегментов: эффективное распределение кредитов на уровне сегментов в обучении с подкреплением для больших языковых моделей

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

May 29, 2025

Авторы: Yiran Guo, Lijie Xu, Jie Liu, Dan Ye, Shuang Qiu

cs.AI

Аннотация

Повышение способности крупных языковых моделей к рассуждению с использованием обучения с подкреплением (RL) остается важной задачей. Существующие подходы в основном используют два контрастных уровня оценки преимуществ: методы на уровне токенов (например, PPO) направлены на предоставление детализированных сигналов преимуществ, но страдают от неточной оценки из-за сложностей в обучении точной модели критика. На другом конце спектра, методы на уровне траекторий (например, GRPO) полагаются исключительно на грубый сигнал преимуществ, основанный на финальной награде, что приводит к неточному распределению заслуг. Для устранения этих ограничений мы предлагаем Segment Policy Optimization (SPO), новый RL-фреймворк, который использует оценку преимуществ на уровне сегментов с промежуточной детализацией, достигая лучшего баланса за счет более точного распределения заслуг по сравнению с методами на уровне траекторий и меньшего количества точек оценки по сравнению с методами на уровне токенов, что позволяет проводить точную оценку преимуществ на основе метода Монте-Карло (MC) без модели критика. SPO включает три компонента с новыми стратегиями: (1) гибкое разделение на сегменты; (2) точная оценка преимуществ сегментов; и (3) оптимизация политики с использованием преимуществ сегментов, включая новую стратегию маскирования вероятностей. Мы также реализуем SPO для двух конкретных сценариев: (1) SPO-chain для коротких цепочек рассуждений (CoT), с использованием нового разделения на основе точек разрыва и оценки преимуществ на основе цепочек, что приводит к улучшению точности на 6-12 процентных пунктов по сравнению с PPO и GRPO на GSM8K. (2) SPO-tree для длинных CoT, с использованием новой оценки преимуществ на основе деревьев, что значительно снижает затраты на оценку MC, обеспечивая улучшение точности на 7-11 процентных пунктов по сравнению с GRPO на MATH500 при оценке в контекстах 2K и 4K. Наш код доступен по адресу 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.