Were RNNs All We Needed? | 起居室老虎

原文：arXiv:2410.01201 · Feng et al., 2024

摘要： Transformer 在序列长度方面的可扩展性限制重新引发了对在训练时可并行化的递归序列模型的兴趣。因此，许多新颖的递归架构被提出，例如 S4、Mamba 和 Aaren，这些模型在性能上达到了相当的水平。在本研究中，我们重新探讨了十多年前提出的传统递归神经网络（RNNs）：LSTM（1997 年）和 GRU（2014 年）。虽然这些模型因需要通过时间反向传播（BPTT）而训练速度较慢，但我们发现，通过移除其隐藏状态在输入、遗忘和更新门中的依赖，LSTM 和 GRU 不再需要进行 BPTT，从而可以高效地并行训练。基于此，我们提出了精简版本（minLSTMs 和 minGRUs），它们具有以下特点：1. 显著减少参数数量，相比传统版本更加轻量化；2. 在训练时完全可并行化，在长度为 512 的序列上训练速度提高了 175 倍。最后，我们展示了这些简化版的旧 RNN 模型在实验性能上与近期的序列模型表现相当。

1. Intro

基于 Transformer 的模型计算复杂度在 token length 上都是平方增长的，限制了 scale up，所以最近人们开始重新关注具有以下特点的模型：

训练时消耗的内存和 token length 成线性关系
推理时消耗常量内存

最近流行的 Mamba、RWKV 以及 Aaren，都采用了 Parallel Prefix Scan Algorithm。这篇文章重新审视之前的 LSTM 和 GRU，他们被时代抛弃的主要原因就是训练时只能通过 BPTT，速度太慢。

2. Background

2.1. LSTM

\begin{align} &f_t = \sigma(\text{Linear}_{d_h}([x_t, h{t-1}]))\\ &i_t = \sigma(\text{Linear}_{d_h}([x_t, h{t-1}]))\\ &\tilde{c}_t = \tanh(\text{Linear}_{d_h}([x_t, h_{t-1}]))\\ &o_t = \sigma(\text{Linear}_{d_h}([x_t, h{t-1}]))\\ &c_t = f_t \odot c_{t-1} + i_t \odot \tilde{c}_t\\ &h_t = o_t \odot \tanh(c_t)\end{align}

其中， $\odot$ 表示向量的按元素乘法， $t$ 是当前时间步， $h_t$ 是输出的隐藏状态， $[x_t, h_{t-1}]$ 是 $x_t$ 与 $h_{t-1}$ 的拼接， $d_h$ 是隐藏状态的维度， $c_t$ 是在序列中维护信息的细胞状态， $\tilde{c}_t$ 是候选细胞状态。LSTM 的输入门 $i_t$ 决定了从候选细胞状态中添加多少新信息，遗忘门 $f_t$ 确定要丢弃多少现有细胞状态中的信息，而输出门 $o_t$ 控制了从细胞状态输出哪些信息。激活函数 $\sigma$ 和 $\tanh$ 用于缩放，确保输出不会发生爆炸或消失。

总参数量 $O(4d_h(d_x + d_h))$ 。

2.2. GRU

\begin{align} &z_t = \sigma(\text{Linear}_{d}([x_t, h{t-1}])) \\ &r_t = \sigma(\text{Linear}_{d}([x_t, h{t-1}])) \\ &\tilde{h}_t = \tanh(\text{Linear}_{d}([x_t, r_t \odot h_{t-1}])) \\ &h_t = (1 - z_t) \odot h_{t-1} + z_t \odot \tilde{h}_t\end{align}

其中， $\tilde{h}_t$ 是候选隐藏状态，代表隐藏状态的潜在新值。GRU 将 LSTM 的输入门和遗忘门合并为一个更新门 $z_t \in (0, 1)$ ，该门控制了多少过去的信息得以保留（即 $1 - z_t$ ），以及从候选隐藏状态中添加多少新信息（即 $z_t$ ）。此外，GRU 省略了 LSTM 的输出门，并引入了重置门 $r_t$ ，该门决定了计算候选隐藏状态时应使用多少过去的信息。

减少了参数量，为 $O(3d_h(d_x + d_h))$ 。

2.3. Parallel Scan

为了替换 Transformer 提出的大部分 RNN or RNN like 算法都采用了 prefix scan algorithm。

prefix scan algorithm：

并行扫描算法是一种并行计算方法，利用结合性操作符 $\oplus$ （如加法和乘法），从 N 个顺序数据点计算出 N 个前缀计算结果。具体而言，该算法从输入数据 $\{u_k\}{k=1}^{N}$ 高效地计算出 $\left\{\sum{i=1}^{k} u_i\right\}_{k=1}^{N}$ 。

特别地，我们可以使用并行扫描算法高效计算如下形式的函数：

$v_t = a_t v_{t-1} + b_t \quad \text{其中} \quad v_t, a_t, b_t \in \mathbb{R}, \quad v_0 \leftarrow b_0 \quad (\text{Heinsen, 2023})$

该方法的输入为 $(a_1, \dots, a_n)$ 和 $(b_0, b_1, \dots, b_n)$ ，并通过并行扫描计算出 $(v_1, \dots, v_n)$ 。

3. Methodology

上面的公式可以拓展到向量形式：

v_t = a_t\odot v_{t-1} + b_t

可以看到 LSTM 和 GRU 里面有长得比较类似的模式。本文简化了 RNN，移除输出范围的限制，并且确保在时间维度上是无关的，描述了 miniGRU 和 miniLSTM。

3.1. A minimal GRU: miniGRU

3.1.1. Step 1: Drop previous hidden state dependencies from gates

GRU：

h_t = (1 - z_t) \odot h_{t-1} + z_t \odot \tilde{h}_t

依赖于上一时刻，并且无法去除这种依赖，所以修改为：

\begin{align*} z_t = \sigma(\text{Linear}_{d_h}([x_t, h{t-1}])) \quad &\rightarrow \quad z_t = \sigma(\text{Linear}_{d_h}(x_t)) \\ \tilde{h}t = \tanh(\text{Linear}_{d_h}([x_t, r_t \odot h_{t-1}])) \quad &\rightarrow \quad \tilde{h}_t = \tanh(\text{Linear}_{d_h}(x_t))\end{align*}

3.1.2. Step 2: Drop range restriction of candidate states

\tilde{h}_t = \tanh(\text{Linear}_{d_h}(x_t)) \quad \rightarrow \quad \tilde{h}_t = \text{Linear}_{d_h}(x_t)

原来的函数是帮助裁剪隐藏状态，但是现在只和自己相关之后不需要做这样的裁剪。

3.1.3. miniGRU

参数从 $O(3d_h(d_x+d_h))$ 下降到 $O(2d_hd_x)$

3.2. A Minimal LSTM: minLSTM

3.2.1. Step 1: Drop previous hidden state dependencies from gates

3.2.2. Step 2: Drop range restriction of candidate states

3.2.3. Step 3: Ensure output is time-independent in scale

对遗忘 gate 和 input gate 做归一化，并且需要保证它们 time independent。

3.2.4. minLSTM

参数从 $O(4d_h(d_x+d_h))$ 下降到 $O(3d_hd_x)$ ，并且支持并行扫描训练速度很快。

4. Were RNNs All We Needed?

4.1. Minimal LSTMs and GRUs are very efficient

传统的 RNN 用 BPTT 的训练效率相当低。min 系列的新 RNN 可以通过并行扫描方法进行并行计算，效率更高。文章重点对比的目标是 Mamba，batch size=64.

Runtime & Memory：

和 sequence length 基本是线性的了。

Other Tasks：

剩下的没有什么好讲的了，感觉实验太 toy 了，正大光明说就用了一张 P100 一张 T4，感觉有点抽象。而且去掉了 Hidden Dim 之后还管自己叫 RNN 也有点奇怪。