如何理解深度學習分布式訓練中的large batch size與learning rate的關系？

雷鋒網 AI科技評論按，本文源自譚旭在知乎問題【如何理解深度學習分布式訓練中的large batch size與learning rate的關系？】下的回答，雷鋒網 AI科技評論獲其授權轉載。

問題詳情：

在深度學習進行分布式訓練時，常常采用同步數據并行的方式，也就是采用大的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

在機器學習優化算法中，GD（gradient descent）是最常用的方法之一，簡單來說就是在整個訓練集中計算當前的梯度，選定一個步長進行更新。GD的優點是，基于整個數據集得到的梯度，梯度估計相對較準，更新過程更準確。但也有幾個缺點，一個是當訓練集較大時，GD的梯度計算較為耗時，二是現代深度學習網絡的loss function往往是非凸的，基于凸優化理論的優化算法只能收斂到local minima，因此使用GD訓練深度神經網絡，最終收斂點很容易落在初始點附近的一個local minima，不太容易達到較好的收斂性能。

另一個極端是SGD（stochastic gradient descent），每次計算梯度只用一個樣本，這樣做的好處是計算快，而且很適合online-learning數據流式到達的場景，但缺點是單個sample產生的梯度估計往往很不準，所以得采用很小的learning rate，而且由于現代的計算框架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]。下一小節將詳細介紹增大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的說法現在還有爭議，暫且引用這個說法），還有可能逃出收斂到更好的地方。

但是lr不能無限制的增大，原因分析如下。深度神經網絡的loss surface往往是高維高度非線性的，可以理解為loss surface表面凹凸不平，坑坑洼洼，不像y=x^2曲線這樣光滑，因此基于當前weight計算出來的梯度，往前更新的learing rate很大的時候，沿著loss surface的切線就走了很大一步，有可能大大偏于原有的loss surface，示例如下圖（a）所示，虛線是當前梯度的方向，也就是當前loss surface的切線方向，如果learning rate過大，那這一步沿切線方向就走了很大一步，如果一直持續這樣，那很可能就走向了一個錯誤的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）中large 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切線跑得太遠，導致收斂出現問題。

同時，有文獻[5]指出，當batchsize變大后，得到好的測試結果所能允許的lr范圍在變小，也就是說，當batchsize很小時，比較容易找打一個合適的lr達到不錯的結果，當batchsize變大后，可能需要精細地找一個合適的lr才能達到較好的結果，這也給實際的large batch分布式訓練帶來了困難。

（3）我們的實驗

最近在考慮分布式訓練NLP相關的深度模型的問題，實驗細節如下，由于某些工作暫時還不方便透露，只提供較為簡略的實驗細節：

模型baseline參數為batch size 32， lr 0.25，最終的accuracy為BLEU score: 28.35。現在進行分布式擴展到多卡并行。

實驗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倍的收斂加速（不考慮系統通信同步所消耗的時間，也就是不考慮系統加速比的情況下）。

深度學習并行訓練能很好的提升模型訓練速度，但是實際使用的過程中會面臨一系列的問題，包括系統層面的架構設計、算法層面的參數調試等，歡迎有興趣的朋友多多探討。

[1] Li M, Zhang T, Chen Y, et al. Efficient mini-batch training for stochastic optimization[C]// Acm Sigkdd International Conference on Knowledge Discovery & Data Mining. ACM, 2014:661-670.

[2] Bottou L, Curtis F E, Nocedal J. Optimization Methods for Large-Scale Machine Learning[J]. 2016.

[3] Dekel O, Gilad-Bachrach R, Shamir O, et al. Optimal distributed online prediction using mini-batches[J]. Journal of Machine Learning Research, 2012, 13(1):165-202.

[4] Krizhevsky A. One weird trick for parallelizing convolutional neural networks[J]. Eprint Arxiv, 2014.

[5] Breuel T M. The Effects of Hyperparameters on SGD Training of Neural Networks[C]., 2015.

[6] Mishkin D, Sergievskiy N, Matas J. Systematic evaluation of CNN advances on the ImageNet[J]. 2016.

[7] Goyal, Priya, et al. "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour." arXiv preprint arXiv:1706.02677 (2017).

[8] Keskar N S, Mudigere D, Nocedal J, et al. On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima[J]. 2016.

[9]Scaling Distributed Machine Learning with System and Algorithm Co-design - Google Search

雷鋒網

雷峰網版權文章，未經授權禁止轉載。詳情見轉載須知。