DUMP: Aprendizaje Automatizado de Plan de Estudios a Nivel de Distribución para Entrenamiento Posterior de LLM Basados en 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

Resumen

Los recientes avances en el post-entrenamiento basado en aprendizaje por refuerzo (RL) han llevado a mejoras notables en los modelos de lenguaje grandes (LLMs), particularmente en la potenciación de sus capacidades de razonamiento para manejar tareas complejas. Sin embargo, la mayoría de los métodos existentes tratan los datos de entrenamiento como un todo unificado, pasando por alto el hecho de que el entrenamiento moderno de LLMs a menudo implica una mezcla de datos provenientes de diversas distribuciones, que varían tanto en su origen como en su dificultad. Esta heterogeneidad introduce un desafío clave: cómo programar de manera adaptativa el entrenamiento a través de las distribuciones para optimizar la eficiencia del aprendizaje. En este artículo, presentamos un marco de aprendizaje curricular fundamentado en la noción de capacidad de aprendizaje a nivel de distribución. Nuestra idea central es que la magnitud de las ventajas de la política refleja cuánto puede beneficiarse aún un modelo de un entrenamiento adicional en una distribución dada. Basándonos en esto, proponemos un marco de aprendizaje curricular a nivel de distribución para el post-entrenamiento de LLMs basado en RL, que aprovecha el principio del Límite Superior de Confianza (UCB) para ajustar dinámicamente las probabilidades de muestreo para diferentes distribuciones. Este enfoque prioriza distribuciones con una ventaja promedio alta (explotación) o un bajo recuento de muestras (exploración), resultando en un programa de entrenamiento adaptativo y teóricamente fundamentado. Instanciamos nuestro marco de aprendizaje curricular utilizando GRPO como el algoritmo de RL subyacente y demostramos su efectividad en conjuntos de datos de razonamiento lógico con múltiples niveles de dificultad y fuentes. Nuestros experimentos muestran que nuestro marco mejora significativamente la velocidad de convergencia y el rendimiento final, destacando el valor de las estrategias curriculares conscientes de la distribución en el post-entrenamiento 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.