DoReMi: Ottimizzare le Miscele di Dati Accelera il Pre-addestramento dei Modelli Linguistici

DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining

May 17, 2023

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

Abstract

Le proporzioni di miscela dei domini di dati di pre-addestramento (ad esempio, Wikipedia, libri, testo web) influenzano notevolmente le prestazioni dei modelli linguistici (LM). In questo articolo, proponiamo il Domain Reweighting with Minimax Optimization (DoReMi), che addestra inizialmente un piccolo modello proxy utilizzando l'ottimizzazione robusta distributiva di gruppo (Group DRO) sui domini per produrre pesi di dominio (proporzioni di miscela) senza conoscere i task downstream. Successivamente, ricampioniamo un dataset con questi pesi di dominio e addestriamo un modello più grande e completo. Nei nostri esperimenti, utilizziamo DoReMi su un modello proxy con 280 milioni di parametri per determinare i pesi di dominio per l'addestramento di un modello da 8 miliardi di parametri (30 volte più grande) in modo più efficiente. Su The Pile, DoReMi migliora la perplessità in tutti i domini, anche quando riduce il peso di un dominio. DoReMi migliora l'accuratezza media few-shot downstream del 6,5% rispetto a un modello di base addestrato utilizzando i pesi di dominio predefiniti di The Pile e raggiunge l'accuratezza di base con 2,6 volte meno passi di addestramento. Sul dataset GLaM, DoReMi, che non ha conoscenza dei task downstream, eguaglia addirittura le prestazioni ottenute utilizzando pesi di dominio ottimizzati sui task 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.