本文提出动态采样预算分配和温度调度机制,通过基于问题难度的资源再分配和维持策略熵的探索能力,显著提升了大型语言模型在数学任务中的强化学习效率和性能,尤其在AIME 2024基准上pass@1和pass@16分别提高5.31%和3.33%。
Reinforcement Learning, Large Language Model, Efficiency, Reasoning, Exploration
Mengqi Liao, Xiangyu Xi, Ruinian Chen, Jia Leng, Yangen Hu, Ke Zeng, Shuai Liu, Huaiyu Wan
Beijing Jiaotong University, Meituan
Generated by grok-3
Background Problem
大型语言模型(LLMs)在复杂任务(如数学和编码)中表现出色,强化学习(RL)因此成为提升其推理能力的重要方法。然而,现有RL方法在训练效率和探索能力上存在两大问题:一是所有问题均分配相同数量的采样预算(rollouts),导致简单问题浪费资源而困难问题采样不足;二是RL训练虽能提高回答精度,但会限制模型的探索能力,可能导致性能低于未经RL训练的基础模型。作者旨在解决这些问题,通过更高效的资源分配和探索机制提升LLMs在复杂任务中的表现。
Method
本文提出了两种主要方法来改进强化学习(RL)在大型语言模型(LLMs)中的应用:
- 动态采样预算分配(Dynamic Rollout Budget Allocation):核心思想是根据问题的难度动态分配采样预算,简单问题减少采样,困难问题增加采样,以提高训练效率。具体步骤包括:首先通过累积平均奖励对问题难度进行建模(难度越高,平均奖励越低),然后基于难度的排名调整每个问题的采样预算(G),确保批次内总预算不变,同时设置最小和最大预算限制(Gmin和Gmax),并随训练迭代逐步调整这些限制,类似于课程学习。
- 温度调度机制(Temperature Scheduler):为避免RL训练中策略熵快速下降导致探索不足,作者提出通过动态调整softmax分布的温度(τ)来维持策略熵的稳定水平,避免引入有害梯度(如熵正则化可能导致的模型崩溃)。具体公式为 ,其中α基于初始熵和当前熵的比值。此外,结合退火机制(Annealing)在训练后期逐步降低目标熵,促进从探索到利用的平衡。
批判性思考:动态采样预算分配的理论依据合理,但其依赖于平均奖励作为难度指标可能存在偏差,例如早期训练中模型能力不足可能导致奖励不准确,从而影响分配的有效性。温度调度机制避免了熵正则化的梯度问题,但其公式假设熵值较小,可能在某些场景下不适用,且退火参数(如η)的选择对结果影响较大,论文中未充分探讨其鲁棒性。
Experiment
实验在数学领域的多个基准数据集上进行,包括AIME 2024、AMC 2023、MATH 500和OlympiadBench,使用DeepSeek-R1-Distill-Qwen 1.5B和7B模型作为基础模型,训练数据集包含10k高质量数据点,验证集为0.5k。实验设置包括批次大小64,默认采样次数(G)为8,训练3个epoch,并重复3次以减少随机性影响。评估指标为pass@1和pass@16,反映模型精度和探索能力。
结果分析:与基线方法GRPO和DAPO相比,作者方法在7B模型上AIME 2024基准的pass@1和pass@16分别提升了5.31%和3.33%,在多个基准上pass@16表现一致优于GRPO,显示出更强的探索潜力。1.5B模型也表现出类似改进。消融实验表明,动态采样预算分配对困难问题(如AIME)的性能提升显著,去除此机制后pass@1下降约3%。温度调度机制有效维持了熵水平,降低了训练过程中的准确率方差,提升了稳定性,但退火机制在简单问题上效果更好,而在困难问题上仅温度调度表现更优。
批判性思考:实验设置较为合理,涵盖了不同规模模型和多个基准,但训练数据规模(10k)和采样次数(G=8)受限于计算资源,可能限制了动态分配机制的潜力。此外,实验未充分探讨不同难度数据集上的细粒度表现,难以判断资源分配是否真正均衡。结果虽显示改进,但与基线的差距在某些基准上较小,且未与其他探索增强方法(如ε-greedy)进行对比,难以评估其相对优越性。退火机制的参数选择(如η=0.9)缺乏充分调优依据,可能影响结果的普适性。
Further Thoughts
本文提出的动态采样预算分配机制启发了一个更广泛的问题:是否可以在其他AI训练范式中(如监督学习或自监督学习)应用类似的基于难度的资源分配策略?例如,在数据集不平衡的情况下,是否可以动态调整训练样本的权重或采样频率以提高模型对困难样本的学习效果?此外,温度调度机制避免了熵正则化的有害梯度问题,这与近期一些关于RL中探索-利用平衡的研究(如meta-gradient方法)有潜在联系,未来可以探索结合meta-learning来动态调整温度或退火参数,以适应不同任务的难度分布。另一个值得思考的方向是,作者提到的基于难度设置不同温度的想法,可能与多智能体系统中基于角色或任务难度的策略调整有相似之处,是否可以借鉴多智能体RL中的协作机制来进一步优化单模型的探索策略?这些跨领域联系可能为本文方法提供更广阔的应用场景和改进空间。