爆款論文提出簡單循環單元SRU：像CNN一樣快速訓練RNN

知識 09-12

選自arXiv

機器之心編輯部

近日，一篇題為《Training RNNs as Fast as CNNs》的 arXiv 論文通過有意簡化狀態計算並展現更多的並行性而提出了一個替代性的 RNN 實現，這一循環單元的運算和卷積層一樣快，並且比 cuDNN 優化的 LSTM 快 5-10x。該實現在諸如分類、問題回答、語言建模上證明了其有效性，並已在 PyTorch 和 CNTK1 中開源。

論文地址：https://arxiv.org/pdf/1709.02755.pdf

由於在並行狀態計算上的內在困難，循環神經網路（RNN）的縮放很差。比如，h_t 的前向計算被阻止直到 h_t?1 的整個計算結束，這對並行計算來說是一個主要的瓶頸。在這項工作中，通過有意簡化狀態計算並展現更多的並行性，我們提出了一個替代性的 RNN 實現。這一循環單元的運算和卷積層一樣快，並且比 cuDNN 優化的 LSTM 快 5-10x。我們在大量應用程序上證明了其有效性，包括分類、問題回答、語言建模、翻譯與語音識別，並在 PyTorch 和 CNTK1 中開源了我們的實現。

1. 介紹

最近深度學習中的很多進展來自日益增長的模型能力和相關性算力。這經常涉及到使用大量超參數設置調試的更大、更深的神經網路。然而，不斷增加的模型尺寸和超參數極大地延長了訓練時間。例如，訓練一個當前最優的翻譯或語音識別系統需要花費數天時間完成（Vaswani et al., 2017; Wu et al., 2016b; Sak et al., 2014）。很明顯，計算能力已經成為了深度學習研究的主要瓶頸。

為了抵消顯著增加的計算量，並行化方法比如 GPU 加速訓練已經被廣泛接受以縮放深度學習 (Diamos et al., 2016; Goyal et al., 201)。而諸如卷積和注意力等操作，雖然適用於多線程/GPU 計算，但循環神經網路仍然不太適應並行化。在典型的實現中，輸出狀態 h_t 一直處於閑置直到 h_t-1 的計算完成。這嚴重限制了獨立計算，拖慢了序列處理的速度。圖 1 舉例展示了 cuCNN 優化的 LSTM(Appleyard et al., 2016)，以及使用 conv2d 的單詞級別卷積的處理時間。即使是優化的相當好的 LSTM 實現也慢了 10 倍，這是相當大的差距。

在這次研究中，我們將介紹一種叫簡單循環單元（SRU）的工具，它比起目前出現的循環實現都要快得多。循環單元簡化了狀態計算，從而表現出了類似 CNN、注意力模型和前饋網路的相同並行性。特別是，雖然內態 c_t 仍然利用以前的狀態 c_t-1 更新，但是在循環步驟中，已經不再依賴於 h_t-1 了。結果，循環單元中所有的矩陣乘法運算可以很輕易在任何維度和步驟中並行化。與 cuCNN 和 conv2d 的實現類似，我們對 SRU 使用 CUDA 級別的最優化方法，將所有元素指向的操作編入一個單一的核函數調用中。如圖 1 所示，我們的方法達到了和其 conv2d 對照相同的速度。

圖 1：使用 cuDNN LSTM 的包含 32 個樣本的批量的平均處理時間（以毫秒為單位），單詞級別的卷積 conv2d，以及我們提出的 RNN 實現。l：每個序列的符號數量，d：特徵維度以及 k：特徵寬度。上述數據基於英偉達 GeForce GTX 1070 GPU 和英特爾 Core i7-7700K 處理器而得出。

2. 方法

在這一章節中我們展示了簡單循環單元（Simple Recurrent Unit／SRU）。我們從一個基本的門控循環神經網路實現開始，接著對加速進行必要的更改。更改可被其他門控循環神經網路採用，並不限於這一特定實例。

2.1 SRU 實現

性能最好的循環神經網路如 LSTM（Hochreiter and Schmidhuber, 1997）和門控循環單元（GRU/Cho et al., 2014）利用神經門控來控制信息流並緩解梯度消失與爆炸問題。讓我們看一個典型實現：

其中 f_t 和 i_t 屬於 Sigmoid 門，它們分別被稱為遺忘門（forget gate）和輸入門。

（下文使用 x?t 表示）為在步驟 t 轉換的輸入。我們選擇共軛表達式 i_t=1—f_t 而簡化表達。x?t 的計算對於不同的 RNN 案例來說也不相同。我們使用最簡單的形式對輸入向量 x?t=W*x_t（Lei et al., 2017; Lee et al., 2017）進行線性變換。最後，將內部狀態 c_t 傳遞給激活函數 g(·) 以計算輸出狀態 h_t = g(c_t)。

我們在實現中還是用了兩個附加特徵。首先我們在循環層之間添加了跳過連接，因為它們訓練深度網路十分有效（He et al., 2016; Srivastava et al., 2015; Wu et al., 2016a）。具體來說，我們採用了高速連接（highway connection/Srivastava et al., 2015），此外輸出狀態 h_t'可以通過以下表達式計算：

其中 r_t 為重設門（reset gate）。第二，為了 RNN 正則化我們除了使用標準的 dropout 外，還使用了變分 dropout（variational dropout/Gal and Ghahramani, 2016）。變分 dropout 在不同的時間步驟 t 上共享 dropout mask。在 RNN 每一個矩陣乘法計算中（即 W*drop(x_t)），mask 需要應用到輸入 x_t。標準的 dropout 是在 h_t 上執行的，即沒有饋送到高速連接的輸出狀態。

2.2 加速循環

現有的 RNN 實現再循環計算中使用了前面的輸出狀態 h_t-1。例如，以往向量可以通過 f_t=σ(W_f*x_t+R_f*(h_t-1)+b_f) 計算得出。但是該式中的 R*h_t-1 會破壞獨立性和並行性，因為隱藏狀態每一個維度都依賴於其他狀態，因此 h_t 的計算必須等到整個 h_t-1 都完成計算。

我們提出了完全 drop 連接，即在 h_t-1 和第 t 步神經門之間的連接。SRU 相關聯的等式在下面給出：

給定輸入向量序列，對於不同的時間步 t=1,...,n 是獨立的，因此可以並行計算所有這些向量。我們的方法和最近提出的 Quasi-RNN（Bradbury et al., 2017）十分相似。當我們在上方 3 到 5 表達式中的線性轉換項 drop h_t-1 時，Quasi-RNN 使用 k-gram conv2d 運算來替代線性項。我們設計出的神經網路的瓶頸在於方程式 3 到 5 中間的三個矩陣乘法。在計算 x?t、f_t 和 r_t 後，方程式 6 和 7 能夠非常迅速和簡潔的執行計算，因為它們的運算都是對應元素之間的操作。

事實上，使用 SRU 訓練一個較深的網路十分簡單，因為每一層都只需要較少的計算力，並有較高的處理速度。

2.3 CUDA 優化

在現存的深度學習庫中，一個簡單的 SRU 實現相比於簡單的 LSTM 實現可快 5 倍。優化 SRU 和 cuDNN LSTM (Appleyard et al., 2016) 相似，但要更簡單一些。

3. 實驗

我們在一系列不同的基準上評估 SRU。這些已選擇的基準具有廣泛的應用場景和計算困難。尤其，我們在文本分類、問題回答、語言建模、機器翻譯和語音識別任務上訓練模型。這些基準的訓練時間從若干分鐘（用於分類）到若干天（用於語音）不等。

3.1 分類

表 1：分類基準上的測試精確度。寬 CNNs 是指使用 3, 4, 5-gram 特徵（即濾波器寬度 3, 4, 5）的語句卷積模型（Kim, 2014）。在沒有標準的訓練開發測試拆分時，我們執行 10 倍的交叉驗證。超過 5 個獨立的測試 SST 結果被平均。所有模型使用帶有默認學習率= 0.001 和權重衰減 = 0 的 Adam 優化器進行訓練。

圖 2：在 6 個分類基準上，LSTM、CNN 和 SRU 前 100 個 epoch 的平均有效準確率（y 軸）。X 軸：與第一次迭代關聯的訓練時間（秒）。時間測定是在 Pytorch 和桌面電腦上完成的，配有單個英偉達 GeForce GTX 1070 GPU、英特爾 Core i7-7700k 處理器、CUDA 8 和 cuDNN 6021。

3.2 問答任務

表 2：不同模型在 SQuAD 上的準確匹配率和 F1 得分。我們也報告了每個 epoch 的整體處理時間、RNN 使用的時間。SRU 有更好的結果，運算速度比 cuDNN LSTM 快了 6 倍。時間測定是在桌面電腦上完成的，配備了單個英偉達 GeForce GTX 1070 GPU 和英特爾 Core i7-7700k 處理器。

3.3 語言模型

表 3：在 PTB 語言模型數據集上的困惑度。對比的模型是使用相似的正則化與學習策略進行訓練的：都使用了 cuDNN LSTM；除了（Zaremba et al., 2014), (Press and Wolf, 2017）模型，都是用了變分 dropout；除了 (Zaremba et al., 2014)，其他模型的輸入和輸出都附上了詞嵌入；所有模型都使用了帶有學習率衰減的 SGD。時間測定是在桌面機器上完成的，配有單個英偉達 GeForce GTX 1070 GPU 和英特爾 Core i7-7700k 處理器。

3.4 機器翻譯

表 4：使用 OpenNMT 系統的英-德翻譯結果，我們展示了參數的總數量與排除詞嵌入之後的參數量。我們的設定對 ht_1 feeding 無效（即 -input_feed 0)，極大的減少了訓練時間。在解碼器與編碼器上每增加一個 LSTM 層，在一次訓練 epoch 上就多花費 23 分鐘，而 SRU 只花費 4 分鐘。時間耗費測量是在單個英偉達 Titan X Pascal GPU 上完成的。

3.5 語音識別