Sélection de données via le contrôle optimal pour les modèles de langage

Data Selection via Optimal Control for Language Models

October 9, 2024

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

cs.AI

Résumé

Ce travail étudie la sélection de données de pré-entraînement de haute qualité à partir de corpus massifs pour améliorer les capacités des Modèles de Langage (LM) pour une utilisation ultérieure. Nous formulons la sélection des données comme un problème de Contrôle Optimal généralisé, qui peut être résolu théoriquement par le Principe du Maximum de Pontryagin (PMP), donnant un ensemble de conditions nécessaires qui caractérisent la relation entre la sélection optimale des données et la dynamique de l'entraînement du LM. Sur la base de ces résultats théoriques, nous introduisons la Sélection de Données basée sur le PMP (PDS), un cadre qui approxime la sélection optimale des données en résolvant les conditions du PMP. Dans nos expériences, nous adoptons PDS pour sélectionner des données à partir de CommmonCrawl et montrons que le corpus sélectionné par PDS accélère l'apprentissage des LMs et améliore constamment leurs performances sur une large gamme de tâches ultérieures à travers diverses tailles de modèles. De plus, les avantages de PDS s'étendent aux modèles de ~400B entraînés sur ~10T jetons, comme en témoigne l'extrapolation des courbes de perte de test selon les Lois d'Échelle. PDS améliore également l'utilisation des données lorsque les données de pré-entraînement sont limitées, en réduisant la demande de données de 1,8 fois, ce qui atténue l'épuisement rapide des corpus disponibles issus du web. Notre code, nos données et nos points de contrôle de modèle peuvent être trouvés sur 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.