DoReMi: Optimizar las mezclas de datos acelera el preentrenamiento de modelos de lenguaje

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

Resumen

Las proporciones de mezcla de los dominios de datos de preentrenamiento (por ejemplo, Wikipedia, libros, texto web) afectan significativamente el rendimiento de los modelos de lenguaje (LM). En este artículo, proponemos Reajuste de Dominios con Optimización Minimax (DoReMi), que primero entrena un modelo proxy pequeño utilizando optimización robusta de distribución grupal (Group DRO) sobre dominios para producir pesos de dominio (proporciones de mezcla) sin conocimiento de las tareas posteriores. Luego, remuestreamos un conjunto de datos con estos pesos de dominio y entrenamos un modelo más grande y de tamaño completo. En nuestros experimentos, utilizamos DoReMi en un modelo proxy de 280 millones de parámetros para encontrar pesos de dominio para entrenar un modelo de 8 mil millones de parámetros (30 veces más grande) de manera más eficiente. En The Pile, DoReMi mejora la perplejidad en todos los dominios, incluso cuando reduce el peso de un dominio. DoReMi mejora la precisión promedio en tareas posteriores de pocos ejemplos en un 6.5% sobre un modelo de referencia entrenado con los pesos de dominio predeterminados de The Pile y alcanza la precisión de referencia con 2.6 veces menos pasos de entrenamiento. En el conjunto de datos GLaM, DoReMi, que no tiene conocimiento de las tareas posteriores, incluso iguala el rendimiento de usar pesos de dominio ajustados en tareas posteriores.

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.