Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Yang Zhao, Kai Xiong, Xiao Ding, Li Du, YangouOuyang, Zhouhao Sun, Jiannan Guan, Wenbin Zhang, Bin Liu, Dong Hu, Bing Qin, Ting Liu

TL;DR

The paper tackles the high computational burden of reinforcement learning fine-tuning for LLMs by introducing UFO-RL, which uses a single-pass uncertainty measure to select data within the model's Zone of Proximal Development (ZPD). By ranking data via a continuous confidence score, UFO-RL identifies fuzzy, intermediate-difficulty samples and trains on only 10% of data, achieving performance comparable to or better than full-data training while reducing overall cost by up to 16×. The method demonstrates robustness and improved generalization across multiple mathematical benchmarks and model scales, addressing both efficiency and learning stability. This approach offers a scalable path to efficient RL fine-tuning with practical impact for large-scale language-model customization and reasoning tasks.

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.

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

TL;DR

The paper tackles the high computational burden of reinforcement learning fine-tuning for LLMs by introducing UFO-RL, which uses a single-pass uncertainty measure to select data within the model's Zone of Proximal Development (ZPD). By ranking data via a continuous confidence score, UFO-RL identifies fuzzy, intermediate-difficulty samples and trains on only 10% of data, achieving performance comparable to or better than full-data training while reducing overall cost by up to 16×. The method demonstrates robustness and improved generalization across multiple mathematical benchmarks and model scales, addressing both efficiency and learning stability. This approach offers a scalable path to efficient RL fine-tuning with practical impact for large-scale language-model customization and reasoning tasks.

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.
Paper Structure (23 sections, 3 equations, 1 figure, 6 tables)

This paper contains 23 sections, 3 equations, 1 figure, 6 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. RL-based Post-training for LLMs
  4. Data-Driven "Less is More"
  5. Fuzzy data is beneficial for the learning process in RL.
  6. Definition of Sampling Accuracy
  7. Experimental Setup
  8. Results and Analysis
  9. Uncertainty-Focused Optimization for Efficient RL Data Selection
  10. Confidence Estimation via Average Log-Softmax
  11. Comparison with Confidence and Sampling Accuracy
  12. Confidence-Based Data Filtering
  13. Experiments and Analyses
  14. Training Setup
  15. Main Results: Performance Evaluation
  16. ...and 8 more sections

Figures (1)

  • Figure 1: Impact of Sampling Accuracy-Based Data Subset Training on Model Learning Efficiency. Evaluation across Qwen2.5 (0.5B-7B), Llama3.1-8B, and Mistral 7B. GSM8K training data was divided into 10 bins by accuracy. Lines represent RL performance from exclusive training on each bin (lowest to highest accuracy), showing the relationship between inferred difficulty and learning efficiency.