如何理解深度學習分布式訓練中的large batch size與learning rate的關系?
本文源自譚旭在知乎問題【如何理解深度學習分布式訓練中的large batch size與learning rate的關系?】下的回答。
問題詳情:
在深度學習進行分布式訓練時,常常采用同步數據并行的方式,也就是采用大的batch size進行訓練,但large batch一般較于小的baseline的batch size性能更差,請問如何理解調試learning rate能使large batch達到small batch同樣的收斂精度和速度?
回答:
最近在進行多GPU分布式訓練時,也遇到了large batch與learning rate的理解調試問題,相比baseline的batch size,多機同步并行(之前有答案是介紹同步并行的通信框架NCCL(譚旭:如何理解Nvidia英偉達的Multi-GPU多卡通信框架NCCL?),有興趣可以查看)等價于增大batch size,如果不進行精細的設計,large batch往往收斂效果會差于baseline的小batch size。因此將自己的理解以及實驗總結如下,主要分為三個方面來介紹:(1)理解SGD、minibatch-SGD和GD,(2)large batch與learning rate的調試關系,(3)我們的實驗。
(1)理解SGD、minibatch-SGD和GD
在機器學習優(yōu)化算法中,GD(gradient descent)是最常用的方法之一,簡單來說就是在整個訓練集中計算當前的梯度,選定一個步長進行更新。GD的優(yōu)點是,基于整個數據集得到的梯度,梯度估計相對較準,更新過程更準確。但也有幾個缺點,一個是當訓練集較大時,GD的梯度計算較為耗時,二是現(xiàn)代深度學習網絡的loss function往往是非凸的,基于凸優(yōu)化理論的優(yōu)化算法只能收斂到local minima,因此使用GD訓練深度神經網絡,最終收斂點很容易落在初始點附近的一個local minima,不太容易達到較好的收斂性能。
另一個極端是SGD(stochastic gradient descent),每次計算梯度只用一個樣本,這樣做的好處是計算快,而且很適合online-learning數據流式到達的場景,但缺點是單個sample產生的梯度估計往往很不準,所以得采用很小的learning rate,而且由于現(xiàn)代的計算框架CPU/GPU的多線程工作,單個sample往往很難占滿CPU/GPU的使用率,導致計算資源浪費。
折中的方案就是mini-batch,一次采用batch size的sample來估計梯度,這樣梯度估計相對于SGD更準,同時batch size能占滿CPU/GPU的計算資源,又不像GD那樣計算整個訓練集。同時也由于mini batch能有適當的梯度噪聲[8],一定程度上緩解GD直接掉進了初始點附近的local minima導致收斂不好的缺點,所以mini-batch的方法也最為常用。
關于增大batch size對于梯度估計準確度的影響,分析如下:
假設batch size為m,對于一個minibatch,loss為:
梯度
整個minibatch的梯度方差為:
由于每個樣本
是隨機從訓練樣本集sample得到的,滿足i.i.d.假設,因此樣本梯度的方差相等,為
等價于SGD的梯度方差,可以看到batch size增大m倍,相當于將梯度的方差減少m倍,因此梯度更加準確。
如果要保持方差和原來SGD一樣,相當于給定了這么大的方差帶寬容量,那么就可以增大lr,充分利用這個方差容量,在上式中添加lr,同時利用方差的變化公式,得到等式
因此可將lr增加sqrt(m)倍,以提高訓練速度,這也是在linear scaling rule之前很多人常用的增大lr的方式[4]。下一小節(jié)將詳細介紹增大lr的問題。
(2)large batch與learning rate
在分布式訓練中,batch size 隨著數據并行的worker增加而增大,假設baseline的batch size為B,learning rate為lr,訓練epoch數為N。如果保持baseline的learning rate,一般不會有較好的收斂速度和精度。原因如下:對于收斂速度,假設k個worker,每次過的sample數量為kB,因此一個epoch下的更新次數為baseline的1/k,而每次更新的lr不變,所以要達到baseline相同的更新次數,則需要增加epoch數量,***需要增加k*N個epoch,因此收斂加速倍數會遠遠低于k。對于收斂精度,由于增大了batch size使梯度估計相較于badeline的梯度更加準確,噪音減少,更容易收斂到附近的local minima,類似于GD的效果。
為了解決這個問題,一個方法就是增大lr,因為batch變大梯度估計更準,理應比baseline的梯度更確信一些,所以增大lr,利用更準確的梯度多走一點,提高收斂速度。同時增大lr,讓每次走的幅度盡量大一些,如果遇到了sharp local minima[8](sharp minima的說法現(xiàn)在還有爭議,暫且引用這個說法),還有可能逃出收斂到更好的地方。
但是lr不能***制的增大,原因分析如下。深度神經網絡的loss surface往往是高維高度非線性的,可以理解為loss surface表面凹凸不平,坑坑洼洼,不像y=x^2曲線這樣光滑,因此基于當前weight計算出來的梯度,往前更新的learing rate很大的時候,沿著loss surface的切線就走了很大一步,有可能大大偏于原有的loss surface,示例如下圖(a)所示,虛線是當前梯度的方向,也就是當前l(fā)oss surface的切線方向,如果learning rate過大,那這一步沿切線方向就走了很大一步,如果一直持續(xù)這樣,那很可能就走向了一個錯誤的loss surface,如圖(b)所示。如果是較小的learning rate,每次只沿切線方向走一小步,雖然有些偏差,依然能大致沿著loss sourface steepest descent曲線向下降,最終收斂到一個不錯的local minima,如圖(c)所示。
同時也可以根據convex convergence theory[2]得到lr的upper bound:lr<1/L,L為loss surface的gradient curve的Lipschitz factor,L可以理解為loss梯度的變化幅度的上界。如果變化幅度越大,L越大,則lr就會越小,如果變化幅度越小,L越小,則lr就可以很大。這和上圖的分析是一致的。
因此,如何確定large batch與learing rate的關系呢?
分別比較baseline和k個worker的large batch的更新公式[7],如下:
這個是baseline(batch size B)和large batch(batch size kB)的更新公式,(4)中l(wèi)arge batch過一步的數據量相當于(3)中baseline k步過的數據量,loss和梯度都按找過的數據量取平均,因此,為了保證相同的數據量利用率,(4)中的learning rate應該為baseline的k倍,也就是learning rate的linear scale rule。
linear scale rule有幾個約束,其中一個約束是關于weight的約束,式(3)中每一步更新基于的weight都是前一步更新過后的weight,因此相當于小碎步的走,每走一部都是基于目前真實的weight計算梯度做更新的,而式(4)的這一大步(相比baseline相當于k步)是基于t時刻的weight來做更新的。如果在這k步之內,W(t+j) ~ W(t)的話,兩者近似沒有太大問題,也就是linear scale rule問題不大,但在weight變化較快的時候,會有問題,尤其是模型在剛開始訓練的時候,loss下特別快,weight變化很快,W(t+j) ~ W(t)就不滿足。因此在初始訓練階段,一般不會直接將lr增大為k倍,而是從baseline的lr慢慢warmup到k倍,讓linear scale rule不至于違背得那么明顯,這也是facebook一小時訓練imagenet的做法[7]。第二個約束是lr不能***的放大,根據上面的分析,lr太大直接沿loss切線跑得太遠,導致收斂出現(xiàn)問題。
同時,有文獻[5]指出,當batchsize變大后,得到好的測試結果所能允許的lr范圍在變小,也就是說,當batchsize很小時,比較容易找打一個合適的lr達到不錯的結果,當batchsize變大后,可能需要精細地找一個合適的lr才能達到較好的結果,這也給實際的large batch分布式訓練帶來了困難。
(3)我們的實驗
最近在考慮分布式訓練NLP相關的深度模型的問題,實驗細節(jié)如下,由于某些工作暫時還不方便透露,只提供較為簡略的實驗細節(jié):
模型baseline參數為batch size 32, lr 0.25,最終的accuracy為BLEU score: 28.35。現(xiàn)在進行分布式擴展到多卡并行。
實驗1:只增加并行worker數(也就相當于增大batch size),lr為baseline的lr0保持不變
可以看到隨著batch的變大, 如果lr不變,模型的精度會逐漸下降,這也和上面的分析相符合。
實驗2:增大batch size,lr相應增大
可以看到通過增加lr到5*lr0(理論上lr應該增加到8倍,但實際效果不好,因此只增加了5倍),并且通過warmup lr,達到和baseline差不多的Bleu效果。最終的收斂速度大約為5倍左右,也就是8卡能達到5倍的收斂加速(不考慮系統(tǒng)通信同步所消耗的時間,也就是不考慮系統(tǒng)加速比的情況下)。
深度學習并行訓練能很好的提升模型訓練速度,但是實際使用的過程中會面臨一系列的問題,包括系統(tǒng)層面的架構設計、算法層面的參數調試等,歡迎有興趣的朋友多多探討。