FastMix: Fast Data Mixture Optimization via Gradient Descent

Summary (Overview)

• FASTMIX is a novel framework that automates data mixture discovery for LLM pre-training and post-training while training only a single proxy model, drastically reducing computational cost compared to prior approaches.
• The core contribution is a reparameterization of mixture selection as a weighted bilevel optimization problem, showing that optimizing mixture ratios is mathematically equivalent to assigning per-source loss weights under uniform source sampling.
• This reformulation makes mixture coefficients fully differentiable, enabling joint, gradient-based optimization of both mixture weights and model parameters via standard optimizers like SGD or Adam.
• FASTMIX achieves state-of-the-art performance with dramatically lower search costs: in pre-training, it reaches 48.2 average score (rank 1 across 14 benchmarks) in just 1.3 GPU-hours—up to 550× faster than RegMix and 55× faster than CLIMB.
• In post-training (SFT), FASTMIX achieves 65.4 average score (5.5 points above the next best method) in 2.2 GPU-hours versus 115+ GPU-hours for CLIMB/RegMix, demonstrating strong cross-domain generalization.

Introduction and Theoretical Foundation

Background and Motivation

The performance of large-scale models depends critically on training data composition. While large and diverse datasets have driven recent advances, identifying the optimal data mixture for pre-training and post-training remains a significant open problem. Popular methods fall into two categories:

1. Manual trial-and-error approaches lack standardization and fail to generalize across settings.
2. Proxy-based methods (RegMix, CLIMB, DoReMi) train many proxy models under different candidate mixtures, incurring prohibitive computational overhead (hundreds of GPU-hours).

Problem Formulation

Data mixture optimization is formally posed as a bilevel optimization problem. Let $D = \{D_1, \ldots, D_k\}$ be a collection of data sources, and let $\alpha \in \mathcal{A} \subset \mathbb{R}^k$ denote mixture weights on the probability simplex ($\alpha_i \geq 0$, $\sum_{i=1}^k \alpha_i = 1$). The training objective under mixture $\alpha$ is $\mathcal{L}_{\text{train}}(D, w \mid \alpha)$, and the target is:

$$\min_{\alpha} \mathcal{L}_{\text{target}}\Big(w^{*}(\alpha)\Big) \quad \text{s.t.} \quad w^{*}(\alpha) = \arg\min_{w} \mathcal{L}_{\text{train}}\Big(D, w | \alpha\Big), \quad \sum_{i=1}^{k} \alpha_{i} = 1, \quad \alpha_{i} \geqslant 0.\tag{1}$$

The crux of the difficulty: unlike model parameters $w$, mixture sampling ratios are non-differentiable, precluding end-to-end backpropagation. Prior approaches resort to greedy heuristics or policy-gradient updates, which are sample-inefficient and scale poorly.

Key Theoretical Insight: Reparameterization

The central theoretical contribution is showing that mixture sampling is equivalent to per-source loss weighting under uniform sampling. Under mixture sampling (first draw source $i \sim \text{Cat}(\alpha)$, then sample $x \sim D_i$):

$$\mathcal{L}_{\text{train}}(D, w \mid \alpha) = \mathbb{E}_{i \sim \mathrm{Cat}(\alpha)} \mathbb{E}_{x \sim D_i}[\ell(x, w)] = \sum_{i=1}^{k} \alpha_{i} \mathcal{L}_{\text{train}}(D_i, w),\tag{3}$$

This yields the fully differentiable reformulation:

$$\min_{\alpha} \mathcal{L}_{\text{target}}\Big(w^{*}(\alpha)\Big) \quad \text{s.t.} \quad w^{*}(\alpha) = \arg\min_{w} \sum_{i=1}^{k} \alpha_{i} \mathcal{L}_{\text{train}}\Big(D_i, w\Big), \quad \sum_{i=1}^{k} \alpha_{i} = 1, \quad \alpha_{i} \geqslant 0.\tag{2}$$

Here $\alpha_i$ serves as a continuous weight scaling each domain's loss, making mixture weights fully differentiable and amenable to gradient-based optimization.

Methodology

Joint Optimization Objective

To improve generalization beyond validation-only search, FASTMIX incorporates two complementary strategies:

1. Entropy-based regularization to prevent mixture collapse: $\mathcal{R}_{\text{entropy}} = \sum_{i=1}^{k} \alpha_i \log \alpha_i$
2. Training loss as auxiliary target to balance validation and training signals

The joint search objective is:

$$\mathcal{L}_{\text{target}}(w) = \ell_{\text{val}}(w) + \beta \mathcal{L}_{\text{train}}(w) + \lambda \sum_{i=1}^{k} \alpha_{i} \log \alpha_{i},\tag{4}$$

where $\beta \geq 0$ and $\lambda \geq 0$ are trade-off hyperparameters. Empirically, $\lambda$ is set small (e.g., $10^{-5}$) and $\beta$ is most effective at moderate values (e.g., 0.1).

Alternating Optimization Procedure

FASTMIX implements an iterative procedure alternating between two steps:

(i) Inner loop (network parameter update) — Given current mixture weights $\alpha^t$, update model parameters for $n_1$ steps:

$$w^{t+1} \leftarrow w^{t} - \eta_{w}^{t} \frac{\partial\Big(\sum_{i=1}^{k} \alpha_{i}^{t} \mathcal{L}_{\mathrm{train}}(D_i, w^{t})\Big)}{\partial w^{t}},\tag{5}$$

(ii) Outer loop (mixture weight update) — Update mixture weights based on validation feedback:

$$\alpha^{t+1} \leftarrow \alpha^{t} - \eta_{\alpha}^{t} \frac{\partial \mathcal{L}_{\mathrm{target}}(w^{t+n_2})}{\partial \alpha^{t}},\tag{6}$$

Closed-Form Gradient (Key Result)

For the special case $n_2 = 1$ with SGD updates, the gradient admits a closed-form solution:

$$\frac{\partial \mathcal{L}_{\text{target}}(w^{t+1})}{\partial \alpha^{t}} = \frac{\partial \mathcal{L}_{\text{target}}(w^{t+1})}{\partial w^{t+1}} \cdot \frac{\partial w^{t+1}}{\partial \alpha_{i}^{t}} = -\eta_{w}^{t} \nabla_{w} \ell_{\text{val}}(V, w^{t+1}) \cdot \nabla_{w} \mathcal{L}_{\text{train}}(D_i, w^{t}),\tag{7}$$

Intuition: The gradient w.r.t. $\alpha_i$ is proportional to the alignment between the validation gradient $\nabla_w \ell_{val}(V, w^{t+1})$ and the training gradient from source $D_i$, \nabla_w \mathcal{L}_{t rain}(D_i, w^t}). If these gradients are aligned (positive dot product), the derivative is negative, so gradient descent increases $\alpha_i$, emphasizing sources that reduce validation loss. If opposed, $\alpha_i$ is decreased. Near-orthogonality yields small updates.

Algorithm

Algorithm 1: FASTMIX Optimization Algorithm
1: Initialize model parameters $w^0$, mixture weights $\alpha^0$, inner-loop duration $n_1$ and outer-loop duration $n_2$.
2: for $t = 0, 1, \ldots, T - 1$ do
3:    if $(t) \mod n_1 
eq 0$ then
4:    // Inner loop: update model parameters
5:    $w^{t+1} \leftarrow w^t - \eta_w^t \frac{\partial [\sum_{i=1}^k \alpha_i^t \mathcal{L}_{train}(D_i, w^t)]}{\partial w^t}$
6:    else
7:    // Outer loop: update mixture weights
8:    $\alpha^{t+1} \leftarrow \alpha^t - \eta_alpha^t \frac{\partial \mathcal{L}_{target}}(w^{t+n_2})}{\partial \alpha^t}$
9:    end if
10: end for
11: Output: the optimized mixture weight $a^{final}$ after the final outer loop update.

Empirical Validation / Results

Pre-training Experiments

Setup: Pile dataset (17 uncopyrighted subsets), 1M-parameter proxy models trained on up to 1B tokens, search target = Pile-cc validation loss. Evaluation: 1B-parameter model trained on 25B tokens, tested on 14 downstream benchmarks.

Key results:

• FASTMIX achieves the highest average performance (48.2) and best average rank (1) across all 14 benchmarks, leading on 9 of 14 individual tasks.
• Search efficiency: 1.3 GPU-hours vs. 71.9 (CLIMB) and 720.5 (RegMix)—a 55× and 550× speedup, respectively.

Post-training (SFT) Experiments

Setup: Qwen2.5-Math-Instruct 7B model, 8 domains (Math, Code, Dialogue, STEM). Search objective: 1:1 weighted sum of GSM8K and gaokao2023en scores. Proxy models: Qwen2.5-1.5B-Instruct (~1B parameters) due to the absence of very small proxies in post-training. Evaluation: MATH, AIME-24, LiveCodeBench-v2, GPQA-Diamond.

Key results:

• FASTMIX achieves 65.4 average score, a 5.5-point lead over the next best method (CLIMB, 59.9).
• Exceptional generalization: Despite optimizing only on math benchmarks (GSM8K, gaokao2023en), FASTMIX achieved the best performance on LiveCodeBench (coding) and GPQA-Diamond (STEM QA)—demonstrating avoidance of overfitting to the optimization signal.
• Search cost: 2.2 GPU-hours vs. 115.9/117.4 for RegMix/CLIMB (a 52× reduction).

Theoretical and Practical Implications

Theoretical Contributions

1. Equivalence theorem: The reparameterization formally establishes that mixture sampling and per-source loss weighting are mathematically equivalent, bridging discrete sampling ratios with continuous differentiable weights.
2. Closed-form gradient: The derivation of Eq. (7) provides an interpretable, computationally efficient gradient that reveals the mechanism of mixture optimization—reallocating mass toward sources whose gradients align with validation improvement.
3. Generalizable framework: The bilevel formulation is applicable across both pre-training and post-training stages, with different proxy model scales.

Practical Implications

1. Scalability: FASTMIX's single-proxy-model approach makes mixture optimization feasible at scales where training hundreds of proxies (RegMix/CLIMB) is prohibitive, particularly in post-training with large (1B+) proxy models.
2. Cross-domain transfer: The method discovers mixtures that generalize beyond the optimization target, suggesting fundamental capability improvement rather than benchmark overfitting.
3. Industrial lessons: The authors share practical insights from industrial development:
   • Non-differentiable metrics: Use differentiable proxies (e.g., SFT loss) instead of black-box estimators (SPSA, finite differences), which are unstable and inefficient.
   • Regularization: Entropy regularization works on clean academic datasets, but industrial robustness is better achieved via strict oversampling ratio constraints (capping up-sampling at 3× original size).
   • Proxy model scale: Small proxies (<0.5B) suffer convergence instability and data-source biases; caution is advised.
   • Search target data: Pre-training and SFT data length discrepancies can cause gradient divergence; concatenating SFT sequences to match pre-training length mitigates this.

Conclusion

FASTMIX introduces an efficient framework for data mixture discovery in large-model training. The key contribution is a weighted bilevel reformulation of mixture selection: via reparameterization, optimizing sampling ratios becomes equivalent to learning per-source loss weights, making mixture coefficients differentiable. This enables joint, gradient-based optimization of both model and mixture using a single proxy model rather than hundreds. Across pre-training and post-training, FASTMIX delivers superior accuracy with orders-of-magnitude lower search cost, making data mixture optimization practical, scalable, and robust for next-generation LLMs.

Future Works

The authors identify several limitations and directions for future exploration:

1. Greediness: The current one-step, short-horizon outer-loop update introduces a degree of greediness, making the algorithm sensitive to data noise.
2. Time-evolving source dynamics: Observed search dynamics reveal that data sources exhibit competitive, time-evolving relationships—some vital early, others dominant after prolonged training. This offers insights into data curriculum design.
3. Extension to data source attribution: FASTMIX could be extended beyond data mixing to serve as a framework for data source attribution, inviting community collaboration.