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

TL;DR

MergeMix introduces a cost-effective method for optimizing data mixtures in large language models by leveraging model merging weights as performance proxies, enabling efficient and scalable data composition tuning.

Contribution

The paper presents MergeMix, a novel approach that uses model merging weights to efficiently optimize data mixing ratios without extensive training or heuristics.

Findings

Achieves comparable or better performance than manual tuning.

Reduces search costs significantly compared to exhaustive methods.

Demonstrates high rank consistency and transferability across models.

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 $\rho > 0.9$) and strong cross-scale transferability, offering a scalable, automated solution for data mixture optimization.