Training RNNs as Fast as CNNs

SRU: 快速递归单元

最新推荐文章于 2022-02-26 19:45:56 发布

翻译最新推荐文章于 2022-02-26 19:45:56 发布 · 707 阅读

文章标签：

#LSTM #RNN #fast-r-cnn #CNN

自然语言处理专栏收录该内容

9 篇文章

订阅专栏

介绍SRU（Simple Recurrent Unit），这是一种优化过的递归神经网络结构，相较于传统RNN，它大大提升了训练速度，特别是在GPU上的并行计算效率。SRU通过简化状态更新过程，并结合CUDA级别的优化技巧，实现了与卷积神经网络相当的计算性能，在多项任务如文本分类、问答系统等应用中表现出高效能。

摘要

RNN的并行性比较差，主要因为它在计算state的时候不能并行，比如要计算输出h(t),它必须依赖于前一步的输出h(t-1)，这个是并行化的瓶颈。
在这篇论文提出一种可选择的RNN结构，它的递归单元可以和卷积层一样快，是cud优化的LSTM的5-10倍。我们通过一序列的实验包括分类，qa，语言模型，翻译以及语音识别来证明这种卷积单元确实是高效的。
论文源码用 PyTorch 和 CNTK实现过 https://github.com/taolei87/sru

简介

最近深度学习的发展主要归功于模型的容量和计算能力的提升上，模型经常通过更深以及更宽的结构来引入更多的超参来实现。不断增长的模型以及参数个数使得计算量会急剧上升。
比如，需要训练一个较好效果的翻译或者语音识别模型，需要花费几天的时间在训练上面。
很显然，计算性能已经变成现在研究的一个很大的瓶颈。现在卷积网络或者attention模型在利用GPU加速上面并行性做的很好，但是rnn在这方面做的不够。之前也有一些工作是在优化LSTM的计算速度上面，但是
和CNN比还是有10倍的速度落后。

在本文我们介绍一种 Simple Recurrent Unit （SRU），比传统的RNN有显著的速度提升。它是通过简化state的计算来实现- 不是每一步的计算都需要依赖上一步的输出的。
简单的说就是复杂的计算，比如矩阵计算比如 forget gate， input gate等计算都不依赖于上一步的输出h(t)，只依赖于当前step的输出x(t)
状态的更新c(t)的更新必须依赖于上一步的状态c(t-1),但是他的计算都是element-wise，很简单的计算。
和cuDNN LSTM 和conv2d 类似，我们也做了 cuda 维度的优化。

模型

SRU 实现

比较流行的RNN结构比如LSTM 和GRU 都是通过gates的机制来控制信息流，下面开始介绍实现：
首先是状态层更新做了一些简化：

$c_t = f_t \bigodot c_{t-1} + i_t \bigodot x_t^1 \\\$
$= f_t \bigodot c_{t-1} + (1 - f_t) \bigodot x_t^1 \\\$
$f_t, i_t$ 分别是和lstm一致的 forget gate 和 input gate，
是一个sigmoid gate， $x_t^1$ 是对输入的x处理变换，
这里我们只是用一个简单的线性变换， $Wx_t$ ，
$i_t$ 的计算方式直接用 $1-f_t$ 来简单化.

在lstm中 $x_t^1, i_t$ 的计算不仅和 $x_t$ 有关，而且和上一步的输出 $h_{t-1}$ 有关。

然后输出状态 $c_t$ 输入一个激活函数g(.)来生成新的输出 $h_t=g(c_t)$

然后对输出的 $h_t$ 做额外的处理：

第一，我们在引入skip connection layer

$h_t^1 = r_t \bigodot h_t + (1-r_t) \bigodot x_t \\\$
$=r_t \bigodot g(c_t) + (1 -r_t) \bigodot x_t \\\$
$r_t$ 代表 reset gate。

第二，我们实现 variational dropout 机制作为标准dropout正则的补充
普通的 droput 是作用在输出 $h_t$ 之后
它是作用在输入 $x_t$ 上面的

速度优化

普通的RNN的gate，比如： forget gate
$ft=\sigma(W_f x_t + R_f h_{t-1} + b_f) \\\$
必须依赖上一步的输出 $h_{t-1}$ ，所以 $h_{t-1}$ 它破坏了并行性，
我们把这种连接也去除了

我们整体介绍下各个步骤：
$x_t^1 = Wx_t \\\$ ————-(3)
$f_t = \sigma(W_fx_t + b_f) \\\$ ——-(4)
$r_t = \sigma(W_rx_t + b_r) \\\$ ——-(5)
$c_t = f_t \bigodot c_{t-1} + (1- f_t) \bigodot x_t^1 \\\$ ——-(6)
$h_t = r_t \bigodot g(c_t) + (1- r_t) \bigodot x_t \\\$ ——-(7)

下3-5 完全可以并行计算的，6-7虽然不能并行，但是他们的计算是很快的，
因为他们都是些element-wise的计算。

CUDA级别优化

简单的SRU优化大概可以相对LSTM做到5倍的速度提升，
下面我们介绍CUDA级别的优化，
1, 所有step的矩阵相乘都可以一起计算，这样可以提高GPU的使用
2. 所有element-wise计算都是可以融合成一个核函数，如果不这么做，比如 + 或者 sigmoid 运算，是不同的函数调用，这个带来额外的底层核函数调用以及数据加载，整体会有额外的开销。