DUMP:面向强化学习大模型后训练的自动化分布式课程学习
DUMP: Automated Distribution-Level Curriculum Learning for RL-based LLM Post-training
April 13, 2025
作者: Zhenting Wang, Guofeng Cui, Kun Wan, Wentian Zhao
cs.AI
摘要
近期,基于强化学习(RL)的后训练方法取得了显著进展,特别是在提升大语言模型(LLMs)处理复杂任务的推理能力方面。然而,现有方法大多将训练数据视为一个整体,忽视了现代LLM训练通常涉及来自不同分布的数据混合——这些数据在来源和难度上均存在差异。这种异质性引入了一个关键挑战:如何自适应地安排跨分布的训练以优化学习效率。本文提出了一种基于分布级可学习性概念的课程学习框架。我们的核心见解是,策略优势的大小反映了模型在特定分布上进一步训练所能获得的收益。基于此,我们为基于RL的LLM后训练设计了一个分布级课程学习框架,该框架利用上置信界(UCB)原理动态调整不同分布的采样概率。该方法优先考虑具有高平均优势(利用)或低样本计数(探索)的分布,从而产生一个自适应且理论依据充分的训练计划。我们以GRPO作为底层RL算法实例化了这一课程学习框架,并在多难度和多来源的逻辑推理数据集上验证了其有效性。实验结果表明,我们的框架显著提高了收敛速度和最终性能,凸显了分布感知课程策略在LLM后训练中的价值。代码:https://github.com/ZhentingWang/DUMP。
English
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.Summary
AI-Generated Summary