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

TL;DR

This work tackles the inefficiency of chain-of-thought (CoT) training in large language models by reframing CoT as a latent-variable problem under an EM-based RAFT-style framework and identifying gradient-variance from uniform sampling as the bottleneck. It introduces Gradient Variance Minimization (GVM) with Dynamic Sample Allocation (GVM-RAFT), a per-prompt budgeting strategy that allocates samples according to acceptance rates $p_i^t$ and gradient magnitudes $G_i$ under a total budget $N$, yielding an unbiased gradient estimator and a closed-form allocation $n_i^t$. Theoretical results establish convergence guarantees under standard smoothness assumptions and quantify the variance reduction via the $ ext{Ω}(k,T)$ term, while practical implementations estimate $p_i^t$ and $G_i$ with a small pre-sampling phase and update budgets periodically. Empirically, GVM improves convergence speed by roughly 2–4× and often enhances final accuracy on mathematical reasoning benchmarks, with demonstrated transferability to RL-style methods such as GRPO. The approach is general, online, and can be integrated into various post-training RL frameworks, offering a principled path to more efficient CoT reasoning training with broader applicability.

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. Our code is available at https://github.com/RLHFlow/GVM.

Figures (7)

• Figure 1: The demonstration of the whole pipeline for GVM. Starting from $\pi_0$, which a supervised fine-tuning (SFT) model, GVM will iteratively conduct the estimation and budget assignment process for $T$ times according to the description in Algorithm \ref{['alg:meta']}. Each iteration could be decomposed into two stages, and the final policy model will be returned after those iterations.
• Figure 2: The assigned sample budget in GVM-RAFT++ with $N^\prime=32,N=32n$ for three iterations.
• Figure 3: The average accuracy of RAFT++ and GVM-RAFT++ with sample sizes 8 and 32, respectively, on Math500, Minerva Math, and Olympiad Bench, with base model Qwen2.5-Math-1.5B.
• Figure 4: The average accuracy of GVM-RAFT++ with different sample sizes for both stage one and stage two on Math500, Minerva Math and Olympiad Bench. According to the trend, though the number of sample sizes in both stage one and stage two increases, the convergence rate stays almost the same.
• Figure 5: The average accuracy of RAFT++, GVM-RAFT++ and GRPO, GVM-GRPO with sample sizes $N^\prime=8, N=8n$ respectively on Math500, Minerva Math and Olympiad Bench. Applying the GVM sample strategy to RL algorithms like GRPO achieves similar results to GVM-RAFT++ compared to vanilla GRPO. Here for vanilla RAFT++ and GRPO, the rollout number per prompt is set to 8 as well.

Theorems & Definitions (13)

• Lemma 1: Unbiased Gradient Estimator
• proof
• Proposition 1
• Remark 1
• Theorem 1: Decreasing rate with smoothness condition.
• Theorem 2: Decreasing rate with smooth and convex condition.
• Theorem 3: Decreasing rate with smooth and strongly convex condition.
• proof : Proof of Proposition \ref{['prop:opt']}
• proof : Proof for Theorem \ref{['thm:smo']}
• proof : Proof for Theorem \ref{['thm:conv']}