鑒于語言模型預(yù)訓(xùn)練成本巨大，因而研究者一直在尋找減少訓(xùn)練時(shí)間和成本的新方向。Adam 及其變體多年來一直被奉為最先進(jìn)的優(yōu)化器，但其會(huì)產(chǎn)生過多的開銷。本文提出了一種簡單的可擴(kuò)展的二階優(yōu)化器 Sophia，在與 Adam 比較中，Sophia 在減少了 50% step 數(shù)量的情況下實(shí)現(xiàn)了與 Adam 相同的驗(yàn)證預(yù)訓(xùn)練損失。

大語言模型（LLM）的能力隨著其規(guī)模的增長而取得了顯著的進(jìn)展。然而，由于龐大的數(shù)據(jù)集和模型規(guī)模，預(yù)訓(xùn)練 LLM 非常耗時(shí)，需要進(jìn)行數(shù)十萬次的模型參數(shù)更新。例如，PaLM 在 6144 個(gè) TPU 上進(jìn)行了為期兩個(gè)月的訓(xùn)練，總共耗費(fèi)大約 1000 萬美元。因此，提高預(yù)訓(xùn)練效率是擴(kuò)展 LLM 規(guī)模的一個(gè)主要瓶頸。

本文來自斯坦福大學(xué)的研究者撰文《 Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training 》，文中提出了 Sophia（Second-order Clipped Stochastic Optimization）輕量級二階優(yōu)化器，旨在通過更快的優(yōu)化器提高預(yù)訓(xùn)練效率，從而減少達(dá)到相同預(yù)訓(xùn)練損失所需的時(shí)間和成本，或者在相同預(yù)算下實(shí)現(xiàn)更好的預(yù)訓(xùn)練損失。

論文地址：https://arxiv.org/pdf/2305.14342.pdf

Sophia 優(yōu)化器使用隨機(jī)估計(jì)作為 Hessian 矩陣對角線的 pre-conditioner，并采用剪切（clipping）機(jī)制來控制最壞情況下的參數(shù)大小更新。在像 GPT-2 這樣的預(yù)訓(xùn)練語言模型上，Sophia 與 Adam 相比，在減少了 50% step 數(shù)量的情況下實(shí)現(xiàn)了相同的驗(yàn)證預(yù)訓(xùn)練損失。

由于 Sophia 可以維持每個(gè) step 內(nèi)的的內(nèi)存和所用時(shí)間，這相當(dāng)于總計(jì)算量減少了 50%，wall-clock 時(shí)間減少了 50%(參見圖 1 (a) 和 (b))。此外，根據(jù)擴(kuò)展定律（從模型大小的角度），在 125M 到 770M 的模型上，Sophia 相對于 Adam 更具優(yōu)勢，并且隨著模型大小的增加，Sophia 和 Adam 差距也在增加（在 100K step 內(nèi)）（圖 1（c））。特別的，在 540M 參數(shù)模型上（100K step），Sophia 和具有 770M 參數(shù)模型上的 Adam 實(shí)現(xiàn)了相同的驗(yàn)證損失。需要注意的是，后者需要多 達(dá)40％的訓(xùn)練時(shí)間和 40％的推理成本。

這項(xiàng)研究也得到了大家的認(rèn)可。英偉達(dá)人工智能科學(xué)家 Jim Fan 表示：「多年來，有無數(shù)論文推翻了 Adam 優(yōu)化器，不知道 Sophia 會(huì)不會(huì)是保留到最后的那個(gè)，這是一個(gè)可擴(kuò)展的二階優(yōu)化器， 其偽代碼只有 13 行，在 GPT-2 （預(yù)訓(xùn)練）上比 Adam 快了 2 倍，我很想試試這個(gè)優(yōu)化器！」

論文作者之一、斯坦福大學(xué)助理教授馬騰宇表示：「（從發(fā)布之初，）Adam 可以說是一個(gè) 9 歲的優(yōu)化器，是訓(xùn)練 LLM 的首選，如 GPT-3、OPT、 LLAMA 等。而我們的新研究 Sophia 是一個(gè)嶄新的優(yōu)化器，在 LLM 上比 Adam 快了 2 倍。只需要你多寫幾行代碼，你的成本就能從 200 萬美元降至 100 萬美元（如果擴(kuò)展定律成立的話）。」

下面我們看看該優(yōu)化器具體是如何實(shí)現(xiàn)的。

方法介紹

至于該研究的動(dòng)機(jī)，作者表示 Adam 對于異構(gòu)曲率（heterogeneous curvatures）的適應(yīng)性不足。另一方面，vanilla Newton 方法在凸函數(shù)中具有最優(yōu)的 pre-conditioner，但對于負(fù)曲率和 Hessian 的快速變化容易受到影響。基于這些見解，該研究設(shè)計(jì)了一種新的優(yōu)化器 Sophia，它比 Adam 更適應(yīng)異構(gòu)曲率，比 Newton 方法更能抵抗非凸性和 Hessian 的快速變化，并且還使用了成本較低的 pre-conditioner。

方法理論方面，在時(shí)間步長 t 上，該研究用 θ_t 表示參數(shù)。在每個(gè) step 上，該研究從數(shù)據(jù)分布中采樣一個(gè)小批次，計(jì)算小批次損失，并用 L_t (θ_t) 表示。g_t 表示 L_t (θ_t) 的梯度，即。設(shè) m_t 為 EMA（ exponential moving average ）的梯度，則更新的分子為 m_t ← β_1m_t?1 + (1 ? β_1) g_t 。

Sophia 使用基于對角 Hessian 的 pre-conditioner，根據(jù)參數(shù)維度的曲率直接調(diào)整更新的大小。為了減少開銷，該研究僅在每 k 個(gè)step內(nèi)（現(xiàn)實(shí)中 k = 10）估計(jì)一次 Hessian。在時(shí)間步 t 上，估計(jì)器返回小批次損失的 Hessian 對角線的估計(jì)。每 k 個(gè)step更新一次 EMA，得到對角 Hessian 估計(jì)的以下更新規(guī)則：

該研究只考慮對角 Hessian 的正項(xiàng)，并在更新中引入按坐標(biāo)裁剪，更新規(guī)則改寫為：

對角 Hessian 估計(jì)器

該研究引入了兩個(gè)對角 Hessian 估計(jì)器，它們的內(nèi)存和運(yùn)行時(shí)間成本都與計(jì)算梯度相似。估計(jì)器分別為 Hutchinson 無偏估計(jì)器以及 GNB（ Gauss-Newton-Bartlett ） 估計(jì)器。偽代碼如下所示：

實(shí)驗(yàn) 

研究將使用 Hutchinson 估計(jì)器和 GNB 估計(jì)器的算法分別稱為 Sophia-H 和 SophiaG。本文用 GPT-2 評估了 Sophia 的自回歸語言建模，模型尺寸從 125M 到 770M 不等。結(jié)果表明，Sophia 在 step、總計(jì)算量和所有模型大小的 wall-clock 時(shí)間方面比 AdamW 和 Lion 快 2 倍。此外，擴(kuò)展定律更有利于 Sophia 而不是 AdamW。

實(shí)驗(yàn)語言建模設(shè)置

該實(shí)驗(yàn)在 OpenWebText 上訓(xùn)練自回歸模型。遵循 GPT-2 的標(biāo)準(zhǔn)協(xié)議，將上下文長度設(shè)置為 1024。使用只有解碼器的 Transformer，模型參數(shù)量分別為 125M (小型)、355M (中型) 和 770M (大型)。

基線：研究主要比較 Sophia 和 Adam。Adam 采用解耦權(quán)重衰減 (AdamW)，這是語言建模任務(wù)中主要使用的優(yōu)化器，而 Lion 是通過符號搜索發(fā)現(xiàn)的一階自適應(yīng)優(yōu)化器。所有優(yōu)化器都進(jìn)行了很好的調(diào)整。權(quán)重衰減被設(shè)置為 0.1，β_1 = 0.9， β_2 = 0.95。對于 Lion，使用 β_1 = 0.95 和 β_2 = 0.98。125M 和 355M 模型是在 10 個(gè) A5000 GPU 上訓(xùn)練的，而 770M 模型是在 8 個(gè) A100 GPU 上訓(xùn)練的。

評估：研究人員使用每個(gè)優(yōu)化器對模型進(jìn)行 100K、200K 或 400K 個(gè)step的預(yù)訓(xùn)練評估，以比較速度。值得注意的是，與標(biāo)準(zhǔn)一樣，LR 調(diào)度取決于預(yù)先指定的總目標(biāo)step 數(shù)，如圖 5 (a) 所示。這使得同一優(yōu)化器的損失曲線在不同的 step 數(shù)下是不同的，因?yàn)榭?step 數(shù)較少的 LR 調(diào)度會(huì)更早地衰減 LR。本文主要評估了這些模型在 OpenWebText 上的 log 困惑度，并繪制了損失曲線。此外，該研究還報(bào)告了 SuperGLUE 上的上下文學(xué)習(xí)結(jié)果，然后對 5 個(gè)提示的結(jié)果取平均值。

實(shí)驗(yàn)結(jié)果

圖 4 展示了相同 step 數(shù) (100K) 下 OpenWebText 上的驗(yàn)證損失曲線 (token 級 log 困惑度)。與 AdamW 和 Lion 方法相比，本文所提出來的方法獲得了更好的驗(yàn)證損失。隨著模型大小的增加，Sophia 和基線之間的差距也變得更大。Sophia-H 和 Sophia-G 在 355M 模型上的驗(yàn)證損失都小 0.04 (圖 4 (b))。

同樣 100K step，Sophia-H 在 770M 模型上的驗(yàn)證損失小了 0.05 (圖 4，(c))。可以看出，這是個(gè)明顯的改進(jìn)，因?yàn)楦鶕?jù)該機(jī)制中的擴(kuò)展定律和圖 5 中的結(jié)果，損失 0.05 的改進(jìn)相當(dāng)于實(shí)現(xiàn)相同驗(yàn)證損失的 step 數(shù)或總計(jì)算量的改進(jìn)的雙倍。

Sophia 在 step 數(shù)、總計(jì)算時(shí)間和 wall-clock 時(shí)間方面快了兩倍。Sophia 對驗(yàn)證損失的改進(jìn)在于減少 step 數(shù)或總計(jì)算量。在圖 1 (a) 和 (b) 和圖 5 中，通過比較達(dá)到相同驗(yàn)證損失水平所需的 step 數(shù)或總計(jì)算量來評估優(yōu)化器。從圖 1 (a) 和 (b) 中可以看出，與 AdamW 和 Lion 相比，Sophia-H 和 Sophia-G 在不同的模型尺寸下實(shí)現(xiàn)了 2 倍的加速。

擴(kuò)展定律更有利于 Sophia-H 而不是 AdamW。在圖 1 (c) 中，該研究繪制了預(yù)訓(xùn)練 100K step 的不同大小模型的驗(yàn)證損失。Sophia 和 AdamW 之間的差距隨著模型的擴(kuò)大而增大。此外，Sophia-H 訓(xùn)練的 540M 模型比 AdamW 訓(xùn)練的 770M 模型的損失更小。Sophia-H 訓(xùn)練的 355M 模型與 AdamW 訓(xùn)練的 540M 模型的損失相當(dāng)。

評估下游任務(wù)的小樣本 (SuperGLUE)。如圖 6 所示，驗(yàn)證損失的改善也使得下游任務(wù)準(zhǔn)確率的提高。在預(yù)訓(xùn)練 step 數(shù)相同的情況下，使用 Sophia 預(yù)訓(xùn)練的 GPT-2 medium 和 GPT-2 large 在大多數(shù)子任務(wù)上具有更好的少樣本準(zhǔn)確率。此外，用 Sophia-H 預(yù)訓(xùn)練的模型與用 AdamW 預(yù)訓(xùn)練的模型具有相當(dāng)?shù)男颖緶?zhǔn)確率。

分析

比較 wall-clock 時(shí)間與計(jì)算量。表 1 比較了每一個(gè) step 的總計(jì)算量 (TFLOPs) 和 A100 GPU 上的 wall-clock 時(shí)間。本文報(bào)告了每個(gè) step 的平均時(shí)間，Hessian 計(jì)算花費(fèi)的時(shí)間的總計(jì)算。較小的批量大小，即每 10 個(gè) step 以計(jì)算對角 Hessian 估計(jì)，Hessian 計(jì)算占總計(jì)算量的 6%，與 AdamW 相比，整體 wall-clock 時(shí)間開銷小于 5%。在內(nèi)存使用方面，優(yōu)化器 m 和 h 兩個(gè)狀態(tài)，這導(dǎo)致了與 AdamW 相同的內(nèi)存開銷。

在 30M 模型上，執(zhí)行網(wǎng)格搜索來測試 Sophia-H 對超參數(shù)的敏感性 (圖 7 (c))。所有組合的性能相近，但 β_2 = 0.99 和 ρ = 0.1 的性能最好。此外，這種超參數(shù)選擇可以跨模型大小遷移。對于 125M、355M 和 770M 的所有實(shí)驗(yàn)，都使用了 30M 模型上搜索超參數(shù) ρ = 0.01， β_2 = 0.99。

訓(xùn)練穩(wěn)定性。與 AdamW 和 Lion 相比，Sophia-H 在預(yù)訓(xùn)練中具有更好的穩(wěn)定性。梯度裁剪 (by norm) 是語言模型預(yù)訓(xùn)練中的一項(xiàng)重要技術(shù)。在實(shí)踐中，梯度裁剪觸發(fā)的頻率與訓(xùn)練的穩(wěn)定性有關(guān) —— 如果梯度被頻繁裁剪，迭代可能處于非常不穩(wěn)定的狀態(tài)。圖 7 (a) 比較了 GPT-2  (125M) 觸發(fā)梯度裁剪的 step 比例。盡管所有方法都使用相同的裁剪閾值 1.0，但 Sophia-H 很少觸發(fā)梯度裁剪，而 AdamW 和 Lion 在超過 10% 的 step 中觸發(fā)梯度裁剪。