Data Mixing for Large Language Models Pretraining: A Survey and Outlook

1. 문제 정의 (Problem Definition)

대규모 언어 모델 (LLM) 의 전처리 (Pretraining) 는 방대하고 이질적인 (heterogeneous) 코퍼스를 기반으로 수행되며, 훈련 데이터의 구성은 제한된 컴퓨팅 및 데이터 예산 하에서 모델의 훈련 효율성과 하위 작업 일반화 능력에 결정적인 영향을 미칩니다.

• 핵심 문제: 기존 연구는 주로 '샘플 레벨'의 데이터 선택 (Data Selection) 에 초점을 맞추어 왔으나, 본 논문은 **'도메인 레벨'의 데이터 믹싱 (Data Mixing)**에 주목합니다. 즉, 제한된 예산 내에서 다양한 데이터 소스 (도메인) 간의 샘플링 가중치를 최적화하여 모델 성능을 극대화하는 문제를 다룹니다.
• 수학적 형식화: 데이터 믹싱 최적화는 이중 최적화 (Bilevel Optimization) 문제로 형식화됩니다.
  • 외부 문제 (Outer): 검증 손실 (Validation Loss) 을 최소화하는 최적의 가중치 벡터 $k$를 찾습니다.
  • 내부 문제 (Inner): 주어진 가중치 $k$에 따라 모델 파라미터 $\theta(k)$를 훈련합니다.
  • 난제: 현대 LLM 전처리 환경에서 이 이중 최적화 문제를 직접 해결하는 것은 계산 비용이 너무 커서 비현실적이므로, 기존 방법들은 이를 실용적으로 풀 수 있는 근사치 (Approximation) 를 사용합니다.

2. 방법론 및 분류 체계 (Methodology & Taxonomy)

저자는 기존 연구들을 체계적으로 분류하기 위해 **정적 (Static)**과 동적 (Dynamic) 믹싱이라는 두 가지 주요 축을 제안하고, 이를 다시 구동 메커니즘에 따라 세분화했습니다.

A. 정적 믹싱 (Static Mixing)

훈련 시작 전 하나의 가중치 $k$를 고정하고 전체 훈련 과정에서 이를 유지하는 방식입니다.

1. 규칙 기반 (Rule-Based):

   • 특징: 추가 학습 비용이 거의 없으며, 휴리스틱 규칙이나 간단한 알고리즘에 의존합니다.
   • 예시: 균일 샘플링 (Uniform), 비례 샘플링 (Proportional), 온도 기반 샘플링 (Temperature-based, $T$ 파라미터 조절), UniMax (반복 횟수 상한 설정).
   • 장단점: 구현이 간단하고 견고하지만, 최적의 믹싱을 보장하지는 못합니다. 최근에는 포트폴리오 선택 이론 (UtiliMax, MEDU) 을 적용하여 더 정교한 규칙 기반 방법도 등장했습니다.

2. 학습 기반 (Learning-Based):

   • 특징: 보조 모델 (Proxy Model) 을 훈련하여 최적의 믹싱을 학습합니다. 비용은 높지만 성능이 우수합니다.
   • 서브 분류:
     • 프록시 최적화 기반 (Proxy Optimization-based): 프록시 모델 자체를 최적화 도구로 사용합니다. (예: DoReMi - 도메인별 손실 차이를 기반으로 가중치 조정, DoGE - 타겟 도메인 일반화 능력을 기반으로 가중치 조정).
     • 예측 기반 (Prediction-based): 모델 성능과 데이터 믹싱 간의 함수 관계 $f(k)$를 학습합니다.
       • 명시적 예측 (Explicit): 수학적 함수 (Scaling Laws, 다항식 등) 를 정의하여 적합합니다. (예: DML, AutoScale, BiMix).
       • 암시적 예측 (Implicit): 블랙박스 함수로 간주하고 회귀 모델 (LightGBM, Gaussian Process 등) 로 학습합니다. (예: RegMix, MixMin, MDE, ADMIRE).

B. 동적 믹싱 (Dynamic Mixing)

훈련 과정에서 모델의 상태에 따라 시간 $t$에 따라 가중치 $k_t$를 업데이트하는 방식입니다.

1. 적응형 (Adaptive):

   • 특징: 외부 모델 없이 훈련 중 생성된 내부 신호 (손실, 기울기, 학습 속도 등) 를 직접 사용하여 믹싱을 업데이트합니다.
   • 예시: ODM (다중 팔 밴딧 문제 접근), ADO (스케일링 법칙 기반 학습 잠재력 추정), Velocitune/GRAPE (학습 속도 느린 도메인 우선), PiKE (기울기 노름과 분산 기반).
   • 장점: 계산 오버헤드가 낮고 통합이 용이합니다.

2. 외부 주도 (Externally Guided):

   • 특징: 외부 컨트롤러 (보통 강화학습 에이전트 또는 예측 모델) 가 훈련 신호를 처리하여 믹싱 업데이트를 지시합니다.
   • 예시: AC-ODM, Data Mixing Agent (Actor-Critic 구조), TiKMiX (영향 함수 기반).
   • 장점: 더 복잡한 최적화 알고리즘과 장기적 피드백을 활용할 수 있어 성능이 높을 수 있으나, 컨트롤러 구축 비용이 발생합니다.

3. 주요 기여 (Key Contributions)

1. 체계적인 개요 및 분류: 데이터 믹싱 연구의 분산된 문헌을 통합하여, 정적/동적 및 규칙/학습 기반의 계층적 분류 체계 (Taxonomy) 를 최초로 제시했습니다.
2. 이중 최적화 관점의 명확화: 데이터 믹싱 문제를 이중 최적화 문제로 형식화하고, 각 방법론이 이를 어떻게 근사화하여 실용화했는지 이론적 배경을 설명했습니다.
3. 성능 - 비용 트레이드오프 분석: 각 방법론의 장단점을 '성능 향상'과 '추가 비용 (계산/데이터)'의 관점에서 분석하여, 연구자들이 상황에 맞는 방법을 선택할 수 있는 기준을 마련했습니다.
4. 미래 방향 제시: 현재 해결되지 않은 문제들과 향후 연구 방향을 구체적으로 제시했습니다.

4. 결과 및 논의 (Results & Discussion)

• 성능 비교: 학습 기반 방법 (특히 예측 기반 및 동적 방법) 은 단순한 규칙 기반 방법보다 제한된 예산 하에서 더 나은 모델 성능을 달성하는 것으로 나타났습니다.
• 전송 가능성 (Transferability) 의 한계: 학습 기반 방법들은 특정 모델 아키텍처, 데이터 구성, 검증 세트에 강하게 의존하는 경향이 있어, 다른 설정으로 전환 시 성능이 급격히 떨어지는 문제가 존재합니다.
• 평가의 비표준화: 연구 간 비교가 어렵습니다. 도메인 정의, 최적화 신호 (손실 vs 하위 작업 지표), 비용 측정 기준 (토큰 수, FLOPs, 시간) 등이 연구마다 상이하여 통일된 벤치마크가 부재합니다.

5. 향후 과제 및 전망 (Challenges & Future Directions)

논문은 다음과 같은 주요 도전 과제와 탐구 방향을 제시합니다:

1. 제한된 전송 가능성 해결:

   • 도메인 정의의 세분화: 단순한 '데이터 소스' 기반 분할을 넘어, 주제 (Topic) 나 기술 (Skill) 기반의 더 세밀하고 의미 있는 도메인 분할 필요.
   • 최적화 목표 및 검증 세트의 통일: 다양한 훈련 단계 (Pretraining, SFT, DPO) 와 목적에 맞는 통일된 프레임워크 및 대표성 있는 검증 세트 구축 필요.

2. 평가 프로토콜 및 벤치마크 표준화:

   • 비용 (Cost) 과 성능 (Performance) 을 동시에 고려하는 통일된 평가 지표와 벤치마크 마련이 시급함.

3. 성능 - 비용 트레이드오프의 극복:

   • 이론적 기반 강화: 데이터 믹싱에 특화된 이론적 프레임워크 (예: 정밀한 스케일링 법칙) 가 부족함. 이를 보완할 새로운 이론 개발 필요.
   • 효율적인 학습자 모델: 프록시 모델이나 컨트롤러로 사용할 최적의 경량 모델 아키텍처 탐색.

4. 탐구적 방향 (Exploratory Directions):

   • 더 세분화된 데이터 믹싱: 도메인 레벨에서 샘플 레벨로 granularity 를 낮추거나, 중간 단계 (Intermediate granularity) 를 모색.
   • 역설계 (Inverse Design): 폐쇄형 모델의 훈련 데이터 믹싱을 모델 자체 (토크나이저 병합 리스트 등) 를 통해 역추적 (Inference) 하는 연구.
   • 파이프라인 인식 (Pipeline-Aware) 믹싱: 전처리, SFT, 정렬 (Alignment) 등 전체 훈련 파이프라인을 하나의 연속된 과정으로 간주하여 단계 간 조율된 믹싱 전략 수립.

6. 의의 (Significance)

본 논문은 LLM 전처리를 위한 데이터 믹싱 연구에 대한 최초의 포괄적인 개요 (Survey) 로서, 다음과 같은 의의를 가집니다:

• 공통의 언어 제공: 분산된 연구들을 통합된 문제 정의와 분류 체계 하에 배치하여, 연구자들 간의 소통과 비교를 용이하게 함.
• 방법론적 나침반: 새로운 방법론을 설계할 때 고려해야 할 성능 - 비용 트레이드오프와 설계 선택지 (정적/동적, 규칙/학습) 에 대한 명확한 가이드 제공.
• 미래 연구의 방향성 제시: 단순한 성능 향상을 넘어, 이론적 기반 강화, 표준화, 그리고 전체 파이프라인 최적화를 위한 연구 방향을 제시하여 해당 분야의 성숙도를 높이는 데 기여함.