DoReMi : Optimiser les mélanges de données accélère le pré-entraînement des modèles de langage

DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining

May 17, 2023

Auteurs: 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

Résumé

Les proportions de mélange des domaines de données de pré-entraînement (par exemple, Wikipédia, livres, texte web) influencent grandement les performances des modèles de langage (LM). Dans cet article, nous proposons Domain Reweighting with Minimax Optimization (DoReMi), qui commence par entraîner un petit modèle proxy en utilisant l'optimisation robuste de distribution par groupe (Group DRO) sur les domaines pour produire des poids de domaine (proportions de mélange) sans connaissance des tâches en aval. Nous rééchantillonnons ensuite un ensemble de données avec ces poids de domaine et entraînons un modèle plus grand et complet. Dans nos expériences, nous utilisons DoReMi sur un modèle proxy de 280 millions de paramètres pour trouver des poids de domaine afin d'entraîner un modèle de 8 milliards de paramètres (30 fois plus grand) de manière plus efficace. Sur The Pile, DoReMi améliore la perplexité dans tous les domaines, même lorsqu'il réduit le poids d'un domaine. DoReMi améliore la précision moyenne en few-shot en aval de 6,5 % par rapport à un modèle de référence entraîné avec les poids de domaine par défaut de The Pile et atteint la précision de référence avec 2,6 fois moins d'étapes d'entraînement. Sur le jeu de données GLaM, DoReMi, qui n'a aucune connaissance des tâches en aval, parvient même à égaler les performances obtenues avec des poids de domaine ajustés sur les tâches en aval.

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.