Daten Auswahl mittels Optimaler Steuerung für Sprachmodelle

Data Selection via Optimal Control for Language Models

October 9, 2024

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

cs.AI

Zusammenfassung

Diese Arbeit untersucht die Auswahl hochwertiger Vor-Trainingsdaten aus umfangreichen Korpora, um die Fähigkeiten von Sprachmodellen (LMs) für den nachgelagerten Gebrauch zu verbessern. Wir formulieren die Datenauswahl als ein verallgemeinertes Optimalsteuerungsproblem, das theoretisch mit dem Maximumsprinzip von Pontryagin (PMP) gelöst werden kann. Dies liefert eine Reihe von notwendigen Bedingungen, die die Beziehung zwischen optimaler Datenauswahl und den Trainingsdynamiken von LMs charakterisieren. Basierend auf diesen theoretischen Ergebnissen stellen wir die PMP-basierte Datenauswahl (PDS) vor, ein Framework, das die optimale Datenauswahl approximiert, indem es die PMP-Bedingungen löst. In unseren Experimenten verwenden wir PDS, um Daten aus CommonCrawl auszuwählen und zeigen, dass das von PDS ausgewählte Korpus das Lernen von LMs beschleunigt und kontinuierlich deren Leistung bei einer Vielzahl von nachgelagerten Aufgaben über verschiedene Modellgrößen hinweg steigert. Darüber hinaus erstrecken sich die Vorteile von PDS auf ~400B-Modelle, die auf ~10T Tokens trainiert wurden, wie durch die Extrapolation der Testverlustkurven gemäß den Skalierungsgesetzen belegt wird. PDS verbessert auch die Datennutzung, wenn die Vor-Trainingsdaten begrenzt sind, indem der Datenbedarf um das 1,8-fache reduziert wird, was der schnellen Erschöpfung verfügbarer web-gecrawlt Korpora entgegenwirkt. Unser Code, Daten und Modell-Checkpoints finden sich unter 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.