RegMix: Data Mixture as Regression for Language Model Pre-training

2024-07-01

arXiv.org

AI Summary

To address the challenge of manually optimizing data mixing ratios during large language model (LLM) pretraining, this paper formulates data proportion search as a regression task and proposes a lightweight, automated framework comprising small-model ensemble sampling, performance regression prediction, and Bayesian optimization. Using only 1M-parameter surrogate models, the framework efficiently evaluates data mixture efficacy at the 1B-token scale. Key findings include: (i) web-sourced corpora yield substantially greater downstream task gains than conventional “high-quality” sources (e.g., Wikipedia); (ii) non-intuitive cross-domain interaction effects necessitate automated optimization; and (iii) data mixture benefits are decoupled from model scaling laws. On a 1B-parameter model, our method achieves state-of-the-art performance, outperforming 64 human-designed mixtures. For a 7B-parameter model, it attains superior performance to DoReMi at merely 10% of its computational cost.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Learning to Search
• Natural Language Processing: (Large) Language Models

Application Category

• Web Mining and Content Analysis: Large pretrained models with web data
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

The data mixture for large language model pre-training significantly impacts performance, yet how to determine an effective mixture remains unclear. We propose RegMix to automatically identify a high-performing data mixture by formulating it as a regression task. RegMix trains many small models on diverse data mixtures, uses regression to predict performance of unseen mixtures, and applies the best predicted mixture to train a large-scale model with orders of magnitude more compute. To empirically validate RegMix, we train 512 models with 1M parameters for 1B tokens to fit the regression model and predict the best data mixture. Using this mixture we train a 1B parameter model for 25B tokens (i.e. 1000x larger and 25x longer) which we find performs best among 64 candidate 1B parameter models with other mixtures. Furthermore, RegMix consistently outperforms human selection in experiments involving models up to 7B models trained on 100B tokens, while matching or exceeding DoReMi using just 10% of the computational resources. Our experiments also show that (1) Data mixtures significantly impact performance; (2) Web corpora rather than data perceived as high-quality like Wikipedia have the strongest positive correlation with downstream performance; (3) Domains interact in complex ways often contradicting common sense, thus automatic approaches like RegMix are needed; (4) Data mixture effects transcend scaling laws. Our code is available at https://github.com/sail-sg/regmix.

Problem

• Large Language Models
• Data Combination
• Performance Enhancement

Innovation

• RegMix
• Data Combination Optimization
• Efficient Model Pre-training