Data Selection via Optimal Control for Language Models
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.
670739dfd4311c2dd68f7fa5
Yuxian Gu
claimed_verified
2024-10-10T08:01:31.000Z
624ac662102fcdff87be51b9
https://cdn-avatars.huggingface.co/v1/production/uploads/624ac662102fcdff87be51b9/rzNahZFFkp194170tactJ.jpeg
Yuxian Gu
user
t1101675
670739dfd4311c2dd68f7fa6
Li Dong
claimed_verified
2024-10-11T07:44:21.119Z
5df85abada6d0311fd3d5408
/avatars/2331cf703c1b5d3a62e2050b1a6eb108.svg
Li Dong
user
unilm
670739dfd4311c2dd68f7fa7
Hongning Wang
admin_assigned
2024-10-10T09:25:41.531Z
62ec23fcbd19e355478fc584
/avatars/d7c567ef5f20bb3b9905cb5015d11e12.svg
Hongning Wang
user
howang
670739dfd4311c2dd68f7fa8
Yaru Hao
670739dfd4311c2dd68f7fa9
Qingxiu Dong
claimed_verified
2024-10-10T08:01:32.592Z
670740744341dcee459fb990
https://cdn-avatars.huggingface.co/v1/production/uploads/no-auth/66UkZvrAk7fQr5YCylEFk.png
Rosy24
user
Rsy24
670739dfd4311c2dd68f7faa
Furu Wei
admin_assigned
2024-10-10T09:25:28.584Z
6368c512fbfe97c16a40baba
/avatars/1c23bc7c0b6d9225699ce27647623d7a.svg
Furu Wei
user
thegenerality
670739dfd4311c2dd68f7fab
Minlie Huang
Data Selection via Optimal Control for Language Models
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.
A framework based on Pontryagin's Maximum Principle is proposed to optimize data selection for language models, enhancing their performance across various sizes and tasks.
Optimal Control problem Pontryagin's Maximum Principle (PMP) PMP-based Data Selection (PDS) language models data selection CommnonCrawl learning acceleration performance Scaling Laws data utilization