[논문리뷰] Holistic Data Scheduler for LLM Pre-training via Multi-Objective Reinforcement Learning

링크: 논문 PDF로 바로 열기

메타데이터

저자: Chenhao Dang, Jing Ma, Mingjie Liao, et al.

---

1. Key Terms & Definitions (핵심 용어 및 정의)

• HDS (Holistic Data Scheduler): LLM pre-training 과정에서 데이터를 동적으로 혼합하기 위해 제안된 강화학습 기반의 프레임워크입니다.
• ODM (Online Data Mixing): 학습 도중 모델의 상태에 따라 데이터 도메인별 샘플링 가중치를 실시간으로 조정하는 전략입니다.
• SAC (Soft Actor-Critic): HDS에서 데이터 스케줄링 정책을 최적화하기 위해 사용된 off-policy 기반의 강화학습 알고리즘으로, 높은 안정성과 sample efficiency를 제공합니다.
• Multi-Objective Reward: 데이터 품질(Data-driven), 도메인 간 영향력(Gradient-driven), 모델 안정성(Model-driven) 등 세 가지 관점을 통합한 HDS의 핵심 보상 함수입니다.

2. Motivation & Problem Statement (연구 배경 및 문제 정의)

본 논문은 LLM pre-training의 데이터 혼합 전략이 단일 최적화 관점에만 치중되어, 복잡하고 동적인 학습 과정을 충분히 반영하지 못하는 문제를 해결합니다. 기존의 정적(Offline) 데이터 혼합 방식은 학습 중 변화하는 모델의 학습 역동성을 고려하지 못하며, 기존의 동적(Online) 방식인 ODM이나 AC-ODM은 데이터 품질이나 gradient alignment 등 제한적인 보상 신호에만 의존한다는 한계가 있습니다 [Figure 2]. 이러한 단일 관점의 최적화는 상호 보완적인 도메인 간의 시너지 효과를 간과하게 만들어 학습 효율성을 저해합니다. 따라서 저자들은 다차원적인 관점을 고려하는 HDS를 통해 보다 안정적이고 효율적인 pre-training 정책을 학습하고자 합니다.

3. Method & Key Results (제안 방법론 및 핵심 결과)

본 논문은 데이터 스케줄링을 Markov Decision Process (MDP)로 정의하고 SAC 알고리즘을 통해 최적의 domain weight vector를 도출하는 HDS를 제안합니다. 제안된 프레임워크는 state 공간(모델 성능, 학습 속도, weight norm 등 포함)을 관찰하여 domain sampling 확률을 결정하며, 세 가지 요소가 결합된 holistic reward를 통해 agent를 학습시킵니다 [Figure 2], [Figure 3]. The Pile 벤치마크를 활용한 Pythia-1B 모델 학습 실험 결과, HDS는 정적 baseline인 TPW 대비 57% 적은 학습 iteration으로 동일한 validation perplexity에 도달했습니다 [Figure 1]. 또한, AC-ODM 대비 44% 빠른 학습 속도를 기록하며 전반적인 학습 효율성을 획기적으로 개선했습니다. Downstream task 성능에서도 MMLU 0-shot 기준 7.2%, 5-shot 기준 4.0%의 향상을 달성하여 모델의 일반화 능력 강화에도 크게 기여함을 입증했습니다 [Table 1], [Table 2].

4. Conclusion & Impact (결론 및 시사점)

본 논문은 다목적 강화학습을 활용하여 LLM pre-training의 데이터 혼합 전략을 최적화하는 HDS를 성공적으로 구축하였습니다. 이 연구는 모델의 내적 상태를 고려한 동적 스케줄링이 학습 효율과 모델 성능을 동시에 극대화할 수 있음을 실증했습니다. 특히, 제안된 보상 체계는 향후 대규모 모델 pre-training 과정에서 데이터 효율성을 개선하기 위한 중요한 가이드라인으로 활용될 것으로 기대됩니다. 본 연구는 학계뿐만 아니라 막대한 자원이 소요되는 대규모 언어 모델 학습의 경제성 및 환경적 지속가능성 개선에 중요한 시사점을 제공합니다.

---

Part 2: 중요 Figure 정보

Figure 1 — HDS 모델 성능 비교

Figure 2 — HDS 프레임워크 개요

Figure 3 — Actor-Critic 네트워크 아키텍처

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.