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

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

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

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

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

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

這項研究也得到了大家的認可。英偉達人工智能科學家 Jim Fan 表示：「多年來，有無數論文推翻了 Adam 優化器，不知道 Sophia 會不會是保留到最后的那個，這是一個可擴展的二階優化器， 其偽代碼只有 13 行，在 GPT-2 （預訓練）上比 Adam 快了 2 倍，我很想試試這個優化器！」

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

下面我們看看該優化器具體是如何實現的。

方法介紹

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

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

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

該研究只考慮對角 Hessian 的正項，并在更新中引入按坐標裁剪，更新規則改寫為：

對角 Hessian 估計器

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

實驗 

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

實驗語言建模設置

該實驗在 OpenWebText 上訓練自回歸模型。遵循 GPT-2 的標準協議，將上下文長度設置為 1024。使用只有解碼器的 Transformer，模型參數量分別為 125M (小型)、355M (中型) 和 770M (大型)。

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

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

實驗結果

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

同樣 100K step，Sophia-H 在 770M 模型上的驗證損失小了 0.05 (圖 4，(c))。可以看出，這是個明顯的改進，因為根據該機制中的擴展定律和圖 5 中的結果，損失 0.05 的改進相當于實現相同驗證損失的 step 數或總計算量的改進的雙倍。

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

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

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

分析

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

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

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