Gegevensselectie via Optimaal Beheer voor Taalmodellen

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

Samenvatting

Dit werk onderzoekt de selectie van hoogwaardige voorafgaande trainingsgegevens uit enorme corpora om de mogelijkheden van taalmodellen (LM's) voor downstream gebruik te verbeteren. We formuleren de gegevensselectie als een gegeneraliseerd Optimaal Besturingsprobleem, dat theoretisch kan worden opgelost met behulp van het Maximumprincipe van Pontryagin (PMP), wat resulteert in een reeks noodzakelijke voorwaarden die de relatie tussen optimale gegevensselectie en de dynamiek van LM-training karakteriseren. Op basis van deze theoretische resultaten introduceren we PMP-gebaseerde Gegevensselectie (PDS), een raamwerk dat optimale gegevensselectie benadert door de PMP-voorwaarden op te lossen. In onze experimenten passen we PDS toe om gegevens te selecteren uit CommmonCrawl en laten zien dat het door PDS geselecteerde corpus het leren van LM's versnelt en voortdurend hun prestaties verbetert op een breed scala van downstream taken over verschillende modelgroottes. Bovendien strekken de voordelen van PDS zich uit tot ~400B modellen getraind op ~10T tokens, zoals blijkt uit de extrapolatie van de testverliescurves volgens de Schalingswetten. PDS verbetert ook het gebruik van gegevens wanneer de voorafgaande trainingsgegevens beperkt zijn, door de vraag naar gegevens met 1,8 keer te verminderen, wat het snelle uitputten van beschikbare web-gecrawlede corpora vermindert. Onze code, gegevens en modelcontrolepunten zijn te vinden op 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.