計算資源不夠?想快速看到結果?顏水成&林宙辰團隊聯合打造新型高效優化器Adan

來自機器之心

新優化器 Adan 隻需常用優化器的 1/2 到 2/3 的訓練成本 (epochs) 達到模型的 SoTA 性能。

最近,新加坡 Sea AI Lab 聯合北大 Zero LAB 提出新的深度模型優化器 Adan。該優化器可以僅用一半的計算量就能追平甚至超越現有的 SoTA 優化器。在計算量一樣的情況下, Adan 的結果在多個場景(涉及 CV、NLP、RL),多種訓練方式(有監督與自監督),多種網絡結構(ViT、CNN、LSTM、Transformer 等)上,均已超過其默認以及對應的 SoTA 優化器。喜歡本文,點贊支持,歡迎收藏。

更多技術交流,文末加入我們

具體來說,在不改變結構和訓練方式的情況下,Adan 在模型 Swin、ViT、ResNet、ConvNext、MAE、LSTM、BERT、Transformer-XL 以及 RL 中常用的 PPO 算法上,均獲得瞭非平凡的性能提升。

  • 代碼、配置文件和訓練 log 均已開源:http://github.com/sail-sg/Adan
  • 原文鏈接:http://arxiv.org/pdf/2208.06677.pdf

1. 為什麼需要重新考慮深度模型優化器?

自從 Vision Transformer (ViT) 提出,視覺模型的訓練變得越來越復雜,訓練時間也變的越來越長。除瞭湧現出一大批的數據增強技巧,常用的如 MixUp、CutMix、AutoRand,模型訓練的 epoch 也越來越大。在 Imagenet 1k 上,訓練 epoch 從 ResNet 剛提出的 90 epoch 已經增長到瞭常用的 300,甚至針對一些自監督學習的模型,例如 MAE,預訓練的 epoch 已經達到瞭 1.6k epoch。

訓練 epoch 增加意味著訓練時間極大的延長,這樣極大地增加瞭學術研究或工業落地的成本。目前一個普遍的做法是增大訓練的 batch size 並輔助並行訓練以減少訓練時間。但是隨之而來的一個問題就是,大的 batch size 往往意味著 performance 的下降,並且 batch size 越大,情況越明顯。這主要是因為模型參數更新次數的急劇減少,以及較小的梯度噪聲不能很好的幫助優化器快速逃離鞍點和尖銳的局部極小點。

除此之外,大部分優化器並不能忍受較大的學習率,這極大的限制瞭優化器的收斂速度,從而增加訓練成本。

同時,還有一個困擾研究者的問題:在小數據集小 epoch 上有效的超參 / 模型,往往在大數據集大 epoch 會失效。是否存在一個能直接在大數據上快速驗證超參 / 模型的優化方法?

以上的這些問題,均指向一條解決路徑:能否給出一個新的優化器,能以較少的訓練成本達到模型的SoTA性能,即提升深度模型優化器收斂速度。

2. 優化器加速需要一個新的沖量范式

由於 Adam 的流行,最近幾年提出的深度模型優化器均沿用著 Adam 的沖量范式:重球法,

其中 g_k 是隨機噪聲,m_k 是 moment,eta 是學習率。

Adam 隻是將 m_k 的更新由累積形式換成瞭平均形式,並沿用 RMSProp 提出的學習率放縮技巧。

ced7490afd7be4f5d6ab8d43533d4052

然而隨著 Adam 訓練原始 ViT 失敗,AdamW 漸漸的變成瞭訓練 ViT 甚至 ConvNext 首選。但是 AdamW 並沒有改變 Adam 中的沖量范式,因此在當 batch size 超過 4,096 的時候,AdamW 訓練出的 ViT 的性能會急劇下降。Adam 類型優化器中的沖量范式似乎並不能在大 batch size 和較強的數據增強下高效的逃離鞍點和壞的局部極小點,尤其是對於 ViT 這種尖銳局部極小點較多的模型。

3. Nesterov 沖量算法

基於上述原因,我們應當為深度模型優化器考慮一種新的沖量范式。在傳統凸優化領域,有一個與重球法其名的沖量技巧 ---Nesterov 沖量算法:

Nesterov 沖量算法在強凸問題上,擁有比重球法更快的理論收斂速度,並且在許多實際問題上也展現出較大的速度優勢。同重球法不同的是,Nesterov 算法不在原點計算梯度,而是利用沖量找到一個外推點,在該點算完梯度以後再進行梯度累積。

外推點能幫助 Nesterov 算法提前感知周圍的 landscape。這種特性直覺上似乎更適應 ViT 的訓練,因為他並不是簡單的依靠沖量去走出尖銳局部極小點,而是通過提前觀察周圍的梯度,調整更新的方向。

基於此,我們可以考慮用 Nesterov 沖量算法去替代目前 Adam 類優化器中的沖量范式。但是在這之前需要解決一個問題:Nesterov 算法需要在外推點計算梯度,在原點更新,如何避免多次模型參數重載和有效進行 back-propagation (BP)?

為瞭解決這個問題,我們首先對 Nesterov 進行改寫:

04fea7ab8d7d247c29b6947d26870e3f

可以證明,改寫的 Nesterov 沖量算法與原算法等價,兩者的迭代點可以相互轉化,且最終的收斂點相同。

可以看到,通過引入梯度的差分項,我們已經可以避免手動的參數重載和人為的在外推點進行 BP。接下來,我們就可以利用改寫的 Nesterov 沖量算法去替換原有的 Adam 類算法的沖量范式瞭。

4. Adan 優化器與收斂性分析

同其他的 adaptive 類型的優化器類似,我們也將 m_k 的更新由累積形式換成瞭平均形式,並使用二階項對學習率進行放縮:

至此,我們已經得到瞭 Adan 算法的基礎版本。

1)梯度差分的沖量

可以發現,m_k 的更新將梯度與梯度的差分耦合在一起,但是在實際場景中,我們往往需要對物理意義不同的兩項進行單獨處理,因此因此我們引入梯度差分的沖量 v_k:

我們可以對梯度的沖量和其差分的沖量設置不同的沖量 / 平均系數。

2)改進的權重衰減

對於帶 L2 權重正則的目標函數,目前較流行的 AdamW 優化器通過對 L2 正則解耦,在最近流行的網絡上(例如 ViT 和 ConvNext)獲得瞭較好的性能。但是 AdamW 所用的解耦方法偏向於啟發式,目前並不能得到其收斂的理論保證。

基於對 L2 正則解耦的思想,嚴格地從優化的角度,我們也給 Adan 引入解耦的權重衰減策略。目前 Adan 的每次迭代可以看成是在最小化優化目標 F 的某種一階近似:

其中

是基於 n_k 的距離度量。由於 F 中的 L2 權重正則過於簡單且光滑性很好,以至於我們不需要對其進行一階近似。因此,我們可以隻對訓練 loss 進行一階近似而忽略 L2 權重正則,那麼 Adan 的最後一步迭代將會變成:

有趣的是,我們可以發現 AdamW 的更新準則

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

3)Adan 優化器

將上述兩個改進結合進 Adan 的基礎版本,我們可以得到如下的 Adan 優化器。

這裡我們給出瞭一個重啟優化器的條件,即當某種條件滿足的時候,Adan 會重置沖量。這是一個在隨機優化領域很常見的技巧,Adan 也給出這樣的接口。需要說明的是,為瞭效率,我們並沒有在實驗中啟用重啟條件,盡管重啟優化器能帶來一定的性能提升(見論文,表 12)。

4)收斂性分析

這裡我們跳過繁復的數學表達式,隻給出結論:在給定或未給定 Hessian-smooth 條件的兩種情況下,Adan 優化器的收斂速度在非凸隨機優化問題上均能達到已知的理論下界。

並且該結論在帶有解耦的權重衰減策略時仍然成立。目前為止,帶有解耦的權重衰減策略的優化器的收斂性分析仍然是個未解的問題。Adan 的收斂性分析為該問題的解決提供瞭一個值得嘗試的思路。

5. 實驗結果

CV 場景

1)有監督學習

  • ViT 模型:針對 ViT 模型,我們分別在傳統的 ViT 和傑出的 Swin 結構上,測試瞭 Adan 的性能。

並且在 ViT-small 上在不同的 epoch 下,詳細的測試瞭 Adan 和其他優化器的性能。

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

  • CNN 模型:除瞭較難訓練的 ViT 模型,我們也在尖銳局部極小點相對較少的 CNN 模型上也測試瞭 Adan 的性能 --- 包括經典的 ResNet 與較先進的 ConvNext。結果如下:

ab660109240e0e8a170aa06d85f2b3e6

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

2)無監督學習

在無監督訓練框架下,我們在最新提出的 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 的網絡訓練中,也表現出較大的潛力。

6. 結論與展望

Adan 優化器為目前的深度模型優化器引入瞭新的沖量范式。相較於之前的優化器,Adan 給出的范式更適合目前的強數據增強大 batch size 的場景。實驗顯示,Adan 僅需 1/2-2/3 的計算量就能追平或者超越現有的 SoTA 優化器。

Adan 在多個場景(涉及 CV,NLP,RL),多個訓練方式(有監督與自監督),多種網絡結構(ViT,CNN,LSTM,Transformer 等)上,均展現出非平凡的性能優勢。理論角度,Adan 優化器的收斂速度在非凸隨機優化上也已經達到瞭理論下界。

技術交流群

建瞭機器學習算法技術交流群!想要進交流群、獲取資料的同學,可以直接加微信號:mlc2060。加的時候備註一下:研究方向 +學校/公司+知乎,即可。然後就可以拉你進群瞭。

強烈推薦大傢關註 機器學習社區 知乎賬號和 機器學習社區 微信公眾號,可以快速瞭解到最新優質文章。

算法交流、求職內推、算法競賽、幹貨資訊匯總、與 10000+來自港科大、北大、清華、中科院、CMU、騰訊、百度等名校名企開發者互動交流~

推薦文章

擴散模型背後數學太難瞭,啃不動?谷歌用統一視角講明白瞭

ECCV22 | CMU提出首個快速知識蒸餾的視覺框架:80.1%精度,訓練加速30%

CVPR22 最新132篇論文分方向整理|包含目標檢測、圖像處理、醫學影像等28個方向

李宏毅《機器學習》國語課程(2022)來瞭!

CVPR 2022 最新 65 篇論文分方向整理,方向包含:目標檢測、動作識別、人群計數等方向

CVPR 2022 | CNN自監督預訓練新SOTA:上交、Mila、字節聯合提出具有層級結構的圖像表征自學習新框架

NAM: 一種新的註意力計算方式,無需額外的參數!

關於機器學習模型可解釋性算法的匯總

谷歌提出新模型 FLASH,讓 Transformer 模型的效率倍增!訓練成本暴減!

阿裡、SFU提出通用QuadTree Attention,復雜度變線性,性能還更強!ICLR 2022已接收

學習視覺和語言的多粒度對齊?字節提出新多模態預訓練方法 X-VLM:代碼已開源!

超越ConvNeXt,VAN用普通卷積,登頂Backbone性能巔峰!

北大《深度強化學習中文版》.pdf 開放下載!

吳恩達:告別,大數據

AAAI 2022 | 時間序列相關論文一覽(附原文源碼)

我刪掉瞭Transformer中的這幾層,性能反而變好瞭

深度學習中的 Attention 機制總結與代碼實現(2017-2021年)

一文全覽機器學習建模流程(Python代碼)

吳恩達:28張圖全解深度學習知識

PyTorch優化神經網絡的17種方法

深度梳理:機器學習算法模型自動超參數優化方法匯總

趕快收藏,PyTorch 常用代碼段合集真香

聊聊愷明大神MAE的成功之處

何凱明團隊又出新論文!北大、上交校友教你用ViT做遷移學習

大道至簡,何愷明新論文火瞭:Masked Autoencoders讓計算機視覺通向大模型

有瞭這個機器學習畫圖神器,論文、博客都可以事半功倍瞭!

发表回复

相关推荐

爆肝測評:激發活力還是威脅健康?15款功能飲料測評告訴你!

點擊上方藍色字體,關註我們溫馨提示:本文約3760字,閱讀時間11分鐘關於功能飲料,小匠先為大傢總結本文的主要內容,感興趣...

· 20秒前

天婦羅,太簡單瞭

天婦羅,在日本十分流行,大多數人認為天婦羅就是油炸蝦,不是這樣的,它是油炸食物的統稱.基督教在齋期禁吃肉,葡萄牙人就吃魚代...

· 1分钟前

古代艾條灸之“回旋灸”

艾條灸又稱艾卷灸,艾卷灸法最早見於明代朱權的《壽域神方》卷三,灸陰證:“用紙實卷艾,以 紙隔之點穴,於隔紙上用力實按之,...

· 5分钟前

解剖事件:重庆红衣男孩复活,这世界上真的有死亡复活术吗?

今天和大家分享的这个事件相信大部分地曾经听说过,

· 8分钟前

Excel函數之數學函數

從小學到大學,數學總是伴我們成長,在Excel中當然也不例外,今天咱們就來聊聊Excel中的幾個常用的數學函數。1、ROUND函數ROU...

· 14分钟前