기술 요약: OP-MIX (On-Policy Mix)

1. 문제 제기

데이터 믹싱 (data mixing) 은 언어 모델 (LM) 학습 중 서로 다른 데이터 도메인에 대한 최적의 확률 분포를 결정하는 과정으로, 모델 품질에 있어 결정적인 요소입니다. 그러나 기존 데이터 믹싱 방법들은 전체 LM 수명 주기 (lifecycle) 의 맥락에서 다음과 같은 중대한 한계에 직면해 있습니다:

• 정적 도메인 가정: 대부분의 방법은 고정된 데이터 도메인 집합을 가정하며, 새로운 데이터셋 (예: 지시어 튜닝 코퍼스) 이 동적으로 유입되는 지속적 학습 (continual learning) 의 현실을 반영하지 못합니다.
• 오프-폴리시 프록시에 대한 의존: 기존 접근법들은 종종 후보 혼합물에 대해 별도의 작은 "프록시" 모델을 학습시켜 성능을 외삽합니다. 이는 계산 비용이 많이 들고 도메인 수에 따라 확장성이 떨어지며, 특히 학습 후기 단계 (중간 학습, 지시어 튜닝) 에서 프록시의 학습 역학이 대상 모델과 달라지는 분포 이동 (distributional shift) 으로 인해 성능이 저하됩니다.
• 수명 주기의 단편화: 사전 학습, 지속적 중간 학습, 지속적 지시어 튜닝을 모두 효과적으로 처리하는 통합된 알고리즘은 존재하지 않습니다.

본 논문은 데이터 믹싱이 근본적으로 온라인 의사결정 문제이며, 별도의 프록시 모델에 의존하지 않고 진화하는 데이터에 적응할 수 있는 단일 통합 솔루션이 필요하다고 주장합니다.

2. 방법론: OP-MIX

저자들은 훈련 실행의 조합적 폭발과 오프-폴리시 프록시의 편차를 피하기 위해 두 가지 핵심 통찰을 활용하여 최적의 데이터 믹싱 비율을 추정하는 OP-MIX (On-Policy Mix) 알고리즘을 제안합니다.

핵심 메커니즘

1. LoRA 를 통한 온-폴리시 프록시: 별도의 작은 모델을 학습시키는 대신, OP-MIX 는 현재 베이스 모델로부터 각 데이터 도메인당 단일 저랭크 적응 (LoRA) 어댑터를 학습합니다. 이를 통해 프록시가 실제 학습 중인 모델의 학습 역학과 매개변수 상태를 반영하는 "온-폴리시" 상태를 유지하도록 보장합니다.
2. 혼합 시뮬레이션을 위한 선형 보간: 재학습 없이 다양한 데이터 혼합을 시뮬레이션하기 위해, OP-MIX 는 학습된 LoRA 어댑터의 가중치를 선형적으로 보간합니다. 최근의 선형 모드 연결성 (linear mode connectivity) 에 관한 발견에 기반하여, 전체 모델을 혼합할 때의 손실 표면 (loss surface) 은 LoRA 어댑터를 보간함으로써 근사할 수 있다고 가정합니다. 이를 통해 알고리즘은 전체 학습 실행 대신 저렴한 순전파 (forward passes) 를 통해 후보 혼합물을 평가할 수 있습니다.

알고리즘 워크플로우

$K$개의 새로운 도메인이 유입되는 지속적 학습 단계에서:

1. 어댑터 학습: 각 새로운 도메인 ($D_{m+k}$) 에 대해 하나의 LoRA 어댑터를 학습하고, 이전 모든 도메인의 혼합을 나타내는 하나의 "이전" 어댑터 ($D_{old}$) 를 학습합니다. 이는 모두 현재 베이스 모델 $\theta_{base}$에서 초기화됩니다.
2. 보간 및 평가: $K$-심플렉스 (K-simplex) 내에서 보간 점 $\alpha$를 샘플링합니다. LoRA 가중치를 선형 결합하여 병합된 어댑터를 형성합니다: $\theta_{\alpha} = \alpha_{old}\theta_{old} + \sum \alpha_{k}\theta_{D_{m+k}}$. 이러한 병합된 어댑터의 도메인별 손실을 평가합니다.
3. 회귀 및 최적화: 평가된 손실에 로그-선형 회귀 모델을 적합시켜 보간 경로 전반에 걸친 매끄러운 손실 표면을 추정합니다. 추정된 손실을 최소화하는 최적의 믹싱 가중치 $\alpha^*$를 찾기 위해 정규화 최적화 문제를 풉니다 (이때 사전 분포를 향한 KL-발산 정규화 항이 포함됩니다).
4. 최종 학습: 계산된 최적의 혼합 가중치를 베이스 모델에 적용하여 전체 학습 실행을 수행합니다.

3. 주요 기여

본 논문은 세 가지 주요 기여를 주장합니다:

1. 최초의 범용 데이터 믹싱 알고리즘: OP-MIX 는 새로운 도메인으로 데이터 혼합을 확장하면서도 별도의 프록시 모델을 피할 수 있는 최초의 방법입니다. 이는 데이터가 진화함에 따라 지속적인 믹싱을 가능하게 하여, 단계별 알고리즘의 필요성을 제거합니다.
2. 수명 주기를 아우르는 최첨단 성능: OP-MIX 의 단일 인스턴스가 사전 학습, 지속적 중간 학습, 지속적 지시어 튜닝에서 최첨단 성능을 달성하여, 서로 다른 단계가 서로 다른 믹싱 전략을 필요로 하지 않음을 입증합니다.
3. 지속적 학습에서의 효율성: 지속적 지시어 튜닝에서 OP-MIX 는 정교한 지속적 학습 기법인 온-폴리시 자기 증류 (self-distillation) 와 동등한 성능을 내면서도 95% 적은 계산 자원을 사용합니다. 또한 데이터 믹싱과 학습 목표 (예: 자기 증류) 는 추가적인 개선을 위해 결합될 수 있는 직교적인 개선 축임을 보여줍니다.

4. 실험 결과

저자들은 150M 에서 7B 파라미터에 이르는 모델을 사용하여 세 가지 단계에서 OP-MIX 를 평가했습니다.

• 사전 학습: OP-MIX 는 믹싱 없이 학습한 경우 (경험적 위험 최소화) 대비 평균 퍼플렉시티를 6.3% 개선했습니다. 이는 가장 좋은 베이스라인 (MergeMix) 과 동등한 성능을 내면서 14% 적은 계산 자원을 사용했으며, 모든 규모에서 OLMix 보다 5~6% 더 우수한 성능을 보였습니다.
• 지속적 중간 학습: OP-MIX 는 전체 재학습 ("하늘선" 베이스라인) 과 거의 동등한 성능을 내면서 최대 66% 적은 계산 자원을 사용했습니다. 또한 표준 지속적 SFT 및 고정 데이터 재생 베이스라인에 비해 재앙적 망각 (catastrophic forgetting) 을 현저히 완화했습니다.
• 지속적 지시어 튜닝: Qwen2.5-7B-Instruct 에 적용했을 때, OP-MIX 는 95% 적은 계산 자원으로 자기 증류 파인튜닝 (SDFT) 과 비교 가능한 성능을 달성했습니다. OP-MIX 와 SDFT 를 결합하면 61.9% 정확도로 전체적으로 가장 좋은 결과가 나와, 두 기법의 상승 효과가 입증되었습니다.
• 효율성 프론티어: 모든 설정에서 OP-MIX 는 기존 베이스라인을 파레토 우위 (Pareto-dominates) 하여, 엄격하게 낮은 계산 비용 (FLOPs) 으로 더 좋거나 동등한 성능을 달성했습니다.

5. 중요성과 주장

본 논문은 언어 모델 학습을 사전 학습, 중간 학습, 파인튜닝과 같은 서로 다른 단계로 나누는 전통적인 관점이 인위적이라고 주장합니다. OP-MIX 는 학습이 데이터로부터의 단일 연속적인 학습 과정이라는 통합된 관점을 제시합니다.

LoRA 와 선형 모드 연결성을 활용함으로써, OP-MIX 는 적응형 데이터 믹싱을 계산적으로 불가능한 오프라인 검색에서 가벼운 반복 가능한 작업으로 변환하여 새로운 데이터가 도착할 때마다 실행할 수 있게 합니다. 저자들은 이 접근법이 기존 연구의 두 가지 주요 한계, 즉 새로운 데이터셋을 통합하지 못한다는 점과 대상 모델의 역학에서 벗어난 오프-폴리시 프록시 모델에 의존한다는 점을 극복했다고 결론지었습니다.

참고: 본 논문은 한계를 인정하며, 특히 사전 학습/중간 학습은 530M 파라미터까지, 지시어 튜닝은 7B 파라미터까지 실험이 제한되어 있어 프론티어 규모 (70B 이상) 의 동작은 검증되지 않았음을 명시합니다. 또한 매우 많은 수의 도메인 (10~100 개) 에 대한 방법의 동작은 아직 규명되지 않았다고 덧붙입니다.