다중 목적 강화학습 기반 LLM 사전학습용 Holistic Data Scheduler

사전학습 데이터의 출처 구성과 혼합 전략이 LLM 성능과 학습 효율을 결정한다. HDS는 데이터 혼합을 에이전트가 실시간으로 최적화하도록 하여 학습 스텝과 연산 비용을 크게 줄이고 downstream 성능을 동시에 향상시켰다.

출발점과 기존 한계: LLM 사전학습에서는 서로 다른 도메인들의 혼합비가 학습 효율과 최종 성능에 직접적인 영향을 준다. 기존의 온라인 데이터 믹싱(ODM, AC-ODM 등)은 단일 관점(예: 도메인별 손실 또는 gradient alignment)에 기반한 보상 신호에 의존해 도메인 간 상호작용이나 학습 단계별 데이터 난이도 변화를 충분히 고려하지 못했다. 이로 인해 학습 초기 단계에서 지나치게 어려운 데이터가 과다 샘플링되거나, 상호 보완적 도메인을 활용하지 못해 수렴 속도가 떨어졌다.

해결 원리: HDS는 데이터 혼합을 연속 제어 문제로 모델링하고, 에이전트가 도메인 가중치(확률 심플렉스)를 직접 출력하도록 한다. 상태는 도메인별 누적 샘플 수, 도메인별 검증 손실과 손실 변화, 선택된 레이어의 weight L2 norm 및 변화 등으로 구성되어 에이전트가 모델의 현재 학습 단계와 내부 안정성을 파악하도록 한다. 보상은 세 축으로 구성된다: (1) gradient alignment — 입력: 도메인 i의 gradient g_i와 다른 도메인들의 gradient 합 Σ_{j≠i} g_j → 연산: 내적 ⟨g_i, Σ_{j≠i} g_j⟩ → 결과: 스칼라값 → 의미: 양수이면 도메인 i의 업데이트가 다른 도메인에 유익하다는 신호, (2) scheduled lexical diversity — 입력: 배치의 MTLD → 연산: 정규화 및 t'와 결합(r_diversity = t' / (MTLD_norm + ε)) → 결과: 시간에 따른 curriculum 보상 → 의미: 초기에는 낮은 MTLD(단순 텍스트) 보상, 후반에는 높은 MTLD(복잡 텍스트) 보상, (3) model stability — 입력: 선택된 레이어들의 L2-norm 변화 Δ||ω||_2 → 연산: 역수 1/(|Δ||ω||_2|+ε), 상한치 적용 → 결과: 스칼라 보상 → 의미: 급격한 파라미터 변화를 억제해 안정적 수렴 유도.

달라지는 점(구체적 변화): 이 구성은 에이전트가 '어떤 도메인을, 언제 더 많이 샘플링할지'를 학습 단계별로 조정하게 해 학습 초기에는 단순·반복 데이터로 기초를 다지고, 중·후반에는 고다양성 도메인으로 전환하는 커리큘럼을 자동 형성한다. 실험에서 이 접근은 TPW 대비 동일 perplexity 달성에 필요한 스텝을 약 57% 줄였고, downstream(MMLU 0-shot) 성능을 0.26915로 끌어올렸다.

전체 접근 방식: HDS는 온라인 데이터 믹싱을 Markov Decision Process로 정식화하여 에이전트가 연속 행동 공간(확률 심플렉스)에서 도메인 가중치 a^{t+1}을 출력하도록 한다. 환경은 LLM의 학습 프로세스 자체이며, 에이전트는 현재 상태 s^t(샘플 수, 스텝, 도메인별 검증 손실 및 델타, 선택된 레이어의 L2 norm 및 그 변화)를 관찰해 stochastic policy π_{θ_A}(a|s)를 통해 샘플링 확률을 결정한다. 이 행동으로 구성된 배치 B^t로 LLM을 업데이트하고 보상 r^t를 계산해 transition을 리플레이 버퍼에 쌓아 SAC 에이전트를 업데이트한다.

핵심 메커니즘/알고리즘 상세: 보상은 세 성분의 가중합으로 구성된다. (1) Inter-domain influence r_align: 도메인 i의 gradient g_i = ∇{θ_M} ℒ(θ_M, B_i)와 다른 도메인들의 gradient 합의 내적 ⟨g_i, Σ{j≠i} g_j⟩을 계산해 도메인의 cross-domain 유용성을 측정한다(입력: gradient 벡터들 → 연산: 내적 → 결과: 스칼라 alignment score → 의미: 양수면 긍정적 전이). (2) Scheduled lexical diversity r_diversity: 배치별 MTLD를 MTLD_norm으로 정규화하고 현재 상대 스텝 t' = t/T_total을 이용해 r_diversity = t'/(MTLD_norm + ε)로 계산한다(입력: MTLD → 연산: 정규화 및 역수 결합 → 결과: 시간 의존적 보상 → 의미: 초기에는 단순 텍스트 보상, 후반에는 복잡 텍스트 보상). (3) Model stability r_stability: 선택된 레이어들의 L2-norm 변화량 | ||ω^t||_2 − ||ω^{t-1}||_2 |을 계산해 역수를 취하고 상한치로 캡하여 안정성 보상을 얻는다(입력: L2-norm 변화 → 연산: 역수 및 상한 적용 → 결과: 스칼라 → 의미: 급격한 파라미터 변화를 억제).

학습·구현 세부사항: SAC 에이전트는 twin critics와 actor, learnable entropy temperature α_ent을 사용해 업데이트된다. critic 목표값은 y = r + γ(min_k Q'{Ck}(s',a') − α_ent log π(a'|s'))이며, critic은 MSE로, actor는 α_ent log π − min_k Q를 최소화해 갱신된다. 네트워크 아키텍처는 입력 선형투영(256) → 8개의 Transformer encoder block(hidden=512) → 4-layer MLP 출력 구조로, 각 네트워크 약 5M 파라미터이며 에이전트 전체 파라미터는 약 26.5M이다. 하이퍼파라미터 예: w_align=1, w_diversity=10, w_stability=10, r_stability 상한 5, N_ac=2. 데이터 샘플링: action a^{t+1}에 따라 도메인 D_i를 확률적으로 선택하고 UNIFORM(D_i)에서 미니배치를 샘플링해 인스턴스 수준 분포 P{a^t}를 구성한다.

메인 벤치마크 결과: Pythia-1B를 The Pile(50B tokens, 41,667 steps)로 학습한 실험에서 HDS는 validation perplexity 감소 곡선이 가장 가파르게 나타났다. HDS는 AC-ODM 대비 동일한 최종 perplexity를 약 44% 적은 스텝으로 달성했고, TPW 대비 동일한 perplexity 달성에 필요한 스텝을 약 57% 절감했다(실제 Steps: TPW 41667 → HDS 17917). 최종 스텝(41,667)에서 HDS는 TPW 대비 검증 perplexity 13.6% 감소, ODM 대비 8.9% 감소, AC-ODM 대비 5.3% 감소를 기록했다. Downstream MMLU 성능은 0-shot 0.26915, 5-shot 0.31064로 보고되었다(Table 1), 0-shot에서 AC-ODM(0.25146) 대비 절대 +0.01769 포인트 개선을 달성했다.

Ablation study 결과: 보상 성분 제거 실험에서 r_align 제거가 가장 큰 성능 저하를 초래했고(r_align은 핵심 요인), r_diversity 제거가 수렴 속도에 큰 영향을 미쳤다. r_stability 제거 시에도 성능 저하가 관찰되었으나 영향도는 상대적으로 작았다(그럼에도 통합 모델이 최저 perplexity를 기록).

효율성/속도 분석: 에이전트 통합에 따른 시간 오버헤드는 미미했다(시간/스텝 2.47s → 2.49s, <1% 증가). HDS는 TPW 기준 동일 perplexity 달성 시 총 학습 시간에서 2.21x의 speedup을 보였다(Table 3). 모델 스케일 실험(Pythia-12B, 25B tokens)에서도 HDS가 ODM 대비 모든 체크포인트에서 더 낮은 perplexity를 기록했고, 마지막 체크포인트(20,832 steps)에서 HDS 4.89 vs ODM 7.32로 약 33% 상대 개선을 보였다(Table 5).

아키텍처 구조: 에이전트(Actor/Critic)는 입력 선형 투영(dim=256) → 8개의 Transformer encoder block(hidden=512) → 4-layer MLP로 구성되어 상태의 고차원 상호의존성을 모델링한다. 각 네트워크는 약 5M 파라미터이며 전체 에이전트 파라미터는 약 26.5M이다. LLM은 Pythia 계열의 decoder-only Transformer(실험 주요: Pythia-1B, 추가 스케일 실험: Pythia-12B)이다.

핵심 메커니즘의 수학적/알고리즘적 기반: 문제는 MDP(𝒮, 𝒜, 𝒫, ℛ, γ)로 정식화된다. 상태 s^t = (n^t, t, l^t, Δl^t, ||ω^t||2, ||Δω^t||2)로 구성된다. 행동 a^t는 도메인 가중치 벡터 a ∈ Δ^K(확률 심플렉스). r_align 계산: 입력 g_i^t(도메인 i의 gradient)와 Σ{j≠i} g_j^t → 연산: 내적 ⟨g_i^t, Σ{j≠i} g_j^t⟩ → 결과: 스칼라 alignment score → 의미: 양수는 cross-domain 전이가 긍정적임을 나타낸다. r_diversity 계산: 입력 MTLD(B_i^t) → 연산: 정규화 MTLD_norm 및 t' = t/T_total 결합 → 결과: r_diversity = t'/(MTLD_norm + ε) → 의미: 시간 가중 커리큘럼. r_stability 계산: 입력 ||ω^t||_2 변화 → 연산: 역수 및 상한 적용 → 결과: 안정성 보상.

Prior work 대비 차별점: 기존 ODM은 도메인별 손실을 보상으로 사용하거나 MAB 관점으로 접근했고, AC-ODM는 gradient alignment를 도입했으나 단일 관점 보상이 주가 되었다. HDS는 세 가지 관점의 보상을 통합해 도메인 선택을 시간에 따라 조율하는 커리큘럼 효과와 교차 도메인 시너지를 동시에 취한다.

구현 및 학습 세부사항: 트레이닝 하드웨어는 Intel Xeon Platinum 8468 CPU와 NVIDIA H800 80GB GPU 8대. 시퀀스 길이 1024, 글로벌 배치 1152(micro-batch 8 per GPU, grad accumulation 18). Warm-up 833 steps 동안 The Pile 원래 가중치를 소음 N(0,0.02)로 사용 후 HDS 활성화. SAC 하이퍼파라미터: w_align=1, w_diversity=10, w_stability=10, r_stability cap=5, N_ac=2. 액터·크리틱 업데이트는 SAC 표준 손실을 사용하며 타깃 네트워크는 Polyak averaging으로 갱신된다.

스케일 가이드: 에이전트 파라미터는 LLM 파라미터의 약 0.3%~1.5% 범위 권장(≈0.5%에서 비용 대비 성능 균형이 우수). r_align 계산 시 더 깊은(후반) 레이어 선택이 약간 더 유리한 것으로 보고되었다.