Skip to content
Sagan

Paper

Gradient Extrapolation-Based Policy Optimization

Unreadunread

AI summary

The authors propose GXPO, a new way to update language models during reinforcement learning (like when teaching a model to reason through math problems). Standard RL updates the model using only the current gradient, while full multi-step lookahead is too expensive. GXPO approximates a longer lookahead using just three backward passes: it takes two quick optimizer steps, measures how the gradient changes, predicts where the model would land after K steps, moves partway there, and applies a correction. When the prediction becomes unstable it falls back to standard single-step updates.

Main takeaways:

  • Improves accuracy on Qwen2.5 and Llama math reasoning by +1.65 to +5.00 points over standard GRPO
  • Achieves up to 4× step speedup and 2.33× wall-clock speedup while using only three backward passes per update
  • Works as a plug-in replacement for GRPO—uses the same rollouts, rewards, and advantages, no extra sampling needed
  • Automatically switches back to single-step updates when the extrapolation becomes unreliable