HGRN2: Gated Linear RNNs with State Expansion

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.

Rating

6

Confidence

2

Ethics flag

1

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

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.

Rating

6

Confidence

3

Ethics flag

1

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). * 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.

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.

Rating

6

Confidence

4

Ethics flag

1

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

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

Rating

6

Confidence

3

Ethics flag

1

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. - 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. - 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). - 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 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? - 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)]

Response to Rebuttal

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.

Response

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

Response to Reviewer HShg

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.