UFO-RL: Uncertainty-Focused Optimization for Efficient Reinforcement Learning Data Selection

Ting Liu (Harbin Institute of Technology) · Yang Zhao (Harbin Institute of Technology) · Li Du (BAAI) · Bin Liu (dxm) · Wenbin Zhang (DXM) · Bing Qin (Harbin Institute of Technology) · Kai Xiong (Harbin Institute of Technology) · Xiao Ding (Harbin Institute of Technology) · Yangou Ouyang (Harbin Institute of Technology) · Zhouhao Sun (Harbin Institute of Technology) · Jiannan Guan (Harbin Institute of Technology) · Dong Hu (Du Xiaoman Technology (Beijing) Co., Ltd.)

computational efficiencydata difficulty assessmentdata evaluation accelerationfine-tuning strategy generalization capabilitiesmulti-samplingpolicy optimization reinforcement learningsingle-pass computationtargeted data selectiontraining data selection training stability uncertainty estimationzone of proximal development

A primary impediment to scaling reinforcement learning (RL) for large language model (LLM) training is the substantial computational cost, predominantly arising from the necessity of multi-sampling for policy optimization and evaluation. This underscores the critical yet challenging nature of efficient training data selection. Drawing inspiration from the Zone of Proximal Development (ZPD) theory, which posits that learners acquire knowledge more effectively from tasks of intermediate difficulty, we hypothesize that LLMs exhibit optimal learning from data they have not yet mastered but demonstrate the potential to comprehend. Conventional methodologies for assessing data difficulty or informativeness typically rely on computationally intensive multi-sampling or iterative procedures. To address this limitation, we introduce UFO-RL (**U**ncertainty-**F**ocused **O**ptimization for **R**einforcement **L**earning), a novel framework that employs a computationally efficient single-pass uncertainty estimation technique to identify informative training instances. This method, requiring only a single forward pass and obviating the need for iterative next-token computation, achieves a significant acceleration (up to 185$\times$) in data evaluation compared to multi-sampling approaches. UFO-RL leverages this efficient metric to select data within the model's estimated ZPD for training. Extensive experimentation across diverse LLMs and mathematical benchmarks demonstrates that training with a mere 10\% of the data, carefully selected by UFO-RL, yields performance comparable to or even surpassing that of full-data training. Furthermore, this targeted data selection results in up to a 16$\times$ reduction in overall training time, concurrently enhancing training stability and improving generalization capabilities. Thus, UFO-RL presents a practical and highly efficient strategy for scaling RL fine-tuning of LLMs by focusing learning efforts on the most informative and valuable data, thereby mitigating the computational bottlenecks associated with traditional RL training.