深度梯度壓縮:降低分散式訓練的通訊頻寬

機器之心發表於2017-11-19

ICLR 2018 將在當地時間明年 4 月 30 日~5 月 3 日在加拿大溫哥華舉辦。目前正處於評審階段的會議論文也正陸續公開。近日,機器之心發現了這樣一篇有望極大改善分散式訓練效率的研究論文。以下是我們對該論文的摘要介紹。

大規模分散式訓練可以提升越來越深和越來越大的模型的訓練效率(Chilimbi et al., 2014; Xing et al., 2015; Moritz et al., 2015; Zinkevich et al., 2010)。分散式訓練廣泛使用了同步隨機梯度下降(SGD)。通過增加訓練節點的數量和利用資料並行的優勢,可以極大減少在同等規模訓練資料上的前向-反向通過的總計算時間。但是,梯度交換(gradient exchange)的計算成本很高,很佔用節省下來的計算時間(Li et al., 2014; Wen et al., 2017),尤其是對於計算通訊比(computation-to-communication ratio)很低的迴圈神經網路(RNN)而言。因此,網路頻寬就變成了用於擴充套件分散式訓練的顯著瓶頸。當在移動裝置上執行分散式訓練時,比如 federated learning(McMahan et al., 2016; Konecny et al., 2016)),這種頻寬問題甚至還會變得更加嚴重。在移動裝置上完成訓練是一個引人矚目的目標,因為這能提供更好的隱私性和更好的個性化(Google, 2017),但移動裝置也面臨著一些關鍵的問題:網路頻寬甚至還更低、網路連線斷斷續續以及移動資料套餐價格高昂。

深度梯度壓縮(Deep Gradient Compression/DGC)是通過壓縮梯度的方式來解決通訊頻寬問題。為了確保無損於準確度,DGC 在梯度稀疏化(gradient sparsification)之上使用了動量校正(momentum correction)和區域性梯度裁剪(local gradient clipping)來維持模型的表現水平。DGC 還使用了動量因子掩蔽(momentum factor masking)和 warmup training 來克服由通訊減少所導致的過時問題(staleness problem)。

我們在範圍廣泛的任務、模型和資料集上對 DGC 進行了實驗驗證,其中有:用於影像分類的 CNN(使用了 Cifar10 和 ImageNet)、用於語言建模(使用了 Penn Treebank)和語音識別(使用了 Librispeech Corpus)的 RNN。這些實驗表明:可以將梯度壓縮 600 倍而不造成準確度損失,這比之前的研究成果(Aji & Heafield, 2017)高一個數量級。

深度梯度壓縮:降低分散式訓練的通訊頻寬

圖 1:深度梯度壓縮可以減少通訊時間、提高可擴充套件性和加速分散式訓練


我們通過僅傳送重要梯度的方式減少了通訊頻寬(稀疏更新)。我們使用了梯度幅度作為重要性的簡單指示:只有大於一個閾值的梯度才會被傳輸。為了避免丟失資訊,我們在本地累積梯度的其它部分。最後,這些梯度會變得足夠大,從而得到傳輸。因此,我們會立即傳送大梯度,但最終會隨時間傳送完所有的梯度。演算法 1 展示了這種方法。

深度梯度壓縮:降低分散式訓練的通訊頻寬

演算法 1:在節點 k 上的梯度稀疏化


深度梯度壓縮:降低分散式訓練的通訊頻寬

圖 2:帶有動量校正的梯度稀疏化


深度梯度壓縮:降低分散式訓練的通訊頻寬

表 1:深度梯度壓縮中的技術


我們首先在影像分類任務上檢查了深度梯度壓縮。圖 3(a) 和 3(b) 是 ResNet-110 在 Cifar10 上使用 4 個節點得到的 Top-1 準確度和訓練損失。其梯度稀疏度是 99.9%(僅有 0.1% 是非零的)。Gradient Dropping(Aji & Heafield, 2017)(紅色)的學習曲線由於梯度過時問題而比基準更差。使用動量校正(黃色),學習曲線的收斂稍微快了一些,準確度也接近基準得多。使用動量因子掩蔽和 warmup training(藍色),梯度過時問題被消除了,而且學習曲線也緊跟基準。表 2 給出了詳細的準確度。在使用深度梯度壓縮時,ResNet-110 的準確度得到了完全保持。

深度梯度壓縮:降低分散式訓練的通訊頻寬

圖 3:ResNet 在影像分類任務中的學習曲線(梯度稀疏度為 99.9%)

深度梯度壓縮:降低分散式訓練的通訊頻寬

表 2:在 Cifar10 資料集上訓練的 ResNet-110


深度梯度壓縮:降低分散式訓練的通訊頻寬

表 3:在 ImageNet 資料集上的梯度壓縮比的比較


深度梯度壓縮:降低分散式訓練的通訊頻寬

圖 4:在 PTB 資料集上的 LSTM 語言模型的困惑度和訓練損失(梯度稀疏度為 99.9%)


深度梯度壓縮:降低分散式訓練的通訊頻寬

表 4:使用 4 個節點的語言建模和語音識別的訓練結果


對於語言建模,圖 4 給出了在梯度稀疏度為 99.9% 時,使用 4 個節點訓練得到的語言模型的困惑度和訓練損失。使用深度梯度壓縮所得到的訓練損失非常接近基準,驗證困惑度也一樣。從表 4 可以看到,深度梯度壓縮將梯度壓縮了 462 倍,而困惑度僅稍微有點損失。

深度梯度壓縮:降低分散式訓練的通訊頻寬

圖 5:在 AN4 上的 5 層 LSTM 的 WER(詞錯率)和訓練損失(梯度稀疏度為 99.9%)

對於語音識別,圖 5 給出了在梯度稀疏度為 99.9% 時,使用 4 個節點在 AN4 資料集上訓練得到的 5 層 LSTM 的詞錯率和訓練損失。對於深度梯度壓縮中的技術所帶來的提升,這個學習曲線表現出了和在影像網路中類似的結果。表 4 給出了在 LibriSpeech 測試資料集上的詞錯率(WER)結果,其中 test-clean 包含乾淨的語音,而 test-other 是有噪聲的語音。使用深度梯度壓縮訓練的模型在乾淨語音和有噪聲語音上都得到了更好的識別能力,即使梯度的規模被壓縮了 608 倍。

深度梯度壓縮:降低分散式訓練的通訊頻寬

圖 6:深度梯度壓縮可以提升分散式訓練的速度和可擴充套件性。每個訓練節點有 4 塊英偉達 Titan XP GPU 和一個 PCI 開關。


論文:深度梯度壓縮:降低分散式訓練的通訊頻寬(DEEP GRADIENT COMPRESSION: REDUCING THE COMMUNICATION BANDWIDTH FOR DISTRIBUTED TRAINING)

深度梯度壓縮:降低分散式訓練的通訊頻寬

論文連結: https://openreview.net/pdf?id=SkhQHMW0W

摘要:大規模分散式訓練需要顯著的通訊頻寬來進行梯度交換,這會限制多節點訓練的可擴充套件性並且需要昂貴的高頻寬網路基礎設施。當在移動裝置上進行分散式訓練(federated learning)時,情況甚至更加糟糕,會遭遇更高的延遲、更低的吞吐量和斷斷續續的糟糕連線的問題。在這篇論文中,我們發現在分散式 SGD 中的梯度交換的 99.9% 都是冗餘的,並且提出了可以極大降低通訊頻寬的深度梯度壓縮(DGC)。為了在這種壓縮過程中保持準確度,DGC 使用了 4 種方法:動量校正、區域性梯度裁剪、動量因子掩蔽和 warmup training。我們已經將深度梯度壓縮應用到了影像分類、語音識別和語言建模任務上,使用了 Cifar10、ImageNet、Penn Treebank 和 Librispeech Corpus 等多個資料集。在這些場景中,深度梯度壓縮實現了 270 倍到 600 倍的梯度壓縮比,並且沒有給準確度造成損失;將 ResNet-50 的梯度規模從 97MB 減少到了 0.35MB,對於 DeepSpeech 則將梯度規模從 488MB 壓縮到了 0.74MB。深度梯度壓縮能夠在廉價的商品級 1Gbps 乙太網上實現大規模分散式訓練,並能促進在移動裝置上的分散式訓練。

相關文章