PaperlayerView exact revision on arXiv

Source-linked AI summary

DUMP: Automated Distribution-Level Curriculum Learning for RL-based LLM Post-training

Zhenting Wang, Guofeng Cui, Yu-Jhe Li, Kun Wan, Wentian Zhao

arXiv:2504.09710v2cs.LGcs.CL

TL;DR

RL post-training commonly combines data from diverse distributions, while many methods do not adaptively schedule training across them. DUMP uses advantage-based UCB sampling to allocate training across distributions, and experiments report faster convergence and stronger final performance than uniform sampling on logic-reasoning tasks.

  • Problem

    Many RL post-training pipelines treat heterogeneous data distributions uniformly or use static curricula, leaving adaptive distribution-level scheduling underexplored.

    Source: §1 · pp. 1–2

  • Method

    DUMP uses expected absolute policy advantage as a distribution-level learnability signal and UCB scores to adaptively adjust sampling probabilities.

    Source: §1 · p. 2

  • Results

    DUMP consistently improves convergence and final performance over baselines in experiments on logic-reasoning distributions.

    Source: §5 · p. 9

  • Takeaways & Limitations

    The findings highlight the value of distribution-aware curriculum learning for RL-based LLM post-training.

    Source: §5 · p. 9

    The experiments cover only LLMs at 7B scale and do not evaluate multimodal LLMs due to computational constraints.

    Source: D Limitations · p. 15

Abstract

from arXiv · show

Recent advances in reinforcement learning (RL)-based post-training have led to notable improvements in large language models (LLMs), particularly in enhancing their reasoning capabilities to handle complex tasks. However, most existing methods treat the training data as a unified whole, overlooking the fact that modern LLM training often involves a mixture of data from diverse distributions-varying in both source and difficulty. This heterogeneity introduces a key challenge: how to adaptively schedule training across distributions to optimize learning efficiency. In this paper, we present a principled curriculum learning framework grounded in the notion of distribution-level learnability. Our core insight is that the magnitude of policy advantages reflects how much a model can still benefit from further training on a given distribution. Based on this, we propose a distribution-level curriculum learning framework for RL-based LLM post-training, which leverages the Upper Confidence Bound (UCB) principle to dynamically adjust sampling probabilities for different distrubutions. This approach prioritizes distributions with either high average advantage (exploitation) or low sample count (exploration), yielding an adaptive and theoretically grounded training schedule. We instantiate our curriculum learning framework with GRPO as the underlying RL algorithm and demonstrate its effectiveness on logic reasoning datasets with multiple difficulties and sources. Our experiments show that our framework significantly improves convergence speed and final performance, highlighting the value of distribution-aware curriculum strategies in LLM post-training. Code: https://github.com/ZhentingWang/DUMP.

1 Introduction

RL post-training draws on heterogeneous data distributions, but existing pipelines often sample them uniformly or use static curricula. DUMP schedules distributions dynamically using advantage-based learnability estimates and UCB scores.

1 Introduction

  • Heterogeneous post-training data varies in source, difficulty, counts, formats, and objectives, while many pipelines still sample distributions uniformly or use static curricula.The introduction cites Tülu 3 [7] and Seed-Thinking v1.5 as examples of diverse or controllably generated data.Source: §1 · pp. 1–2
  • Automated distribution-level scheduling remains underexplored relative to instance-level curricula and heuristic multi-stage schedules [10] [11].The challenge is identifying signals of each distribution’s current learnability and using them stably and efficiently.Source: §1 · p. 2
  • DUMP uses policy-advantage magnitude as a learnability signal and UCB scores to adaptively guide training across data distributions.The framework is instantiated with GRPO and is described as compatible with advantage-based RL algorithms.Source: §1 · p. 2

2 Background

RL-based post-training optimizes model behavior using reward signals, and GRPO uses group-relative feedback to update the policy. Existing LLM curricula often focus on sample difficulty or fixed training stages rather than adaptive distribution-level scheduling.

2 Background

  • RL plays a central role in LLM post-training; RL-post-trained OpenAI o1, o3, and DeepSeek-R1 significantly outperform LLMs without RL post-training on challenging mathematics and coding benchmarks, including AIME and Codeforces [18] [4].PPO [23] is widely used because its clipped objective stabilizes training by preventing large policy updates, and it remains a strong baseline in many LLM alignment settings.Source: §2 · p. 2
  • GRPO [3] uses group-based feedback: it samples candidate outputs for each prompt and normalizes each output’s reward relative to the group.Its clipped surrogate objective uses normalized advantages and a KL penalty to keep the updated policy close to the reference policy.Source: §2 · p. 3
  • Prior LLM curricula include instance-level preference-pair ranking and fixed heuristic stages [10] [11], while automated distribution-level scheduling remains underexplored.The cited fixed-stage strategies use manually specified difficulty labels and do not adapt to learning progress.Source: §2 · p. 3

3 Method

DUMP estimates distribution-level learnability from expected absolute advantages and uses a bandit-inspired UCB curriculum to adjust sampling during RL post-training. Its algorithm tracks recent advantage statistics and sample counts, updates distribution weights, and trains on the resulting batches.

3.1 Measuring Learnability via Absolute Advantage

  • Expected absolute advantage is proposed as a proxy for how much a data distribution can help the model improve.Averaging absolute advantages captures signal in both directions and avoids cancellation between positive and negative advantages; small magnitudes may also reflect tasks that are too difficult or noisy to learn from effectively.Source: §3.1 · pp. 3–4

3.2 Formalizing Distribution-Level Curriculum Learning as Multi-armed Bandit

  • The curriculum is formalized as a multi-armed bandit in which each distribution’s learnability is its reward, and UCB balances observed advantage against exploration of under-sampled distributions.DUMP softmaxes the scores to produce sampling probabilities rather than selecting only one distribution, preserving partial exploration.Source: §3.2 · p. 4

3.3 Algorithm

  • DUMP initializes per-distribution advantage windows, sample counts, and uniform sampling weights, then updates the statistics from rollout batches.The algorithm retains recent observations and computes UCB scores from mean absolute advantage and an exploration bonus tied to sample count.Source: §3.2 · p. 5
  • At each training step, DUMP converts UCB scores to sampling probabilities with a temperature-controlled softmax and updates the model using an RL algorithm such as GRPO.Lower temperature makes sampling sharper; higher temperature produces a smoother, more exploratory curriculum.Source: §3.2 · p. 5, §3.3 · p. 5

4 Experiments and Results

Across K&K, diverse logic, and math settings, DUMP improves test rewards over uniform sampling, while experiments also show exploration and online advantage signals shaping its curriculum. The evaluation uses GRPO across multiple reasoning datasets and difficulty-defined distributions.

4.1 Experiments Setup

  • Experiments use GRPO with Qwen2.5 models across K&K, logic-reasoning, and math settings, including distributions partitioned by puzzle characters, reasoning steps, or AIME years.The settings include 12 K&K distributions, 15 logic distributions, and five math distributions.Source: §4.1 · p. 6

4.2 Effectiveness of DUMP

  • Average test answer reward rises from 0.90 to 1.17 on 15 diverse logic distributions with DUMP.The reported gains include AR-LSAT (-0.70 to -0.52) and K&K 7 Characters (0.56 to 1.02).Source: Table 1 · p. 7, §4.2 · p. 7
  • Average test answer reward rises from -0.59 to -0.43 on diverse math distributions, with the largest reported gain on AIME 1994–2004 (-1.50 to -1.02).The evaluation uses GSM-8K and AIME subsets grouped by competition years.Source: Table 2 · p. 7, §4.2 · p. 7
  • On the 9-character K&K distribution, DUMP exceeds 0.5 test reward while uniform sampling remains below 0.0.Across the 12 K&K distributions, DUMP is reported to converge faster and outperform the uniform baseline, especially on 6–12 characters.Source: §4.2 · p. 7, Figure 1 · p. 8

4.3 Ablation Study on the Sampling Strategy

  • With temperature 0.1, soft sampling outperforms greedy UCB selection on the 13- and 14-character K&K tasks, where greedy rewards are -0.91 and -1.38 versus -0.66 and -1.16.Greedy selection tends to lock onto one distribution early; the reported ablation attributes its underperformance to lack of exploration.Source: §4.3 · pp. 7–8

4.4 Analyzing the Automated Curriculum by DUMP

  • DUMP samples simpler 3–5-character puzzles early, then increasingly prioritizes harder 10–14-character puzzles as training progresses, without a manually specified order.The progression emerges from empirical advantage signals and adapts to online training dynamics.Source: §4.4 · pp. 8–9, Figure 2 · p. 9

5 Conclusion

DUMP uses expected absolute advantage as a learnability signal and UCB sampling to balance exploitation and exploration across distributions. The paper reports improved convergence and final performance over baselines.

  • DUMP treats expected absolute advantage as a learnability signal and uses UCB-based sampling to balance exploitation and exploration across heterogeneous distributions.The framework formalizes scheduling as a multi-armed bandit.Source: §5 · p. 9
  • Experiments report that DUMP consistently improves convergence and final performance over baselines.Source: §5 · p. 9

A Proof for Theorem 3.1

The proof links expected absolute advantage to the learning signal a distribution contributes, making it a proxy for that distribution’s potential to improve the model. This link relies on assumptions about policy-gradient norms across distributions.

A Proof for Theorem 3.1

  • Each sampled output’s advantage is estimated as its reward minus a prompt-dependent baseline, such as the group’s mean reward.This quantity is used in the policy-gradient expression for methods such as PPO or GRPO.Source: A Proof for Theorem 3.1 · p. 13
  • The proof assumes policy-gradient norms are bounded and vary slowly across distributions, so expected absolute advantage is the dominant term in gradient magnitude.The expected gradient norm is bounded below in terms of this signal.Source: A Proof for Theorem 3.1 · p. 13
  • Expected absolute advantage approximates a distribution’s learning-signal magnitude and, therefore, how much training on it can improve model parameters.The proof identifies it as a suitable curriculum signal.Source: A Proof for Theorem 3.1 · p. 13

B Theoretical Justification for UCB-Based Distribution Scheduling

The analysis treats distributions as bandit arms whose rewards are expected absolute advantages, and uses UCB to balance estimated learnability against exploration. Under stated reward and stationarity assumptions, the regret grows logarithmically with training steps and linearly with inverse suboptimality gaps.

B Theoretical Justification for UCB-Based Distribution Scheduling

  • The bandit formulation treats each data distribution as an arm, with mean absolute advantage as its reward and observed batch advantages as empirical rewards.UCB selects the distribution with the highest upper confidence bound, combining empirical mean reward with an exploration term based on sampling count.Source: B Theoretical Justification for… · pp. 13–14
  • The guarantee assumes per-distribution absolute advantages are i.i.d. and bounded, while expected rewards remain approximately stationary over a local training window.In practice, advantages can be clipped or normalized to satisfy the boundedness condition.Source: B Theoretical Justification for… · pp. 13–14
  • Regret scales logarithmically in T and linearly in 1/∆j under the stated UCB analysis.The bound concerns the gap between the learner’s cumulative learnability gain and that of the best fixed distribution.Source: B Theoretical Justification for… · p. 14

C Comparison to Heuristic Curriculum

The paper contrasts DUMP’s adaptive scheduling with heuristic curricula that use fixed schedules over distributions. It argues that heuristics require prior knowledge and cannot respond to the model’s changing training needs.

C Comparison to Heuristic Curriculum

  • Prior heuristic curricula prescribe fixed distribution schedules and require prior knowledge of relative difficulty and learnability [11] [10].The paper notes that such knowledge may be unavailable or misleading across diverse domains.Source: C Comparison to Heuristic Curriculum · p. 14
  • Unlike static heuristics, DUMP adjusts sampling priorities during training using policy advantages as a measure of model performance.This lets it focus on distributions considered most beneficial at each training stage.Source: C Comparison to Heuristic Curriculum · p. 15
  • Because standardized heuristics are lacking for the task suite, the paper benchmarks DUMP against uniform sampling and adaptive baselines.Source: C Comparison to Heuristic Curriculum · p. 15

D Limitations

The experiments evaluate DUMP only with LLMs, not multimodal LLMs, because of computational constraints, and use models at the 7B scale.

D Limitations

Loading 2504.09710v2…