DUMP: Apprendimento Automatico a Livello di Distribuzione per il Post-Addestramento di LLM Basati su RL

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

April 13, 2025

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

cs.AI

Abstract

I recenti progressi nel post-addestramento basato sull'apprendimento per rinforzo (RL) hanno portato a notevoli miglioramenti nei grandi modelli linguistici (LLM), in particolare nel potenziamento delle loro capacità di ragionamento per gestire compiti complessi. Tuttavia, la maggior parte dei metodi esistenti tratta i dati di addestramento come un insieme unificato, trascurando il fatto che l'addestramento moderno degli LLM spesso coinvolge una miscela di dati provenienti da distribuzioni diverse, che variano sia per origine che per difficoltà. Questa eterogeneità introduce una sfida cruciale: come pianificare in modo adattivo l'addestramento attraverso le distribuzioni per ottimizzare l'efficienza dell'apprendimento. In questo articolo, presentiamo un framework di apprendimento curriculare basato sul concetto di apprendibilità a livello di distribuzione. La nostra intuizione principale è che l'entità dei vantaggi della politica riflette quanto un modello possa ancora beneficiare di ulteriore addestramento su una determinata distribuzione. Sulla base di ciò, proponiamo un framework di apprendimento curriculare a livello di distribuzione per il post-addestramento RL degli LLM, che sfrutta il principio dell'Upper Confidence Bound (UCB) per regolare dinamicamente le probabilità di campionamento per diverse distribuzioni. Questo approccio privilegia le distribuzioni con un vantaggio medio elevato (sfruttamento) o un basso numero di campioni (esplorazione), producendo una pianificazione di addestramento adattiva e teoricamente fondata. Istanziamo il nostro framework di apprendimento curriculare utilizzando GRPO come algoritmo RL sottostante e ne dimostriamo l'efficacia su dataset di ragionamento logico con multiple difficoltà e fonti. I nostri esperimenti mostrano che il nostro framework migliora significativamente la velocità di convergenza e le prestazioni finali, evidenziando il valore delle strategie curriculari consapevoli della distribuzione nel post-addestramento degli LLM. Codice: 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.