DUMP: RL 기반 LLM 사후 학습을 위한 자동화된 분산 수준 커리큘럼 학습

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 훈련이 다양한 분포(출처와 난이도가 모두 다름)의 데이터 혼합을 포함한다는 사실을 간과하고 있습니다. 이러한 이질성은 학습 효율성을 최적화하기 위해 분포 간 훈련을 어떻게 적응적으로 스케줄링할 것인가라는 핵심적인 도전 과제를 제기합니다. 본 논문에서는 분포 수준 학습 가능성 개념에 기반한 원칙적인 커리큘럼 학습 프레임워크를 제시합니다. 우리의 핵심 통찰은 정책 이점의 크기가 주어진 분포에서 모델이 추가 훈련을 통해 얼마나 더 이익을 얻을 수 있는지를 반영한다는 것입니다. 이를 바탕으로, RL 기반 LLM 사후 훈련을 위한 분포 수준 커리큘럼 학습 프레임워크를 제안하며, 이는 상한 신뢰 구간(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.