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

TL;DR

UFO-RL introduces an efficient uncertainty-based data selection method for reinforcement learning with LLMs, significantly reducing training time while maintaining or improving performance by focusing on data within the model's potential comprehension zone.

Contribution

The paper presents UFO-RL, a novel framework that uses single-pass uncertainty estimation for data selection, enabling faster and more effective RL training of LLMs.

Findings

Up to 185x faster data evaluation.

Training with 10% of selected data matches or exceeds full-data performance.

Reduces training time by up to 16x while improving stability.

Abstract

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.