다중 목적 강화학습 (Multi-Objective Reinforcement Learning)을 통한 LLM 사전 학습용 통합 데이터 스케줄러

소스의 다양성과 혼합 전략(mixing strategy)에 의해 결정되는 학습 데이터의 구성은 거대 언어 모델 (Large Language Model, LLM) 사전 학습 (pre-training)의 초석입니다. 학습 중에 데이터 혼합을 적응적으로 조정하는 기술인 온라인 데이터 혼합 (Online Data Mixing, ODM)은 효율성을 개선하기 위한 유망한 방향으로 부상했습니다. 그러나 기존 방법들은 단일한 최적화 관점에 의존한다는 제약이 있으며, 이는 복잡한 LLM 사전 학습이 여러 차원에서 동적인 데이터 구성을 고려해야 할 필요성을 근본적으로 간과합니다. 이러한 한계를 극복하기 위해, 우리는 새로운 온라인 데이터 혼합 프레임워크인 통합 데이터 스케줄러 (Holistic Data Scheduler, HDS)를 소개합니다. HDS는 데이터 스케줄링 문제를 연속 제어 공간 (continuous control space)에서의 강화학습 (reinforcement learning) 문제로 공식화하며, 고차원 정책 공간 (policy space)을 탐색하는 데 있어 안정성과 샘플 효율성 (sample efficiency)을 갖춘 Soft Actor-Critic (SAC) 알고리즘을 활용합니다. HDS의 핵심에는 세 가지 중요한 관점을 통합하는 새로운 다중 목적 통합 보상 함수 (multi-objective, holistic reward function)가 있습니다: 품질을 위한 데이터 기반 보상 (data-driven reward), 도메인 간 영향력을 포착하는 손실 기반 보상 (loss-driven reward), 그리고 가중치 규범 (weight norms)에 기반한 모델 기반 보상 (model-driven reward)입니다. 우리의 설계를 검증하고 최적의 구성을 결정하기 위해, 다양한 크기의 LLM을 대상으로 체계적인 실험을 수행했습니다. The Pile 벤치마크에서 HDS는 차순위 방법보다 44% 적은 학습 반복 (training iterations)만으로 다음 최선 방법의 최종 검증 퍼플렉시티 (validation perplexity)에 도달했습니다. 또한, MMLU 0-shot 태스크에서 7.2%의 향상을 달성함과 동시에 다른 벤치마크에서도 일관된 이득을 보여주었으며, 이는 학습 효율성과 최종 모델 성능을 모두 향상시키는 능력을 입증합니다.