Scaling Laws for Optimal Data Mixtures

TL;DR

This paper introduces a universal scaling law framework that predicts model performance based on data mixture and size, enabling optimal domain weight selection without extensive trial-and-error in large-scale model training.

Contribution

It develops a systematic, predictive approach to determine optimal data mixtures for training large models across multiple domains using scaling laws.

Findings

Scaling laws accurately predict loss across different model sizes and data mixtures.

The laws generalize to new data mixtures and scales with minimal small-scale training.

Optimal domain weights can be derived efficiently, reducing trial-and-error in data selection.

Abstract

Large foundation models are typically trained on data from multiple domains, with the data mixture--the proportion of each domain used--playing a critical role in model performance. The standard approach to selecting this mixture relies on trial and error, which becomes impractical for large-scale pretraining. We propose a systematic method to determine the optimal data mixture for any target domain using scaling laws. Our approach accurately predicts the loss of a model of size $N$ trained with $D$ tokens and a specific domain weight vector $h$. We validate the universality of these scaling laws by demonstrating their predictive power in three distinct and large-scale settings: large language model (LLM), native multimodal model (NMM), and large vision models (LVM) pretraining. We further show that these scaling laws can extrapolate to new data mixtures and across scales: their parameters can be accurately estimated using a few small-scale training runs, and used to estimate the performance at larger scales and unseen domain weights. The scaling laws allow to derive the optimal domain weights for any target domain under a given training budget ($N$,$D$), providing a principled alternative to costly trial-and-error methods.