訓練ViT和MAE減少一半計算量！ Sea和北大聯合提出高效優化器Adan，深度模型都能用

自Google提出Vision Transformer(ViT)以來，ViT漸漸成為許多視覺任務的預設backbone。憑藉著ViT結構，許多視覺任務的SoTA都得到了進一步提升，包括影像分類、分割、偵測、辨識等。

然而，訓練ViT並非易事。除了需要較複雜的訓練技巧，模型訓練的計算量往往也比之前的CNN大很多。近日，新加坡Sea AI LAB (SAIL) 和北大ZERO Lab的研究團隊共同提出新的深度模型優化器Adan，該優化器可以僅用一半的計算量就能完成ViT的訓練。

論文連結：https://arxiv.org/pdf/2208.06677.pdf

程式碼連結：https://github.com/sail-sg/Adan

此外，在計算量相同的情況下， Adan在多個場景（涉及CV、NLP、RL）、多種訓練方式（有監督與自監督）和多種網絡結構/演算法（Swin、ViT、 ResNet、ConvNext、MAE、LSTM、BERT、Transformer-XL、PPO演算法）上，均獲得了效能提升。

程式碼、設定檔、訓練log都已開源。

深度模型的訓練範式與最佳化器

隨著ViT的提出，深度模型的訓練方式變得越來越複雜。常見的訓練技巧包括複雜的資料增強（如MixUp、CutMix、AutoRand）、標籤的處理（如label smoothing和noise label）、模型參數的移動平均、隨機網路深度、dropout等。伴隨著這些技巧的混合運用，模型的泛化性與魯棒性均得到了提升，但是隨之而來的便是模型訓練的計算量變得越來越大。

在ImageNet 1k上，訓練epoch數從ResNet剛提出的90已經成長到了訓練ViT常用的300。甚至針對一些自監督學習的模型，例如MAE、ViT，預先訓練的epoch數已經達到了1.6k。訓練epoch增加意味著訓練時間極大的延長，急劇增加了學術研究或工業落地的成本。目前一個普遍的解決方案是增大訓練的batch size並輔助並行訓練以減少訓練時間，但是伴隨的問題便是，大的batch size往往意味著performance的下降，並且batch size越大，情況越明顯。

這主要是因為模型參數的更新次數隨著batch size的增加在急遽減少。目前的優化器並不能在複雜的訓練範式下以較少的更新次數實現對模型的快速訓練，這進一步加劇了模型訓練epoch數的成長。

因此，是否存在一種新的優化器能在較少的參數更新次數情況下更快更好地訓練深度模型？在減少訓練epoch數的同時，也能緩解batch size增加帶來的負面影響？

被忽略的衝量

要想加速最佳化器的收斂速度，最直接的方法就是引入衝量。近年來提出的深度模型最佳化器皆沿用Adam中所使用的衝量範式－重球法：

##其中g_k是隨機噪聲，m_k是moment，eta是學習率。 Adam將m_k的更新由累積形式換成了移動平均的形式，並引入二階moment（n_k）對學習率進行放縮，即：

然而隨著Adam訓練原始ViT失敗，它的改進版本AdamW漸漸地變成了訓練ViT甚至ConvNext的首選。但是AdamW並沒有改變Adam中的衝量範式，因此在當batch size超過4,096的時候，AdamW訓練出的ViT的表現會急劇下降。

在傳統凸優化領域，有一個與重球法齊名的衝量技巧－Nesterov衝量演算法：

Nesterov衝量演算法在光滑且一般凸的問題上，擁有比重球法更快的理論收斂速度，並且理論上也能承受更大的batch size。同重球法不同的是，Nesterov演算法不在當前點計算梯度，而是利用衝量找到外推點，在該點算完梯度以後再進行沖量累積。

外推點能幫助Nesterov演算法提前感知當前點周圍的幾何資訊。這種特性使得Nesterov衝量更適合複雜的訓練範式和模型結構（如ViT），因為它並不是單純地依靠過去的衝量去繞過尖銳的局部極小點，而是透過提前觀察周圍的梯度，調整更新的方向。

儘管Nesterov衝量演算法擁有一定的優勢，但是在深度優化器中，卻鮮有被應用與探索。其中一個主要的原因是Nesterov演算法需要在外推點計算梯度，在當前點更新，期間需要多次模型參數重載以及需要人為地在外推點進行back-propagation (BP)。這些不便利性大大限制了Nesterov衝量演算法在深度模型最佳化器中的應用。

Adan優化器

透過結合改寫的Nesterov衝量與自適應最佳化演算法，並引入解耦的權重衰減，可以得到最終的Adan優化器。利用外推點，Adan可以提前感知周圍的梯度訊息，從而有效率地逃離尖銳的局部極小區域，以增加模型的泛化性。

1) 自適應的Nesterov衝量

#為了解決Nesterov衝量演算法中多次模型參數重載的問題，研究人員首先對Nesterov進行改寫：

#可以證明，改寫的Nesterov衝量演算法與原始演算法等價，兩者的迭代點可以相互轉化，且最終的收斂點相同。可以看到，透過引入梯度的差分項，已經可以避免手動的參數重載和人為地在外推點進行BP。

將改寫的Nesterov衝量演算法同自適應類優化器結合－將m_k的更新由累積形式替換為移動平均形式，並使用二階moment對學習率進行放縮：

至此已經得到了Adan的演算法的基礎版本。

2) 梯度差分的衝量

#可以發現，m_k的更新將梯度與梯度的差分耦合在一起，但是在實際場景中，往往需要對物理意義不同的兩項進行單獨處理，因此研究人員引入梯度差分的衝量v_k：

這裡對梯度的衝量和其差分的衝量設定不同的衝量/平均係數。梯度差分項可以在相鄰梯度不一致的時候減緩優化器的更新，反之，在梯度方向一致時，加速更新。

3) 解耦的權重衰減

對於L2權重正則的目標函數，目前較流行的AdamW優化器透過對L2正則與訓練loss解耦，在ViT和ConvNext上獲得了較好的性能。但是AdamW所用的解耦方法偏向啟發式，目前並不能得到其收斂的理論保證。

基於對L2正規解耦的思想，也為Adan引入解耦的權重衰減策略。目前Adan的每次迭代可以看成是在最小化優化目標F的某種一階近似：

由於F中的L2權重正規過於簡單且光滑性很好，以至於不需要對其進行一階近似。因此，可以只對訓練loss進行一階近似而忽略L2權重正則，那麼Adan的最後一步迭代將會變成：

##有趣的是，可以發現AdamW的更新準則是Adan更新準則在學習率eta接近0時的一階近似。因此，可從proximal 算子的角度給Adan甚至AdamW給予合理的解釋而不是原來的啟發式改進。

4)   Adan優化器

#將2）和3）兩個改良結合進Adan的基礎版本，可以得到如下的Adan優化器。

Adan結合了自適應優化器、Nesterov衝量以及解耦的權重衰減策略的優點，能承受更大的學習率和batch size，以及可以實現對模型參數的動態L2正則。

5)   收斂性分析

#這裡跳過繁複的數學分析過程，只給出結論：

定理：在給定或未給定Hessian-smooth條件的兩種情況下，Adan優化器的收斂速度在非凸隨機最佳化問題上均能達到已知的理論下界，且該結論在帶有解耦的權重衰減策略時仍然成立。

實驗結果

一、CV場景

1）有監督學習——ViT模型

針對ViT模型，研究人員分別在ViT和Swin結構上，測試了Adan的表現。

可以看到，例如在ViT-small、ViT-base、Swin-tiny以及Swin-base上，Adan只消耗了一半的計算資源就獲得了同SoTA優化器接近的結果，並且在同樣的計算量下，Adan在兩種ViT模型上均展現出較大的優勢。

此外，也在大batch size下測試了Adan的效能：

##可以看到，Adan在各種batch size下都表現得不錯，且相對於專為大batch size設計的優化器（LAMB）也具有一定的優勢。

2）有監督學習－CNN模型

除了較難訓練的ViT模型，研究人員也在尖銳局部極小點相對較少的CNN模型上也測試了Adan的性能－包括經典的ResNet與較先進的ConvNext。結果如下：

可以觀察到，不管是ResNet還是ConvNext，Adan均能在大約2/3訓練epoch以內獲得超越SoTA的性能。

3) 無監督學習

在無監督訓練框架下，研究人員在最新提出的MAE上測試了Adan的表現。結果如下：

同有監督學習的結論一致，Adan僅消耗了一半的計算量就追平甚至超過了原來的SoTA優化器，當訓練epoch越小，Adan的優勢就越明顯。

二、NLP場景

#1) 有監督學習

在NLP的有監督學習任務上，分別在經典的LSTM以及先進的Transformer-XL上觀察Adan的表現。

Adan在上述兩個網路上，都表現出一致的優越性。而對於Transformer-XL，Adan在一半的訓練步數內就追平了預設的Adam優化器。

2)   無監督學習

為了測試Adan在NLP場景下無監督任務上的模型訓練情況。研究人員從頭開始訓練BERT：在經過1000k的預訓練迭代後，在GLUE資料集的7個子任務上測試經過Adan訓練的模型效能，結果如下：

Adan在所測試的7個字句分類任務上均展現出較大的優勢。值得一提的是，經過Adan訓練的BERT-base模型，在一些子任務上（例如RTE、CoLA以及SST-2）的結果甚至超過了Adam訓練的BERT-large.

#三、RL場景

研究人員將RL常用的PPO演算法裡的優化器替換為了Adan，並在MuJoCo引擎中的4個遊戲上測試了Adan的性能。在4個遊戲中，以Adan作為網路優化器的PPO演算法，總是能獲得較高的reward。

Adan在RL的網路訓練中，也展現出較大的潛力。

結論與展望

Adan優化器為目前的深度模型最佳化器引入了新的衝量範式。在複雜的訓練範式下以較少的更新次數實現對模型的快速訓練。

實驗顯示，Adan只需1/2-2/3的運算量就能追平現有的SoTA優化器。

#

Adan在多個場景（涉及CV、NLP、RL）、多個訓練方式（有監督與自監督）和多種網路結構（ViT、CNN、LSTM、Transformer等）上，均展現較大的性能優勢。此外，Adan優化器的收斂速度在非凸隨機優化上也達到了理論下界。

以上是訓練ViT和MAE減少一半計算量！ Sea和北大聯合提出高效優化器Adan，深度模型都能用的詳細內容。更多資訊請關注PHP中文網其他相關文章！