Learning-Zone Energy: Online Data Selection for Efficient RL Post-Training

TL;DR

Learning-Zone Energy (LZE) is an online data selection method for reinforcement learning in large language models that efficiently focuses training on the most informative prompts, reducing compute while maintaining or improving performance.

Contribution

LZE introduces a theoretically grounded, fully online scoring framework that dynamically concentrates training on the model's active learning frontier, improving efficiency over uniform sampling methods.

Findings

Retains only 40% of training data per step while matching or surpassing full-data baselines.

Achieves +45.9% out-of-distribution gains on AIME25 and +18.2% on AMC23.

Reduces training FLOPs by an estimated 36%.

Abstract

Reinforcement Learning (RL) post-training has emerged as the dominant paradigm for eliciting mathematical reasoning in Large Language Models (LLMs), yet prevailing techniques such as GRPO and DAPO distribute rollout and gradient budgets nearly uniformly across prompts, squandering compute on samples that are already mastered or remain far beyond the model's current capability. To address this fundamental inefficiency, we propose Learning-Zone Energy (LZE), a theoretically grounded, fully online data selection framework that concentrates computation on the model's active learning frontier. At its core, we define a closed-form Learning-Zone Energy Score that fuses three complementary signals, an initial-difficulty anchor, a normalized outcome-uncertainty term, and a pass-rate momentum, into a single scalar that is provably aligned with the expected magnitude of group-relative policy gradient updates. A forward pruner with replay further reduces wall-clock time cost by skipping rollout generation for persistently solved prompts while periodically checking for forgetting. Evaluated on Qwen-family models (1.5B-8B) across GSM8K, MATH and DAPO-MATH, our method retains only 40% of the training data per step yet matches or surpasses full-data baselines, with especially pronounced out-of-distribution gains on AIME25 (+45.9%) and AMC23 (+18.2%), alongside an estimated 36% reduction in training FLOPs. Our code is available at https://github.com/Stellaris167/LZE.