ASAPP与MIT联合发文:我们该如何训练RNN才能使其像CNN那样快?

图：pixabay

原文来源：arxiv

作者：Tao Lei Yu Zhang

「机器人圈」编译：嗯~阿童木呀、多啦A亮

相信大家都知道，由于在并行化其状态计算时所存在的固有困难，循环神经网络的规模状态相对来说比较差。例如，直到ht-1的整个计算完成之后，ht的前向计算才开始执行，这就是并行计算村子的主要瓶颈问题。在这项工作中，我们提出了一种可替代的RNN实现方法，即故意简化状态计算并展现更多的并行性。而我们在此处所提出的循环单元的运行速度与卷积层一样快，甚至要比经cuDNN优化的LSTM快5至10倍。我们展示了该单元在广泛应用中所展现的有效性，包括分类，问答，语言建模，翻译以及语音识别。而我们将开源用PyTorch和CNTK1实现的代码资源。

现如今，在深度学习中所取得的一系列最新进展大都得益于日益增加的模型容量和相关的计算。这通常涉及使用更大更深的网络，而这往往是通过广泛的超参数设置进行调整得到的。但是，不断增长的模型大小和超参数，大大增加了训练时间。例如，训练一个最先进的翻译或语音识别系统往往需要几天的时间才能完成（Vaswani等人于2017提出，Wu等人于2016提出，Sak等人于2014提出）。很显然的一点是，计算已经成为深度学习研究领域的主要瓶颈问题。

为了应对显著增加的计算需求，诸如GPU加速训练这样的并行化已被广泛应用于扩展深度学习领域中（Diamos等人于2016提出; Goyal等人于2017提出）。虽然诸如卷积和注意这样的操作非常适合于多线程/ GPU计算，但是循环神经网络仍然不太适合并行化。在一个典型的实现中，输出状态ht的计算将会保持暂停，直到ht-1的整个计算完成。这个约束阻碍了独立的计算，并且大大地减缓了序列处理的进程。图1展示了经cuDNN优化的LSTM（Appleyard等，2016）和使用conv2d的词级卷积的处理时间。它们之间的区别是非常明显的——即使是经过特别优化的LSTM的实现速度也要慢10倍以上。

图1：使用cuDNN LSTM、单次级别卷积conv2d和建议的RNN实现的一批32个样本的平均处理时间（以毫秒为单位）。l：每个序列的token数，d：特征维度和k：特征宽度。报道的数字是基于具有Nvidia GeForce GTX 1070 GPU和Intel Core i7-7700K处理器的PyTorch。

在这项工作中，我们介绍了简单循环单元（SRU），其运行明显快于现有的循环实现。循环单元简化了状态计算，从而暴露出与CNN、注意力和前馈网络相同的并行性。具体来说，虽然内部状态ct的更新仍然使用先前的状态ct-1，但是在循环步骤中对ht-1的依赖性已经下降。因此，可以在不同的维度和步骤中容易地并行化循环单元中的所有矩阵乘法（即gemm）和元素操作。类似于cuDNN LSTM和conv2d的实现，我们通过将所有元素操作编译成单个内核函数调用来执行对SRU的CUDA级优化。如图1所示，我们的实现达到了与其conv2d配对物相同的速度。

当然，不能提供可比较或更好的准确性的替代实现将不具有适用性。为此，我们评估SRU广泛的应用范围包括分类、问答、语言建模、翻译和语音识别。实验结果证实了SRU的有效性——与这些任务的循环（或卷积）基线模型相比，它实现了更好的性能，同时能够更快地进行训练。

这项工作提供了简单的循环单元（SRU），其运行速度与CNN一样快速的模块，可轻松扩展到10层以上。 我们对NLP和语音识别任务进行了广泛的评估，证明了这个循环单元的有效性。我们开源实现方法，以促进未来NLP和深度学习的研究。

论文全文链接：https://arxiv.org/pdf/1709.02755.pdf

PyTorch代码实现：https://github.com/taolei87/sru