Source-linked AI summary
Efficient Online Data Mixing For Language Model Pre-Training
Alon Albalak, Liangming Pan, Colin Raffel, William Yang Wang
TL;DR
Pretraining data strongly influences language-model performance, while existing selection and mixing methods can be costly or unable to adapt during training. ODM uses a multi-armed-bandit algorithm to optimize data-mixing proportions online from training loss. It reaches the next-best method's final perplexity with 19% fewer iterations and improves 5-shot MMLU performance by 1.9% relative accuracy with negligible wall-clock overhead.
Problem
Existing data-selection methods are computationally expensive, while fixed data-mixing proportions cannot adapt to changing pretraining dynamics.
Method
ODM models data domains as multi-armed-bandit arms and uses per-domain training loss as an efficient reward to update mixing proportions online.
Results
ODM reaches the final perplexity of the next-best method with 19% fewer iterations and improves 5-shot MMLU performance by 1.9% relative accuracy over DoReMi-50k.
Takeaways & Limitations
Online data mixing improves training efficiency and downstream performance while adding only negligible computational overhead.
Takeaways & Limitations
ODM's homogeneous micro-batches can cause higher initial perplexity, a weakness that can be alleviated by decreasing micro-batch size at added technical cost.
Abstract
from arXiv · showhide
The data used to pretrain large language models has a decisive impact on a model's downstream performance, which has led to a large body of work on data selection methods that aim to automatically determine the most suitable data to use for pretraining. Existing data selection methods suffer from slow and computationally expensive processes, a problem amplified by the increasing size of models and of pretraining datasets. Data mixing, on the other hand, reduces the complexity of data selection by grouping data points together and determining sampling probabilities across entire groups. However, data mixing proportions are typically fixed before training and therefore cannot adapt to changing training dynamics. To address these limitations, we develop an efficient algorithm for Online Data Mixing (ODM) that combines elements from both data selection and data mixing. Based on multi-armed bandit algorithms, our online approach optimizes the data mixing proportions during training. Remarkably, our method trains a model that reaches the final perplexity of the next best method with 19\% fewer training iterations, and improves performance on the 5-shot MMLU benchmark by 1.9% relative accuracy, while adding negligible wall-clock time during pretraining.
1 Introduction
Pretraining data strongly affects language-model performance, but existing data-mixing approaches either rely on costly model training or fixed sampling weights. ODM addresses these limitations by adapting domain proportions online with minimal overhead and achieves better training efficiency and downstream performance.
- Pretraining data can materially affect downstream results, with similarly sized models showing substantially different leaderboard scores despite nearly identical architectures and training processes.This motivates improving how pretraining data is selected and mixed.
- Data-mixing weights are often chosen heuristically or require training multiple models, and fixed weights cannot adapt to changing training dynamics.The cited limitations apply to The Pile and DoReMi-style approaches.
- ODM treats each data domain as a multi-armed-bandit arm and uses per-domain training loss as a reward to adapt sampling toward data with more information to learn.The reward requires no additional forward or backward passes beyond standard training.
- 19% fewer iterations reach the final validation perplexity of the next-best method, while 5-shot MMLU improves by 3% relative accuracy over the baseline.The method adds only 0.000007% computational overhead during pretraining.
2 Online Data Mixing (ODM)
ODM treats each data domain as a multi-armed-bandit arm and updates sampling proportions online using information-oriented rewards, allowing adaptation during pretraining.
- Online Data Mixing (ODM): ODM formulates online data mixing as a multi-armed bandit that updates domain sampling probabilities at every training iteration.Unlike traditional mixing, which fixes π(D_i) before training, ODM redefines it online to adapt to changing dynamics.
- Bandit adaptation: ODM uses an Exp3-style policy with decaying exploration and a moving-average reward so recent observations receive greater weight than older ones.The moving average is the main deviation from standard Exp3’s cumulative estimated reward.
- Reward function: The reward is the current training loss for a sampled dataset, favoring higher-perplexity data while adding no extra forward or backward passes.Because perplexity is the exponentiated loss, higher loss is treated as greater expected information gain.
- Algorithm: At each training turn, the policy samples datasets, accumulates their losses across gradient steps, updates model parameters, and then uses sampled-group rewards to update the next policy.A 1% warmup period keeps the policy stationary while the model trains, reducing early loss variance.
3 Experimental Setup
Experiments pretrain a 1-billion-parameter transformer on The Pile and compare ODM with static and DoReMi-based mixing weights using perplexity and 5-shot MMLU.
- Training: Experiments train a 1-billion-parameter decoder-only transformer on The Pile’s 22 domains, using 50 billion tokens and gradient accumulation across 8 GPUs.The total batch size is 480 sequences with sequence length 1024.
- Evaluation: Evaluation measures held-out perplexity for each domain and average performance, plus 5-shot multiple-choice accuracy across 57 MMLU tasks.Each MMLU task uses five in-context examples.
- Baselines: Baselines include The Pile’s original weights and DoReMi weights recalculated for the experimental 50k tokenizer, alongside the original 256k-tokenizer DoReMi weights.The tokenizer distinction matters because DoReMi weights may vary with tokenizer vocabulary.
4 Findings and analysis.
ODM improves training efficiency and downstream performance while adapting data mixing during training, although its sampling policy can produce early and domain-specific trade-offs.
- Main results: 19% fewer iterations reach the final perplexity of DoReMi-50k, while ODM’s final validation perplexity is 4.8% lower than TPW, 2.4% lower than DoReMi-50k, and 4.9% lower than DoReMi-256k.ODM also reaches TPW’s final performance with 30% fewer iterations, and Figure 1 reports the validation perplexity comparison.
- Main results: 3% higher 5-shot MMLU accuracy than TPW and 1.9% higher accuracy than DoReMi-50k demonstrate improved downstream performance.Table 1 reports average 5-shot MMLU accuracy for ODM and the baseline methods.
- Main results: DoReMi-256k has 0.7% worse average test perplexity than TPW, whereas DoReMi-50k improves over TPW by 2.6%, indicating weak transfer of DoReMi weights across models.Figure 3 compares held-out test perplexity and average perplexity across methods.
- Domain-level effects: Across domains, TPW is best on 7 and worst on 7, while DoReMi-50k is in the middle on 17/22 domains and best on only three.These buckets compare individual-domain test perplexity under different data-mixing objectives.
- Domain-level effects: ODM performs worst on two web-text domains, but this decreased performance does not appear to harm downstream performance.The paper identifies this as a domain-level trade-off rather than a demonstrated downstream failure.
- Sampling policy: ODM’s sampling distribution moves closer to uniform, with most domain ratios stabilizing by approximately 40000 iterations while Github continues decreasing and BookCorpus2 and HackerNews continue increasing.The largest increases are for PhilPapers, HackerNews, and BookCorpus2; the largest decreases are for Github, ArXiv, and PubMed_Central.
- Limitations: ODM begins with higher validation perplexity because homogeneous micro-batches create an early weakness that smaller batches could alleviate, although smaller batches reduce GPU utilization and slow wall-clock time.The paper suggests mixing domains within micro-batches during warm-up as a better solution.
5 Conclusion
The paper formulates online data mixing as a multi-armed bandit problem and uses information-gain rewards to train models with lower held-out perplexity in fewer iterations. Training loss makes the approach computationally efficient, adding only trivial wall-clock overhead.
- Conclusion: ODM formulates online data mixing as a multi-armed bandit problem and uses a reward designed to maximize information gain.The approach updates data mixing online rather than relying on fixed proportions.
- Conclusion: ODM achieves lower held-out perplexity in fewer iterations than baseline methods while adding 0.000007% additional wall-clock time to training.The computational efficiency comes from using training loss as the reward.
A Model Configuration
The experiments use a 1-billion-parameter language model with a 1024-token sequence length, 16 layers, hidden size 2048, and 16 attention heads.
- Model architecture: The model has 1 billion parameters, sequence length 1024, 16 layers, hidden size 2048, and 16 attention heads.It uses rotary positional embeddings and the GPT-NeoX-20B tokenizer.
- Optimization: Training uses FlashAttention, Adam, linear warmup for 1000 iterations, and cosine learning-rate decay from 2.5e-4 to 2.5e-5.The learning rate starts at 2.5e-5, warms to 2.5e-4, and decays back to 2.5e-5.