【像训练CNN一样快速训练RNN】全新RNN实现，比优化后的LSTM快10倍

2017 年 9 月 12 日 新智元

1新智元报道

来源：ArXiv

作者：文强

【新智元导读】如何有效训练RNN是一个活跃的研究领域，有很多方法，但还没有哪种表现出了明显的优势，因此也让今天要介绍的这项工作值得注意。来自ASAPP公司和MIT的两位研究人员提出了一种名为“简单循环单元”（Simple Recurrent Unit，SRU）的结构，对现有门控单元做了调整，简化了状态计算的过程，从而展现出了与CNN、注意力和前馈网络相同的并行性。实验结果表明，SRU训练速度与CNN一样，并在图像分类、机器翻译、问答、语音识别等各种不同任务中证明了有效性。

项目已经开源：https://github.com/taolei87/sru

先看论文摘要，有个大概的了解：

标题非常直接，也是很多人都想实现的——《像训练CNN一样快速训练RNN》：

摘要

RNN因其状态计算固有的特性难以并行化因而很难扩展。例如，的前向计算要到的计算完成后才能开始，这构成了并行计算的主要瓶颈。在这项工作中，我们提出了一种RNN的替代实现，简化了状态计算过程，展现更多的并行性。我们所提出的循环单元，运行速度与卷积层一样快，比cuDNN优化的LSTM快5-10倍。我们展示了这种循环单元在广泛应用中的有效性，包括分类、问答、语言建模、翻译和语音识别。我们开源了在PyTorch和CNTK中的实现。

简单循环单元SRU，简化状态计算，速度与CNN一样快

近来深度学习取得的许多进展都源于模型容量的增加和计算力的相应提升。模型容量增大，通常会涉及使用更大、更深的网络，而这些网络又需要复杂的超参数设置和调整。因此，不断增大的模型和超参数数量也大大增加了训练时间。

显然，计算力已经成为深度学习研究的一大主要瓶颈。作为应对，研究人员开始深入挖掘并行计算的潜力，很多人使用GPU加速训练来扩展深度学习。不过，虽然诸如卷积和注意力的运算非常适合于多线程/GPU计算，但是循环神经网络（RNN）仍然不太适合并行化。在典型的RNN实现中，输出状态的计算需要等到计算完成后才能开始。这阻碍了独立计算，并大大减慢了序列处理的速度。

图1展示了cuDNN优化后的LSTM和使用conv2d的字级卷积的处理时间。可以看出，两者区别非常明显，即使是优化后的LSTM，运行速度也可能慢10倍多。

图1：cuDNN优化后的LSTM和使用conv2d的字级卷积的处理时间：即使是优化后的LSTM，运行速度也可能慢10倍多

于是，作者提出了“简单循环单元”（Simple Recurrent Unit，SRU），并在论文中表示其计算速度明显快于现有的循环实现。SRU简化了状态计算的过程，从而展现出了与CNN、注意力和前馈网络相同的并行性。

具体说，虽然SRU的内部状态ct的更新仍然与前一个状态ct-1有关，但是在循环步骤中不再依赖于。因此，SRU中的所有矩阵乘法（即gemm）和元素方面的操作可以在不同的维度和步骤中实现并行化。

SRU实现：增加highway连接和变分dropout

那么，SRU是怎么实现的呢？

作者指出，当前性能最佳的RNN，比如LSTM和GRU，都使用神经门（neural gate）控制信息流，缓解梯度消失（或爆炸）的问题。

因此，他们在此基础上，对门控单元进行了调整。具体说，作者新增加了两个特征：首先，他们在循环层之间增加了highway连接，因为此前的研究已经证明，像highway连接这样的skip connections，在训练深度网络时非常有效；其次，在将RNN正则化时，他们在标准的dropout外，增加了变分dropout，变分dropout在时间步长t与dropout使用相同的mask。

然后，作者将和步骤t的神经门之间的连接全部丢弃，以这种方法来加速循环计算。相比之下，现有的RNN在实现时是要使用先前的输出状态的。

经过这样的改善后，对于现有的深度学习库，SRU已经可以实现超过5倍的加速。接下来，作者还进行了CUDA级的优化，并在一系列不同的基准上进行测试，评估SRU的效果。

实验评估结果：图像分类、语音识别、机器翻译等任务，实现更好性能的同时，训练速度更快

作者还在分类、问答、语言建模、翻译和语音识别等一系列不同任务中评估了SRU。

实验结果证实了SRU的有效性——与这些任务的循环（或卷积）基准模型相比，SRU在实现更好性能的同时，训练速度也更快。