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

TL;DR

MergeMix tackles the costly problem of optimizing mid-training data mixtures for large language models by recasting it as a weight-space merging task. It trains domain-specific experts on limited data, then learns a performance surface that maps merging weights to downstream capabilities, enabling a utility-driven optimization without full retraining. The approach is theoretically grounded in first-order equivalence between data-gradient accumulation and weight interpolation, with second-order terms bounded, and empirically validated on 8B and 16B parameter models where it matches or surpasses manual tuning at over 100× lower search cost. The method achieves high rank consistency (Spearman ρ>0.9) and strong cross-scale transfer, offering a scalable, automated path to enhancing targeted capabilities while reducing the cost and effort of data-mix exploration.

Abstract

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.

Figures (9)

• Figure 1: Cost-performance efficiency analysis. (Top) Comparison of estimated computational costs (log scale) and downstream benchmark accuracy across different data mixing strategies. Details on cost estimation are provided in Appendix \ref{['app:cost']}. (Bottom) Conceptual illustration of the search dynamics. Conventional methods require increasing computational investment to asymptotically approach the optimal performance zone through iterative trials or fitting. In contrast, MergeMix leverages weight-space merging as an proxy, identifying the optimal mixtures with minimal cost.
• Figure 2: Rank consistency between model merging and data mixture training. The high correlation indicates that weight interpolation accurately predicts the relative ranking of data mixtures. We also present the value of $\lambda$ for each configure in percent.
• Figure 3: Performance trend comparison between model merging and actual mixture tuning on math and code benchmarks.
• Figure 4: Training dynamics comparison across five domains. The light-colored curves (Pre-Anneal) track the performance of models trained with a constant learning rate. The dark-colored curves (Annealed) represent the performance after applying learning rate annealing (simulated by merging the most recent 20B tokens). The horizontal dashed lines denote the final performance by merging the top-16 checkpoints. The model trained with MergeMix-derived ratios consistently matches or outperforms the strong manually tuned baseline.
• Figure 5: Performance comparison between the MergeMix-optimized mixture, uni-domain baselines (100% single-domain data), and an aggressive specialization mixture.