DUMP: Geautomatiseerd Curriculum Leren op Distributieniveau voor RL-gebaseerde LLM Post-training

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

Samenvatting

Recente vooruitgang in reinforcement learning (RL)-gebaseerde na-training heeft geleid tot aanzienlijke verbeteringen in grote taalmmodellen (LLM's), met name in het versterken van hun redeneervaardigheden om complexe taken aan te kunnen. De meeste bestaande methoden behandelen de trainingsgegevens echter als een geheel, waarbij over het hoofd wordt gezien dat moderne LLM-training vaak een mix van gegevens uit diverse distributies omvat, die variëren in zowel bron als moeilijkheidsgraad. Deze heterogeniteit introduceert een belangrijke uitdaging: hoe het trainingsschema adaptief over distributies te plannen om de leer efficiëntie te optimaliseren. In dit artikel presenteren we een principieel curriculumlerenraamwerk gebaseerd op het concept van distributieniveau-leerbaarheid. Onze kerninzicht is dat de omvang van beleidsvoordelen weerspiegelt hoeveel een model nog kan profiteren van verdere training op een gegeven distributie. Op basis hiervan stellen we een distributieniveau-curriculumlerenraamwerk voor RL-gebaseerde LLM-na-training voor, dat het Upper Confidence Bound (UCB)-principe benut om de steekproefkansen voor verschillende distributies dynamisch aan te passen. Deze aanpak geeft prioriteit aan distributies met een hoog gemiddeld voordeel (exploitatie) of een laag steekproefaantal (exploratie), wat resulteert in een adaptief en theoretisch onderbouwd trainingsschema. We concretiseren ons curriculumlerenraamwerk met GRPO als het onderliggende RL-algoritme en demonstreren de effectiviteit ervan op logisch redeneerdatasets met meerdere moeilijkheidsgraden en bronnen. Onze experimenten tonen aan dat ons raamwerk de convergentiesnelheid en de uiteindelijke prestaties aanzienlijk verbetert, wat de waarde van distributiebewuste curriculumstrategieën in LLM-na-training benadrukt. 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.