Learning at the Right Pace: Adaptive Data Scheduling Improves LLM Reinforcement Learning

summary

Video file (mp4)

The gist

The gist The Adaptive Data Scheduling (ADS) framework, a dual-level data scheduling system, improves Large Language Model reinforcement learning by replacing uniform sampling with an adaptive

In short

The Adaptive Data Scheduling (ADS) framework improves Large Language Model reinforcement learning by replacing uniform data sampling with a dual-level adaptive system. ADS first organizes data into semantic clusters and then adaptively samples between these clusters based on their success rates. It also selects specific 'policy-boundary' samples within those clusters to guide the policy toward its capability limits, leading to better accuracy across different reasoning tasks.

Key concepts

Semantic Clustering
This process organizes training data into groups based on shared patterns and meanings derived from the base policy's embeddings. It creates semantically coherent collections of prompts that share underlying conceptual structures, allowing the system to understand the relationships between different types of problems.
Inter-Cluster Distribution
ADS uses a formula to adapt how often samples are drawn from each semantic cluster during training. This distribution prioritizes clusters that have shown higher success rates so the policy focuses its learning efforts on areas where it is currently performing best, ensuring progress is made efficiently.
Policy-Boundary Samples
These are specific data points located near the current limits of what the model can achieve. By tracking samples near a 50% success rate within a cluster, ADS identifies informative examples that lie right at the edge of the policy's capability, helping it learn how to improve its performance in challenging areas.

Terminology used across episodes

This episode discusses

The paper

Learning at the Right Pace: Adaptive Data Scheduling Improves LLM Reinforcement Learning · Read on arXiv

Zicheng Xu, Ruixuan Zhang, Yu-Neng Chuang, Xiuyi Lou, Hoang Anh Duy Le, Oren Gal, Alexander S. Szalay, Ìhao Xu, ľ Guanchu Wang†

Johns Hopkins University · Rice University · University of Haifa · Workato · University of North Carolina at Charlotte

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom.

Jane: And I'm Jane, here with Lu, Meng and our in-house large language model, Lalam.

Tom: Today's paper: "Learning at the Right Pace".

Jane: The gist The Adaptive Data Scheduling (ADS) framework, a dual-level data scheduling system,

Tom: First, who's behind it and why it matters.

Title and authors: Tom: The title itself, "Learning at the Right Pace: Adaptive Data Scheduling Improves LLM Reinforcement Learning," tells you that the core idea is about timing and matching the data to what the model can actually handle at that moment.

Jane: It’s about moving away from a fixed schedule where you just feed everything in evenly, toward a system that adjusts based on two things: what kind of patterns are in the data, and how capable the AI is right now.

Lu: The authors are taking this idea of adapting to semantic structure and policy capability and building it into a practical framework called Adaptive Data Scheduling, which they describe as having two levels of scheduling.

Meng: So instead of one static way to pick samples from the whole dataset, they are suggesting a more dynamic system that considers both the meaning of the prompt and where the model is currently struggling or succeeding.

Lalam: The goal is to replace that uniform sampling with something smarter, organizing data into groups based on shared patterns and then adapting how often we pull from those groups based on their success rate.

The paper's summary: Tom: They start by saying existing methods like GRPO, which uses relative outcome comparisons within a group, rely on a naive, unstructured data sampling strategy where you just maximize the expected loss over prompts drawn uniformly from the entire dataset.

Jane: That uniform distribution is what they find is the problem because it doesn't account for where samples lie relative to the policy’s current capability boundary, which leads to advantage collapse.

Lu: They propose organizing training data into semantic clusters based on shared patterns and meanings first, and then maintaining an adaptive inter-cluster distribution by updating the probability of selecting each cluster according to its cluster-level success rate.

Meng: So if a certain group of prompts—say, those about physics—is succeeding well, the system should start pulling more from that physics group before shifting focus elsewhere.

Lalam: And they also have intra-cluster scheduling which focuses on selecting policy-boundary samples within those clusters by tracking empirical success rates near fifty percent.

The paper's improvements: Tom: One of the key improvements they highlight is that this dual-level structure allows for inter-cluster distribution that consolidates progress at the cluster level, meaning it helps solidify learning on specific semantic patterns before moving to harder concepts.

Jane: And at the sample level, they focus on policy-boundary samples, which are those near where the model's current capability is shifting, because those rollout groups are likely to contain both correct and incorrect responses.

Lu: They define a difficulty score d i for each sample based on the negative log-likelihood of its reference solution under the base policy, and then they use an online rollout feedback rule to update a mini-cluster M m using samples near fifty percent success rate.

Meng: That sounds like it’s trying to avoid getting stuck with prompts that are either too easy or way too hard, which is something I worry about when we deploy models at scale.

Lalam: They show this approach works across different objective designs, meaning it's RL objective-agnostic and can improve methods with different reinforcement learning objectives like GRPO, OPD, DAPO, and GSPO.

Conclusion: Tom: So to wrap up the Adaptive Data Scheduling paper, they show that both the inter-cluster distribution and intra-cluster scheduling are necessary; removing either one makes the performance worse than using the full ADS framework.

Jane: They also showed that this dynamic pacing produces what they call a "spiral reward increase," which reflects how it continuously adapts to where the policy's capability boundary is shifting during training.

Lu: It’s important to remember that this offline preprocessing stage, which computes those sample embeddings and difficulty scores, only happens once before training starts and doesn't add extra per-step overhead during the actual policy optimization.

Meng: That upfront cost for preprocessing has to be manageable if we want to use this in a production environment where we're running things continuously.

Lalam: Overall, Learning at the Right Pace: Adaptive Data Scheduling Improves LLM Reinforcement Learning replaces uniform sampling with an adaptive distribution over semantic clusters and policy-boundary sample selection to pace RL posttraining according to both semantic structure and policy capability.

Tom: It’s a solid framework for making sure we are training on the most useful data available at any given point in the process.

Jane: It really changes how we think about pacing model learning after initial training, moving beyond just feeding it more data.

More episodes

← Home