Selección de Datos a través del Control Óptimo para Modelos de Lenguaje

Data Selection via Optimal Control for Language Models

October 9, 2024

Autores: Yuxian Gu, Li Dong, Hongning Wang, Yaru Hao, Qingxiu Dong, Furu Wei, Minlie Huang

cs.AI

Resumen

Este trabajo investiga la selección de datos de pre-entrenamiento de alta calidad de corpus masivos para mejorar las capacidades de los Modelos de Lenguaje (LMs) para su uso posterior. Formulamos la selección de datos como un problema generalizado de Control Óptimo, que puede resolverse teóricamente mediante el Principio del Máximo de Pontryagin (PMP), produciendo un conjunto de condiciones necesarias que caracterizan la relación entre la selección óptima de datos y la dinámica del entrenamiento de LM. Basándonos en estos resultados teóricos, presentamos la Selección de Datos basada en PMP (PDS), un marco que aproxima la selección óptima de datos mediante la resolución de las condiciones de PMP. En nuestros experimentos, adoptamos PDS para seleccionar datos de CommmonCrawl y demostramos que el corpus seleccionado por PDS acelera el aprendizaje de LMs y mejora constantemente su rendimiento en una amplia gama de tareas posteriores en varios tamaños de modelos. Además, los beneficios de PDS se extienden a modelos de ~400B entrenados en ~10T tokens, como lo demuestra la extrapolación de las curvas de pérdida de prueba de acuerdo con las Leyes de Escala. PDS también mejora la utilización de datos cuando los datos de pre-entrenamiento son limitados, al reducir la demanda de datos en 1.8 veces, lo que mitiga el rápido agotamiento de los corpus rastreados en la web disponibles. Nuestro código, datos y puntos de control del modelo se pueden encontrar en https://github.com/microsoft/LMOps/tree/main/data_selection.

English

This work investigates the selection of high-quality pre-training data from massive corpora to enhance LMs' capabilities for downstream usage. We formulate data selection as a generalized Optimal Control problem, which can be solved theoretically by Pontryagin's Maximum Principle (PMP), yielding a set of necessary conditions that characterize the relationship between optimal data selection and LM training dynamics. Based on these theoretical results, we introduce PMP-based Data Selection (PDS), a framework that approximates optimal data selection by solving the PMP conditions. In our experiments, we adopt PDS to select data from CommmonCrawl and show that the PDS-selected corpus accelerates the learning of LMs and constantly boosts their performance on a wide range of downstream tasks across various model sizes. Moreover, the benefits of PDS extend to ~400B models trained on ~10T tokens, as evidenced by the extrapolation of the test loss curves according to the Scaling Laws. PDS also improves data utilization when the pre-training data is limited, by reducing the data demand by 1.8 times, which mitigates the quick exhaustion of available web-crawled corpora. Our code, data, and model checkpoints can be found in https://github.com/microsoft/LMOps/tree/main/data_selection.