HGRN2: Gated Linear RNNs with State Expansion
Authors: Zhen Qin
Songlin Yang
Weixuan Sun
Xuyang Shen
Dong Li
Weigao Sun
Yiran Zhong
Venue: COLM (2024)
Keywords: linear attention, long convolution, linear RNN, Efficient LM, Linear complexity
TL;DR: HGRN2 improves upon HGRN by introducing an outer product-based gating mechanism, enhancing memory capacity with better performance.

Abstract: Hierarchically gated linear RNN (HGRN) has demonstrated competitive training speed and performance in language modeling while offering efficient inference. However, the recurrent state size of HGRN remains relatively small, limiting its expressiveness. To address this issue, we introduce a simple outer product-based state expansion mechanism, which significantly enlarges the recurrent state size without introducing any additional parameters. This enhancement also provides a linear attention interpretation for HGRN2, enabling hardware-efficient training. Our extensive experiments verify the advantage of HGRN2 over HGRN consistently across different settings and comptetive to other recurrent models.

Published as a conference paper at COLM 2024
HGRN2: Gated Linear RNNs with State Expansion
1 'Zhnn Qin', 2Songlin Yang", Weixuan Sun, "Xuyang Shen, Dongg Li, Weigao Sun, 3 Yiran Zhong"
Tap MIT CSAIL POpenNI.PLab, Shanghai AlLab
attps///gtthu.coo/ppallllP//a//
Abstract
Hierarchically gated linear RNN (HGRN, Qin et al. 2023c) has demonstrated competitive training speed and performance in language modeling while offering efficient inference. However, the recurrent state size of HGRN remains relatively small, limiting its expressiveness. To address this issue, we introduce a simple outer product-based state expansion mechanism, which significantly enlarges the recurrent state size without introducing any additional parameters. This enhancement also provides a linear attention interpretation for HGRN2, enabling hardware-efficient training. Our exten- sive experiments verify the advantage of HGRN2 over HGRN consistently across different settings and comptetive to other recurrent models.
1 Introduction
Large language models (LLMs) have achieved significant empirical success in recent years. However, serving Transformer-based LLMs is costly due to the expensive KV cache man- agement. Recurrent neural networks (RNNs), on the other hand, offer linear inference complexity with constant state size, making them ideal for serving. Consequently, there is substantial interest in studying parallelizable linear recurrent models, such as linear RNNs (Peng et 2023 Orvieto al. 2023 Qin al. 2023c De al. 2024), linear attention (Sun et 2023 Oin et al. 2023b] Yang et al. 2023 2024 Arora et al.) 2024), and state space
models (Gu et al. 2022a Smith et al. 2023 Gu & Dao 2023
Dao Gu 2024).
RNNs have a fixed recurrent state size to encode all historical information. Therefore, it is important for RNNs to (i) utilize the fixed-sized states effectively and (ii) increase the recurrent state size to enhance memory capacity. Recent improvements in linear RNNs follow this approach, incorporating techniques such as data-dependent decays and state expansion.
Data-dependent decays (also known as forget gates) are crucial for RNNs (van der West huizen & Lasenby 2018), allowing them to selectively retain useful information while erasing irrelevant information. This enables the fixed-size recurrent state to store only important information more efficiently. HGRN (Qin et al. 2023c) first emphasized the importance of data-dependent decays for linear RNNs. Many recent linear recurrent models, such Mamba (Gu & Dao 2023), Gated Linear Attention (GLA, Yang et al. 2023), Griffin (De et al. 2024), and RWKV (Peng et al. 2024), also employ data-dependent decays.
However, HGRN did not increase the recurrent state size, which is greatly restricted by limited memory capacity. This limitation prevents it from achieving LLaMa-like (Touvron) et al. 2023a b) language modeling performance, as noted in Qin et al. (2024). Recent state- of-the-art linear recurrent models, such as Mamba, GLA, and RWKV-6, have addressed this issue by employing state-expansion techniques. These techniques significantly increase the recurrent state size and thereby enhance memory capacity, which has been shown to be crucial for language modeling performance and directly correlated with retrieval ability
(Arora et al. 2024).
Corresponding author. Email: zhongyirandgmail.com Equal contributions.
1

---

Published as a conference paper at COLM 2024
In this work, we propose HGRN2, which aims to increase the recurrent state size for HGRN while retaining both parameter and training efficiency. We first explore structured matrices to expand the state size directly in parameter-effictee that this approach improves language modeling performance but still encounters training inefficiencies, which limit the scaling of the recurrent state size. Inspired by linear attention, we then explore using a non-parametric outer product-based state expansion mechanism. This approach allows for efficient scaling of the recurrent state size during training without introducing additional parameters. Due to the matrix multiply form of linear attention, we can leverage the hardware-efficient linear attention training algorithm described in Yang et al. (023); Qin et al. 2024) for large-scale experiments. As result, HGRN2 can be regarded as an improved parameterization of GLA.
manner. Empirically, we found
We extensively evaluate HGRN2 across various tasks, demonstrating that it consistently outperforms HGRN1 in multiple domains. In language modeling, we show HGRN2 to be highly competitive compared to other subquadratic efficient models.
2
Background
2.1 Gated linear RNN
Given input x Nxd, where the sequence length is N and the model dimension is d, minimalist gated linear recurrent layer (Martin & Cundy, 2018) transforms the input x into hidden states h E Nxd and the output y E Rxdx defined below:
= a (Ux; + bb,,,
i, = т (Vx: + b,),
0: (Wx: + bw),
(1)
h g: h + (1- ii
yt h @ Ot,
where @ denotes element-wise product; 0 is the sigmoid function, and T is nonlinear activation function (we choose to use SiLU)) is the input vector; gr and o; are the forget gate and output gate, respectively. The input gate is tied to the forget gate as 1 - gi, common approach used in many gated RNNs such as GRU (Chung al.) 2014).
2.2 HGRN QQn et al. 2023c)
Compared to Eq. 1 HGRN makes two adjustments: (i) complex-valued recurrence and (ii) forget gates with monotonically increased lower bound values from bottom layers to upper layers.
For (i), similar to the findings in Gu & Dao 2023 and De et al. 2024 empirically found that complex-valued recurrence is not necessary, as shown in lable I The reason why HGRN found it useful is due to state expansion: the complex-valued recurrent state is twice the size of that in the real-valued recurrent state. If we directly expand the real-valued recurrent state size from d to 2d, the language modeling performance on the Wikitext-103 corpus is even better. Therefore, we only consider the real-valued recurrence thereafter.
Table 1: Comparison of real HGRN and complex HGRN. We found that real HGRN with twice the state size performs better than complex HGRN in Wikil03 language modeling.
Method
State size
PPL(val)
PPL(test)
Params (M)
Complex HGRN1
2d
24.14
24.82
46.25
Real HGRNI
d
25.34
26.12
46.24
Real HGRN1
2d
24.04
24.64
45.46
For (ii), suppose the total number of layers is L. HGRN introduces data-independent
learnable matrix I E Rixd where I; represents the lowest values of the forget gate for the i-th layer at all time steps. HGRN argues that this lower bound should be monotonically increasing from bottom to top, encouraging the bottom layers to model short-term local
2

---

Published as a conference paper at COLM 2024
dependencies and the upper layers to model long-term dependencies. To enforce this monotonicity, HGRN uses the cumulative softmax operator cumax (Shen et al. 2018
= cumax (T) =
dim = 0), dim = 0) E RLdd 8 = [ E IR
To prevent the lower bound from reaching one in the highest layer, HGRN subtracts all values by 88, making the lower bound for the first layer zero. After obtaining the lower bound values, the forget gate gr learns the residuals instead, resulting in the new forget gate f,:
= B² + (1- B)) o g h = f, の h (1 @
where the superscript indicates the layer index. This additive lower bound approach has been shown to mitigate the issue of saturated gates (Gu et al. 2020).
3 Method
3.1
Explorations of state expansion methods
The goal of this work is to scale the size of the HGRN recurrent state from to nd, where the state expansion ratio. However, if we use the original parameterization in Eq. [ the matrices U, V, W will have dimensions d x nd, which becomes very parameter inefficient when n is large. Ideally, the number of parameters should be around d² as in the original case for each projection. To achieve this, we first consider using structured matrices (e.g., low-rank matrices) to replace the dense projection matrix Rd -> RRnd, as described in Table
2
Table 2: Parameter Efficient State Expansion (PESE) methods using Einstein Summation notation. Blue represents the input, Black represents data-independent weights, and Red represents the output. We list the Einstein Summation for low-rank (LR), group linear transformation (GLT), group linear transformation with interaction (GLTI), Khatri-Rao product (KRP), and Kronecker product (KP).
Method
Equation
Parameter #
Naive
d, nd nd d, r rnd + nd
nd2
LR
dr(n + 1) Rs d2
d= (ne) + ne e + n
GLT
d2
GLTT
d= (n e) - n
d? + n?
11 ne d, -> nd
KRP
d,n nd
nd
KP
d, d, n + nd
d2 +
After obtaining the expanded i, we feed them into element-wise gated linear recurrent layers in Eq. 1 and Eq. [ resulting in the output vector y: RRa To project the expanded dimension back to the original dimension, we simply sum over the dimension corresponding иτ
The results are shown in Table 3 performance, with the low-rank matrix performing the best among these candidates.
We found that state expansion generally improves
3

---

Published as a conference paper at COLM 2024
However, these methods face training inefficiency issues, as they require conducting element-wise lin- ear recurrence in high dimensions (i.e., nd). Since these element-wise operations cannot leverage tensor cores (a fast matrix multiplication unit GPUs), the dramatically increasing FLOPs and 1/O costs signifi- cantly slow down training when n is large. We notice that this is similar to the case Mamba which re- quires relatively small expansion ratio (i.e., = 16) and custom I/O-efficient CUDA implementation to achieve reasonably fast running speed.
Table 3: PESE Ablation. Abla- tion studies on various parameter- efficient methods, as described in Ta- ble Each model was trained 10 billion tokens from the Pile dataset.
Method
Params (M)
Xímr
5.16
380
Ximr++
4.62
386
HGRNI
5.10
379
385
4.77
386
In the next subsection, we explore an alternative strat- egy that does not replace the dense projection ma- trices with structured ones but instead changes the element-wise gating operations in Eq to other ma- trix/vector operations similar to those used in linear attention. This approach allows for more efficient training.
GLT
5.06
386
GLTII
4.83
386
KRP
3.08
386
KP
5.06
386
HGRN2
4.79
385
4.73
385
128
4.62
385
3.2 HGRN2
Output
Output
HGRU2
Linear
LayerNorm
GLU
Recurrent Computing
B
HGRU2
9t
1
SiLU
Sigmoid
Linear
Linear
Linear
Input
Input
Figure 1: Network Structure of HGRN2. Each HGRN2 layer includes token mixer layer, HGRU2, and a channel mixer, GLU. HGRU2 employs recurrent computation as described in Eq. 3 where i, is the input state, g: is the forget gate, o, is the output gate, and B¹ is the lower bound for layer i.
The modification from HGRN1 to HGRN2 is simple yet effective. For the input gate, HGRN2 replaces the element-wise product with the outer product for state expansion. Consequently, h, E Rdddd and HGRN2 first diagonalizes the forget gate vector and uses the matrix dot product to update the hidden state. For the output gate, HGRN2 replaces the element-wise product with matrix-vector multiplication to project the expanded state back to the original dimension. The recurrent equation of HGRN2 is as follows:
h; h Diag(f)) +i, @ (1 (1-f,) f,) E Radd
(3)
y; = h, o, E Rd
where Diag denotes the diagonalization of vectors, represents the matrix dot product, and indicates the outer product.
TTouugh Mamba has an attention mechanism (Ali et al 2024) similar to that in linear attention, the
attention computation cannot be expressed as a matrix multiplication like linear attention, and thus does not facilitate tensor core-based acceleration, well acknowledged in Mamba2 (Dao & Gu
2024
4

---

Published as a conference paper at COLM 2024
4.90
24.4
4.85
24.2
24.0
23.8
4.75
4.70
4.65
23.2
1
16
32
64
128
4
8
16
64
128
Expand Ratio
Expand Ratio
Figure 2: Expand Ratio (Head Dimension) Ablation. We tested the relationship between PPL (Perplexity) and the expand ratio on the Wikitext-103 (Merity et al. 2017) dataset (left) and subset of the Pile (Gao et al. 2020) dataset (right).
Multihead Variant. The complexity of recurrence increases dramatically from O(Ndd) to O O(BN22) due to state expansion. To address this, we introduce multihead variant of HGRN (similar to that in linear attenttonn such that the complexity is reduced to O(BNd for the number of heads H, effectively making the state size d² H, i.e,, the expansion ratio d, d/ H i We conducted an ablation study the expansion ratio (or head dimension) 7,, as shown in Figure 2 The results show that state expansion significantly improves language modeling performance. However, when the head dimension (i.e., state expansion ratio) exceeds 128, the performance gain diminishes. To balance computational cost and performance, we chose d; 128 for the main experiments.
Comparison to GLA. is important to note that the recurrence form in HGRN2 is identical to that of GLA (Yang et al. (023), except for the specific parameterization. We list the correspondences between the two parameterizations in Table As shown, the output gate in HGRN2 corresponds to the query in GLA, while the output gate in GLA is omitted in HGRN2. The key vector in GLA corresponds to the input gate in HGRN2, which is tied to the forget gate, thereby saving parameters.
Table 4: The correspondence between HGRN2 and GLA is as follows.
HGRN2
GLA
(output gate)
q (query vector)
1 f (input gate)
(key vector)
(input vector)
v (value vector)
(forget gate)
a (forget gate)
(output gate)
Hardware-Efficient Training. Due to its computational structure's similarity to GLA, can directly leverage their chunkwise algorithm and highly optimized kernels for hardware- efficient large-scale training. For more details, refer readers their paper.
Concluding Remarks. Although HGRN2 shares many similarities with GLA, believe that HGRN2 offers a unique perspective distinct from linear attention, originating from the approach of gated linear RNNs. For instance, may not be immediately clear from the perspective of linear attention why key vectors should be constrained within the range of (0, become quite intuitive when starting from the gated linear RNN framework and exploring state expansion.
1) or why the key vector and forget gate value should sum one. However, these concepts
2See Bolya et al. (002 for detailed complexity analysis.

---

Published as a conference paper at COLM 2024
4 Experiments
4.1 MQAR
Setting. Multi-Query Associative Recall (MQAR) (Arora et al. 2023) is enhanced version of the synthetic induction head dataset (Fu al. 2023), designed to test the in-context associative recall ability of subquadratic models. Arora al. (2023) found strong correlations between MQAR accuracy and language modeling performance. Our experimental setting strictly follows the original paper Our hyperparameter sweep included the following ranges: expansion ratio E (64, 128] and learning rate E {le 5, 5, le 4, 5e 4, le 3,5e - -, 3,1e-22 2}.
As shown in Fig- 3 HGRN2 significantly outperforms HGRN1 across various model dimensions, demonstrating the benefits of state expansion in improving memory capacity and, consequently, in-context recall ability.
Result.
Sequence Length: 64
Sequence Length: 128
Sequence Length: 256
Sequence Length: 512
Accuracy
0.75
0.50
HGRN1 HGRN2
0.00
64
128
256 dimension
512
64
128
256
512
64
128
256
512
64
128
256
512
dimension
dimension
dimension
Figure 3: Results on MQAR, where the x-axis represents the model dimension and the y-axis represents accuracy. The task becomes more challenging as the sequence length increases. HGRN2 outperforms HGRN1 in all scenarios.
Language modeling
Table Results on Wikitext-103.
4.2.1 Wikitext-103
Model
PPL (val)
PPL (test)
Params (M)
Transformer
24.40
24.78
44.65
Setting. For the Wikitext-103 experiment, we followed the configuration of HGRN1 to vali- date the performance of 44M models against wide range of subquadratic models: FLASH (Hua et al. 2022), 1+elu (Katharopoulos et al. 2020), Performer (Choromanski et al., 2021), COS-
FLASH
25.92
26.70
42.17
1+elu
27.44
28.05
44.65
Performer
62.50
63.16
44.65
cosFormer
26.53
44.65
Syn(D)
31.31
32.43
46.75
Syn(R)
gMLP
28.08
29.13
47,83
54
38.34
39.66
45.69
Former (Qin et 2022bb,, (D), Syn(R) (Ta al. (021a), gMLP (Liu et al. 2021 S4 (Gu at
DSS
39.39
41.07
45.73
CSS
29.61
30.74
43.84
RWKV4
46.23
2022a), DSS Gupta & Berant 2022), RWKV-v4
LRU
29.86
31.12
46.24
al., 2023). (Peng et al. 2023), , LRU (Orvieoo
23.98
Mamba
22.58
23.19
44.39
HGRNI (Qin et al. 2023c), TNN (Qin et al. 2023a), and Mamba (Gu & Dao 2023). All re- ported results are from our own runs under the
HGRNI
24.14
24.82
46.25
HGRN2
23.10
23.73
44.66
same settings.
Result. Table shows the results. HGRN2 clearly outperforms HGRNI but slightly underperforms Mamba.
4.2.2 Slimpajama
We conducted language modeling experiments with 1.3B and 2.7B parameters the Slimpajama dataset (Soboleva et al. 2023), using the FLASHLINFAAATTENTTU (Yang &
Zhang 2024 codebase for training. The results, shown in Table demonstrate that
hhtpp: //gttp://github... wtt.e/cc.o/ecacaeeee htt.r///....eeeeeeeaaarrrr/oooo.
Model checkpoints available at atttpt httpee/eteee///gggaaa....eeee//e

---

Published as a conference paper at COLM 2024
HGRN2 consistently outperforms other competitive linear recurrent models across three model scales. This suggests that HGRN2 provides a superior parameterization compared to GLA, both models share identical recurrent structure.
Table 6: Slimpajama language modeling results.
Lamb.
Wiki.
ARC₂ ppl
ARC. acc
Hella. aCCn
Lamb. aCCn
PIQA acc
Wino. acc
Avg
ppl
acc
1.3B parameters with 100B
training
tokens
Transformer++
15.3
17.1
54.1
27.1
49.3
47.0
70.3
54.9
50.5
Mamba
16.5
18.2
57.3
26.6
48.1
43.4
69.5
53.7
49.8
RetNet
15.4
57.4
27.9
50.3
44.6
71.7
50.6
GLA
15.4
17.6
55.4
27.7
49.0
46.4
69.9
54.0
50.4
HGRN2
11.8
16.9
58.1
28.1
51.8
49.4
71.4
52.3
51.9
2.7B parameters with 100B training tokens
Transformert+
10.7
15.2
59.8
27.5
54.2
52.3
72.7
56.2
53.8
Mamba
13.6
15.9
60.7
29.8
53.9
46.4
72.8
53.9
52.9
RetNet
11.9
15.8
59.6
28.1
54.0
49.6
72.3
53.8
52.9
GLA
12.4
15.5
59.2
29.9
54.0
50.4
71.7
55.7
53.5
HGRN2
8.8
14.6
60.8
30.3
58.7
55.4
73.0
54.2
55.4
4.2.3 The Pile
We also conducted experiments on the Pile dataset. First, we trained 150M, 350M, and 1B HGRN1 and HGRN2 models for 100B tokens, and the results are shown in Table 7 We observe that HGRN2 consistently outperforms HGRN1.
Table 7: Comparison between HGRN1 and HGRN2 on Commonsense Reasoning Tasks.
Model
Bn Params
Bn Tokens
PIQA
Hella.
Wino.
ARC-e
ARC-c
OBQA
AVG
HGRN1
0,15
100
65.02
33.33
50.20
46.68
23.81
28.60
41.27
HGRN2
0.15
100
66.43
35.44
51.70
46.63
24.32
28.40
42.15
HGRN1
0.35
100
66.70
38.12
51.70
49.20
25.26
30.60
43.60
HGRN2
0.39
100
69.97
46.16
52.72
53.58
23.98
32.40
46.47
HGRN1
100
70.89
48.02
51.62
55.64
27.90
31.60
47.61
HGRN2
1
100
74.16
54.85
56.12
58.71
27.22
34.00
50.84
Next, we scaled the token horizon to 300B and trained strong baseline models, Mamba and LLaMA, under the same settings for comparison. We also compared them against several open-sourced language models, such as OPT (Zhang al. 2022), Pythia (Biderman al.
2023), BLOOM (Scao et al. 2022), and RWKV-4 (Peng et al. 2023). We found that HGRN2 performs competitively with Mamba, LLaMA, and other open-sourced LLMs.
To evaluate long-context abilities, we conducted tests on SCROLLS (Shaham et al. 2022 and found that HGRN2 exhibits better scaling behavior compared to Mamba, indicating stronger long-context capabilities, potentially due to its larger recurrent state size. However, we also observed that the 7B HGRN2 model is still not as strong as the LLaMA model, suggesting that the scaling behavior of linear models for long-context modeling remains an area for further study.
To test the retrieval ability of our trained 3B models, we ran the easy mode of the Needle in Haystack Test. LLaMA almost achieves perfect retrieval performance for evaluation
'nn this mode Shen 2024 Shen al, 2024), both the question and answer (QA pair) are embedded within a lengthy text, challenging the model to locate and respond to the query. This mode is particularly suitable for base models without instruction tuning, In contrast, the standard mode only places the answer within the long context, requiring the model to understand the question and find the relevant answer.
7

---

Published as a conference paper at COLM 2024
Table 8: Comparison between HGRN2 and other open-sourced language models, alongside strong baseline models (LLaMA and Mamba re-trained under the same settings), on Com- monsense Reasoning Tasks. indicates our own trained model.
Token
PIQA
OBQA
OPT
0.35
300
64.58
36.69
52.49
44.02
23,89
28.20
41.65
Pythia
0.40
300
67.08
40.52
53.59
51.81
24.15
29.40
44.43
BLOOM
0.56
350
64.09
46.97
52.80
47.35
29.38
28.20
42.23
RWKV-4
Llama!
0.4
350
67.19
38.75
52.19
49.24
23.72
30.00
43.51
Mamba'
0.4
300
67.90
40.74
52.72
33.07
24.74
31.20
45.06
HGRN2
51.78
GPT-Nee
1.3
300
71.11
48.93
54,93
56.19
33.60
48.44
OPT
1.3
300
71.71
53.70
59.35
37.24
29.69
50.82
Pythia
1.4
300
70.67
47.18
53.51
36.99
26.88
31.40
47.77
BLOOM
RWKV-4
1.5
72.36
52.48
54.62
60.48
29.44
34.00
50.56
Llama'
1.0
300
69.97
47.04
52.72
57.07
26.18
32.60
47.93
Mamba'
56.35
HGRN2
1.0
300
7166
49.52
54.38
60,27
28.07
33.40
49.55
OPT
2.7
300
73.83
61.01
35.20
53.79
Pythia
2,8
300
74.10
59,31
59.91
64.14
33.02
35.60
54.35
BLOOM
70.57
RWKV-4
3.0
72.42
58.75
57,30
62.92
35.15
36.20
53.79
Llama'
3.0
350
73.18
57 88
59,59 99.59
63.93
33.51
35.40
53.93
Mamha'
HGRN2
3.0
300
74.10
61.48
58.64
65.61
34.47
35.60
54.98
Ilama*
HGRN2
7.0
300
76.50
66.96
61.40
69.02
36.86
38.00
58.12
Table 9: Performance Comparison on SCROLLS. R-1/2/L stand for parameter size, tokens, and rouge-l/rouge-/rougee respectively.
Model
Params
Token
GovRep R-1/2/L
SumSer R-1/2/L
QMSum R-1/2/L
Qspr
Nity F1
QALT EM
CNLI EM
Avg
t
Bn
Llama
0.4
300
82/3.5/6.2
11.3/16/87
107/21/99
17.8
15.4
28.0
13.9
10.5
Mamba
300
82/24/6.2
11.2/18/89
9.3/16/844
25.8
HGRN2
0.4
300
15.3/355/109
7.4/0.86.2
83/1.2/7.4
12.4
10.9
26.4
31.5
10.9
Llama
300
12.9/3.1/9.4
9.5/08/7.7
10.9/22/994
22,8
16.0
28.4
9,9
11.0
Mamba
300
15.2/42/106
12.3/16/94
139/3.1/11.7
18.3
14.7
26.7
9.1
11.6
HGRN2
300
149/422/10.
11.4/1/992
109/23//9.
16.2
15.1
27.8
10.6
11.1
Llama
11.2/49/8.1
11.9/19993
16.1/4.3/129
Mamba
3.0
300
21.3/666/139
13.2/2.0/10.1
15.0/3.2/123
22.1
17.9
28.8
24.0
14.7
HGRN2
300
21.7/6.6/14.1
14.6/2.1/0.8
12.5/2.710.6
25.4
18.8
28.9
31.9
15.4
Llama
7.0
300
17.4/73/114
129/1&/10.0
14./337/1188
32.4
22.3
33.8
10.0
14.6
HGRN2
14.9/52/10.2
15./24/111
14./3.0/118
30.1
lengths no greater than the training length. As shown in Figure HGRN2 and Mamba still face difficulties in retrieval tasks; however, HGRN2 outperforms Mamba due its larger state size, enabled by linear attention-styled state expansion.
Figure 4: Easy mode Needle in Haystack Test on 3B models: Mamba (left) and HGRN2 (right). The evaluation context length 16K, and the models were trained on sequence length of 8K.
8

---

Published as a conference paper at COLM 2024
4.3 Long Range Arena
Table 10: Results on LRA.
indicates the results reported by Alonso et al. 2024
Model
ListOps
Text
Retrieval
Image
Pathfinder
Path-X
AVG
Transformer
38.37
61.95
80.69
40.57
65.26
47.81
cosFormer
36.50
67.70
83.15
51.23
71.96
51.76
FLASH
38.70
64.10
86.10
47.40
70.25
51.09
S4
59.60
86.82
90.90
88.65
94.20
96.35
86.09
TNN
61.04
87.90
90,97
88.24
93.00
96.10
86.21
S5
62.15
89.31
91.40
88.00
95.33
98.56
87.45
Mega
63.14
90.43
91.25
90.44
96.01
97.98
88.21
SGConv
61.45
89.20
91.11
87.97
95.46
97.83
87.17
LRU
60.20
89.40
89.90
89.00
95.10
94.20
86.30
Mamba
38.02
82.98
72.14
69.82
69.26
67.32
66.59
Griffin
32.34
71.75
66,58
61.15
73.38
69.53
62.45
HGRNI
59.95
88.14
94.23
88.69
92.92
97.50
86.91
HGRN2
60.52
88.97
95.07
89.33
93.95
98.12
87.66
Setting. Long Range Arena [Tay et al. 2021b) is benchmark designed to assess model's ability handle long-range dependencies. We used HGRN1's configuration and compared it with existing methods, as shown below.
Result. Table 10 shows the results. HGRN2 outperforms HGRN1, while Mamba and Griffin failed to achieve high accuracy on this benchmark.
4.4 Image Modeling
Setting. For the image classification task, followed the configuration of HGRN1 and trained on ImageNet-1k, comparing it with TNN and the vanilla transformer.
Result. Table shows the results. HGRN2 outperforms HGRNI with a similar parameter size, while also demonstrating an advantage over previous TNN (Qin et al. 2023a) and DeiT models [Touvron et al. 2021).
Table 11:
Performances comparison of image classification on ImageNet-1k, HGRN2 per- forms favorably compared to competing methods with similar parameter sizes.
DeiT-Tiny
DeiT-Small
Model
Top-1 Acc
Params (M)
Top-1 Acc
Params (M)
DeiT
72.20
5.7
79.90
22.0
TNN
72.29
6.4
79.20
23.4
HGRN1
74.40
6.1
80.09
23.7
HGRN2
75.39
6.1
80.12
23.8
Related work
Linear recurrent models. Linear recurrent models mainly include linear RNNs, state-space models, and linear attention. State-space models (SSMs) are gaining great attention since the seminal work S4 (Gu et al. (002aa) and its more efficient diagonalized version (Gu et al.. 2022b). Despite excellent performance in the LRA benchmark, it has been shown to have interior performance in language modeling. Gating mechanisms have been shown to be crucial in improving SSMs' language modeling performance (Mehta et al. 2023 Wang et al.

---

Published as a conference paper at COLM 2024
2022:
Gu & Dao 2023). Gupta et al. (2022) build the connection between SSM and linear Orvieto et al. (2023) proposes a linear RNN layer (i.e., LRU) inspired by SSMs. Peng
RNN.
et al. 2023) successtully scale linear RNN models to billions of parameters for the first time.
For linear attention models, their language modeling performance has been underperform- ing softmax attention for long time. Several improvements have been proposed to bridge
the performance gap: (i) incorporating the forgetting mechanism (Peng et al. 2021
Schlag
et al.) 2021 Sun et al., 2023: Qin et al. 2023b; Yang et al., 2023 Peng et al. 2024), (11 using
local attention (Qin et al. 2022a Zhang et al. [023] Arora et al. 2024; Ren et al., 2024),
(iii) using higher-order polynomial feature map (Arora et al. 2024 Kacham et al. 2023) to
make the resulting attention distribution more sharp (Zhang et al. 2024), (IV) using more
expressive yet efficient recurrent update rule (Schlag a022 Yang et al. 2024 Liu al.
2024
Sun et al. 2024a).
Gated linear recurrence. Martin & Cundy (2018) first proposed a minimal gated linear recurrent layer and showed how to the parallel scan algorithm to train linear RNNs in sequence-level parallel. Qin et al. (2023c) is largely based on this work with several adaptations and highlights the importance of data-dependent decay. De et al. (2024) build their model on LRU (Orvieto et al. 2023) and replace data-independent decays with data- dependent ones. They further use slidin-wiidow attention boost the performance. These models are limited in recurrent state size.
Gated recurrent models with matrix-valued recurrent state have been investigated in the literature of Neural Turing Machine (NTM Graves et al. 2014) and linear Transformer (Katharopoulos et al.) 2020). In NTM, the number of memory slots can be regarded as the state expansion ratio discussed in this work. NTM also included data-dependent decays in the form of erase vectors. However, NTM is hard to parallelize and thus slow to train in practice. The linear transformer is known to have the recurrent form (Katharopoulos et al., 2020) and is known to be closely related fast weight programming (FWP Schlag et al.
2021). Gated FWPs have been investigated Schcla & Schmidhuuber ; Zhang &
Zhou (017)), and have recently been revisited in Peng et al. (2021)) Mao
2022); Yang et al.
(2023)
Katsch (2023); Pramanik et al. (2023). In particular, Yang et al. (2023) proposed
hardware-efficient training algorithm for these types of models.
More recently, Mamba2 (Dao & Gu 2024), XLSTM (Beck al., 2024), and Gated Retention (Sun et al.) 2024b) have shown that sharing data-dependent decays across different dimen- stons within the same head is effective. This approach improves efficiency over GLA because intra-chunk computations are more amenable to tensor core-based matrix multiplication acceleration, at the cost of sacrificing the fine-grainedness of decays. In GLA GLA/HGRN2, each head dimension has its own decay rate, whereas in Mamba2 / xI SLSTM/CO GGGted Retention, all dimensions share the decay under single head. an interesting question to study how much improvement fine-gra tinengraieed decay will bring.
Conclusion
In this work, we propose HGRN2, an enhancement HGRN QQnn et al. 2023c using an outer product-based state expansion mechanism inspired by linear attention, enabling efficient training. Experiments across multiple tasks validate the advantages of HGRN2 over HGRN1.
Acknowledgement
We thank Yu Zhang for conducting some language modeling experiments and for the valuable discussions.
10

---

Published as a conference paper at COLM 2024
References
Ameen Ali, Itamar Zimerman, and Lior Wolf. The hidden attention of mamba models. 2024.
URL https: swwanticschoolar http://ccccoaaaoo/ooorr..saa..ee U
(222rrp81::2282482222 268248520
Carmen Amo Alonso, Jerome Sieber, and Melanie Nicole Zeilinger. State space mod-
els as foundation models: A control theoretic overview. 2024. URL https://api
semanticscholar. wttp/oogg/////ooo Summafiiccclollrr.oog//rrruuuuuu : 268681121
Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou,
Atri Rudra, and Christopher Ré. Zoology: Measuring and improving recall in efficient language models. ar Xii::2312.04927, 2023.
Simran Arora, Sabri Byuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. Simple linear attention language models balance the recall-hroughhuut tradeoff. CoRR, abs 1862..86688 2024. doi: 10.48550/ARXI //doi. org//0../1//.....oo 48550/arXiv.... 2402. 18668
2402.18668.
Maximilian Beck, Korbinian Poppel, Markus Spanring, Andreas Auer, Oleksandra Prud- nikova, Michael K Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochre- iter. xlstm: Extended long short-term memory, ArXiv, abs/ abs/2405.04517, 2024. URL /tttp:://ap..smmmatiicooo. www.seeesccicccooolee..
Stella Biderman, Hailey Schoelkopf, Quentin G. Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar van der Wal. Pythia: A suite for analyzing large language models across training and scaling. ArXiv, abs /bbs/230400133,, 2023.
URL https://app www.etiiscooosneeeeeeeeeeeeeonee org/ hetg/orrpusID-- (02rpusD::55991880 257921893
Daniel Bolya, Cheng- Yang Fu, Xiaoliang Dai, Peizhao Zhang, and Judy Hoffman. Hydra attention: Efficient attention with many heads. In ECCV Workshops, 2022. URL https
s/api..emanticcchol... org Trrg/oorpusID ID: : 252284084
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, An- dreea Gane, Tamás Sarlós, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J. Colwell, and Adrian Weller. Rethinking atten- tion with performers. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. Reveeevi, 2021.
Junyoung Chung, Çaglar Gülçehre, KyungHyun Cho, and Yoshua Bengio. Empirical evalu- ation of gated recurrent neural ortw on sequence modeling. CoRR, abs/ 1412..355, 2014. http::tttt//i.cccto
1552
Tri Dao and Albert Gu. Transformers are ssms:
Generalized models and efficient algorithms 2024. URL https:://aapt
through structured state space duality. ArXiv,
semanticscholar. htt.:/scchh.aaaaooorrorrroooeeoe
Trrg/CorpusID: 27019977777
Soham De, Samuel L. Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, Guillaume Desjardins, Arnaud Doucet, David Budden, Yee Whye Teh, Razvan Pascanu, Nando de Freitas, and Caglar Gulcehre. Griffin: Mixing gated linear recurrences with local attention for efficient language models. ArXiv, abs/2402.119427 2024. URLhtttss:/// https://api
semanticscha
wwt. Trgg/CorpusID.. 268091246
Daniel Y. Fu, Tri Dao, Khaled Kamal Saab, Armin W. Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.ned. 2023.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800gb dataset of diverse text for language modeling, ar preprint Xii: .0007,, 2020.
11

---

Published as a conference paper at COLM 2024
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. ArXiv, abs/1410.5401, 2014. URL httpss/////p semanticscholar
www.llarr.rrgCopppcsiiii http:////ppp.emmmmmccccccccarooa 15299054
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. 2023.
Albert Gu, Çaglar Gülçehre, Thomas Paine, Matt Hoffman, and Razvan Pascanu. Improving the gating mechanism of recurrent neural networks. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 3800-3809. PMLR, 2020. URL http 7/proceedings. mlr. pptsss/v119/gu200aaa html
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In The Tenth International Conference Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022a.
Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré. On the parameterization and initialization of diagonal state space models. ArXiv, (bss/2206.1188 2022b. URL https //api. s/api.ssemanticscho.... orw.org//oooo..ooooooooooo.o werg/CorpusID-... :2999538775
Ankit Gupta and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces. ArXio, 2203/2233,,333,, 2022. URL https://app..seeee www.eesscsaanncccchcooooeoeo http:s// org/
CorpusID:2247621199
Ankit Gupta, Harsh Mehta, and Jonathan Berant. Simplifying and understanding state space models with diagonal linear mnss ArXiv, abs/ 20122.22,,,,,8, 2022. URL https
//api.semanticsccoolaar org/CorrpssD:225112222
Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc Le. Transformer quality in linear time. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato (eds.), International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pp. 9099-9117. PMLR, 2022.
Praneeth Kacham, Vahab Mirrokni, and Peilin Zhong. Polysketchformer: Fast transformers via sketching polynomial kernels, 2023.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pp. 5156-5165. PMLR, 2020.
Tobias Katsch. Gateloop: Fully data-controlled linear recurrence for sequence modeling. ArXio, abs/2231.00197. 2023.
Bo Liu, Rui Wang, Lemeng Wu, Yihao Feng, Peter Stone, and Qian Liu. Longhorn: State space models amortized online learners. 2024. URL http:::/pppp...eeeete htt.///ppp..nmannicccchchoolll httphppp://ppp..eaatttcccccooo
orr//oorusID:2771310
Hanxiao Liu, Zihang Dai, David R. So, and Quoc V. Le. Pay attention to mips, 2021.
Huanru Henry Mao. Fine-tuning pre-trained transformers into decaying fast weights. In Proceedings of the 2022 Conference Empirical Methods in Natural Language Processing, pp. 10236-10242, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. doi: 10.18653/v1
Eric Martin and Chris Cundy. Parallelizing linear recurrent neural nets over sequence length. In 6th International Conference on Learning Representations, ICLR 2018, Vancourer, BC, Canada, April 30- May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018.
Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. Reveewei,, 2023.
12

---

Published as a conference paper at COLM 2024
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. 5th International Conference on Learning Representations, ICLR, Toulon, France, 2017.
Antonio
Orvieto, Samuel L. Smith, Albert Gu, Anushan Fernando, Çaglar Gülgehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pp. 26670-26688 26698. PMLR, 2023. URL htRpsttps:tm http:s:s//proceedingge mlr press/v202/ orvieto23a. html
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran G. V., Xuzheng He, Haowen Hou, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Xiangru Tang, Bolun Wang, Johan S. Wind, Stanislaw Wozniak, Ruichong Zhang, Zhenyuan Zhang, Qihang Zhao, Peng Zhou, Jian Zhu, and Rui-Jie Zhu. RWKV: reinventing mnns for the transformer era. CoRR, 2tee/2335..0000 2023. doi: 10.48550/ ARXTV
Bo Peng, Daniel Goldstein, Quentin Anthony, Alon Albalak, Eric Alcaide, Stella Biderman, Eugene Cheah, Teddy Ferdinan, Haowen Hou, Przemys aw Kazienko, G Kranthiki- ran, Jan Koco'n, Bartlomiej Koptyra, Satyapriya Krishna, Ronald McClelland, Niklas Muennighoff, Fares Obeid, Atsushi Saito, Guangyu Song, Haoqin Tu, Stanislaw Wo zniak, Ruichong Zhang, Bingchen Zhao, Qihang Zhao, Peng Zhou, Jian Zhu, and Ruijie Zhu. Eagle and finch: Rwkv with matrix-valued states and dynamic recurrence. ArXiv, abs/2404.000 2555.4444. 2024. URL htttps://apl www.stsiseoooenneeeeeee http://ccccccoo.ooooooorooosoo www.Corppooo. :690010055
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. arv preprint ar a'Xiv:2103.02144 2021.
Subhojeet Pramanik, Esraa Elelimy, Marlos C. Machado, and Adam White. Recurrent linear transformers. CoRR, abs hbbb//223..17199 2023.
Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. The devil in linear transformer. Xiv preprint arXiv:2210. 2022a.
Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong. cosformer: Rethinking softmax in attention. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OevnRevie,,net, 2022b.
Zhen Qin, Xiaodong Han, Weixuan Sun, Bowen He, Dong Li, Dongxu Li, Yuchao Dai, Lingpeng Kong, and Yiran Zhong. Toeplitz neural network for sequence modeling, In The Eleventh International Conference on Learning Representations (ICLR), 2023a. URL https://openreviert www.:/ewwwweeoooooooo/eooaomaaaa
Zhen Qin, Dong Li, Weigao Sun, Weixuan Sun, Xuyang Shen, Xiaodong Han, Yunshen Wei, Baohong Lv, Fei Yuan, Xiao Luo, et al. Scaling transnormer to 175 billion parameters. arxiv preprint ar 119999377.14995, 2023b.
Zhen Qin, Songlin Yang, and Yiran Zhong. Hierarchically gated recurrent neural network for sequence modeling. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10- 16, 2023, 2023c. URL httt:://pperrrs .ips cc halppppppllleeeeeeepppllppppppee Abbetactbbbsbattc CoaeettCooooefeeee html
Zhen Qin, Weigao Sun, Dong Li, Xuyang Shen, Weixuan Sun, and Yiran Zhong. Lightniig attention-2: A free lunch for handling unlimited sequence lengths in large language models. 2024.
13

---

Published as a conference paper at COLM 2024
Liliang Ren, Yang Liu, Yadong Lu, Yelong Shen, Chen Liang, and Weizhu Chen. Samba: Sim- ple hybrid state space models for efficient unlimited context language modeling, ArXiv, abs/2406.07522,, 2024. URL U/L_haggss//pp..ssmmmttttccooaao
semanticscholarr ort./Oorpus1D::2733022244
Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana lli',, Daniel Hesslow, Roman Castagn'e e,, Alexandra Sasha Luccioni, François Yvon, Matthias Gallé, Jonathan Tow, Alexander M. Rush, Stella Biderman, Albert Webson, Pawan Sasanka Ammana- manchi, Thomas Wang, Benoit Sagot, Niklas Muennighoff, Albert Villanova del Moral, Olatunji Ruwase, Rachel Bawden, Stas Bekman, Angelina McMillan-Major, Tz Beltagy, Huu Nguyen, Lucile Saulnier, Samson Tan, Pedro Ortiz Suarez, Victor Sanh, Hugo Lau- renccon, Yacine Jernite, Julien Launay, Margaret Mitchell, Colin Raffel, Aaron Gokaslan, Adi Simhi, Aitor Soroa Etxabe, Alham Fikri Aji, Amit Alfassy, Anna Rogers, Ariel Kreis- berg Nitzav, Canwen Xu, Chenghao Mou, Chris C. Emezue, Christopher Klamm, Colin Leong, Daniel Alexander van Strien, David Ifeoluwa Adelani, Dragomir R. Radev, Ed- uardo Gonz' alez Ponferrada, Efrat Levkovizh, Ethan Kim, Eyal Natan, Francesco De Toni, Gérard Dupont, Germán Kruszewski, Giada Pistilli, Hady ElSahar, Hamza Benyam- ina, Hieu Trung Tran, Ian Yu, Idris Abdulmumin, Isaac Johnson, Itziar Gonzalez-Dioss Javier de la Rosa, Jenny Chim, Jesse Dodge, Jian Zhu, Jonathan Chang, Jorg Frohberg, Josephine Tobing, Joydeep Bhattacharjee, Khalid Almubarak, Kimbo Chen, Kyle Lo, Lean- dro von Werra, Leon Weber, Long Phan, Loubna Ben Allal, Ludovic Tanguy, Manan Dey, Manuel Romero Muñoz, Maraim Masoud, Mar'ia Grandury, Mario vSavsko, Max Huang, Maximin Coavoux, Mayank Singh, Mike Tian-Jian Jiang, Minh Chien Vu, Mohammad A. Jauhar, Mustafa Ghaleb, Nishant Subramani, Nora Kassner, Nurulaqilla Khamis, Olivier Nguyen, Omar Espejel, Ona de Gibert, Paulo Villegas, Peter Henderson, Pierre Colombo, Priscilla Amuok, Quentin Lhoest, Rheza Harliman, Rishi Bommasani, Roberto L'opez, Rui Ribeiro, Salomey Osei, Sampo Pyysalo, Sebastian Nagel, Shamik Bose, Shamsuddeen Has- san Muhammad, Shanya Sharma, S. Longpre, Somaich Nikpoor, S. Silberberg, Suhas Pai, Sydney Zink, Tiago Timponi Torrent, Timo Schick, Tristan Thrush, Valentin Danchev, Vassilina Nikoulina, Veronika Laippala, Violette Lepercq, Vrinda Prabhu,, Zaid Alyafeai, Zeerak Talat, Arun Raja, Benjamin Heinzerling, Chenglei Si, Elizabeth Salesky, Sabrina J. Mielke, Wilson Y. Lee, Abheesht Sharma, Andrea Santilli, Antoine Chaffin, Arnaud Stiegler, Debaiyoti Datta, Eliza Szczechla, Gunjan Chhablani, Han Wang, Harshit Pandey, Hendrik Strobelt, Jason Alan Fries, Jos Rozen, Leo Gao, Lintang Sutawika, M Saiful Bari, Maged S. Al-Shaibani, Matteo Manica, Nihal V. Nayak, Ryan Teehan, Samuel Albanie, Sheng Shen, Srulik Ben-David, Stephen H. Bach, Taewoon Kim, Tali Bers, Thibault Févry, Trishala Neeraj, Urmish Thakker, Vikas Raunak, Xiang Tang, Zheng-Xin Yong, Zhiqing Sun, Shaked Brody, Y Uri, Hadar Tojarieh, Adam Roberts, Hyung Won Chung, Jaesung Tae, Jason Phang, Ofir Press, Conglong Li, Deepak Narayanan, Hatim Bourfoune, Jared Casper, Jeff Rasley, Max Ryabinin, Mayank Mishra, Minjia Zhang, Mohammad Shoeybi, Myriam Peyrounette, Nicolas Patry, Nouamane Tazi, Omar Sanseviero, Patrick von Platen, Pierre Cornette, Pierre Francoois Lavall' ee, Rémi Lacroix, Samyam Rajbhandari, Sanchit Gandhi, Shaden Smith, Stéphane Requena, Suraj Patil, Tim Dettmers, Ahmed Baruwa, Amanpreet Singh, Anastasia Cheveleva, Anne-Laure Ligozat, Arjun Subramo- nian, Aur' eeie N'ev'eol, Charles Lovering, Daniel H Garrette, Deepak R. Tunuguntla, Ehud Reiter, Ekaterina Taktasheva, Ekaterina Voloshina, Eli Bogdanov, Genta Indra Winata, Hailey Schoelkopf, Jan Jan-Christoph Kalo, Jekaterina Novikova, Jessica Zosa Forde, Xiangru Tang, Jungo Kasai, Ken Kawamura, Liam Hazan, Marine Carpuat, Miruna Clin- ciu, Najoung Kim, Newton Cheng, Oleg Serikov, Omer Antverg, Oskar van der Wal, Rui Zhang, Ruochen Zhang, Sebastian Gehrmann, Shachar Mirkin, S. Osher Pais, Ta- tiana Shavrina, Thomas Scialom, Tian Yun, Tomasz Limisiewicz, Verena Rieser, Vitaly Protasov, Vladislav Mikhailov, Yada Pruksachatkun, Yonatan Belinkov, Zachary Bam- berger, Zdenvek Kasner, Zdeněk Kasner, Amanda Pestana, Amir Feizpour, Ammar Khan, Amy Faranak, Ananda Santa Rosa Santos, Anthony Hevia, Antigona Unldreaj, Arash Aghagol, Arezoo Abdollahi, Aycha Tammour, Azadeh HajiHosseini, Bahareh Behroozi, Benjamin Ayoade Ajibade, Bharat Kumar Saxena, Carlos Muñoz Ferrandis, Danish Con- tractor, David M. Lansky, Davis David, Douwe Kiela, Duong Anh Nguyen, Edward Tan, Emi Baylor, Ezinwanne Ozoani, Fatim Tahirah Mirza, Frankline Ononiwu, Habib Rezane- jad, H.A. Jones, Indrani Bhattacharya, Irene Solaiman, Irina Sedenko, Isar Nejadgholi, Jan Passmore, Joshua Seltzer, Julio Bonis Sanz, Karen Fort, Lívia Dutra, Mairon Samagaio,
14

---

Published as a conference paper at COLM 2024
Maraim Elbadrii, Margot Mieskes, Marissa Gerchick, Martha Akinlolu, Michael McKenna, Mike Qiu, Muhammed Ghauri, Mykola Burynok, Nafis Abrar, Nazneen Rajani, Nour Elkott, Nourhan Fahmy, Olanrewaju Samuel, Ran An, R. P. Kromann, Ryan Hao, Samira Alizadeh, Sarmad Shubber, Silas L.. Wang, Sourav Roy, Sylvain Viguier, Thanh-Cong Le, Tobi Oyebade, Trieu Nguyen Hai Le, Yoyo Yang, Zach Nguyen, Abhinav Ramesh Kashyap, Alfredo Palasciano, Alison Callahan, Anima Shukla, Antonio Miranda-Escalada, Ayush Kumar Singh, Benjamin Beilharz, Bo Wang, Caio Matheus Fonseca de Brito, Chenxi Zhou, Chirag Jain, Chuxin Xu, Clémentine Fourrier, Daniel Le' on Perin' an, Daniel Molano, Dian Yu, Enrique Manjavacas, Fabio Barth, Florian Fuhrimann, Gabriel Altay, Giyaseddin Bayrak, Gully Burns, Helena U. Vrabec, Iman I.B. Bello, Isha Dash, Ji Soo Kang, John Srorgi,onas, Golde, Jose David Posada, Karthi Sivaraman, Lokesh Bulchandani, Lu Liu, Luisa Shinzato, Madeleine Hahn de Bykhovetz, Maiko Takeuchi, Marc Pàmies, María An- drea Castillo, Marianna Nezhurina, Mario Sanger, Matthias Samwald, Michael Cullan, Michael Weinberg, M Wolf, Mina Mihaljccic, Minna Liu, Moritz Freidank, Myungsun Kang, Natasha Seelam, Nathan Dahlberg, Nicholas Michio Broad, Nikolaus Muellner, Pascale Fung, Patricia Haller, Patrick Haller, Renata Eisenberg, Robert Martin, Rodrigo Canalli, Rosaline Su, Ruisi Su, Samuel Cahyawijaya, Samuele Garda, Shlok S Deshmukh, Shubhanshu Mishra, Sid Kiblawi, Simon Ott, Sinee Sang-aroonsiri, Srishti Kumar, Stefan Schweter, Sushil Pratap Bharati, Tanmay Laud, Théo Gigant, Tomoya Kainuma, Wojciech Kusa, Yanis Labrak, Yashasvi Bajaj, Y. Venkatraman, Yifan Xu, Ying Xu, Yu Xu, Zhee Xao Tan, Zhongli Xie, Zifan Ye, Mathilde Bras, Younes Belkada, and Thomas Wolf. Bloom: A 176b-parameter open multilingual language model. ArXiv, abs/2211 1.01100, 2022. URL https::/app..ss.ee http://app..smmannccccho..aa.a / htt.:/oo......og/ropppsi/////222
Imanol Schlag and Jürgen Schmidhuber. Gated fast weights for on-the-fly neural program generation. 2017. URL. httpp://appp ses.mmtteccciiool....oorrrrrooaa
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th Pr International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 9366- PMLR, 2021.
Uri Shaham, Elad Segal, Maor Ivgi, Avia Efrat, Ori Yoran, Adi Haviv, Ankit Gupta, Wenhan Xiong, Mor Geva, Jonathan Berant, and Omer Levy. SCROLLS: Standardized CompaRison over long language sequences. In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang (eds.), Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 12007-12021, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. doi: 10.18653 www///22..mmm/mmaii...m scllathology.org http:/22222.mmmmpmmeeeeee emnlp-main. 823
URL https://
Xuyang Shen. Limtest needleinahaystack hfmodel: Support huggingface model do simple retrieval from Iim models at various context lengths to measure accuracy, 2024. URL HFModel (ttpaaassttca_hhhoeelelll
htttps://githbb.com http:aaaa..accaaaaagggaaaaaaeeee
Xuyang Shen, Dong Li, Ruitao Leng, Zhen Qin, Weigao Sun, and Yiran Zhong. Scaling laws for linear complexity language models, 2024. utL.stts:://arx/ org/abs/2406 16690
Yikang Shen, Shawn Tan, Alessandro Sordoni, and Aaron C. Courville. Ordered neurons: Integrating tree structures into recurrent neural networks. ArXiv, abs/1810.09536, 2018. URL hollhttps://// //ap::::///i... semanti Trg//orpus 53034786
Jimmy T. H. Smith, Andrew Warrington, and Scott W. Linderman. Simplified state space layers for sequence modeling. In The Eleventh International Conference on Learning Represen- fations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenKeview.net, 2023.
Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey. SlimPajama: A 627B token cleaned and deduplicated version of RedPajama, 2023.
Yu Sun, Xinhao Li, Karan Dalal, Jiarui Xu, Arjun Vikram, Genghan Zhang, Yann Dubois, Xinlei Chen, Xiaolong Wang, Sanmi Koyejo, Tatsunori Hashimoto, and Carlos Guestrin. Learning to (learn at test time): Rnns with expressive hidden states. 2024a. URL https
15

---

Published as a conference paper at COLM 2024
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. arv preprint ar arXi:22307.08621 2023.
Yutao Sun, Li Dong, Yi Z Shaohan Huang, Wenhui Wang, Shuming Ma, Quanlu Zhang, Jianyong Wang, and Furu Wei. You only cache once: Decoder-decoder architectures for language models. ArXiv, /bb00240005554, 2024b. httpss////pp. http:///aapisemaniiiiiiicc sewanticscoolaa. (
Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng. Synthesizer: Rethinking self-attention in transformer models, 2021a.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. In 9th International Conference Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021b. URL https:///peerreiiee
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Herve Jegou. Training data-efficient image transformers &amp; distillation through attention. In International Conference on Machine Learning, volume 139, pp. 10347-10357, July 2021.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint ar arXi::2302.13971 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023b.
Jos van der Westhuizen and Joan Lasenby. The unreasonable effectiveness of the forget gate. CoRR, 1bbs/1804.04849 2018.
Junxiong Wang, Jing Nathan Yan, Albert Gu, and Alexander M. Rush. Pretraining without attention. CoRR, abs/ /abbs/2211 2212.10544, 2022.
Songlin Yang and Yu Zhang. FLA: A Priton-Based Library for Hardware-Efficient Imple- mentations of Linear Attention Mechanism, January 2024. URL htt.ttts::////// https:///tttttii com/ susscaongliin/flass linear- attention
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficieet training. CoRR, abs /3122.035,, 2023. doi: 10.48550/ ARXIV22312.0663 URL
48550//r11. 2312 06635
Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transformers with the delta rule over sequence length. arXiv preprint arXio:2406.06484, 2024.
Jun Zhang, Shuyang Jiang, Jiangtao Feng, Lin Zheng, and Lingpeng Kong. Linear attention via orthogonal memory, 2023.
16

---

Published as a conference paper at COLM 2024
Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Ré. The hedgehog & the porcupine: Expressive linear attentions with softmax mimicry, 2024.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona T. Diab, Xian Li, Xi Victoria Lin, Todor Mihaylov, Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Punit Singh Koura, Anjali Sridhar, Tianlu Wang, and Luke Zettlemoyer. Opt: Open pre-trained transformer language models. ArXio, abs/2205.01068, 2022. URL htthtsss://app www.sssseaannccchiiaeee.ee http:///w..pp...aaaa....oeoo/or/ wrg/CorpusII........ :444499922 248446692
Wei Zhang and Bowen Zhou. Learning to update auto-associative memory in recurrent
neural networks for improving sequence memorization. ArXiv, abs//170.064/ 2017.
URL sem.aiiissmmmnnicccclollaa Www.aaarrroooooooeee http://a......rrooouuui::::::550
A
Appendix
A.1 Experiment Configurations
Table 12 provides detailed setups for both Auto-regressive Language Modeling (ALM) and ImageNet (IM) experiments, focusing on the Wikilext-103 and ImageNet-1k datasets, respectively. The ALM experiments utilize Byte Pair Encoding (BPE) with a vocabulary size of 50,265 and sequence length of 512, featuring total batch size of 128 and 50,000 updates. The ImageNet experiments differentiate between 6 million and 23 million parameter models, with total batch sizes of 1024 and 2048, both running for 300 epochs but with differing warm-up periods. Optimization strategies vary between Adam for ALM and AdamW for IM, with specific learning rate schedulers and hyperparameters tailored to each model's scale. Additional configurations outline variations in model complexity, ranging from 0.15 to 2.9 million parameters, adjusting layers, hidden dimensions, and GPUs used, aiming to comprehensively explore model performance across scales and setups.
Table 12:
Comprehensive Configurations of the Model and Training Procedures for HGRN2 Experiments. "Total batch size" means batch_per-gpu x update_fred ffeeq x num-gpus; "LLM" stands for Autoregressive Language Model; "IM" stands for Image Modeling.
ALM
IM(6M)
IM(23M)
Dataset
1ext-exx
ImageNet-lk
ImageNet-Ik
Tokenizer method
BPE
Src Vocab size
50265
Sequence length
512
Total batch size
128
1024
2048
Number of updates/epoxt epochs
50k updates
300 epochs
300 epochs
Warmup steps/ /poocs
4k steps
20 epochs
10 epochs
Peak learning rate
5e-4
7.5e-4
7.5e-4
Learning rate scheduler
Inverse sqrt
Cosine
Cosine
Optimizer
Adam
AdamW
AdamW
Adam
1e-8
1e-8
1e-8
Adam (f1,82)
(0.9, 0.98)
(0.9, 0.98)
(0.9, 0.98)
Weight decay
0.1
0.05
0.1
Gradient clipping
5.0
5.0
Table 13: Model Configurations
Params
Layers 15
Hidden Dim
Exp. Ratio
L. R.
Batch Size
SeqLen
GPUs
0.15
768
128
3.00E-04
26
2048
0.385
26
1024
128
3.00E-04
15
2048
1
18
2048
128
3,00E-04
10
2048
16
2.9
36
2560
128
3.00E-04
36
2048
17


Reviews:
[{"review_id": "MPLxN5K1qV", "reviewer": "Reviewer_m3yg", "summary": "The paper proposes HGRN2 (Hierarchically Gated Linear RNN2), which is an improvement over the previous [HGRN](https://arxiv.org/abs/2311.04823). The main changes of HGRN2 are state expansion of the recurrent state from d to nd, which is done through some clever matrix multiplication and diagonalization tricks (exact details in Sec 3.2). The paper evaluates on various tasks such as Multi-Query Associative Recall (MQAR), language modeling on WikiText (perplexity eval), downstream common sense evals, and long-range arena. It shows that HGRN2 outperforms HGRN1 on most tasks, and it also performs a number of other models such as Mamba, OPT, Pythia, etc. \n\nAlso, a quick timeline here, in case it helps other reviewers/ACs who are wondering why the paper did not discuss baseline X:\n- COLM deadline: March 29\n- [Jamba](https://arxiv.org/abs/2403.19887) release: March 28 (probably too late to include in the paper) \n- [Infini-attention](https://arxiv.org/abs/2404.07143) release: April 10\n- [RecurrentGemma](https://arxiv.org/abs/2404.07839) release: April 11", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": true, "final_justification": "", "rebuttal": "Q1. Incremental improvements over the previous work HGRN.\n\nA1. The primary contribution of this work is to recognize and emphasize the significance of state expansion and to offer  a simple yet elegant solution for expanding the recurrent state. In comparison to complex network structure modifications, we believe that our adjustment to HGRN is more foundational and addresses the core weakness of HGRN. Additionally, we provide analysis that supports our conclusion and offer comprehensive evaluations of state expansion strategies. Given these factors, we are confident that HGRN2 is a solid work for submission to a full conference.\n\nQ2. The transformer models that the paper compares with are weak.\n\nA2. The performance of a model depends on two factors: the model's capacity and token consumption. To make a fair comparison based on model capacity, the number of tokens consumed should be similar. We have provided these baselines because their token consumption (around 300B) is relatively close to our token consumption (around 100B). In contrast, updated transformer models such as Gemma-2B, StableLM2-1.6B, and phi2 consume around or more than 1T tokens, which we believe would make the comparison highly unfair. For a fair comparison result, please refer to Reviewer cvJL's A3.\n\nQ3. More long-context evals.\n\nA3. Please refer to Reviewer 9TbB's A1.\n\nQ4. Why are your Mamba results in Table 6 different form the results reported in the Mamba paper? \n\nA4. We are using the same setup as described in [1] and training our models with 100 billion tokens. To ensure a fair comparison, we are referencing the Mamba results from [1]. The original Mamba results, as outlined in the Mamba paper, are derived from a model trained with 300 billion tokens.\n\nQ5. Why is Mamba there for the 1B models but not there for the 3B models?\n\nA5. Our model only used 100 billion tokens, while the original Mamba 3B used 300 billion tokens, making the comparison unfair. For a fair comparison, we trained on 10 billion tokens from the same dataset. The results are presented in the table below, showing that HGRN2 has a lower perplexity (ppl) at 3 billion.\n\n| Method | PPL |\n| --- | --- |\n| Llama-3b | 7.129 |\n| HGRN2-3b | 6.989 |\n| Mamba-3b | 7.056 |\n\nCitation.\n1. Yang et. Gated Linear Attention Transformers with Hardware-Efficient Training.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- I feel like linear models are generally becoming quite popular these days, so this paper is very much aligned with the trend of linear attention research in 2024.\n- Extensive evaluation setup, covering a wide range of tasks and settings. Good comparisons with other models. I especially liked Tables 6 and 7, where a lot of models were compared against, including both transformer models and other linear models/state-space models.\n- The paper evaluated on actual natural language tasks! I'm very happy to see this, since I've seen a lot of papers in this area that report only validation loss/perplexity, so it's nice to see this paper actually cover downstream natural language tasks and commonsense tasks.", "reasons_to_reject": "- Incremental improvements over the previous work HGRN (This is not to say that it's not good work; I just feel like it might be more suitable for a blog/technical report rather than a full conference submission).\n- The transformer models that the paper compares with are weak -- OPT, BLOOM, and Pythia are known to be pretty bad models. For smaller sizes, there are stronger models such as Gemma-2B, StableLM2-1.6B, and phi2. I suspect that if the paper compares with these models, then these will probably perform better than HGRN2, since the weak models like BLOOM already perform quite close to it. \n   - As a consequence of above: If the models fail to beat already existing good transformer models, then what's the whole point to begin with?\n- I'm also not very convinced about long-context evals. I feel like long range arena is somewhat outdated and that a lot of the tasks feel like toy tasks and not very related to natural-language. These days there are a lot of newer evals such as needle-in-the-haystack style evals or natural language evals like [SCROLLS](https://arxiv.org/abs/2201.03533).", "questions_to_authors": "- Why are your Mamba results in Table 6 different form the results reported in the [Mamba](https://arxiv.org/pdf/2312.00752) paper? \n   - Also, why is Mamba there for the 1B models but not there for the 3B models? The original paper/release had both 1B and 3B models.\n\n\n\n[Note: Increased my score after rebuttals (see discussion below)]", "extra_scores": {}}, {"review_id": "qgXKfa0NjJ", "reviewer": "Reviewer_cvJL", "summary": "The paper presents a modification to an existing recurrent neural network architecture called hierarchically gated linear RNN. The goal is to improve the expressiveness that is limited by its small state size. This is a niche topic, but given the increased interest to RNN-like alternatives for Transformers, standard RNN architectures are of interest to the community. \n\nSpecifically, they achieve the state size expansion by employing outer product operation on the hidden state and then sum the expanded representations to a single vector. In terms of originality, this technique is not new as it has been previously used in several linear attention works although it is interesting to see it applied in this standard RNN architecture.\n\nEvaluation covers a number of tasks including multi-query associative recall, language modeling, commonsense reasoning, and image classification. The main baseline is the hierarchical gated RNN and other efficient architectures such as Mamba, RWKV, Griffin, and others. Even though coverage is fair, the results would have been more convincing if it covered more on downstream generative tasks.\n\nOverall, it's a focused contribution with strong empirical focus but it has a limited scope and it lacks methodological novelty.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Q1. The novelty of this work\n\nA1. The primary contribution of this work is to recognize and emphasize the significance of state expansion and to offer  a simple yet elegant solution for expanding the recurrent state. In comparison to complex network structure modifications, we believe that our adjustment to HGRN is more foundational and addresses the core weakness of HGRN. Additionally, we provide analysis that supports our conclusion and offer comprehensive evaluations of state expansion strategies.  In terms of structure modification, our model leverateges parameter sharing for keys and decay rates in linear attention and eliminates the output gate of linear attention, making the overall architecture concise.\n\nQ2. Restrictive for the scope of the study.\n\nA2. Thank you for your question. In an additional experiment, we want to show that state expansion works equally well for general linear RNNs. We compared the perplexity (ppl) at different expansion ratios on the WikiText-103 dataset for a basic linear RNN(i.e. HGRN2 without the lower bound). We found that as the expansion ratio increases, the perplexity decreases.\n\n| Expand ratio | PPL(val) | PPL(test) |\n| --- | --- | --- |\n| 1 | 25.17 | 26.01 |\n| 2 | 24.56 | 25.41 |\n| 4 | 24.27 | 25.05 |\n| 8 | 23.86 | 24.51 |\n| 16 | 23.69 | 24.28 |\n| 32 | 23.28 | 23.93 |\n| 64 | 23.1 | 23.73 |\n| 128 | 23.21 | 23.76 |\n\nQ3. More evaluation result.\n\nA3. For a fair compare, we trained Llama and HGRN2 on the same dataset with 300B tokens sampled from the SlimPajama and Pile datasets, and it can be seen that HGRN2 performs better on most tasks.\n| Model | PIQA | HS | WG | ARC-E | ARC-C | OBQA | CSR-AVG | DROP | MMLU |\n| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |\n| HGRN2-385m | 68.01 | 40.37 | 52.25 | 53.66 | 24.32 | 31.20 | 44.97 | 4.16 | 25.07 |\n| Llama-385m | 67.19 | 38.75 | 52.17 | 49.24 | 23.72 | 30.00 | 43.51 | 3.95 | 27.08 |\n| HGRN2-1b | 71.60 | 49.45 | 53.91 | 60.40 | 28.24 | 33.20 | 49.47 | 3.42 | 25.47 |\n| Llama-1b | 69.97 | 47.04 | 52.72 | 57.07 | 28.16 | 32.60 | 47.93 | 5.88 | 24.9 |\n| HGRN2-3b | 74.37 | 61.49 | 58.41 | 65.49 | 34.56 | 36.00 | 55.05 | 5.20 | 25.33 |\n| Llama-3b | 73.18 | 57.88 | 59.59 | 63.93 | 31.40 | 34.00 | 53.33 | 3.92 | 27.52 |\n| HGRN2-7b | 76.77 | 66.81 | 60.93 | 69.15 | 36.43 | 38.00 | 58.02 | 5.67 | 32.99 |\n| Llama-7b | 75.19 | 64.39 | 61.88 | 67.55 | 35.41 | 35.00 | 56.57 | 4.19 | 32.2 |", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- Straight-forward, clear exposition and a simple method that address the capacity issue in hierarchically gated RNNs. It can also be applied to any gated-recurrent network.\n- Addresses a problem that is not typical in the present literature; this attempt is forward looking and led to interesting findings. For instance, the improved gated RNN actually outperforms stronger efficient architectures based on Transformers. \n- The claims made in the introduction are supported well by empirical evidence; the main target of comparison is the hierarchical gated RNN.", "reasons_to_reject": "- The proposed method for addressing the capacity issue is from prior work and its methodological novelty is limited. \n- The evaluation focuses mainly on improving hierarchical gated RNN only that is restrictive for the scope of the study.", "questions_to_authors": "- I was wondering if you can report scores on some of the general benchmarks such as MMLU, BBH or DROP that large language models are typically being evaluated on. This would help show that the method can generalize more broadly.", "extra_scores": {}}, {"review_id": "phQjiKoE2E", "reviewer": "Reviewer_HShg", "summary": "HGRN2 presents extensions to HGRN by introducing a simple outer product-based state expansion mechanism, resulting in a larger recurrent state without introducing any additional parameters.  This corresponds to a linear attention interpretation of HGRN2 and hardware efficient training.  Mamba and GLA both use state expansion techniques.  HGRN has data dependent decays.  HGRN2 is evaluated on language modeling, image classification and long range arena.  The following sections describe background on the GRN and the HGRN.\n\nThe authors then explore parameter efficient state expansion methods, as shown in Table 2.  Section 3.2 then describes HGRN2, which does not use any of the approaches in Table 3.  They introduce a multi-headed HGRN as with linear attention.  Figure 2 shows better LM performance as the expansion ratio is expanded.  They compare HGRN to GLA and note all the similarities to the point that they can use their implementation.  They then perform various language modeling experiments, including WikiText and also long range arena.\n\nThere are some interesting experiments in this paper, and the updated HGRN could be an interesting contribution.  I have a few questions below and concerns about the experiments as described below.\n\n* This paper would be much more impactful if it were planned to release the code for this?\n* The Wikitext experiment needs more details.  For example, why wasn’t Mamba evaluated for perplexity?  Also, does perplexity hold for larger models (50M is relatively small).\n* The approach seems pretty similar in to GLA in that they can share the same implementation and am concerned about differentiation from this work.  Can the authors expand on their description in the paper?\n* Why wasn’t Mamba 2.8B in Table 6?  And perhaps Phi2 as well, which is more performant than the ones you put in the table.  \n* And for long range arena in Table 7, why not also evaluate on some newer benchmarks like Qasper?  But also, the table is deceiving, as the context length on which the model was trained isn’t listed in the table.  What context length was the transformer in the table trained on?  And HGRN2?  And did you try Rope extension for transformer models?  https://arxiv.org/html/2402.13753v1\n\nSmall things:\n* Table 1 shows real HGRN1 with 1.5m more parameters, which seems like a lot and could be the reason it’s better than the complex version.  I don’t think this changes the result, but it’s sort of an unfair comparison.\n* Would a larger expand ratio have (e.g., 256) have performed better?  Or were there diminishing returns?\n* It would be nice if there was some other baseline (e.g., Transformer) in Fig. 3.\n* Fig 4 (Right) is labeled HRGN instead of HGRN.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": true, "final_justification": "", "rebuttal": "Q1. Are there plans to open-source the code?\n\nA1. Yes, we plan to open-source the code very soon.\n\nQ2. More about Wikitext experiment and larger model result.\n\nA2. We have updated the results of Mamba on the WikiText-103 dataset, as shown in the table below; additionally, we tested Mamba on 10 billion tokens in a subset of the Pile and observed that Hgrn2 outperforms Mamba.\n\nWikitext-result:\n| Method | PPL(val) | PPL(test) | Params |\n| --- | --- | --- | --- |\n| HGRN2 | 23.1 | 23.73 | 44.66 |\n| Mamba | 22.58 | 23.19 | 44.39 |\n\n10b token result, the same as figure 4.\n| Method | PPL |\n| --- | --- |\n| Llama-3b | 7.129 |\n| HGRN2-3b | 6.989 |\n| Mamba-3b | 7.056 |\n\nQ3. Relation to GLA and the novelty of HGRN2.\n\nA3. For most part, please refer to Reviewer cvJL' A1. In comparison with GLA, it is noteworthy that the parameters of keys and decay rates in linear attention can be shared. To the best of our knowledge, this has not been proposed previously in the literature. Our model also eliminates the output gate of linear attention, making the overall architecture more concise.\n\nQ4. More about table 6.\n\nA4. In future updates, we will add more results. For the reason, please refer to Reviewer m3yg'A1.\n\nQ5. More about Table 7.\n\nA5. The transformer results for the table are sourced from [1]. All models are trained from scratch, so there is no issue of extrapolation. In the future, we will also denote the sequence length for each task.\n\nQ6. More long context result.\n\nA6. Please refer to Reviewer 9TbB's A1.\n\nQ7. Update Table 1.\n\nA7. We have updated Table 1 to provide a version with fewer parameters, expand 2. As you can see, the performance is still better.\n\n| Method | State size | PPL(val) | PPL(test) | Params |\n| --- | --- | --- | --- | --- |\n| HGRN1 | d | 24.14 | 24.82 | 46.25 |\n| Real HGRN1 | 2d | 24.04 | 24.64 | 45.46 |\n\nQ8. Larger expand ratio.\n\nA8. We provide the results of expand 256 in the table below. As seen, the performance is better but the improvement is not significant. As a compromise, we chose an expand ratio of 128 for lm.\n\n| expand ratio | PPL(val) | PPL(test) | \n| --- | --- | --- | \n| 128 | 23.1 | 23.73 | \n| 256 | 23.41 | 23.79 | \n\nQ9. Transformer Baseline for MQAR.\n\nA9. We will provide it in future versions.\n\nQ10. Fig 4 typo.\n\nA10. We will fix this issue in the revised paper.\n\nCitations:\n[1] Yi Tay, et. Long Range Arena : A Benchmark for Efficient Transformers . In *International Conference on Learning Representations*.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "The paper presents incremental extensions to HGRN1 and some performance gains as a result over HGRN1.", "reasons_to_reject": "* The results aren't well-described and at times inconsistent (see above).\n* The approach seems to not differentiate much from GLA after performing the new state expansion, which makes the contributions somewhat limited.", "questions_to_authors": "See above in summary.", "extra_scores": {}}, {"review_id": "GMzrAWv9nV", "reviewer": "Reviewer_9TbB", "summary": "The paper introduces HGRN2 which brings a simple but in practice significant improvements over the well performing HGRN model, an RNN based SOTA language model. \n\nThe paper is clear and easy to follow. The paper focuses on efficient state expansion to improve RNN model performance and the first part of section 3 motivates the state expansion. \n\nNot only the approach is tested and compared with the most natural baseline HGRN1. The paper also test on various benchmarks (LGA, wiki, etc.) and compare with extensive number of models including Llama. The empirical results thus look convincing; the paper looks like a good contribution on RNN language models as alternatives to the main stream transformer based language models.", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Q1.  Compared transformer family models are a bit outdated.\n\nA1. Firstly, the network structures of Llama1/2/3 are almost identical, with the only difference being the tokens consumed. Therefore, the Transformer structure we provide remains state-of-the-art. We will add more results in future updates. As a clarification, the reason we did not include many updated models is because of the significant difference in token consumption. For example, Phi2 consumes 1 T tokens, while HGRN2 only consumes 100b, making the comparison unfair. As a supplement, we trained Llama and HGRN2 on the same dataset with 300B tokens sampled from the SlimPajama and Pile datasets, and it can be seen that HGRN2 performs better on most tasks. For the most result, please refer to Reviewer cvJL's A3. We also provide result on Scroll dataset:\n\n| Model | GovRep | SumScr | QMSUM | Qspr | QALT | CNLI |\n| --- | --- | --- | --- | --- | --- | --- |\n| HGRN2-385m | 15.78/3.72/11.26 | 7.13/0.76/6.03 | 8.03/1.28/6.99 | 12.39 | 26.46 | 31.53 |\n| Llama-385m | 7.51/2.21/5.92 | 9.96/1.26/7.76 | 6.32/0.78/5.71 | 17.48 | 27.95 | 13.5 |\n| HGRN2-1b | 20.26/5.28/13.66 | 11.83/1.72/9.11 | 9.65/1.74/8.82 | 19.82 | 27.61 | 11.38 |\n| Llama-1b | 7.53/2.13/6.19 | 9.34/1.33/7.31 | 7.82/1.17/6.4 | 18.99 | 27.85 | 26.9 |\n| HGRN2-3b | 21.89/6.7/14.2 | 15.02/2.15/10.83 | 12.48/2.61/10.46 | 25.2 | 28.81 | 31.92 |\n| Llama-3b | 9.77/2.95/7.27 | 10.39/1.49/8.1 | 7.72/1.46/6.27 | 28.8 | 30.58 | 20.15 |\n\nQ2. Provide some scaling law.\n\nA2. We sampled 300 billion tokens from the SlimPajama and Pile datasets and then retrained the HGRN2 model. We evaluated the model's performance on commonsense reasoning tasks and observed that the model's performance improves as more tokens are consumed.\n\n| Model | Tokens | PIQA | HS | WG | ARC-E | ARC-C | OBQA | AVG |\n| --- | --- | --- | --- | --- | --- | --- | --- | --- |\n| hgrn2_385m | 150 | 68.17 | 38.33 | 50.91 | 51.56 | 25.60 | 30.80 | 44.23 |\n| hgrn2_385m | 300 | 67.74 | 40.32 | 51.78 | 54.21 | 24.83 | 31.20 | 45.01 |\n| hgrn2_1b | 150 | 70.24 | 46.31 | 53.99 | 57.28 | 27.30 | 33.40 | 48.09 |\n| hgrn2_1b | 300 | 71.65 | 49.52 | 54.38 | 60.27 | 28.07 | 33.40 | 49.55 |\n| hgrn2_3b | 160 | 73.88 | 57.32 | 57.85 | 63.72 | 33.87 | 36.60 | 53.87 |\n| hgrn2_3b | 300 | 74.10 | 61.48 | 58.64 | 65.61 | 34.47 | 35.60 | 54.98 |\n| hgrn2_7b | 150 | 74.81 | 62.30 | 61.25 | 65.95 | 36.26 | 36.60 | 56.20 |\n| hgrn2_7b | 300 | 76.50 | 66.96 | 61.40 | 69.02 | 36.86 | 38.00 | 58.12 |", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "The paper follows recent trend to leverage RNN mechanisms to enable capturing long range dependencies, overcoming some inherent limitations of tranformer architectures. The paper is clear to communicate its contributions and situate well its core contributions within existing works from both recent RNN developments as well as transformers with respect to extending the range. The paper covers related works appropriately (e.g. relate to early state expansions techniques such as NTM).\n\nThe techniques proposed by the paper is very focused and can be implemented relatively easily. The simple changes and the significant gain of the proposal makes the techniques particularly attractive.", "reasons_to_reject": "Although the paper has done arguably extensive comparisons over existing models, I find that the transfomer family models in the comparison are a bit outdated. For example, the transformer models in Table 6 and as well as the whole paper are around 2022-2023 while the field has made improvements to the architecture in tranformers and even Llama 2 is arguably a bit outdated (released July 2023). However, given that it is not the main focus of the current paper, this should be considered as a minor flaw of the paper.\n\nThe paper scales as large as 3B and it is still a relatively small model in today's standard. I understand the cost concern but some initial scaling law plot should be considered, since the question is how well the model scales is central for the actual adoption.", "questions_to_authors": "I suppose that the performance in the paper refers to performance in test set, right? It will be good to make such points clear in the paper.", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper presents an expansion of the recurrent state in HGRN, unlocking better performance as demonstrated in the paper and in extended evaluations presented throughout the rebuttal. As the community explore new sub-quadratic transformer alternatives, I find the ideas and evidence in this paper very timely and important to present at COLM, despite the reviewers' lack of enthusiasm. The authors have also clearly made a heroic effort to address the reviewers' original concerns (more natural evals, robust comparison with transformers, and scaling trends), which I believe greatly improves the paper and merits its acceptance to COLM.

[comments from the PCs] Please revise the paper to include the additional content from the discussion period as noted by the AC.
Decision: Accept