爆款论文提出简单循环单元：像CNN一样快速训练RNN（附开源代码）

选自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 相关联的等式在下面给出：

给定输入向量序列 {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

@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 实现语音识别的代码：https://github.com/taolei87/sru/tree/master/speech