PaperlayerView exact revision on arXiv

Source-linked AI summary

Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance

Jiasheng Ye, Peiju Liu, Tianxiang Sun, Jun Zhan, Yunhua Zhou, Xipeng Qiu

arXiv:2403.16952v2cs.CLcs.AIcs.LG

TL;DR

LLM pretraining requires choosing proportions across interacting data domains, but existing heuristics do not provide reliable advance estimates of mixture effects. The paper fits data mixing laws and nests them with scaling laws to predict large-scale performance from small experiments, achieving a RedPajama result comparable to 48% more default-mixture training steps. Its scope includes non-monotonic mixture effects and practical fitting concerns such as implicit-domain complexity.

  • Problem

    Existing heuristic and qualitative methods make it difficult to predict how candidate data mixtures will affect model performance before training.

  • Method

    The paper fits functional data mixing laws on sampled mixtures and nests them with training-step and model-size scaling laws to predict target-scale performance.

  • Results

    48% more steps: the optimized RedPajama mixture achieved performance comparable to the default mixture for a 1B model trained on 100B tokens.

  • Takeaways & Limitations

    Data mixing laws enable low-cost quantitative tuning of unseen mixture proportions and can inform capability balancing and continual-training data schedules.

  • Takeaways & Limitations

    Mixture effects can be non-monotonic, and implicit-domain extensions raise concerns about increasing fitting samples and choosing the number of implicit domains.

Abstract

from arXiv · show

Pretraining data of large language models composes multiple domains (e.g., web texts, academic papers, codes), whose mixture proportions crucially impact the competence of outcome models. While existing endeavors rely on heuristics or qualitative strategies to tune the proportions, we discover the quantitative predictability of model performance regarding the mixture proportions in function forms, which we refer to as the data mixing laws. Fitting such functions on sample mixtures unveils model performance on unseen mixtures before actual runs, thus guiding the selection of an ideal data mixture. Furthermore, we propose nested use of the scaling laws of training steps, model sizes, and our data mixing law to enable predicting the performance of large models trained on massive data under various mixtures with only small-scale training. Moreover, experimental results verify that our method effectively optimizes the training mixture of a 1B model trained for 100B tokens in RedPajama, reaching a performance comparable to the one trained for 48% more steps on the default mixture. Extending the application of data mixing laws to continual training accurately predicts the critical mixture proportion that avoids catastrophic forgetting and outlooks the potential for dynamic data schedules

1 Introduction

LLM pretraining mixes multiple domains whose proportions affect model capabilities, yet ideal mixtures remain difficult to select with existing heuristic or qualitative methods. The paper introduces data mixing laws and a scaling-law pipeline to predict unseen-mixture performance before expensive target-scale training.

  • Multiple-domain pretraining data can exhibit facilitation, independence, or conflict, making mixture proportions important for balancing model capabilities.
  • Existing approaches commonly use heuristics or qualitative optimization, making their effects difficult to predict before the actual training run.
  • Data mixing laws model validation loss as a function of mixture proportions, enabling prediction on unseen mixtures from evaluated samples.
  • The proposed pipeline nests scaling laws for training steps and model sizes with data mixing laws to predict large-scale performance from affordable experiments.
  • A RedPajama 1B model trained on 100B tokens reached performance comparable to the default mixture trained for 48% more steps, while continual-training experiments predicted proportions avoiding catastrophic forgetting.

2 Background

The paper frames mixture optimization within autoregressive language-model pretraining and neural scaling laws. It seeks a quantitative loss function over domain proportions that can identify a mixture achieving minimum validation loss.

  • Pretraining samples domain indices according to mixture proportions, then trains an autoregressive language model through next-token prediction.
  • Model evaluation computes loss on validation data after training on the sampled multi-domain dataset.
  • Neural scaling laws describe how loss varies with factors such as model size, training data, training steps, and computation through power laws.
  • The study formalizes mixture optimization as predicting loss for any domain-proportion vector under fixed model size and training steps.

3 The proportions of data mixtures influence model losses in a quantitatively predictable way

The paper finds that model losses vary with data-mixture proportions through predictable functional relationships, despite multi-variable and nonmonotonic challenges. These data mixing laws accurately estimate unseen mixtures and support aggregation across domains, including implicit validation domains.

  • Data mixing laws describe predictable functional relationships between model losses and domain-mixture proportions.The study frames these relationships as the central object of analysis.
  • Multi-domain mixtures introduce K −1 proportion variables, while nonmonotonic losses complicate function-form identification.A lopsided mixture may minimize loss, so the relationship need not decrease monotonically with any single domain proportion.
  • 3.1 Pilot Study on Domain Losses under Two-domain Mixtures: 70M and 160M models trained on five Github–Pile-CC mixtures showed log-scale domain losses linearly related to domain proportions after subtracting a shared constant.The models were trained for 30k steps, totaling 30B tokens, and evaluated separately on Github and Pile-CC.
  • 3.2 Extension to Domain Losses Trained on Multi-domain Mixtures: For multiple domains, candidate functions were constrained by compatibility with the two-domain law and symmetry across domain variables.These principles produce exponential terms for each training domain combined through commutative operations.
  • 3.2 Extension to Domain Losses Trained on Multi-domain Mixtures: M4 provided reliable predictions with fewer coefficients than M1, and Eqn. 7 accurately estimated unseen domain-loss samples.Overall validation loss is obtained by explicitly aggregating domain losses according to validation-domain proportions.
  • 3.3 Predicting Language Modeling Performance of Any Validation Mixture: Domain interactions include unrelated, mutually facilitative, and conflicting relationships, while sparse parameters reduce sample requirements below quadratic growth.Implicit aggregation performs better when the assumed number of implicit domains is at least the actual number, and regularization helps exploit sparsity.

4 Nested scaling laws predict losses trained on various mixtures using only small-scale experiments

Nested scaling laws use small-scale experiments to predict losses for large models trained on massive data across mixtures. The resulting pipeline effectively optimizes a 1B RedPajama model’s mixture, while mixture rankings can vary with scale.

  • Motivation: Fitting data mixing laws directly at target scales is prohibitively expensive because it requires many models across mixtures for each target size and dataset.Small-scale scaling-law extrapolation addresses this cost.
  • Pipeline: The pipeline first extrapolates training-step losses, then model-size losses, and finally predicts losses for target mixtures using the data mixing law.It trains smaller models on candidate mixtures and nests the resulting scaling laws.
  • Pipeline: The method uses decomposed power laws because directly fitting the complete two-variable Chinchilla law was empirically unstable.The decomposition separately models training steps and model size under fixed conditions.
  • Experiment: Mixture rankings can change with model size and token count, so an optimum found at experimental scales may be suboptimal at the target scale.This motivates predicting directly at the intended scale rather than transferring a small-scale optimum unchanged.
  • Experiment: The optimized RedPajama mixture matches default-mixture performance using 73% of the steps and reaches performance requiring 48% more default-mixture steps.The experiment targets a 1B model trained on 100B tokens and validates on the Pile.
  • Experiment: The optimized mixture reaches the lowest losses among mixtures tested at the same model size and data budget, with plausibly accurate loss predictions.Actual-run rankings were consistent with the predicted rankings.

5 Application to Continual Pretraining Outlooks Data Schedule Driven by Data Mixing Laws

Data mixing laws also apply to continual pretraining, where models begin from pretrained parameters and combine original data with a target domain. They predict a target-data proportion that preserves original-domain performance while improving the target domain.

  • Application: Continual pretraining starts from pretrained parameters rather than random initialization and can inject updated knowledge into an existing model.The paper studies a mixture of original pretraining data and upcoming target-domain data.
  • Application: Training Pythia-70M on four Pile–Python mixtures for 10B tokens tests whether the data mixing law fits continual-pretraining losses.The fitted law is evaluated on losses for both domains.
  • Mixture optimization: A too-large target-data proportion can hurt performance on the original domain during continual pretraining.The relevant original-domain performance is measured by losses on the Pile.
  • Mixture optimization: The predicted critical proportion preserves performance on the original domain while improving performance on the target domain.Figure 10 reports similar performance to the pre-continual-training model on the Pile.
  • Outlook: Successful continual-training predictions suggest a possible basis for designing dynamic data schedules across multiple training stages.The paper frames infinitely many stages as the limiting case of a dynamic schedule.

6 Related Work

Related work situates the paper within efforts to curate pretraining data and characterize neural scaling behavior. These lines of work motivate optimizing both data quality and the effects of controllable training factors.

  • Curating pretraining data: Pretraining-data curation seeks high-quality datasets through source selection, deduplication, and document filtering.Existing filters use features ranging from surface characters to semantic similarity.
  • Scaling laws: Scaling laws relate model performance to controllable factors such as model size and the amount of training data.They are studied as functional relationships describing scaling behavior.

7 Discussions

The paper frames data mixing laws as a quantitative way to predict losses for unseen training mixtures and reduce the cost of mixture tuning. It connects this framework to broader efforts to curate pretraining data and characterize scaling behavior.

  • Data mixing laws quantify how model losses vary with training-mixture proportions.
  • The framework estimates performance on unseen mixtures before training, enabling lower-cost data-mixture tuning alongside scaling laws.
  • Pretraining-data curation selects and filters multiple domains because pretraining strongly affects LLM knowledge and computational cost.

B Limitations and Discussions

The discussion identifies unresolved limitations involving domain definitions, prediction errors, joint scaling laws, dynamic curation, and theoretical understanding. These boundaries motivate more sophisticated quantitative approaches to data engineering.

  • On the clarification of domains: The predefined training domains are not well-defined, while implicit validation-domain aggregation is less interpretable and may raise accuracy concerns.
  • On the error analyses: Experiment design affects fitting quality, and nested scaling laws can accumulate errors across prediction steps.
  • On joint laws of multiple factors: A joint law covering training steps, model size, and mixture proportions could reveal factor synergies and reduce the number of fitting samples through coefficient sharing.
  • On dynamic data curating: The current pipeline selects fixed mixture proportions, leaving dynamic proportions and curricula as future extensions.
  • On theoretical understandings: The data mixing laws are empirical findings whose theoretical grounding in training dynamics remains open.

C The ranking of data mixtures depend on model sizes and training steps.

Mixture rankings vary with model size and training duration, so a mixture that performs well at small scale may not remain best at large scale. Across 20 mixtures, only partial ranking agreement persists.

  • Relative performance fluctuates across model sizes and training steps despite a relatively consistent overall trend.
  • A mixture that performs better at small scales does not always perform better at large scales.
  • The longest common partial-order sequence spans 10 mixtures across model sizes and 11 across training steps.
  • Figure 11 compares rankings for 20 RedPajama mixtures validated on the Pile, varying model size in panel A and training steps for 70M models in panel B.

D.1 Model Training

The experiments use small-scale models and sampled mixtures to fit nested scaling laws and data mixing laws, then extrapolate performance across training steps, model sizes, and mixture proportions. The reported validations support the feasibility of this prediction pipeline.

  • The pipeline nests power laws for training steps and model sizes to predict language-modeling performance at larger scales.
  • Scaling laws of training steps: Training-step scaling laws fit 70M models through 30k steps and extrapolate performance to 100k steps with mean absolute error 0.02.
  • Scaling laws of model sizes: Model-size scaling laws fit 70M, 160M, and 305M models and predict 410M performance with mean absolute error 0.003.
  • Sampling mixture proportions: Fitting mixture laws requires sampled mixtures that cover the proportion space; the method samples candidates by distributing losses approximately evenly.
  • Implicit domain aggregation: The mixture law predicts domain losses in one layer and aggregates them into overall validation loss in a second layer.
  • Connections between implicit domain aggregation and MLP: The mixture-law computation graph is equivalent to an MLP with exponential activation and softmax output aggregation.

F.1 Prediction Results on More Domains

Data mixing laws generalize beyond the main experimental domains. On Wikipedia, ArXiv, and StackExchange mixtures, predicted and observed losses were consistent.

  • Experimental setup: The evaluation uses Wikipedia, ArXiv, and StackExchange, three domains different from those in the main experiments.The models are 70M models trained for 10k steps.
  • Prediction results: Predicted and observed losses are consistent across different mixtures.This result is reported in Figure 15 for domain-loss prediction using seven fitting mixtures and three validation mixtures.
  • Conclusion: The results confirm that data mixing laws also work on domains beyond the main experiments.

F.2 Data Mixing Laws Make no Domain-Independent Assumptions

The data mixing law models interactions among domain proportions rather than assuming independent domain effects, while nested scaling laws reduce the cost of predicting large-model performance. Experiments report accurate extrapolation across training steps and model sizes and estimate substantial computation savings.

  • Domain interactions: Linear combinations of domain proportions enter the exponent, producing interaction terms between different mixture proportions.Thus, the formulation does not make a domain-independent assumption that domains affect losses independently.
  • Domain interactions: Correlated-domain experiments mix Wikipedia with CommonCrawl and ArXiv to test whether the law models overlapping-domain effects.The two synthetic domains each contain 50% Wikipedia and 50% of another source.
  • Nested scaling-law pipeline: Nested scaling laws train multiple mixtures on smaller models and shorter runs, then predict target performance at larger model sizes and training durations.The computation depends primarily on how far the scaling laws extrapolate accurately.
  • Computation: State-of-the-art scaling-law predictions can reduce the relevant computation fraction to 1/100–1/1000, implying an expected extra training cost of 1/5–1/50 with 20 mixtures.
  • Scaling-law validation: 40k steps accurately predict performance through 250k steps, a 6.25x extrapolation, while sub-100M models accurately predict performance up to 7.25B parameters, a 72.5x extrapolation.Combining the training-step and model-size scaling laws may achieve 450x scaling.
  • Baselines and mixture optimization: The comparison evaluates DoGE and DoReMi mixtures on RedPajama with validation on the Pile, including DoGE’s universal and out-of-distribution settings.DoGE’s out-of-distribution mixture is obtained by training a 160M proxy model for 10k steps and averaging proportions over training.

G Loss Prediction Results with Nested Scaling Laws

Nested scaling laws provide plausible predictions of overall and domain-specific validation losses for larger models and longer training. The optimized mixtures generally improve performance across domains while preserving the ability to target different capabilities.

  • Prediction results: Nested scaling-law predictions provide plausible relative loss estimates for overall validation data and different validation domains.The predictions extrapolate from fitting data using 70M–410M models trained for 30B tokens to 1B models trained for 100B tokens.
  • Prediction results: Optimized mixtures perform better in most validation domains.
  • Capability-specific evaluation: Overall loss supports optimizing aggregate performance, while domain losses expose model capabilities across different aspects.
  • Capability-specific evaluation: Tuning data mixtures can improve specific model capabilities without sacrificing others.
Loading 2403.16952v2…