Data Selection via Optimal Control for Language Models (2410.07064v2)

Abstract: 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 helps mitigate the quick exhaustion of available web-crawled corpora. Our code, model, and data can be found at https://github.com/microsoft/LMOps/tree/main/data_selection.

• The paper introduces PMP-based Data Selection (PDS), a novel framework that optimizes pre-training data selection for language models.
• The methodology formulates selection as an Optimal Control problem, reducing data demand by 1.8 times and accelerating training.
• Empirical validation shows consistent improvements across tasks for models up to ~400B parameters, demonstrating robust scalability.

Data Selection via Optimal Control for LLMs

This paper presents a novel approach to enhancing LLM (LM) performance through careful selection of pre-training data. The research formulates data selection as an Optimal Control problem, leveraging Pontryagin's Maximum Principle (PMP) to derive conditions for optimal data selection. This establishes a theoretical foundation for understanding the relationship between data and LM training dynamics.

The authors introduce PMP-based Data Selection (PDS), which applies these theoretical insights to select high-quality pre-training data from massive corpora like CommonCrawl. Empirical results demonstrate that PDS not only accelerates model learning but also consistently enhances downstream performance across various model sizes, indicating the framework's robustness and scalability. Specifically, the benefits extend even to models as large as ~400B, trained on ~10T tokens, showcasing the approach's applicability to modern, large-scale LMs.

A significant advantage of PDS is its ability to improve data utilization when pre-training data is limited. The paper reports a data demand reduction by 1.8 times, addressing concerns over the quick depletion of available high-quality web-crawled corpora. This aspect of the research has practical implications for the efficient use of data, especially as the LM community confronts constraints on accessible corpora.

Theoretical Foundation and Methodology

The paper outlines how data selection can be approached using Optimal Control theory, where control variables (data points) in a dynamic system (pre-training process) are optimized to achieve desired outcomes (low downstream loss). Through PMP, the authors derive necessary conditions that include a "target vector" integrating LM training dynamics and downstream task performance, guiding the selection of valuable data.

The authors implement PDS by approximating the optimal data selection solution defined by PMP. The procedure involves solving the equation system from PMP in a proxy environment, training a small LM to predict data quality scores, and employing these scores to select the final pre-training corpus.

Experimental Validation

In experiments, PDS leads to approximately two times faster pre-training for a 1.7B parameter model and yields consistent improvements in both LLMing tasks and a range of downstream tasks. Extrapolation based on the Scaling Law predicts persistent benefits for extremely large models, illustrating PDS's scalability.

The research includes a thorough analysis of PDS, comparing it to conventional pre-training and other offline data selection methods. PDS consistently shows superior performance, supporting the claim that incorporating LM training dynamics into data selection results in higher-quality data selection compared to traditional pattern-based or deduplication methods.

Implications and Future Directions

The framework put forth in this paper has significant implications. Practically, it allows for more efficient pre-training by reducing computational costs and data usage. Theoretically, it proposes a rigorous, control-theoretic approach to data selection, potentially influencing future research and methodologies in AI and LM training.

Future work could explore the application of Optimal Control to other areas of machine learning beyond data selection. Additionally, the integration of this approach with online data selection methods, which dynamically adjust during training, might yield further improvements in model efficiency and performance. Continued investigation into the transferability of data quality information across different model sizes and architectures also offers a promising avenue for research.

In summary, this paper provides a comprehensive and theoretically grounded method for optimizing data selection in LM pre-training, with both theoretical sophistication and practical impact.