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を実証的に検証するため、512個の100万パラメータのモデルを10億トークンの異なる混合で訓練し、回帰モデルを適合させて最適な混合を見つけた。この混合を用いて、10億パラメータのモデルを250億トークン（つまり1000倍大きく、25倍長く）訓練し、他の混合で訓練された64個の候補モデルの中で最高の性能を示すことを確認した。さらに、本手法は人間による選択よりも優れた性能を示し、DoReMiと同等またはそれ以上の結果を達成しながら、計算予算の10%しか使用しない。実験結果から、(1) データ混合は性能に大きな影響を与え、単一タスクの性能変動は最大14.6%に及ぶこと、(2) Wikipediaのような高品質とされるデータよりも、Webコーパスが下流タスクの性能と最も強い正の相関を持つこと、(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.