RegMix: Mistura de Dados como Regressão para Pré-treinamento de Modelos de Linguagem

Resumo

A mistura de dados para o pré-treinamento de grandes modelos de linguagem impacta significativamente o desempenho, no entanto, como determinar uma mistura eficaz permanece incerto. Propomos o RegMix para identificar automaticamente uma mistura de dados de alto desempenho formulando-a como uma tarefa de regressão. O RegMix envolve treinar um conjunto de pequenos modelos com misturas de dados diversas e ajustar um modelo de regressão para prever seu desempenho dado suas respectivas misturas. Com o modelo de regressão ajustado, simulamos a mistura mais bem classificada e a utilizamos para treinar um modelo em larga escala com ordens de magnitude mais computacionais. Para validar empiricamente o RegMix, treinamos 512 modelos com 1M parâmetros para 1B de tokens de diferentes misturas para ajustar o modelo de regressão e encontrar a mistura ótima. Utilizando essa mistura, treinamos um modelo com 1B de parâmetros para 25B de tokens (ou seja, 1000 vezes maior e 25 vezes mais longo), o qual descobrimos ter o melhor desempenho entre 64 modelos candidatos com 1B de parâmetros e outras misturas. Além disso, nosso método demonstra desempenho superior em comparação com a seleção humana e alcança resultados que correspondem ou superam o DoReMi, enquanto utilizando apenas 10% do orçamento computacional. Nossos experimentos também mostram que (1) as misturas de dados impactam significativamente o desempenho, com variações de desempenho de tarefas únicas de até 14,6%; (2) corpora da web, em vez de dados percebidos como de alta qualidade como a Wikipedia, têm a mais forte correlação positiva com o desempenho subsequente; (3) os domínios interagem de maneiras complexas, muitas vezes contradizendo o senso comum, assim abordagens automáticas como o RegMix são necessárias; (4) os efeitos da mistura de dados transcendem as leis de escalonamento, e nossa abordagem captura a complexidade ao considerar todos os domínios juntos. Nosso código está disponível em 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.

RegMix: Mistura de Dados como Regressão para Pré-treinamento de Modelos de Linguagem

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

Resumo

Support