SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$^{1}$ Zhen Qin*,  $^{2}$ Songlin Yang*,  $^{1}$ Yiran Zhong
$^{1}$ OpenNLPLab, Shanghai Artificial Intelligence Laboratory,  $^{2}$ MIT CSAIL
https://github.com/OpenNLPLab/HGRN
Transformers have surpassed RNNs in popularity due to their superior abilities in parallel training and long-term dependency modeling. Recently, there has been a renewed interest in using linear RNNs for efficient sequence modeling. These linear RNNs often employ gating mechanisms in the output of the linear recurrence layer while ignoring the significance of using forget gates within the recurrence. In this paper, we propose a gated linear RNN model dubbed Hierarchically Gated Recurrent Neural Network (HGRN), which includes forget gates that are lower bounded by a learnable value. The lower bound increases monotonically when moving up layers. This allows the upper layers to model long-term dependencies and the lower layers to model more local, short-term dependencies. Experiments on language modeling, image classification, and long-range arena benchmarks showcase the efficiency and effectiveness of our proposed model. The source code is available at https://github.com/OpenNLPlab/HGRN.
Sequence modeling is a fundamental problem in various domains such as natural language processing [12, 43, 44, 61, 64], time series analysis [84], computer vision [3, 13, 45, 74], and audio processing [1, 18, 73]. Prior to the invention of Transformers [81], RNN and its variants were the primary selections of architectures for sequence modeling, and have been widely used in machine translation [6], stock price prediction [68], weather forecasting [65], speech recognition [51], and etc.
RNNs have two main drawbacks: slow sequential training and limited capability in modeling long-term dependencies. With the swift development of deep learning and the pervasive use of GPUs, these drawbacks prevent it from flourishing in modern long-sequence modeling tasks. Meanwhile, Transformers [81] have rapidly gained popularity and now dominate various research areas in sequence modeling due to their better abilities in parallel training and long-term dependency modeling. However, Transformer's quadratic time complexity makes long sequence modeling expensive. On the other hand, RNN offers linear complexity and serves as an ideal choice for long sequence modeling. This works aim to addressing these RNN drawbacks, revitalizing their applicability in long-sequence modeling tasks.
To address the training inefficiency problem, we turn to more efficient RNN variants that employ element-wise linear recurrence (ELR) relations [48]. ELR provides two main advantages: (i) By removing nonlinearities in the recurrence, it enables parallelized training. (ii) By assuming independence between distinct hidden states, it enables efficient hidden state updates (through element-wise product instead of matrix multiplication) [20, 40]. Notably, ELR has been used in many modern linear RNN models, including the diagonalized versions of structured state-space models [21] (S4) [20, 26, 71] and RWKV [55]. In recent advancements, numerous studies have incorporated
*Equal contribution. ☑ Indicates corresponding author (Email address: zhongyiran@gmail.com).
gating mechanisms into the outputs of linear recurrence layers [11, 46, 49, 55, 82], similar to the output gates in LSTMs and leading to considerable performance gains. However, most current studies overlook the significance of the forget gate, which is often regarded as the most important gate in LSTMs [19, 80]. In this work, we underscore the importance of employing forget gates in linear RNNs and adopt gated linear RNNs for both efficiency and high performance.
To effectively capture long-term dependencies in gated RNNs, it is crucial to maintain high forget gate values close to one [23]. However, gates in saturated regimes (i.e., close to zero or one) suffer from the gradient vanishing issue [23]. Moreover, if all forget gate values are close to one, RNNs will not be able to effectively forget irrelevant information, compromising their ability to model short-term dependencies. To address these challenges, we introduce Hierarchically Gated Recurrent Units(HGRU). In HGRU, we add an additive learnable value, referred to as the lower bound, to the original forget gate value, effectively mitigating the issue of saturated gates [23] by pushing gate activations away from the saturated regimes. Furthermore, we design the lower bounds to increase monotonically as we move up the layers of the RNN. This ensures that the forget gate values in the lower layers remain relatively small, enabling the necessary forgetting of past information for modeling short-term dependencies. In contrast, in the uppermost layer, the forget gate values approach one, facilitating the effective modeling of long-term dependencies. Our proposed model has proven to be highly efficient and effective, as demonstrated by its outstanding performance in language modeling, image classification, and long-range arena benchmarks.
Efficient token mixing for sequence modeling. [83] abstracts self-attention (SA) as token mixing, thereby transforming the Transformer architecture into MetaFormer. MetaFormer comprises essential components such as token mixer, channel mixer, residual connections, and LayerNorm. This abstraction highlights that the success of Transformers does not solely rely on SA but rather on the holistic integration of these components. Notably, token mixers can be replaced with simpler alternatives like pooling layers without compromising the model's performance in the context of vision transformer. For sequence modeling tasks, [29] provides a comprehensive analysis and discussion of different token mixing strategies. Two prominent contenders, long convolution and linear recurrence, show promise as replacements for SA modules in long sequence modeling due to their superior asymptotic time complexity and competitive performances. In long convolution models [14, 41, 57, 59], the kernel size matches the input sequence length, enabling a broader context compared to traditional convolutions. Training is accomplished using the efficient  $\mathcal{O}(n\log n)$  fast Fourier transforms (FFT) algorithm. However, long convolutions face challenges such as the need for causal convolution inference, which requires caching all historical computations similar to the key-value (KV) cache in SA. This can lead to memory limitations when processing long sequences. Moreover, the inference complexity of long convolutions remains higher than that of RNNs. These factors make linear RNNs a more suitable alternative to replace SA in long-sequence modeling. TransNormerLLM [61] scales efficient token mixing in large language models to achieve competitive performance and superior training and inference efficiency compared to transformer-based models.
Element-wise linear recurrence. The slower training speeds of traditional RNNs can be attributed to two main factors: (i) The updating of the hidden state involves full matrix multiplication. (ii) The presence of nonlinearity within the recurrence prevents parallel computation. To tackle the first issue, [40] introduced a simplified interaction between hidden states. This allowed the hidden state update to be performed using an element-wise product instead of full matrix multiplication. They demonstrated that this approach is notably fast when the (nonlinear) recurrence for each dimension is fused within a single CUDA kernel. Likewise, for the linear case, diagonalized versions of S4 [20, 26] have also exhibited speed improvements over S4 by leveraging element-wise recurrence. Regarding the second challenge, the ability to capture nonlinear dependencies on past data can be achieved by stacking multiple linear recurrence layers interleaved with nonlinear MLP blocks. This indicates the potential to eliminate nonlinearity, as suggested by [4, 25, 48]. Empirical support for this strategy's effectiveness came later, as demonstrated by [11, 20, 24, 53, 55, 71]. [52] further highlighted that such an architecture still possesses Universal Approximator properties, thus justifying the employment of linear recurrence. By excluding nonlinearity, [48, 71] showed that the parallel scan algorithm can be used for parallel training.
Linear recurrence can be broadly categorized into exponential moving averages (EMA) and gating schemes, as noted by [48]. The key difference is whether the decay rate is data-dependent. Models such as S4 [21], S4D [20], MEGA [46], RWKV [55], and LRU [53] utilize the EMA approach, where the decay rate is static for all time steps (i.e., data-independent), while our model uses a data-dependent dynamic decay rate through the use of the forget gate. We remark on the importance of incorporating a data-dependent decay rate, which is largely ignored by current works in linear RNNs. Although liquid S4 [28] uses a dynamic transition matrix (which amounts to a data-dependent decay rate), it employs a limited form for FFT-based training. Our model does not have the convolutional view and thus cannot use FFT for training but allows the use of parallel scan.
The field of linear Transformers and linear RNNs exhibits a close relationship. [34] shows that linear Transformers can be reformulated as RNNs during auto-regressive decoding, revealing similarities to the update rules observed in fast weight additive outer products [66, 67]. These updates can be seen as a special case of element-wise linear recurrence, where forget gate values are consistently set to one across time and hidden states are two-dimensional. However, this formulation in linear Transformers lacks the ability to forget irrelevant information, resulting in the attention dilution issue [60]. To address this limitation, [66] introduced the delta rule to forget values associated with the current write key by removing the corresponding value before adding the new value. Alternatively, [47, 56] proposed gating mechanisms similar to those in gated RNNs to facilitate the forgetting of irrelevant information.
Long-term dependencies in RNNs. RNNs fall short in long-term dependency modeling, which is commonly attributed to the gradient vanishing issue. Three methods are typically applied to mitigate this issue. (i) Gating mechanisms [9, 17, 23, 30, 70], which are believed to be crucial to the success of LSTMs, use additive (instead of multiplicative) hidden state update rules to improve gradient flow. (ii) Regularizing or initializing the eigenvalues of the recurrent weight matrix (close) to one via identity matrices [38] or unitary matrices [2]. In the diagonal linear RNN case, the eigenvalues coincide with the element-wise decay rates, and LRU [53] uses randomized linear algebra techniques to initialize eigenvalues to be close to one. [53] also interestingly points out that many modern state-space models use a very small time step value on initialization for discretization, resulting in eigenvalues or decay rates close to one. (iii) Adding skip connections between distant time steps to allow shortcuts for gradient flow [5, 8, 37]. Our approach combines (i) and (ii), which improves gating mechanisms with a regularized dynamic decay rate that approaches one in the upper layer.
3.1 Architecture overview
Our proposed Hierarchically Gated Recurrent Network (HGRN) is depicted in Figure 1. It has multiple stacked layers, each of which consists of a token mixing module HGRU and a channel mixing module GLU (Gated Linear Unit [69]).
3.2 HGRU exploration
We begin with a simple gated linear recurrent layer, which is defined as:
$$
\begin{array}{l} \mathbf {f} _ {t} = \operatorname {S i g m o i d} \left(\mathbf {x} _ {t} \mathbf {W} _ {f} + \mathbf {b} _ {f}\right) \in \mathbb {R} ^ {1 \times d}, \\ \mathbf {i} _ {t} = \text {S i g m o i d} \left(\mathbf {x} _ {t} \mathbf {W} _ {i} + \mathbf {b} _ {i}\right) \in \mathbb {R} ^ {1 \times d}, \\ \end{array}
$$
$$
\mathbf {c} _ {t} = \operatorname {S i L U} \left(\mathbf {x} _ {t} \mathbf {W} _ {t} + \mathbf {b} _ {z}\right) \in \mathbb {R} ^ {1 \times d}, \tag {1}
$$
$$
\begin{array}{l} \mathbf {h} _ {t} = \mathbf {f} _ {t} \odot \mathbf {h} _ {t - 1} + \mathbf {i} _ {t} \odot \mathbf {c} _ {t} \in \mathbb {R} ^ {1 \times d}, \\ \mathbf {h} _ {0} = \mathbf {0} \in \mathbb {R} ^ {1 \times d}, \\ \end{array}
$$
where  $\odot$  denotes the element-wise product. Following the terminology used in the RNN literature, we refer to  $\mathbf{f}_t$  and  $\mathbf{i}_t$  as the forget and input gates, respectively. It is worth noting that  $\mathbf{f}_t$  and  $\mathbf{i}_t$  depend only on  $\mathbf{x}_t$  and not on  $\mathbf{h}_{t - 1}$ .
This characteristic enables the use of the parallel scan algorithm [48, 71], otherwise it is infeasible. We then make the following changes toward our final HGRU step by step.
Algorithm 1 Recurrent Computing
$$
\begin{array}{l} \begin{array}{l} 1 \colon \operatorname {I n p u t}: \mathbf {c} _ {t} \in \mathbb {C} ^ {1 \times d}, \mu_ {t}, \theta , \gamma^ {k} \in \\ \mathbb {R} ^ {\mathrm {f} \times d}, t = 1, \ldots , n, k = 1, \ldots , H. \end{array} \\ 2: \operatorname {I n i t}: \mathbf {h} = \mathbf {0} \in \mathbb {C} ^ {1 \times d}, \mathbf {H} \in \mathbb {C} ^ {n \times d}. \\ 3: \text {f o r} t = 1 \text {t o} n \text {d o} \\ 4: \quad \text {b e g i n} \\ 5: \quad \lambda_ {t} = \gamma^ {k} + (1 - \gamma^ {k}) \odot \mu_ {t}. \\ 6: \quad \mathbf {h} = \lambda_ {t} \exp (i \theta) \mathbf {h} + (1 - \lambda_ {t}) \mathbf {c} _ {t}. \\ 7: \quad [ \mathbf {H} ] _ {t} = \mathbf {h}. \\ 8: \quad \text {e n d} \\ 9: \text {r e t u r n} \mathbf {H}. \\ \end{array}
$$
Figure 1: Illustration of the neural architecture. Each HGRN layer consists of a token mixer HGRU and a channel mixer GLU. HGRU employs linear recurrence in the complex domain:  $\mathbf{h}_t = \lambda_t\odot \exp (i\theta)\odot \mathbf{h}_{t - 1} + (1 - \lambda_t)\odot \mathbf{c}_t$ . Here  $c_{t}$  is the input vector,  $\theta$  is the rotation angle,  $\mu_{t}$  is the output of the original forget gate,  $\gamma^k$  is the lower bound of the  $k$ th layer,  $\lambda$  is the resulting data dependent decay rate:  $\lambda_{t} = \gamma^{k} + (1 - \gamma^{k})\odot \mu_{t}$ .
Complex-valued recurrence. For linear RNNs with static decay rates, it is common to perform eigendecompositions on the recurrent weight matrix to achieve element-wise linear recurrence. However, if only real-valued eigenvalues are allowed, it restricts the range of the recurrent weight matrix to be symmetric, limiting the model's expressiveness. To overcome this limitation, linear RNNs often employ complex-valued eigenvalues to enhance the model's expressive power [20, 26, 27, 32, 53]. Motivated by this, we extend our model to consider  $\mathbf{h}_t, \mathbf{i}_t, \mathbf{c}_t \in \mathbb{C}^{1 \times d}$  as complex values. For the input  $c_t$ , we parameterize its real and imaginary parts separately as follows:
$$
\operatorname {R e} \left(\mathbf {c} _ {t}\right) = \operatorname {S i L U} \left(\mathbf {x} _ {t} \mathbf {W} _ {c r} + \mathbf {b} _ {c r}\right) \in \mathbb {R} ^ {1 \times d},
$$
$$
\operatorname {I m} \left(\mathbf {c} _ {t}\right) = \operatorname {S i L U} \left(\mathbf {x} _ {t} \mathbf {W} _ {c i} + \mathbf {b} _ {c i}\right) \in \mathbb {R} ^ {1 \times d}.
$$
Regarding the forget gate values, we find it convenient to use the exponential representation of complex numbers and parameterize  $\mathbf{f}_t$  as follows:  $\mathbf{f}_t = \lambda_t\odot \exp (i\theta_t)$ . Here,  $i^2 = -1$ ,  $\lambda_t,\theta_t\in \mathbb{R}^d$  and  $\exp (i\theta_{t}) = \cos \theta_{t} + \sin \theta_{t}i$ . The magnitude argument  $\lambda_{t}$  determines the intensity of remembering historical information, while the phase argument  $\theta_{t}$  determines the oscillation frequencies. We find that parameterizing  $\theta_{t}$  in a data-independent manner is preferable, as it allows for a clear interpretation of encoding relative position information (see next subsection for more discussions), which is reminiscent of Rotary Positional Embedding (RoPE) [72]. We shared  $\theta$  across times steps, i.e.,  $\mathbf{f}_t = \lambda_t\odot \exp (i\theta)$ , initialize  $\theta$  as RoPE does, but make it learnable like LRPE [63].
Lower bound on forget gate values. Since the intensity of remembering information is only related to the magnitude argument  $\lambda_{t}$ , we focus on how to add a lower bound to  $\lambda_{t}$ . As mentioned earlier, we want to set a monotonically increasing lower bound on the forget gate (magnitude) values. Inspired by ON-LSTM [70], we employ the cummax activation function to achieve this. Concretely, we allocate  $\Gamma \in \mathbb{R}^{H\times d}$  to parameterize lower bounds independently for all hidden states, where  $H$  is the number of layer. Assuming the layer index is  $k$ , we have the following calculations:
$$
\mathbf {P} = (\operatorname {S o f t m a x} (\boldsymbol {\Gamma}, \dim = 0) \in \mathbb {R} ^ {H \times d},
$$
$$
\gamma^ {k} = [ \mathrm {C u m s u m} (\mathbf {P}, \dim = 0) ] _ {k} \in \mathbb {R} ^ {1 \times d}.
$$
Here we define  $[\mathrm{Cumsum}(\mathbf{x})]_k = (\sum_{i=1}^k x_i) - x_1$  to prevent the highest layer's lower bound from being one as we still want the ability to forget irrelevant information.
We remark that there is a difference in the use of cummax between our model and ON-LSTM. In ON-LSTM, cummax is applied to the hidden state dimension within a single layer, while in our case, we apply cummax on the layer dimension across different layers to enable upper layers to model long-range dependencies.
Finally,  $\lambda_{t}$  in the  $k$ -th layer is parameterized as follows:
$$
\mu_ {t} = \operatorname {S i g m o i d} \left(\mathbf {x} _ {t} \mathbf {W} _ {\mu} + \mathbf {b} _ {\mu}\right) \in \mathbb {R} ^ {1 \times d},
$$
$$
\lambda_ {t} = \gamma^ {k} + (1 - \gamma^ {k}) \odot \mu_ {t} \in \mathbb {R} ^ {1 \times d}.
$$
Comparing to before (i.e., without lower bounds), to achieve the same forget rate value  $\bar{\gamma}$  closed to one,  $\mu_t$  will be pushed away from the Sigmoid activation function's saturated regions (i.e., near one),
$$
\mu_ {t} = \frac {\bar {\gamma} - \gamma^ {k}}{1 - \gamma^ {k}} <   \bar {\gamma},
$$
thereby mitigating the gradient vanishing issue [23] and making gradient-based optimization easier.
Tying input and forget gates. To reduce the number of parameters, it is common to use leaky units, i.e., tying the input gate with the forget gate using  $\mathbf{i}_t = 1 - \mathbf{f}_t$ , which has a close relationship to the discretization of continuous-time system [75] and exponential moving average [33], and has been proven effective empirically [9, 19]. To allows for a clear interpretation of encoding relative position information, we only apply this strategy on the magnitude argument:
$$
\mathbf {h} _ {t} = \lambda_ {t} \odot \exp (i \theta) \odot \mathbf {h} _ {t - 1} + (1 - \lambda_ {t}) \odot \mathbf {c} _ {t} \in \mathbb {C} ^ {1 \times d}. \tag {2}
$$
Output gates and projection. The addition of gates to the output of the recurrence layer has been shown to be effective in state-space models [11, 46, 49, 82]. Motivated by this, we incorporate an output gate before performing the output projection as follows and get HGRU:
$$
\mathbf {g} _ {t} = \operatorname {S i g m o i d} \left(W _ {g} \mathbf {x} _ {t} + b _ {g}\right) \in \mathbb {R} ^ {1 \times 2 d},
$$
$$
\mathbf {o} _ {t} ^ {\prime} = \operatorname {L a y e r N o r m} \left(\mathbf {g} _ {t} \odot [ \operatorname {R e} \left(\mathbf {h} _ {t}\right), \operatorname {I m} \left(\mathbf {h} _ {t}\right) ]\right) \in \mathbb {R} ^ {1 \times 2 d}, \tag {3}
$$
$$
\mathbf {o} _ {t} = \mathbf {o} _ {t} ^ {\prime} \mathbf {W} _ {o} + \mathbf {b} _ {o} \in \mathbb {R} ^ {1 \times d}.
$$
3.3 Token mixing perspective of HGRU
We provide the token mixing perspective of HGRU similar to [32]. Expanding Equation 2, we have:
$$
\mathbf {h} _ {t} = \sum_ {s = 1} ^ {t} (1 - \lambda_ {s}) \left[ \prod_ {k = s + 1} ^ {t} \lambda_ {k} \exp (i \theta) \right] \mathbf {c} _ {s} = \sum_ {s = 1} ^ {t} (1 - \lambda_ {s}) \left[ \prod_ {k = s + 1} ^ {t} \lambda_ {k} \right] \exp (i (t - s) \theta) \mathbf {c} _ {s} \tag {4}
$$
Written in matrix form, we have:
$$
\mathbf {H} = \left[ \begin{array}{c} \mathbf {h} _ {1} \\ \vdots \\ \vdots \\ \mathbf {h} _ {n} \end{array} \right], \mathbf {A} = \left[ \begin{array}{c c c c} 1 - \lambda_ {1} & 0 & \dots & 0 \\ (1 - \lambda_ {1}) \lambda_ {2} \exp (i \theta) & 1 - \lambda_ {2} & & \vdots \\ \vdots & \vdots & \ddots & 0 \\ (1 - \lambda_ {1}) [ \prod_ {s = 2} ^ {n} \lambda_ {k} ] \exp (i (n - 1) \theta) & \dots & \dots & 1 - \lambda_ {n} \end{array} \right], \mathbf {C} = \left[ \begin{array}{c} \mathbf {c} _ {1} \\ \vdots \\ \vdots \\ \mathbf {c} _ {n} \end{array} \right] \tag {5}
$$
So the token mixing module can be formed as follows:
$$
\mathbf {H} = \mathbf {A C}. \tag {6}
$$
Note that the token mixing matrix  $\mathbf{A}$  can be decomposed into two parts  $\mathbf{A} = \mathbf{\Lambda} \odot \mathbf{\Theta}$ :
$$
\boldsymbol {\Lambda} = \left[ \begin{array}{c c c c} 1 - \lambda_ {1} & 0 & \dots & 0 \\ (1 - \lambda_ {1}) \lambda_ {2} & 1 - \lambda_ {2} & & \vdots \\ \vdots & \vdots & \ddots & 0 \\ (1 - \lambda_ {1}) [ \prod_ {s = 2} ^ {n} \lambda_ {k} ] & \dots & \dots & 1 - \lambda_ {n} \end{array} \right], \boldsymbol {\Theta} = \left[ \begin{array}{c c c c} 1 & 0 & \dots & 0 \\ \exp (i \theta) & 1 & & \vdots \\ \vdots & \vdots & \ddots & 0 \\ \exp (i (n - 1) \theta) & \dots & \dots & 1 \end{array} \right] \tag {7}
$$
This decomposition means that the Token mixing matrix  $\Lambda$  can be decoupled into two independent modules, where  $\Lambda$  models the long-distance dependency and  $\Theta$ , a Toeplitz matrix, models the relative positional relationship and enhanced expressiveness. Note that if  $\Theta$  depends on the input, then the matrix  $\Lambda$  will no longer be a Toeplitz matrix, thus unable to capture relative position information. It can be also viewed as a RoPE-enhanced attention mechanism:  $\Lambda$  corresponds to the attention matrix but the attention score here is the cumulative product of data-dependent decay rates;  $\Theta$  directly corresponds to RoPE.
Table 1: Results on Wikitext-103 (TNN[59]'s setting).  $\downarrow$  means lower is better.
Model	PPL
(val)↓	PPL
(test)↓	Params
(M)
Attn-based
Transformer [81]	24.40	24.78	44.65
FLASH [10]	25.92	26.70	42.17
1+elu [35]	27.44	28.05	44.65
Performer [7]	62.50	63.16	44.65
cosFormer [62]	26.53	27.06	44.65
MLP-based
Syn(D) [76]	31.31	32.43	46.75
Syn(R) [76]	33.68	34.78	44.65
gMLP[42]	28.08	29.13	47.83
RNN-based
S4 [22]	38.34	39.66	45.69
DSS [26]	39.39	41.07	45.73
GSS [49]	29.61	30.74	43.84
RWKV [55]	24.31	25.07	46.23
LRU [53]	29.86	31.12	46.24
FFT-based
TNN [59]	23.98	24.67	48.68
Ours
HGRN	24.14	24.82	46.25
We conduct a comparative analysis between our proposed HGRN and four widely adopted sequence modeling structures, i.e., attention-based, MLP-based, FFT-based, and state-space-based. We evaluate HGRN on the WikiText-103 dataset [50] and the Pile [15] dataset for autoregressive language modeling, as well as the length extrapolation ability. To assess the accuracy and efficiency of our model in handling long-term dependencies, we utilize the LRA benchmark [78]. Additionally, we showcase the robustness of HGRN in computer vision task on the ImageNet-1k dataset.
We implement our models in Pytorch [54] and train them on 8 Nvidia A100 GPUs. For HGRN, we found that fusing element-wise recurrence into a single CUDA kernel results in fast running speed in practice. [48] also found that unless the sequence length is sufficiently large, the parallel scan's implementation is not necessarily faster than the sequential scan. As such, we use a CUDA-based sequential scan for implementation; however, our model has the potential to model very long sequences through the use of a parallel scan.