RegMix: Смесь данных как регрессия для предварительного обучения языковой модели

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

July 1, 2024

Авторы: Qian Liu, Xiaosen Zheng, Niklas Muennighoff, Guangtao Zeng, Longxu Dou, Tianyu Pang, Jing Jiang, Min Lin

cs.AI

Аннотация

Смесь данных для предварительного обучения крупных языковых моделей значительно влияет на производительность, однако как определить эффективную смесь остается неясным. Мы предлагаем RegMix для автоматической идентификации высокопроизводительной смеси данных, формулируя это как задачу регрессии. RegMix включает в себя обучение набора небольших моделей с разнообразными смесями данных и подгонку регрессионной модели для прогнозирования их производительности с учетом их соответствующих смесей. С помощью подогнанной регрессионной модели мы симулируем лучшую смесь и используем ее для обучения модели крупного масштаба с порядком большим вычислительным потенциалом. Для эмпирической проверки RegMix мы обучаем 512 моделей с 1M параметрами для 1B токенов различных смесей для подгонки регрессионной модели и находим оптимальную смесь. Используя эту смесь, мы обучаем модель с 1B параметрами для 25B токенов (т.е. в 1000 раз больше и на 25 раз дольше), которая, как мы обнаружили, показывает лучшие результаты среди 64 кандидатских моделей с 1B параметрами и другими смесями. Кроме того, наш метод демонстрирует превосходную производительность по сравнению с человеческим выбором и достигает результатов, соответствующих или превосходящих DoReMi, используя лишь 10% бюджета на вычисления. Наши эксперименты также показывают, что (1) Смеси данных значительно влияют на производительность с вариациями производительности однозадачных моделей до 14.6%; (2) Веб-корпуса, а не данные, воспринимаемые как высококачественные, такие как Википедия, имеют наиболее сильную положительную корреляцию с последующей производительностью; (3) Домены взаимодействуют сложным образом, часто противореча общему смыслу, поэтому необходимы автоматические подходы, такие как RegMix; (4) Эффекты смеси данных превосходят законы масштабирования, и наш подход улавливает сложность, рассматривая все домены вместе. Наш код доступен по адресу https://github.com/sail-sg/regmix.

English

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 involves training a set of small models with diverse data mixtures and fitting a regression model to predict their performance given their respective mixtures. With the fitted regression model, we simulate the top-ranked mixture and use it 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 of different mixtures to fit the regression model and find the optimal 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. Further, our method demonstrates superior performance compared to human selection and achieves results that match or surpass DoReMi, while utilizing only 10% of the compute budget. Our experiments also show that (1) Data mixtures significantly impact performance with single-task performance variations of up to 14.6%; (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, and our approach captures the complexity by considering all domains together. Our code is available at https://github.com/sail-sg/regmix.