DUMP: Автоматизированное обучение с распределением сложности на уровне учебного плана для RL-обучения крупных языковых моделей после тренировки

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

April 13, 2025

Авторы: Zhenting Wang, Guofeng Cui, Kun Wan, Wentian Zhao

cs.AI

Аннотация

Недавние достижения в области посттренинга на основе обучения с подкреплением (RL) привели к значительным улучшениям в работе крупных языковых моделей (LLM), особенно в усилении их способностей к рассуждению для выполнения сложных задач. Однако большинство существующих методов рассматривают обучающие данные как единое целое, игнорируя тот факт, что современное обучение LLM часто включает данные из различных распределений, различающихся как по источнику, так и по сложности. Эта неоднородность создает ключевую проблему: как адаптивно планировать обучение по различным распределениям для оптимизации эффективности обучения. В данной статье мы представляем принципиальную структуру обучения по учебному плану, основанную на понятии обучаемости на уровне распределения. Наше ключевое наблюдение заключается в том, что величина преимуществ политики отражает, насколько модель может еще выиграть от дальнейшего обучения на данном распределении. На основе этого мы предлагаем структуру обучения по учебному плану на уровне распределения для посттренинга LLM на основе RL, которая использует принцип Верхней Доверительной Границы (UCB) для динамической корректировки вероятностей выборки для различных распределений. Этот подход отдает приоритет распределениям с либо высоким средним преимуществом (эксплуатация), либо низким количеством выборок (исследование), что приводит к адаптивному и теоретически обоснованному графику обучения. Мы реализуем нашу структуру обучения по учебному плану с использованием GRPO в качестве базового алгоритма RL и демонстрируем ее эффективность на наборах данных для логического рассуждения с различными уровнями сложности и источниками. Наши эксперименты показывают, что наша структура значительно улучшает скорость сходимости и итоговую производительность, подчеркивая ценность стратегий обучения по учебному плану, учитывающих распределение, в посттренинге LLM. Код: 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.