Accelerating RL for LLM Reasoning with Optimal Advantage Regression

We thank the reviewer for their detailed comments.

the active area of research is generalising those to problems that include non-linearities and non-convexity.

We acknowledge this point, and would like to clarify that our theoretical analysis is exactly on nonconvex and nonlinear cases. More specifically, Theorem 1 addresses general policy classes, which are nonconvex and nonlinear. Consequently, the loss function $\widehat{l}_t$ is also nonconvex and nonlinear with respect to $\pi$. To handle this, we introduce the decoupling coefficient (Definition 1), which is effective when dealing with nonconvexity and nonlinearity.

We also present a case study on log-linear policy classes to illustrate our theoretical results more clearly. We would like to clarify that log-linear policy classes and the corresponding loss functions are still “nonconvex and nonlinear” with respect to the parameter $\theta$. We designed two specific algorithms (utilizing FTPL and OGD) for this case and they achieve sub-linear regret. We will further highlight our focus on nonconvexity in the paper.

Weaknesses 1

We assume policy class realizability to simplify presentation, as almost all the existing theoretical works do. That said, our analysis can be generalized to misspecified policy classes easily. Assume that the approximation error of the policy class is as follows: $\min_{\pi\in\Pi, x, y} |\ln\pi(y|x)-\ln\pi^{\star}(y|x)|\leq C_{\mathsf{approx}}.$ Let $\widetilde{\pi}$ denote the policy in $\Pi$ that is closest to $\pi^{\star}$. Then we have $\widehat{l} ^t(\widetilde{\pi}) \leq 2\widehat{l} ^t(\pi^{\star}) + 2\beta^2C^2_{\mathsf{approx}}.$ Substituting the above inequality to the proof in Appendix H, we know that working with misspecified policy classes is equivalent to adding $\beta^2C^2_{\mathsf{approx}}$ to the regret. All the results still hold, with just an additional approximation error.

Weaknesses 2

We would like to point out that in empirical works, with clipping techniques, such log ratios can be controlled and not blow up. That is also the reason why the existing theoretical works on DPO-type algorithms have similar assumptions.

Weaknesses 3

First we would like to clarify that our bound depends on $min(\exp(1/\beta), v_{ref})$. Therefore, if $\beta$ is fixed (which is typically the setting in practice), our result would not be worsened by small $v_{ref}$. On the other hand, a well known and widely applied technique in LLM training is filtration, i.e., filtering out the prompts that are too difficult or too easy. With this process, $v_{ref}$ will be reasonably large during RL training, preventing $1/v_{ref}$ from blowing up.

Weaknesses 4

First, Fisher matrix assumption is not necessary in terms of achieving sublinear regret. Corollary 1 does not need any assumptions on Fisher matrix. The reason why we introduce Fisher matrix assumption is that we want to claim that our algorithm can enjoy a faster convergence when the problem is regular and well conditioned. That said, in our experiments, we used OGD and our algorithm works pretty well. This demonstrates that even if the Fisher matrix can be ill conditioned sometimes, the problem is often regular and OGD can work in practice.

Weaknesses 5

Through a rough calculation, the constants in Theorems 1 and 2 can essentially be bounded by 16 and 64, respectively. For Theorem 1, we omitted four constants of 2 in Lines 680, 681, 686, and 690. For Theorem 3, the constant on the right-hand side of the inequality in Line 651 can be bounded by 256, which implies that the constant in Line 655 is $4\sqrt{256} = 64$. Notice that the actual constants will be much smaller than these theoretical values, especially considering that we relaxed these constant terms for analytical convenience in analysis. Even so, it suffices to say that these constants are not massive.

those assumptions are standard and do not apply to non-convex losses, deep networks, and transformers

We agree that these assumptions are standard, but we respectively disagree that they do not apply to these settings. First, realizability does not depend on whether the function classes are convex. Instead, more complicated function classes like deep networks and transformers typically have a higher representation capacity and thus realizability is easier to satisfy (at least the approximation error would be smaller.) Second, other assumptions have their origins in practice. The log ratio is bounded because of clipping, and $v_{ref}$ is not small because of filtration. In addition, these techniques work well in practice, which further support the validity of these assumptions. Since these assumptions are standard in existing works, we do not assume unreasonable assumptions to simplify analysis or use other works’ results directly.

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

We respectfully disagree with this claim. As the reviewer has noted, our assumptions are standard in existing works, indicating that we do not introduce additional, unreasonable assumptions merely to reuse existing theory. On the contrary, our analysis introduces several novel techniques and results. We summarize them below for clarity:

• Getting rid of policy completeness assumption. In our analysis we only need policy realizability. In contrast, other on-policy policy optimization approaches including PG, PPO and REBEL rely on a stronger condition called policy completeness. We avoid such a condition by designing novel loss functions, which directly regress to the optimal policy rather than next-step policies similar to policy iteration.
• Getting rid of exploration mechanisms. Our method does not incorporate any explicit exploration mechanism such as optimism in the face of uncertainty or reset to some exploratory distribution as in existing works. This is a significant improvement because exploration is necessary in general RL settings. We achieve this by noticing that in the context of LLMs, there exists a close connection between the novel loss functions we proposed and the KL divergence between the learned policy and the optimal policy (illustrated in step 2 of the proof in Appendix H). This is a new finding. Leveraging this observation, we managed to bound the suboptimality of the learned policy efficiently.
• New findings on nonconvex no-regret learning with OGD. In the case study of log-linear policy classes, we find that OGD provably converges with the optimal rate. This is a new finding because the loss functions are still nonconvex with log-linear policy classes and OGD does not work for nonconvex losses in general. We achieve such a result not by making unreasonable assumptions, but by establishing a connection between the gradient of the loss and the distance from the optimal parameter (Lemma 3 in Appendix G), which is a new finding in this area. This connection enables us to efficiently bound the distance between the learned policy and optimal policy. Notably, the techniques used in Appendix G are totally different from the ones we use in Appendix H (Theorem 1), which is the reason why we can obtain a faster convergence rate.

Weaknesses 6

We agree that coding benchmarks are an important and challenging class of reasoning tasks. However, we did not include them in this work due to the significant computational overhead involved—specifically, reward computation typically requires running unit tests in isolated environments, which introduces substantial cost and engineering complexity. Despite this, we believe $A^\star$-PO remains well-suited for coding tasks, as it only requires one sample per prompt during online training, and thus demands fewer reward queries than approaches like GRPO that rely on multiple generations. This efficiency advantage would likely persist in the coding domain, and we are excited to explore this direction in future work.

Weaknesses 7

Thank you for the suggestion. We include wall-clock training times for each method in Figure 2 and Tables 1 and 2. As described in Appendix A.2 and A.4, our experiments were run on machines with 4 H100 GPUs (for Sections 4.1 and 4.2) and 8 H100 GPUs (for Section 4.3). Therefore, GPU hours can be directly computed by multiplying the reported training times by the number of GPUs used. To further clarify the training efficiency of $A^\star$-PO, we include a rate-of-convergence analysis in the supplemental material, showing training and generation time versus reward. We will move this analysis into the appendix in the camera-ready version to make this information more accessible.

Weaknesses 8

We appreciate the opportunity to clarify. While we did not include error bars in the current submission, all experiments were run with fixed seeds, making our results exactly reproducible. Furthermore, in large-scale RL for LLMs, empirical results tend to be stable across runs due to the high sample count per iteration—often involving hundreds of prompt-response pairs per update and training over thousands of steps. To further reduce variance, we report results on large-scale datasets and, for smaller benchmarks such as AMC 23, AIME 24/25, and HMMT Feb 24/25, we compute the average performance over 32 generations per prompt. This setup ensures statistical reliability even in the absence of formal error bars. We are confident that the observed performance improvements are robust and not due to chance, as they are consistent across models, datasets, and evaluation metrics.

We thank the reviewer for their constructive and thoughtful comments.

How would the estimation of $V^\star$ perform with non-binary or denser reward signals (e.g., with process-supervision)?

Our method is compatible with non-binary and dense rewards, as we make no assumptions about the reward values. While our experiments focus on the bandit setting with terminal rewards, the value estimation naturally extends to dense or process-level supervision. Consider a setting with a prompt $x$ and a response with $H$ steps $y_{1:H}$. Then the optimal value at step $h$ can be expressed as $V^\star(x, y_{1:h}) = \beta \ln \mathbb{E} {\pi_{ref}} \left[\exp( \sum_{t > h} r(x,y_{1:t}) / \beta )\right]$ where $r(x,y_{1:t})$ is the process reward for step $t$. Thus, our framework remains applicable under denser reward feedback. We leave the empirical exploration under these settings for future work.

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?

Theoretically, using multiple samples per prompt has no effect on the expected loss as $\mathbb{E}{x \sim \rho, y \sim \pi(\cdot|x)}[\ell(x, y)] = \mathbb{E}{x \sim \rho,(y_1, \dots, y_K) \sim \pi(\cdot|x)}\left[\frac{1}{K} \sum_{k=1}^K \ell(x, y_k)\right]$ holds for any loss function $\ell$ and any $K$. However, in practice, generating multiple samples per prompt—as in GRPO—increases training time and memory usage. Our method is more efficient because it achieves comparable performance with only a single generation per prompt.

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?

Yes, the filtering is applied before the online stage: we discard prompts for which none of the $N$ offline samples from $\pi_{\text{ref}}$ yield a correct answer. While this strategy could be applied to PPO or GRPO, those methods do not have offline sampling and would need to perform additional inference solely for filtering, which adds further overhead to already more complex training pipelines. In contrast, our method naturally performs offline sampling for value estimation, so filtering introduces negligible extra cost and offers a natural efficiency bonus.

Have the authors ablated the effect of different values of $\beta$ for the two stages?

Yes, this is addressed in Appendix D. As $\beta_1 \to \infty$, the estimated $V^\star(x)$ approaches the expected reward under $\pi_{\text{ref}}$; as $\beta_1 \to 0$, it converges to the $\text{Pass@}N$ performance. These properties are theoretically proven in Appendix E.

We choose a larger $\beta_1$ for the offline stage to provide a smooth and stable value estimate, and a smaller $\beta_2$ during online training to relax the KL constraint and promote reward maximization. Our ablations show that overly optimistic or pessimistic estimates of $V^\star(x)$ can both harm performance. For $\beta_2$, we found that $10^{-3}$ consistently gives strong results; larger or smaller values would degrade the performance.

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?

Lower KL divergence indicates that the trained policy stays closer to the reference policy while still achieving strong performance. This is significant because the KL-regularized RL objective aims to optimize reward while staying minimally divergent from the base model. Our results show that $A^\star$-PO induces significantly less change to the base model than PPO/GRPO for comparable or better performance—suggesting better stability and less catastrophic forgetting.

Indeed, the KL divergence can be influenced by the choice of $\beta$ in baseline methods. However, we tuned $\beta$ for all methods to maximize the validation accuracy to ensure a fair comparison, and $A^\star$-PO consistently achieves the smallest KL divergence across model sizes and datasets.

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.

This is an excellent suggestion. Intuitively, prompts with intermediate success rates (e.g., around 50% accuracy under $\pi_{\text{ref}}$) have higher variance in their value estimates, and would benefit from more samples to improve estimation accuracy. In contrast, very easy or very hard prompts (i.e., near 0% or 100% success) require fewer samples, as their value estimates are more stable. While we have not explored adaptive selection of $N$ in this work, it is indeed a promising direction.

How well does a model post-trained with $A^\star$-PO retain its original abilities, i.e., does it exhibit more/less catastrophic forgetting?

While we did not explicitly test for catastrophic forgetting, the significantly lower KL divergence to the base model suggests that $A^\star$-PO better preserves the original capabilities compared to PPO and GRPO. This is an important benefit for applications requiring general-purpose LLMs with minimal task-specific overfitting.

We sincerely thank the reviewer for the thoughtful and constructive feedback.

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).

Our key contribution is to decouple value estimation from online learning by estimating the optimal value function offline using samples from the reference policy. This avoids the need for explicit critics or multi-sample estimation of the current policy’s value, which are required by PPO and GRPO and substantially increase both algorithmic complexity and training cost. While we understand the curiosity about online estimation under our setup, such an approach would revert to GRPO or PPO, thereby obscuring the benefits and novelty of our method. Moreover, as discussed in Section 2, the optimal value function under the KL-regularized RL objective is defined with respect to the reference policy $\pi_{\text{ref}}$, not the learned policy $\pi$. This provides a theoretical justification for our offline estimation approach and highlights the limitations of online estimation in this setting.

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.

We fully agree that evaluating on more complex, long-context reasoning tasks is important. To address this, we included experiments in Section 4.3 using the DeepSeek-R1 distilled Qwen-1.5B model with a 16k context length which is a setting for long-form reasoning. As shown in Table 2, $A^\star$-PO achieves the best performance on AIME 2025 and HMMT February 2024, and performs comparably to PPO and GRPO on AIME 2024 and HMMT February 2025. These results demonstrate that $A^\star$-PO effectively scales to long-context reasoning while remaining significantly more efficient in terms of training time and memory.

We sincerely thank the reviewer for the thoughtful and detailed feedback.

If 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.

Thank you for the suggestion. We will highlight the key theoretical innovations more explicitly in the main text in the camera-ready version. With the additional page for camera ready, we plan to expand on the technical challenges and proof introduced in the Appendix.

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.

This is a great observation. Indeed, using different values of $\beta$ leads to different KL-regularized objectives and corresponding optimal value functions. Our design choice is driven by empirical findings.

As $\beta \to \infty$, the KL-regularized objective becomes more conservative, and $V^\star(x)$ converges to the expected reward under the reference policy. Conversely, as $\beta \to 0$, $V^\star(x)$ approximates the $\text{Pass@}N$ performance of $\pi_{\text{ref}}$. The proofs are shown in Appendix E.

We use a larger $\beta_1$ in the offline stage to estimate $V^\star(x)$ and a smaller $\beta_2$ in the online stage. Our intuition is that using a larger $\beta_1$ results in a more stable and realistic estimate of the achievable value after policy training, avoiding overly optimistic targets. During the online stage, we use a smaller $\beta_2$ to loosen the KL constraint, allowing the policy to better optimize reward. Appendix D presents an ablation showing that both overly optimistic or pessimistic estimates of $V^\star(x)$ can degrade final performance.

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.

We are grateful for your interest and are happy to elaborate on the theoretical contributions. Our analysis introduces several novel techniques and findings:

• Getting rid of policy completeness assumption. In our analysis we only need policy realizability. In contrast, other on-policy policy optimization approaches including PG, PPO and REBEL rely on a stronger condition called policy completeness. We avoid such a condition by designing novel loss functions, which directly regress to the optimal policy rather than next-step policies similar to policy iteration.
• Getting rid of exploration mechanisms. Our method does not incorporate any explicit exploration mechanism such as optimism in the face of uncertainty or reset to some exploratory distribution as in existing works [1, 2, 3, 4, 5] This is a significant improvement because exploration is necessary in general RL settings. We achieve this by noticing that in the context of LLMs, there exists a close connection between the novel loss functions we proposed and the KL divergence between the learned policy and the optimal policy (illustrated in step 2 of the proof in Appendix H). Leveraging this observation, we managed to bound the suboptimality of the learned policy efficiently.
• New findings on nonconvex no-regret learning with OGD. In the case study of log-linear policy classes, we find that OGD provably converges with the optimal rate. This is a new finding because OGD doesn’t work for nonconvex losses in general. We achieve such a result by establishing a connection between the gradient of the loss and the distance from the optimal parameter (Lemma 3 in Appendix G). This connection enables us to efficiently bound the distance between the learned policy and optimal policy. Notably, the techniques used in Appendix G are totally different from the ones we use in Appendix H (Theorem 1), which is the reason why we can obtain a faster convergence rate.

[1] Xie, T., Foster, D. J., Krishnamurthy, A., Rosset, C., Awadallah, A., & Rakhlin, A. (2024). Exploratory preference optimization: Harnessing implicit q*-approximation for sample-efficient rlhf. arXiv preprint arXiv:2405.21046.

[2] Cen, S., Mei, J., Goshvadi, K., Dai, H., Yang, T., Yang, S., ... & Dai, B. (2024). Value-incentivized preference optimization: A unified approach to online and offline rlhf. arXiv preprint arXiv:2405.19320.

[3] Zhang, S., Yu, D., Sharma, H., Zhong, H., Liu, Z., Yang, Z., ... & Wang, Z. (2024). Self-exploring language models: Active preference elicitation for online alignment. arXiv preprint arXiv:2405.19332.

[4] Bai, C., Zhang, Y., Qiu, S., Zhang, Q., Xu, K., & Li, X. (2025). Online preference alignment for language models via count-based exploration. arXiv preprint arXiv:2501.12735.

[5] Chen, M., Chen, Y., Sun, W., & Zhang, X. (2025). Avoiding $\mathbf {exp (R_ {max})}$ scaling in RLHF through Preference-based Exploration. arXiv preprint arXiv:2502.00666.