FastMix: Fast Data Mixture Optimization via Gradient Descent

Excited to share our team's (Tencent Hunyuan, HKU, CUHK) work from last year: "Fast Data Mixture Optimization via Gradient Descent", which has been accepted to ICLR 2026!

TL;DR: We reformulate the data sampling ratio assignment into a bilevel optimization problem corresponding to loss proportions under uniform sampling, proving their equivalence in expectation. Consequently, the mixture coefficients naturally enter the PyTorch computation graph, allowing them to be solved directly via gradient descent.

This represents our deep exploration into the field of data mixture optimization. Historically, optimizing data mixtures has faced numerous challenges. Existing approaches either rely on predefined heuristics, which lack flexibility, or require training a large number of small proxy models to simulate the process—a highly resource-intensive and inefficient endeavor that also suffers from a potential scale gap.

Our proposed framework, FastMix, aims to break through these limitations.

The core highlight of FastMix is reformulating the data sampling ratio assignment into a bilevel optimization problem linked to the loss proportions under uniform sampling. In principle, optimizing the mixture ratios is equivalent to appropriately allocating loss weights for each data source under uniform source sampling. In this way, the mixture coefficients are directly integrated into a differentiable iterative optimization objective, enabling highly efficient, gradient-based optimization for both the data mixture and the model.

To implement this, FastMix constructs an approximate iterative optimization process. During iterations, it updates the model parameters on data sampled according to the current mixture ratios (inner loop) while simultaneously updating the mixture ratios based on validation feedback (outer loop) in an alternating fashion (Figure 2).

In our experiments, FastMix consistently matches or outperforms the strongest baselines across both pre-training and post-training stages, while drastically reducing search costs. During pre-training, it is up to 550× faster than the state-of-the-art method RegMix; during post-training, the time cost is reduced by 52× compared to RegMix.

Of course, there are still many exciting open questions worth exploring (see Section 4.3). We sincerely hope this paper can provide new perspectives and inspiration for researchers working in related fields.

The code is now open-source! For those interested, please check out our repository here 👉 [https://github.com/hrtan/fastmix] to dive deeper. We warmly welcome discussions, collaborations, and growing together!