機器之心報道
機器之心編輯部
鑒于語言模型預訓練成本巨大，因而研究者一直在尋找減少訓練時間和成本的新方向。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 中觸發梯度裁剪。