훈련과 검색의 분리: 대규모 언어 모델 사전 훈련을 위한 모델 병합 기반 데이터 혼합 확장

Decouple Searching from Training: Scaling Data Mixing via Model Merging for Large Language Model Pre-training

January 31, 2026

저자: Shengrui Li, Fei Zhao, Kaiyan Zhao, Jieying Ye, Haifeng Liu, Fangcheng Shi, Zheyong Xie, Yao Hu, Shaosheng Cao

cs.AI

초록

효과적인 데이터 혼합 방식을 결정하는 것은 대규모 언어 모델(LLM) 사전 학습의 핵심 요소로, 모델이 일반적인 능력과 수학 및 코드 같은 어려운 과제 수행 능력 간의 균형을 유지해야 합니다. 그러나 기존 접근법은 신뢰할 수 없는 소규모 프록시 실험에 의존하거나 지나치게 비용이 많이 드는 대규모 탐색이 필요하기 때문에 최적의 혼합 방식을 규명하는 것은 여전히 해결 과제로 남아 있습니다. 이를 해결하기 위해 우리는 모델 병합을 활용하여 최적의 데이터 비율을 예측하는 새로운 프레임워크인 DeMix(Decouple Searching from Training Mix)를 제안합니다. DeMix는 샘플링된 모든 혼합 방식에 대해 프록시 모델을 훈련하는 대신, 후보 데이터셋으로 대규모 구성 요소 모델을 훈련하고 가중 모델 병합을 통해 데이터 혼합 프록시를 도출합니다. 이 패러다임은 탐색과 훈련 비용을 분리하여 추가 훈련 부담 없이 무제한의 샘플링된 혼합 방식을 평가할 수 있게 하므로 더 많은 탐색 시도를 통해 더 나은 혼합 방식을 발견할 수 있도록 합니다. 대규모 실험을 통해 DeMix가 충분성, 정확성 및 효율성 간의 트레이드오프를 깨고, 더 낮은 탐색 비용으로 더 높은 벤치마크 성능을 보이는 최적의 혼합 방식을 얻을 수 있음을 입증했습니다. 또한 개방형 연구를 촉진하기 위해 검증된 혼합 방식과 고품질 사전 학습 데이터로 구성된 포괄적인 22T 토큰 데이터셋인 DeMix Corpora를 공개합니다. 우리의 코드와 DeMix Corpora는 https://github.com/Lucius-lsr/DeMix에서 확인할 수 있습니다.

English

Determining an effective data mixture is a key factor in Large Language Model (LLM) pre-training, where models must balance general competence with proficiency on hard tasks such as math and code. However, identifying an optimal mixture remains an open challenge, as existing approaches either rely on unreliable tiny-scale proxy experiments or require prohibitively expensive large-scale exploration. To address this, we propose Decouple Searching from Training Mix (DeMix), a novel framework that leverages model merging to predict optimal data ratios. Instead of training proxy models for every sampled mixture, DeMix trains component models on candidate datasets at scale and derives data mixture proxies via weighted model merging. This paradigm decouples search from training costs, enabling evaluation of unlimited sampled mixtures without extra training burden and thus facilitating better mixture discovery through more search trials. Extensive experiments demonstrate that DeMix breaks the trade-off between sufficiency, accuracy and efficiency, obtaining the optimal mixture with higher benchmark performance at lower search cost. Additionally, we release the DeMix Corpora, a comprehensive 22T-token dataset comprising high-quality pre-training data with validated mixtures to facilitate open research. Our code and DeMix Corpora is available at https://github.com/Lucius-lsr/DeMix.