Title: Optimizing Chain-of-Thought Reasoners via Gradient Variance Minimization in Rejection Sampling and RL
Venue: NeurIPS 2025
Accepted: yes
Keywords: Large Language Models; Reinforcement Learning; Math Reasoning; Gradient Variance Minimization
Abstract: Chain-of-thought (CoT) reasoning in large language models (LLMs) can be formalized as a latent variable problem, where the model needs to generate intermediate reasoning steps. While prior approaches such as iterative reward-ranked fine-tuning (RAFT) have relied on such formulations, they typically apply uniform inference budgets across prompts, which fails to account for variability in difficulty and convergence behavior. This work identifies the main bottleneck in CoT training as inefficient stochastic gradient estimation due to static sampling strategies. We propose GVM-RAFT, a prompt-specific Dynamic Sample Allocation Strategy designed to minimize stochastic gradient variance under a computational budget constraint. The method dynamically allocates computational resources by monitoring prompt acceptance rates and stochastic gradient norms, ensuring that the resulting gradient variance is minimized. Our theoretical analysis shows that the proposed dynamic sampling strategy leads to accelerated convergence guarantees under suitable conditions. Experiments on mathematical reasoning show that GVM-RAFT achieves a 2-4x speedup and considerable accuracy improvements over vanilla RAFT. The proposed dynamic sampling strategy is general and can be incorporated into other reinforcement learning algorithms, such as GRPO, leading to similar improvements in convergence and test accuracy.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper makes an important contribution to optimizing CoT inference training for large language models (LLMs). The core innovation is to introduce dynamic resource allocation into the EM framework, and to address the bottleneck of static sampling strategies in existing methods (such as RAFT) by minimizing the gradient variance.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths：\n\n1） The GVM algorithm (Algorithm 2) is proposed to dynamically allocate the sample budget by monitoring the prompt acceptance rate and gradient norm to minimize the gradient variance (Equation 7). This solves the static sampling limitations of methods such as RAFT and significantly improves resource efficiency.\n\n2)  The paper provides a complete convergence analysis, proving that under the conditions of smoothness and convexity, GVM can ensure the monotonic decrease of the objective function. This fills the gap in the lack of theoretical support for existing RAFT work.\n\n3) The superiority of GVM is verified on multiple mathematical datasets.\n\n\nWeaknesses:\nAlthough the paper makes significant contributions, this paper have some issues:\n\n1) ​​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.\n\n2） 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).\n\n3) 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.\n\n4) ​​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.\n\n5) ​​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).\n\nQUESTIONS:\nPlease see the Weaknesses part.\n\nLIMITATIONS:\nNo", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-23"}

comments: SUMMARY:
This study identifies inefficient stochastic gradient estimation caused by static sampling strategies as the primary bottleneck in Chain-of-Thought (CoT) training. To address this challenge, the authors propose GVM-RAFT - a Prompt-Specific Dynamic Sample Allocation Strategy that optimizes stochastic gradient variance under computational budget constraints. The proposed algorithm is supported by theoretical guarantees, and comprehensive experiments demonstrate its consistent performance superiority over existing methods.

STRENGTHS AND WEAKNESSES:
Strengths:
1. This work identifies the main bottleneck in CoT training as inefficient stochastic gradient estimation due to static sampling strategies, and proposes a new algorithm to address this issue. I believe this is an important work.
2. The authors provide theoretical guarantees for the proposed algorithm, and conduct experiments to show the outperformance over existing methods.

Weakness:
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.
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?

QUESTIONS:
See the weaknesses.

LIMITATIONS:
Yes
TITLE: Official Review
CLARITY: 3
IMPACT: 4
ORIGINALITY: 4
DATE: 2025-07-03

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents an improved version of the iterative reward-ranked fine-tuning (RAFT) method.\n\nRAFT is a large language model fine-tuning method in which, at each step of fine-tuning, a reward model is used to select only the high-reward samples (i.e., prompts and their corresponding responses). This approach has been shown to effectively improve model fine-tuning and reduce noise in the training data.\n\nThis work argues that the RAFT method has limitations due to its static rejection of samples and its neglect of the coverage behavior across different prompts. To address this, the authors propose a dynamic rejection sampling method that dynamically allocates the sampling budget and rejection decisions per prompt. The paper provides both theoretical and empirical analyses of the method.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n\n1. The problem addressed in this paper is an important real-world challenge. Post-training of large language models (LLMs) requires effective fine-tuning strategies. Efficiently rejecting low-quality samples while reducing variance is critical for improving fine-tuning outcomes. Existing methods have not explored dynamic approaches in this area, and this paper offers a valuable incremental advancement in that direction.\n\n2. The paper provides both empirical and theoretical analyses of the proposed dynamic sampling allocation method, which enhances the quality and clarity of the work.\n\n### Weaknesses\n\n1. 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.\n\n2. 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.\n\nQUESTIONS:\n1. Why are the experiments only conducted on small-scale LLMs and methods?\n   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.\n\n2. When introducing theoretical concepts, can you provide more intuition behind them?\n   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.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-04"}