다중 목적 강화 학습을 통한 LLM 사전학습을 위한 총체적 데이터 스케줄러

초록

훈련 데이터의 구성은 데이터 소스의 다양성과 혼합 전략에 의해 결정되며, 이는 대규모 언어 모델(LLM) 사전 훈련의 초석이다. 훈련 중 데이터 혼합을 적응적으로 조정하는 기법인 온라인 데이터 혼합(ODM)은 효율성을 개선하기 위한 유망한 방향으로 부상했다. 그러나 기존 방법은 단일 최적화 관점에 의존한다는 한계를 지니며, 이는 복잡한 LLM 사전 훈련이 다차원에서 동적 데이터 구성을 고려해야 할 필요성을 근본적으로 간과한다. 이러한 한계를 극복하기 위해, 우리는 새로운 온라인 데이터 혼합 프레임워크인 HDS(Holistic Data Scheduler)를 소개한다. HDS는 데이터 스케줄링 문제를 연속 제어 공간에서의 강화 학습 문제로 정식화하고, 고차원 정책 공간 탐색에서의 안정성과 샘플 효율성을 위해 SAC(Soft Actor-Critic) 알고리즘을 활용한다. HDS의 핵심에는 데이터 기반 품질 보상, 도메인 간 영향을 포착하는 손실 기반 보상, 가중치 노름에 기반한 모델 기반 보상이라는 세 가지 중요한 관점을 통합하는 새로운 다중 목표의 전체적 보상 함수가 자리한다. 설계의 타당성을 검증하고 최적 구성을 결정하기 위해, 우리는 다양한 크기의 LLM에 대해 체계적인 실험을 수행했다. The Pile 벤치마크에서 HDS는 다음으로 우수한 방법의 최종 검증 퍼플렉서티(perplexity)에 도달하면서 훈련 반복 횟수를 44% 줄였다. 또한 MMLU 제로샷 태스크에서 7.2%의 성능 향상을 달성하고 다른 벤치마크에서도 일관된 개선을 보여, 훈련 효율성과 최종 모델 성능을 모두 향상시킬 수 있는 능력을 입증한다.

English

The composition of training data, governed by the diversity of sources and their mixing strategy, is a cornerstone of Large Language Model (LLM) pre-training. Online Data Mixing (ODM), the technique of adaptively adjusting data mixtures during training, has emerged as a promising direction to improve efficiency. However, existing methods are constrained by their reliance on a singular optimization perspective, which fundamentally overlooks the need for complex LLM pre-training to consider the dynamic data composition from multiple dimensions. To overcome this limitation, we introduce the Holistic Data Scheduler (HDS), a novel online data mixing framework. HDS formulates the data scheduling challenge as a reinforcement learning problem in a continuous control space and leverages the Soft Actor-Critic (SAC) algorithm for its stability and sample efficiency in exploring the high-dimensional policy space. At the core of HDS lies a novel multi-objective, holistic reward function that integrates three critical perspectives: a data-driven reward for quality, a loss-driven reward capturing inter-domain influence, and a model-driven reward based on weight norms. To validate our design and determine its optimal configuration, we conducted systematic experiments on LLMs of various sizes. On The Pile benchmark, HDS reaches the final validation perplexity of the next best method with 44% fewer training iterations. Furthermore, it achieves a 7.2% improvement on the MMLU 0-shot task along with consistent gains on other benchmarks, showcasing its ability to enhance both training efficiency and final model capability.