谷歌提出最新引數優化方法Adafactor,已在TensorFlow中開源

AI前線發表於2019-03-04

谷歌提出最新引數優化方法Adafactor,已在TensorFlow中開源

策劃編輯 | Natalie
編譯 | 郝毅
編輯 | Emily
AI 前線導讀:在深度學習訓練過程中,需要對大量引數進行優化,目前主流的優化方法主要有 SGD、Adam、RMSProp 等。儘管這些方法已經被廣泛的使用,但在許多方面仍需改進。近日,谷歌團隊提出一種新的優化方法 Adafactor,針對目前主流優化方法存在的問題進行改進,對比其他優化演算法在機器翻譯任務上取了的目前最好的結果。本文對這一成果進行了解讀,這是 AI 前線的第 29 篇論文導讀。

更多幹貨內容請關注微信公眾號“AI 前線”,(ID:ai-front)
簡介

深度學習的訓練過程需要優化大量的引數,基於梯度的優化方法是目前用於訓練深度神經網路的主要方法。例如最簡單的隨機梯度下降演算法 (SGD),SGD 就是在一個 minibatch 內讓引數沿著損失函式的負梯度方向走一定的步長。在這個基礎之上,許多基於梯度的自適應方法應運而生,例如使用累計平方梯度的 Adagra、使用指數平均的 RMSProp 方法、使用自適應矩估計的 Adam 演算法以及 Adadelta 等。

在使用大量引數訓練神經網路時,假設我們的計算資源是一定量的,如果訓練過程能夠佔用更少的記憶體,則可以使用更加複雜的模型進行訓練,這樣模型就能獲得更好的效果,也就是記憶體的用量直接決定了模型的大小。而上述優化演算法在帶來更好效能的同時,也需要大量的記憶體進行計算,尤其是自適應優化演算法需要更多的記憶體來儲存額外的引數。例如 Adam 演算法對於每一個引數都需要額外兩個指數衰減值,造成了三倍的記憶體消耗。此外,自適應演算法引數更新的時容易出現迭代過程不穩定的情況。

本文從記憶體消耗、引數更新以及迭代步長三個方面對 Adam 演算法進行了改進,並與 Adam 方法在大型機器翻譯任務上進行了比較,Adafactor 在使用更少記憶體的同時獲得了更好的效果。

Adam 演算法的回顧

Adam 這個名字來源於 adaptive moment estimation,自適應矩估計。概率論中矩的含義是:如果一個隨機變數 X 服從某個分佈,X 的一階矩是 E(X),也就是樣本平均值,X 的二階矩就是 E(X^2),也就是樣本平方的平均值。Adam 演算法根據損失函式對每個引數的梯度的一階矩估計和二階矩估計動態調整針對於每個引數的學習速率。

Adam 也是基於梯度下降的方法,但是每次迭代引數的學習步長都有一個確定的範圍,不會因為很大的梯度導致很大的學習步長,引數的值比較穩定。下面是 Adam 演算法的虛擬碼:

谷歌提出最新引數優化方法Adafactor,已在TensorFlow中開源

Adam 演算法的每部迭代過程主要是以下兩步驟:

  1. 從訓練集中隨機抽取一些樣本以及相關輸出

  2. 計算梯度和誤差,更新 mt 和 vt 以及梯度,並計算引數更新。

步長一般設定為常量,但是近期一些工作認為使用線性衰減的方法設定步長可以得到更好的優化效果。本篇論文中,作者使用倒數平方根衰減獲得更穩定的結果與 Adafactor 方法進行了對比。

記憶體消耗
使用低秩分解

在深度學習方面,普遍存在一個“潛規則”,那就是更大的模型往往能帶來更好的效果。例如語言模型和機器翻譯任務,甚至在一些政治模型也有數十億的引數。優化方法使用的額外引數就會帶來嚴重的記憶體佔用問題。

處理一個大型矩陣的時候,我們習慣使用低秩畢竟的方法替換這個大型的矩陣來解決一些運算上的問題。比如最廣為人知的低秩逼近方法——奇異值分解 (SVD)。本文的作者為了尋求更合適的方法,從非負矩陣分解中尋找靈感。除了 F 範數外,另一種比較常用的代價函式是 K-L 散度,也被稱為 I- 散度。對於非負的標量輸入,I- 散度可以表示為下式:

谷歌提出最新引數優化方法Adafactor,已在TensorFlow中開源

使用該代價函式,我們可以最小化所有的元素級散度,讓它們服從部件級非負約束:

谷歌提出最新引數優化方法Adafactor,已在TensorFlow中開源

對於上面這個方程組,我們無法直接獲得 k 階分解,但是在一階分解時,可以求得其解析解。然後利用一階解的性質對模型引數進行優化,演算法過程如下:

谷歌提出最新引數優化方法Adafactor,已在TensorFlow中開源

類似 Adafactor 的方法最早是由 Hinton 團隊在論文 2017 年 ICLR 的論文”Outrageously large neural networks: The sparsely-gated mixture-of-experts layer”提出。在 2014 年,Gupta 團隊使用嵌入向量的方法對 Adagrad 的累積項進行平均,同樣實現了減少記憶體的使用。在效能相同的情況下,Adafactor 可以將記憶體的佔用從 O(nm) 減小到 O(n+m)。

刪除動量項

在 Adam 演算法中,每個引數都需要兩個額外的累計項。上面的演算法對第二個累積項進行了分解,減少了其記憶體佔用。為了移除第一個累積項,作者首先嚐試了直接將 beta1 設定為 0。

通過實驗結果可以看出,在沒有 warmup 的情況下,沒有動量的優化演算法會很不穩定,下圖是實驗結果的對比。BLEU 是機器翻譯領域的一種評價準則,全稱是 Bilingual Evaluation Understudy,簡單的理解就是分數越高表示翻譯結果越接近正確的人工翻譯結果。

谷歌提出最新引數優化方法Adafactor,已在TensorFlow中開源

引數更新
使用整流的方法約束更新

Reddi 團隊的文章“On the convergence of adam and beyond. International Confer- ence on Learning Representations”談到了 Adam 演算法中使用快速的二階矩估計衰減 (即 beta2 設定的較低) 會導致不收斂的問題。本文的作者在實驗中也發現了這個問題,並且較慢的衰減 (即較高的 beta2) 會導致訓練的不穩定,有無預訓練過程對結果影響很大,實驗結果見下表:

谷歌提出最新引數優化方法Adafactor,已在TensorFlow中開源

對於這一現象,作者是這樣解釋的:較慢的衰減速率意味著二階矩估計是基於更早的梯度的原因。如果模型優化的很快,就會導致很高的估計誤差,這就導致了引數更新時候產生的錯誤。本文中,作者使用了均方根值來檢測模型的訓練過程中是否出現了這個問題。

谷歌提出最新引數優化方法Adafactor,已在TensorFlow中開源

作者繪製了 RMS 曲線來說明訓練過程中不同的衰減對訓練效果的影響。從下圖可以看出,beta 設定為 0.999 時,RMS 值時大時小,訓練過程不穩定,這就是由過時的二階矩估計導致的。

谷歌提出最新引數優化方法Adafactor,已在TensorFlow中開源

為了避免過時的二階矩估計帶來的引數更新,作者使用了一種整流的方法對更新值進行修正,對於超過設定閾值的更新,會等比例的進行縮小,函式如下:

谷歌提出最新引數優化方法Adafactor,已在TensorFlow中開源

在以往的研究中,有許多直接修正梯度的方法。但是對於 SGD 方法,函式的下降方向就是負梯度方向,所以直接修正梯度會對訓練過程造成負面影響。同樣,在許多自適應方法中使用這種做法,標準化的更新方向仍然會超過使用者設定的閾值,這是由於每個引數都有額外的縮放因子。在本文中,作者對需要更新的引數本身進行標準化而不是隻使用梯度。

增加衰減引數

本文的作者受“On the convergence of adam and beyond. International Confer- ence on Learning Representations”一文啟發,使用逐漸增加衰減的方法修改 beta1 和 beta2,具體演算法如下:

谷歌提出最新引數優化方法Adafactor,已在TensorFlow中開源

使用該演算法後,beta2 在 t=1 的時刻為 0,然後逐漸趨向於設定的 beta2 的值。除此之外作者還提出了其他的具體演算法實現衰減的逐漸增加。

迭代步長

不同於傳統的固定步長的方法,本文提出了一種相對步長的策略。針對向量形式的權重和矩陣形式的權重分別有不同的優化方法,實驗證明使用相對步長優化演算法可以獲得更好的表現,下面是演算法的實驗結果。

谷歌提出最新引數優化方法Adafactor,已在TensorFlow中開源

Adafactor 演算法

結合上述三條策略,完整的 Adafactor 演算法如下所示,其中,演算法 4 是針對矩陣形式的引數設計的優化演算法,演算法 5 則是為向量形式的引數所設計。值得一提的是 Adafactor 演算法已經在最新的 TensorFlow 中開源,使用者可以通過 TensorFlow 的 optimize 直接呼叫 Adafactor,原始碼地址:

https://github.com/tensorflow/tensor2tensor/blob/80b2f7300ce2021b9ca0a7877b2670bbdea2b2bf/tensor2tensor/utils/optimize.py

谷歌提出最新引數優化方法Adafactor,已在TensorFlow中開源

谷歌提出最新引數優化方法Adafactor,已在TensorFlow中開源

論文原文:https://arxiv.org/pdf/1804.04235.pdf


更多幹貨內容請關注微信公眾號“AI 前線”,(ID:ai-front)


相關文章