HGRN2: Gated Linear RNNs with State Expansion

Motivation

• Fixed recurrent states must encode all historical information, making effective state use and increased state size important for recurrent-model capacity.State-expansion techniques have been associated with improved language-modeling performance and retrieval ability.Source: §1 · pp. 1–2

Background

• Data-dependent decays selectively retain useful information and erase irrelevant information, and HGRN-based models use them to improve fixed-state storage efficiency.HGRN emphasized this mechanism for linear RNNs, followed by models including Mamba, GLA, Griffin, and RWKV-6.Source: §1 · p. 1

Contribution

• HGRN2 uses a non-parametric outer product-based state expansion to enlarge HGRN’s recurrent state without additional parameters.The design also gives HGRN2 a linear-attention interpretation that supports hardware-efficient training.Source: §1 · p. 2

Results

• HGRN2 is evaluated across multiple tasks and consistently outperforms HGRN1, while remaining highly competitive with other subquadratic efficient models in language modeling.The motivation arises from the cost of Transformer KV-cache management and the linear inference complexity of recurrent models.Source: §1 · pp. 1–2

2.1 Gated linear RNN

• A gated linear RNN transforms an input sequence into hidden states and outputs using forget, input, and output gates.The input gate is tied to the forget gate as 1 − g_t, while element-wise products update the hidden state.Source: §2.1 · p. 2

2.2 HGRN (Qin et al., 2023c)

• HGRN modifies the basic recurrence with complex-valued states and forget gates whose lower bounds increase monotonically from bottom to upper layers.The lower-bound schedule encourages lower layers to model short-term dependencies and upper layers to model long-term dependencies.Source: §2.2 · p. 2
• HGRN constructs layerwise forget-gate lower bounds with cumulative softmax and subtracts the first-layer value so the first lower bound is zero.The forget gate then learns residuals above these lower bounds.Source: §2.2 · pp. 2–3
• Real-valued HGRN with twice the state size performs better than complex HGRN on WikiText-103 language modeling.This comparison motivates using real-valued recurrence after accounting for the complex state’s doubled size.Source: §2.2 · p. 2, Table 1 · p. 3
• The additive lower-bound parameterization is reported to mitigate saturated gates.The superscript in the formulation indicates the layer index.Source: §2.2 · p. 3

3.1 Explorations of state expansion methods

• State expansion scales HGRN’s recurrent state from d to nd, but dense projections would become parameter-inefficient as the expansion ratio grows.The target is to keep projection parameters around d^2 per projection.Source: §3.1 · p. 3
• Structured state-expansion methods improve performance, with low-rank matrices performing best among the evaluated candidates.The comparison is reported in the PESE ablation trained on 10 billion Pile tokens.Source: §3.1 · p. 3, Table 3 · p. 4
• High-dimensional element-wise recurrence causes training inefficiency because it increases FLOPs and I/O costs without leveraging GPU tensor cores.The authors relate this constraint to Mamba1’s use of a relatively small expansion ratio and custom CUDA implementation.Source: §3.1 · p. 4

3.2 HGRN2

• HGRN2 replaces the input gate’s element-wise product with an outer product, stores a d × d hidden state, and uses matrix-vector multiplication for output projection.The recurrent update diagonalizes the forget gate and applies matrix dot products to update the expanded state.Source: §3.2 · p. 4
• Multihead HGRN2 reduces recurrence complexity from O(BNd2) to O(BNd2/H), making the effective state size d2/H.The main experiments use head dimension dh = 128 because gains diminish beyond 128.Source: §3.2 · p. 5
• HGRN2’s recurrence is identical to GLA’s apart from parameterization, enabling direct use of chunkwise algorithms and optimized kernels for hardware-efficient training.HGRN2 ties its input gate to the forget gate, while omitting GLA’s output gate, thereby saving parameters.Source: §3.2 · p. 5

4.1 MQAR

• HGRN2 significantly outperforms HGRN1 across MQAR model dimensions, demonstrating improved in-context recall with state expansion.MQAR evaluates associative recall, and the task becomes harder as sequence length increases.Source: §4.1 · p. 6, Figure 3 · p. 6

4.2.1 Wikitext-103

• HGRN2 clearly outperforms HGRN1 on Wikitext-103 but slightly underperforms Mamba under matched experimental settings.The comparison uses 44M-parameter models evaluated against a broad set of subquadratic models.Source: §4.2.1 · p. 6, Table 5 · p. 6

4.2.2 Slimpajama

• HGRN2 consistently outperforms other competitive linear recurrent models across three Slimpajama model scales.The experiments use 1.3B and 2.7B models and report results from the FLASHLINEARATTENTION codebase.Source: §4.2.2 · pp. 6–7, Table 6 · p. 7

4.2.3 The Pile

• On the Pile, HGRN2 outperforms HGRN1, competes with Mamba and LLaMA, and shows better SCROLLs scaling than Mamba but remains weaker than LLaMA at 7B.In the easy Needle in a Haystack test, both HGRN2 and Mamba face difficulties, although HGRN2 performs better.Source: §4.2.3 · pp. 7–8, Table 7 · p. 7, Table 8 · p. 8

4.3 Long Range Arena

• HGRN2 outperforms HGRN1 on Long Range Arena, while Mamba and Griffin fail to achieve high accuracy.Long Range Arena assesses the ability to handle long-range dependencies.Source: §4.3 · p. 9, Table 10 · p. 9

4.4 Image Modeling

• HGRN2 outperforms HGRN1 at similar parameter size and has an advantage over previous TNN and DeiT models on ImageNet-1k.The comparison is reported for image classification against TNN and the vanilla transformer.Source: §4.4 · p. 9, Table 11 · p. 9

Linear recurrent models

• Linear recurrent models include linear RNNs, state-space models, and linear attention, with gating improving state-space and language-modeling performance.Prior work connects state-space models and linear RNNs and scales linear RNNs to billions of parameters.Source: §5 · pp. 9–10
• Linear attention has narrowed its language-modeling gap through forgetting mechanisms, local attention, higher-order feature maps, and more expressive recurrent updates.These approaches target sharper attention distributions or stronger recurrent dynamics while retaining efficiency.Source: §5 · p. 10

Gated linear recurrence

• Gated linear recurrences emphasize data-dependent decay and parallel scan training, but existing models remain limited in recurrent state size.This limitation motivates expanding the recurrent state while preserving efficient recurrence.Source: §5 · p. 10
• Matrix-valued recurrent states and gated fast-weight models provide precedents for state expansion, but Neural Turing Machines are difficult to parallelize and slow to train.Linear Transformer recurrence and fast-weight programming connect matrix-valued states with hardware-efficient training research.Source: §5 · p. 10
• Mamba2, xLSTM, and Gated Retention share decay values within heads to improve tensor-core efficiency, trading away fine-grained per-dimension decays.GLA and HGRN2 use separate decay rates for each head dimension, whereas these newer models share one decay per head.Source: §5 · p. 10

• HGRN2 uses outer product-based state expansion inspired by linear attention to enhance HGRN while enabling efficient training.Source: §6 · p. 10
• Experiments across multiple tasks validate HGRN2’s advantages over HGRN1.Source: §6 · p. 10

• Autoregressive language modeling uses BPE with a 50,265-word vocabulary, sequence length 512, batch size 128, and 50,000 updates on WikiText-103.Source: §A.1 · p. 17
• ImageNet experiments use 6-million- and 23-million-parameter models with batch sizes of 1024 and 2048 for 300 epochs.Source: §A.1 · p. 17
• The configurations vary model complexity from 0.15 to 2.9 million parameters while adjusting layers, hidden dimensions, and GPU usage.Source: §A.1 · p. 17

• The 1B HGRN2 model consistently outperforms the 150M and 385M models in training loss.Source: §A.2 · p. 18, Figure 5 · p. 18
• All three HGRN2 models consume approximately 100 billion tokens, and loss decreases as model size increases.Source: Figure 5 · p. 18