Data Selection via Optimal Control for Language Models
Authors: name: Yuxian Gu, name: Li Dong, name: Hongning Wang, name: Yaru Hao, name: Qingxiu Dong, name: Furu Wei, name: Minlie Huang

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.