如何更好地繼續(xù)預訓練（Continue PreTraining）

寫在前面

預訓練（Pretraining）是一個非常消耗資源的工作，尤其在 LLM 時代。隨著LLama2的開源，越來越多人都開始嘗試在這個強大的英文基座模型上進行中文增強。但，我們如何才能保證模型在既學到「中文知識」的情況下，又不丟掉原有的「英文知識」呢？

今天給大家?guī)硪黄?Continue Pretraining 的論文（來自何枝大佬，知乎@何枝），Continual Pre-Training of Large Language Models: How to (re)warm your model?

知乎：https://zhuanlan.zhihu.com/p/654463331
paper//arxiv.org/pdf/2308.04014.pdf

1.實驗設定

作者使用一個 410M 大小的模型 Pythia，已經(jīng)在 Pile數(shù)據(jù)上進行過預訓練，然后在下游數(shù)據(jù)集 SlimPajama 上進行微調。

論文中直接使用 Loss 作為評估指標，即 Loss 越小，說明在上游（或下游）任務中的效果越強。

Pythia:https://huggingface.co/EleutherAI/pythia-410m-v0
Pile:https://huggingface.co/datasets/EleutherAI/pile
SlimPajama:https://www.cerebras.net/blog/slimpajama-a-627b-token-cleaned-and-deduplicated-version-of-redpajama

2. 關鍵結論

2.1 warmup 的步數(shù)不會影響最終性能

warmup 是一種 finetune 中常用的策略，指學習率從一個很小的值慢慢上升到最大值。那么，這個「慢慢上升」的階段持續(xù)多久是最好的呢？

作者分別使用訓練步數(shù)的：0%, 0.5%, 1%, 2% 這 4 種不同預熱步數(shù)來進行實驗：

從上圖中可以看到：當模型經(jīng)過「充分」訓練后，不管多長的預熱步數(shù)最后的性能都差不多。

但，這種前提是「充分訓練」，如果只看訓練前期的話，使用更長的預熱步數(shù)（黃色的線）。無論是「上游任務」還是「下游任務」，模型的 Loss 都要比其他預熱步數(shù)要低（下游學的快，上游忘的慢）。

2.2 學習率越大，下游任務越好，上游任務越差

為了探究學習率對學習效果的影響，作者使用了 4 種不同的最大學習率進行對比實驗，

此外，還比對了和從頭訓練（from scratch）模型的效果：

從圖中可以看到：經(jīng)過充分訓練后，學習率越大（紫色），下游性能最好，上游性能最差（忘得最多）。同樣，我們看前期訓練，盡管紫色線條在最后的 loss 是最低的，但在前期 loss 會增加的非常大，隨后下降。

PS：解釋一下這里為什么這么關注訓練前期，是因為在真實訓練中，我們可能不一定會增強圖中所示的 250B 這么多的 tokens，尤其是在模型參數(shù)很大的情況中。所以，當資源不允許充分訓練的情況下，較小的學習率和較長的 warmup 步數(shù)可能是一個不錯的選擇。

此外，圖中還能看出：未經(jīng)過預訓練的模型（藍色）無論是上游任務還是下游任務，都不如預訓練過的模型效果。

這鼓勵我們今天在進行訓練任務時，最好選擇一個已經(jīng)過預訓練的模型上繼續(xù)訓練（以利用其中的先驗知識）。

2.3 在初始預訓練中使用 Rewarmup 會損傷性能

盡管 warmup 策略在 Finetune 和 Continue Pretraining 中都起到了更好的效果（相較于常量學習率），但是，這建立在「切換了訓練數(shù)據(jù)集（數(shù)據(jù)分布）」的前提下。

作者做了一個實驗，不切換數(shù)據(jù)集，而是繼續(xù)在之前的「預訓練數(shù)據(jù)集（The Pile）」上繼續(xù)訓練：

從圖中結果可以發(fā)現(xiàn)：無論使用多大學習率的 warmup 策略，效果都不如使用常量學習率。

這進一步證明，在原數(shù)據(jù)集上使用 warmup 接著訓練會造成性能損傷，學習率越大則損傷越大，且這種損傷是無法在后續(xù)的訓練中被找回的。

PS：這里提示我們，當預訓練中遇到了訓練中斷需要繼續(xù)訓練時，我們應該在重新開始訓練時將學習率恢復到中斷之前的狀態(tài)（無論是數(shù)值還是衰減率）。

3. 實驗限制

作者在論文的最后貼出了得出上述結論的一些局限性。

3.1 上下游數(shù)據(jù)分布類似

因為實驗中選用的上游數(shù)據(jù)集 [Pile] 和下游數(shù)據(jù)集 [SlimPajama] 中存在一部分的數(shù)據(jù)重疊，

所以導致上下游數(shù)據(jù)的分布是比較相似的，但在我們真實的訓練任務中，上下游數(shù)據(jù)的差異可能會比這大的多。

3.2 模型規(guī)模較小

論文中所使用的模型規(guī)模為 410M，這和今天人們 7B 起步的 LLM 規(guī)模相差甚遠。

不過該團隊打算在接下來的工作中繼續(xù)在 3B 和 7B 的規(guī)模下進行嘗試，期待他們最后的實驗結論。