Were RNNs All We Needed?

Abstract: The introduction of Transformers in 2017 reshaped the landscape of deep learning. Originally proposed for sequence modelling, Transformers have since achieved widespread success across various domains. However, the scalability limitations of Transformers - particularly with respect to sequence length - have sparked renewed interest in novel recurrent models that are parallelizable during training, offer comparable performance, and scale more effectively. In this work, we revisit sequence modelling from a historical perspective, focusing on Recurrent Neural Networks (RNNs), which dominated the field for two decades before the rise of Transformers. Specifically, we examine LSTMs (1997) and GRUs (2014). We demonstrate that by simplifying these models, we can derive minimal versions (minLSTMs and minGRUs) that (1) use fewer parameters than their traditional counterparts, (2) are fully parallelizable during training, and (3) achieve surprisingly competitive performance on a range of tasks, rivalling recent models including Transformers.

• The paper introduces minLSTM and minGRU, simplified RNNs that remove state dependencies in gates to allow full parallel training.
• The methodology eliminates non-linear range restrictions and leverages parallel prefix scan, achieving up to 1361× faster training on long sequences.
• Empirical results demonstrate competitive performance in language modeling, selective copying, and reinforcement learning despite an 88% increase in memory usage.

The paper simplifies traditional Recurrent Neural Networks (RNNs), specifically Long Short-Term Memory (LSTM) networks and Gated Recurrent Units (GRUs), to create minimal versions named minLSTMs and minGRUs that can be parallelized during training. These simplified models use fewer parameters than their traditional counterparts while achieving competitive performance on various sequence modeling tasks.

The paper begins by revisiting the landscape of sequence modeling, dominated by RNNs for two decades before the advent of Transformers. The inherent sequential nature of RNNs limited parallelization, making them computationally inefficient for long sequences. Transformers, introduced in 2017, enabled parallel training through self-attention but suffer from quadratic computational complexity with respect to sequence length. This limitation has sparked renewed interest in parallelizable recurrent models that scale more efficiently. Recent methods such as state-space models, linearized attention, and linear recurrent neural networks have shown promise in addressing these scalability issues.

The authors focus on LSTMs and GRUs as early examples of input-dependent recurrent models. By removing dependencies on previous states from the gates of these models, they enable parallel training. Further simplification leads to minLSTMs and minGRUs, which use fewer parameters, are fully parallelizable, and achieve competitive performance. Implementations of minGRU and minLSTM are provided in plain PyTorch.

The paper reviews traditional RNNs, highlighting their suitability for tasks involving sequential data but also noting the challenges related to vanishing and exploding gradients. The LSTM is defined by the following equations:

• $\Gamma_o = \sigma(W_{io}x_t + U_{ho}h_{t-1} + b_o)$
  • $\Gamma_o$ is the output gate.
  • $\sigma$ is the sigmoid function.
  • $W_{io}$ is the weight matrix for the input $x_t$.
  • $U_{ho}$ is the weight matrix for the previous hidden state $h_{t-1}$.
  • $b_o$ is the bias vector for the output gate.
• $\Gamma_f = \sigma(W_{if}x_t + U_{hf}h_{t-1} + b_f)$
  • $\Gamma_f$ is the forget gate.
  • $W_{if}$ is the weight matrix for the input $x_t$.
  • $U_{hf}$ is the weight matrix for the previous hidden state $h_{t-1}$.
  • $b_f$ is the bias vector for the forget gate.
• $\Gamma_i = \sigma(W_{ii}x_t + U_{hi}h_{t-1} + b_i)$
  • $\Gamma_i$ is the input gate.
  • $W_{ii}$ is the weight matrix for the input $x_t$.
  • $U_{hi}$ is the weight matrix for the previous hidden state $h_{t-1}$.
  • $b_i$ is the bias vector for the input gate.
• $$\tilde{C}_t = \tanh(W_{ic}x_t + U_{hc}h_{t-1} + b_c)$$
  • $\tilde{C}_t$ is the candidate cell state.
  • $\tanh$ is the hyperbolic tangent function.
  • $W_{ic}$ is the weight matrix for the input $x_t$.
  • $U_{hc}$ is the weight matrix for the previous hidden state $h_{t-1}$.
  • $b_c$ is the bias vector for the candidate cell state.
• $$C_t = \Gamma_f \odot C_{t-1} + \Gamma_i \odot \tilde{C}_t$$
  • $C_t$ is the cell state.
  • $\odot$ denotes element-wise multiplication.
• $h_t = \Gamma_o \odot \tanh(C_t)$
  • $h_t$ is the hidden state.

The GRU, a simplification of the LSTM, uses only two gates and a single state. The GRU equations are:

• $\Gamma_u = \sigma(W_{iu}x_t + U_{hu}h_{t-1} + b_u)$
  • $\Gamma_u$ is the update gate.
  • $W_{iu}$ is the weight matrix for the input $x_t$.
  • $U_{hu}$ is the weight matrix for the previous hidden state $h_{t-1}$.
  • $b_u$ is the bias vector for the update gate.
• $\Gamma_r = \sigma(W_{ir}x_t + U_{hr}h_{t-1} + b_r)$
  • $\Gamma_r$ is the reset gate.
  • $W_{ir}$ is the weight matrix for the input $x_t$.
  • $U_{hr}$ is the weight matrix for the previous hidden state $h_{t-1}$.
  • $b_r$ is the bias vector for the reset gate.
• $$\tilde{h}_t = \tanh(W_{ih}x_t + U_{hh}(\Gamma_r \odot h_{t-1}) + b_h)$$
  • $\tilde{h}_t$ is the candidate hidden state.
  • $W_{ih}$ is the weight matrix for the input $x_t$.
  • $U_{hh}$ is the weight matrix for the previous hidden state $h_{t-1}$.
  • $b_h$ is the bias vector for the candidate hidden state.
• $$h_t = (1 - \Gamma_u) \odot h_{t-1} + \Gamma_u \odot \tilde{h}_t$$

The paper reviews the parallel prefix scan algorithm and its application in computing recurrence relations of the form $v_t = a_t v_{t-1} + b_t$.

The methodology section details the simplification of GRUs and LSTMs to enable training via parallel scan. For the minGRU, the first step involves dropping previous state dependencies from the gates. The GRU's hidden state recurrence is $$h_t = (1 - z_t) \odot h_{t-1} + z_t \odot \tilde{h}_t$$, where $z_t = \sigma(W_{iz}x_t + U_{hz}h_{t-1} + b_z)$ and $$\tilde{h}_t = \tanh(W_{ih}x_t + U_{hh}(r_t \odot h_{t-1}) + b_h)$$. To enable parallel scan, the dependencies on $h_{t-1}$ are removed, simplifying the equations to $z_t = \sigma(W_{iz}x_t + b_z)$ and $\tilde{h}_t = \tanh(W_{ih}x_t + b_h)$.

The second step involves dropping the range restriction of candidate states. The hyperbolic tangent function, which restricts the range of hidden states, is removed, further simplifying the model. The resulting minGRU equations are $$h_t = (1 - z_t) \odot h_{t-1} + z_t \odot \tilde{h}_t$$, $z_t = \sigma(W_{iz}x_t + b_z)$, and $\tilde{h}_t = W_{ih}x_t + b_h$. The minGRU requires $O(2d_h d_x)$ parameters, compared to GRU's $O(3d_h (d_x + d_h))$.

For the minLSTM, the first step involves dropping previous state dependencies from the gates. Similar to the minGRU, the hidden state dependencies are removed from the input, forget, and candidate cell state equations. The second step involves dropping the range restriction of candidate states, removing the hyperbolic tangent function. The third step involves simplifying the scaling of the output by dropping the output gate. The resulting minLSTM equations are $C_t = f_t \odot C_{t-1} + i_t \odot \tilde{C}_t$, $f_t = \sigma(W_{if}x_t + b_f)$, $i_t = \sigma(W_{ii}x_t + b_i)$, and $\tilde{C}_t = W_{ic}x_t + b_c$. The minLSTM requires $O(3d_h d_x)$ parameters compared to LSTM's $O(4d_h (d_x + d_h))$.

The paper then presents empirical results, comparing the minimal versions with their traditional counterparts and modern sequence models. The runtime for sequence lengths of $512$ for minLSTM, minGRU, and Mamba were $2.97$, $2.72$, and $2.71$ milliseconds respectively. For a sequence with length $4096$, the runtime were $3.41$, $3.25$, and $3.15$ respectively. For a sequence length of $512$, minGRUs and minLSTMs were $175 \times$ and $235 \times$ faster per training step than GRUs and LSTMs on a T4 GPU. The improvement is even more significant as sequences grow in length with minGRUs and minLSTMs being $1324 \times$ and $1361 \times$ faster for a sequence length of $4096$. The minimal variants use $\sim 88\%$ more memory compared to their traditional counterparts. minLSTM and minGRU are able to solve the Selective Copying task, achieving performance comparable to S6 and surpassing other modern baselines. In reinforcement learning tasks, minLSTM and minGRU outperform Decision S4 and achieve performance competitive with Decision Transformer, Aaren, and Mamba. In language modeling tasks, minGRU, minLSTM, Mamba, and Transformers achieved comparable test losses of $1.548$, $1.555$, $1.575$, and $1.547$ respectively.

The related work section provides an overview of recent efficient recurrent sequence models, categorized into deep state-space models, recurrent versions of attention, and parallelizable RNNs.

The paper concludes by highlighting the parallel training enabled by removing gate dependencies on previous states. The minimal versions offer fewer parameters, full parallelizability, and competitive performance. The authors suggest a reevaluation of simpler foundational models like LSTM and GRU.

The limitations section acknowledges the hardware constraints that impacted the scale of the experiments, including the use of older GPUs with limited memory. Gradient accumulation was used to accommodate memory limitations, reducing the effective batch size and slowing down training.