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

论文地址：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 Schmi