DoReMi: Otimização de Misturas de Dados Acelera o Pré-treinamento de Modelos de Linguagem

DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining

May 17, 2023

Autores: Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy Liang, Quoc V. Le, Tengyu Ma, Adams Wei Yu

cs.AI

Resumo

As proporções de mistura dos domínios de dados de pré-treinamento (por exemplo, Wikipedia, livros, texto da web) afetam significativamente o desempenho dos modelos de linguagem (LM). Neste artigo, propomos o Rebalanceamento de Domínios com Otimização Minimax (DoReMi), que primeiro treina um pequeno modelo proxy usando otimização robusta de distribuição em grupos (Group DRO) sobre domínios para produzir pesos de domínio (proporções de mistura) sem conhecimento das tarefas downstream. Em seguida, reamostramos um conjunto de dados com esses pesos de domínio e treinamos um modelo maior e de tamanho completo. Em nossos experimentos, usamos o DoReMi em um modelo proxy com 280 milhões de parâmetros para encontrar pesos de domínio para treinar um modelo com 8 bilhões de parâmetros (30 vezes maior) de forma mais eficiente. No conjunto The Pile, o DoReMi melhora a perplexidade em todos os domínios, mesmo quando reduz o peso de um domínio. O DoReMi melhora a precisão média few-shot downstream em 6,5% em relação a um modelo de linha de base treinado usando os pesos de domínio padrão do The Pile e atinge a precisão da linha de base com 2,6 vezes menos etapas de treinamento. No conjunto de dados GLaM, o DoReMi, que não tem conhecimento das tarefas downstream, chega a igualar o desempenho do uso de pesos de domínio ajustados para tarefas downstream.

English

The mixture proportions of pretraining data domains (e.g., Wikipedia, books, web text) greatly affect language model (LM) performance. In this paper, we propose Domain Reweighting with Minimax Optimization (DoReMi), which first trains a small proxy model using group distributionally robust optimization (Group DRO) over domains to produce domain weights (mixture proportions) without knowledge of downstream tasks. We then resample a dataset with these domain weights and train a larger, full-sized model. In our experiments, we use DoReMi on a 280M-parameter proxy model to find domain weights for training an 8B-parameter model (30x larger) more efficiently. On The Pile, DoReMi improves perplexity across all domains, even when it downweights a domain. DoReMi improves average few-shot downstream accuracy by 6.5% over a baseline model trained using The Pile's default domain weights and reaches the baseline accuracy with 2.6x fewer training steps. On the GLaM dataset, DoReMi, which has no knowledge of downstream tasks, even matches the performance of using domain weights tuned on downstream tasks.