언어 모델을 위한 최적 제어를 통한 데이터 선택

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의 하류 사용 능력을 향상시키는 것을 조사합니다. 우리는 데이터 선택을 일반화된 최적 제어 문제로 정의하며, Pontryagin의 최대 원리(PMP)에 의해 이론적으로 해결될 수 있습니다. 이는 최적 데이터 선택과 LM 훈련 역학 사이의 관계를 특징 짓는 필요 조건 집합을 제공합니다. 이론적 결과를 기반으로 우리는 PMP 기반 데이터 선택(PDS)을 소개합니다. 이는 PMP 조건을 해결하여 최적 데이터 선택을 근사하는 프레임워크입니다. 실험에서 우리는 PDS를 채택하여 CommmonCrawl에서 데이터를 선택하고, PDS로 선택된 말뭉치가 LM의 학습을 가속화시키고 다양한 모델 크기에 걸쳐 다양한 하류 작업에서 지속적으로 성능을 향상시키는 것을 보여줍니다. 또한 PDS의 이점은 스케일링 법칙에 따라 테스트 손실 곡선의 추정에 의해 확인된 것처럼, ~10T 토큰으로 훈련된 ~400B 모델에까지 확장됩니다. 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.