ICLR 2018 | Deep Gradient Compression: Reducing the Communication Bandwidth for Distributed Training

為了降低大規模分布式訓練時的通信開銷，作者提出了一種名為深度梯度壓縮(Deep Gradient Compression, DGC)的方法。DGC通過稀疏化技術，在每次迭代時只選擇發送一部分比較“重要”的梯度元素，以達到降低整個訓練過程通信量的目的。為了保證使用DGC后模型的精度，作者還使用了幾種花里胡哨的技術，包括動量修正(momentum correction)、本地梯度裁剪(local gradient cliping)、動量因子遮蔽(momentum factor masking)和預訓練(warmup training)。

梯度稀疏化

為了降低訓練過程中的通信量，我們可以讓每個節點在每次迭代中只發送那些“重要”的梯度。因為“不重要”的梯度元素對模型參數更新的貢獻比較小，所以就可以不用發送這些對更新貢獻較小的梯度。那么問題來了，我們如何知道梯度對參數更新的貢獻呢，換句話說，如何評估梯度的重要性？作者在這里使用了一種啟發式方法——以梯度元素的大小是否超過某個閾值來判斷該元素的重要性——當然，這也是目前最常用的方法。為了防止丟失大量的信息，我們可以把每次迭代中沒有超過閾值的小梯度元素存起來，在下次迭代中加回到原始梯度向量中。隨着訓練的進行，這些較小的梯度元素會累加地越來越大，直至在以后的某次迭代中超過閾值，被節點發送出去。

令\(F(w)\)是我們想要優化的損失函數，同步的分布式SGD算法在\(N\)個訓練節點上會進行如下更新：

\[\begin{aligned} F(w) &= \frac{1}{|\mathcal{X}|}\sum_{x\in \mathcal{X}}f(x,w)\\ w_{t+1} &= w_t -\eta\frac{1}{Nb}\sum_{k=1}^N\sum_{x\in\mathcal{B}_{k,t}}\triangledown f(x,w_t) \end{aligned}\tag{1}\label{1} \]

其中\(\mathcal{X}\)是訓練數據集，\(w\)是神經網絡的權值，\(f(x,w_t)\)是由\(x\in \mathcal{X}\)計算的損失值，\(\eta\)是學習率，\(\mathcal{B}_{k,t}\)是第\(t\)次迭代中第\(k\)個節點上讀取的一個batch的數據樣本，每個batch的大小為\(b\)。考慮將\(w\)拉直后在第\(i\)個位置上的權值\(w^{(i)}\)，經過\(T\)輪迭代后，我們有：

\[w_{t+T}^{(i)} = w_{t}^{(i)} - \eta T · \frac{1}{NbT}\sum_{k=1}^1\left( \sum_{\tau=0}^{T-1} \sum_{x\in \mathcal{B}_{k,t+\tau}} \triangledown^{(i)} f(x,w_{t+\tau})\right) \tag{2}\label{2} \]

等式\(\ref{2}\)表明本地梯度累加可以看作將batch size從\(Nb\)增加到\(NbT\)，其中\(T\)是兩次稀疏更新的間隔，即每進行\(T\)次迭代就發送一次\(w^{(i)}\)的梯度。這是作者說等式\(\ref{2}\)滿足學習率縮放規則，學習率\(\eta T\)中的\(T\)和批量大小\(NbT\)中的\(T\)相互抵消了。說實話，這里有點沒太看懂。個人理解的學習率縮放指的是學習率與batch size等比例增加，比如batch size從128變為1024，那么學習率應該變為原來的8倍。搞不清楚這里出現的兩個\(T\)是什么意思。。。

本地梯度累加

動量修正

在\(N\)個節點上使用標准的動量SGD進行分布式訓練的過程如下所述：

\[u_t = mu_{t-1}+\sum_{k=1}^N(\triangledown_{k,t}),\quad w_{t+1}=w_t -\eta u_t \tag{3}\label{3} \]

這里\(m\)是動量項，\(\triangledown_{k,t}\)是梯度\(\frac{1}{Nb}\sum_{x\in \mathcal{B}_{k,t}}\triangledown f(x,w_t)\)的簡寫形式。考慮將權重\(w\)拉直后第\(i\)個位置上的元素\(w^{(i)}\)，經過\(T\)輪迭代后，\(w^{(i)}\)的變化為：

\[w_{t+T}^{(i)} = w_{t}^{(i)}-\eta \left[···+\left(\sum_{\tau=0}^{T-2}m^\tau \right)\triangledown^{(i)}_{k,t+1} +\left(\sum_{\tau=0}^{T-1}m^\tau \right)\triangledown^{(i)}_{k,t} \right] \tag{4}\label{4} \]

如果動量SGD直接使用稀疏梯度進行更新（算法1的第15行），那么整個更新過程就不再等價於等式\(\ref{3}\)，而是變成了：

\[\begin{aligned} v_{k,t} &= v_{k,t-1}+\triangledown_{k,t}\\ u_t &= mu_{t-1}+\sum_{k=1}^{N}sparse(v_{k,t})\\ w_{t+1} &= w_t-\eta u_t \end{aligned}\tag{5}\label{5} \]

這里，\(v_{k,t}\)是節點\(k\)上的梯度累加和，一旦累加結果大於閾值，它將會被編碼處理然后發送出去，以參與\(u_t\)的更新。隨后，節點\(k\)上的累加結果\(v_{k,t}\)在sparse()中通過掩碼被清空。經過\(T\)輪稀疏更新后，權值\(w^{(i)}\)變為：

\[w_{t+T}^{(i)}=w_{t}^{(i)} - \eta \left(···+\triangledown_{k,t+1}^{(i)}+\triangledown_{k,t}^{(i)}\right)\tag{6}\label{6} \]

等式\(\ref{6}\)相比於等式\(\ref{4}\)少了\(\sum_{\tau=0}^{T-1}m^\tau\)這一項，這就導致收斂速率的降低。

在圖中，等式\(\ref{4}\)會從點A優化到點B，但是由於每個節點上的本地梯度會累加，等式\(\ref{4}\)會到達點C。當梯度稀疏性很高時，那些“不重要”梯度的更新間隔\(T\)會顯著增加，這就會導致模型性能的下降。為了避免上述情況，我們需要對等式\(\ref{5}\)進行動量修正，以令其與等式\(\ref{3}\)等價。如果我們等式\(\ref{3}\)中的速度\(u_t\)看成“梯度”，那么等式\(\ref{3}\)的第二項可以看成針對“梯度”\(u_t\)的標准SGD算法。因此，我們可以在局部累積速度\(u_t\)而不是實際梯度\(\triangledown_{k,t}\)從而將等式\(\ref{5}\)變成與等式\(\ref{3}\)相近的形式：

\[\begin{aligned} u_{k,t} &= mu_{k,t-1} +\triangledown_{k,t} \\ v_{k,t} &= v_{k,t-1}+u_{k,t} \\ w_{t+1} &= w_{t}-\eta\sum_{k=1}^Nsparse(v_{k,t}) \end{aligned}\tag{7}\label{7} \]

這里前兩項是修正后的局部梯度累加，累加結果\(v_{k,t}\)用於隨后的稀疏化和通信。通過對局部累加的簡單修改，我們可以由等式\(\ref{7}\)推導出等式\(\ref{4}\)中的累加折扣因子\(\sum_{\tau =0}^{T-1} m^\tau\)，如圖中(b)所示。注意，動量修正只對更新方程進行調整，並不會引入任何超參數。

本地梯度裁剪

梯度裁剪被廣泛地用於防止梯度爆炸。該方法會在梯度的L2范數之和超過某一閾值時對梯度進行重縮放。一般地，從所有節點進行梯度聚合之后執行梯度裁剪。因為我們會在每個訓練節點的每次迭代中獨立地累加梯度，所以我們會在將本次梯度\(G_t\)加到前一次的累加梯度\(G_{t-1}\)之前進行梯度裁剪。如果所有\(N\)個節點具有相同的梯度分布，那么我們將閾值縮放\(N^{-\frac{1}{2}}\)。在實踐中，我們發現局部梯度裁剪與標准梯度裁剪在訓練中的行為非常相似，這表明我們的假設在實際數據中是有效的。正如我們將在第4節中看到的那樣，動量修正和局部梯度裁剪有助於將AN4語料庫中的單詞錯誤率從14.1％降低到12.9％，而訓練曲線更接近帶動量的SGD。

梯度陳舊性問題

因為延遲了小梯度的更新，所以當這些更新確實發生時，它們已經變得陳舊了。在作者的實驗中，當梯度稀疏度為99.9％時，大多數參數每600到1000次才迭代更新一次，這種陳舊性會降低收斂速度和模型性能。為了解決這個問題，作者使用了動量因子遮蔽和預訓練等技術。

動量因子遮蔽

根據文獻[1]中結論，異步SGD會產生一個隱式動量(implicit momentum)，從而導致收斂變慢。本地梯度累加跟異步GD存在相似性：不能及時更新梯度產生staleness。文獻[1]發現負動量(negative momentum)能一定程度抵消隱式動量的效果，提高收斂速度。本文作者采用了一種類似的方法，如果累加梯度\(v_{k,t}\)大於閾值（即本次迭代將會進行數據傳輸和權重更新），那么就將\(v_{k,t}\)和\(u_{k,t}\)中對應元素清零，從而防止陳舊的動量影響模型權重的更新：

\[Mask \leftarrow |v_{k,t}| \gt thr, \quad v_{k,t}\leftarrow v_{k,t} \odot \urcorner Mask, \quad u_{k,t}\leftarrow u_{k,t}\odot \urcorner Mask \]

預訓練

在訓練的早期，網絡權重會迅速地變化，梯度會比較稠密。因此，在訓練剛開始時，我們需要使用一個較小地學習率來減緩神經網絡的權重變化速率以及增加梯度的稀疏性。這里作者說在預訓練階段按照指數速率控制梯度的稀疏性：75%、93.75%、98.4375%、99.6%、99.9%，搞不懂是怎么手動控制梯度稀疏性的。。。

參考文獻
[1] Mitliagkas, Ioannis, et al. "Asynchrony begets momentum, with an application to deep learning." 2016 54th Annual Allerton Conference on Communication, Control, and Computing (Allerton). IEEE, 2016.