阅读论文 《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
机器之星的翻译