爆款論文提出簡單循環單元SRU：像CNN同樣快速訓練RNN（附開源代碼）

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

選自arXiv

機器之心編譯

機器之心編輯部

論文地址：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 相關聯的等式在下面給出：

給定輸入向量序列 {x_1, · · · , x_n}，{x˜t,f_t, r_t} 對於不一樣的時間步 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 語音識別

表 5：不一樣神經模型的詞錯率。注意這裏報告的速度值是基於 SRU 在 CNTK 上的簡單實現。沒有表現出 CUDA 級別的最優化。

4. 結論

該項工做提出了 SRU，這是一個如同 CNN 同樣快的循環模塊，且易於擴展到超過 10 層。咱們在 NLP 與語音識別任務上對其進行了大量測試，驗證了其效率，並在 Github 上開源了咱們實現以助力將來 NLP 與深度學習的研究。

Pytorch 源代碼

在如下內容中，咱們介紹了 SRU 的 Pytorch 源代碼。

項目地址：https://github.com/taolei87/sru

引用

論文：Training RNNs as Fast as CNNs

引用

論文：Training RNNs as Fast as CNNs

@article{lei2017sru,
  title={Training RNNs as Fast as CNNs},
  author={Lei, Tao and Zhang, Yu},
  journal={arXiv preprint arXiv:1709.02755},
  year={2017}
}
複製代碼

要求

• GPU 和 CUDA 8
• PyTorch
• CuPy
• pynvrtc

經過 pip install -r requirements.txt 安裝以上需求。Cupy 和 pynvrtc 須要在運行時把 CUDA 代碼編譯到一個可調用的函數中。

示例

SRU 的使用相似於 nn.LSTM。

import torch
from torch.autograd import Variable
from cuda_functional import SRU, SRUCell

# input has length 20, batch size 32 and dimension 128
x = Variable(torch.FloatTensor(20, 32, 128).cuda())

input_size, hidden_size = 128, 128

rnn = SRU(input_size, hidden_size,    
    num_layers = 2,          # number of stacking RNN layers
    dropout = 0.0,           # dropout applied between RNN layers
    rnn_dropout = 0.0,       # variational dropout applied on linear transformation
    use_tanh = 1,            # use tanh or identity activation
    bidirectional = False    # bidirectional RNN ?
)
rnn.cuda()

output, hidden = rnn(x)      # forward pass

# output is (length, batch size, hidden size * number of directions)
# hidden is (layers, batch size, hidden size * number of directions)
複製代碼

保證 cuda_functional.py 和共享庫 cuda/lib64 能被 system,eg 找到。

export LD_LIBRARY_PATH=/usr/local/cuda/lib64
export PYTHONPATH=path_to_repo/sru
複製代碼

• SRU 實現分類任務的代碼：https://github.com/taolei87/sru/tree/master/classification
• SUR 實現問答任務的代碼：https://github.com/taolei87/sru/tree/master/DrQA
• SRU 語言模型：https://github.com/taolei87/sru/blob/master/language_model/train_lm.py
• 機器翻譯（還未附上）
• SRU 實現語音識別的代碼：github.com/taolei87/sr…