阅读论文 《Train RNN as fast as CNN》

Train RNN as fast as CNN

@(NLP)[cuDNN-LSTM, SRU]

ABSTRACT

作者提出了一个简单循环单元的架构（SRU），能够轻易的实现并行化。SRU像CNN一样快，比优化后的LSTM实现快5到10倍。

In SRU, the majority of computation for each step is independent of the recurrence and can be easily parallelized. SRU is as fast as a convolutional layer and 5-10x faster than an optimized LSTM implementation.

Architecture

左边是通用的RNN结构，右边是SRU

• 在传统的结构中，每次输入step xt $x_t$, t=1,…,n $t=1,\dots ,n$ 依赖于前一步。这阻碍了并行化。
• SRU的设计，使得处理输入的过程能与其他输入独立（大灰块），用相对轻量级的计算做递归组合（小灰块）。主要的计算（虚线框内）能被轻易的并行化。

Average processing time

使用cuDNN LSTM和字级别的2D卷积，提出的SRU在32个样本每一批的平均处理时间（毫秒）
- D：特征维数
- K：特征宽度
- l $l$：每一个序列的令牌数

METHOD

Formula

SRU的基础形式包含一个遗忘门，给出一个在t时间给出 xt $x_t$输入，我们可以计算一个装换 x˜t ${\tilde{x}}_t$和一个遗忘门 ft $f_t$

x˜t=Wxt $${\tilde{x}}_t=W{x}_t$$

ft=σ(Wfxt+bf) $$f_t=\sigma (W_f{x}_t+b_f)$$

计算只依赖 xt $x_t$ ，使得能让所有的时间步并行化， 遗忘门能调制用来产生输出的状态的 ht $h_t$ 的内部状态 ct $c_t$ .

ct=ft⊙ct−1+(1−ft⊙x˜t) $$c_t=f_t\odot c_{t-1}+(1-f_t\odot {\tilde{x}}_t)$$

ht=g(ct) $$h_t=g(c_t)$$

其中 g(⋅) $g(\cdot )$ 是一个激活函数，用来产生 ht $h_t$ 状态。使用公路连接（SRistava等人，2015），并添加类似于遗忘门 ft $f_t$ 计算的复位门 rT $r_T$ 。复位门用于计算输出状态 ht $h_t$ 作为内部状态 g(ct) $g(c_t)$ 和输入 xt $x_t$ 的组合。
完整结构如下：

ct=ft⊙ct−1+(1−ft⊙x˜t)(1) $$\begin{array}{}\text{(1)}& c_t=f_t\odot c_{t-1}+(1-f_t\odot {\tilde{x}}_t)\end{array}$$

ht=g(ct)(2) $$\begin{array}{}\text{(2)}& h_t=g(c_t)\end{array}$$

rt=σ(Wrxt+br)(3) $$\begin{array}{}\text{(3)}& r_t=\sigma (W_r{x}_t+b_r)\end{array}$$

ct=ft⊙ct−1+(1−ft)⊙x˜t(4) $$\begin{array}{}\text{(4)}& c_t=f_t\odot c_{t-1}+(1-f_t)\odot {\tilde{x}}_t\end{array}$$

ht=rt⊙g(ct)+(1−rt)⊙xt(5) $$\begin{array}{}\text{(5)}& h_t=r_t\odot g(c_t)+(1-r_t)\odot x_t\end{array}$$

We set g() = tanh for all our experiments, unless specified otherwise

完全抛弃第 t $t$步的门计算和第 t=1…n $t=1\dots n$步之间的联系。给定输入序列 x1,…,xn $x_1,\dots ,x_n$,得到 t=1…n $t=1\dots n$时对应的 x˜t,ft,rt ${\tilde{x}}_t,f_t,r_t$,能进行并行计算

CUDA-LEVEL Optimization

优化SRU和cuDNN LSTM是如何优化的一样。SRU公式允许两个优化：
- 第一个： 把公式 1−3 $1-3$中的在批处理中的矩阵运算优化：

UT=⎛⎝⎜WWfWr⎞⎠⎟[x1,x2,…,xn] $$U^T=\left(\begin{array}{c}W\\ W_f\\ W_r\end{array}\right)\left[\begin{array}{c}{x}_1,x_2,\dots ,x_n\end{array}\right]$$

- 第二个：序列的所有元素操作可被融合成一个核函数，并在隐藏状态D的维度上并行化。

Algorithm I

论文原文：Train RNN as fast as CNN
机器之星的翻译