HGRN2: Gated Linear RNNs with State Expansion

Authors: zhen  qin
songlin  yang
weixuan  sun
xuyang  shen
dong  li
weigao  sun
yiran  zhong
mit  csail
shanghai ai opennlplab
  lab
Venue: colmweb_org_COLM_2024_Conference
Published: 2024-03-22 07:58:54

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.

Main paper: 
 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 management. 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 al., 2023; Orvieto et al., 2023; Qin et al., 2023c; De et al., 2024) , linear attention (Sun et al., 2023; Qin 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 Westhuizen & 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 as Mamba (Gu & Dao, 2023) , Gated Linear Attention (GLA, Yang et al. 2023) , Griffin (De et al., 2024) , and RWKV-6 (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 stateof-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) .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 a parameter-efficient manner. Empirically, we found 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. (2023) ; Qin et al. (2024) for large-scale experiments. As a result, HGRN2 can be regarded as an improved parameterization of GLA.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. 
gated linear rnn: 
 Given input x ∈ R N×d , where the sequence length is N and the model dimension is d, a minimalist gated linear recurrent layer (Martin & Cundy, 2018) transforms the input x into hidden states h ∈ R N×d and the output y ∈ R N×d , as defined below:EQUATIONwhere ⊙ denotes element-wise product; σ is the sigmoid function, and τ is a nonlinear activation function (we choose to use SiLU); i t is the input vector; g t and o t are the forget gate and output gate, respectively. The input gate is tied to the forget gate as 1g t , a common approach used in many gated RNNs such as GRU (Chung et al., 2014) . (Qin 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) , we empirically found that complex-valued recurrence is not necessary, as shown in Table 1 . 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. β := cumax(Γ) = cumsum(softmax(Γ, dim = 0), dim = 0) ∈ R L×d , β i = [β] i ∈ R d .To prevent the lower bound from reaching one in the highest layer, HGRN subtracts all β values by β 0 , making the lower bound for the first layer zero. After obtaining the lower bound values, the forget gate g t learns the residuals instead, resulting in the new forget gate f t :EQUATIONwhere 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) . 
explorations of state expansion methods: 
 The goal of this work is to scale the size of the HGRN recurrent state from d to nd, where n is the state expansion ratio. However, if we use the original parameterization in Eq. 1, the matrices U, V, W will have dimensions d × nd, which becomes very parameter inefficient when n is large. Ideally, the number of parameters should be around d 2 , 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 R d → R nd , 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, d nd → nd nd 2 LR d, d r, r nd → nd dr(n + 1) ≈ d 2 GLT d = (n e) → n e n e, n e d → n d d 2 GLTI d = (n e) → n e n e, n e d, n n → nd d 2 + n 2 KRP d, n d → nd nd KP d, d d, n → nd d 2 + nAfter obtaining the expanded g, i, o, we feed them into element-wise gated linear recurrent layers as in Eq. 1 and Eq. 2, resulting in the output vector y t ∈ R n×d . To project the expanded dimension back to the original dimension, we simply sum over the dimension corresponding to n.The results are shown in Table 3 . We found that state expansion generally improves performance, with the low-rank matrix performing the best among these candidates. However, these methods face training inefficiency issues, as they require conducting element-wise linear recurrence in high dimensions (i.e., nd). Since these element-wise operations cannot leverage tensor cores (a fast matrix multiplication unit on GPUs), the dramatically increasing FLOPs and I/O costs significantly slow down training when n is large. We notice that this is similar to the case in Mamba1 , which requires a relatively small expansion ratio (i.e., n = 16) and a custom I/O-efficient CUDA implementation to achieve a reasonably fast running speed.In the next subsection, we explore an alternative strategy that does not replace the dense projection matrices with structured ones but instead changes the element-wise gating operations in Eq.1 to other matrix/vector operations similar to those used in linear attention. This approach allows for more efficient training. 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 t ∈ R d×d , 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:EQUATIONwhere Diag denotes the diagonalization of vectors, • represents the matrix dot product, and ⊗ indicates the outer product. = d h = d/H. 2We conducted an ablation study on the expansion ratio (or head dimension) n = d H , 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 h = 128 for the main experiments.Comparison to GLA. It is important to note that the recurrence form in HGRN2 is identical to that of GLA (Yang et al., 2023) , except for the specific parameterization. We list the correspondences between the two parameterizations in Table 4 . 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. Hardware-Efficient Training. Due to its computational structure's similarity to GLA, we can directly leverage their chunkwise algorithm and highly optimized kernels for hardwareefficient large-scale training. For more details, we refer readers to their paper.Concluding Remarks. Although HGRN2 shares many similarities with GLA, we believe that HGRN2 offers a unique perspective distinct from linear attention, originating from the approach of gated linear RNNs. For instance, it may not be immediately clear from the perspective of linear attention why key vectors should be constrained within the range of (0, 1) or why the key vector and forget gate value should sum to one. However, these concepts become quite intuitive when starting from the gated linear RNN framework and exploring state expansion. 
mqar: 
 Setting. Multi-Query Associative Recall (MQAR) (Arora et al., 2023) is an enhanced version of the synthetic induction head dataset (Fu et al., 2023) , designed to test the in-context associative recall ability of subquadratic models. Arora et al. (2023) found strong correlations between MQAR accuracy and language modeling performance. Our experimental setting strictly follows the original paper3 . Our hyperparameter sweep included the following ranges: expansion ratio ∈ {64, 128} and learning rate ∈ {1e -5, 5e -5, 1e -4, 5e -4, 1e -3, 5e -3, 1e -2}.Result. As shown in Fig. 3 Setting. For the Wikitext-103 experiment, we followed the configuration of HGRN1 to validate the performance of 44M models against a wide range of subquadratic models: FLASH (Hua et al., 2022) , 1+elu (Katharopoulos et al., 2020 ), Performer (Choromanski et al., 2021) , cos-Former (Qin et al., 2022b) , Syn(D), Syn(R) (Tay et al., 2021a) , gMLP (Liu et al., 2021) , S4 (Gu et al., 2022a) , DSS (Gupta & Berant, 2022) , RWKV-v4 (Peng et al., 2023) , LRU (Orvieto et al., 2023) , HGRN1 (Qin et al., 2023c) , TNN (Qin et al., 2023a) , and Mamba (Gu & Dao, 2023) . All reported results are from our own runs under the same settings.Result. Table 5 shows the results. HGRN2 clearly outperforms HGRN1 but slightly underperforms Mamba.We conducted language modeling experiments with 1.3B and 2.7B parameters on the Slimpajama dataset (Soboleva et al., 2023) , using the FLASHLINEARATTENTION (Yang & Zhang, 2024) codebase for training. 4 The results, shown in Table 6 , demonstrate that HGRN2 consistently outperforms other competitive linear recurrent models across three model scales. This suggests that HGRN2 provides a superior parameterization compared to GLA, as both models share an identical recurrent structure. 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. 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 et al., 2022) , Pythia (Biderman et 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 a Haystack Test. 5 LLaMA almost achieves perfect retrieval performance for evaluation Setting. Long Range Arena (Tay et al., 2021b ) is a benchmark designed to assess a model's ability to 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.Setting. For the image classification task, we followed the configuration of HGRN1 and trained it on ImageNet-1k, comparing it with TNN and the vanilla transformer.Result. Table 11 shows the results. HGRN2 outperforms HGRN1 with a similar parameter size, while also demonstrating an advantage over previous TNN (Qin et al., 2023a) and DeiT models (Touvron et al., 2021) . (Gu et al., 2022a) and its more efficient diagonalized version (Gu et al., 2022b) . Despite excellent performance in the LRA benchmark, it has been shown to have inferior 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., 2022; Gu & Dao, 2023) . Gupta et al. (2022) build the connection between SSM and linear RNN. Orvieto et al. (2023) proposes a linear RNN layer (i.e., LRU) inspired by SSMs. Peng et al. (2023) successfully scale linear RNN models to billions of parameters for the first time.For linear attention models, their language modeling performance has been underperforming softmax attention for a 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) , (ii) using local attention (Qin et al., 2022a; Zhang et al., 2023; 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 et al., 2021; Yang et al., 2024; Liu et al., 2024; Sun et al., 2024a) .Gated linear recurrence. Martin & Cundy (2018) first proposed a minimal gated linear recurrent layer and showed how to use 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 datadependent ones. They further use sliding-window attention to 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 to fast weight programming (FWP Schlag et al. More recently, Mamba2 (Dao & Gu, 2024) , xLSTM (Beck et al., 2024) , and Gated Retention (Sun et al., 2024b) have shown that sharing data-dependent decays across different dimensions 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/HGRN2, each head dimension has its own decay rate, whereas in Mamba2/xLSTM/Gated Retention, all dimensions share the decay under a single head. It is an interesting question to study how much improvement fine-grained decay will bring. 
conclusion: 
 In this work, we propose HGRN2, an enhancement of HGRN (Qin 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. 
 
 Appendix: Though 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 GPU acceleration, as well acknowledged in Mamba2(Dao & Gu, 2024).See Bolya et al. (2022) for more detailed complexity analysis.https://github.com/HazyResearch/zoologyModel checkpoints are available at https://huggingface.co/fla-hub.In this mode(Shen, 2024;Shen et 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. 


Decision: 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.

Reviews: review: paper_summary: paper_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. 

Also, a quick timeline here, in case it helps other reviewers/ACs who are wondering why the paper did not discuss baseline X:
- COLM deadline: March 29
- [Jamba](https://arxiv.org/abs/2403.19887) release: March 28 (probably too late to include in the paper) 
- [Infini-attention](https://arxiv.org/abs/2404.07143) release: April 10
- [RecurrentGemma](https://arxiv.org/abs/2404.07839) release: April 11
strength_weakness: strength_weakness: - 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).
- 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. 
   - 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?
- 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: questions: - Why are your Mamba results in Table 6 different form the results reported in the [Mamba](https://arxiv.org/pdf/2312.00752) paper? 
   - 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.



[Note: Increased my score after rebuttals (see discussion below)]
paper_summary: paper_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. 

Specifically, 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.

Evaluation 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.

Overall, it's a focused contribution with strong empirical focus but it has a limited scope and it lacks methodological novelty.
strength_weakness: strength_weakness: - The proposed method for addressing the capacity issue is from prior work and its methodological novelty is limited. 
- The evaluation focuses mainly on improving hierarchical gated RNN only that is restrictive for the scope of the study.
questions: questions: - 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.
paper_summary: paper_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.

The 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.

There 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.

* This paper would be much more impactful if it were planned to release the code for this?
* 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).
* 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?
* 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.  
* 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

Small things:
* 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.
* Would a larger expand ratio have (e.g., 256) have performed better?  Or were there diminishing returns?
* It would be nice if there was some other baseline (e.g., Transformer) in Fig. 3.
* Fig 4 (Right) is labeled HRGN instead of HGRN.
strength_weakness: strength_weakness: * The results aren't well-described and at times inconsistent (see above).
* The approach seems to not differentiate much from GLA after performing the new state expansion, which makes the contributions somewhat limited.
questions: questions: See above in summary.
paper_summary: paper_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. 

The 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. 

Not 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.
strength_weakness: strength_weakness: 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.

The 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: questions: 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.

Comments: title: Looking forward to your reply
Response
Looking forward to your reply
Response to Reviewer m3yg
Response to Reviewer m3yg
Response to Reviewer HShg
Response
Response to Rebuttal
Rebuttal by Authors
comment: Dear Reviewer m3yg,

Could you please let us know if our response has addressed your concerns? If you have any further questions, please feel free to raise them at any time.
Under the assumption that the above are added to the final version of the paper, including ideally a Mamba result, I will raise my score.
Dear Reviewer HShg,

Could you please let us know if our response has addressed your concerns? If you have any further questions, please feel free to raise them at any time.
We rechecked the experiment and made sure that full context was used, and listed the actual sequence length below.

| GovRep | SumScr | QMSUM | Qspr | QALT | CNLI |
|--------|--------|-------|------|------|------|
| 65k    | 23k    | 30k   | 21k  | 8.7k | 6.4K |
Thank you for your engaging and helpful suggestions! We will ensure that Table 2 is included in the next version of the paper. We are pleased to hear that our discussion has increased your confidence in our work. Could you kindly consider revising your score upwards in response?
Thanks for these results. This looks great! I think your Table 2 above should 100% be on the paper. I believe these results are more convincing than both Tables 5 and 6 in the original paper, mostly because these new results present the fairest comparison in terms of compute FLOPs and dataset quality. These numbers give me more confidence me that the method works.
Makes sense, yeah, I believe this one is pretty consistent with what I would expect for a truncated version. Were you able to double check that the initial version was indeed using the full context?

Thanks for providing these numbers!
As a supplement, we also tested the truncated version, and the results are as follows. It can be seen that the overall results are quite similar.

| Model      | GovRep           | SumScr           | QMSUM            | Qspr  | QALT  | CNLI  |
|------------|------------------|------------------|------------------|-------|-------|-------|
| HGRN2-385m | 15.33/3.54/10.91 | 7.35/0.76/6.17   | 8.32/1.22/7.4    | 12.36 | 26.37 | 31.53 |
| Llama-385m | 8.32/3.49/6.27   | 11.67/1.59/8.96  | 10.99/2.18/9.63  | 18.03 | 28.09 | 13.5  |
| HGRN2-1b   | 20.44/5.32/13.78 | 11.25/1.71/8.83  | 10.45/1.91/9.33  | 19.86 | 27.56 | 11.38 |
| Llama-1b   | 7.52/3.59/5.84   | 11.0/1.64/8.43   | 14.49/3.5/11.83  | 19.27 | 27.9  | 26.9  |
| HGRN2-3b   | 21.7/6.62/14.09  | 14.55/2.13/10.79 | 12.48/2.69/10.58 | 25.41 | 28.86 | 31.92 |
| Llama-3b   | 11.19/4.87/8.12  | 11.87/1.79/9.22  | 15.88/4.14/12.76 | 28.79 | 30.49 | 20.15 |
Q1. About Mamba result.

A1. According to our experimental results, on Wikitext-103, mamba outperforms Hgrn2; as the model size increases (to 3b scale), Hgrn2 outperforms mamba. We will add the results of mamba in future versions.

Q2. More about table 6.

A2. For the fairest comparison, we re trained llama and hgrn2 on the same dataset for 300b tokens, as shown in the follow table. Due to time constraints, we did not train mamba. However, based on the existing results, Hgrn2 outperformed Llama. Training of mamba will be conducted in future.

| Model | PIQA | HS | WG | ARC-E | ARC-C | OBQA | AVG |
| --- | --- | --- | --- | --- | --- | --- | --- |
| HGRN2-385m | 68.01 | 40.37 | 52.25 | 53.66 | 24.32 | 31.20 | 44.97 |
| Llama-385m | 67.19 | 38.75 | 52.17 | 49.24 | 23.72 | 30.00 | 43.51 |
| HGRN2-1b | 71.60 | 49.45 | 53.91 | 60.40 | 28.24 | 33.20 | 49.47 |
| Llama-1b | 69.97 | 47.04 | 52.72 | 57.07 | 28.16 | 32.60 | 47.93 |
| HGRN2-3b | 74.37 | 61.49 | 58.41 | 65.49 | 34.56 | 36.00 | 55.05 |
| Llama-3b | 73.18 | 57.88 | 59.59 | 63.93 | 31.40 | 34.00 | 53.33 |
| HGRN2-7b | 76.77 | 66.81 | 60.93 | 69.15 | 36.43 | 38.00 | 58.02 |
| Llama-7b | 75.19 | 64.39 | 61.88 | 67.55 | 35.41 | 35.00 | 56.57 |

Q3. Add context length to table 7.

A3. In the upcoming versions, we will add the sequence length in the table. For now, we provide the sequence length for each task.

| Data | ListOps | Text | Retrieval | Image | Pathfinder | Path-X |
| --- | --- | --- | --- | --- | --- | --- |
| Seqlen | 2k | 4k | 4k | 1k | 1k | 16k |

Q4. The reason for diminishing returns for the 256 expand ratio

A4. This may be related to the sequence length; our current sequence length is 512, and at this sequence length, 128 is likely sufficient for Wikitext-103.
* For Mamba, it appears to be better than HGRN2 on perplexity, is that right?  If you agree, then and put the Mamba result in the paper, that would be helpful.
* For table 6, what happens with more tokens (it looks like you only compared at 100B token size for Mamba).  Do the results hold at larger token sizes (e.g., 300B)?
* For table 7, please add context length to the paper.  Context length does matter, as linear models will do well with short context, but potentially not long context, and could be one of the advantages of approaches like these.
* For 256 expand ratio, why do you think there is diminishing returns?

If you can address these, I would consider raising my score.
Thanks for the response and the additional efforts. 

- Besides the technical component of state expansion itself, I agree that the study has novel aspects and I'd encourage the authors to emphasize them in the final version. 

- To support the claim regarding effectiveness and applicability in RNNs more broadly, I'd suggest providing empirical evidence different architectures such as LSTM, GRU, standard RNN. 

- Regarding the downstream performance, my concern has been addressed. It's encouraging that the improvement holds as the model size increases.
Hmm, in that case, then I find these SCROLLS results a bit strange. I believe the performance of Llama models would usually crash towards 0 if they are evaluated beyond the context length they were trained for (due to the positional encoding). Just to double-check: do you explicitly set the context length to be something very large like 100k? Otherwise, I believe lm-eval-harness might be doing truncation automatically to 2048 (or to the length that the model was trained for) (see [the max_length property here](https://github.com/EleutherAI/lm-evaluation-harness/blob/main/lm_eval/models/huggingface.py#L389-L400))

(I also need to double-check again. Feel free to correct me if I'm wrong)
Thanks for your comments.  We feed in the entire context without truncation using the lm-eval-harness framework.
Thanks for your comments. We will address your concerns as follows:

1. To fairly compare with other methods, we agree that we need to make sure we are using the same dataset with the same consumed tokens. Therefore, we report the results of retraining HGRN2, Mamba, and Llama on the same dataset with the same hyperparameters. The results are shown below:

| Method | PPL |
| --- | --- |
| Llama-3b | 7.129 |
| HGRN2-3b | 6.989 |
| Mamba-3b | 7.056 |

HGRN2 achieves lower ppl than Llama and Mamba in the controlled experiment. (In fact, Mamba is not a peer-reviewed method. Therefore, we do not need to compare with it.)

2. For downstream evals, we trained Llama and HGRN2 on the same dataset with 300B tokens sampled from the SlimPajama and Pile datasets. It can be seen that HGRN2 performs consistently better.

| Model | PIQA | HS | WG | ARC-E | ARC-C | OBQA | AVG |
| --- | --- | --- | --- | --- | --- | --- | --- |
| HGRN2-385m | 68.01 | 40.37 | 52.25 | 53.66 | 24.32 | 31.20 | 44.97 |
| Llama-385m | 67.19 | 38.75 | 52.17 | 49.24 | 23.72 | 30.00 | 43.51 |
| HGRN2-1b | 71.60 | 49.45 | 53.91 | 60.40 | 28.24 | 33.20 | 49.47 |
| Llama-1b | 69.97 | 47.04 | 52.72 | 57.07 | 28.16 | 32.60 | 47.93 |
| HGRN2-3b | 74.37 | 61.49 | 58.41 | 65.49 | 34.56 | 36.00 | 55.05 |
| Llama-3b | 73.18 | 57.88 | 59.59 | 63.93 | 31.40 | 34.00 | 53.33 |
| HGRN2-7b | 76.77 | 66.81 | 60.93 | 69.15 | 36.43 | 38.00 | 58.02 |
| Llama-7b | 75.19 | 64.39 | 61.88 | 67.55 | 35.41 | 35.00 | 56.57 |
Thanks for your responses. I think at the core of my concerns is really just trying to make sure that the method is something that's indeed competitive with other models out there. 

Some of my concerns have been addressed, but I still don't feel entirely convinced. 
- Figure 4 is good, but not enough. I think loss is a good indicator *if* paired with other downstream evals (which you indeed do, but see point below)
- I understand your point about other models being trained with 300B tokens and your model only being trained with 100B tokens, but this to me isn't very convincing. Just because your 100B-token model is worse but close to other 300B-token models, it does not necessarily imply that your model will be better when trained at 300B. For example, based on Table 7, HGRN2-3B performs quite poorly at some benchmarks such as Winogrande and ARC-C, as compared to other 3B models. I understand this is probably because it's trained on 100B tokens, but there's no guarantee that going to 300B will fully close those gaps.
- There's also the issue of datasets. If the models that you are comparing against were trained on datasets that are worse than the Pile, then it's not really a fair comparison -- your model will naturally perform better simply because of the dataset. For instance, I found a version of [RWKV4 which was trained on the Pile](https://huggingface.co/BlinkDL/rwkv-4-pile-3b) which gets 59.63 on HellaSwag, which is higher than the results that you reported in the paper for RWKV. I think as much as possible you should be comparing with models that have been trained on the same dataset. 

Other concerns:
- A5: "Our model only used 100 billion tokens, while the original Mamba 3B used 300 billion tokens, making the comparison unfair." --> But in the paper you seem to compare with other 3B-300B token models though (e.g. OPT, Pythia, Bloom), so how is Mamba any different?
I'm not Reviewer 9TbB, but just wanted to ask a quick question about the long-context evals here. A lot of the benchmarks in scrolls are long context. When you performed the evaluations, did you truncate to the context length of the trained model like what they did in the [LongLLaMA](https://arxiv.org/abs/2309.16039) paper? Or did you just feed in the entire context without truncation?