Source-linked AI summary
FastMix: Fast Data Mixture Optimization via Gradient Descent
Haoru Tan, Sitong Wu, Yanfeng Chen, Jun Xia, Ruobing Xie, Bin Xia, Xingwu Sun, Xiaojuan Qi
TL;DR
Choosing effective data mixtures for large-scale pre-training and post-training remains difficult and costly. FASTMIX uses weighted bilevel, gradient-based optimization with a single proxy model, outperforming baselines across settings while substantially reducing search cost.
Problem
Identifying effective data mixtures for large-scale pre-training and post-training remains challenging because existing manual and proxy-based searches scale poorly.
Method
FASTMIX reformulates mixture selection as weighted bilevel optimization, jointly learning mixture coefficients and model parameters through gradient-based inner and outer updates.
Results
48.2 average score and rank 1 across 14 pre-training benchmarks, with 1.3 GPU-hours of search, while FASTMIX also outperforms baselines in post-training.
Takeaways & Limitations
FASTMIX makes data mixture optimization more practical and scalable by achieving strong accuracy with substantially lower search cost.
Takeaways & Limitations
FASTMIX can fail when long pre-training sequences and shorter SFT data produce substantially divergent gradients.
Abstract
from arXiv · showhide
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)
1 Introduction
FASTMIX addresses the challenge of identifying effective training-data mixtures by jointly optimizing mixture coefficients and model parameters while training only one proxy model. It improves scalability over prior approaches and achieves strong pre-training results with substantially lower computational cost.
- Large-model performance depends critically on training data, yet identifying optimal mixtures for pre-training and post-training remains challenging.
- Manual trial-and-error and proxy-based methods scale poorly because proxy-based searches require training many proxy models, creating prohibitive computational overhead.RegMix and CLIMB provide strong generalization and stability but incur this search cost.
- FASTMIX automates mixture discovery by training only a single proxy model and jointly optimizing mixture coefficients with model parameters through weighted bilevel optimization.This avoids predefined heuristics and resource-intensive simulations.
- 48.2 top average score and rank 1 across 14 benchmarks, best on 9, were achieved in pre-training with 1.3 GPU-hours.FASTMIX was evaluated across model scales and tasks in both pre-training and post-training.
- ˆ550 faster than RegMix and ˆ55 than CLIMB, FASTMIX outperformed baselines in pre-training at a fraction of the computational cost.
2 Related Work
Related work primarily explores proxy-based or dynamic data-mixture optimization. FASTMIX retains proxy-based reliability while reducing search cost and improving generalization.
- Data mixture optimization: Large-model progress relies on strategically mixing data across diverse sources, including languages, modalities, and difficulty levels.The resulting data-mixture problem affects both pre-training and post-training.
- Proxy-based methods: Proxy-based methods train proxy models under candidate mixtures, then use proxy performance to infer optimal sampling ratios.DoReMi adjusts domain weights using relative losses, while RegMix trains hundreds of proxies across ratios and fits a regression model to extrapolate the optimum.
- Proxy-based methods: CLIMB improves proxy-based efficiency by iteratively refining the search region, reducing the number of proxy models required.Other work studies whether mixtures found on small models transfer effectively to larger models.
- Dynamic methods: Dynamic methods eliminate a separate search phase by adjusting mixtures during training, with IDEAL dynamically rebalancing data using influence-function estimates.IDEAL estimates domain contributions to downstream performance before rebalancing training data.
- Comparison and FASTMIX: Proxy-based methods generally deliver stronger and more stable performance than dynamic approaches, but require substantial computational cost.FASTMIX cuts search time from hundreds of GPU-hours to nearly one while achieving higher efficiency and stronger generalization.
3 FastMix
FASTMIX formulates data-mixture selection as a bilevel optimization problem and makes mixture weights differentiable by reparameterizing them as per-source loss weights. It jointly updates model parameters and mixture coefficients, while using entropy and training-loss terms to improve generalization beyond validation-only search.
- Bilevel formulation: FASTMIX defines mixture optimization as bilevel learning: an inner loop trains model parameters under mixture weights, while an outer loop minimizes validation loss.The mixture weights lie on the probability simplex, and the outer objective evaluates validation loss at approximately optimized model parameters.
- Differentiable formulation: Reparameterizing stochastic source sampling as uniformly sampled per-source losses weighted by α_i makes the bilevel objective mathematically equivalent and fully differentiable.Each mixture coefficient continuously scales one source’s training loss, enabling gradient-based optimization with optimizers such as SGD or Adam.
- Generalization objectives: FASTMIX improves robustness by combining entropy regularization, which discourages collapsed mixtures, with an auxiliary training-loss target alongside validation performance.The entropy coefficient λ is empirically set to a small value such as 10^-5, while β is most effective at moderate values such as 0.1.
- Iterative optimization: FASTMIX alternates inner SGD updates of model parameters under current mixture weights with outer updates based on validation feedback after n2 training iterations.The procedure assigns larger weights to sources whose training contributes more to improving validation performance.
- Gradient estimation: For n2 = 1 with SGD, the validation gradient with respect to each mixture coefficient is determined directly by the corresponding source training loss.The resulting gradient can be interpreted through alignment between the validation gradient and each source’s training gradient.
4 Experiments
FASTMIX is evaluated for data-mixture optimization in both pre-training and post-training, where it achieves the strongest reported downstream performance and efficiency. The experiments also identify practical limitations involving non-differentiable targets, long outer-loop horizons, small proxy models, and search-target data.
- Experimental setup: FASTMIX is evaluated against expert tuning, proxy-based methods including DoReMi, RegMix, and CLIMB, and dynamic methods across pre-training and post-training.Pre-training uses 17 uncopyrighted Pile subsets and small proxy models trained on up to 1B tokens.
- Pre-training results: 48.2 average performance and rank 1 across 14 downstream benchmarks exceed CLIMB (47.5) and RegMix (47.2) in pre-training.The results are attributed to a single proxy model and gradient-based optimization, alongside reduced search overhead.
- Post-training results: 65.4 average performance and rank 1 across four post-training benchmarks exceed CLIMB (59.9) by 5.5 points.Although automated methods optimize using mathematics benchmarks, FASTMIX also performs best on LiveCodeBench and GPQA-Diamond.
- Limitations and practical considerations: Non-differentiable validation metrics make direct optimization unsuitable, while finite differences and SPSA rarely converge, especially on industrial datasets.The proposed response is to use a differentiable proxy objective.
- Limitations and practical considerations: Restricting the outer-loop horizon n_2 to 1 avoids intractable mixture-weight gradients and the memory, speed, and stability problems of BPTT or finite differences.This restriction yields a closed-form gradient and more stable, efficient optimization behavior.
- Limitations and practical considerations: Small surrogate models below 0.5B can produce unstable, noisy mixing ratios, while industrial search targets differ from the Pile-cc validation-loss setup.The passages caution that these issues may limit transfer from academic experiments to industrial development.
5 Conclusion
FASTMIX is an efficient framework for discovering data mixtures for large-model training by reformulating mixture selection as weighted bilevel optimization. This enables differentiable, joint optimization of model parameters and mixture coefficients using a single proxy model rather than hundreds.
- 5 Conclusion: FASTMIX discovers data mixtures efficiently for large-model training.The passage identifies efficient data-mixture discovery as the framework’s purpose.
- 5 Conclusion: A weighted bilevel reformulation makes mixture coefficients differentiable by equating sampling-ratio optimization with learning per-source loss weights.This equivalence follows from a reparameterization of mixture selection.
- 5 Conclusion: Joint gradient-based optimization of the model and mixture uses a single proxy model rather than hundreds.The differentiable formulation enables simultaneous optimization of both components while reducing the number of proxy models required.
6 Future Works
FastMix’s one-step, short-horizon outer-loop updates introduce greediness and sensitivity to data noise. Its optimization also reveals competitive, time-evolving source dynamics, with source importance changing across training stages.
- Limitations: FastMix’s one-step, short-horizon outer-loop updates introduce greediness and make the algorithm somewhat sensitive to data noise.This identifies a limitation and motivates exploring less greedy, longer-horizon update mechanisms.
- Future exploration: Data sources exhibit competitive, time-evolving search dynamics, with some vital early and others becoming dominant only after prolonged training.The observed shifts in source importance suggest that data-mixture optimization should account for training-stage dynamics.