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

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 trains many small models on diverse data mixtures, uses regression to predict performance of unseen mixtures, and applies the best predicted mixture 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 to fit the regression model and predict the best data 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. Furthermore, RegMix consistently outperforms human selection in experiments involving models up to 7B models trained on 100B tokens, while matching or exceeding DoReMi using just 10% of the computational resources. Our experiments also show that (1) Data mixtures significantly impact performance; (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. Our code is available at https://github.com/sail-sg/regmix.

한 줄 요약

RegMix는 데이터 믹스 선택을 회귀 문제로 전환하여 대규모 언어 모델의 최적 학습 데이터를 자동으로 결정하는 방법이다.

핵심 기여도

RegMix는 512개의 1M 파라미터 소형 모델을 1B 토큰으로 학습하여 회귀 모델을 핏팅하고, 1B 파라미터 대형 모델에 최적의 데이터 믹스를 적용함.
1B 파라미터 모델을 25B 토큰으로 학습한 결과, 64개 후보 중 가장 높은 성능을 보임.
7B 파라미터 모델까지 확장한 실험에서 인간의 데이터 선택을 지속적으로 초과함.
DoReMi와 유사한 성능을 10%의 계산 자원으로 달성함.

핵심 아이디어

RegMix는 대규모 언어 모델의 데이터 믹스를 자동화하기 위해 **회귀 문제**로 접근한다. 기존 연구는 대형 모델을 직접 사용하거나, 소형 프록시 모델을 오랜 시간 학습하여 데이터 믹스를 조정하는 방식이었으나, 이는 계산 비용이 높고 확장성이 낮았다. RegMix는 **랭크 불변성**(rank invariance)이라는 가정을 바탕으로, 소형 모델을 다양한 데이터 믹스로 학습한 후 그 성능을 회귀 모델에 핏팅하여, **미래의 데이터 믹스 성능을 예측**한다. 이는 병렬 학습이 가능하며, 계산 효율성이 향상된다. 핵심 통찰은 **데이터 믹스의 상대적 순위는 모델 크기와 토큰 수에 관계없이 일정하다**는 점이다.

기술적 접근법

**RegMix 알고리즘**:
512개의 1M 파라미터 소형 모델을 1B 토큰으로 학습.
각 모델은 고유한 데이터 믹스로 학습됨.
소형 모델의 성능을 기반으로 회귀 모델을 핏팅하여, **미래의 데이터 믹스 성능을 예측**.
예측된 최적 믹스로 1B 파라미터 대형 모델을 25B 토큰으로 학습.
**데이터셋**: CommonCrawl, arXiv, GitHub, BookCorpus 등.
**하이퍼파라미터**: 1M 파라미터 모델은 1B 토큰, 1B 파라미터 모델은 25B 토큰으로 학습됨.

주요 결과

**1B 파라미터 모델**에서 RegMix 믹스는 64개 후보 중 가장 높은 성능을 보임.
**7B 파라미터 모델**에서 RegMix는 인간 선택을 지속적으로 초과함.
**DoReMi 대비 성능**: 유사한 성능을 10%의 계산 자원으로 달성.
**CommonCrawl**는 Wikipedia보다 다운스트림 성능과 강한 양의 상관관계를 보임.
**도메인 간 복잡한 상호작용**이 존재하며, 이는 인간 직관과 반대되는 경우가 많음.

의의 및 한계

RegMix는 대규모 언어 모델의 데이터 믹스 선택을 **자동화하고 계산 효율성을 높인다**. 특히, 데이터 믹스가 성능에 **14.6%까지 영향을 미친다는 점**에서, 이는 기존 스케일링 법칙을 넘어서는 새로운 통찰을 제공한다. 또한, 도메인 간 복잡한 상호작용을 자동으로 파악할 수 있어, 데이터 큐레이션의 신뢰성을 높인다. 그러나, RegMix는 **회귀 모델의 예측 정확도에 의존**하며, 일부 도메인의 중요도를 과소평가할 수 있는 한계가 있다. 또한, **회귀 모델 핏팅에 필요한 소형 모델 학습 비용**은 여전히 상당하다.

실용적 활용

RegMix는 대규모 언어 모델 개발 과정에서 **데이터 큐레이션의 자동화**를 가능하게 하며, 특히 **다양한 도메인의 복잡한 상호작용을 고려해야 하는 연구 및 산업적 상황**에 유용하다. 예를 들어, 대형 언어 모델을 개발하는 기업이나 연구소에서 데이터 믹스 최적화를 효율적으로 수행할 수 있다.