动态分配计算资源,让大模型推理更快更准。
Optimizing Chain-of-Thought Reasoners via Gradient Variance Minimization in Rejection Sampling and RL
- 根据提示接受率动态调整采样数量,降低梯度方差。
- 数学推理任务中速度提升2-4倍,准确率显著提高。
- 适用于各类强化学习方法,适合追求高效推理的研究者。
大语言模型的思维链(CoT)推理可形式化为潜在变量问题,需生成中间推理步骤。现有方法如迭代奖励排序微调(RAFT)虽基于此框架,但对所有提示使用统一推理预算,忽略难度与收敛行为差异。本文发现,静态采样策略导致梯度估计效率低下是主要瓶颈。提出GVM-RAFT,一种提示相关的动态采样分配策略,在计算预算约束下最小化随机梯度方差。该方法通过监控提示接受率与梯度范数,动态分配资源,实现梯度方差最小化。理论分析表明,在合适条件下,该策略可加速收敛。实验显示,在数学推理任务中,GVM-RAFT相比原始RAFT实现2-4倍提速并显著提效。该动态采样策略通用性强,可集成至其他强化学习算法(如GRPO),带来类似收敛与测试准确率提升。代码已开源:https://github.com/RLHFlow/GVM。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。