寫在前面

大型語言模型在預(yù)訓(xùn)練過程中，如何選取數(shù)據(jù)的混合比例（利用較少的Tokens來實現(xiàn)較小的Loss從而加速預(yù)訓(xùn)練過程）是一個復(fù)雜和關(guān)鍵的問題。手動確認數(shù)據(jù)集中各個組成的比例是不可擴展的，并且很可能不是最優(yōu)選擇。

今天給大家介紹一個用回歸任務(wù)解決大模型數(shù)據(jù)混合問題的方法-RegMix。其核心思想是，利用不同的數(shù)據(jù)混合比例先訓(xùn)練多個小模型并獲取其結(jié)果，在利用這些樣本訓(xùn)練一個回歸模型，再遍歷所有比例利用回歸模型找到最優(yōu)的數(shù)據(jù)混合比例，最后用最優(yōu)數(shù)據(jù)混合比例訓(xùn)練更大的語言模型。

Paper: https://arxiv.org/abs/2407.01492
Github: https://github.com/sail-sg/regmix

通過訓(xùn)練512個1M的小模型，擬合回歸模型，找到top64的數(shù)據(jù)混合比例，訓(xùn)練1B模型，最優(yōu)數(shù)據(jù)混合比例訓(xùn)練的模型的驗證集loss也是最低。

方法

整體流程如上圖所示，

• 生成隨機數(shù)據(jù)混合比例，按照比例采用混合數(shù)據(jù)并訓(xùn)練小模型；
• 利用數(shù)據(jù)混合比例作為特征值，模型訓(xùn)練的目標(biāo)值作為標(biāo)簽，擬合回歸模型；
• 在模擬更大數(shù)據(jù)混合比例空間，利用回歸模型預(yù)測最佳目標(biāo)值，以獲取最佳混合比例；
• 使用模擬出的最佳混合比例的數(shù)據(jù)訓(xùn)練更大的模型。

訓(xùn)練小模型時越多越好，但為了節(jié)約成本需要盡量減少小模型訓(xùn)練次數(shù)，那么在初始化數(shù)據(jù)混合比例時就需要時多樣化的，并且每個數(shù)據(jù)領(lǐng)域需要都存在極端值，數(shù)據(jù)采用過程主要是基于Tokens（chunk-level）分布的狄利克雷分布來實現(xiàn)。

詳見：mixture_config/synthesize_mixture.py

同時在擬合回歸模型時，采用了線性回歸和LightGBM兩種回歸模型。

結(jié)果

數(shù)據(jù)集采用Pile dataset中不涉及版權(quán)的17個數(shù)據(jù)集，如下表所示，

512個1M小模型在1B Tokens訓(xùn)練得到的回歸模型，與在25B Tokens數(shù)據(jù)下訓(xùn)練的1B模型，排序具有97.12%的高相關(guān)性，如下表所示，

同時訓(xùn)練次數(shù)要比訓(xùn)練的總Token數(shù)要重要，更影響回歸模型的效果，并且采用LightGBM建模要比線性回歸建模要好。

PS：跟作者@乾神交流過，512個樣本訓(xùn)練回歸模型會不會數(shù)據(jù)量太少，乾神說他們做過1024的實驗，但并回歸模型效果無明顯提高，并且從成本考慮，那么512最佳。

不同的數(shù)據(jù)混合比例對下游任務(wù)結(jié)果影響較大，在Lambada數(shù)據(jù)集上最好和最差的效果相差14.6%，如下表所示，

同時發(fā)現(xiàn)了一個與傳統(tǒng)理解不一致的結(jié)果，一般我們?nèi)蝿?wù)維基數(shù)據(jù)質(zhì)量很高，是評估大型語言模型最具代表性的數(shù)據(jù)集。但實驗結(jié)果發(fā)現(xiàn)，網(wǎng)絡(luò)數(shù)據(jù)集上評估的效果，更能體現(xiàn)模型在下游任務(wù)上的好壞，如下圖所示，可以發(fā)現(xiàn)Pile-CC數(shù)據(jù)集作為驗證時損失值與下游任務(wù)的相關(guān)性更強。

并且RegMix可以發(fā)現(xiàn)各領(lǐng)域數(shù)據(jù)之間是如何相互作用的，數(shù)據(jù)領(lǐng)域之間復(fù)雜的相互作用利用人類固有經(jīng)驗很難直接區(qū)分。

本文轉(zhuǎn)載自 ??NLP工作站??，作者：劉聰NLP