Optimizing Chain-of-Thought Reasoners via Gradient Variance Minimization
  in Rejection Sampling and RL
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.
6819a63c64ae18f1b60a5c43
Jiarui Yao
claimed_verified
2025-05-06T08:32:24.344Z
66f8689725464a7989b75845
/avatars/43a61a528c5779103eaf5687ba44ee14.svg
Jiarui Yao
user
FlippyDora
6819a63c64ae18f1b60a5c44
Yifan Hao
6819a63c64ae18f1b60a5c45
Hanning Zhang
admin_assigned
2025-05-06T08:59:20.459Z
6470e0f1cfd57849519033a5
/avatars/7ffefee3e36a4e37b9f4510bc6b689d1.svg
Hanning Zhang
user
HanningZhang
6819a63c64ae18f1b60a5c46
Hanze Dong
admin_assigned
2025-05-06T08:59:27.200Z
63a3ff69f91ad3ea5703841d
/avatars/69227c4bce01d33747c1377b6f9672db.svg
Hanze Dong
user
hendrydong
6819a63c64ae18f1b60a5c47
Wei Xiong
admin_assigned
2025-05-06T08:59:34.381Z
6319b29809baf858241f05de
/avatars/29eef2c52814abea82e2aa9bf37a7f9c.svg
Xiong
user
WeiXiong
6819a63c64ae18f1b60a5c48
Nan Jiang
admin_assigned
2025-05-06T08:59:47.400Z
64b8922ca1827cc8d04ae919
/avatars/0aaa83e3d09a82434e1d6af724aaa485.svg
Nan Jiang
user
nanjiang
6819a63c64ae18f1b60a5c49
Tong Zhang
Optimizing Chain-of-Thought Reasoners via Gradient Variance Minimization
  in Rejection Sampling and RL
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.
GVM-RAFT, a dynamic sampling strategy for chain-of-thought reasoning in large language models, improves convergence and accuracy by adaptively allocating computational resources.
chain-of-thought large language models latent variable problem iterative reward-ranked fine-tuning stochastic gradient estimation dynamic sample allocation prompt-specific computational budget constraint prompt acceptance rates stochastic gradient norms GRPO reinforcement learning