Selezione dei Dati tramite Controllo Ottimale per Modelli Linguistici

Data Selection via Optimal Control for Language Models

October 9, 2024

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

cs.AI

Abstract

Questo lavoro investiga la selezione di dati di pre-training di alta qualità da corpora massicci per potenziare le capacità dei Modelli Linguistici (LM) per utilizzi successivi. Formuliamo la selezione dei dati come un problema di Controllo Ottimale generalizzato, che può essere risolto teoricamente dal Principio del Massimo di Pontryagin (PMP), producendo un insieme di condizioni necessarie che caratterizzano la relazione tra la selezione ottimale dei dati e la dinamica di addestramento del LM. Sulla base di questi risultati teorici, introduciamo la Selezione dei Dati basata su PMP (PDS), un framework che approssima la selezione ottimale dei dati risolvendo le condizioni del PMP. Nei nostri esperimenti, adottiamo PDS per selezionare dati da CommmonCrawl e mostriamo che il corpus selezionato da PDS accelera l'apprendimento dei LM e migliora costantemente le loro prestazioni su una vasta gamma di compiti successivi attraverso varie dimensioni del modello. Inoltre, i benefici di PDS si estendono a modelli di ~400B addestrati su ~10T token, come dimostrato dall'estrapolazione delle curve di perdita di test secondo le Leggi di Scalabilità. PDS migliora anche l'utilizzo dei dati quando i dati di pre-training sono limitati, riducendo la domanda di dati di 1,8 volte, il che mitiga il rapido esaurimento dei corpora disponibili raccolti da web-crawled. Il nostro codice, dati e checkpoint dei modelli possono essere trovati su 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.