UFO-RL: Uncertainty-Focused Optimization for Efficient Reinforcement Learning Data Selection
This addresses the problem of high computational costs in RL fine-tuning for LLMs, offering a practical and efficient data selection strategy that is incremental but impactful for scaling applications.
The paper tackles the computational expense of scaling reinforcement learning for large language models by introducing UFO-RL, a framework that uses efficient single-pass uncertainty estimation to select informative data, achieving up to 185x faster data evaluation and reducing training time by up to 16x while maintaining performance with only 10% of the data.
Scaling RL for LLMs is computationally expensive, largely due to multi-sampling for policy optimization and evaluation, making efficient data selection crucial. Inspired by the Zone of Proximal Development (ZPD) theory, we hypothesize LLMs learn best from data within their potential comprehension zone. Addressing the limitation of conventional, computationally intensive multi-sampling methods for data assessment, we introduce UFO-RL. This novel framework uses a computationally efficient single-pass uncertainty estimation to identify informative data instances, achieving up to 185x faster data evaluation. UFO-RL leverages this metric to select data within the estimated ZPD for training. Experiments show that training with just 10% of data selected by UFO-RL yields performance comparable to or surpassing full-data training, reducing overall training time by up to 16x while enhancing stability and generalization. UFO-RL offers a practical and highly efficient strategy for scaling RL fine-tuning of LLMs by focusing learning on valuable data.