Scaling Laws for Optimal Data Mixtures

Motivation

• Domain weights determine each domain’s training proportion and significantly affect model performance, but are typically chosen through ad-hoc trial and error.The introduction identifies the lack of a principled selection method as a gap.Source: §1 · p. 2

Contribution

• The laws model loss as a predictable function of domain weights interacting with model parameters and training tokens, validated in LLM, NMM, and LVM pretraining.Experiments reached 7B parameters and 150B tokens for LLMs, 8B parameters and 160B tokens for NMMs, and 1B parameters with 330B tokens for LVMs.Source: §1 · p. 2

Scaling laws that extrapolate

• Fitted scaling laws extrapolate from small-scale runs to accurately predict large-model loss under unseen domain weights.For NMMs, they quantify how increasing text data reduces text loss.Source: §1 · p. 2

Optimal domain weights estimation

• Minimizing the fitted loss estimate yields optimal domain weights, providing a principled alternative to trying different weights and selecting the best.The paper illustrates this approach with the average loss of NMMs.Source: §1 · p. 2

2.1 Problem setup

• For a target domain that may differ from the training domains, the goal is to predict its loss after training a model of size N on D tokens with domain weights h.The weights are positive and sum to one; the framework covers text, image, and multimodal domains.Source: §2.1 · p. 3

2.2 Scaling laws derivation

• The additive law makes only its bias term depend on h, whereas the joint law additionally models scale-related terms as functions of h while keeping αh and βh constant.The additive law has 5+2k parameters; the joint law has 5+4k parameters.Source: §2.2 · pp. 3–4
• The joint law models interactions among model size, token count, and domain weights, making optimal weights compute-dependent; the additive law’s optimum is independent of size and tokens.The joint law can recover the additive law, and fitted training-run error is always lower for the joint law, though it often reduces estimation error significantly in practice.Source: §2.2 · p. 4

2.3 Fitting the scaling laws

• To fit laws with up to 5+2k or 5+4k parameters, the method uses random initial-parameter search and Basin-hopping with L-BFGS rather than repeated L-BFGS initializations.For the multimodal Interleaved-domain example, Basin-hopping converges faster, requiring fewer L-BFGS calls to reach low fitting loss.Source: §2.3 · p. 4, Figure 2 · p. 4

3.1 Pretraining of large language models (LLMs)

• LLM experiments use Llama-style transformers from 186M to 7B parameters and seven unfiltered SlimPajama domains; smaller analyses also use GPT2-style models and up to eight Pile domains.The Llama-style setup uses rotary positional embeddings, SwiGLU, and RMSNorm; GPT2-style models range from 90M to 3B parameters.Source: §3.1 · p. 5

3.2 Pretraining of native multimodal models (NMMs)

• NMMs use an early-fusion transformer that interleaves text and image tokens, trained on text-only, interleaved multimodal, and paired image-caption data.Images are resized to 224×224 and patchified at 14×14; data sources include DCLM, Obelics, DFN, COYO, and HQITP (Shukor et al., 2025).Source: §3.2 · p. 5

3.3 Pretraining of large vision models (LVMs)

• LVM experiments follow AIMv2’s late-fusion vision-encoder and multimodal-decoder recipe, training autoregressively on image and text tokens across four image-caption domains.The domains comprise noisy Internet alt-text, two high-quality datasets with 134 million and 400 million samples, and synthetic recaptioned data.Source: §3.3 · p. 5
• Figure 3 compares observed and predicted average loss across LLM, NMM, and LVM pretraining, with scaling laws fitted on small models and extrapolated to larger ones.Per-domain MRE is reported in Table 2.Source: Figure 3 · p. 6

3.4 Implementation details

• Model scale is varied by changing transformer hidden dimension while holding layer count fixed; most experiments use a constant learning-rate schedule to sample multiple token counts per run.The authors also validate extrapolation with cosine learning-rate scheduling in section 6.Source: §3.4 · p. 6

Experimental setup

• The laws are fitted on small models trained with few tokens and validated on larger models with more tokens across LLM, NMM, and LVM pretraining.The experiments use 7 text domains for LLMs, 3 multimodal domains for NMMs, and 4 domains for LVMs.Source: §4 · p. 6

Results

• Predicted losses closely match observed losses for both joint and additive laws, and mean relative errors are consistently low, with the joint law improving over the additive law.The comparison covers larger model sizes; Table 2 reports MRE results.Source: §4 · pp. 6–7, Table 2 · p. 7
• MRE varies across LLM domains, ranging from 0.31% on C4 to 4.45% on Wikipedia.Source: §4 · p. 7

FLOPs count

• One small-scale run costs approximately 6ND FLOPs, and the reported large-scale runs have comparable computational cost.The paper notes that large-scale runs could use far more tokens, increasing their cost.Source: §4 · p. 7

Optimal domain weights estimation

• After fitting a scaling law, the method selects simplex-constrained domain weights to minimize target loss using mirror descent; multiple targets can be combined through a weighted average loss.Task importance can be represented with importance weights in the summed objective.Source: §5 · p. 7
• Optimal weights are scale-independent under the additive law but scale-dependent under the joint law, while small-scale fits can guide mixture choices for a target model size and token budget.The paper presents this as a principled alternative to ad hoc mixture selection.Source: §5 · pp. 7–8

NMM results

• Estimated NMM mixtures consistently outperformed uniform, prior-work (Shukor et al., 2025), and sampled mixtures; additive and joint laws performed similarly, and mixtures generalized to larger models.The authors describe additive laws as a more practical baseline because they use the same optimal mixture across runs.Source: §5 · pp. 8–9

LLM results

• On downstream evaluations, the OpenHermes-loss-optimized 7B model was overall better than the other models; OpenHermes and training-domain losses favored their respective optimized weights.The Table 3 caption reports that the model optimized for OpenHermes performs best; Figure 4 reports the losses on OpenHermes and the training domains.Source: §5 · p. 8, Figure 4 · p. 8, Table 3 · p. 9

LVM results

• On AIMv2, the 1B model trained with weights optimized for average loss across four domains performed better than the uniform-weight model.The weights were estimated from fitted scaling laws.Source: §5 · p. 9

• Additional LLM experiments use the Pile dataset with the number of domains varying from k = 4 to k = 8.Source: §6 · p. 9

Optimal domain weights behavior when scaling FLOPs

• As D increases, interleaved data becomes less important, while bigger models rely more on text; the additive law captures average behavior across scales.These optimal-weight trends are predicted by the joint scaling law for multimodal data.Source: Figure 5 · p. 9, §6 · p. 10

Other scaling laws formulas

• Table 4 reports average MRE: additive reduces MRE versus simple especially for LLM, full extrapolates worse than joint despite lower train MRE, while for LVM simple beats additive and full improves on joint.The full law is more expressive than the joint law and adds 2k +1 parameters.Source: Table 4 · p. 10, §6 · pp. 10–11

Only 10-20 runs are needed to fit the scaling laws

• About 10 runs suffice for NMM and 4-domain LLMs, versus about 20 for 6- and 8-domain LLMs; with very few runs, the additive law has slightly lower eval MRE.For 6- and 8-domain LLMs, the fit uses a single model size and omits dependence on N.Source: §6 · pp. 9–10

Cosine learning rate scheduler

• With cosine decay, on the 1B model, additive law averages 0.76% MRE and joint law 0.54%; estimated optimal weights resemble those from a constant learning rate.For the additive law, the constant-rate optimal weights are [0.34, 0.17, 0.32, 0.17].Source: §6 · p. 10

Asymptotic behavior

• Asymptotically, the bias term captures target-distribution entropy and the shift KL(p∥q(h∗)), while the expected log-likelihood ratio measures distance from the model trained on the optimal mixture.The passage frames the loss as cross-entropy between the target distribution p and the model output q(h).Source: §6 · p. 11

Scaling laws

• Scaling-law studies span language, vision, diffusion, sparse experts, and multimodal models, but Shukor et al. (2025) analyze native multimodal models with a fixed pretraining mixture.Earlier work also examines parameter allocation, model width and depth, distillation, sparsity, expert counts, and routing strategies.Source: §7 · p. 11

Scaling laws for data mixtures

• Ye et al. (2024) and Ge et al. (2024) model loss versus h at fixed (N, D), unlike this work’s joint law over (N, D, h).Their approaches generally focus on a single modality and models below 1B parameters; for fixed (N, D), this work reports better extrapolation to unseen mixtures (section B).Source: §7 · p. 11

Data mixture selection

• Data-mixture selection commonly relies on costly trial and error, while alternatives use heuristics, small models to predict performance, or auxiliary models to rank training data.Heuristics adjust ratios by domain data size or target-task distribution; small predictors take the mixture as input, and auxiliary models select high-quality data.Source: §7 · pp. 11–12

Limitations

• The study focuses on pretraining and assumes no data repetition and fixed mixture weights, limiting its coverage of scarce sources and dynamic curricula.The law predicts generic target loss, which is known to correlate with downstream task performance; future work could predict that performance directly and address continual pretraining and finetuning.Source: §8 · p. 12

Broader impact

• A few small-scale runs can yield useful mixture coefficients for larger models, reducing pretraining cost, energy use, and CO2 emissions.The authors note that extensive trial-and-error mixture selection can be prohibitively expensive and that better models may result in the long run.Source: §8 · p. 12

Conclusion

• Small-scale optimal mixture coefficients transfer to larger models and budgets, significantly improving over naive grid search across language, multimodal, and vision pretraining.The proposed law predicts arbitrary target-domain loss from mixture coefficients and compute budget (N, D).Source: §8 · p. 12

A.1 LLM pretraining

• LLM experiments use LLAMA-style architectures; Pile-domain analyses use Brown et al. (2020) architectures with 24 layers and varying latent dimensions across model scales.Hyperparameters are described in Tables 5 and 7.Source: §A.1 · p. 17

A.2 Multimodal pretraining

• Multimodal models range from 100M to 8B parameters, scale width at a constant 24-layer depth, and use causal attention for text and bidirectional attention for images.The implementation follows Shukor et al. (2025) and uses the L3M code base (El-Nouby* et al., 2025).Source: §A.2 · p. 17
• Multimodal training uses fully decoupled AdamW with β_1 = 0.9, β_2 = 0.95, weight decay of 1 × 10−4, and model-size-adjusted learning rates determined empirically.Learning rates generally decrease as model size increases.Source: §A.2 · p. 17
• Each multimodal batch has 2,000 samples and 2 million tokens with a 1,024-token context; training clips gradients at 1.0 and uses 1,000 warmup iterations followed by a constant learning rate.The constant schedule reduces the number of experiments.Source: §A.2 · p. 17
• Vision inputs use (14, 14) patches with Random Resized Crop at 224px and scale range [0.4, 1.0], plus Random Horizontal Flip at 50%..Source: §A.2 · p. 17
• Training uses bfloat16, FSDP (Zhao et al., 2023), activation checkpointing, gradient accumulation, and sequence packing to minimize padding in image-captioning data..Source: §A.2 · p. 17
• Evaluation uses held-out interleaved Obelics, image-caption HQITP, and text-only DCLM subsets, following Hoffmann et al. (2022), Aghajanyan et al. (2023), and Abnar et al. (2025).The setup assesses generalization across diverse data types.Source: §A.2 · p. 17

B Detailed extrapolation results

• Detailed per-domain and per-model-size MREs are reported for each experiment and scaling law.The passage also names the laws of Ye et al. (2024) and Liu et al. (2024).Source: B Detailed extrapolation results · p. 18, §B.1 · p. 18

B.1 Comparison to

• Ye et al. (2024) model loss as a function of h at fixed N and D, while Liu et al. (2024) model it as linear in h.The paper contrasts these forms with its own scaling law, which takes N, D, and h as inputs.Source: §B.1 · p. 18
• With 25 mixtures used for fitting at N = 200m and D = 8B tokens, our formula gives systematically better training errors and usually better test errors on the remaining 59 mixtures.The authors report difficulty fitting the M3 law, which Ye et al. (2024) also report underperforms.Source: §B.1 · p. 18

B.2 Comparison to Ge et al. (2024)

• Ge et al. (2024) predict loss on domain i using its own weight h_i and token count D, but not the weights of other domains.The authors argue this assumes all other domains have the same effect, although some may help the task and others may not.Source: §B.2 · p. 21
• The Ge et al. comparison fits laws on 25 mixtures from the LLM experiment at a fixed model size of N = 200M.The comparison evaluates MRE on the remaining test set and compares against the proposed additive law at fixed model scale.Source: §B.2 · p. 21
• Our law achieves significantly lower training and testing MRE than Ge et al. (2024)’s law in the LLM comparison.The authors attribute the improvement to accounting for all other domains.Source: §B.2 · p. 21

• The objective accommodates arbitrary target-domain weights w and is optimized with mirror descent; when training and target domains coincide, it maps w to h*(w, N, D).Source: §C.1 · p. 21

Behavior at the corners

• For a 1.3B model trained on 10B tokens, pure-domain targets typically favor the corresponding simplex corner, supporting h*_i≈1 when the target is D_i and domain overlap is limited.Source: §C.1 · p. 22, Figure 7 · p. 23

Accumulation point of asymptotes

• As model size N or data D tends to infinity, optimal-mixture asymptotes reach an accumulation point that depends on their relative growth rates.For multimodal data, proportional and square-root scaling show little difference, while holding either quantity constant or growing it logarithmically changes the accumulation point.Source: §C.1 · p. 22, Figure 10 · p. 25

Optimal weights for the additive scaling law

• Table 15 reports optimal LLM domain weights; for OpenHermes, the predicted Wikipedia weight of 0 was artificially set to 1%.Source: §C.1 · p. 21, Table 15 · p. 23

Fixed-points

• On the 4-domain Pile, target mixtures generally differ from their optimal training mixtures; no fixed points appear except at simplex corners, suggesting direct training on the desired mixture is usually suboptimal.Source: Figure 9 · p. 24, Figure 8 · p. 24