DoReMi: Optimalisatie van Datamengsels Versnelt de Voorafgaande Training van Taalmodellen

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

Samenvatting

De mengverhoudingen van domeinen in de vooraf getrainde data (bijv. Wikipedia, boeken, webtekst) hebben een grote invloed op de prestaties van taalmodel (LM). In dit artikel stellen we Domain Reweighting with Minimax Optimization (DoReMi) voor, waarbij eerst een klein proxy-model wordt getraind met behulp van group distributionally robust optimization (Group DRO) over domeinen om domeingewichten (mengverhoudingen) te produceren zonder kennis van downstream taken. Vervolgens hermonsteren we een dataset met deze domeingewichten en trainen we een groter, volledig model. In onze experimenten gebruiken we DoReMi op een proxy-model met 280M parameters om domeingewichten te vinden voor het trainen van een model met 8B parameters (30x groter) op een efficiëntere manier. Op The Pile verbetert DoReMi de perplexiteit in alle domeinen, zelfs wanneer het een domein onderwaardeert. DoReMi verbetert de gemiddelde few-shot downstream nauwkeurigheid met 6,5% ten opzichte van een basismodel dat is getraind met de standaard domeingewichten van The Pile en bereikt de baseline-nauwkeurigheid met 2,6x minder trainingsstappen. Op de GLaM-dataset, waarbij DoReMi geen kennis heeft van downstream taken, komt het zelfs overeen met de prestaties van het gebruik van domeingewichten die zijn afgestemd op downstream taken.

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.