계층적 데이터 균형 최적화로 LLM 파인튜닝 혁신

초록

본 논문은 대규모 언어 모델(LLM)을 다중 데이터셋으로 파인튜닝할 때 발생하는 전역·국부 데이터 불균형과 이질성을 동시에 해결하는 계층적 균형 최적화(HBO) 프레임워크를 제안한다. 외부 최적화와 내부 샘플링을 bilevel 구조로 설계하고, 전역 Actor와 각 데이터셋별 로컬 Actor가 모델의 학습 상태(그라디언트 L2 norm, 퍼플렉시티 비율)를 보상으로 받아 샘플링 확률을 동적으로 조정한다. 세 가지 LLM 백본(Llama‑3.1‑8B, Qwen2.5‑7B, EuroLLM‑9B)과 9개 다언어·다태스크 벤치마크에서 기존 샘플링 전략을 크게 능가함을 실험적으로 입증한다.

상세 분석

HBO는 기존 데이터 균형 방법이 전역(데이터셋 간) 수준에만 초점을 맞추고, 각 데이터셋 내부의 난이도·품질 차이를 무시한다는 한계를 정확히 짚고 있다. 이를 해결하기 위해 논문은 두 단계의 Actor를 도입한다. 전역 Actor는 전체 데이터셋 집합을 N개의 서브셋으로 보고, 각 서브셋의 샘플링 확률 p_global(i)를 정책 네트워크 ψ_global을 통해 학습한다. 로컬 Actor는 각 서브셋 i를 다시 M_i개의 난이도 그룹(쉬움‑어려움)으로 분할하고, 그룹별 확률 p_local(i,j)를 별도의 정책 네트워크 ψ_local(i)로 최적화한다. 두 Actor 모두 REINFORCE 기반의 정책 그래디언트를 사용하며, 보상 함수는 모델의 현재 학습 상태를 직접 반영한다. 전역 보상 R_global(i)는 해당 서브셋에서 계산된 그라디언트 L2 노름으로 정의되어, 학습이 아직 충분히 진행되지 않은 데이터에 더 높은 가중치를 부여한다. 로컬 보상 R_local(i,j)는 파인튜닝된 모델과 원본 모델의 퍼플렉시티 비율로, 특정 난이도 그룹에서 상대적 성능 향상이 큰 경우 해당 그룹의 샘플링 비율을 증가시킨다. 이러한 설계는 “학습이 필요한 데이터”와 “학습이 충분히 된 데이터”를 자동으로 구분해 자원을 재분배한다는 점에서 혁신적이다.

알고리즘 1은 전체 학습 루프를 명확히 제시한다. 매 스텝마다 전역 샘플링 → 로컬 샘플링 → 배치 업데이트가 이루어지고, 지정된 주기(F_global, F_local)마다 각각의 보상을 계산해 정책 파라미터 ψ_global, ψ_local를 업데이트한다. 정책 네트워크는 2‑layer fully‑connected 구조로 경량화돼 전체 학습 시간에 약 15 % 정도만 추가 오버헤드를 발생시킨다. 또한, 데이터 전처리 단계에서 SuperFiltering 기법을 이용해 각 서브셋을 난이도 4그룹으로 나누어 로컬 Actor가 보다 정밀한 조정을 할 수 있게 한다.

실험에서는 세 가지 백본 모델을 동일한 하이퍼파라미터 설정 하에 다국어(MMMLU, XCOPA, XStoryCloze, XNLI, MGSM)와 다태스크(MMLU, MultiFin‑EN, GSM8K, MedMCQA) 시나리오에 적용했다. HBO는 기존의 static balancing, temperature‑scaled sampling, 그리고 최신 동적 샘플링 기법(예: Curriculum‑based, RL‑based)보다 평균 2.3 ~ 4.1 %의 정확도 향상을 기록했다. 특히, 어려운 그룹에 대한 성능이 크게 개선되었으며, 쉬운 예시를 완전히 배제하지 않고 적절히 활용함으로써 전반적인 일반화가 향상되는 현상이 관찰되었다.

분석 파트에서는 전역 Actor가 초기 학습 단계에서 대규모 데이터셋을 집중적으로 샘플링하고, 학습이 진행될수록 균형을 맞춰가는 “사이클형” 패턴을 시각화했다. 로컬 Actor는 난이도 그룹별로 샘플링 비율을 동적으로 조정하며, 특히 중간 난이도 그룹에 대한 비중이 상승하는 경향을 보였다. Ablation 실험에서는 보상 함수를 각각 제거하거나, 로컬 Actor를 비활성화했을 때 성능이 급격히 감소함을 확인해 두 보상의 상호 보완성을 입증했다.

한계점으로는 보상 계산이 추가적인 전방/후방 연산을 필요로 하며, 매우 대규모 데이터(수억 건)에서는 보상 추정 비용이 증가할 수 있다는 점을 언급한다. 또한, 정책 네트워크가 단순히 2‑layer MLP이므로 복잡한 데이터 특성을 완전히 포착하지 못할 가능성도 있다. 향후 연구에서는 메타‑러닝 기반의 보상 예측 모델이나, 샘플링 정책을 트랜스포머 형태로 확장해 더 정교한 균형을 추구할 여지가 있다.

전반적으로 HBO는 LLM 파인튜닝 시 데이터 불균형과 이질성을 동시에 다루는 최초의 계층적 동적 샘플링 프레임워크로, 실험적 검증과 분석을 통해 그 효용성을 충분히 입증하였다.