论文

EPIG-Tree:梯度高效强化学习的计算最优分支

EPIG-Tree: Compute-Optimal Branching for Gradient-Efficient Reinforcement Learning

模型训练强化学习

摘要

基于奖励的语言模型强化学习,以组相对策略优化 (GRPO) 为例,将整个随机轨迹分解为单个标量奖励。这是干净且可扩展的,但它探索和分配奖励的效率低下:轨迹可能包含许多因果决策、恢复尝试和环境随机事件,但每个词元或动作都继承一个轨迹级别的优势。我们研究基于树的轨迹采样构建作为策略梯度估计的计算分配问题。我们的中心主张是,分支不应该放置在策略仅仅不确定的地方,而应该放置在附加分支最能减少每单位计算的策略梯度不确定性的地方。从局部策略梯度随机变量的总方差法则分解中,我们得出两个分配法则:新分支减少决策不确定性,而重复后缀轨迹采样减少连续不确定性。生成的 EPIG-Tree 分数使用已计算的 rollout 来分配分支。它估计占用率和得分加权值的不确定性,以及后缀定律 $n_e \propto w_e \|\nabla_θ\log π(a_e|h_e)\| σ_e / \sqrt{c_e}$。根据经验,EPIG 降低了克隆状态控制中的梯度 MSE,在 13 个环境扫描的所有 9 个密集连续控制环境中获胜,并近乎完美地恢复了参考梯度方向,并且它改进了相对于熵分支的冻结 LLM 梯度校准。在在线单轮数学中,树本地信用击败了平坦的 GRPO,而分支放置仅次于词元级信用分配。在在线多轮 Wordle 中,EPIG 获得了最高的最终胜率(0.850),超越了早期饱和于 0.790 的平坦 GRPO,并随着训练的进行而出现熵分支,证实了梯度估计优势转移到了有状态的大动作设置。