如何加速神經語言模型訓練？東北大學小牛翻譯團隊有妙招

雷鋒網AI科技評論按：本文為東北大學自然語言處理實驗室（小牛翻譯團隊）成員肖桐、李垠橋為雷鋒網AI科技評論撰寫的獨家稿件，得到了其指點和審核，雷鋒網在此表示感謝。

團隊簡介：東北大學自然語言處理實驗室（小牛翻譯團隊）由姚天順教授於1980年創立，長期從事機器翻譯及語言分析方面的研究。小牛翻譯團隊所研發的NiuTrans系統開源版和小牛翻譯企業版已經被來自  60 余個國家的2000多個機構使用，並支持44種語言的翻譯。團隊在CWMT、NTCIR等評測中取得多項第一、第二的成績，發表AI、IJCAI、AAAI、ACL等期刊及會議論文40餘篇，並於2016年獲得中文信息處理領域最高獎錢偉長中文信息科學技術獎（一等獎）。

一、神經語言模型簡介

語言作為人類之間進行溝通交流的橋樑，在歷史發展的長河中扮演著極其重要的角色。而隨著人工智慧時代的到來，人們更加希望語言不僅僅能夠在人和人之間進行信息的傳遞，也能讓機器「理解」人所說的話。而在這個過程中十分重要的一個環節就是給人所說的語言建立一個科學有效的模型，具體來說就是讓機器有能力從眾多文字序列中找到合適的字元組合表達語義，與人類進行正常的溝通。

上圖的例子中，我們將語言模型作為一個黑盒，它能夠對三個輸入的「句子」進行一個評判，判斷其語言學上的合理性，這便是語言模型要做的工作。具體來說，它通過鏈式法則對整個文字序列作為句子的可能性進行量化，對每個詞語出現在當前位置（指歷史單詞出現後該詞語出現）的概率求乘積得到整句話的概率。

那要如何知道每個詞出現時的概率呢？在神經語言模型中，其整體結構如上圖所示（圖為循環神經語言模型[1]），輸入層將離散表示的辭彙轉化成連續空間上的詞向量，不同辭彙之間的向量距離能夠反映詞語之間的相似性。隱藏層將輸入層傳遞來的詞向量進行更深層次的表述。輸出層結合隱藏層所傳遞來的信息對可能出現在下一個位置的詞進行預測，獲得辭彙表中每個詞的預測概率，有了它我們就能夠按照鏈式法則對句子的概率進行預測了。

二、更快地訓練神經語言模型

隨著在語音識別、機器翻譯等領域對語言模型性能需求的日益增長，我們迫切地希望能夠進一步改善神經語言模型的性能。一般來說使用更大量的訓練數據來訓練更複雜的模型結構將會獲得更優的模型性能，但與此同時帶來的問題在於其訓練所花費的時間消耗也相應地變長了。面對這種情況，如何有效地減少訓練神經語言模型佔用的時間成為了亟待解決的問題。

那麼如何加快神經語言模型的訓練呢？在處理器的發展過程中，前期主要通過提升硬體主頻來提升設備的運算速度，而後期由於受到物理環境等客觀因素的限制，逐漸將方向轉到多核並行上來提升計算速度。類似於該過程，在神經網路的訓練加速中也主要有兩種方法，一種是通過提高硬體的性能來對訓練進行加速，另一種是通過增加參與運算的設備數量來達到加速的目的。前者就是對目前來說的一種主流思想，這種方法通過逐漸將神經網路中矩陣運算相關的操作由CPU向GPU遷移，達到加速的目的。後者是我們要研究的重點，希望能在更多的設備上對神經語言模型進行訓練。

三、多設備並行的訓練方法及優化

面對著更多的訓練數據，我們要如何在多台設備上進行網路的訓練呢？一種常見的方法就是數據並行[2]，該方法由Jeffrey Dean等人提出，其核心思想就是將訓練數據分散到多台設備上同時進行訓練，並維護同一套網路參數。具體的過程為，在網路中設置一個參數伺服器，用來保存全局的網路參數，其餘每台設備在計算好自己的權重梯度後將其發送給參數伺服器，由參數伺服器將梯度更新到全局最新的權重中來，最後將最新的參數返回給設備伺服器完成一個minibatch的更新。

上述方法中存在兩種參數同步的策略，一種是同步地進行參數更新，另一種為非同步。二者的區別在於參數伺服器向每個設備返回更新後的參數的時機不同。前者需要等待收集到全部設備上發送來的梯度後再進行權重的更新，而後者每當參數伺服器收到一個梯度後便將其更新到自身權重中，並把當前權重值返回給設備，過程如上圖所示。由於非同步的方式加速效果有限，同時會對模型收斂產生一定程度的負面影響，因此在實際情況中一般較少使用。

理想情況下，由於數據並行將數據分散到了多台設備上進行訓練，因此系統訓練速度應該同設備數量的增長而線性增大，但實際結果並不盡如人意。我們在一台機器上使用多台GPU設備對神經語言模型進行訓練，速度圖像如上圖所示，並沒有達到所期待的幾張卡速度就翻幾倍的現象。造成這個問題的原因是什麼呢？

我們可以看到，在數據並行的過程中，權重矩陣以及梯度需要在設備之間進行頻繁的數據傳輸，不論是機器內部設備之間的傳遞還是機器之間的，這部分的傳輸速度相對設備的計算性能均相去甚遠。也就是說，雖然我們使用了多台設備進行協同訓練，成功地將運算分散到多台設備上並行執行，但是在這個過程中為了維護同一套網路參數，反而在設備間的數據傳輸上投入了不小的精力，導致整體的速度提升受到制約。在之前提到的四卡實驗中，數據傳輸所花費的時間就佔據了整個訓練耗時的50%左右，給系統並行效率造成嚴重的負面影響。針對上述問題，眾多科研人員也紛紛提出各種策略，希望能降低多設備訓練中數據傳輸耗時對並行效果造成的不良影響，這裡主要介紹三種常見的優化方法。

第一種加速策略叫做All-Reduce[3]，它將原本中央化的傳輸結構改為環形，取消了參數伺服器的設置，讓設備和設備之間以一種有向環的結構進行數據的傳輸，使得系統內帶寬的利用率更高，數據傳輸更加高效。

第二種方式叫做double-buffering，通過雙緩衝區協同工作的方法，將神經網路訓練過程中梯度計算和數據傳遞兩部分進行時間上的重疊，使得單位時間能夠執行更多的minibatch。

此外還有一種更加簡單粗暴的方式叫做模型平均[4]，該方法在每一輪的訓練過程中，不進行設備間的參數同步，只有當每一輪訓練結束後才進行參數的平均。這樣做有效降低了數據傳輸的頻次從而達到了加速的效果，對於帶寬較低的物理環境（如設備處於多台機器上，使用乙太網進行連接）常常使用它對系統進行優化。

四、基於採樣的方式進行數據傳輸

面對數據傳輸在並行中的大量耗時，我們希望能進一步減小這部分所花費的時間。我們知道傳輸時間等於所要傳輸的數據量除以系統中的帶寬，因此在不對數據並行模式進行修改的前提下，我們只能通過兩種手段對其進行加速，分別是提升傳輸速度以及減小數據傳輸的數量。All-Reduce的方法就屬於第一類，通過數據傳輸結構的變化提升了傳輸性能。那麼我們換一種思想，能否從第二個角度，也就是減少所要傳輸的數據量來對數據傳輸時間進行優化呢？

我們擁有有限的帶寬，希望能儘可能快地完成權重梯度的傳遞，那麼一種方式就是對所要傳輸的梯度矩陣進行採樣，使得需要被傳輸的數據量更少，這樣花費在這部分的時間也會相應地降低。定下了整體的優化方向之後，接下來的問題就是我們要按照哪種策略來進行採樣呢？

簡單一句話，「集中力量辦大事」。在向參數伺服器發送梯度的時候，從完整的梯度矩陣中抽取出對提高神經語言模型性能更有幫助的部分進行傳輸，減少傳輸的數據量，從而達到降低時間消耗的目的，如上圖所示。

這種方法針對神經語言模型的不同層提出了不同的採樣策略。以輸出層為例，權重矩陣為W，其大小為v*h，v為詞表大小，h為隱藏層節點數目。我們可以看到當網路在前向傳播的過程中，權重的每一行都將與隱藏層的輸出向量進行點乘，得到辭彙表中某一詞的預測概率。也就是說，輸出層權重的每一行都和詞表中對應位次的詞相對應，這就是我們要對該層權重進行採樣的一大重要依據。為了能讓網路更快地收斂，我們讓訓練集中頻繁出現的詞所對應的行向量具有更高的概率被抽取出來，傳遞給參數伺服器進行更新。具體的選取策略為：

其中，Vbase為在當前minibatch中出現的詞，Vα為從辭彙表中選擇頻繁出現的若干詞，Vβ為從辭彙表中隨機抽取的詞以保證系統具有良好的魯棒性，在測試集上更穩定。同樣，在輸入層和隱藏層也有各自不同的採樣策略。

舉個例子，如下圖所示在訓練階段，GPU1根據自身權重矩陣計算好它的權重梯度，然後通過採樣的方法從中抽取第一行和第三行傳遞給參數伺服器，同理GPU2也執行了類似的操作，抽取第二和第三行進行傳輸。參數伺服器GPU3收到兩台設備傳遞過來的權重梯度後對它們按行進行累加，得到最終的梯度矩陣，將其更新到網路參數中。在這個過程中由於兩台設備均未採樣到第四行的梯度向量，因此在更新的時候並未對其進行更新。

前述方法中我們使用採樣的方式對需要傳輸的梯度進行抽取，儘可能頻繁地去更新那些對網路性能幫助較大的部分。但其實既然我們已經決定捨棄掉了這部分梯度，為了進一步獲得速度上的提升，完完全全可以不去計算它們。我們都知道神經網路的訓練需要經過前向傳播和反向傳播兩個步驟以獲得權重的梯度，這裡我們為它們建立兩套不同的採樣方式。

還是以輸出層為例，在前向傳播的過程中，我們依據採樣方式A捨棄掉一部分進行計算，接下來在梯度的反向傳輸過程中按照採樣方式B再次捨棄掉部分梯度的傳遞。方式B採樣的方式同數據傳輸中的採樣策略相同，使用Vall對行向量進行篩選，而方式A中我們從辭彙表中又隨機抽取出若干單詞加入到Vall中，這種方式在反向傳播過程中引入了一些雜訊辭彙，能夠有效提升模型的魯棒性。

針對上述兩種方法我們在PTB、FBIS以及Xinhua的數據集上進行測試，結果如上圖所示，其中baseline為非同步訓練的結果，可以看出加速效果有限。在三個任務集上對網路進行分散式訓練，我們可以看到將採樣的方法應用到梯度計算中的加速效果最佳，在使用4張卡的情況下，相對非同步的方法可提速2.1倍左右，相對單卡可提升4.2倍左右。

在三個數據集上訓練LSTM語言模型[5]得到的性能如上圖所示（指標為困惑度），我們可以看到在使用了採樣的方式後模型性能與基線相差不大。下圖給出了在Xinhua的數據集上，校驗集困惑度隨訓練時間變化的曲線圖，其中在4張卡的環境下使用了採樣的方式後，可以看到其收斂性明顯優於基線，在第五輪訓練中困惑度得到明顯下降，同時在8-10輪左右的時候收斂。

在本章的開始處我們提到，在不對數據並行模式進行修改的前提下數據傳輸的加速方式主要有兩種，提升傳輸效率以及減少傳輸數量，二者分別對應我們提到的All-Reduce和採樣的兩種方法。我們在實驗中對比了兩種方式的加速效果，如下圖所示，All-Reduce的方式相對基線在四張卡上可提速約25%，在基於採樣的方式下提速比為41%左右。由於二者對於加速策略的出發點不同，因此理論上可以同時使用兩種方式對系統進行加速。

通過以上幾個實驗我們可以看到，基於採樣的方法在多設備上訓練神經語言模型可以有效減少在數據傳輸中的時間消耗，為在更多設備上訓練提供了可能。該方法主要包括數據傳輸中的採樣以及梯度計算中採樣，在4張GPU設備上訓練，相對非同步的方式可提速約2.1倍，相對單卡訓練可提速4.2倍左右。該方法已被成功應用到小牛翻譯核心引擎中，對於其所支持的包括漢、英、日、韓等44種語言在內的翻譯系統，多設備系統的並行訓練性能提升明顯。

註：本文主要內容參考IJCAI 2017論文Fast Parallel Training of Neural Language Models，部分實驗來自CWMT 2017論文基於數據並行的神經語言模型多卡訓練分析。

參考文獻

[1]Mikolov T, Karafiát M, Burget L, et al. Recurrent neural network based language model[C]// INTERSPEECH 2010, Conference of the International Speech Communication Association, Makuhari, Chiba, Japan, September. DBLP, 2010:1045-1048.

[2] Dean J, Corrado G S, Monga R, et al. Large scale distributed deep networks[C]. International Conference on Neural Information Processing Systems. Curran Associates Inc. 2012:1223-1231.

[3] Wang L, Wu W, Bosilca G, et al. Efficient Communications in Training Large Scale Neural Networks[J]. 2017.

[4] Hansen L K, Salamon P. Neural Network Ensembles[J]. IEEE Transactions on Pattern Analysis & Machine Intelligence, 1990, 12(10):993-1001.

[5] Surhone L M, Tennoe M T, Henssonow S F. Long Short Term Memory[J]. Betascript Publishing, 2010.

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！