通过拒绝采样与强化学习中的梯度方差最小化优化思维链推理器
机器学习
2025-05-06 v1 人工智能
计算与语言
摘要
大型语言模型(LLMs)中的思维链(CoT)推理可以被形式化为一个隐变量问题,其中模型需要生成中间的推理步骤。虽然先前的方法,如迭代奖励排序微调(RAFT),依赖于此类形式化,但它们通常对所有提示应用统一的推理预算,这未能考虑到难度和收敛行为的可变性。本工作识别出思维链训练的主要瓶颈是由于静态采样策略导致的低效随机梯度估计。我们提出了GVM-RAFT,一种针对特定提示的动态样本分配策略,旨在计算预算约束下最小化随机梯度方差。该方法通过监控提示接受率和随机梯度范数来动态分配计算资源,确保最终的梯度方差最小化。我们的理论分析表明,所提出的动态采样策略在合适的条件下能带来加速收敛的保证。在数学推理上的实验表明,GVM-RAFT相比普通的RAFT实现了2-4倍的加速和显著的准确率提升。所提出的动态采样策略具有通用性,可以整合到其他强化学习算法中,如GRPO,从而在收敛速度和测试准确率上带来类似的改进。我们的代码可在https://github.com/RLHFlow/GVM获取。
引用
@article{arxiv.2505.02391,
title = {Optimizing Chain-of-Thought Reasoners via Gradient Variance Minimization in Rejection Sampling and RL},
author = {Jiarui Yao and Yifan Hao and Hanning Zhang and Hanze Dong and Wei Xiong and Nan Jiang and Tong Zhang},
journal= {arXiv preprint arXiv:2505.02391},
year = {2025}
}