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

Optimizing data mixtures is is essential for unlocking the full potential of 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 MergeMix, a novel appraoch 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 $\rho > 0.9$) and strong cross-scale transferability, offering a scalable, automated solution for data mixture optimization.

Large Language Models (LLMs) learn from massive amounts of data, but the "recipe"—the specific mix of mathematics, computer code, and general language—is what truly defines their intelligence. Currently, finding the perfect data recipe is incredibly expensive because researchers typically have to train entire models from scratch just to test if a single combination works. We introduce MergeMix, a method that finds the ideal data mixture at a fraction of the cost. Our key discovery is that instead of training many models on different data mixes, we can train small "experts" in specific areas (like a "math expert" and a "coding expert") and then mathematically blend their internal settings. We found that this blending process—known as model merging—acts as a high-fidelity mirror: it accurately predicts how a model would have performed if it had been trained on a real mixture of that data. This approach reduces the cost of finding the best data recipe by over 100 times while achieving performance that matches or even beats traditional, expensive manual tuning. By turning a massive computing challenge into a simple mathematical blending task, MergeMix makes training powerful, specialized AI more efficient and accessible.