DoReMi: Die Optimierung von Datenmischungen beschleunigt das Vortraining von Sprachmodellen

DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining

May 17, 2023

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

Zusammenfassung

Die Mischungsverhältnisse der Domänen der Vortrainingsdaten (z. B. Wikipedia, Bücher, Webtexte) haben einen erheblichen Einfluss auf die Leistung von Sprachmodellen (Language Models, LMs). In diesem Artikel schlagen wir Domain Reweighting with Minimax Optimization (DoReMi) vor, bei dem zunächst ein kleines Proxy-Modell mithilfe von Group Distributionally Robust Optimization (Group DRO) über Domänen trainiert wird, um Domänengewichte (Mischungsverhältnisse) ohne Kenntnis der nachgelagerten Aufgaben zu erzeugen. Anschließend wird ein Datensatz mit diesen Domänengewichten neu abgetastet und ein größeres, vollständiges Modell trainiert. In unseren Experimenten verwenden wir DoReMi auf einem Proxy-Modell mit 280 Millionen Parametern, um Domänengewichte für das Training eines Modells mit 8 Milliarden Parametern (30-mal größer) effizienter zu ermitteln. Auf The Pile verbessert DoReMi die Perplexität über alle Domänen hinweg, selbst wenn eine Domäne geringer gewichtet wird. DoReMi steigert die durchschnittliche Few-Shot-Genauigkeit bei nachgelagerten Aufgaben um 6,5 % im Vergleich zu einem Baseline-Modell, das mit den Standard-Domänengewichten von The Pile trainiert wurde, und erreicht die Baseline-Genauigkeit mit 2,6-mal weniger Trainingsschritten. Auf dem GLaM-Datensatz erreicht DoReMi, das keine Kenntnis der nachgelagerten Aufgaben hat, sogar die Leistung von Domänengewichten, die auf nachgelagerten Aufgaben optimiert wurden.

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.