Efficient Online Data Mixing For Language Model Pre-Training

2023/12/05 by Alon Albalak, Albalak, Alon, Liangming Pan +5 · 30 citations
Computer Science · #Computation and Language (cs.CL) #FOS: Computer and information sciences #Machine Learning (cs.LG) #Machine Learning and Data Classification #Natural Language Processing Techniques #Topic Modeling

paper · pdf · doi:10.48550/arxiv.2312.02406

openalex publication_date 2023/12/05 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

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.

Cited by

• A Data-Centric Approach to Generalizable Speech Deepfake Detection
• ADAPT: Learning Task Mixtures for Budget-Constrained Instruction Tuning
• Fast Data Attribution for Text-to-Image Models
• Sampling and Loss Weights in Multi-Domain Training
• Revisiting Multilingual Data Mixtures in Language Model Pretraining
• Shift is Good: Mismatched Data Mixing Improves Test Performance
• DaMo: Data Mixing Optimizer in Fine-tuning Multimodal LLMs for Mobile Phone Agents
• BLISS: A Lightweight Bilevel Influence Scoring Method for Data Selection in Language Model Pretraining
• Interactive Training: Feedback-Driven Neural Network Optimization
• Learning to See Before Seeing: Demystifying LLM Visual Priors from Language Pre-training
• MixtureVitae: Open Web-Scale Pretraining Dataset With High Quality Instruction and Reasoning Data Built from Permissive-First Text Sources
• Train Once, Answer All: Many Pretraining Experiments for the Cost of One
• Explaining Data Mixing Scaling Laws
• TiKMiX: Take Data Influence into Dynamic Mixture for Language Model Pre-training
• DynamixSFT: Dynamic Mixture Optimization of Instruction Tuning Collections
• Data Mixing Optimization for Supervised Fine-Tuning of Large Language Models
• Language Models Improve When Pretraining Data Matches Target Tasks
• Scaling Laws for Optimal Data Mixtures
• FineInstructions: Scaling Synthetic Instructions to Pre-Training Scale
• AutoMixAlign: Adaptive Data Mixing for Multi-Task Preference Optimization in LLMs
• MoDoMoDo: Multi-Domain Data Mixtures for Multimodal LLM Reinforcement Learning
• Chameleon: A Flexible Data-mixing Framework for Language Model Pretraining and Finetuning
• AC-ODM: Actor--Critic Online Data Mixing for Sample-Efficient LLM Pretraining
• Rethinking Data Mixture for Large Language Models: A Comprehensive Survey and New Perspectives
• GRAPE: Optimize Data Mixture for Group Robust Multi-target Adaptive Pretraining
• Efficient Data Selection at Scale via Influence Distillation
• Merge to Mix: Mixing Datasets via Model Merging
• HBO: Hierarchical Balancing Optimization for Fine-Tuning Large Language Models
• A Component-Based Survey of Interactions between Large Language Models and Multi-Armed Bandits
• MixAtlas: Uncertainty-aware Data Mixture Optimization for Multimodal LLM Midtraining

Related

• AC-ODM: Actor--Critic Online Data Mixing for Sample-Efficient LLM Pretraining
• Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance
• SampleMix: A Sample-wise Pre-training Data Mixing Strategey by Coordinating Data Quality and Diversity
• Optimizing Pretraining Data Mixtures with LLM-Estimated Utility
• Rethinking Data Mixture for Large Language Models: A Comprehensive Survey and New Perspectives