Title: Accelerating RL for LLM Reasoning with Optimal Advantage Regression
Venue: NeurIPS 2025
Accepted: yes
Keywords: reinforcement learning

large language models

reasoning
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 overhead 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^\star$-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^\star$, 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^\star$-PO achieves competitive performance across a wide range of mathematical reasoning benchmarks, while reducing training time by up to 2$\times$ and peak memory usage by over 30\% compared to PPO, GRPO, and REBEL. Implementation of $A^\star$-PO can be found at https://github.com/ZhaolinGao/A-PO.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces A⋆-PO (Policy Optimisation via Optimal Advantage Regression), a two-stage RL procedure for post-training LLMs on reasoning tasks. The work offers a neat regression-based alternative to critic-heavy RL algorithms, with solid speed-ups and \"formal guarantees\", but those guarantees require many unrealistic assumptions that may not hold for frontier-scale LLMs.\n\nI would like to clarify a point before I move on to the questions. it is known that under sufficient assumptions, one can get vanishing regrets; however the active area of research is generalising those to problems that include non-linearities and non-convexity.\n\nSTRENGTHS AND WEAKNESSES:\nThe paper has a nice flow to it and reads well, till the theoretical section which I felt was not very well presented. I feel for the authors due to the constraints of space, but more discussions on the assumptions and their impact are indeed needed. I would like to ask the authors some questions: \n\n1. When are policy classes not misspecified? Because if they are, then the realizability assumption would have issues. Please discuss the approximation error or clearly detail which architectures make the assumption plausible and why? \n\n2. In assumption 2, can the log-ratio explode? I thought it could, e.g, if we pair a small model with a larger reference one? Is this correct? Can the authors please help me understand? If so, I think the theory should handle how to renormalise? \n\n3. In assumption 3, v_ref can be very low, especially on some hard problems. This might make the bound vacuous, can the authors please characterise performance as v_ref tends to zero? Maybe an adaptive sampling procedure might help to remedy this issue? \n\n4. The Fisher assumption also doesn't really hold in many high-dimensional cases with deep networks - in fact, many times the Fisher matrix ends up ill-conditioned. Maybe the authors could add a regularisation (which will change the update rules) and study and quantify the value of such regularisers needed in practice. \n\n5. The Gap hides a massive constant. Can the authors please bound those constants or give ballpark numbers so that those aren't massive, like they are not > 1?  \n\nI don't mean to be annoying, but given that I come from a theoretical background, the value of the theoretical section in this paper is minimal from my point of view; those assumptions are standard and (almost always and this is a well known fact) do not apply to non-convex losses, deep networks, and transformers. In other words, it is like taking a theory that had been developed for some types of problems and adding it to a new type of problem, but making the relevant assumptions so that the new type of problem meets the first type. \n\n6. I would like to see more experiments related to coding benchmarks as well; they also require formal reasoning. Can the authors please report those if possible? \n\n7. Please publish wall-clock times per stage and GPU hours for all baselines.  \n\n8. There are no error bars, and the authors clearly say NO to the statistical significance question. I am not sure how to take that? Like are the results reported in the paper by luck? \n\nI am happy to improve my scores of the above questions have been answered.\n\nQUESTIONS:\nPlease see above\n\nLIMITATIONS:\nBriefly mentioned", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work proposes $A^*$-PO, an RL technique to finetune reasoning capabilities with improved sample efficiency. It achieves this by first learning an optimal advantage function offline on generations produced by the base reference policy, then during online training, regressing the log ratio against it. The authors demonstrate the improved efficiency compared to PPO/GRPO with only one generation per prompt during online training.\n\nSTRENGTHS AND WEAKNESSES:\n**Strength**\n- The proposed method is novel and theoretically grounded. The paper theoretically and empirically supports the improved sample complexity.\n- The experiments are well-designed and comprehensive, including evaluations on GSM8K, MATH, and long-context benchmarks, as well as showing generalization to other benchmarks after training.\n\n**Weaknesses**\nI am not an RL expert, and would defer constructive criticism to more experienced reviewers. I mainly have a few questions:\n- How would the estimation of $V^*$ perform with non-binary or denser reward signals (e.g., with process-supervision)?\n- The authors state that during online training, only one sample is required to regress on the estimated optimal advantage, what would be the impact if multiple samples were generated?\n- Clarifying question: In Fig 4, is filtering applied during the online stage, i.e., problems that did not find a correct solution (as measured in Stage 1) are not used for online training? Would this technique also apply to PPO/GRPO, or is this something specific to your method?\n- Have the authors ablated the effect of different values of $\\beta$ for the two stages?\n\nQUESTIONS:\nThese are some specific questions I had in addition to the ones I listed above:\n- In Table 1, what is the significance of showing lower KL-divergence from the base policy? Also, would this not be also impacted by the choice of $\\beta$ in baseline methods?\n- Have the authors considered adaptively changing $N$ based on problem difficulty? For simpler problems, a larger $N$ might not be necessary to estimate the advantage function, whereas for harder problems, it could be more crucial.\n- How well does a model post-trained with $A^{*}$-PO retain its original abilities, i.e., does it exhibit more/less catastrophic forgetting?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work focuses on the topic of using reinforcement learning to enhance the LLMs' reasoning capability.  Previous RL methods either rely on training a critic network (e.g., PPO) or generating multiple samples per prompt (e.g., GRPO) during online finetuning. This work, based on the property of the optimal value function under the KL-regularized policy, proposes a two-stage method, where the first stage offline sampling is performed directly via the reference policy to estimate the optimal value function, and the second stage performs online training with on-policy updates using a simple least-square loss and one response per prompt. The effectiveness of the proposed method, i.e., A*PO is verified with experimental results, showing that it achieves comparable results with previous methods, while speeding up training and reducing memory usage. Furthermore, theoretical results are also provided, providing the theoretical soundness of a more general version of A*PO, without strong assumptions or using complex exploration strategies but achieving nice convergence results and yielding attractive properties.\n\nSTRENGTHS AND WEAKNESSES:\nThis work is a nice exploration of the more efficient methods in RL finetuning of LLM:\n\n- The proposed method is based on theoretical properties of the optimal value under the KL-regularized target, in particular, discovering that it can be estimated via sampling through the reference policy\n\n- The designed algorithm is clear, intuitive, and easy to implement. It tackles the problem that existing methods use extra resources (both computation and memory) to estimate the value function. The pre-computed value estimation from the offline sampling can save resource and time in online fine-tuning.\n\n- The experimental results are comprehensive in terms of metrics. It helps a lot to visualize the comparison of time and memory used by different methods.\n\n- The theoretical results are also solid. I really appreciate the nice connection between empirical designs and theoretical understandings. The results themselves can be valuable in understanding the KL-regularized RL/contextual bandits problem. \n\n- The overall presentation is clear and well-organized.\n\nIf I have to nit-picking, I would hope the authors provide more illustrations of the novelty or innovation behind the theoretical discussions. Currently, the results are well-illustrated, but the underlying proof technics are largely deferred to the appendix without many highlights on the challenges/novelties there.\n\nQUESTIONS:\nThere are a few clarification questions that I have in mind to discuss with the authors:\n\n- At the beginning of Section 4, there is a discussion on using two different KL-regularization coefficient during two stages. I am wishing to understand more about the choices and the reason behind. In particular, if I understand it correctly, using different beta essentially lead to different optimal values, which then damages the usefulness of the offline value estimation for the online phase.\n\n- Regarding the theoretical innovation concern raised above, I am also looking forward to learning from the authors more about the underlying and probably hidden contributions there.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

comments: SUMMARY:
This paper proposes a way to significantly reduce the cost of online RL by decoupling value estimation. Specifically, for a given RL prompt dataset, they collect offline rollouts using a reference model to serve as the value baseline. Then during online RL, they use only 1 rollout per prompt and simplify the policy objective to a plain regression on the advantage (using MSE loss). Experiments are done on Qwen models and math domain.

STRENGTHS AND WEAKNESSES:
### Strengths:
- Easy to plug into existing pipelines. Compared to other methods that modify value or advantage estimators, regressing offline advantages is a fresh idea.
- The writing is clear and easy to follow.
- The theory section gives a solid in-principle analysis, which is nice to see.

### Weaknesses:
- The key novelty — using offline value estimation — needs more ablations. For example, I would really like to see how it performs if they just estimate value online with the same setup (i.e. just keep updating \pi in the process how will the value baseline change).
- Experiments are limited and focus only on short-context (non-thinking) models. I totally get the resource constraints (they do not even have complete baselines for thinking models), but it still undercuts the strength of the paper’s message.

QUESTIONS:
See the weaknesses part point 1.

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-07-03