RegMix: Mezcla de Datos como Regresión para el Preentrenamiento de Modelos de Lenguaje

RegMix: Data Mixture as Regression for Language Model Pre-training

July 1, 2024

Autores: Qian Liu, Xiaosen Zheng, Niklas Muennighoff, Guangtao Zeng, Longxu Dou, Tianyu Pang, Jing Jiang, Min Lin

cs.AI

Resumen

La mezcla de datos para el preentrenamiento de modelos de lenguaje a gran escala impacta significativamente el rendimiento, pero aún no está claro cómo determinar una mezcla efectiva. Proponemos RegMix para identificar automáticamente una mezcla de datos de alto rendimiento formulándola como una tarea de regresión. RegMix implica entrenar un conjunto de modelos pequeños con diversas mezclas de datos y ajustar un modelo de regresión para predecir su rendimiento dadas sus respectivas mezclas. Con el modelo de regresión ajustado, simulamos la mezcla mejor clasificada y la usamos para entrenar un modelo a gran escala con órdenes de magnitud más de capacidad computacional. Para validar empíricamente RegMix, entrenamos 512 modelos con 1 millón de parámetros para 1 billón de tokens de diferentes mezclas para ajustar el modelo de regresión y encontrar la mezcla óptima. Usando esta mezcla, entrenamos un modelo de 1 billón de parámetros para 25 billones de tokens (es decir, 1000 veces más grande y 25 veces más largo), que encontramos que tiene el mejor rendimiento entre 64 modelos candidatos de 1 billón de parámetros con otras mezclas. Además, nuestro método demuestra un rendimiento superior en comparación con la selección humana y logra resultados que igualan o superan a DoReMi, utilizando solo el 10% del presupuesto computacional. Nuestros experimentos también muestran que (1) Las mezclas de datos impactan significativamente el rendimiento, con variaciones de rendimiento en tareas individuales de hasta el 14.6%; (2) Los corpus web, en lugar de datos percibidos como de alta calidad como Wikipedia, tienen la correlación positiva más fuerte con el rendimiento en tareas posteriores; (3) Los dominios interactúan de maneras complejas que a menudo contradicen el sentido común, por lo que se necesitan enfoques automáticos como RegMix; (4) Los efectos de la mezcla de datos trascienden las leyes de escalado, y nuestro enfoque captura la complejidad al considerar todos los dominios juntos. Nuestro código está disponible en https://github.com/sail-sg/regmix.

English

The data mixture for large language model pre-training significantly impacts performance, yet how to determine an effective mixture remains unclear. We propose RegMix to automatically identify a high-performing data mixture by formulating it as a regression task. RegMix involves training a set of small models with diverse data mixtures and fitting a regression model to predict their performance given their respective mixtures. With the fitted regression model, we simulate the top-ranked mixture and use it to train a large-scale model with orders of magnitude more compute. To empirically validate RegMix, we train 512 models with 1M parameters for 1B tokens of different mixtures to fit the regression model and find the optimal mixture. Using this mixture we train a 1B parameter model for 25B tokens (i.e. 1000x larger and 25x longer) which we find performs best among 64 candidate 1B parameter models with other mixtures. Further, our method demonstrates superior performance compared to human selection and achieves results that match or surpass DoReMi, while utilizing only 10% of the compute budget. Our experiments also show that (1) Data mixtures significantly impact performance with single-task performance variations of up to 14.6%; (2) Web corpora rather than data perceived as high-quality like Wikipedia have the strongest positive correlation with downstream performance; (3) Domains interact in complex ways often contradicting common sense, thus automatic approaches like RegMix are needed; (4) Data mixture effects transcend scaling laws, and our approach captures the complexity by considering all domains together. Our code is available at https://github.com/sail-sg/regmix.