LLM RLVR 가속화를 위한 저랭크 최적화 경로 모델링

Low-rank Optimization Trajectories Modeling for LLM RLVR Acceleration

최근, 대규모 언어 모델(LLM)에 대한 검증 가능한 보상을 활용한 강화 학습(RLVR)은 모델 성능을 크게 향상시키는 효과적인 학습 패러다임으로 부상했지만, 모델이 광범위한 탐색과 학습을 수행하도록 유도해야 하므로 상당한 계산 비용이 발생하며, 이는 주요 과제가 되고 있습니다. 훈련 단계를 줄이기 위해 기존 연구에서는 모델 파라미터를 선형 보간하는 방법을 사용합니다. 그러나 RLVR 훈련 과정에서 모델 파라미터 업데이트의 동역학은 아직 충분히 이해되지 않고 있습니다. 본 연구에서는 RLVR 훈련 과정에서의 LLM의 진화를 더 자세히 조사하기 위해 실험을 진행한 결과, 모델의 랭크-1 부분 공간이 선형적으로 변하지 않으며, 특히 LoRA 훈련 과정에서 원래 파라미터에 대한 지배력이 더욱 강화되는 것을 확인했습니다. 이러한 통찰력을 바탕으로, 우리는 저랭크 파라미터 경로를 비선형적으로 모델링하고 추정하는 새로운 프레임워크인 **N**onlinear **Ext**rapolation of low-rank trajectories (**NExt**)를 제안합니다. 구체적으로, 우리는 먼저 LoRA를 사용하여 모델을 훈련하고, 여러 훈련 단계에서 파라미터 차이의 랭크-1 부분 공간을 추출합니다. 이 추출된 부분 공간은 이후의 비선형 추정에 사용됩니다. 그런 다음, 추출된 랭크-1 부분 공간을 사용하여 RLVR 훈련 과정에서 파라미터 업데이트 경로를 모델링할 수 있는 예측 모델을 훈련하고, 예측-확장 과정을 통해 모델 파라미터를 추정하여 RLVR을 가속화합니다. NExt의 효과와 견고성을 입증하기 위해 종합적인 실험을 수행했습니다. 본 연구의 방법은 약 37.5%의 계산 비용을 절감하는 동시에 다양한 RLVR 알고리즘 및 작업과 호환됩니다. 저희는 코드를 다음 주소에서 공개합니다: https://github.com/RUCAIBox/NExt.

Recently, scaling reinforcement learning with verifiable rewards (RLVR) for large language models (LLMs) has emerged as an effective training paradigm for significantly improving model capabilities, which requires guiding the model to perform extensive exploration and learning, leading to substantial computational overhead and becoming a key challenge. To reduce the number of training steps, Prior work performs linear extrapolation of model parameters. However, the dynamics of model parameter updates during RLVR training remain insufficiently understood. To further investigate the evolution of LLMs during RLVR training, we conduct empirical experiments and find that the rank-1 subspace of the model does not evolve linearly, and its dominance over the original parameters is further amplified during LoRA training. Based on the above insights, we propose the \textbf{N}onlinear \textbf{Ext}rapolation of low-rank trajectories (\textbf{NExt}), a novel framework that models and extrapolates low-rank parameter trajectories in a nonlinear manner. Concretely, we first train the model using LoRA and extract the rank-1 subspace of parameter differences at multiple training steps, which is then used for the subsequent nonlinear extrapolation. Afterward, we utilized the extracted rank-1 subspace to train a predictor, which can model the trajectory of parameter updates during RLVR, and then perform the predict-extend process to extrapolate model parameters, achieving the acceleration of RLVR. To further study and understand NExt, we conduct comprehensive experiments that demonstrate the effectiveness and robustness of the method. Our method reduces computational overhead by approximately 37.5\% while remaining compatible with a wide range of RLVR algorithms and tasks. We release our code in https://github.com/RUCAIBox/NExt.