為了提高LLM底座的通用能力，通常預訓練數據都會包含有各種領域的數據，比如The Pile [2] 就是一個800GB大小的，涵蓋了22個不同領域的常用預訓練數據集，如Fig 1所示。對于LLM預訓練而言（采用next token prediction的自回歸方式進行預訓練），不同領域數據的配比很重要，之前的工作大部分基于啟發式的方法拍定配比（比如均勻采樣，或者根據不同領域數據的數據量大小進行采樣），由于不同領域數據的學習難度各不相同，這些啟發式的配比方法不能保證LLM在預訓練階段充分利用數據。本文旨在利用一個小規模的代理模型（280M參數量），通過Group DRO (Distributionally Robust Optimization) [3-4] [腳注1] 的方式去尋找最佳的數據配比，然后在全量參數的LLM（8B參數量）上進行預訓練，通過這種方式，作者發現預訓練速度能加速2.6倍，并且最終能獲得6.5%的效果提升，最終整個方法被稱之為 DoReMi （Domain Reweighting with Minimax Optimization），接下來讓我們具體看看這個方法是怎么實施的。
Fig 1. The Pile數據集中包含有22個不同領域的數據。
首先，我們先理解一個點：LLM在通過自回歸進行預訓練的時候，其最理想狀態應該是能擬合不同領域數據的分布，也就是各個領域數據的困惑度都最小化 [腳注2]。因此，如果我們能在“監控”預訓練過程中的各個領域數據中，通過損失大小確定擬合程度最差的樣本來自于哪個領域，然后適當提高這個領域數據的采樣率，那么理論上LLM就能有更多機會去學習較難擬合的領域的數據分布，從而以更小的訓練代價實現更好的擬合效果。當然，完整的LLM參數量動輒10億（B）規模，訓練一次的成本較高，我們假設參數量較小的（億級別，如280M）LLM模型的預訓練趨勢和完整參數量（8B）的LLM模型是類似的，因此可以共享相同的數據配比，我們稱此處的較小參數量模型為代理模型，通過代理模型我們能找出最佳數據配比，然后給完整參數模型進行預訓練，整個過程可以參考Fig 2.示意。
由于需要判斷擬合程度最差的樣本，這是一個“比較”的過程，比較的一方是代理模型（proxy model），而比較的另一方則是一個參考模型（reference model） ，參考模型指的是在默認的啟發式數據配比下，用和代理模型同樣參數量的模型（也就是280M），預訓練得到的模型。
Fig 2. DoReMi的基本框架示意圖，通過一個代理模型對預訓練數據集合中不同領域數據的采樣率進行重新參數化，得到最佳數據配比后去訓練最終的LLM。
至此，我們需要了解的DoReMi的基本思路已經介紹完了，我們開始深入細節。首先先進行形式化表示，假設預訓練數據集中具有k個域的數據，表示為Di,i=1,?,k，每個域的權重記為α∈Δk，可以理解為在k個域中的采樣概率分布，也就是∑ikαi=1。那么訓練數據的分布可以表示為：
(1)Pα=∑i=1kαi⋅unif(Di)
其中unif(D)=1|D|∑x∈Dδx是一個在數據集合D上的均勻分布，如果在x=x′ 的時候 δx(x′)=1，否則δx(x′)=0。整個DoReMi的過程，可以描述為：
第一步，訓練一個參考模型
首先采用初始化的數據配比αref先訓練一個280M參數量的參考模型pref，此處的數據配比可以根據啟發式的方式選取，比如采用數據集中不同領域數據的數據量比例作為配比，見Fig 1.的baseline一列。
第二步，通過Group DRO的方式訓練一個代理模型并且得到新的域權重
有了參考模型pref之后，就可以開始訓練代理模型和學習新的域權重了（也即是學習得到α¯），整個過程采用DRO-LM的框架 [5] ，采用Group DRO優化器進行訓練，其中的θ為代理模型的參數。整個框架的公式如公式(2)所示，不難看出，這是一個minmax過程，其優化目標正如我們前文討論的，優化（體現在min目標上）各域上最差擬合程度樣本（體現在max目標上）的損失，我們著重講解下這個公式。
(2)minθmaxα∈ΔkL(θ,α):=∑i=1kαi⋅[1∑x∈Di|x|∑x∈Di?θ(x)−?ref(x)]
其中的?θ(x)=−log?pθ(x)和?ref(x)=−log?pref(x)為代理模型和參考模型的負對數似然概率 [腳注3]，|x|是當前樣本x的token數量。?θ(x)−?ref(x) 可視為是超額損失（excess loss），這個損失度量了對于代理模型而言，當前樣本x的可優化空間。超額損失越大，說明該樣本需要更多訓練才能達到參考模型的效果；超額損失越小，則有兩種可能，第一種可能是?ref很大，這表明了這個是一個高熵的樣本，其生成會趨向于均勻分布，因此難以學習，第二種可能是?θ很小，這說明對于樣本x而言，代理模型已經學習得足夠好了，屬于簡單樣本，這兩種情況下都需要適當減少權重，減少該域樣本的采樣。首先從這個minmax的內層目標（也即是max）看起，此時代理模型不進行參數更新，優化項是α∈Δk， 就是根據超額損失的大小去啟發式更新權重，然后看到外層的min目標，此時優化項是代理模型參數θ，也即是從內層找到了最大的超額損失后，嘗試讓代理模型去擬合這個超額損失。通過交替地進行最小最大優化，從而訓練得到一個代理模型和新的域權重α¯。 最終的域權重，則是從每一步的訓練過程中進行累計平均得到，即是α¯=1T∑t=1Tαt。
這個過程，用算法描述的形式表述則是：
輸入：各個域的數據D={D1,?,Dk}，訓練步數T，batch size大小b和更新步長η ，平滑系數c∈[0,1]，本實現采用的c=10−3。
- 初始化代理模型參數θ0
- 初始化域權重α0=1k1∈Rk
對于t從1到T開始循環：
- 從 Pu 中采樣一個大小為 b 的小批量B={x1,?,xb}，其中u=1k1
- 令|x|為樣本x的token長度
- 計算每個域i∈{1,?,k}的超額損失（?∗,j(x)是第j個token的損失），此處的max(?,0)就是在實現公式(2)里面提到內層
max過程，需要保證超額損失的非負性，其中的?∗,j(x)=−log?p∗(xj|x1,?,xj−1)。(3)λt[i]←1|x|∑x∈B∩Di∑j=1|x|max(?θt−1,j(x)−?ref,j(x),0) - 啟發式地更新域權重（指數更新）： α′←αt−1⊙exp?(ηλt)
- 更新歸一化和平滑的域權重：αt←(1−c)αt′∑i=1kαt′[i]+cu,αt∈Rk， 此處采用平滑更新的方式，是希望整個域權重更新的過程更加平滑（因為每次只見到了當前batch的數據，因此可能存在噪聲），最好是在先驗分布u的基礎上進行增量更新。
- 使用目標L(θt−1,αt)更新代理模型的參數θt（可以采用Adam優化器）。
結束循環
返回：α¯=1T∑t=1Tαt
第三步，用新的權重訓練一個完整的LLM
采用第二步得到的α¯構建新的訓練分布Pα¯，從中采樣數據去預訓練最終需要的完整參數量的LLM（本實驗中是8B）。
迭代式的DoReMi
第一輪迭代的DoReMi的αref是啟發式得到的，不是最優的選擇，因此整個DoReMi過程是可以迭代進行的，即是當第一輪迭代中得到了α¯1之后，可以將第二輪迭代的αref:=α¯1，然后重復整個DoReMi的過程，直到域權重收斂為止，在本文收斂的定義是||α¯−αref||∞<10−3，此處的無窮范數即是α¯和αref差的最大值。從作者的經驗看，在GLaM數據集上只需要3輪迭代即可收斂。
介紹完了整個DoReMi的操作過程，我們看下實驗結果。作者是在The Pile和GLaM [6] 這兩個預訓練數據集上進行預訓練的，The Pile的情況前文介紹了，GLaM是一個具有8個域的文本數據集，由于這個數據集的域權重根據下游任務的效果而調節得到，因此GLaM的域權重可以視為是真實標簽，由此檢驗DoReMi的域權重調整結果。從Fig 1.中可以看出，經過DoReMi后，在Pile數據集上不同領域的權重已經有了很大的變化，有些域的權重存在大幅度的下降，然而如果我們看到DoReMi (280M -> 8B) 在Pile數據集上保留驗證集上所有域的困惑度，如Fig 3.所示，則會發現DoReMi在所有域上的困惑度都是明顯下降。這并不是很符合直覺，因為某些域（如Arxiv、PubMed central）的權重下降了很多，意味著LLM預訓練過程中采樣到這些域數據的幾率下降了，為什么還能得到困惑度的下降呢？
一種可能性是，正如在前文討論的，超額損失低的那部分樣本都是難樣本（接近均勻分布）或者簡單樣本，簡單樣本容易學習，不需要那么多樣本支持學習，而難樣本則由于難度太高，即便提高了樣本量也無法學習出來，因此降低采樣也沒帶來效果損失。并且很幸運的，提高了其他中度難度樣本的采樣比例后，讓模型泛化能力得到了進步，因此在各個域數據上的表現都沒折損，都有提升。
Fig 3. 對比基線，采用了DoReMi后在Pile數據集得到所有域的保留驗證集上都得到了困惑度的明顯下降。
讓我們看看DoReMi在下游任務上的表現，從Fig 4. (a) 中能發現，在利用The Pile預訓練集的情況下，采用了DoReMi后，在所有的訓練步中，性能（下游任務指標，用的是精準匹配的比例，筆者覺得可能類似ROUGE的指標）都能得到持續的提升。從Fig 4. (b) 是采用GLaM數據集作為預訓練的結果，有以下結論：
- 采用多輪迭代的DoReMi（round 1 vs round 2），采用多輪迭代的效果會持續比單輪的好。
- 采用了單輪的DoReMi效果沒有基線好，可能是由于GLaM數據集本身的域只有8個，DoReMi的發揮空間不如The Pile數據集。而采用了多輪DoReMi的效果能超越基線，可能說明對于域比較少的數據集，需要多輪迭代才能得到較好效果。
- 采用多輪迭代的DoReMi，其效果接近最佳權重（通過下游任務調優得到）的效果。
再看到Fig 4. (c)， 這是在GLaM數據集中多輪迭代DoReMi的權重和最佳權重的對比，可以發現采用了DoReMi后，的確權重都趨向于了最佳權重了，這也證實了DoReMi的合理性和有效性。
Fig 4. DoReMi在下游任務中的模型性能對比
作者同樣做了消融試驗，去回答DoReMi中幾個關鍵問題：
- Q1：是否挑選最難的樣本或者最簡單的樣本，而不是超額損失最大的樣本效果會更好？A1：當前的挑選標準是超額損失，即是?θ(x)−?ref(x)，如果挑選標準變成最難樣本，即是?θ(x)，或者挑選最簡單樣本，即是−?ref(x)，試驗效果會如何呢？見Fig 5. 右側所示，我們發現采用了最簡單樣本或者最難樣本的試驗，效果都不如DoReMi，并且采用了最簡單樣本的方案，對比基線都遠遠落后。這說明了，學習最簡單樣本明顯不是一個好主意，這使得大模型的底座能力缺失嚴重，單純學習最難的樣本也能提升LLM的能力，但是光是關注最難的樣本，而忽略了中等難度的樣本，則也不是最優的方案。這個也和前面分析困惑度的試驗結論遙相呼應了。
Fig 5. 在 The Pile 數據集上訓練的模型的平均下游準確度。左側: 在 DoReMi 中將參考/代理模型的規模從 70M 增加到 280M，可以提高 8B 主模型的下游準確度，但這一趨勢在 1B 代理模型中并未繼續。我們假設 Group DRO 優化器在更大規模的代理模型中表現不佳。右側: 僅針對最難或最容易的領域
- Q2：提高代理模型的參數尺寸，是否能獲得最后效果的提升？A2：考慮采用不同的代理模型尺寸，如70M、150M、280M、1B參數量，而最終模型的尺寸仍然是8B，是否會觀察到scaling law呢？如Fig 5. 左側圖片所示，適當提升代理模型尺寸（
70M -> 150M -> 280M）可以提高最終模型的效果，但是當代理模型尺寸達到一定大小（如1B）后，反而出現了性能下降。因此對于代理模型而言，也并不是尺寸越大越好。 - Q3：如果代理模型達到和最終模型同樣的尺寸，代理模型和最終模型的效果對比如何？A3：這個問題其實也很符合直覺，代理模型和最終模型采用的采樣策略是不同的（損失重參數化 vs 重采樣）。作者嘗試將代理模型和最終模型的參數量都設置為相同（為了試驗對比公平），然后對比基線、DoReMi (x -> x)和代理模型的表現，如Fig 6所示，我們發現采用了代理模型的表現都低于最終的主模型，并且隨著模型尺寸增大，性能差別則越大。并且在1B規模的代理模型中，甚至性能還不如基線（但是其DoReMi結果還是比基線好），這意味即便代理模型沒有訓練得很好，在整個DoReMi體系下仍然能提升最終模型的效果。
Fig 6. DoReMi 主模型和相同規模代理模型的困惑度，盡管 1B 代理模型的質量相對較低，但由此產生的領域權重仍然能夠改善主模型。
Reference
[1]. Xie, Sang Michael, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy S. Liang, Quoc V. Le, Tengyu Ma, and Adams Wei Yu. "Doremi: Optimizing data mixtures speeds up language model pretraining." Advances in Neural Information Processing Systems 36 (2024). aka DoReMi
[2]. Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The pile: An 800gb dataset of diverse text for language modeling. arXiv, 2020. aka The Pile
[3]. Arkadi Nemirovski, Anatoli Juditsky, Guanghui Lan, and Alexander Shapiro. Robust stochastic approximation approach to stochastic programming. SIAM Journal on optimization, 19(4):1574–1609, 2009.
[4]. Shiori Sagawa, Pang Wei Koh, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. In International Conference on Learning Representations (ICLR), 2020.
[5]. Yonatan Oren, Shiori Sagawa, Tatsunori Hashimoto, and Percy Liang. Distributionally robust language modeling. In Empirical Methods in Natural Language Processing (EMNLP), 2019. aka DRO-LM
[6]. Nan Du, Yanping Huang, Andrew M. Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, M. Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, Barret Zoph, Liam Fedus, Maarten Bosma, Zongwei Zhou, Tao Wang, Yu Emma Wang, Kellie Webster, Marie Pellat, Kevin Robinson, K. Meier-Hellstern, Toju Duke, Lucas Dixon, Kun Zhang, Quoc V. Le, Yonghui Wu, Zhifeng Chen, and Claire Cui. GLaM: Efficient scaling of language models with mixture-of-experts. arXiv, 2021. aka GLaM
腳注區域：
[腳注1]: Group DRO 的關鍵在于通過最小化最壞情況下的損失來優化領域權重，從而使得模型在所有領域上都能達到較好的性能。
[腳注2]: 在自然語言處理（NLP）中，困惑度（Perplexity）是評估語言模型性能的一個重要指標。它衡量模型對測試數據的預測能力，具體計算方法如下：對于一個給定的詞序列 W=(w1,w2,?,wT)，困惑度 PP(W) 的計算公式為：PP(W)=P(w1,w2,?,wT)−1T，其中的P(w1,?,wT)為語言模型建模的聯合概率分布，而T為序列長度。低困惑度: 如果語言模型能夠準確預測詞序列，那么它給出的聯合概率會較高，從而導致困惑度較低。這意味著模型對測試數據的預測能力較強。高困惑度: 如果模型對詞序列的預測能力較差，給出的聯合概率會較低，導致困惑度較高
[腳注3]: 此處計算的方法是，給定一個樣本x={x1,x2,?,xN}，其中N表示序列長度，xi表示token，那么p(xi|xi−1,xi−2,?,x1)就是當前token xi被預測正確的概率。通過求對數似然和，能得到log?p(x)，也即是對每個token的對數概率進行加和。這代表了這個序列x被當前語言模型采樣出來的概率。