轉(zhuǎn)載自| 新智元
【導(dǎo)讀】換個(gè)優(yōu)化器，計(jì)算量少一半涮瞻。
自Google提出Vision Transformer(ViT)以來鲤拿，ViT漸漸成為許多視覺任務(wù)的默認(rèn)backbone。憑借著ViT結(jié)構(gòu)署咽，許多視覺任務(wù)的SoTA都得到了進(jìn)一步提升近顷，包括圖像分類、分割宁否、檢測(cè)窒升、識(shí)別等。
然而慕匠，訓(xùn)練ViT并非易事异剥。除了需要較復(fù)雜的訓(xùn)練技巧，模型訓(xùn)練的計(jì)算量往往也較之前的CNN大很多絮重。近日冤寿，新加坡Sea AI LAB (SAIL) 和北大ZERO Lab的研究團(tuán)隊(duì)共同提出新的深度模型優(yōu)化器Adan歹苦，該優(yōu)化器可以僅用一半的計(jì)算量就能完成ViT的訓(xùn)練。
論文鏈接：https://arxiv.org/pdf/2208.06677.pdf
代碼鏈接：https://github.com/sail-sg/Adan
此外督怜，在計(jì)算量一樣的情況下殴瘦，Adan在多個(gè)場(chǎng)景（涉及CV、NLP号杠、RL）蚪腋、多種訓(xùn)練方式（有監(jiān)督與自監(jiān)督）和多種網(wǎng)絡(luò)結(jié)構(gòu)/算法（Swin、ViT姨蟋、ResNet屉凯、ConvNext、MAE眼溶、LSTM悠砚、BERT、Transformer-XL堂飞、PPO算法）上灌旧，均獲得了性能提升。
代碼绰筛、配置文件枢泰、訓(xùn)練log均已開源。
1
『深度模型的訓(xùn)練范式與優(yōu)化器』
隨著ViT的提出铝噩，深度模型的訓(xùn)練方式變得越來越復(fù)雜衡蚂。常見的訓(xùn)練技巧包括復(fù)雜的數(shù)據(jù)增強(qiáng)（如MixUp、CutMix骏庸、AutoRand）讳窟、標(biāo)簽的處理（如label smoothing和noise label）、模型參數(shù)的移動(dòng)平均敞恋、隨機(jī)網(wǎng)絡(luò)深度、dropout等谋右。伴隨著這些技巧的混合運(yùn)用硬猫，模型的泛化性與魯棒性均得到了提升，但是隨之而來的便是模型訓(xùn)練的計(jì)算量變得越來越大改执。
在ImageNet 1k上啸蜜，訓(xùn)練epoch數(shù)從ResNet剛提出的90已經(jīng)增長(zhǎng)到了訓(xùn)練ViT常用的300。甚至針對(duì)一些自監(jiān)督學(xué)習(xí)的模型辈挂，例如MAE衬横、ViT，預(yù)訓(xùn)練的epoch數(shù)已經(jīng)達(dá)到了1.6k终蒂。訓(xùn)練epoch增加意味著訓(xùn)練時(shí)間極大的延長(zhǎng)蜂林，急劇增加了學(xué)術(shù)研究或工業(yè)落地的成本遥诉。目前一個(gè)普遍的解決方案是增大訓(xùn)練的batch size并輔助并行訓(xùn)練以減少訓(xùn)練時(shí)間，但是伴隨的問題便是噪叙，大的batch size往往意味著performance的下降矮锈，并且batch size越大，情況越明顯睁蕾。
這主要是因?yàn)槟Ｐ蛥?shù)的更新次數(shù)隨著batch size的增加在急劇減少苞笨。當(dāng)前的優(yōu)化器并不能在復(fù)雜的訓(xùn)練范式下以較少的更新次數(shù)實(shí)現(xiàn)對(duì)模型的快速訓(xùn)練，這進(jìn)一步加劇了模型訓(xùn)練epoch數(shù)的增長(zhǎng)子眶。
因此瀑凝，是否存在一種新的優(yōu)化器能在較少的參數(shù)更新次數(shù)情況下更快更好地訓(xùn)練深度模型？在減少訓(xùn)練epoch數(shù)的同時(shí)臭杰，也能緩解batch size增加帶來的負(fù)面影響粤咪？
2
『被忽略的沖量』
要想加速優(yōu)化器的收斂速度，最直接的方法便是引入沖量硅卢。近年提出的深度模型優(yōu)化器均沿用著Adam中使用的沖量范式——重球法：
其中g(shù)_k是隨機(jī)噪聲射窒，m_k是moment，eta是學(xué)習(xí)率将塑。Adam將m_k的更新由累積形式換成了移動(dòng)平均的形式脉顿，并引入二階moment（n_k）對(duì)學(xué)習(xí)率進(jìn)行放縮，即：
然而隨著Adam訓(xùn)練原始ViT失敗点寥，它的改進(jìn)版本AdamW漸漸地變成了訓(xùn)練ViT甚至ConvNext的首選艾疟。但是AdamW并沒有改變Adam中的沖量范式，因此在當(dāng)batch size超過4,096的時(shí)候敢辩，AdamW訓(xùn)練出的ViT的性能會(huì)急劇下降蔽莱。
在傳統(tǒng)凸優(yōu)化領(lǐng)域，有一個(gè)與重球法齊名的沖量技巧——Nesterov沖量算法：
Nesterov沖量算法在光滑且一般凸的問題上戚长，擁有比重球法更快的理論收斂速度盗冷，并且理論上也能承受更大的batch size。同重球法不同的是同廉，Nesterov算法不在當(dāng)前點(diǎn)計(jì)算梯度仪糖，而是利用沖量找到一個(gè)外推點(diǎn)，在該點(diǎn)算完梯度以后再進(jìn)行沖量累積迫肖。
外推點(diǎn)能幫助Nesterov算法提前感知當(dāng)前點(diǎn)周圍的幾何信息锅劝。這種特性使得Nesterov沖量更加適合復(fù)雜的訓(xùn)練范式和模型結(jié)構(gòu)（如ViT），因?yàn)樗⒉皇菃渭兊匾揽窟^去的沖量去繞開尖銳的局部極小點(diǎn)蟆湖，而是通過提前觀察周圍的梯度故爵，調(diào)整更新的方向。
盡管Nesterov沖量算法擁有一定的優(yōu)勢(shì)隅津，但是在深度優(yōu)化器中诬垂，卻鮮有被應(yīng)用與探索劲室。其中一個(gè)主要的原因就是Nesterov算法需要在外推點(diǎn)計(jì)算梯度，在當(dāng)前點(diǎn)更新剥纷，期間需要多次模型參數(shù)重載以及需要人為地在外推點(diǎn)進(jìn)行back-propagation (BP)痹籍。這些不便利性極大地限制了Nesterov沖量算法在深度模型優(yōu)化器中的應(yīng)用。
3
『Adan優(yōu)化器』
通過結(jié)合改寫的Nesterov沖量與自適應(yīng)優(yōu)化算法晦鞋，并引入解耦的權(quán)重衰減蹲缠，可以得到最終的Adan優(yōu)化器。利用外推點(diǎn)悠垛，Adan可以提前感知周圍的梯度信息线定，從而高效地逃離尖銳的局部極小區(qū)域，以增加模型的泛化性确买。
[if !supportLists]1)?[endif]自適應(yīng)的Nesterov沖量為了解決Nesterov沖量算法中多次模型參數(shù)重載的問題斤讥，研究人員首先對(duì)Nesterov進(jìn)行改寫：
可以證明，改寫的Nesterov沖量算法與原算法等價(jià)湾趾，兩者的迭代點(diǎn)可以相互轉(zhuǎn)化芭商，且最終的收斂點(diǎn)相同〔蟛可以看到铛楣，通過引入梯度的差分項(xiàng)，已經(jīng)可以避免手動(dòng)的參數(shù)重載和人為地在外推點(diǎn)進(jìn)行BP艺普。
將改寫的Nesterov沖量算法同自適應(yīng)類優(yōu)化器相結(jié)合——將m_k的更新由累積形式替換為移動(dòng)平均形式簸州，并使用二階moment對(duì)學(xué)習(xí)率進(jìn)行放縮：
至此已經(jīng)得到了Adan的算法的基礎(chǔ)版本。
[if !supportLists]2)?[endif]梯度差分的沖量
可以發(fā)現(xiàn)歧譬，m_k的更新將梯度與梯度的差分耦合在一起岸浑，但是在實(shí)際場(chǎng)景中，往往需要對(duì)物理意義不同的兩項(xiàng)進(jìn)行單獨(dú)處理瑰步，因此研究人員引入梯度差分的沖量v_k：
這里對(duì)梯度的沖量和其差分的沖量設(shè)置不同的沖量/平均系數(shù)矢洲。梯度差分項(xiàng)可以在相鄰梯度不一致的時(shí)候減緩優(yōu)化器的更新，反之缩焦，在梯度方向一致時(shí)读虏，加速更新。
[if !supportLists]3)?[endif]解耦的權(quán)重衰減
對(duì)于帶L2權(quán)重正則的目標(biāo)函數(shù)舌界，目前較流行的AdamW優(yōu)化器通過對(duì)L2正則與訓(xùn)練loss解耦，在ViT和ConvNext上獲得了較好的性能泰演。但是AdamW所用的解耦方法偏向于啟發(fā)式呻拌，目前并不能得到其收斂的理論保證。
基于對(duì)L2正則解耦的思想睦焕，也給Adan引入解耦的權(quán)重衰減策略藐握。目前Adan的每次迭代可以看成是在最小化優(yōu)化目標(biāo)F的某種一階近似：
由于F中的L2權(quán)重正則過于簡(jiǎn)單且光滑性很好靴拱，以至于不需要對(duì)其進(jìn)行一階近似。因此猾普，可以只對(duì)訓(xùn)練loss進(jìn)行一階近似而忽略L2權(quán)重正則袜炕，那么Adan的最后一步迭代將會(huì)變成：
有趣的是，可以發(fā)現(xiàn)AdamW的更新準(zhǔn)則是Adan更新準(zhǔn)則在學(xué)習(xí)率eta接近0時(shí)的一階近似初家。因此偎窘，可從proximal 算子的角度給Adan甚至AdamW給出合理的解釋而不是原來的啟發(fā)式改進(jìn)。
[if !supportLists]4)?[endif]Adan優(yōu)化器
將2）和3）兩個(gè)改進(jìn)結(jié)合進(jìn)Adan的基礎(chǔ)版本溜在，可以得到如下的Adan優(yōu)化器陌知。
Adan結(jié)合了自適應(yīng)優(yōu)化器、Nesterov沖量以及解耦的權(quán)重衰減策略的優(yōu)點(diǎn)掖肋，能承受更大的學(xué)習(xí)率和batch size仆葡，以及可以實(shí)現(xiàn)對(duì)模型參數(shù)的動(dòng)態(tài)L2正則。5)?收斂性分析
這里跳過繁復(fù)的數(shù)學(xué)分析過程志笼，只給出結(jié)論：定理：在給定或未給定Hessian-smooth條件的兩種情況下沿盅，Adan優(yōu)化器的收斂速度在非凸隨機(jī)優(yōu)化問題上均能達(dá)到已知的理論下界，并且該結(jié)論在帶有解耦的權(quán)重衰減策略時(shí)仍然成立纫溃。
4
『實(shí)驗(yàn)結(jié)果』
[if !supportLists]一腰涧、[endif]CV場(chǎng)景1）有監(jiān)督學(xué)習(xí)——ViT模型
針對(duì)ViT模型，研究人員分別在ViT和Swin結(jié)構(gòu)上皇耗，測(cè)試了Adan的性能南窗。
可以看到，例如在ViT-small郎楼、ViT-base万伤、Swin-tiny以及Swin-base上，Adan僅僅消耗了一半的計(jì)算資源就獲得了同SoTA優(yōu)化器接近的結(jié)果呜袁，并且在同樣的計(jì)算量下敌买，Adan在兩種ViT模型上均展現(xiàn)出較大的優(yōu)勢(shì)。此外阶界，也在大batch size下測(cè)試了Adan的性能：
可以看到虹钮，Adan在各種batch size下都表現(xiàn)得不錯(cuò)，且相對(duì)于專為大batch size設(shè)計(jì)的優(yōu)化器（LAMB）也具有一定的優(yōu)勢(shì)膘融。
[if !supportLists]2）[endif]有監(jiān)督學(xué)習(xí)——CNN模型
除了較難訓(xùn)練的ViT模型芙粱，研究人員也在尖銳局部極小點(diǎn)相對(duì)較少的CNN模型上也測(cè)試了Adan的性能——包括經(jīng)典的ResNet與較先進(jìn)的ConvNext。結(jié)果如下：
可以觀察到氧映，不管是ResNet還是ConvNext春畔，Adan均能在大約2/3訓(xùn)練epoch以內(nèi)獲得超越SoTA的性能。
[if !supportLists]3)?[endif]無監(jiān)督學(xué)習(xí)
在無監(jiān)督訓(xùn)練框架下，研究人員在最新提出的MAE上測(cè)試了Adan的表現(xiàn)律姨。其結(jié)果如下：
同有監(jiān)督學(xué)習(xí)的結(jié)論一致振峻，Adan僅消耗了一半的計(jì)算量就追平甚至超過了原來的SoTA優(yōu)化器，并且當(dāng)訓(xùn)練epoch越小择份，Adan的優(yōu)勢(shì)就越明顯扣孟。
[if !supportLists]二、[endif]NLP場(chǎng)景
[if !supportLists]1)?[endif]有監(jiān)督學(xué)習(xí)
在NLP的有監(jiān)督學(xué)習(xí)任務(wù)上荣赶，分別在經(jīng)典的LSTM以及先進(jìn)的Transformer-XL上觀察Adan的表現(xiàn)凤价。
Adan在上述兩種網(wǎng)絡(luò)上，均表現(xiàn)出一致的優(yōu)越性讯壶。并且對(duì)于Transformer-XL料仗，Adan在一半的訓(xùn)練步數(shù)內(nèi)就追平了默認(rèn)的Adam優(yōu)化器。
[if !supportLists]2)?[endif]無監(jiān)督學(xué)習(xí)
為了測(cè)試Adan在NLP場(chǎng)景下無監(jiān)督任務(wù)上的模型訓(xùn)練情況伏蚊。研究人員從頭開始訓(xùn)練BERT：在經(jīng)過1000k的預(yù)訓(xùn)練迭代后立轧，在GLUE數(shù)據(jù)集的7個(gè)子任務(wù)上測(cè)試經(jīng)過Adan訓(xùn)練的模型性能，結(jié)果如下：
Adan在所測(cè)試的7個(gè)詞句分類任務(wù)上均展現(xiàn)出較大的優(yōu)勢(shì)躏吊。值得一提的是氛改，經(jīng)過Adan訓(xùn)練的BERT-base模型，在一些子任務(wù)上（例如RTE比伏、CoLA以及SST-2）的結(jié)果甚至超過了Adam訓(xùn)練的BERT-large.
[if !supportLists]三胜卤、[endif]RL場(chǎng)景
研究人員將RL常用的PPO算法里的優(yōu)化器替換為了Adan，并在MuJoCo引擎中的4個(gè)游戲上測(cè)試了Adan的性能赁项。在4個(gè)游戲中葛躏，用Adan作為網(wǎng)絡(luò)優(yōu)化器的PPO算法，總能獲得較高的reward悠菜。
Adan在RL的網(wǎng)絡(luò)訓(xùn)練中舰攒，也表現(xiàn)出較大的潛力。
5
『結(jié)論與展望』
Adan優(yōu)化器為目前的深度模型優(yōu)化器引入了新的沖量范式悔醋。在復(fù)雜的訓(xùn)練范式下以較少的更新次數(shù)實(shí)現(xiàn)對(duì)模型的快速訓(xùn)練摩窃。
實(shí)驗(yàn)顯示，Adan僅需1/2-2/3的計(jì)算量就能追平現(xiàn)有的SoTA優(yōu)化器芬骄。
Adan在多個(gè)場(chǎng)景（涉及CV猾愿、NLP、RL）账阻、多個(gè)訓(xùn)練方式（有監(jiān)督與自監(jiān)督）和多種網(wǎng)絡(luò)結(jié)構(gòu)（ViT蒂秘、CNN、LSTM淘太、Transformer等）上姻僧，均展現(xiàn)出較大的性能優(yōu)勢(shì)观挎。此外，Adan優(yōu)化器的收斂速度在非凸隨機(jī)優(yōu)化上也已經(jīng)達(dá)到了理論下界段化。
參考資料：
