DUMP : Apprentissage curriculaire automatisé au niveau de la distribution pour le post-entraînement des LLM basés sur l'apprentissage par renforcement

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

April 13, 2025

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

cs.AI

Résumé

Les récents progrès dans l'apprentissage par renforcement (RL) appliqué au post-entraînement ont conduit à des améliorations notables des grands modèles de langage (LLM), en particulier dans le renforcement de leurs capacités de raisonnement pour traiter des tâches complexes. Cependant, la plupart des méthodes existantes considèrent les données d'entraînement comme un ensemble unifié, négligeant le fait que l'entraînement moderne des LLM implique souvent un mélange de données provenant de distributions diverses, variant à la fois en source et en difficulté. Cette hétérogénéité introduit un défi majeur : comment planifier de manière adaptative l'entraînement à travers ces distributions pour optimiser l'efficacité de l'apprentissage. Dans cet article, nous présentons un cadre d'apprentissage curriculaire fondé sur la notion d'apprenabilité au niveau des distributions. Notre idée centrale est que l'amplitude des avantages de la politique reflète à quel point un modèle peut encore bénéficier d'un entraînement supplémentaire sur une distribution donnée. Sur cette base, nous proposons un cadre d'apprentissage curriculaire au niveau des distributions pour le post-entraînement des LLM basé sur le RL, qui exploite le principe de la borne supérieure de confiance (UCB) pour ajuster dynamiquement les probabilités d'échantillonnage pour différentes distributions. Cette approche privilégie les distributions présentant soit un avantage moyen élevé (exploitation) soit un faible nombre d'échantillons (exploration), aboutissant à un plan d'entraînement adaptatif et théoriquement fondé. Nous instancions notre cadre d'apprentissage curriculaire avec GRPO comme algorithme de RL sous-jacent et démontrons son efficacité sur des ensembles de données de raisonnement logique avec plusieurs niveaux de difficulté et sources. Nos expériences montrent que notre cadre améliore significativement la vitesse de convergence et les performances finales, mettant en lumière la valeur des stratégies curriculaires conscientes des distributions dans le post-entraînement des LLM. Code : 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.