NeurIPS 2025

Accelerating RL for LLM Reasoning with
Optimal Advantage Regression
Kianté Brantley1∗, Mingyu Chen2, Zhaolin Gao3,
Jason D. Lee4, Wen Sun3, Wenhao Zhan4, Xuezhou Zhang2
1Harvard University, 2Boston University, 3Cornell University, 4Princeton University
Abstract
Reinforcement learning (RL) has emerged as a powerful tool for fine-tuning large
language models (LLMs) to improve complex reasoning abilities. However, state-
of-the-art policy optimization methods often suffer from high computational over-
head and memory consumption, primarily due to the need for multiple generations
per prompt and the reliance on critic networks or advantage estimates of the current
policy. In this paper, we propose A⋆-PO, a novel two-stage policy optimization
framework that directly approximates the optimal advantage function and enables
efficient training of LLMs for reasoning tasks. In the first stage, we leverage
offline sampling from a reference policy to estimate the optimal value function V ⋆,
eliminating the need for costly online value estimation. In the second stage, we
perform on-policy updates using a simple least-squares regression loss with only a
single generation per prompt. Theoretically, we establish performance guarantees
and prove that the KL-regularized RL objective can be optimized without requiring
complex exploration strategies. Empirically, A⋆-PO achieves competitive perfor-
mance across a wide range of mathematical reasoning benchmarks, while reducing
training time by up to 2× and peak memory usage by over 30% compared to PPO,
GRPO, and REBEL[Gao et al., 2024a]. Implementation of A⋆-PO can be found at
https://github.com/ZhaolinGao/A-PO.
1
Introduction
Recent advances in large language models (LLMs), including OpenAI-o1 [OpenAI, 2024], DeepSeek-
R1 [DeepSeek-AI, 2025], and Kimi-1.5 [Team, 2025], have demonstrated remarkable reasoning
capabilities. These models excel at producing long Chain-of-Thought (CoT)[Wei et al., 2023,
DeepSeek-AI, 2025, Zeng et al., 2025] re...