Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Entropy-Gated Selective Policy Optimization:Token-Level Gradient Allocation for Hybrid Training of Large Language Models

Yuelin Hu, Zhengxue Cheng, Wei Liu, Li Song

TL;DR

EG-SPO addresses instability and inefficiency in hybrid SFT+RL training for complex reasoning by introducing token-level entropy gating with advantage-aware gradients. It couples a three-stage process—SFT warm-up, per-token entropy computation during RL rollouts, and entropy-gated PPO updates where high-entropy tokens receive full updates and low-entropy tokens receive $\phi(p)$-attenuated updates—while preserving the advantage term $A_t$ to prevent reinforcement of confident errors. The Predictive Entropy Module enables adaptive, token-wise routing based on $H(y_t)$ and per-sequence quantiles, improving learning efficiency and exploration without trajectory mismatch. Empirical results on AIME, AMC, and MATH show gains of up to 3.8 percentage points over CHORD-$\phi$ with only about 3.4% extra compute, underscoring the value of token-level uncertainty and advantage-aware credit assignment in hybrid LLM training.

Abstract

Hybrid training methods for large language models combine supervised fine tuning (SFT) on expert demonstrations with reinforcement learning (RL) on model rollouts, typically at the sample level. We propose Entropy Gated Selective Policy Optimization (EGSPO), a three stage framework that extends sample level mixing with token level gradient modulation. Stage 1, SFT expert learning, establishes a reliable warm up policy using expert demonstrations with a pure SFT loss. Stage 2, RL rollout generation, samples trajectories from the current policy and computes per token predictive entropy. Stage 3, the EGSPO mechanism, applies entropy gated gradient allocation: a predictive entropy module routes high entropy tokens to full PPO updates to encourage exploration, and low entropy tokens to attenuated PPO updates to reduce variance and preserve knowledge. Critically, both branches incorporate the advantage function A_t, ensuring that incorrect trajectories receive consistent negative learning signals and preventing reinforcement of confident errors. EGSPO achieves consistent improvements on mathematical reasoning benchmarks, with gains of 3.8 percent on AIME and 2.9 percent on MATH over the CHORD phi baseline, while incurring only 3.4 percent additional computational overhead.

Entropy-Gated Selective Policy Optimization:Token-Level Gradient Allocation for Hybrid Training of Large Language Models

TL;DR

EG-SPO addresses instability and inefficiency in hybrid SFT+RL training for complex reasoning by introducing token-level entropy gating with advantage-aware gradients. It couples a three-stage process—SFT warm-up, per-token entropy computation during RL rollouts, and entropy-gated PPO updates where high-entropy tokens receive full updates and low-entropy tokens receive -attenuated updates—while preserving the advantage term to prevent reinforcement of confident errors. The Predictive Entropy Module enables adaptive, token-wise routing based on and per-sequence quantiles, improving learning efficiency and exploration without trajectory mismatch. Empirical results on AIME, AMC, and MATH show gains of up to 3.8 percentage points over CHORD- with only about 3.4% extra compute, underscoring the value of token-level uncertainty and advantage-aware credit assignment in hybrid LLM training.

Abstract

Hybrid training methods for large language models combine supervised fine tuning (SFT) on expert demonstrations with reinforcement learning (RL) on model rollouts, typically at the sample level. We propose Entropy Gated Selective Policy Optimization (EGSPO), a three stage framework that extends sample level mixing with token level gradient modulation. Stage 1, SFT expert learning, establishes a reliable warm up policy using expert demonstrations with a pure SFT loss. Stage 2, RL rollout generation, samples trajectories from the current policy and computes per token predictive entropy. Stage 3, the EGSPO mechanism, applies entropy gated gradient allocation: a predictive entropy module routes high entropy tokens to full PPO updates to encourage exploration, and low entropy tokens to attenuated PPO updates to reduce variance and preserve knowledge. Critically, both branches incorporate the advantage function A_t, ensuring that incorrect trajectories receive consistent negative learning signals and preventing reinforcement of confident errors. EGSPO achieves consistent improvements on mathematical reasoning benchmarks, with gains of 3.8 percent on AIME and 2.9 percent on MATH over the CHORD phi baseline, while incurring only 3.4 percent additional computational overhead.
Paper Structure (26 sections, 8 equations, 1 figure, 3 tables)

This paper contains 26 sections, 8 equations, 1 figure, 3 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Reinforcement Learning for Language Models
  4. Supervised Fine-Tuning for Reasoning
  5. Hybrid Training Approaches
  6. Methodology
  7. Overview: Three-Stage EG-SPO Framework
  8. Avoiding Trajectory Mismatch
  9. Advantage-Aware Gradient Modulation
  10. Predictive Entropy Module (Stage 2-3)
  11. Joint Gradient Update (Stage 3)
  12. The $\phi$ Weighting Function
  13. Training Procedure
  14. Experimental Setup
  15. Datasets and Tasks
  16. ...and 11 more sections

Figures (1)

  • Figure 1: Overview of our three-stage EG-SPO training framework. Stage 1 (SFT Expert Learning): Expert demonstrations train a warm-up policy using pure SFT loss ($\sim$20% samples). Stage 2 (RL Rollout Generation): Current policy generates model rollouts and computes per-token entropy. Stage 3 (EG-SPO Main Mechanism): The Predictive Entropy Module routes high-entropy tokens to full PPO updates (encouraging exploration) and low-entropy tokens to $\phi$-attenuated PPO (reducing variance, preserving knowledge). Both branches retain advantage $A_t$, ensuring advantage-aware gradients that avoid reinforcing confident errors.