最近閱讀了《A Theory on Adam Instability in Large-Scale Machine Learning 》這篇論文。比較全面的闡述了100B以上的大模型預訓練中出現loss spike的原因(loss 突然大幅度上漲)，并介紹了一些可能的解決辦法。論文寫的非常精彩，但整體上有點散和深，我嘗試著站在工業立場上把它串一下

突刺是什么

首先介紹一下什么是loss spike：

loss spike指的是預訓練過程中，尤其容易在大模型（100B以上）預訓練過程中出現的loss突然暴漲的情況

圖片

如圖所示模型訓練過程中紅框中突然上漲的loss尖峰 loss spike的現象會導致一系列的問題發生，譬如模型需要很長時間才能再次回到spike之前的狀態（論文中稱為pre-explosion），或者更嚴重的就是loss再也無法drop back down，即模型再也無法收斂

PaLM和GLM130b之前的解決辦法是找到loss spike之前最近的checkpoint，更換之后的訓練樣本來避免loss spike的出現。

突刺成因分析

這篇論文（以下稱本文）對loss spike的出現原因做了十分詳細的分析，最后認為預訓練使用的Adam優化器是導致這個現象出現的重要原因之一

Adam算法是牛頓下降法的一個迭代逼近 

圖片

一切顯得十分完美，但是理想很豐滿，現實很骨感，收斂過程并不是一帆風順的

圖片

非穩態

圖片

中間態

圖片

穩態 

進入正態分布的穩態之后，理想的更新參數變化趨勢應該是方差越來越小，所有更新參數逐漸向0靠近。這應該是一個單向的過程，即穩定的單峰狀態（unimodal）不會再次進入非穩定的雙峰狀態(bimodal)，但事實并非如此，更新參數會再次進入非穩定的雙峰狀態

本文在理論層面做了研究和解釋，從中心極限定理（可以結合道爾頓板實驗理解）出發，認為隨機事件的疊加進入單峰的正態分布的必要條件之一是各個隨機事件事件之間應該是相互獨立的，但是梯度變化以及更新參數的變化并不能特別好的滿足獨立性這一條件，而這一點恰恰是導致更新參數振蕩，loss spike出現以及loss 不收斂的重要原因之一

圖片

造成梯度變化不獨立的原因（1、淺層參數長時間不更新2、batch太大，后期梯度更新趨于平穩） 上述的理論有些晦澀，本文作者可能也了解這一點，之后開始直接點題，結合實驗觀察拋出了重要現象和結論

即訓練過程中loss spike的出現與：梯度更新幅度， 大小，batch大小這三個條件密切相關

本文作者對loss spike出現時模型的前后變化做了仔細拆解，發現下列一系列連續現象的出現導致了loss spike：

圖片

1. 當前模型處在穩態（健康狀態），即單峰的正態分布狀態，并且梯度值 ，此時loss平穩，訓練過程正常

圖片

2.模型淺層(embedding層)梯度  ，這一般是由于訓練一段時間之后，淺層的語義知識表示此時一般已經學習的較好。但此時深層網絡（對應復雜任務）的梯度更新還是相對較大

圖片

3.一段時間淺層(embedding層)梯度  之后會導致  ，  。此時趨于0。因此導致淺層參數得不到更新（也對應于上述參數更新事件不獨立的原因）

圖片

4.此時雖然淺層(embedding層)參數長時間不更新，但是深層的參數依然一直在更新。長時間這樣的狀態之后，batch之間的樣本分布變化可能就會直接導致淺層(embedding層)再次出現較大的梯度變化（可以想象成一個水壩蓄水太久終于被沖開了。至于小模型為什么不會出現這種情況，推測是小模型函數空間小，無法捕獲樣本的分布變化，越大規模的模型對樣本之間不同維度的特征分布變化越敏感），此時 ， 再次集中在 附近（此時 , ），變成雙峰的非穩定狀態，本文提到了淺層(embedding層)這種突然的參數變化可能造成模型的連鎖反應進而出現loss spike的現象（這也對應了更換樣本重新訓練有可能會減少loss spike的出現頻率，實際上就是選擇分布變化較小的樣本，減小淺層梯度變換幅度）

圖片

5.這個階段模型處于非穩態，梯度變化幅度較大，每一次的梯度變化和更新參數變化事件之間又出現了一定的獨立性，因此經過一定的時間之后模型有可能再次進入穩態，loss再次drop back down（注意，本文著重提了這個再次drop back down并不是一定出現的，也很有可能loss長期處于flat狀態，再也無法收斂）

因此我們得出一些結論，loss spike的出現和淺層的梯度更新幅度， 大小密切相關（batch大小帶來的相關性問題倒是顯得沒那么大說服力），實際上就是淺層網絡參數突然進入到了之前長時間不在的狀態與模型深層參數當前的狀態形成了連鎖反應造成了模型進入非穩態。同時一般情況即使出現loss spike也會自動回復到正常狀態，但也有可能再也不會

突刺解法

本文最后提到了防止loss spike出現的一些方法：

1.如之前提到的PaLM和GLM130B提到的出現loss spike后更換batch樣本的方法（常規方法，但是成本比較高）

2.減小learning rate，這是個治標不治本的辦法，對更新參數的非穩態沒有做改進

3.減小 大小。或者直接把 設為0，重新定義

在等于0時候的值（這應該是個值得嘗試的辦法）

值得一提的是智譜華章在本文發表之前，在去年的GLM130B訓練時似乎也觀察到了淺層梯度變化和loss spike相關這一現象（GLM-130B: An Open Bilingual Pre-trained Model），他采取的是把淺層梯度直接乘以縮放系數 來減小淺層梯度更新值

出自130b

圖片

其實這塊我有個自己的想法，和是否也可以做衰減，隨著訓練過程逐漸減小，來避免loss spike的現象

另外假設我們能一次性加載所有樣本進行訓練（實際上不可能做到），是否還會出現loss spike的現象

最后目前流行的fp8，fp16混合訓練，如果upscale設置的過小，導致梯度在進入優化器之前就下溢，是不是會增加淺層梯度長時間不更新的可能性，進而增加loss spike的出現的頻率。（這么看來似乎提升upscale大小以及優化 大小是進一步提升模型效果的一個思路）