FastMix: Fast Data Mixture Optimization via Gradient Descent

2026-06-12

AI Summary

This work addresses the challenge of determining optimal data mixing ratios in large language model pretraining and post-training, which traditionally relies on manual heuristics or expensive trial-and-error procedures. The authors propose an efficient, automated method that formulates ratio selection as a bilevel optimization problem and, for the first time, equivalently transforms it into uniform sampling with source-loss weighting embedded within a differentiable objective amenable to end-to-end gradient-based optimization. By employing an approximate alternating update strategy—iteratively optimizing model parameters in the inner loop and mixture coefficients in the outer loop—the approach requires training only a single proxy model guided by validation feedback. Experiments demonstrate that the method consistently outperforms existing baselines in both pretraining and post-training settings while substantially reducing computational search costs.

Abstract

While large and diverse datasets have driven recent advances in large models, identifying the optimal data mixture for pre-training and post-training remains a significant open problem. We address this challenge with FASTMIX, a novel framework that automates data mixture discovery while training only a single proxy model. Instead of relying on predefined heuristics or resource-intensive simulations, FASTMIX jointly optimizes mixture coefficients and model parameters, substantially improving efficiency and scalability over prior approaches. At the core of FASTMIX is a reformulation of mixture selection as a bilevel optimization problem. Under this reformulation, we show that optimizing mixture ratios is mathematically equivalent to assigning per-source loss weights under uniform source sampling. This embeds the mixture coefficients directly into the differentiable iterative optimization objective, enabling efficient, gradient-based optimization of both mixture and model. To solve the optimization problem, FASTMIX implements an approximate iterative optimization procedure, alternating between (i) updating model parameters on data sampled according to current mixture ratios (inner loop) and (ii) updating mixture ratios based on validation feedback (outer loop). Across pre- and post-training, FASTMIX outperforms baselines while drastically reducing search cost. Code (https://github.com/hrtan/fastmix)

Problem

• data mixture optimization
• pre-training
• post-training
• optimal data selection
• large language models

Innovation

• data mixture optimization
• bilevel optimization
• gradient-based optimization
• proxy model
• loss weighting