Were RNNs All We Needed?

## Abstract

Abstract The scalability limitations of Transformers regarding sequence length have renewed interest in recurrent sequence models that are parallelizable during training. As a result, many novel recurrent architectures, such as S4, Mamba, and Aaren, have been proposed that achieve comparable performance.
In this work, we revisit traditional recurrent neural networks (RNNs) from over a decade ago: LSTMs (1997) and GRUs (2014).
While these models were slow due to requiring to backpropagate through time (BPTT), we show that by removing their hidden state dependencies from their input, forget, and update gates, LSTMs and GRUs no longer need to BPTT and can be efficiently trained in parallel.
Building on this, we introduce minimal versions (minLSTMs and minGRUs) that (1) use significantly fewer parameters than their traditional counterparts and (2) are fully parallelizable during training ( 175 × 175\times faster for a sequence of length 512 512 ). Lastly, we show that these stripped-down versions of decade-old RNNs match the empirical performance of recent sequence models.

## 1 Introduction

Over the past few years, Transformers (Vaswani et al., ) have been the dominant architecture in many areas, leading to advancements in tasks like machine translation (Devlin et al., ), text generation (Brown et al., ), and more.
However, Transformers have a quadratic computational complexity in the sequence length, making them prohibitively expensive for long sequences, especially in low-resource settings.
As such, numerous works have investigated the design of more efficient alternatives that achieve competitive performance with that of Transformers.
Recently, there has been a renewed interest in recurrent sequence models that can be trained efficiently processing their context in parallel.
These models (1) during training require only linear memory in the sequence length and (2) at inference time are rolled out recurrently token-by-token, requiring only constant memory.
As a result, these models can scale to significantly longer sequences than Transformers(^1^11The title of this paper pays tribute to the original Transformers paper, “Attention is All You Need”.).

A family of efficiently trainable recurrent sequence models that has recently gained much traction is that of state-space models, specifically the recently proposed Mamba (Gu & Dao, ). Mamba (S6) is a state-space model that differentiates itself from prior works by leveraging input-dependent transitions.
The recent success of Mamba and the proposals of many new variants of state-space models has led to several survey papers (Wang et al., ; Patro & Agneeswaran, ; Qu et al., ).
Another extensively explored group of methods is those based on attention.
Peng et al. () proposed a linear attention model that can be written recurrently while being trained in parallel.
Feng et al. () showed that softmax attention (and Transformers) can be viewed as a recurrent neural network (RNN).
Building on their RNN formulation of attention, they proposed Aaren, a softmax attention model, that can be computed in parallel for efficient training or unrolled sequentially as an RNN for efficient inference.
Although many recurrent models have been proposed with vastly different architectures, these recent state-of-the-art methods are all efficiently trainable using the same algorithm – the parallel prefix scan algorithm (Blelloch, ).

Inspired by the striking algorithmic similarities between the numerous recently proposed sequence models, we revisit LSTMs (Hochreiter & Schmidhuber, ) and GRUs (Cho et al., ) from a modern lens.
As traditional RNNs from over a decade ago, LSTMs and GRUs are only computable sequentially and require to backpropagate through time (BPTT) during training.
As such, LSTMs and GRUs were far too slow to scale beyond a few hundred tokens, resulting in their deprecation.
Revisiting these models, we show that by removing hidden state dependencies from their input, forget, and update gates, LSTMs and GRUs no longer need to BPTT and can be trained efficiently using the parallel scan algorithm.
Building on this, we simplify LSTMs and GRUs further by removing their constraints on output range, (i.e., their use of $\mathrm{tanh}$) and ensuring their output is time-independent in scale.
These steps result in minimal versions (minLSTMs and minGRUs) that (1) use significantly fewer parameters than their traditional counterpart and (2) are trainable in parallel ($175\times$ faster for a context length of $512$).
Finally, we show that these stripped-down versions of decade-old RNNs match the empirical performance of recent sequence models.

## 2 Background

In this section, we review recurrent neural networks (RNNs).
RNNs are recurrent sequence models that maintain a hidden state across time steps, capturing temporal dependencies. As such, RNNs are particularly suitable for sequence modelling settings such as those involving time series, natural language processing, and other sequential tasks where context from previous steps informs the current prediction. Vanilla RNNs (Elman, ), however, struggle with issues of vanishing and exploding gradients, limiting their ability to learn long-term dependencies.

### 2.1 LSTM

Addressing this limitation, Hochreiter & Schmidhuber () introduced Long Short-Term Memory (LSTM) networks. LSTMs are enhanced RNNs designed to mitigate the vanishing gradient problem, allowing the model to learn long-term dependencies. LSTMs are computed as follows:

$$ $\displaystyle{\bm{f}}_{t}$ $\displaystyle=\sigma({\mathrm{Linear}}_{d_{h}}([{\bm{x}}_{t},{\bm{h}}_{t-1}]))$ $\displaystyle{\bm{i}}_{t}$ $\displaystyle=\sigma({\mathrm{Linear}}_{d_{h}}([{\bm{x}}_{t},{\bm{h}}_{t-1}]))$ $\displaystyle\tilde{{\bm{c}}}_{t}$ $\displaystyle=\mathrm{tanh}({\mathrm{Linear}}_{d_{h}}([{\bm{x}}_{t},{\bm{h}}_{t-1}]))$ $\displaystyle{\bm{o}}_{t}$ $\displaystyle=\sigma({\mathrm{Linear}}_{d_{h}}([{\bm{x}}_{t},{\bm{h}}_{t-1}]))$ $\displaystyle{\bm{c}}_{t}$ $\displaystyle={\bm{f}}_{t}\odot{\bm{c}}_{t-1}+{\bm{i}}_{t}\odot\tilde{{\bm{c}}}_{t}$ $\displaystyle{\bm{h}}_{t}$ $\displaystyle={\bm{o}}_{t}\odot\mathrm{tanh}({\bm{c}}_{t})$ $$

where $\odot$ represents an element-wise multiplication of vectors,
$t$ is the current timestep,
${\bm{h}}_{t}$ is the outputted hidden state,
$[{\bm{x}}_{t},{\bm{h}}_{t-1}]$ represents the concatenation of ${\bm{x}}_{t}$ with ${\bm{h}}_{t-1}$,
$d_{h}$ is the size of the hidden state,
${\bm{c}}_{t}$ is a cell state that maintains information over the sequence, and
$\tilde{{\bm{c}}}_{t}$ is the candidate cell state to be added,
${\bm{i}}_{t}$, ${\bm{f}}_{t}$, and ${\bm{o}}_{t}$ are gating mechanisms.
The input gate ${\bm{i}}_{t}$ controls how much new information from the candidate cell state is added.
The forget gate ${\bm{f}}_{t}$ determines the proportion of information in the cell gate to discard. The output gate ${\bm{o}}_{t}$ decides what information from the cell state should be outputted.
The $\sigma$ and $\tanh$ are used for scaling to ensure that the output does not explode/vanish.
An LSTM module maintains both a cell and a hidden state and, in total, contains $O(4d_{h}(d_{x}+d_{h}))$ parameters.

### 2.2 GRU

Simplifying LSTM, Cho et al. () introduced Gated Recurrent Unit (GRU) which only uses two gates and a single state instead of LSTM’s three gates and two states (hidden and cell state). GRU’s reduced complexity leads to faster training and inference times while achieving competitive performance in many tasks. GRUs are computed as follows:

$$ $\displaystyle{\bm{z}}_{t}$ $\displaystyle=\sigma({\mathrm{Linear}}_{d}([{\bm{x}}_{t},{\bm{h}}_{t-1}]))$ $\displaystyle{\bm{r}}_{t}$ $\displaystyle=\sigma({\mathrm{Linear}}_{d}([{\bm{x}}_{t},{\bm{h}}_{t-1}]))$ $\displaystyle\tilde{{\bm{h}}}_{t}$ $\displaystyle=\mathrm{tanh}({\mathrm{Linear}}_{d}([{\bm{x}}_{t},{\bm{r}}_{t}\odot{\bm{h}}_{t-1}]))$ $\displaystyle{\bm{h}}_{t}$ $\displaystyle=(1-{\bm{z}}_{t})\odot{\bm{h}}_{t-1}+{\bm{z}}_{t}\odot\tilde{{\bm{h}}}_{t}$ $$

where $\tilde{{\bm{h}}}_{t}$ is the candidate hidden state that represents a potential new value for the hidden state.
GRU combines LSTM’s forget and input gates into a single update gate ${\bm{z}}_{t}\in(0,1)$ which decides how much of the past information to carry forward (i.e., $1-{\bm{z}}_{t}$) and how much new information from the candidate hidden state to add (i.e., ${\bm{z}}_{t}$).
Additionally, LSTM’s output gate is removed and instead, a reset gate ${\bm{r}}_{t}$ is added that controls how much past information is used in computing the candidate hidden state.
GRU reduces the total number of parameters and computations, requiring only $O(3d_{h}(d_{x}+d_{h}))$ parameters.
However, GRUs and LSTMs are only computable sequentially.
As a result, during training they require backpropagating their gradients through time (BPTT), requiring linear training time and greatly limiting their ability to scale to long contexts.

### 2.3 Parallel Scan

Due to this limitation, Transformers replaced LSTMs and GRUs as the defacto sequence modelling method for years by leveraging parallelization during training.
However, Transformers have a quadratic complexity in the sequence length, limiting their ability to scale to long contexts.
Recently, a resurgence of many new recurrent models have been proposed as replacements for Transformers that achieve comparable performance and are trainable in parallel, while avoiding the BPTT issue that traditional RNNs (e.g., LSTMs and GRUs) faced.
Although many different architectures have been proposed, many of these models are efficiently trained using the parallel prefix scan algorithm (Blelloch, ).

The parallel scan algorithm is a parallel computation method for computing $N$ prefix computations from $N$ sequential data points via an associative operator $\oplus$ (e.g., $+$ and $\times$).
The algorithm efficiently computes $\{\bigoplus_{i=1}^{k}u_{i}\}_{k=1}^{N}$ from $\{u_{k}\}_{k=1}^{N}$.
In particular, we can apply the parallel scan method for efficiently computing a popular family of functions: $v_{t}=a_{t}v_{t-1}+b_{t}$ where $v_{t},a_{t},b_{t}\in\mathbb{R}$ and $v_{0}\leftarrow b_{0}$ (Heinsen, ). The method takes as input $a_{1},\ldots,a_{n}$ and $b_{0},b_{1},\ldots,b_{n}$ and computes via parallel scans $v_{1},\ldots,v_{n}$.

## 3 Methodology

Naturally, the aforementioned algorithm also extends to vectors: ${\bm{v}}_{t}={\bm{a}}_{t}\odot{\bm{v}}_{t-1}+{\bm{b}}_{t}$ where $\odot$ is the element-wise multiplication.
Interestingly, we can see that the GRU and LSTM state recurrences resemble the vector formulation.
In this section, we show that GRUs and LSTMs are trainable via parallel scan by simplifying and removing several hidden state dependencies from their various gates.
Building on this, we further simplify these RNNs by removing their constraints on output range, (i.e., $\mathrm{tanh}$) and ensuring the outputs are time-independent in scale.
Combining the steps, we describe minimal versions of GRUs and LSTMs (minGRUs and minLSTMs) that are trainable via parallel scan and perform comparably to Transformers and recently proposed sequence methods.

### 3.1 A Minimal GRU: minGRU

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

Revisiting GRU’s hidden state recurrence which works as follows:

$$ ${\bm{h}}_{t}=(\bm{1}-{\bm{z}}_{t})\odot{\bm{h}}_{t-1}+{\bm{z}}_{t}\odot\tilde{{\bm{h}}}_{t}$ $$

We can observe that the recurrence resembles the aforementioned parallel scan’s formulation where ${\bm{a}}_{t}\leftarrow(\bm{1}-{\bm{z}}_{t})$, ${\bm{b}}_{t}\leftarrow{\bm{z}}_{t}\odot\tilde{{\bm{h}}}_{t}$, and ${\bm{v}}_{t}\leftarrow{\bm{h}}_{t}$. However, ${\bm{z}}_{t}$ and $\tilde{{\bm{h}}}_{t}$ are dependent on previous hidden states ${\bm{h}}_{t-1}$, i.e., ${\bm{z}}_{t}=\sigma({\mathrm{Linear}}_{d_{h}}([{\bm{x}}_{t},{\bm{h}}_{t-1}]))$ and $\tilde{{\bm{h}}}_{t}=\mathrm{tanh}({\mathrm{Linear}}_{d_{h}}([{\bm{x}}_{t},r_{t}\odot{\bm{h}}_{t-1}]))$.
As a result, it is not possible to apply the parallel scan as is since the algorithm’s inputs ${\bm{a}}_{1},\ldots,{\bm{a}}_{n}$ and ${\bm{b}}_{1},\ldots,{\bm{b}}_{n}$ are conditional on already knowing its outputs ${\bm{h}}_{1},\ldots,{\bm{h}}_{n-1}$.

We can remedy this by simplifying GRUs, removing their previous hidden state (i.e., ${\bm{h}}_{t-1}$) dependencies. Specifically, the changes are as follows:

$$ $\begin{aligned} {\bm{z}}_{t}&=\sigma({\mathrm{Linear}}_{d_{h}}([{\bm{x}}_{t},{\bm{h}}_{t-1}]))\\ {\bm{r}}_{t}&=\sigma({\mathrm{Linear}}_{d_{h}}([{\bm{x}}_{t},{\bm{h}}_{t-1}]))\\ \tilde{{\bm{h}}}_{t}&=\mathrm{tanh}({\mathrm{Linear}}_{d_{h}}([{\bm{x}}_{t},{\bm{r}}_{t}\odot{\bm{h}}_{t-1}]))\\ \end{aligned}\quad\Rightarrow\quad\begin{aligned} {\bm{z}}_{t}&=\sigma({\mathrm{Linear}}_{{\bm{d}}_{h}}({\bm{x}}_{t}))\\ \tilde{{\bm{h}}}_{t}&=\mathrm{tanh}({\mathrm{Linear}}_{{\bm{d}}_{h}}({\bm{x}}_{t}))\\ \end{aligned}$ $$

By removing the dependence on ${\bm{h}}_{t-1}$ from the candidate hidden state $\tilde{{\bm{h}}}_{t}$, the reset gate ${\bm{r}}_{t}$ that would control ${\bm{h}}_{t-1}$ weight is also no longer needed and is removed.
Without the dependencies on previous hidden states, the inputs to the algorithm ${\bm{a}}_{1},\ldots,{\bm{a}}_{n}$ and ${\bm{b}}_{1},\ldots,{\bm{b}}_{n}$ are all easily computed in parallel and can thus be used to compute ${\bm{h}}_{1},\ldots,{\bm{h}}_{n}$ efficiently via the parallel scan.

#### 3.1.2 Step 2: Drop range restriction of candidate states

In GRU’s hidden state recurrence, the proportion carried over from the previous hidden state ($\mathbf{1}-{\bm{z}}_{t}$) and the amount added for the new candidate hidden state (${\bm{z}}_{t}$) sum to $1$. As a result, the scale of GRU’s hidden state value is time-independent.
Instead, the scale of its hidden state depends on that of its candidate hidden states $\tilde{{\bm{h}}}_{t}$.
The hyperbolic tangent function ($\mathrm{tanh}$) plays a crucial role in LSTMs and GRUs, restricting the range of (candidate) hidden states, i.e., $\tilde{{\bm{h}}}_{t},{\bm{h}}_{t}\in(-1,1)^{d_{h}}$.
The $\mathrm{tanh}$ helps stabilize the training and mitigates vanishing gradients that result from applying sigmoid ($\sigma$) activations to linear transformations of the hidden state (e.g., ${\bm{z}}_{t}=\sigma({\mathrm{Linear}}_{d_{h}}([{\bm{x}}_{t},{\bm{h}}_{t-1}]))$).
In the previous step, these hidden state dependencies were removed.
As such, we can simplify GRU further by removing the range restriction ($\mathrm{tanh}$) on the (candidate) hidden states as follows:

$$ $\begin{aligned} \tilde{{\bm{h}}}_{t}&=\mathrm{tanh}({\mathrm{Linear}}_{d_{h}}({\bm{x}}_{t}))\\ \end{aligned}\quad\Rightarrow\quad\begin{aligned} \tilde{{\bm{h}}}_{t}&={\mathrm{Linear}}_{d_{h}}({\bm{x}}_{t})\\ \end{aligned}$ $$

#### 3.1.3 minGRU

Combining the two simplification steps results in a minimal version of GRU (minGRU):

$$ $\Rightarrow$ $$

The resulting model is significantly more efficient than the original GRU (1) requiring only $O(2d_{h}d_{x})$ parameters instead of GRU’s $O(3d_{h}(d_{x}+d_{h}))$ parameters where $d_{x},d_{h}$ corresponds to the sizes of $x_{t}$ and $h_{t}$ respectively. In terms of training, minGRU (2) can be trained in parallel using the parallel scan algorithm, speeding up training significantly. In Section 4.1, we show that this corresponded to a $175\times$ speedup in training steps for a sequence length of $512$ on a T4 GPU.
The parameter efficiency gains are also significant. Typically, in RNNs, state expansion is performed (i.e., $d_{h}=\alpha d_{x}$ where $\alpha\geq 1$) allowing the models to more readily learn features from their inputs.
minGRU uses approximately $33\%,22\%,17\%,$ or $13\%$ of parameters compared to GRU when $\alpha=1,2,3,$ or $4$ respectively.

### 3.2 A Minimal LSTM: minLSTM

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

Revisiting LSTMs, we focus on their cell state recurrence which works as follows:

$$ ${\bm{c}}_{t}={\bm{f}}_{t}\odot{\bm{c}}_{t-1}+{\bm{i}}_{t}\odot\tilde{{\bm{c}}}_{t}\\$ $$

Similar to GRU’s hidden state, we can see that LSTM’s cell state recurrence resembles the aforementioned parallel scan’s formulation ${\bm{v}}_{t}={\bm{a}}_{t}\odot{\bm{v}}_{t-1}+{\bm{b}}_{t}$ where ${\bm{a}}_{t}\leftarrow{\bm{f}}_{t}$, ${\bm{b}}_{t}\leftarrow{\bm{i}}_{t}\odot\tilde{{\bm{c}}}_{t}$, and ${\bm{v}}_{t}\leftarrow{\bm{c}}_{t}$.
However, ${\bm{f}}_{t}$, ${\bm{i}}_{t}$ and $\tilde{{\bm{c}}}_{t}$ are dependent on the previous hidden state ${\bm{h}}_{t}$. As such, LSTM’s cell state recurrence is unable to apply the parallel scan algorithm as is. We can address this in a similar fashion to GRU by removing their hidden state dependencies as follows:

$$ $\begin{aligned} {\bm{f}}_{t}&=\sigma({\mathrm{Linear}}_{d_{h}}([{\bm{x}}_{t},{\bm{h}}_{t-1}]))\\ {\bm{i}}_{t}&=\sigma({\mathrm{Linear}}_{d_{h}}([{\bm{x}}_{t},{\bm{h}}_{t-1}]))\\ \tilde{{\bm{c}}}_{t}&=\mathrm{tanh}({\mathrm{Linear}}_{d_{h}}([{\bm{x}}_{t},{\bm{h}}_{t-1}]))\\ \end{aligned}\quad\Rightarrow\quad\begin{aligned} {\bm{f}}_{t}&=\sigma({\mathrm{Linear}}_{d_{h}}({\bm{x}}_{t}))\\ {\bm{i}}_{t}&=\sigma({\mathrm{Linear}}_{d_{h}}({\bm{x}}_{t}))\\ \tilde{{\bm{c}}}_{t}&=\mathrm{tanh}({\mathrm{Linear}}_{d_{h}}({\bm{x}}_{t}))\\ \end{aligned}$ $$

#### 3.2.2 Step 2: Drop range restriction of candidate states

Similar to GRUs, LSTMs leverage the hyperbolic tangent function ($\mathrm{tanh}$) to restrict the range of its states between $(-1,1)$.
LSTMs apply the range restriction twice: once when computing the candidate cell state and once computing its hidden state.
In this step, we drop both as follows:

$$ $\begin{aligned} \tilde{{\bm{c}}}_{t}&=\mathrm{tanh}({\mathrm{Linear}}_{d_{h}}({\bm{x}}_{t}))\\ {\bm{h}}_{t}&={\bm{o}}_{t}\odot\mathrm{tanh}({\bm{c}}_{t})\\ \end{aligned}\quad\Rightarrow\quad\begin{aligned} \tilde{{\bm{c}}}_{t}&={\mathrm{Linear}}_{d_{h}}({\bm{x}}_{t})\\ {\bm{h}}_{t}&={\bm{o}}_{t}\odot{\bm{c}}_{t}\\ \end{aligned}$ $$

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

In many sequence modelling settings (e.g., text generation), the optimization objective/target is time-independent in scale.
Recall LSTM’s cell state recurrence ${\bm{c}}_{t}={\bm{f}}_{t}\odot{\bm{c}}_{t-1}+{\bm{i}}_{t}\odot\tilde{{\bm{c}}}_{t}$ where ${\bm{i}}_{t},{\bm{f}}_{t}\in(0,1)^{d_{h}}$, and GRU’s hidden state recurrence(^2^22A superscript is added to differentiate GRU’s hidden state from LSTM’s.), ${\bm{h}}^{GRU}_{t}=(\bm{1}-{\bm{z}}_{t})\odot{\bm{h}}^{GRU}_{t-1}+{\bm{z}}_{t}\odot\tilde{{\bm{h}}}^{GRU}_{t}$ where ${\bm{z}}_{t}\in(0,1)^{d_{h}}$.
GRUs retain $(\bm{1}-{\bm{z}}_{t})\in(0,1)$ of the previous hidden state and add ${\bm{z}}_{t}$ of the new candidate state.
Since these proportions sum to $\bm{1}$, the model ensures its outputs (i.e., hidden states) are time-independent in scale.
In contrast, LSTM’s forget and input gates are computed independently (e.g., ${\bm{f}}_{t},{\bm{i}}_{t}\rightarrow\bm{1}$ or ${\bm{f}}_{t},{\bm{i}}_{t}\rightarrow\bm{0}$), making its cell states time-dependent in scale(^3^33For example, ${\bm{c}}_{t}\rightarrow{\bm{c}}_{0}+\sum_{i=1}^{t}\tilde{{\bm{c}}}_{t}$ when ${\bm{f}}_{1:t},{\bm{i}}_{1:t}\rightarrow 1$, growing in scale as the sequence length increases.) and optimization more difficult.
As such, we ensure LSTM’s output is time-independent in scale.

To do so, we can simply normalize the two gates, i.e., ${\bm{f}}^{\prime}_{t},{\bm{i}}^{\prime}_{t}\leftarrow\frac{f_{t}}{{\bm{f}}_{t}+{\bm{i}}_{t}},\frac{{\bm{i}}_{t}}{{\bm{f}}_{t}+{\bm{i}}_{t}}$, ensuring that ${\bm{f}}^{\prime}_{t}+{\bm{i}}^{\prime}_{t}=\bm{1}$ and the scale of LSTM’s cell state is time-independent.
Ensuring that the hidden state is time-independent in scale, we also drop the output gate ${\bm{o}}_{t}$ which scales the hidden state.
Without the output gate, the normalized hidden state is equal to the cell state, i.e., ${\bm{h}}_{t}={\bm{o}}_{t}\odot{\bm{c}}_{t}\,\Rightarrow\,{\bm{h}}_{t}={\bm{c}}_{t}$, making having both a hidden and cell state unnecessary.
As such, we drop the cell state as well.
In summary, the modifications are as follows:

$$ $\begin{aligned} {\bm{h}}_{t}&={\bm{o}}_{t}\odot{\bm{c}}_{t}\\ {\bm{o}}_{t}&=\sigma({\mathrm{Linear}}_{d_{h}}({\bm{x}}_{t}))\\ {\bm{c}}_{t}&={\bm{f}}_{t}\odot{\bm{c}}_{t-1}+{\bm{i}}_{t}\odot\tilde{{\bm{c}}}_{t}\\ \tilde{{\bm{c}}}_{t}&={\mathrm{Linear}}_{d_{h}}({\bm{x}}_{t})\\ \end{aligned}\quad\Rightarrow\quad\begin{aligned} {\bm{h}}_{t}&={\bm{f}}^{\prime}_{t}\odot{\bm{h}}_{t-1}+{\bm{i}}^{\prime}_{t}\odot\tilde{{\bm{h}}}_{t}\\ \tilde{{\bm{h}}}_{t}&={\mathrm{Linear}}_{d_{h}}({\bm{x}}_{t})\\ {\bm{f}}^{\prime}_{t},{\bm{i}}^{\prime}_{t}&\leftarrow\frac{{\bm{f}}_{t}}{{\bm{f}}_{t}+{\bm{i}}_{t}},\frac{{\bm{i}}_{t}}{{\bm{f}}_{t}+{\bm{i}}_{t}}\\ \end{aligned}$ $$

Notably, GRUs do not need this step as their outputs are already time-independent in scale.

#### 3.2.4 minLSTM

Combining the three steps results in a minimal version of LSTM (minLSTM):

$$ $\Rightarrow$ $$

The minimal version (minLSTM) is significantly more efficient (1) requiring only $O(3d_{h}d_{x})$ parameters compared to LSTM’s $O(4d_{h}(d_{x}+d_{h}))$. Furthermore, minLSTM (2) can be trained in parallel using the parallel scan algorithm, speeding up training significantly. For example, in Section 4.1, we found that minLSTM corresponded to a $235\times$ speedup for a sequence of length $512$ compared to LSTM on a T4 GPU. In terms of parameter efficiency, minLSTM uses only $38\%,25\%,19\%,$ or $15\%$ of parameters compared to LSTM when $\alpha=1,2,3,$ or $4$ respectively where $d_{h}=\alpha d_{x}$.

## 4 Were RNNs All We Needed?

In this section, we compare the minimal versions (minLSTMs and minGRUs) with their traditional counterparts (LSTMs and GRUs) and modern sequence models.
Pseudocode, PyTorch implementation, and detailed information regarding the experiment setup are available in the Appendix.

### 4.1 Minimal LSTMs and GRUs are very efficient

Figure: Figure 1: Training runtime (left), speedup (middle), and memory footprint (right) on a T4 GPU for a batch size of $64$. In the training runtime plot (left), minGRU, minLSTM, and Mamba lines overlap. These methods are approximately the same in training runtime.
Refer to caption: figs/Runtime-Seq_Len.png

At test time, recurrent sequence models are rolled out sequentially, making their inferences efficient.
Instead, the bottleneck of traditional RNNs is their training which requires linear training time (backpropagating through time) which resulted in their eventual deprecation.
The renewed interest in recurrent sequence models is due to many new architectures being efficiently trained in parallel (Gu et al., ).
In this section, we compare the resources required to train the traditional RNNs (LSTM and GRU), their minimal versions (minLSTM and minGRU), and a recent state-of-the-art sequence model.
In particular, we focus on the comparison with Mamba (Gu & Dao, ) which has seen significant popularity recently.
For these experiments, we consider a batch size of $64$ and vary the sequence length.
We measure the total runtime and memory complexity of performing a forward pass through the models, computing a loss, and computing gradients via a backward pass.

Runtime. In terms of runtime (see Figure 1 (left)), the simplified versions of LSTM and GRU (minLSTM and minGRU) Mamba achieve similar runtimes.
Averaging over $100$ runs, 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.
In contrast, the traditional RNN counterparts (LSTMs and GRUs) required a runtime that scaled linearly with respect to sequence length.
For a sequence length of $512$, minGRUs and minLSTMs were $175\times$ and $235\times$ faster per training step (see Figure 1 (middle)) 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$.
As such, in a setting where minGRU would take a day to finish training for a fixed number of epochs, its traditional counterpart GRU could take over $3$ years.

Memory.
By leveraging a parallel scan algorithm to compute the outputs in parallel efficiently, minGRU, minLSTM, and Mamba create a larger computational graph, thus needing more memory compared to traditional RNNs (see Figure 1 (right)).
The minimal variants (minGRU and minLSTM) use $\sim 88\%$ more memory compared to their traditional counterpart.
Mamba uses $56\%$ more memory compared to minGRU.
In practice, however, runtime is the bottleneck when training RNNs.

Effect of removing ${\bm{h}}_{t-1}$.
The original LSTM and GRU compute their various gates using their inputs ${\bm{x}}_{t}$ and previous hidden states ${\bm{h}}_{t-1}$.
These models leverage their time-dependent gates to learn complex functions.
However, minLSTM and minGRU’s training efficiencies are achieved by dropping their gates’ dependencies on the previous hidden states ${\bm{h}}_{t-1}$.
As a result, minLSTM and minGRU’s gates are dependent only on their inputs ${\bm{x}}_{t}$, resulting in a simpler recurrent module.
As such, the gates of a model consisting of a single layer of minLSTM or minGRU are time-independent due to being conditioned on time-independent inputs ${\bm{x}}^{(1)}_{1:n}$.

**Table 1: Comparison of the number of layers on the Selective Copying Task (Gu & Dao, ).**
| Model | # Layers | Accuracy |
| --- | --- | --- |
| MinLSTM | 1 | 37.6 ± 2.0 |
| 2 | 85.7 ± 5.8 |  |
| 3 | 96.0 ± 2.8 |  |
| MinGRU | 1 | 37.0 ± 2.3 |
| 2 | 96.8 ± 3.2 |  |
| 3 | 99.5 ± 0.2 |  |

However, in deep learning, models are constructed by stacking modules.
Although the inputs to the first layer ${\bm{x}}^{(1)}_{1:n}$ is time-independent, its outputs ${\bm{h}}^{(1)}_{1:n}$ are time-dependent and are used as the inputs to the second layer, i.e., ${\bm{x}}^{(2)}_{1:n}\leftarrow{\bm{h}}^{(1)}_{1:n}$.
As such, beginning from the second layer onwards, minLSTM and minGRU’s gates will also be time-dependent, resulting in the modelling of more complex functions.
In Table 1, we compare the performance of the models with varying numbers of layers on the Selective Copying Task from the Mamba paper (Gu & Dao, ).
We can immediately see the impact of the time dependencies: increasing the number of layers to $2$ or more drastically increases the model’s performance.

Training Stability.
Another effect of the number of layers is increased stability with decreased variance in the accuracy as the number of layers increases (see Table 1).
Furthermore, although minLSTM and minGRU both solve the Selective Copying task, we can see that minGRU is an empirically more stable method than minLSTM, solving the task with more consistency and lower variance.
minLSTM discards old information and adds new information, controlling the ratio with two sets of parameters (forget and input gate).
During training, the two sets of parameters are tuned in different directions, making the ratio harder to control and optimize.
In contrast, minGRU’s discarding and adding of information is controlled by a single set of parameters (update gate), making it easier to optimize.

**Table 2: Selective Copy Task. minLSTM, minGRU, and Mamba’s S6 (Gu & Dao, ) are capable of solving this task. Other methods such as S4, H3, and Hyena at best only partially solve the task.**
| Model | Layer | Accuracy |
| --- | --- | --- |
| H3 | Hyena | 30.1 |
| Mamba | Hyena | 28.4 |
| S4 | S4 | 18.3 |
| H3 | S4 | 57.0 |
| Mamba | S4 | 56.4 |
| S4 | S6 | 97.0 |
| H3 | S6 | 99.7 |
| Mamba | S6 | 99.8 |
| minGRU | minGRU | 99.5 ± 0.2 |
| minLSTM | minLSTM | 96.0 ± 2.8 |

### 4.2 Minimal LSTMs and GRUs perform well

In the previous section, we showed the significant efficiency gains achieved by simplifying traditional RNNs.
Here, we explore the empirical performance aspect of these minimal versions of LSTMs and GRUs compared to several popular sequence models.

Selective Copy. We consider the long-range Selective Copying task from the Mamba paper (Gu & Dao, ).
Unlike the original Copying task (Arjovsky et al., ), the Selective Copying task’s input elements are randomly spaced relative to their output, making the task harder.
To solve the task, models are required to perform content-aware reasoning, memorizing relevant and filtering out irrelevant tokens.

In Table 2, we compare the simplified versions of LSTMs and GRUs (minLSTM and minGRU) against well-known recurrent sequence models that can trained in parallel: S4 (Gu et al., ), H3 (Fu et al., ), Hyena (Poli et al., ), and Mamba (S6) (Gu & Dao, ).
The results for these baselines are quoted from the Mamba paper.
Out of all of these baselines, only S6 from Mamba’s paper is capable of solving this task.
minGRU and minLSTM are also capable of solving the Selective Copying task, achieving comparable performance to S6 and outperforming all other baselines.
LSTMs and GRUs leverage content-aware gating mechanisms, making these minimal versions sufficient for solving this task that many popular sequence models fail to solve.

**Table 3: Reinforcement Learning results on the D4RL (Fu et al., ) datasets. We report the expert normalized returns (higher is better), following (Fu et al., 2020), averaged across five random seeds. The minimal versions of LSTM and GRU, minLSTM and minGRU outperform Decision S4 (David et al., ) and perform comparably with Decision Mamba (Ota, ), (Decision) Aaren (Feng et al., ) and Decision Transformer (Chen et al., ).**
| Dataset | DT | DS4 | DAaren | DMamba | minLSTM | minGRU |
| --- | --- | --- | --- | --- | --- | --- |
| HalfCheetah-M | 42.6 | 42.5 | 42.2 | 42.8 | 42.7 ± 0.7 | 43.0 ± 0.4 |
| Hopper-M | 68.4 | 54.2 | 80.9 | 83.5 | 85.0 ± 4.4 | 79.4 ± 8.2 |
| Walker-M | 75.5 | 78.0 | 74.4 | 78.2 | 72.0 ± 7.5 | 73.3 ± 3.3 |
| HalfCheetah-M-R | 37.0 | 15.2 | 37.9 | 39.6 | 38.6 ± 1.1 | 38.5 ± 1.1 |
| Hopper-M-R | 85.6 | 49.6 | 77.9 | 82.6 | 88.5 ± 4.7 | 90.5 ± 0.9 |
| Walker-M-R | 71.2 | 69.0 | 71.4 | 70.9 | 69.7 ± 10.7 | 72.8 ± 8.9 |
| HalfCheetah-M-E | 88.8 | 92.7 | 75.7 | 91.9 | 85.4 ± 1.7 | 86.3 ± 0.5 |
| Hopper-M-E | 109.6 | 110.8 | 103.9 | 111.1 | 110.3 ± 1.6 | 109.7 ± 2.7 |
| Walker-M-E | 109.3 | 105.7 | 110.5 | 108.3 | 110.3 ± 0.5 | 110.3 ± 0.4 |
| Average | 76.4 | 68.6 | 75.0 | 78.8 | 78.1 | 78.2 |

Reinforcement Learning. Next, we consider the MuJoCo locomotion tasks from the D4RL benchmark (Fu et al., ).
Specifically, we consider the three environments: HalfCheetah, Hopper, and Walker.
For each environment, the models are trained on three datasets of varying data quality: Medium (M), Medium-Replay (M-R), and Medium-Expert (M-E).

In Table 3, we compare minLSTM and minGRU with various Decision Transformer variants, including the original Decision Transformer (DT) (Chen et al., ), Decision S4 (DS4) (David et al., ), Decision Mamba (Ota, ), and (Decision) Aaren (Feng et al., ).
The baseline results are retrieved from the Decision Mamba and Aaren papers.
minLSTM and minGRU outperform Decision S4 and achieve performance competitive with Decision Transformer, Aaren, and Mamba.
Unlike other recurrent methods, Decision S4 is a model whose recurrence transitions are not input-aware, affecting their performance.
In terms of average score across the $3\times 3=9$ datasets, minLSTM and minGRU outperform all the baselines except for Decision Mamba where the difference is marginal.

Figure: Figure 2: Language Modelling results on the Shakespeare dataset. Minimal versions of decade-old RNNs (LSTMs and GRUs) performed comparably to Mamba and Transformers. Transformers required $\sim 2.5\times$ more training steps to achieve comparable performance, overfitting eventually.
Refer to caption: figs/Shakespeare_learning_curve.png

Language Modelling. Finally, we consider a language modelling task. In this setting, we train a character-level GPT on the works of Shakespeare using the nanoGPT (Karpathy, ) framework.
In Figure 2, we plot the learning curves with a cross-entropy loss comparing the proposed minimal LSTM and GRU (minLSTM and minGRU) with Mamba and Transformers.
We found that minGRU, minLSTM, Mamba, and Transformers achieved comparable test losses of $1.548$, $1.555$, $1.575$, and $1.547$ respectively.
Mamba performed slightly worse than the other models but trained faster, particularly in the early stages, achieving its best performance at $400$ steps while minGRU and minLSTM continued training until $575$ and $625$ steps respectively.
In contrast, Transformers trained significantly slower, requiring $2000$ steps ($\sim 2.5\times$) more training steps than minGRU to achieve comparable performance, making it significantly slower and more resource-intensive to train (quadratic complexity compared to minGRU, minLSTM, and Mamba’s linear complexity).

## 5 Related Work

In this section, we provide a discussion of the similarities and differences between existing recurrent sequence models and the simplified versions of LSTMs and GRUs (minLSTM and minGRU).

State-Space Models (SSMs). Although Mamba (Gu & Dao, ) and state-space models have gained significant popularity recently, the steps towards the recent success of Mamba began years ago.
Gu et al. () first proposed a discretized structured state-space model.
Gu et al. () scaled the idea up, introducing S4.
The success of S4 became the basis for many future works (Gu et al., ; Gupta et al., ; Hasani et al., ; Smith et al., ) and state-space model applications in language (Mehta et al., ), audio (Goel et al., ), and more.
Recently, Mamba was a significant breakthrough in SSM, outperforming previous methods and garnering substantial attention.
A major novelty in Mamba was the proposal of S6, a state-space model whose transition matrices are input-dependent (i.e., $A_{t}$ and $B_{t}$ are functions of $x_{t}$).
In contrast, earlier state-space model transition matrices were input-independent, limiting their expressivity.
The success of Mamba and state-space models led to the writing of several survey papers (Wang et al., ; Patro & Agneeswaran, ; Qu et al., ).

Recurrent Versions of Attention. Another direction that proposed efficient recurrent sequence models is that of attention.
Building on variations of linear attention (Katharopoulos et al., ), several papers have introduced recurrent versions that can be computed in parallel.
Notably, Sun et al. () and Qin et al. () introduced variants that use an input-independent gating mechanism (decay factor).
More recently, Katsch () and Yang et al. () proposed linear attention variants that use input-dependent gating.
Feng et al. () showed softmax attention can be viewed as an RNN and proposed a recurrent model based on their RNN formulation.

Parallelizable RNNs.
Alternatively, several papers have proposed RNNs that can be trained efficiently in parallel.
Orvieto et al. () proposed an RNN that leverages complex diagonal recurrences and an exponential parameterization.
Beck et al. () proposed various enhancements to LSTM such as exponential gating, covariance update rule, and a normalizer state.

Although these three directions of designing efficient recurrent sequence models have proposed vastly different architectures, the core recurrent component of these models is remarkably similar.
For example, although state-space models are typically written as ${\bm{h}}_{t}={\bm{A}}_{t}{\bm{h}}_{t-1}+{\bm{B}}_{t}{\bm{x}}_{t}$, in practice, the transition matrices are typically diagonal for efficiency reasons.
As such, Mamba’s S6 (Gu & Dao, ) can be viewed as ${\bm{h}}_{t}=\mathrm{diag}({\bm{A}}_{t})\odot{\bm{h}}_{t-1}+\mathrm{diag}({\bm{B}}_{t})\odot{\bm{x}}_{t}$ where ${\bm{A}}_{t}$ and ${\bm{B}}_{t}$ are functions of ${\bm{x}}_{t}$.
In contrast, consider the minimal version of GRU ${\bm{h}}_{t}=(\bm{1}-{\bm{z}}_{t})\odot{\bm{h}}_{t-1}+{\bm{z}}_{t}\odot\mathrm{Linear}_{d_{n}}({\bm{x}}_{t})$ and minimal version of LSTM ${\bm{h}}_{t}={\bm{f}}^{\prime}_{t}\odot{\bm{h}}_{t-1}+{\bm{i}}^{\prime}_{t}\odot\mathrm{Linear}_{d_{n}}({\bm{x}}_{t})$.
The recurrences of these models are similar.
The major difference between these minimal RNNs, Mamba’s S6, and other models is how their transitions (e.g., ${\bm{z}}_{t},{\bm{i}}^{\prime}_{t},{\bm{f}}^{\prime}_{t},{\bm{A}}_{t},$ and ${\bm{B}}_{t}$) are computed from the input token ${\bm{x}}_{t}$.

Parallel Scan. Generalizing across the families of methods (including minLSTM and minGRU), these recent sequence models can be viewed as members of the same family of functions trainable via a parallel scan: ${\bm{v}}_{t}={\bm{a}}_{t}\odot{\bm{v}}_{t-1}+{\bm{b}}_{t}$ (see Section 2.3) where ${\bm{a}}_{t}$ and ${\bm{b}}_{t}$ are functions of the input token ${\bm{x}}_{t}$.
Improving upon the parallel scan algorithm, several models (Yang et al., ; Gu & Dao, ) such as Mamba have proposed specialized hardware-efficient methods that leverage GPU’s memory hierarchy to reduce high I/O costs and speed up training.
In our work, we implemented minLSTM and minGRU in plain PyTorch.
However, due to the structural similarities in recurrences amongst the numerous methods that leverage parallel scan, many techniques such as chunking that apply to one work for speeding up training can also apply to others such as minGRU and minLSTM.

Parameter Initializations.
Unrolling the recurrences of these new recurrent sequence models over time often results in their outputs and gradients vanishing/exploding (Wang et al., ) due to time dependency in their output’s scale.
To ensure model stability, the parameters of many models such as state-space models are initialized according to special distributions (Gu et al., , ; Orvieto et al., ).
In contrast, we found that minLSTM and minGRU are already stable using the default PyTorch initialization.
Unlike SSMs, minLSTM and minGRU’s outputs are time-independent in scale, avoiding potential instabilities.

## 6 Conclusion

In this work, we revisited RNNs from over a decade ago: LSTMs and GRUs.
We show that these models are trainable via the parallel scan algorithm by removing their hidden state dependencies from their gates.
Simplifying these models further, we removed their constraints on output range and ensured their output was time-independent in scale.
These steps result in their minimal versions (minLSTM and minGRU).
Empirically, we showed that minLSTM and minGRU (1) address the computational limitations of their traditional counterparts and (2) are as computationally efficient as Mamba, a popular recent state-of-the-art recurrent sequence model, and (3) are competitive in performance with recent sequence models.
Considering the strong empirical performance of these simplified RNNs and their fundamental similarities with many recently proposed recurrent sequence methods, we question ”Were RNNs all we needed?”

## Limitations

Our experiments were run on P100 (16 GBs) and T4 (16 GBs) GPUs. Due to computation limitations, our experiments are smaller in scale compared to works such as Mamba (Gu & Dao, ) which leveraged A100 80GB GPUs.
To fit the selective copy task on the GPU, we leveraged gradient accumulation for training, splitting the standard batch size in half and slowing training significantly.
Nonetheless, we hypothesize that these conclusions generalize to larger-scale settings due to the fundamental similarities between the minimal RNNs (minLSTM and minGRU) and many recent sequence methods.

## Appendix A Implementation Details: Vanilla Version

In this section, we provide the pseudocode and equivalent PyTorch code for minGRU and minLSTM.
When performing repeated multiplications such as in many recurrent sequence models, numerical instabilities are common, especially during training.
As such, we trained using a log-space implementation (see Section B) for improved numerical stability.

### A.1 Pseudocode: Vanilla Version

#### A.1.1 minGRU: A Minimal GRU

Figure: Algorithm 1 Sequential Mode: Minimal Version of GRU (minGRU)

#### A.1.2 minLSTM: A Minimal LSTM

Figure: Algorithm 3 Sequential Mode: Minimal Version of LSTM (minLSTM)

### A.2 PyTorch Code: Vanilla Version

#### A.2.1 minGRU: A Minimal GRU

Figure: Listing 1: Sequential Mode: Minimal Version of GRU (minGRU)

Figure: Listing 2: Parallel Mode: Minimal Version of GRU (minGRU)

#### A.2.2 minLSTM: A Minimal LSTM

Figure: Listing 3: Sequential Mode: Minimal Version of LSTM (minLSTM)

Figure: Listing 4: Parallel Mode: Minimal Version of LSTM (minLSTM)

## Appendix B Implementation Details: Log-Space Version (Additional Numerical Stability)

In this section, we detail the log-space version of minLSTM and minGRU for improved numerical stability. During training, the parallel modes are used to avoid backpropagation through time (BPTT), speeding up the training time significantly. At inference time, the sequential modes are used.

### B.1 Parallel Scan: Log-Space Implementation

Recall that, the parallel scan’s objective is to compute ${\bm{h}}_{1:t}$ where ${\bm{h}}_{k}={\bm{a}}_{k}\odot{\bm{h}}_{k-1}+{\bm{b}}_{k}$.
In code, the vanilla parallel scan function would take as input: coefficients ${\bm{a}}_{1:t}$ and values ${\bm{b}}_{0:t}$. The function then outputs ${\bm{h}}_{1:t}$.
For numerical stability, we consider a log-space implementation which takes as input $\log({\bm{a}}_{1:t})$ and $\log({\bm{b}}_{0:t})$ instead and outputs ${\bm{h}}_{1:t}$. The code for the parallel scan in log-space is included below and is based on the code by Heinsen ().

Figure: Listing 5: Parallel scan based on Heinsen (). This function computes ${\bm{h}}_{1:t}$ given log coefficients $\log({\bm{a}}_{1:t})$ and log values $\log({\bm{b}}_{0:t})$.

### B.2 Pseudocode: Log-Space Version

For maximal numerical stability, we rewrite the log-space versions of minGRU and minLSTM.

#### B.2.1 minGRU: A Minimal GRU

Recall minGRU’s recurrence is as follows ${\bm{h}}_{t}\leftarrow(\bm{1}-{\bm{z}}_{t})\odot{\bm{h}}_{t-1}+{\bm{z}}_{t}\odot\tilde{{\bm{h}}}_{t}$.
As such, ${\bm{a}}_{t}\leftarrow(\bm{1}-{\bm{z}}_{t})$ and ${\bm{b}}_{t}\leftarrow{\bm{z}}_{t}\odot\tilde{{\bm{h}}}_{t}$ where ${\bm{z}}_{t}=\sigma({\bm{k}}_{t})$ and ${\bm{k}}_{t}=\mathrm{Linear}_{d_{h}}(x_{t})$.
As a result, $\log({\bm{a}}_{t})\leftarrow\log(\bm{1}-{\bm{z}}_{t})$ and $\log({\bm{b}}_{t})\leftarrow\log({\bm{z}}_{t})+\log(\tilde{{\bm{h}}}_{t})$. We can break down these down as follows:

$$ $\displaystyle\log({\bm{z}}_{t})$ $\displaystyle=\log(\sigma({\bm{k}}_{t}))$ $\displaystyle=\log\left(\frac{1}{1+\exp(-{\bm{k}}_{t})}\right)$ $\displaystyle=-\mathrm{Softplus}(-{\bm{k}}_{t})$ $\displaystyle\log(\bm{1}-{\bm{z}}_{t})$ $\displaystyle=\log\left(\frac{\exp(-{\bm{k}}_{t})}{1+\exp(-{\bm{k}}_{t})}\right)$ $\displaystyle=\log\left(\frac{1}{1+\exp({\bm{k}}_{t})}\right)$ $\displaystyle=-\mathrm{Softplus}({\bm{k}}_{t})$ $$

where ${\bm{k}}_{t}=\mathrm{Linear}_{d_{h}}({\bm{x}}_{t})$. However, we need to compute $\log(\tilde{{\bm{h}}})_{t}$ which is inconvenient if $\tilde{{\bm{h}}}_{t}$ has some negative values.
We could use complex numbers and a complex number version of the parallel scan, but this would result in the parallel scan increasing in complexity.
Instead, we propose to ensure that $\tilde{{\bm{h}}}_{t}>0$. This be can done in a variety of ways.
In our experiments, we added a continuous activation function $g$ replacing $\tilde{{\bm{h}}}_{t}\leftarrow\mathrm{Linear}_{d_{h}}(x_{t})$ with $\tilde{{\bm{h}}}_{t}\leftarrow g(\mathrm{Linear}_{d_{h}}(x_{t}))$ where $g(x)=\begin{cases}x+0.5,&\text{if }x\geq 0\\
\sigma(x),&\text{otherwise}\end{cases}$ and its log: $\log(g(x))=\begin{cases}\log(x+0.5),&\text{if }x\geq 0\\
-\mathrm{Softplus}(-x),&\text{otherwise}\end{cases}$.

At inference time, the sequential mode (Algorithm 5) is used.
During training, the parallel mode (Algorithm 6) is used.

Figure: Algorithm 5 Sequential Mode: Minimal Version of GRU (minGRU) trained in log-space

#### B.2.2 minLSTM: A Minimal LSTM

We also derive minLSTM’s log-space formulation as well. Recall minLSTM’s recurrence is as follows ${\bm{h}}_{t}\leftarrow{\bm{f}}^{\prime}_{t}\odot{\bm{h}}_{t-1}+{\bm{i}}^{\prime}_{t}\odot\tilde{{\bm{h}}}_{t}$.
As such, ${\bm{a}}_{t}\leftarrow{\bm{f}}^{\prime}_{t}$ and ${\bm{b}}_{t}\leftarrow{\bm{i}}^{\prime}_{t}\odot\tilde{{\bm{h}}}_{t}$.
As a result, $\log({\bm{a}}_{t})\leftarrow\log({\bm{f}}^{\prime}_{t})$ and $\log({\bm{b}}_{t})\leftarrow\log({\bm{i}}^{\prime}_{t})+\log(\tilde{{\bm{h}}}_{t})$.

$$ $\displaystyle\log({\bm{f}}^{\prime}_{t})$ $\displaystyle=\log\left(\frac{{\bm{f}}_{t}}{{\bm{f}}_{t}+{\bm{i}}_{t}}\right)$ $\displaystyle=\log\left(\frac{1}{1+\frac{{\bm{i}}_{t}}{{\bm{f}}_{t}}}\right)$ $\displaystyle=-\log\left(1+\frac{{\bm{i}}_{t}}{{\bm{f}}_{t}}\right)$ $\displaystyle=-\log\left(1+\exp\left(\log\left(\frac{{\bm{i}}_{t}}{{\bm{f}}_{t}}\right)\right)\right)$ $\displaystyle=-\mathrm{Softplus}\left(\log\left(\frac{{\bm{i}}_{t}}{{\bm{f}}_{t}}\right)\right)$ $\displaystyle=-\mathrm{Softplus}\left(\log({\bm{i}}_{t})-\log({\bm{f}}_{t})\right)$ $$

Recall that ${\bm{i}}_{t}$ and ${\bm{f}}_{t}$ are computed via sigmoid. In other words, ${\bm{i}}_{t}=\sigma({\bm{k}}_{t})$ and ${\bm{f}}_{t}=\sigma({\bm{p}}_{t})$ where ${\bm{k}}_{t}=\mathrm{Linear}_{d_{h}}({\bm{x}}_{t})$ and ${\bm{p}}_{t}=\mathrm{Linear}_{d_{h}}({\bm{x}}_{t})$.
Furthermore, recall in minGRU’s derivation we showed that $\log(\sigma({\bm{k}}_{t}))=-\mathrm{Softplus}(-{\bm{k}}_{t})$
Using this, we can simplify the computation as follows:

$$ $\displaystyle\log({\bm{f}}^{\prime}_{t})$ $\displaystyle=-\mathrm{Softplus}\left(\log(\sigma({\bm{k}}_{t}))-\log(\sigma({\bm{p}}_{t}))\right)$ $\displaystyle=-\mathrm{Softplus}\left(\mathrm{Softplus}(-{\bm{p}}_{t})-\mathrm{Softplus}(-{\bm{k}}_{t}))\right)$ $$

Similarly, we also get that:

$$ $\log({\bm{i}}^{\prime}_{t})=-\mathrm{Softplus}\left(\mathrm{Softplus}(-{\bm{k}}_{t})-\mathrm{Softplus}(-{\bm{p}}_{t}))\right)\\$ $$

Combining these derivations, we get the parallel mode (Algorithm 8) for efficient training.

Figure: Algorithm 7 Sequential Mode: Minimal Version of LSTM (minLSTM) trained in log-space

### B.3 PyTorch Code: Log-Space Version

Figure: Listing 6: The continuous function $g$ ensures that $\tilde{{\bm{h}}}_{t}\leftarrow g(\mathrm{Linear}_{d_{h}}(x_{t}))$ is positive.

#### B.3.1 minGRU: A Minimal GRU

Figure: Listing 7: Sequential Mode: Minimal Version of GRU (minGRU) trained in log-space

Figure: Listing 8: Parallel Mode: Minimal Version of GRU (minGRU) for training in log-space

#### B.3.2 minLSTM: A Minimal LSTM

Figure: Listing 9: Sequential Mode: Minimal Version of LSTM (minLSTM) trained in log-space

Figure: Listing 10: Parallel Mode: Minimal Version of LSTM (minLSTM) for training in log-space

## Appendix C Detailed Experiment Setup

In this section, we describe the experiment setup in detail.

### C.1 Architecture

In all models, residual connections are added between layers and layer norms are applied before each layer.

Selective Copying. Each layer in the model consisted of (1) either a minLSTM or minGRU layer and (2) a linear layer.

Reinforcement Learning. In this work, we consider the Decision Transformer framework for (Offline) RL. Following prior works (Feng et al., ; Ota, ), we replace the Self-Attention module with our recurrent sequence modules: minLSTM and minGRU respectively.

Language Modelling. Prior works (Gu & Dao, ; Beck et al., ) apply a convolutional layer in addition to their recurrent sequence module.
Following them, a layer of the model consists of an RNN (minLSTM and minGRU), a convolutional layer applied temporally with a kernel size of $4$, and a two-layer MLP.

### C.2 Hyperparameters and general experimental details

Selective Copying. Models are trained for $400,000$ steps with a batch size of $64$, and using early stopping.
The optimizer used is Adam with a learning rate of $3\times 10^{-4}$.
Due to GPU memory limitations, gradient accumulation is performed during training.
Gradients for two batches of size $32$ are accumulated for each gradient update and clipped to $1.0$.
Each model consists of $3$ layers, an input dimension of $64$, and a dropout ratio of $0.1$.
minLSTM and minGRU have an expansion factor of $6$.
Results for the baselines are referenced from the Mamba paper.

Reinforcement Learning.
We follow the hyperparameter settings outlined by Ota (). For Hopper (Medium) and Hopper (Medium-Replay), an embedding dimension of $256$ is used, while all other environments utilize an embedding dimension of $128$. The learning rate is set to $1\times 10^{-4}$ for Hopper (Medium), Hopper (Medium-Replay), and Walker (Medium). For all other environments and datasets, the learning rate is $1\times 10^{-3}$.
The models are optimized using AdamW with a weight decay of $1\times 10^{-4}$ and a linear warmup for $10,000$ steps. Each model consists of $3$ layers and has a dropout ratio of $0.1$. The models are trained for $100,000$ steps with a batch size of $64$. Gradients are clipped to $0.25$. Results for the baselines are referenced from the Mamba and Aaren papers.

Language Modelling. The models are optimized using AdamW with a learning rate of $1\times 10^{-3}$. Each model consists of three layers, a dropout ratio of $0.2$, and an embedding dimension of $384$. Training is done with $5000$ steps using a batch size of 64 and evaluated every $25$ steps. Gradients are clipped to $1.0$. The Transformer is configured with $6$ heads. Mamba uses an SSM state expansion factor of $16$ and a block expansion factor of $2$. Following Mamba, both minLSTM and minGRU utilize an expansion factor of $2$ as well.

### C.3 Datasets

Selective Copying. In this task, the model learns to extract data tokens from a sequence while disregarding noise tokens. Following Gu & Dao (), we consider a vocabulary of $16$ and sequences of length $4096$. Each sequence includes $16$ randomly placed data tokens. The remainder of the tokens are noise.

Reinforcement Learning.
In this setting, we consider continuous control tasks from the D4RL benchmark (Fu et al., ). These tasks based on MuJoCo comprise of three environments with dense rewards: HalfCheetah, Hopper, and Walker. For each environment, three different datasets are considered that have varying level represent varying levels of data quality:

- •
Medium (M): One million timesteps generated by a policy scoring about one-third of an expert policy’s score.
- •
Medium-Replay (M-R): A replay buffer from an agent trained to perform like the Medium policy.
- •
Medium-Expert (M-E): One million timesteps from the Medium policy combined with one million from an expert policy.

Following Fu et al. (), reported scores are normalized such that $100$ represents an expert policy performance.

Language Modelling. In this setting, we consider the Shakespeare dataset, comprising a collection of text data derived from the works of William Shakespeare. The training and testing data consists of $1,003,854$ and $111,540$ tokens respectively.

## Peer Reviews

confidence: 4
contribution: 4
initial_rating: 3
presentation: 3
questions: Under "Parameter Initializations" you state that minLSTM and minGRU are stable with default initialization, in contrast to the specialized initializations in prior linear RNN papers. Did you perform any studies using different initialization methods?

When trying to replicate minGRU using the log-space pseudocode I am struggling to get the sequential and parallel modes to match, could this pseudocode be checked through? (update: I managed to do it using the cited paper, but I think the pseudocode provided is incorrect)
soundness: 4
strengths: The core idea of creating minimal versions of LSTM and GRU models for efficient parallel training is compelling.
I believe this contribution is novel and could be highly useful.
The benchmark results are encouraging and the comparisons made with other models seem appropriate.
summary: The paper introduces minimal Long Short-Term Memory (minLSTM) and minimal Gated Recurrent Unit (minGRU) models.
The authors modify the traditional LSTM and GRU models by removing dependencies on prior hidden states from the gating mechanisms, enabling parallelization through a method similar to approaches in linear RNNs and state-space models.
The changes to the gating mechanisms significantly alter how minLSTMs and minGRUs are expected to function when compared with LSTMs and GRUs, but it allows for efficient computation without the limitations that result from back-propagation through time.
Comparisons are made with state-of-the-art models, including the state-space model Mamba. 
Empirical results demonstrate that the minLSTM and minGRU models achieve competitive performance on the selective copy task, multiple reinforcement learning tasks, and a language modeling task based on the Shakespeare dataset.
weaknesses: As the authors point out, computational restrictions prevent them from providing large-scale experiments. I would be curious to see performance on additional benchmarks, such as WikiText103, Pile or the long range arena. However, I still believe the submission is strong without them. 

----

Small typo: line 370, should read '... recurrent sequence models that can *be* trained in parallel ...'

in B.1, parallel_scan_log: log_x0_plus_b_star is not defined, should this be log_h0_plus_b_star?

in B.3.1 the pseudocode takes x_t but the code uses x 

For the purposes of reproducibility, providing anonymized code in the supplementary material would strengthen the submission.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. **Layer Depth in Computational Comparisons:** (Please refer to Weaknesses) Were the computational benchmarks (Figure 1) conducted with models having comparable quantitative performances? The paper could benefit from a comparison on computational complexity normalized by model performances or depth. 

2. **Parameter Matching for Competitors:** Related to previous question. When taking results from other papers (as reported in the Appendix), were the competitor models matched in terms of parameters and model size? Variability in model configurations could impact the validity of performance comparisons.

3. **Choice of Datasets:** The model is mainly tested on very simple datasets (mostly synthetic) that are not very representative for the current literature (see [3]). The authors should comment on why this choice taken.  If the model cannot scale to large datasets due to its inner working and memory requirements (as reported in the Limitation section), this is something that hinder the paper contributions.  For instance, the Long Range Arena benchmark has became a standardize framework to test sequence models performances, that could help also in identifying the model ability to preserve long term dependencies -- that is something extremely requested by current literature. The removal of time dependecy questions the model ability to perform in such benchmarks, which I believe should be better investigated. Thus,  at least a discussion on this is required. 

4. **Theoretical Comparison with xLSTM [1]:** Could the authors elaborate on the theoretical distinctions and advantages of minLSTM over xLSTM, especially in terms of efficiency and sequence modeling?
soundness: 3
strengths: - The approach effectively repurposes older RNNs by leveraging simplifications that enable parallel training, providing an interesting contrast to complex, modern architectures.

- The proposed models significantly reduce training time, achieving speed improvements of up to 175x for sequence lengths of 512, which is a notable practical advantage.

- This work challenges the abandonment of RNNs in favor of more recent architectures and suggests potential for simpler, more interpretable models in long-sequence processing.
summary: The paper investigates simplified versions of traditional recurrent neural networks (RNNs), specifically LSTMs and GRUs, adapting them for efficient training by removing their hidden state dependencies, enabling parallel training. The proposed **minLSTM** and **minGRU** models retain the functional structure of their predecessors but omit time dependencies in their gates, reducing parameter count and improving computational efficiency. Empirical tests indicate these models achieve comparable performance to contemporary state-of-the-art sequence models, suggesting that streamlined versions of older RNN architectures may offer viable alternatives in sequence modeling tasks.
weaknesses: *Insufficient Comparison Context:* From the current text, it is not clear whether the computational comparisons (Figure 1) were carried on considering the fact that the proposed models potentially require more layer to achieve competitor quantitative performances on tasks (Table 1). This fact potentially skews the results if competitors require fewer layers for comparable performance.

*Limited Dataset Representativeness:* The model's evaluation primarily relies on synthetic, simplified datasets, which are not representative of those used in current literature benchmarks (see Section 7/Table 4 in  [3]). This limits insights into the model’s scalability and ability to handle real-world data complexities (and it is not clear whether this was caused by memory requirements drawbacks). Moreover, the removal of hidden state dependencies from gates raises concerns regarding the model ability to preserve long range dependendencies - something that should be investigated. 

*Literature Comparisons and references:* The paper references but does not thoroughly compare against xLSTM, a similar recent work aimed at improving LSTM performance. A more detailed theoretical (at least) comparison with xLSTM would strengthen the paper's argument for minLSTM's contributions. Moreover, the paper is inspired and puts emphasis (even in the title) on the recent **resurgence** of RNNs - a concept that was highligthed by recent surveys [2,3] that could help the reader better contextualize the work.  In the aforementioned surveys, other very related works [4,5] are described that should be at least mentioned 

[1] Beck, Maximilian, et al. "xLSTM: Extended Long Short-Term Memory." arXiv preprint arXiv:2405.04517 (2024).

[2] Tiezzi, Matteo, et al. "On the resurgence of recurrent models for long sequences: Survey and research opportunities in the transformer era." arXiv preprint arXiv:2402.08132 (2024).

[3] Tiezzi, Matteo, et al. "State-Space Modeling in Long Sequence Processing: A Survey on Recurrence in the Transformer Era." arXiv preprint arXiv:2406.09062 (2024).

[4] Bradbury, James, et al. "Quasi-recurrent neural networks." arXiv preprint arXiv:1611.01576 (2016).

[5] Martin, Eric, and Chris Cundy. "Parallelizing linear recurrent neural nets over sequence length." arXiv preprint arXiv:1709.04057 (2017).

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 3
questions: Can the authors provide a detailed comparison with Martin & Cundy (2018), specifically addressing how minGRU differs from GILR, if at all, and what novel contributions this work makes beyond the prior work? Additionally, it would be valuable to understand why these ideas are being revisited now and what new insights are gained.

How do the proposed new architectures fare at larger data sizes? How do these models perform on standard language modeling benchmarks? Results on datasets with training sequences longer than 1000 steps would be particularly good to see. Understanding the scaling behavior on modern language modeling tasks is crucial for assessing the practical value of this approach.

Finally, more detailed inference speed comparisons would strengthen the paper. How do minGRU and minLSTM compare with Mamba and traditional LSTM/GRU implementations across different sequence lengths and batch sizes?
soundness: 3
strengths: The paper's technical exposition is clear, offering good motivation for removing parts of the RNN architectures. The authors provide thorough motivation for each architectural decision, especially regarding the time-independence properties of their models. The connection to parallel scan algorithms makes sense.

The architectural innovations demonstrate consideration of practical implementation challenges. The rescaling mechanism in minLSTM ensures time-independence while maintaining model expressivity. The authors have given careful thought to numerical stability issues, providing both standard and log-space formulations of their approach. It is nice to see the PyTorch implementation details, which help to make the algorithm very concrete.

The experimental validation covers a range of tasks and provides evidence of training speed advantages wrt traditional RNNs. The performance on the selective copying task is particularly impressive, demonstrating the model's ability to handle a particular type of long-range dependencies.
summary: The paper introduces minimal variants of LSTM and GRU (minLSTM and minGRU) that enable parallel training via the parallel scan algorithm. Key modifications include removing hidden state dependencies and rescaling mechanisms in the LSTM. The authors demonstrate competitive performance with modern architectures on several tasks.
weaknesses: The most significant concern is the paper's relationship to prior work, specifically Martin & Cundy (ParalleIizing Linear Recurrent Neural Nets over Sequence Length, ICLR 2018), which appears to have developed very similar ideas. The proposed minGRU architecture appears mathematically equivalent to their GILR architecture when \(g_t = 1 - z_t\) and \(i_t = \tilde{h}_t\) in the authors' notation. The parallel scan approach for training is also very similar. While the LSTM rescaling mechanism appears novel, this substantial overlap significantly reduces the paper's claimed novelty.

The experimental evaluation, while broad, has several limitations. The benefits over baselines are modest in most tasks, with the exception of the selective copying task, which seems carefully chosen to showcase the method's strengths. The language modeling evaluation relies solely on the Shakespeare dataset, which is both too small by modern standards. This choice of dataset is not representative of real-world language modeling challenges, which typically do not have overfitting as a major worry.

The methodology would benefit from more thorough analysis in several areas. The authors could strengthen their contribution by providing more extensive ablation studies on their architectural choices. The paper would also benefit from a clearer discussion of scaling behavior, to indicate if this architecture could be a potential replacement for transformers or Mamba at the larger scales of modern models.

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 4
questions: If my arguments above are valid, I don't have any question that can change my opinion of the paper.
soundness: 4
strengths: The paper is well written and easy to follow. It convincingly demonstrates the speed benefits of min LSTM and GRU over their sequential counterparts.
summary: Transformers are expensive to run on long sequences due to their quadratic complexity. Traditional RNNs such as LSTMs and GRUs have cheaper inference costs, but they cannot be parallelized during training. This paper shows how we can minimally modify these architectures to train them with parallel scans. This greatly reduces training runtime, while still maintaining good empirical performance (as demonstrated on small scale reinforcement learning and language modelling tasks).
weaknesses: I am afraid that the paper is half a year to one year late. Both the xLSTM (Beck et al. 2024) and the MatMul-free (Zhu et al. 2024) papers leverage a similar insight to the one of this paper, that is removing everything in LSTM/GRUs that cannot be parallelized during training. Those papers demonstrate the effectiveness of this approach at a larger scale (more than 10x the number of parameters). In addition to that, detailed comparison to these works is missing. For these reasons, I think the paper requires heavy modifications (e.g. more detailed experiments, at a larger scale) before being accepted.

Meta-review: This paper introduces minLSTM and minGRU, simplified variants of traditional LSTM and GRU models, designed to enable parallel training by removing hidden state dependencies. These modifications result in models that are computationally efficient and achieve competitive performance on various tasks, rivaling recent models including transformers, while using fewer parameters and being fully parallelizable during training. Another contribution of the paper comes through its framing, and the attempt to spark a discussion within the community about the evolution of sequence modeling and the importance of reassessing simpler RNNs. The reviewers had differing opinions about the value this paper adds to that discussion in light of existing literature. Additionally, some reviewers raised concerns about the paper's novelty and about the empirical evaluation, questioning scalability and whether the results would be representative of modern language modeling tasks. Altogether, I do not believe this paper is appropriate for publication at ICLR at this time.

Decision: Reject