MergeMix: 학습 중 데이터 혼합 최적화를 위한 학습 가능한 모델 병합 기법

MergeMix: Optimizing Mid-Training Data Mixtures via Learnable Model Merging

대규모 언어 모델(LLM)의 잠재력을 최대한 활용하기 위해서는 데이터 혼합을 최적화하는 것이 필수적이지만, 최적의 구성을 찾는 것은 휴리스틱 기반의 시행착오 또는 비용이 많이 드는 프록시 학습에 의존하기 때문에 계산적으로 매우 어렵습니다. 이러한 문제를 해결하기 위해, 본 논문에서는 모델 병합 가중치를 고품질의 저비용 성능 지표로 활용하여 최적의 데이터 혼합 비율을 효율적으로 결정하는 새로운 방법인 extbf{MergeMix}를 제안합니다. MergeMix는 최소한의 토큰으로 도메인별 전문가 모델을 학습하고, 이들의 병합 가중치를 다운스트림 벤치마크에 맞게 최적화하여, 전체 규모의 학습 비용 없이 데이터 혼합의 성능을 최적화합니다. 8B 및 16B 파라미터 모델에 대한 광범위한 실험 결과, MergeMix는 철저한 수동 튜닝과 비교하거나 능가하는 성능을 달성하면서도 검색 비용을 획기적으로 줄이는 것으로 나타났습니다. 또한, MergeMix는 높은 순위 일관성(Spearman $ρ > 0.9$) 및 강력한 크기 간 전이성을 보여주며, 데이터 혼합 최적화를 위한 확장 가능하고 자동화된 솔루션을 제공합니다.

Optimizing data mixtures is essential for unlocking the full potential of large language models (LLMs), yet identifying the optimal composition remains computationally prohibitive due to reliance on heuristic trials or expensive proxy training. To address this, we introduce \textbf{MergeMix}, a novel approach that efficiently determines optimal data mixing ratios by repurposing model merging weights as a high-fidelity, low-cost performance proxy. By training domain-specific experts on minimal tokens and optimizing their merging weights against downstream benchmarks, MergeMix effectively optimizes the performance of data mixtures without incurring the cost of full-scale training. Extensive experiments on models with 8B and 16B parameters validate that MergeMix achieves performance comparable to or surpassing exhaustive manual tuning while drastically reducing search costs. Furthermore, MergeMix exhibits high rank consistency (Spearman $ρ> 0.9$) and strong cross-scale transferability, offering a scalable, automated solution for data mixture optimization.