一、結論寫在前面

論文標題：RegMix: Data Mixture as Regression for Language Model Pre-training

論文鏈接：??https://arxiv.org/pdf/2407.01492??

代碼：??https://github.com/sail-sg/regmix??

大型語言模型預訓練中的數據配比對性能有顯著影響，但如何確定有效的配比仍不明確。論文提出REGMIX，通過將數據配比問題形式化為回歸任務，自動識別高性能的數據配比。REGMIX包括使用多種數據配比訓練一組小型模型，并擬合一個回歸模型來預測給定各自配比的模型的性能。利用擬合的回歸模型，論文模擬排名最高的配比，并使用它來訓練一個大規模模型，其計算量是之前的幾個數量級。

為了實證驗證REGMIX，論文訓練了512個具有1M參數的模型，用于1B個不同配比的token，以擬合回歸模型并找到最佳配比。使用這種配比，論文訓練了一個1B參數的模型，用于25B個token（即比之前大1000倍，長25倍），論文發現它在64個候選的1B參數模型中表現最佳，這些模型使用了其他配比。

為了驗證REGMIX，論文使用不同的數據配比訓練了1M和1B參數的模型。通過在1Btoken上訓練512個1M參數的模型，論文發現REGMIX優化的數據配比在64個1B參數模型（比1M參數大1000倍）和訓練時間長25倍（25Btoken）的模型中表現更優，如圖1所示。此外，使用REGMIX優化的數據配比產生的模型性能優于人工選擇，并且與旗艦DoReMi方法[64]的性能相當，盡管它需要的總計算量更少，并允許并行訓練。論文還發現：(1) 數據配比顯著影響下游性能，導致單任務性能差異高達14.6%；(2) 通用網絡語料庫（如CommonCrawl），而非維基百科，與下游任務性能提升的正相關性最強；(3) 領域間的交互復雜且常常與直覺相悖，突顯了像REGMIX這樣的自動化方法的必要性；(4) 數據配比效應超越了縮放定律，REGMIX通過綜合考慮所有領域來捕捉其復雜性。    

二、論文的簡單介紹

2.1 論文的背景

大規模公開數據集的可用性已成為創建大型語言模型（LLMs）的關鍵因素。大多數數據可在互聯網上獲取，包括學術論文（例如arXiv）、書籍（例如古騰堡計劃）和代碼配比(例如GitHub)。在創建首批LLMs之一，GPT-3 時，作者已經認識到選擇最佳訓練數據的重要性，因此他們決定增加維基百科的采樣比例，因其被認為具有高質量。然而，這種手動數據選擇不具備可擴展性，可能導致次優選擇。隨著用于LLM預訓練的數據規模和多樣性不斷增長，確定最優數據配比變得愈發具有挑戰性。這引發了一個關鍵的研究問題：論文如何以可擴展且高效的方式選擇最優數據配比？

先前的工作[64, 16, 2]采用小規模模型（“代理模型”）來預測大規模語言模型的領域權重。這些工作通過訓練具有大量token(例如，1000)，有時甚至與用于訓練LLM的數量相同，并通過監控訓練動態來動態調整分配策略。然而，隨著用于預訓練LLM的訓練數據不斷增長，這些方法變得效率低下。為當前模型（如Llama-3）訓練代理模型，需要使用多達15T token，這在當前方法下可能過于昂貴且過于緩慢，不值得。    

圖1：左：論文假設數據配比在不同模型大小和訓練token數量下的秩不變性。利用這一假設，論文使用在較少token上訓練的小模型來預測用于訓練具有更多token的大型模型的有效數據配比。右：通過訓練512 x 1M模型，論文的方法在訓練64 x 1B模型之前識別出最佳數據配比。預測的最佳數據配比，用紅色星號表示，實現了最低的驗證損失。

論文認為在有限數量的token上訓練小型模型足以預測用于大型語言模型（LLM）訓練的有效數據配比。論文的關鍵假設是數據配比的秩不變性，即在不同模型大小和訓練token數量下，數據配比對模型性能影響的相對排名是一致的。在這一假設下，關鍵挑戰在于從近乎無限數量的潛在數據配比中發掘出排名最高的數據配比。為此，論文將數據配比選擇視為回歸任務。論文不是通過每一種可能的配比來窮盡地訓練小型模型，而是僅訓練一組小型模型，每個模型具有獨特的數據配比?；谶@些模型的性能及其配比，論文擬合一個回歸模型來預測其他數據配比的性能。論文的方法比先前的工作更具可擴展性，因為它允許并行訓練小型代理模型，而不是長時間訓練單個模型。此外，回歸模型提供了關于領域交互的洞察，有助于理解和數據整理。

圖2：使用Hacker News、GitHub和Philpapers作為訓練域，以StackExchange域上的損失作為目標（，展示論文的方法。通過使用小規模代理模型訓練日志擬合回歸模型，并用于預測模擬空間內最佳數據配比，從而直接預測大規模語言模型預訓練的數據配比

2.2 REGMIX：數據配比作為回歸

如圖2所示，論文的方法涉及四個關鍵步驟：（1）生成隨機數據配比并在此基礎上訓練小規模代理模型。（2）使用表1：Pile數據集概覽，其中因版權問題不再可用的數據集以灰色標記。在論文的實驗中，論文使用可獲得的17個領域來研究語言模型預訓練的數據配比作為特征，目標值作為標簽，擬合一個線性回歸模型。（3）在更大規模上模擬數據配比空間，并利用回歸模型識別目標值的最佳配比。（4）使用模擬的最佳數據配比訓練一個大規模模型。    

2.2.1 訓練小規模代理模型

第一步是在多個不同的數據配比上訓練一組小規模代理模型。為了減少所需的運行次數，論文的目標是選擇一系列涵蓋每個領域從0%到100%極端權重的多樣化數據配比。論文通過基于詞頻分布的Dirichlet分布來實現這一點，這使論文能夠采樣廣泛的范圍值，并將回歸模型暴露于各種極端情況。同時，基于詞頻分布確保了整體數據配比在統計上反映了數據的可用性。例如，這可以防止任何詞頻低于1%的領域被過度強調，這在大型訓練中是不可行的，因為該領域沒有足夠的可用詞頻。在實踐中，論文將詞頻分布乘以0.1到5.0之間的值，以構建各種稀疏和近似均勻的分布，然后將這些分布向量作為Dirichlet分布的超參數alpha。

在訓練小規模代理模型幾步之后，論文可以獲得幾個訓練良好的小模型。例如，在論文的主要實驗中，每個代理模型包含1M參數，并在1B詞頻上進行訓練。然后，論文可以選擇在領域或基準上評估這些訓練好的模型，以獲得論文想要優化的目標值。通常，目標值可以是某個領域的損失，如圖2所示的StackExchange領域。一旦論文獲得了這些目標值，論文就可以使用數據配比作為特征，目標值作為標簽來擬合一個回歸模型。

2.2.2 擬合回歸模型

第二步是使用數據配比作為特征，目標值作為標簽來擬合一個回歸模型?；貧w任務是一個傳統的監督學習任務，涉及基于輸入特征X=( x_1, x_2, ..., x_n )預測連續目標變量y。目標是找到一個函數f，該函數最好地將輸入特征映射到目標變量，使得y=f(X)+epsilon，其中epsilon表示數據中的誤差或噪聲。輸入特征X對應于數據配比的領域權重，而目標變量y是論文想要優化的值。使用這些數據，論文訓練回歸模型，學習一個函數來預測基于任意數據配比的目標值，而無需進一步訓練。    

線性回歸。線性回歸模型在回歸分析中應用廣泛。它假設輸入特征與目標變量之間存在線性關系，可以表示為：

其中，omega_0 是截距，omega=(omega_1, ..., omega_n) 是與相應輸入特征 x_1, ..., x_n 相關的系數。系數 omega 通常使用普通最小二乘法等技術進行估計，旨在最小化預測值與實際值之間的殘差平方和。

表2展示了論文基于512x 1M模型在1Btoken上訓練的結果來擬合回歸模型，并在未見過的數據配比上評估1M、60M和1B參數模型的表現。皮爾遜相關系數 r 和均方誤差（MSE）衡量損失預測性能，而斯皮爾曼相關系數 rho 比較預測值與實際值的排名。

LightGBM回歸。LightGBM 是一種強大的梯度提升算法，可用于回歸和分類任務。在回歸的背景下，LightGBM 學習一組決策樹來預測目標變量。該過程由基于梯度的優化算法指導，該算法最小化指定的損失函數（例如均方誤差)。此外，LightGBM 被設計為高效且可擴展，適合大型數據集。

2.2.3 模擬與預測

一旦論文訓練了回歸模型，論文就可以高效地探索整個可能的數據配比空間。通過使用訓練好的模型預測每個潛在數據配比的目標值，論文可以快速識別產生最佳目標值的輸入。這種基于模擬的優化相對廉價，因為模擬和回歸預測在計算上都很快。例如，對1,000,000個數據配比進行預測僅需不到10個CPU秒。    

2.2.4 大規模模型訓練

在通過模擬識別出最佳數據配比后，論文將排名靠前的數據配比推廣到包含更多token的大規模模型訓練中。如圖 2所示，論文直接使用最佳數據配比來訓練更大的模型。實際上，為了提高回歸預測的穩健性，論文選擇前100個配比并將其平均作為大規模訓練的數據配比。

2.3 回歸預測評估

論文評估REGMIX預測未見數據配比效果的能力。首先，論文使用小型（即1M參數）模型的訓練工件擬合回歸模型，并評估小型模型上的損失預測性能。然后，為了驗證論文的排名不變性假設，論文測試學習到的回歸模型在不同模型大小和token數量上預測排名的能力。

2.3.1 實驗設置

數據集和模型。論文使用表1中描繪的Pile數據集的domain進行實驗。由于版權問題，論文使用HuggingFace上不違反版權問題的17個子集。論文考慮線性回歸和LightGBM回歸模型，其中目標變量y設置為Pile-CC領域的驗證損失。

訓練與評估。該回歸模型利用512個512x 1M模型的訓練成果，基于10億個token進行擬合，并在256個未見過的數據配比體上對1M、60M模型（各自以10億個token訓練）以及64個未見過的數據配比體上對1B模型（各自以250億個token訓練）進行評估。

評估指標。論文采用三種不同指標來基準測試論文的回歸模型：(1) 斯皮爾曼等級相關系數（ρ）是一種非參數度量方法，用于評估兩個排序變量之間關聯的強度和方向。(2) 皮爾遜相關系數（I）則是衡量兩個變量間線性關系的相關性。（3）均方誤差（MSE）是一種常用指標，用于通過測量預測值與實際值之間的平均平方差來評估回歸模型。

圖 3：Spearman 秩相關系數 rho 在不同訓練token數量和不同代理模型數量下，線性回歸（左）和 LightGBM 回歸（右）預測秩與真實秩之間的相關性圖。如圖所示，增加代理模型數量顯著提升 rho，而增加更多訓練token則收益遞減。          

2.3.2 實驗結果

模型大小間的強相關性。如表 2 所示，LightGBM 模型在所有三個指標上均優于線性回歸模型，并且在對具有更多訓練token的大型模型進行評估時，其優勢變得更加明顯。同時，使用 1B token訓練的 1M 模型在未見過的 1B 模型與 25B token配比上達到 97.12% 的高相關性，直接驗證了論文的秩不變性假設。

代理模型數量優于訓練token數量。在相同的小規模訓練 FLOPs 預算下，論文可以增加token數量（即訓練token數量）或代理模型數量。因此，論文研究哪種方法會產生更好的性能。如圖 3 所示，增加代理模型的訓練token在約 0.25B token后飽和。相比之下，增加代理模型數量持續提升性能，尤其是對于 LightGBM 模型。值得注意的是，在 0.2B token上訓練的 512 個模型的性能超過了在 0.8B token上訓練的 128 個模型，表明增加代理模型數量比在超過某個token閾值后增加訓練token數量更為有效。

2.4 在下游任務上的評估

論文將方法應用于現實下游任務，以展示其有效性。在評估時，論文排除了根據先前工作 [36] 和論文的觀察報告的性能跟蹤中表現出較大性能差異 (例如，RIL) 的特定基準。最終，論文選擇了以下基準作為下游任務：Social IQA 、HellaSwag 、PiQA 、OpenBookQA 、Lambada 、SciQ 、ARC Easy 、COPA 、RACE、LogiQA 、QQP、WinoGrande 和 MultiRC。這些基準涵蓋了多樣化的任務，能夠全面評估 REGMIX 在現實世界中的影響。對于每個基準，如果 lm-eval-anrnss 提供歸一化準確率，則使用該指標，否則使用常規準確率。

2.4.1 數據配比顯著影響下游性能

最初，論文訓練了64個模型，每個模型具有10億參數，使用不同的數據配比。每個模型都在Pile數據集[18]的250億個詞元上進行訓練，詞元的分配基于其相應領域的權重。表3展示了每個下游任務中最差和最佳模型的性能。報告的性能是從零樣本到五樣本評估的平均值，使用lm-eval-harness評估框架[19, 4]進行評分。論文發現數據配比顯著影響下游性能，其中最大的性能\triangle在Lambada任務上達到14.6。這突顯了研究最優數據配比的重要性。    

表3：論文實驗了64個模型，每個模型具有10億參數，訓練于不同的數據配比，并在各種基準上評估它們的性能。每個任務的報告性能是從零樣本到五樣本設置的平均分數，遵循Muennighoff等人[41]。

圖4：64個10億參數模型在各領域驗證損失與下游性能的相關性。注意，在計算相關性時，論文取損失值的負值，因為這使得可視化更直觀。這同樣適用于圖6。    

2.4.2 網絡語料庫對下游任務性能的提升最為顯著

論文在圖4（a）中展示了論文64個1B模型在不同領域的驗證損失與其在各種下游任務上的表現之間的相關性。在可視化之前，論文假設維基百科（en）子集上的驗證損失會與大多數下游任務顯示出強烈的關聯，因為它是一個高質量的數據集，并且許多下游任務源自維基百科文本。同樣，先前的工作通常將WikiText [38]作為標準基準來指示語言模型的性能。

然而，出乎意料的是，Pile-CC數據集上的驗證損失與大多數下游任務顯示出最強的相關性。例如，HellaSwag任務與Pile-CC驗證損失之間的相關系數非常接近1.0。這一意外結果挑戰了傳統的假設，即WikiText是評估深度分布式模型（DDMs）最具代表性的數據集。此外，這一結果與先前研究[17, 24]的發現相符，這些研究發現網絡數據集上的驗證損失與下游性能密切相關。

此外，論文分析了模型在C4100Domain驗證集上的損失之間的相關性，該驗證集取自C4數據集[47]，并被認為與Pile-CC共享相似的分布，因為它們都源自CommonCrawl語料庫。由于CommonCrawl是一個多樣化領域的集合，論文原本預期每個領域的損失與下游任務之間的相關性會有所不同。然而，令人驚訝的是，超過85%的領域都與Pile-CC表現出非常強的相關性（完整的相關性圖表見附錄D）。這一點在www.ign.com域名上得到了很好的例證，如圖4（b）所示，它與Pile-CC的整體相關性圖表非常相似。這也表明，Pile-CC與下游任務性能之間的高相關性可能歸因于它對各種主題和領域的廣泛覆蓋。

表4 不同數據選擇方法的性能比較。Human指的是The Pile[18]中提出的權重，Pile-CC Only表示僅在Pile-CC組件上進行訓練，而DoReMi則是指Xie et al.[64]中的權重。每個任務報告的性能是五個不同運行中從O-shot到5-shot設置的平均分數及標準差。論文估算了達到訓練數據配比所需計算量（以FLOPs衡量）。對于每個任務，顯著優于Human基線的分數以粗體突出顯示，顯著性由Cohen's確定。為了全面評估，論文還報告了使用LightEval的評估結果，遵循Penedo et al.[44]在附錄C.4中的設置。LightEval結果顯示，REGMIX略優于DoReMi和Pile-CC Only。    

2.4.3 通過REGMIX的數據配比提升下游性能

先前的工作表明，數據配比方法能夠通過使用較少的訓練token實現更小的驗證損失（或困惑度）來加速LLM預訓練[64]。然而，一個關鍵問題是應該優化哪個驗證損失？最直觀的方法，也是先前工作所采用的方法，是盡量減少所有領域的損失。然而，基于論文對100萬條訓練日志的研究，論文發現這在實踐中幾乎是不可能實現的。沒有任何數據配比能夠同時超越所有領域驗證損失上的人工選擇。這表明，簡單地最小化所有領域的損失很可能是不可行的。因此，論文選擇優化Pile-CC驗證損失，以期在下游任務上實現普遍性能提升，因為它與下游性能顯示出最高的相關性。

論文實施了兩種方法來確定數據配比。第一種方法依賴于人類的直覺。由于Pile-CC與其自身分布應該是最接近的匹配，論文假設僅在Pile-CC上進行預訓練可能會比基線產生更好的性能。第二種方法利用REGMIX，以Pile-CC驗證損失作為目標變量。論文采用LightGBM來預測能夠最小化Pile-CC驗證損失的數據配比。

論文比較了論文提出的方法與強基線的性能，包括由人類為Pile [18]進行的篩選，以及DoReMi [64]。對于DoReMi，論文直接從其報告的最佳域權重中獲取數據配比，并重新歸一化所有可用的\textbfI /域。這可能導致DoReMi的性能相對于原始報告結果有所下降。如表4所示，Pile-CC Only和REGMIX都顯示出與基線相比的強勁性能。在廣泛使用的HellaSwag基準測試中，REGMIX比人類選擇提高了6.8。此外，REGMIX在8項任務性能中擊敗了所有其他三種方法。    

圖5：左：預訓練語料庫中包含Pile-CC的不同方法在Pile-CC上的驗證損失。右：預訓練中排除Pile-CC的情況下，在Pile-CC上的驗證損失。

圖 6：使用線性回歸模型對不同目標域驗證損失與訓練域權重之間相關性的可視化。左側基于 Pile 數據集，右側基于 Stack 數據集。高度相關表明增加訓練域權重對降低目標域驗證損失有積極影響。

在 14 個案例中，Pile-CC Only 表現突出，平均得分最高。Pile-CC Only 的驚人強勁表現進一步強化了論文上一節的結論：網絡語料庫對下游性能有益。最終，REGMIX 超越了表 3 中的最佳模型，證明論文的自動數據配比方法比隨機搜索更高效。

盡管 Pile-C 驗證損失是下游性能的有益指標，但它可能不適用于所有感興趣的任務。有時論文無法假設驗證集與訓練集數據分布相似，而是面臨分布外場景。為驗證論文的方法在分布外場景中的有效性，論文完全排除 Pile-CC 域，并利用剩余域找到最小化 Pile-CC 驗證損失的最優數據配比。如圖 5（右）所示，論文提出的方法仍優于基線方法。這表明 REGMIX 無論目標域是否分布內或分布外都具有魯棒性。在此設置下，論文還提供了回歸評估結果，如圖 5 所示。

2.4.4 域間交互對人類理解具有挑戰性

為了理解不同域之間的影響，論文在圖 6 中可視化了線性回歸模型的系數 ( \omega)。該可視化揭示了各數據域如何相互貢獻，展示了它們之間的復雜交互。論文還展示了基于 Stack 數據集 [28] 訓練的每個 1M 代碼模型的代碼相關性圖。令人驚訝的是，無論是域交互可視化還是代碼相關性圖，都顯示了復雜的關系，這些關系對于人類專家來說難以完全理解。例如，Pile 數據集中的 PhilPapers 域似乎在線性回歸建模下為所有其他域帶來了增益，這是一個挑戰直覺的人類理解的非顯而易見發現。這些可視化突顯了確定最優數據配比的內在復雜性，強調了論文自動化 REGMIX 方法在高效識別高性能配比方面的價值，而非僅依賴人類直覺。    

圖 7：跨多種數據配比的1M訓練日志可視化。

2.4.5 數據配比效應超越縮放定律

近期研究[68, 20]證明了數據配比縮放定律的可行性。然而，論文在第5.4節的發現表明，域權重與驗證損失之間的關系比縮放定律所暗示的要復雜得多。為了可視化這種復雜性，論文在圖7中繪制了所有1M訓練日志的實驗點。如果數據配比的縮放定律成立，論文期望看到所有域之間存在明顯的對數-對數線性關系。然而，論文的結果揭示了一個更為細致的圖景。例如，DM Mathematics域，可能由于其與其他域相比的獨特分布，顯示出損失與域權重之間接近對數-對數線性關系。相比之下，大多數域如Pile-CC顯示出更復雜的模式，預測驗證損失變得不簡單。如圖所示，域間交互似乎很復雜，僅基于域在配比中的權重來預測驗證損失變得具有挑戰性。這些發現表明，盡管縮放定律提供了有價值的見解，但它們可能無法完全捕捉數據配比動態的復雜性。論文的方法通過將整個數據配比作為回歸模型的輸入來應對這一挑戰，提供了一個更全面的框架，用于理解和預測驗證損失，同時兼顧所有域權重。    

本文轉載自 ??AI帝國??，作者： 無影寺