Gradient Extrapolation-Based Policy Optimization

Authors: Ismam Nur Swapnil, Aranya Saha, Tanvir Ahmed Khan et al.

arXiv · PDF

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

Relevance

Not directly related to my persona or midtraining work—this is about optimizing the RL fine-tuning phase after the assistant is already installed. Could be tangentially relevant if I ever wanted to install behaviors via RL rather than supervised fine-tuning, but my current focus is on understanding what gets installed during different training regimes, not on making the RL optimizer itself faster.

Abstract

arXiv:2605.06755v1 Announce Type: new Abstract: Reinforcement learning is widely used to improve the reasoning ability of large language models, especially when answers can be automatically checked. Standard GRPO-style training updates the model using only the current step, while full multi-step lookahead can give a better update direction but is too expensive because it needs many backward passes. We propose Gradient Extrapolation-Based Policy Optimization (GXPO), a plug-compatible policy-update rule for GRPO-style reasoning RL. GXPO approximates a longer local lookahead using only three backward passes during an active phase. It reuses the same batch of rollouts, rewards, advantages, and GRPO loss, so it does not require new rollouts or reward computation at the lookahead points. GXPO takes two fast optimizer steps, measures how the gradients change, predicts a virtual K-step lookahead point, moves the policy partway toward that point, and then applies a corrective update using the true gradient at the new position. When the lookahead signal becomes unstable, GXPO automatically switches back to standard single-pass GRPO. We also give a plain-gradient-descent surrogate analysis that explains when the extrapolation is exact and where its local errors come from. Across Qwen2.5 and Llama math-reasoning experiments, GXPO improves the average sampled pass@1 by +1.65 to +5.00 points over GRPO and by +0.14 to +1.28 points over the strongest SFPO setting, while keeping the active-phase cost fixed at three backward passes. It also achieves up to 4.00x step speedup, 2.33x wall-clock speedup, and 1.33x backward-pass speedup in reaching GRPO's peak accuracy.