面向LLM RLVR加速的低秩优化轨迹建模

最近，针对大语言模型（LLM）进行基于可验证奖励的强化学习（RLVR）已成为有效的训练范式，以显著提高模型能力，这需要引导模型进行大量探索和学习，导致计算开销巨大，成为关键挑战。为减少训练步数，先前工作对模型参数进行线性外推。然而，LLM在RLVR训练期间的参数更新动力学仍不充分理解。为进一步考察LLM在RLVR训练中的演化，我们进行了经验性实验，发现模型的秩1子空间在训练中不呈线性演化，其对原始参数的主导作用在LoRA训练期间进一步被放大。基于上述洞察，我们提出了\textbf{N}onlinear \textbf{Ext}rapolation of low-rank trajectories（\textbf{NExt}），一种以非线性方式建模和外推低秩参数轨迹的新框架。具体而言，我们首先使用LoRA训练模型并提取多个训练步骤中参数差异的秩1子空间，用于后续非线性外推。随后，我们利用提取的秩1子空间训练一个预测器，来建模RLVR期间参数更新的轨迹，然后进行预测-外推过程，对模型参数进行外推，从而实现RLVR的加速。为进一步研究和理解NExt，我们进行了全面的实验，表明该方法有效且稳健。该方法约减少37.5%的计算开销，同时与广泛的RLVR算法和任务兼容。我们在https://github.com/RUCAIBox/NExt发布代码。