Seleção de Dados via Controle Ótimo para Modelos de Linguagem

Data Selection via Optimal Control for Language Models

October 9, 2024

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

cs.AI

Resumo

Este trabalho investiga a seleção de dados de pré-treinamento de alta qualidade a partir de corpora massivos para aprimorar as capacidades de Modelos de Linguagem (LMs) para uso posterior. Formulamos a seleção de dados como um problema generalizado de Controle Ótimo, que pode ser teoricamente resolvido pelo Princípio do Máximo de Pontryagin (PMP), resultando em um conjunto de condições necessárias que caracterizam a relação entre a seleção ótima de dados e a dinâmica de treinamento do LM. Com base nesses resultados teóricos, introduzimos a Seleção de Dados baseada em PMP (PDS), um framework que aproxima a seleção ótima de dados ao resolver as condições do PMP. Em nossos experimentos, adotamos o PDS para selecionar dados do CommmonCrawl e demonstramos que o corpus selecionado pelo PDS acelera a aprendizagem dos LMs e constantemente melhora seu desempenho em uma ampla gama de tarefas posteriores em vários tamanhos de modelo. Além disso, os benefícios do PDS se estendem a modelos de ~400B treinados em ~10T tokens, conforme evidenciado pela extrapolação das curvas de perda de teste de acordo com as Leis de Escalonamento. O PDS também melhora a utilização de dados quando os dados de pré-treinamento são limitados, reduzindo a demanda de dados em 1,8 vezes, o que mitiga o rápido esgotamento dos corpora disponíveis na web. Nosso código, dados e checkpoints de modelo podem ser encontrados em 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.