Logo
FrontierNews.ai

How Peking University's New Framework Is Fixing the Data Problem in AI Reasoning Models

Peking University researchers have released DataFlex-RL, an open-source framework that addresses a fundamental bottleneck in training advanced reasoning models: figuring out which data should be fed to AI systems during reinforcement learning rounds. The tool, developed by the university's DCAI team in collaboration with UCAS and other institutions, provides a modular system for data selection, sample weighting, and domain mixing. After its release, the accompanying paper ranked as the number two paper of the day on Hugging Face, following 591 experiments across multiple models and tasks.

Why Does Data Strategy Matter for Reasoning Models?

As reasoning models and reinforcement learning verification (RLVR) become more central to how large language models are trained, the question of what data to use at each training step has grown increasingly complex. A question that seems impossible for a model to answer initially may become consistently solvable after targeted training. Researchers must constantly decide which data to keep, how much learning weight to assign to different samples, and how to adjust the balance between different domains like mathematics, logic, and science.

The challenge is compounded by the fact that there are often multiple competing data strategies, each with its own implementation details. Changing selection criteria might require rewriting parts of the training code. When comparing performance across different approaches, inconsistencies in reward functions, model versions, and training budgets make it difficult to determine whether improvements come from the data strategy itself or from other configuration changes.

How Does DataFlex-RL Solve This Problem?

DataFlex-RL integrates with existing training pipelines as a plugin, allowing researchers to test different data strategies without repeatedly modifying their training code. The framework reuses signals already generated during reinforcement learning, such as rewards, advantages, and token probabilities, converting them into actionable decisions about which data to include in training.

The framework operates across three distinct dimensions of data strategy:

  • Data Selection: Determines which responses participate in the current training update by filtering based on solve rates within groups, advantage rankings, reward variance, and answer efficiency. For example, the "difffilter" strategy retains only groups where two or three answers are correct, concentrating updates on medium-difficulty questions.
  • Sample Reweighting: Adjusts how much learning weight different samples contribute to model updates, either at the response level or token level. Methods like "softmax" adjust contributions based on advantage magnitude, while others weight samples whose rewards fall within specific ranges.
  • Domain Mixing: Dynamically adjusts the proportion of data from different domains based on historical training feedback. Since learning progress in mathematics, logic, and science tasks may vary, a fixed ratio may not work at every training stage.

By separating scoring from execution, DataFlex-RL allows strategy components to be independently replaced and tested. Different strategies share the same model generation, reward verification, and model update processes, with changes focused solely on the data strategy integration point.

Steps to Implement Dynamic Data Strategies in Your Training Pipeline

  • Integrate as a Plugin: Connect DataFlex-RL to your existing training framework through the VERL plugin system, which allows you to reuse existing training capabilities without rewriting core components.
  • Configure Selection Criteria: Choose which filtering method matches your training goals, such as filtering by solve rate, advantage ranking, or reward variance, depending on whether you want to focus on medium-difficulty or high-variance examples.
  • Set Reweighting Parameters: Decide whether to apply response-level or token-level weighting, and select the normalization approach that minimizes changes to your overall loss scale during updates.
  • Monitor Domain Ratios: Use a sliding window to accumulate training feedback across domains and update target proportions at configured intervals, ensuring both target ratios and actual data distribution are tracked simultaneously.

The framework's modular design means developers can focus their efforts on strategy components while leveraging existing training capabilities. When adding new methods, teams can concentrate on the strategy logic rather than rebuilding the entire training pipeline.

What Makes This Approach Different From Previous Methods?

DataFlex-RL's most direct value is reducing the time researchers spend tweaking training configurations and increasing the time spent verifying whether data strategies actually work. By providing a unified technical architecture, the framework enables fair comparison of different approaches under identical conditions. The 591 experiments conducted across multiple models and tasks in mathematics, logic, and science demonstrate the breadth of evaluation the researchers performed to validate their approach.

The tool is available on GitHub, making it accessible to the broader AI research community. This open-source release reflects a growing trend of researchers sharing not just model weights, but the training infrastructure and methodologies that make advanced models possible. As reinforcement learning becomes more sophisticated and compute-intensive, tools that help researchers make smarter decisions about data allocation could significantly reduce wasted training effort and accelerate progress in reasoning model development.

" }