Optimizing Chain-of-Thought Reasoners via Gradient Variance Minimization in Rejection Sampling and RL

Q 1. Experimental generalization is insufficient: The experiment is based only on the Qwen model series (Qwen2.5-Math-1.5B/7B), and the effect on other mainstream LLMs (such as Llama or GPT architectures) is not verified. This limits the generalizability of the conclusions.

Response: Thanks for the suggestion for adding more generalizability validation experiments. We conducted more experiments with Llama series models, on both RAFT++ and GRPO algorithms. Due to computation limitation, we first conduct experiments with models at 1B and 3B scale, and will continue to do experiments on larger scales like 7B. We report the metric average @ 8 as follows. One iteration in GVM experiments here corresponds to 9 steps of optimization in vanilla experiments. It could be seen GVM achieves a nearly 2 times acceleration for comparable performance.

Q 2. Computational overhead is not quantified: The pre-sampling stage of GVM (Algorithm 2) requires additional samples to estimate, but the paper does not report the computational cost (such as GPU hours or total number of samples).

Response: We’d like to mention that there is a tradeoff for the sampling budget allocation estimation - if the samples used to estimate the difficulty of the problems is large, it takes more time to generate while would return a more accurate estimation; and if the samples used in the estimation stage is small, the estimation would be less accurate but more efficient. In our main experiments, we choose $N^\prime=8$ as the number of responses per prompt in our stage 1 difficulty estimation. On an 8*A100 GPU instance, it takes about 20 minutes additional computation for the GVM stage, which may slightly increase as the training goes due to the response length increasing. Compared to the RL stage which takes about 100 minutes per iteration, the additional overhead introduces 0.2 times the original training time. If we set $N=c\cdot N$ where $c$ is the average samples per prompt in the RL training stage, then larger $c$ will correspond to less additional overhead introduced by the GVM stage.

Q 3. The comparative benchmark is not comprehensive: Table 1 shows that GVM outperforms RAFT++ and GRPO, but it is not compared with the latest RLHF methods (such as online DPO or PPO variants), which may underestimate competing methods.

Response: Our algorithm serves as a plug in algorithm for sampling budget allocation and should be combined with other RL optimization algorithms together for policy model update. GVM is not a new standalone RL algorithm itself, thus we choose the commonly used PPO and GRPO as the underlying RL optimization algorithms. The whole training pipeline could be divided into two stages - with the first one corresponding to GVM sampling budget estimation, and the second one corresponding to the vanilla RL optimization.

Q 4. Limitations of theoretical assumptions: Theorems 1-2 rely on smoothness (1/γ-smooth) and convexity assumptions, but the LLM loss function is often non-convex. The paper does not discuss convergence guarantees in non-convex scenarios, which may overestimate the robustness of the theory.

Response: Thank you for the insightful comment. In fact, Theorem 1 relies only on the smoothness assumption and does not require convexity, so it remains applicable in non-convex settings. We agree that this point should be clarified more explicitly in the main text, and we will revise the manuscript accordingly to avoid any confusion.

Q 5. Fuzzy implementation details: The GVM-RAFT++ loss function (Equation 8) in Section 3.3 introduces importance sampling and clipping, but does not explain the choice of hyperparameters (e.g., ϵ value).

Response: The clipping parameter used in the objective of RAFT++ is $0.2$ in general, while we adopt the clip-higher trick from DAPO [1] with $\epsilon_{\rm high}=0.3$ for experiments like GVM-RAFT++ to mitigate the rapid entropy loss drop. The parameters should be adjusted according to the training dynamic like entropy loss curve. For other parameters, please refer to Table 2 in Appendix B.

[1] Yu, Qiying, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yu Yue, Weinan Dai et al. "Dapo: An open-source llm reinforcement learning system at scale." arXiv preprint arXiv:2503.14476 (2025).

Q 1. For the theoretical guarantees, with a sufficiently large sample size, Ω(k,T) becomes small, and the upper bound for E[L(θKT)−L(θ∗)]−E[L(θ0)−L(θ∗)] is strictly negative. How about the case when the sample size is not large enough? The authors do not provide the details about how large the sample size needs to be to make the RHS strictly negative.

Response: Thanks for pointing this out. o ensure that the right-hand side (RHS) of the bound is strictly negative, we require the following condition to hold for each prompt $x_i$: $\sum_{i=1}^{n} b_i^{-1} \left\lVert E_y \nabla_{\theta} ln P(y, z_i | x_i, \theta) \right\rVert_{2}^{2} \le \left\lVert \sum_{i=1}^{n} E_y \nabla_{\theta} \ln P(y, z_i \mid x_i, \theta) \right\rVert_2^2$ , where $b_i$ is the batch size for $x_i$, and $n_i \ge b_i$ is the total allocated budget for $x_i$. In practice, this condition is not stringent and is often satisfied, since a proper $b_i$ can help reduce the variance and norm of the gradient on the left-hand side. However, in scenarios with extremely small sample sizes, we cannot guarantee the stability or performance of the EM algorithms, as the estimation becomes highly noisy. Nevertheless, compared to the standard EM algorithm, GVM-EM still provides theoretical advantages by dynamically allocating samples to reduce the term $\Omega(k, T)$, which contributes to better convergence behavior even under limited data.

Q 2. For the GVM-GRPO, the authors have not provided the details of how to generalize the design to GVM-GRPO. Is it directly using the Dynamic Sample Allocation Strategy to generate different numbers of answers for different prompts in GRPO?

Response: The sampling allocation strategy is shared when applying to different optimization algorithms like RAFT or GRPO. It is disentangled from the underlying RL training algorithm, and thus could be applied to different algorithms after allocating the sample budgets. In our experiments GVM-RAFT++ and GVM-GRPO, the sampling stage for GVM is exactly the same.

Q 1. The overall originality of the work is incremental, as the core idea of intelligent rejection sampling has already been explored in previous methods such as RAFT. This paper builds on existing approaches rather than introducing a fundamentally new concept.

Response: Thanks for the opinion. The previous works like RAFT are essentially RL algorithms, which focus how to utilize existing data to optimize the policy model. In contrast, our method focuses on improving the sampling strategy of training data, instead of proposing a novel training algorithm. Take RAFT as an example, for a given prompt $x$, it will sample n responses ${y_1,\cdots,y_n}$, and then call a reward model to obtain the rewards for each response ${r_1,\cdots,r_n}$ and select the response with the highest reward to form a training pair $(x,y_{\arg\max_i r_i})$. It assign each prompt $x$ with equal sampling budget $n$, without differentiating different samples based on their characteristics like difficulty and contribution to the model (measured by the gradient). Our GVM algorithm considers both the sample difficulty and gradients, and assign different sampling budget based on them. Then we could feed the samples to different RL algorithms like RAFT or GRPO, etc. It is a two-stage pipeline instead of substituting other algorithms.

Q 2. As noted in the conclusion, the method is only evaluated on a limited set of baselines and models—specifically the Qwen and PPO series. It would be valuable to see this method extended and tested on a broader range of LLMs and training setups.

Response: We select Qwen model series as the base models following the common practice in the line of LLMs math reasoning research, where Qwen serves as one of the strongest and most commonly used open-source base model. As mentioned in question 1, our GVM algorithm serves as a pre-stage of the vanilla RL optimization stage, therefore it could be combined with any RL algorithms. For demonstration purpose, we choose two commonly adopted ones - RAFT and GRPO. To make our conclusion more generalizable, we conduct more experiments on Llama series base models. Due to computation limitation, we first conduct experiments with models at 1B and 3B scale, and will continue to do experiments on larger scales like 7B. We report the metric average @ 8 as follows. One iteration in GVM experiments here corresponds to 9 steps of optimization in vanilla experiments. It could be seen GVM achieves a nearly 2 times acceleration for comparable performance.

Q 3. Why are the experiments only conducted on small-scale LLMs and methods? Adding more justification for this choice, and including experiments on other variants or larger-scale models, would significantly strengthen the paper and increase its impact.

Response: We conduct our main experiments on Qwen 1.5B and 7B base models due to the computation source limitation, and conduct additional experiments on Llama 1B and 3B series during the rebuttal period with an 8*A100 GPU instance. For experiments on larger model scale, we hope to generalize GVM when more computation resources are available.

Q 4. When introducing theoretical concepts, can you provide more intuition behind them? For example, when discussing the optimal budget allocated for each prompt, it would be helpful to explain why the given equation makes sense. While a complete theoretical analysis is valuable, offering some high-level intuition would make the paper more accessible to a broader audience.

Response: In our algorithm, we allocate the budget for each prompt according to Proposition 1, which takes into account both problem difficulty and model efficiency. Specifically, to address problem difficulty, we allocate more budget to harder problems, which associated with larger gradients and lower accept rates during sampling, as this can improve model performance. At the same time, to ensure model efficiency, we avoid allocating excessive budget to extremely difficult problems with near-zero acceptance rates, as this would lead to inefficiency.