寫在前面

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

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

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

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

方法

整體流程如上圖所示，

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

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

詳見：mixture_config/synthesize_mixture.py

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

結(jié)果

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

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

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

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

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

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

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

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