RegMix: 언어 모델 사전 학습을 위한 회귀적 데이터 혼합

RegMix: Data Mixture as Regression for Language Model Pre-training

July 1, 2024

저자: Qian Liu, Xiaosen Zheng, Niklas Muennighoff, Guangtao Zeng, Longxu Dou, Tianyu Pang, Jing Jiang, Min Lin

cs.AI

초록

대규모 언어 모델 사전 학습을 위한 데이터 혼합은 성능에 큰 영향을 미치지만, 효과적인 혼합을 결정하는 방법은 여전히 불분명합니다. 우리는 이를 회귀 작업으로 공식화하여 고성능 데이터 혼합을 자동으로 식별하는 RegMix를 제안합니다. RegMix는 다양한 데이터 혼합으로 소규모 모델 세트를 학습하고, 각 혼합에 따른 성능을 예측하기 위해 회귀 모델을 피팅하는 과정을 포함합니다. 피팅된 회귀 모델을 사용하여 상위 순위의 혼합을 시뮬레이션하고, 이를 통해 훨씬 더 많은 계산 자원을 사용하여 대규모 모델을 학습시킵니다. RegMix를 실증적으로 검증하기 위해, 우리는 1M 매개변수를 가진 512개의 모델을 1B 토큰의 다양한 혼합으로 학습시켜 회귀 모델을 피팅하고 최적의 혼합을 찾았습니다. 이 혼합을 사용하여 1B 매개변수 모델을 25B 토큰(즉, 1000배 더 크고 25배 더 길게)으로 학습시켰으며, 이 모델은 다른 혼합을 사용한 64개의 후보 1B 매개변수 모델 중에서 최고의 성능을 보였습니다. 또한, 우리의 방법은 인간 선택보다 우수한 성능을 보였으며, DoReMi와 동등하거나 더 나은 결과를 달성하면서도 계산 예산의 10%만 사용했습니다. 우리의 실험은 또한 다음과 같은 사실을 보여줍니다: (1) 데이터 혼합은 성능에 큰 영향을 미치며, 단일 작업 성능 변동은 최대 14.6%에 이릅니다; (2) 위키피디아와 같은 고품질 데이터보다 웹 코퍼스가 다운스트림 성능과 가장 강한 양의 상관관계를 보입니다; (3) 도메인은 종종 상식과 반대되는 복잡한 방식으로 상호작용하므로, RegMix와 같은 자동화된 접근 방식이 필요합니다; (4) 데이터 혼합 효과는 스케일링 법칙을 초월하며, 우리의 접근 방식은 모든 도메인을 함께 고려하여 이러한 복잡성을 포착합니다. 우리의 코드는 https://github.com/sail-sg/regmix에서 확인할 수 있습니다.

English

The data mixture for large language model pre-training significantly impacts performance, yet how to determine an effective mixture remains unclear. We propose RegMix to automatically identify a high-performing data mixture by formulating it as a regression task. RegMix involves training a set of small models with diverse data mixtures and fitting a regression model to predict their performance given their respective mixtures. With the fitted regression model, we simulate the top-ranked mixture and use it to train a large-scale model with orders of magnitude more compute. To empirically validate RegMix, we train 512 models with 1M parameters for 1B tokens of different mixtures to fit the regression model and find the optimal mixture. Using this mixture we train a 1B parameter model for 25B tokens (i.e. 1000x larger and 25x longer) which we find performs best among 64 candidate 1B parameter models with other mixtures. Further, our method demonstrates superior performance compared to human selection and achieves results that match or surpass DoReMi, while utilizing only 10% of the compute budget. Our experiments also show that (1) Data mixtures significantly impact performance with single-task performance variations of up to 14.6%; (2) Web corpora rather than data perceived as high-quality like Wikipedia have the strongest positive correlation with downstream performance; (3) Domains interact in complex ways often contradicting common sense, thus automatic approaches like RegMix are needed; (4) Data mixture effects transcend scaling laws, and our approach captures the complexity by considering all domains together. Our code is available at https://github.com/sail-sg/regmix.