PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jiarui Yao, Yifan Hao, Hanning Zhang, Hanze Dong, Wei Xiong, Nan Jiang, Tong Zhang

arXiv:2505.02391v1cs.LGcs.AIcs.CL

TL;DR

Uniform prompt sampling in RAFT overlooks differences in prompt difficulty and convergence, making stochastic gradient estimation a key CoT-training bottleneck. The paper proposes prompt-adaptive sampling to reduce gradient variance under a computational budget. On mathematical reasoning tasks, it reports 2–4× faster convergence and improved final test accuracy.

  • Problem

    RAFT-style uniform sampling overlooks prompt-specific differences in sample efficiency and convergence, while CoT training faces a Monte Carlo ELBO-gradient estimation bottleneck.

    Source: §1 · p. 2

  • Method

    GVM dynamically allocates sampling resources across prompts using prompt hardness and gradient norms to minimize stochastic-gradient variance under a computational budget.

    Source: §1 · p. 3

  • Results

    2–4× speedup in convergence and considerably improved final test accuracy are reported for the approach on mathematical reasoning tasks.

    Source: §1 · p. 3

  • Takeaways & Limitations

    GVM is reported to improve convergence in rejection-sampling pipelines and is generalized to RL algorithms such as GRPO.

    Source: §5 · p. 15

    Experiments use Qwen-series base models, so effectiveness on other base models remains to be verified.

    Source: §5 · p. 15

Abstract

from arXiv · show

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.

1 Introduction

The introduction frames CoT training as latent-variable learning and identifies uniform prompt sampling as a source of inefficient gradient estimation. GVM uses prompt-adaptive sampling, with reported gains in convergence and test accuracy.

1 Introduction

  • RAFT is presented as an EM-style approach to CoT reasoning, with generated responses selected in the E-step and model fine-tuning in the M-step.The introduction describes this as an approximation to the posterior over latent rationales.Source: §1 · pp. 1–2
  • Uniform best-of-n sampling overlooks prompt-specific sample efficiency and convergence behavior, leaving stochastic gradient estimation inefficient.The paper identifies Monte Carlo estimation of the ELBO gradient during the E-step as a key bottleneck.Source: §1 · pp. 2–3
  • GVM allocates sampling resources by prompt hardness and gradient norms to minimize stochastic-gradient variance under computational-budget constraints.The authors also report theoretical convergence guarantees under suitable conditions.Source: §1 · p. 3
  • 2–4× speedup in convergence and considerably improved final test accuracy are reported for the approach on mathematical reasoning tasks.The experiments apply the method to RAFT++ and GRPO.Source: §1 · p. 3

1.1 Related Work

  • Prior budget-allocation work by Tong et al. (2024) addresses offline synthetic-data construction, but lacks dynamic scheduling for iterative online algorithms.The paper distinguishes its approach from that setting by emphasizing dynamic allocation during iterative training.Source: §1.1 · p. 3
  • Related CoT-EM methods differ in focus: TRICE uses MCMC and a prompt–CoT table, while this work schedules rejection-sampling budgets and analyzes convergence.The paper also contrasts its iterative fine-tuning focus with Hu et al. (2024), which focuses on in-context learning.Source: §1.1 · pp. 3–4

2 Problem Formulation and EM Framework

This section formulates CoT reasoning as a latent-variable problem and connects EM updates to rejection sampling. It motivates prompt-specific sampling by the cost of posterior estimation and the uneven acceptance rates induced by uniform budgets.

2 Problem Formulation and EM Framework

  • The model represents a prompt x, intermediate CoT y, and final answer z, and targets minimizing the negative log-likelihood of the correct answer.The setup omits reference CoT rationales, reflecting practice after pre-training or supervised fine-tuning.Source: §2 · p. 4
  • Introducing latent CoT y is motivated by the difficulty of estimating P(z|x) directly; the EM framework iteratively approximates the posterior over latent variables.The joint model class is defined through P(y|x,θ) and P(z|y,θ).Source: §2 · p. 4
  • In the E-step, the posterior over latent rationales is updated to tighten the objective bound; in the M-step, model parameters are optimized with that posterior fixed.The posterior is induced by the model and the observed prompt-answer pair; its entropy is constant during the M-step.Source: §2 · p. 5
  • Computing the posterior and its objective requires summing over all possible CoT rationales, so the paper approximates them by sampling.Rejection sampling approximates the target posterior using the model’s proposal distribution.Source: §2 · p. 5
  • Best-of-n sampling keeps responses with correct answers, connecting rejection sampling to RAFT when final-answer randomness given a rationale is low.The paper notes that the EM–RAFT connection was previously observed by Singh et al. (2023) and Sordoni et al. (2023).Source: §2 · pp. 5–6
  • With a fixed inference budget, uniform sampling can bias accepted examples toward easier prompts because difficult prompts have lower acceptance rates.The paper describes difficult prompts as having proposal distributions concentrated on incorrect rationales, while easier prompts yield more accepted responses.Source: §2 · p. 6
  • The algorithm alternates posterior estimation and prompt-specific sample-budget assignment with rejection sampling, then updates model parameters by gradient descent.It repeats these E- and M-step operations for T iterations and returns the final model.Source: §2 · p. 6

2.1 Notation; 3.2 Theoretical Result

  • The notation section fixes attention on J_Qt(θ) because the entropy term is constant in the M-step; under a finite budget, scheduling can reduce the variance term Ω(k,T).The paper says increasing sample size also reduces this term and points to Algorithm 2 for its scheduling scheme.Source: §2.1 · p. 6

3 Gradient Variance Minimization by Dynamic Sample Allocation

The method allocates prompt-specific inference budgets to address variation in acceptance rates and gradient-estimation variance. It combines a variance-aware schedule with theoretical guarantees and practical extensions.

3 Gradient Variance Minimization by Dynamic Sample Allocation

  • GVM dynamically allocates inference budgets across prompts to reduce gradient-estimator variance under a fixed total sampling budget.The schedule adapts to prompt acceptance rates.Source: §3 · p. 6, §3.1 · pp. 7–8

3.1 Dynamic Inference Budget Scheduling to Minimize Gradient Variance

  • The method constructs an unbiased gradient estimator from accepted samples drawn by rejection sampling for each prompt.Its expectation matches the gradient under the EM objective.Source: §3.1 · p. 7
  • The allocation penalizes sampling on extremely low-acceptance prompts, while a lower bound ensures sufficient accepted samples during training.The allocation is obtained by solving a regularized optimization problem.Source: §3.1 · p. 8

3.2 Theoretical Result

  • Under smoothness, sufficiently large sample sizes guarantee loss decrease; under convexity, an additional negative term indicates faster decrease.The smooth-and-convex result assumes a step size 0 < η ≤ γ/2.Source: §3.2 · pp. 9–10
  • Updating the budget more often can improve training efficiency but raises computational overhead, so update frequency requires balancing both.Less frequent updates reduce computation, but small gradient norms as optimization proceeds can slow training.Source: §3.2 · p. 10

3.3 Practical Implementation

  • In practice, sampled estimates are plugged into the allocation rule; GVM is also described as adaptable to GRPO fine-tuning.For GRPO, the passage introduces token advantages based on each response’s final reward.Source: §3.3 · pp. 10–11
  • Allocated sample budgets rise smoothly as prompt acceptance rates fall, while Lipschitz coefficients also shape the allocation.Figure 2 shows GVM-RAFT++ budgets for N′ = 32 and N = 32n over three iterations.Source: Figure 2 · p. 11, §3.3 · p. 11

4 Experiments and Results

The experiments evaluate GVM sample-budget rebalancing with RAFT++ and GRPO on mathematical reasoning tasks. Results report faster convergence and improved performance, while the ablation finds convergence largely unchanged as sample sizes increase.

4.1 Experiments Setup

  • Experiments use Qwen2.5-Math-1.5B and Qwen2.5-Math-7B with Math-Verify as verifier, and evaluate average accuracy from eight samples.The 1.5B model uses evaluation temperature 1.0; the 7B model uses 0.7.Source: §4.1 · p. 12, §4.2 · p. 12

4.2 Main Results

  • Across five benchmarks, Table 1 reports that GVM-RAFT++ and GVM-GRPO outperform RAFT++ and GRPO without sample-size reweighting.The experiments use Math500, Minerva Math, Olympiad Bench, AIME24, and AMC23.Source: Table 1 · p. 12
  • For Qwen2.5-Math-7B, final performance is comparable to baselines while convergence is faster.The authors attribute the result to budget rebalancing using prompt acceptance rates and Lipschitz coefficients.Source: §4.2 · p. 13
  • With N′=8 and 32, GVM-RAFT++ is about 2× and 4× faster, with about 1.25% and 5% final-accuracy gains, respectively.The speedups are measured in update steps against the corresponding vanilla versions.Source: §4.2 · p. 13, Figure 3 · p. 13
  • At N′=8, 16, and 32, GVM-RAFT++ training rewards are respectively higher, approximately equal, and lower than RAFT++ rewards.The authors suggest larger N′ may shift budget toward harder prompts, which are more likely to be solved at least once.Source: §4.2 · p. 13
  • GVM-GRPO achieves convergence performance similar to GVM-RAFT++, supporting use of the sampling strategy outside the EM pipeline.The comparison applies the same sample-budget rebalancing strategy to GRPO.Source: §4.2 · p. 13, Figure 5 · p. 14

4.3 Ablation Study

  • Increasing stage-one and stage-two sample sizes does not evidently change convergence, suggesting prompt difficulty can be estimated sufficiently with relatively small budgets.The figure’s training-step definition corresponds to 1024 × 4 samples, so it may not map directly to an actual training step.Source: §4.3 · p. 14, Figure 4 · p. 14

5 Conclusion and Discussion

The conclusion presents GVM as a two-stage, prompt-specific sampling and updating strategy, with theoretical guarantees and experiments showing faster convergence. The authors note that broader model and RL-algorithm validation remains future work.

  • GVM assigns prompt-specific sampling budgets through a difficulty-estimation pre-sampling phase followed by parameter optimization.The authors describe the strategy as a fine-grained dynamic allocation and updating method.Source: §5 · p. 15
  • Experiments report faster convergence and sometimes superior final performance under suitable settings across rejection-sampling pipelines and GRPO.The conclusion characterizes this as evidence of potential for adaptive sampling and update-budget reweighting.Source: §5 · p. 15
  • The paper establishes theoretical performance guarantees for this class of two-stage algorithms.Source: §5 · p. 15
  • Experiments use Qwen-series base models, and broader base-model validation and further PPO or Reinforce experiments remain future work.Source: §5 · p. 15

A Author Contributions

The author-contribution statement assigns implementation, experimentation, theory, baseline work, project advice, and computational support across the team.

  • JY implemented the workflow, ran preliminary and scaled experiments, and primarily wrote the experiments sections.Source: A Author Contributions · p. 18
  • YH initiated the project, contributed to the algorithm and pipeline clarification, derived GVM’s theoretical analysis and bounds, and primarily wrote methodology and theory sections.The contributions are described across the adjacent author-contribution passages.Source: A Author Contributions · p. 18
  • HZ conducted experiments and ablations on entropy-loss and model-performance trade-offs, and reviewed related work on mitigating early RL-training entropy loss.Source: A Author Contributions · p. 18
  • HD implemented baselines, ran RAFT++ and GRPO experiments, advised experiment design, and contributed to writing.Source: A Author Contributions · p. 18
  • WX advised on theory, experiment design, and writing while contributing to the codebase, baseline training, and draft revision; NJ and TZ provided resources, guidance, and suggestions.Source: A Author Contributions · p. 18

B Experiment Details

The experiment details describe GVM’s parameter choices, gradient-based budget assignment, and practical allocation procedure, alongside implementation settings for GRPO.

  • In stage two, smaller α or β allocates more sample budget to harder prompts when the other parameter is fixed.The experiments set α=0.001 and β=2.0.Source: B Experiment Details · p. 19
  • The gradient calculation compares lm_head and embed_tokens, while summation rather than sequence-length averaging is chosen for gradient reduction.The authors report that layer choice makes little difference to final sample sizes; summation biases allocation toward lower-accept-rate problems.Source: B Experiment Details · pp. 19–20
  • GVM sorts prompts by descending acceptance rate in stage two and adds samples to lower-acceptance prompts until the total budget N is met.Each reported iteration uses nine update steps because its training data is slightly below the amount needed for ten steps.Source: B Experiment Details · p. 20
  • For GVM-GRPO, the calculated second-stage budget is divided by a fixed size and prompts are copied accordingly, with rollout number set to four.Source: B Experiment Details · p. 20
  • A GVM iteration with N′=8 and N=8n typically takes 90 minutes on a 4×H100 server.Source: B Experiment Details · p. 20

C Notation Table

Table 3 collects the paper’s notation for reference.

D More Theoretical Results and Missing Proofs

The appendix provides additional theoretical conditions and proof steps. The supplied passages identify assumptions and proof techniques, but do not include complete theorem conclusions.

Theorem 3

Proofs for Proposition 1 and Theorem 1

Proof for Theorem 2

E Technical Lemmas

The appendix states Jensen’s inequality for convex functions and notes its application to norms.

Lemma 2 (Jensen’s Inequality)

  • For convex ϕ, Jensen’s inequality gives ϕ(Ew∼p w) ≤ E_w∼p ϕ(w), including the norm by convexity.The proof is referenced to Proposition A.9 of Zhang (2023).Source: E Technical Lemmas · p. 24

F More about Rejection Sampling

The appendix presents a general rejection-sampling pipeline and notes that practical prompt selection uses outcome-based reward signals.

Rejection sampling

  • In practice, outcome-based reward signals decide whether a prompt is selected.Source: F More about Rejection Sampling · p. 24
  • Algorithm 3 samples candidate responses from q(y|xi, zi), draws uniform random values, and uses an acceptance test to build an output set.The displayed algorithm excerpt provides the inputs, sampling steps, and set output, but not the complete acceptance expression.Source: F More about Rejection Sampling · p. 24
  • Figure 7 compares gradient-calculation layer choices and gradient-reduction methods on Qwen2.5-Math-7B.Source: Figure 7 · p. 25

G Case Study

In the showcased cases, vanilla RAFT++ misses correct answers within a fixed attempt budget, while GVM-RAFT++ or GVM-GRPO finds them. The examples include arithmetic, a list problem, and a triangular-array problem.

G Case Study

  • Across the showcased benchmark instances, vanilla RAFT++ does not find correct answers within a fixed attempt budget, while GVM-RAFT++ or GVM-GRPO does.The passages give examples from several benchmarks.Source: G Case Study · p. 26
  • 4 distinct values result from inserting parentheses; RAFT++ reports 3, while GVM-RAFT++ reports 4.The listed values in the GVM-RAFT++ response are 121, 122, 126, and 144.Source: G Case Study · p. 26
  • GVM-GRPO’s response identifies the list as [9, 9, 3, 1, 8] and calculates its sum of squares as 236, matching the problem’s stated answer of 236.Source: G Case Study · pp. 30–31
  • 5 is the stated answer for the 2023rd-row units digit; GVM-RAFT++ reports 5, whereas RAFT++ reports 1.The two model outputs give different conclusions for the same triangular-array question.Source: G Case Study · pp. 31–32
Loading 2505.02391v1…