Выбор данных с помощью оптимального управления для языковых моделей

Data Selection via Optimal Control for Language Models

October 9, 2024

Авторы: Yuxian Gu, Li Dong, Hongning Wang, Yaru Hao, Qingxiu Dong, Furu Wei, Minlie Huang

cs.AI

Аннотация

Данная работа исследует выбор высококачественных данных для предварительного обучения из обширных корпусов для улучшения возможностей языковых моделей (LMs) для последующего использования. Мы формулируем выбор данных как обобщенную задачу оптимального управления, которую можно теоретически решить с помощью Принципа Максимума Понтрягина (PMP), что приводит к набору необходимых условий, характеризующих отношение между оптимальным выбором данных и динамикой обучения LM. Основываясь на этих теоретических результатах, мы представляем PMP-основанный выбор данных (PDS), фреймворк, который приближает оптимальный выбор данных путем решения условий PMP. В наших экспериментах мы используем PDS для выбора данных из CommmonCrawl и показываем, что выбранный PDS корпус ускоряет обучение LMs и постоянно повышает их производительность на широком спектре последующих задач при различных размерах моделей. Более того, выгоды от PDS распространяются на модели размером ~400 млрд, обученные на ~10 трлн токенов, как показано экстраполяцией кривых тестовых потерь в соответствии с законами масштабирования. PDS также улучшает использование данных, когда объем данных для предварительного обучения ограничен, путем сокращения потребности в данных в 1,8 раза, что смягчает быстрое истощение доступных корпусов, полученных из веб-сканирования. Наш код, данные и контрольные точки моделей можно найти по ссылке 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.