DUMP: Aprendizado de Currículo Automatizado em Nível de Distribuição para Pós-treinamento de LLMs Baseados em RL

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

April 13, 2025

Autores: Zhenting Wang, Guofeng Cui, Kun Wan, Wentian Zhao

cs.AI

Resumo

Os recentes avanços no pós-treinamento baseado em aprendizagem por reforço (RL) levaram a melhorias notáveis em modelos de linguagem de grande escala (LLMs), particularmente no aprimoramento de suas capacidades de raciocínio para lidar com tarefas complexas. No entanto, a maioria dos métodos existentes trata os dados de treinamento como um todo unificado, ignorando o fato de que o treinamento moderno de LLMs frequentemente envolve uma mistura de dados provenientes de diversas distribuições, variando tanto em origem quanto em dificuldade. Essa heterogeneidade introduz um desafio fundamental: como agendar de forma adaptativa o treinamento entre as distribuições para otimizar a eficiência do aprendizado. Neste artigo, apresentamos uma estrutura de aprendizado curricular fundamentada na noção de aprendibilidade em nível de distribuição. Nossa principal percepção é que a magnitude das vantagens da política reflete o quanto um modelo ainda pode se beneficiar de um treinamento adicional em uma determinada distribuição. Com base nisso, propomos uma estrutura de aprendizado curricular em nível de distribuição para o pós-treinamento de LLMs baseado em RL, que utiliza o princípio do Limite Superior de Confiança (UCB) para ajustar dinamicamente as probabilidades de amostragem para diferentes distribuições. Essa abordagem prioriza distribuições com alta vantagem média (exploração) ou baixa contagem de amostras (exploração), resultando em um cronograma de treinamento adaptativo e teoricamente fundamentado. Instanciamos nossa estrutura de aprendizado curricular com GRPO como o algoritmo de RL subjacente e demonstramos sua eficácia em conjuntos de dados de raciocínio lógico com múltiplas dificuldades e origens. Nossos experimentos mostram que nossa estrutura melhora significativamente a velocidade de convergência e o desempenho final, destacando o valor de estratégias curriculares conscientes da distribuição no pós-treinamento de LLMs. Código: 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.