Scaling Laws for RNN LLM in Long-Context Scenarios with State Size
Authors: Kai Liu, Jianfei Gao, Kai Chen
Accepted at: acl25

Scaling up the State Size of RNN LLMs for Long-Context Scenarios

Kai Liu1,2 and Jianfei Gao1 and Kai Chen1
1Shanghai AI Laboratory
2Tongji University
{liukai,gaojianfei,chenkai}@pjlab.org.cn

Abstract

The Transformer architecture has become the standard LLM architecture due to its powerful self-attention mechanism. However, it suffers from quadratic computational complexity and linear memory complexity. RNN-based LLMs have been proposed as alternatives. Yet, RNN models struggle in long-context scenarios, making it challenging to replace self-attention with RNNs. We identify the state size as a critical bottleneck, which is significantly smaller than that of Transformers with a basic context length of 2k. However, simply increasing the state size significantly raises the number of parameters and lowers training efficiency. In this paper, we propose an efficient scaling method to scale the state size of RNN models to match the 2k context length of Transformers, with small parameters overhead. Experimental results demonstrate that scaling the state size significantly enhances long-context understanding. Retrieval performance scales almost linearly with state size, with a 454M model featuring an expanded state achieving performance comparable to a 1.47B model on FDA, a recall-intensive task. These findings highlight state scaling as a promising approach for advancing RNN-based LLMs.

1 Introduction

Transformer-based Large Language Models (LLMs) (Achiam et al., 2023) have achieved state-of-the-art performance across various tasks and applications (Liu et al., 2023; Zhang et al., 2024). Their core self-attention mechanism (Vaswani, 2017) excels at capturing long-range dependencies and fine-grained past information. However, its quadratic computational complexity and linear memory requirements limit scalability. To overcome these challenges, RNN models (Peng et al., 2024; Gu and Dao) have been proposed as efficient alternatives, offering linear computational complexity and constant memory usage.

RNN models have recently made significant progress, achieving performance comparable to Transformers on tasks like language modeling and common-sense reasoning (Zuo et al.; Waleffe et al., 2024). This positions RNNs as a promising alternative to Transformers. However, they still struggle with long-context scenarios, particularly in recall-intensive tasks like needle-in-a-haystack (Hsieh et al., 2024), limiting their ability to replace self-attention. Improving long-context understanding remains a key challenge.

Recent RNN models have made notable advancements, particularly in state expansion (Peng et al., 2024; Qin et al.; Sun et al., 2023a). Early RNNs (Peng et al.; Qin et al., 2023) relied on vector states, which were severely limited in memory capacity. To overcome this limitation, newer models (Peng et al., 2024; Qin et al.) introduced matrix states, significantly enhancing memory capacity. Despite these improvements, the state sizes of RNNs remain much smaller than those of Transformers with a 2k context length, making state size a critical bottleneck for long-context understanding. Yet, simply increasing the state size significantly raises the number of parameters and lowers training efficiency, making it challenging to scale RNN models to match Transformers. We propose an efficient method to scale the state size of RNN models to match Transformers with a 2k context length, alleviating the challenges of increased parameter counts and reduced training efficiency.

We evaluate the performance of the scaled RNN models on various tasks, including language modeling, zero-shot common-sense reasoning, and long-context understanding tasks. The results demonstrate that scaling the state size of RNN models significantly enhances their long-context understanding ability. Additionally, we conduct experiments to investigate the scalability of RNN models by examining how increasing the state size affects the maximum recallable context length. Our findings
show that the maximum recallable context length scales nearly linearly with the RNN state size, providing a valuable reference for further scaling of RNN models. These results highlight scaling the state size of RNN models as a promising approach to improving their long-context understanding ability.

In conclusion, our contributions are as follows:

• We propose an efficient method to scale the state size of RNN models to match Transformers with a 2k context length, addressing the challenges of increased parameter counts and reduced training efficiency.

• We evaluate the performance of the scaled RNN models on various tasks, showing that scaling the state size significantly improves the long-context understanding ability of RNN models.

• We conduct experiments to investigate the scalability of RNN state size, providing a valuable reference for further scaling of RNN models.

2 Related Work

2.1 RNN LLMs

RNN LLMs have made significant progress in long-context understanding tasks. In this paper, we focus on the state expansion of RNN models. Early RNN models, such as RWKV4 (Peng et al.) and HGRN (Qin et al., 2023), utilized vector states, where states are updated at each step by a vector. However, this approach has limited memory capacity. Building on linear attention (Katharopoulos et al., 2020), matrix states have been widely adopted to replace vector states. In this approach, states are updated at each step by a matrix, which is the outer product of keys and values, significantly enhancing the memory capacity of RNN LLMs. Recent RNN models, such as RWKV5 (Peng et al., 2024), Mamba2 (Dao and Gu), and others, commonly employ matrix states. Nevertheless, the state size of RNN models remains much smaller than that of self-attention mechanisms with a basic context length of 2k. To the best of our knowledge, this is the first work to scale the state size of RNN LLMs to match Transformers with a 2k context length. We argue that there is considerable room to scale the state size of RNN models, which is crucial for advancing their performance in long-context understanding tasks.

In addition to state expansion, various other methods have been proposed to enhance the performance of RNN LLMs. Data-dependent gates (Yang et al.; Peng et al., 2024; Dao and Gu) enable models to selectively forget irrelevant information while retaining important information. The delta update rule (Yang et al., 2025, 2024) and test-time training (Sun et al.) reformulate recurrent updates as optimization problems, making state updates more efficient. Notably, our scaling method can be integrated with these techniques to further improve the performance of RNN LLMs.

Outside the scope of LLMs, some prior methods have attempted to upscale the state size of linear attention in a parameter-efficient manner. DPFP (Deterministic Parameter-Free Projection) (Schlag et al., 2021) leverages interactions between the features of queries and keys to expand their dimensions. While this approach avoids introducing additional parameters, it results in sparse features. Similarly, LFM (Learnable Feature Map) (Pramanik et al., 2024) proposes using the outer product to expand the state size, which is equivalent to repeating heads with different weights, but this also leads to reduced performance. We compare our method with these approaches in the experiments.

2.2 Scalability

In recent years, research has highlighted the importance of LLM scalability, with scaling laws widely used to analyze this property. (Kaplan et al., 2020; Hoffmann et al., 2022) conducted extensive experiments using dense curve fitting and regression analysis to examine the relationships among performance, model size, dataset size, and other factors, establishing their scaling law as a key reference for guiding LLM scaling. Broader scaling trend analyses, which do not rely on dense curves or regression analysis, have also been conducted to explore LLM scalability under different scenarios and targets, as seen in (Gu and Dao; Dao and Gu; Li et al., 2023).

In this paper, we focus on the scalability of state size in RNN LLMs, a dimension less explored than model size. We assess scalability through coarse scaling trends, given real-world limitations in obtaining precise scaling laws for RNN LLMs, as detailed in Section 6. Consistent with (Ye et al., 2024), we observe that scaling different model components affects performance in distinct ways, and increasing the state size of RNNs is particularly effective for improving performance on long-context understanding tasks.
3 Preliminaries

3.1 Attention
Attention has become the cornerstone of large language models, leveraging queries and keys to compute attention scores, which are then used to generate the output.

The output is computed using the following formula:

\[
\mathbf{O} = \operatorname{softmax}\left( \frac{\mathbf{QK}^T + \mathbf{M}}{\sqrt{d_{qk}}} \right) \mathbf{V}
\] (1)

where \( \mathbf{Q}, \mathbf{K}, \) and \( \mathbf{V} \) represent the queries, keys, and values, respectively. \( d_{qk} \) denotes the dimensionality of the queries and keys, and \( \mathbf{M} \) is the causal attention mask.

During the generation process, key and value caches/states are maintained continuously, causing the state size to grow proportionally with the context length. This mechanism allows the attention model to recall information from earlier contexts effectively. For a context length of \( L \), the average state size per batch during inference is \( L \times H \times (d_{qk} + d_v)/2 \), where \( H \) is the number of attention heads, and \( d_v \) is the dimensionality of the values.

3.2 Linear Attention
To address the quadratic complexity of traditional attention mechanisms, linear attention (Katharopoulos et al., 2020) was proposed. It eliminates the softmax operation and uses feature maps to process \( \mathbf{Q}, \mathbf{K}, \) and \( \mathbf{V} \), approximating the attention score. The computation takes the following form:

\[
\mathbf{O} = (\mathbf{QK}^T) \odot \mathbf{MV}
\] (2)

This can also be reformulated in an RNN-like structure:

\[
\begin{align*}
\mathbf{S}_t &= \mathbf{S}_{t-1} + \mathbf{k}_t^T \mathbf{v}_t = \sum_0^t \mathbf{k}_i^T \mathbf{v}_i, \\
\mathbf{o}_t &= \mathbf{q}_t \times \mathbf{S}_t, \quad \mathbf{S}_t \in \mathbb{R}^{[H \times d_{qk} \times d_v]}
\end{align*}
\] (3)

Here, \( \mathbf{S} \) represents the state matrix, while \( \mathbf{k}_t \) and \( \mathbf{v}_t \) denote the key vector and value vector, respectively. Linear attention reduces the quadratic complexity of self-attention to linear complexity and replaces the growing state size with a fixed-sized state. As a result, linear attention is more efficient than self-attention in terms of both inference speed and memory usage. The state size for linear attention per batch is \( H \times d_{qk} \times d_v \).

Assuming \( d_{qk} = d_v = d \), the ratio of the state sizes between self-attention and linear attention can be expressed as:

\[
\text{Ratio} = \frac{L \times H \times (d_{qk} + d_v)/2}{H \times d_{qk} \times d_v} = \frac{L}{d}
\] (4)

For \( d = 128 \) and \( L = 2048 \), the ratio is 16. This indicates that the state size of linear attention is 16 times smaller than that of self-attention, potentially making the reduced state size a bottleneck for long-context understanding tasks.

3.3 Recent Linear Attention Variants
In recent years, several linear attention variants have been proposed to improve the memory efficiency of linear attention. Among them, Gated Linear Attention (GLA) (Yang et al.) and Gated DeltaNet (Yang et al., 2024) are two notable examples with SOTA performance.

GLA introduces data-dependent gating mechanisms, similar to the forget gate used in LSTMs (Graves and Graves, 2012), to enhance the memory capabilities of linear attention. Its formulation is as follows:

\[
\mathbf{S}_t = \operatorname{diag}(\mathbf{g}_t) \odot \mathbf{S}_{t-1} + \mathbf{k}_t^T \mathbf{v}_t
\] (5)

where \( \mathbf{g}_t \) acts as a gate that enables the model to selectively forget irrelevant information and retain important information, thereby improving the memory retention of linear attention.

Gated DeltaNet builds upon this idea by incorporating the delta rule (Schlag et al.) to update the state matrix. The update rule is expressed as:

\[
\begin{align*}
\mathbf{S}_t &= \alpha_t \mathbf{S}_{t-1} + \beta_t \mathbf{k}_t^T (\mathbf{v}_t - \alpha_t \mathbf{k}_t \mathbf{S}_{t-1}) \\
&= \mathbf{S}_{t-1} (\alpha_t (\mathbf{I} - \beta_t \mathbf{k}_t^T \mathbf{k}_t)) + \beta_t \mathbf{k}_t^T \mathbf{v}_t
\end{align*}
\] (6)

From a test-time training perspective, the state \( \mathbf{S} \) is treated as a weight matrix optimized via online stochastic gradient descent (SGD) with the objective \( Loss(\mathbf{S}_t) = \frac{1}{2} \| \mathbf{k}_t \mathbf{S}_t - \mathbf{v}_t \|^2 \). In this context, \( \alpha_t \) functions as an adaptive weight decay term, while \( \beta_t \) serves as a learning rate. This delta-based update mechanism further strengthens the memory capabilities of linear attention.
Figure 1: Our method for scaling the state size of RNN models includes head-wise expansion and the division of heads into subheads, highlighted by the red dashed-line rectangles. This figure demonstrates the expansion process with an expansion factor of 2 applied to 2 heads.

4 Method

As discussed in the previous section, linear attention employs a fixed-size state matrix to store information from the past, achieving greater computational and memory efficiency compared to self-attention. However, this efficiency comes at a cost: accumulating information over time makes it more challenging for RNN-based LLMs to retrieve past information, limiting their long-context understanding ability. From Equation (3), assuming all keys are normalized, when \( \mathbf{q}_j = \mathbf{k}_j \) is used to retrieve \( \mathbf{v}_j \), the output becomes:

\[
\mathbf{o}_j = \mathbf{q}_j^T \mathbf{S}_t = \sum_{i=0}^{t} (\mathbf{q}_j \mathbf{k}_i^T) \mathbf{v}_i \\
= \mathbf{v}_j + \sum_{i=0, i \neq j}^{t} (\mathbf{q}_j \mathbf{k}_i^T) \mathbf{v}_i \tag{7}
\]

Without the softmax function to polarize attention scores, linear attention suffers from *attention dilution* (Qin et al., 2022). Ideally, all keys should be orthogonal to one another, allowing the attention score to focus entirely on the target key. One straightforward way to improve the model’s retrieval ability is to increase the dimensionality of the keys, which makes it more likely for keys to become orthogonal.

Furthermore, as previously mentioned, there is significant room to scale the state size, which is proportional to \( d_{qk} \). Therefore, scaling the state size with \( d_{qk} \) is a promising direction to enhance the long-context retrieval ability of RNN-based models. However, scaling the state size introduces two key challenges: an increased parameter count and reduced training efficiency. In the following sections, we propose our method to alleviate these challenges effectively. Our entire expansion method is summarized in Figure 1.

4.1 Head-wise Expansion for Parameter Efficiency

Directly scaling \( d_{qk} \) significantly increases the number of parameters. Specifically, \( \mathbf{W}_q, \mathbf{W}_k \in \mathbb{R}^{D \times H d_{qk}} \), and scaling \( d_{qk} \) by a factor of \( E \) results in a proportional \( E \)-fold increase in parameters. This approach is not parameter-efficient. For instance, in a 400M-parameter model, scaling \( d_{qk} \) by a factor of 8 would nearly double the total number of parameters.

To address this inefficiency, we propose a head-wise expansion method. Instead of naively scaling \( d_{qk} \) across all heads, we expand \( d_{qk} \) within each head individually. Queries and keys are first generated using the original \( \mathbf{W}_q \) and \( \mathbf{W}_k \), and then each head is expanded separately. For example, the query expansion process of the \( i \)th head is as follows:

\[
\mathbf{q}_i = \mathbf{x} \times \mathbf{W}_{q_i} \in \mathbb{R}^{1 \times d_{qk}}, \\
\mathbf{q}_i^{\text{expand}} = \mathbf{q}_i \times \mathbf{W}_{q_i}^{\text{expand}} \in \mathbb{R}^{1 \times Ed_{qk}}, \\
\mathbf{W}_{q_i}^{\text{expand}} \in \mathbb{R}^{d_{qk} \times Ed_{qk}}, \\
\mathbf{W}_{q_i} \in \mathbb{R}^{d_{\text{model}} \times d_{qk}}, \quad i \in [0, H] \tag{8}
\]

We summarize the parameter comparison in Table 1. Naive scaling introduces \( E \)-fold more parameters, whereas the head-wise expansion method introduces only \( (1 + \frac{E}{H}) \)-fold more parameters, making it significantly more parameter-efficient. When \( E = H \), the additional \( \mathbf{W}_{q_i}^{\text{expand}} \) has the same size as \( \mathbf{W}_q \). Since \( \mathbf{W}_q \) and \( \mathbf{W}_k \) constitute only a small portion of the total parameters, the overhead is limited—approximately 10%, which is acceptable. For example, in a 400M-parameter model with \( H = 8 \),
<table>
  <tr>
    <th></th>
    <th>Number of Parameters</th>
    <th>Scaling Ratio</th>
  </tr>
  <tr>
    <td>Baseline</td>
    <td>\( d_{\text{model}} \times (d_{qk} \times H) \)</td>
    <td>1 ×</td>
  </tr>
  <tr>
    <td>Naive</td>
    <td>\( d_{\text{model}} \times (d_{qk} \times H) \times E \)</td>
    <td>E ×</td>
  </tr>
  <tr>
    <td>Head-wise</td>
    <td>\( d_{\text{model}} \times (d_{qk} \times (H + E)) \)</td>
    <td>\( (1 + \frac{E}{H}) \times \)</td>
  </tr>
</table>

Table 1: Parameter comparison: The number of parameters for queries or keys using different methods. The parameter count is for a single layer of the model. Our head-wise expansion method is significantly more parameter-efficient than naive scaling.

<table>
  <tr>
    <th>Method</th>
    <th>Device</th>
    <th>E=8</th>
    <th>E=16</th>
  </tr>
  <tr>
    <td>Naive</td>
    <td>H800</td>
    <td>9569</td>
    <td>2933</td>
  </tr>
  <tr>
    <td>Subhead</td>
    <td>H800</td>
    <td>15722</td>
    <td>6554</td>
  </tr>
  <tr>
    <td>Naive</td>
    <td>RTX-3090</td>
    <td>OOM</td>
    <td>OOM</td>
  </tr>
  <tr>
    <td>Subhead</td>
    <td>RTX-3090</td>
    <td>4829</td>
    <td>2658</td>
  </tr>
</table>

Table 2: Training Efficiency Comparison. We compare the training throughput (tokens per second) of the naive and subhead division on two different GPUs. The subhead division significantly improves training efficiency when the state size is large. OOM indicates that the model runs out of shared memory. The speed is measured with 400M-parameter models, with training lengths of 4096 and 1024 on H800 and RTX-3090, respectively.

scaling \( d_{qk} \) by a factor of 8 results in an increase to 454M parameters, representing a manageable overhead.

4.2 Subhead Division for Training Efficiency

In existing RNN kernels (Yang et al., 2024), a single CUDA block is assigned to process a linear attention head and stores the states in shared memory instead of global memory. This approach significantly improves training efficiency. However, each CUDA block has access to only limited shared memory and computational resources. When the expansion factor \( E \) becomes large, the model runs slowly or may even exhaust the shared memory.

To overcome this limitation, we propose splitting a head into multiple subheads, running on multiple CUDA blocks, and then merging them back into a single head. Specifically, an expanded query or key head is divided into \( E \) subheads, with each subhead retaining the same size as the original head. Additionally, RNN-based LLMs commonly employ convolutional layers and feature maps to process queries and keys, which we denote as \( f \). We ensure that \( f \) operates on each subhead independently. The values (\( v_i \)) are shared across all subheads within the same head. The dimensions of other components, such as gates, \( \alpha \), and \( \beta \), are expanded to align with the subheads, as their associated parameters are relatively small. We find that the expanded gates, \( \alpha \), and \( \beta \) provide finer control over the state update, which is important for the model’s performance. The output for each subhead is computed independently, and the final output of the head is obtained by summing the outputs of all subheads:

\[
\mathbf{d}_{ij}^{\text{expand}} = \mathbf{d}_i^{\text{expand}}[:,j:d_{qk}:(j+1)-d_{qk}] \in \mathbb{R}^{1 \times d_{qk}}
\]
\[
\mathbf{o}_i = \sum_{j=0}^{E} \mathrm{RNN}(f(\mathbf{q}_{ij}^{\text{expand}}), f(\mathbf{k}_{ij}^{\text{expand}}), \mathbf{v}_i)
\]

We compare our subhead division with the naive method in terms of training efficiency, as shown in Table 2. Our subhead division doubles the training throughput on H800 when \( E = 16 \). Additionally, it enables the model to run efficiently on less resource-intensive GPUs, such as the RTX-3090.

4.3 Enhancing Nonlinear Capacity

Building on the head-wise expansion method, we further enhance the model’s nonlinear capacity by incorporating a SiLU activation function (Elfwing et al., 2017) prior to the head-wise expansion. Our experiments demonstrate that the SiLU activation function improves the model’s retrieval ability, enabling it to capture more complex patterns.

5 Experiment

5.1 Experimental Setting

We scale the state size of two RNN models, Gated-DeltaNet (Yang et al., 2024) and GLA (Yang et al.), under two configurations: approximately 400M and 1.5B parameters. Both models are derived from Transformer++ (Touvron et al., 2023), replacing the self-attention mechanism with linear attention variants. All base models consist of 24 layers. For the 400M-parameter models, the model dimension is set to 1024 with \( d_{qk} = d_v = 128 \), while for the 1.5B-parameter models, the model dimension is 2048 with \( d_{qk} = d_v = 256 \).

For the 400M-parameter models, we use expansion factors of 2, 4, 8, and 16, corresponding to state sizes equivalent to those of a Transformer with context lengths of 256, 512, 1024, and 2048, respectively. For the 1.5B-parameter models, the expansion factors are 2, 4, and 8, corresponding to state sizes equivalent to those of a Transformer with context lengths of 512, 1024, and 2048, respectively. Additionally, we compare the effects of scaling the
number of layers versus scaling the state size while maintaining the same total parameter count. In the tables below, “-L” indicates scaling the number of layers, while “-E” indicates scaling the state size. We provide more details in Appendix A.3.

5.2 Common Sense Reasoning

First, we evaluate the performance of scaled RNN models and Transformer++ (Touvron et al., 2023) on language modeling and zero-shot common-sense reasoning tasks. While improving performance on these tasks is not our primary objective, increasing the state size of RNNs yields modest gains. With a comparable number of parameters, scaling the state size achieves performance similar to scaling the number of layers. At the 400M parameter scale, Gated DeltaNet-E8 outperforms Gated DeltaNet-L by an average of 0.2%.

5.3 Real-World Recall-Intensive Tasks

We evaluate the performance of scaled RNN models and Transformer++ on real-world retrieval tasks, including FDA (Arora et al., 2023), SWDE (Lockard et al., 2019), and SQuAD (Rajpurkar et al., 2018), with maximum context lengths of 2723, 2222, and 928, respectively.

Finding 1: Scaling the state size of RNNs significantly enhances recall ability, especially for long-context tasks. Our results demonstrate that increasing the state size of RNNs substantially improves recall performance, particularly for tasks with longer contexts. Expanding the state size of RNN models from \(2 \times\) to \(8 \times\) consistently yields performance gains. At the 1.5B scale, Gated DeltaNet-E8 outperforms the base Gated DeltaNet by an average of 11.73%, highlighting the effectiveness of state size scaling in enhancing long-context retrieval ability.

In contrast, scaling the number of layers has a more limited impact. At the same 1.5B scale, Gated DeltaNet-L achieves only a 3.8% improvement on average compared to the baseline model. Gated DeltaNet-E8 with 454M parameters achieves better performance than Gated DeltaNet with 1.47B parameters on FDA. Furthermore, the performance gains across the three tasks follow the order FDA > SWDE > SQuAD, indicating that scaled RNN models are particularly effective for tasks involving longer contexts.

Finding 2: The update rule plays a critical role in recall ability scaling. At the 400M scale, Gated DeltaNet-E8 achieves a 9.36% improvement, whereas GLA achieves only a 3.46% average gain. This underscores the importance of the update rule in scaling retrieval ability. Specifically, the delta rule employed by Gated DeltaNet proves more effective than the gated linear attention mechanism used in GLA.

5.4 Long Context Understanding

To further assess the performance of scaled RNN models, we evaluate Gated DeltaNet with 1.5B parameters on LongBench (Bai et al., 2024), a specialized benchmark designed for long-context understanding. The maximum context length in LongBench is up to 64k tokens. The results of the scaled RNN models are presented in Table 5. Our findings indicate that expanding the state size of RNN models consistently enhances performance across tasks. Specifically, Gated DeltaNet-E8 achieves an average improvement of 5.74% compared to the baseline model and outperforms Gated DeltaNet-L by 3.16% on average. Furthermore, due to the limited extrapolation capabilities of Transformer++, its performance on LongBench is inferior to that of Gated DeltaNet. We further compare with Transformer++-2k, which is a Transformer++ model using sliding window attention with a window size of 2048. The results show that Gated DeltaNet-E8 outperforms Transformer++-2k by 2.79% on average with a smaller state size\(^1\).

5.5 State-Recall Context Scaling Trend

To quantify the scalability of RNNs on long-context tasks, we propose the state-recall context scaling trend. The maximum recall context length is defined as the longest context length at which RNNs maintain recall accuracy above a specified threshold. To evaluate this, we test our models on S-NIAH (Single Needle-In-A-Haystack) (Hsieh et al., 2024) with UUIDs as needle values. We measure accuracy with progressively increasing haystack context lengths, using a step size of 512. Thresholds are set between 0.5 and 0.8, incremented by 0.1. The resulting maximum recall context lengths for various state sizes are presented in Figure 2.

Finding 3: The maximum recall context length of RNNs scales linearly with state size. The results in the figure reveal a near-linear relationship between the maximum recall context length and the state size, indicating that the recall ability of RNNs can be significantly improved by scaling the

\(^1\)The computation of the state size for Transformer++ with sliding window attention is detailed in Appendix A.4.
<table>
  <tr>
    <th>Model</th>
    <th>State</th>
    <th>Params</th>
    <th>Wiki ↓</th>
    <th>LMB ↓</th>
    <th>LMB ↑</th>
    <th>PIQA ↑</th>
    <th>Hella ↑</th>
    <th>Wino ↑</th>
    <th>ARC-e ↑</th>
    <th>ARC-c ↑</th>
    <th>SiQA ↑</th>
    <th>BoolQ ↑</th>
    <th>Avg ↑</th>
  </tr>
  <tr>
    <td>Transformer++</td>
    <td>/</td>
    <td>400M</td>
    <td>24.24</td>
    <td>32.49</td>
    <td>34.64</td>
    <td>64.69</td>
    <td>30.97</td>
    <td>52.33</td>
    <td>45.79</td>
    <td>18.94</td>
    <td>36.23</td>
    <td>55.29</td>
    <td>42.36</td>
  </tr>
  <tr>
    <td>Transformer++</td>
    <td>/</td>
    <td>1.47B</td>
    <td>14.31</td>
    <td>9.74</td>
    <td>52.53</td>
    <td>72.36</td>
    <td>43.55</td>
    <td>59.12</td>
    <td>60.14</td>
    <td>26.79</td>
    <td>40.43</td>
    <td>61.41</td>
    <td>52.04</td>
  </tr>
  <tr>
    <td>GLA</td>
    <td>1×</td>
    <td>400M</td>
    <td>26.32</td>
    <td>36.60</td>
    <td>31.54</td>
    <td>63.66</td>
    <td>30.85</td>
    <td>49.41</td>
    <td>45.12</td>
    <td>19.62</td>
    <td>37.67</td>
    <td>59.91</td>
    <td>42.22</td>
  </tr>
  <tr>
    <td>GLA-L</td>
    <td>1.2×</td>
    <td>455M</td>
    <td>25.41</td>
    <td>32.99</td>
    <td>33.98</td>
    <td>65.67</td>
    <td>31.45</td>
    <td>52.64</td>
    <td>47.22</td>
    <td>20.39</td>
    <td>37.72</td>
    <td>60.21</td>
    <td>43.66</td>
  </tr>
  <tr>
    <td>GLA-E2</td>
    <td>2×</td>
    <td>412M</td>
    <td>26.34</td>
    <td>35.56</td>
    <td>33.28</td>
    <td>64.69</td>
    <td>31.07</td>
    <td>51.30</td>
    <td>44.95</td>
    <td>19.20</td>
    <td>36.39</td>
    <td>58.78</td>
    <td>42.46</td>
  </tr>
  <tr>
    <td>GLA-E4</td>
    <td>4×</td>
    <td>427M</td>
    <td>26.07</td>
    <td>32.53</td>
    <td>34.43</td>
    <td>64.42</td>
    <td>31.04</td>
    <td>49.96</td>
    <td>45.54</td>
    <td>19.28</td>
    <td>37.31</td>
    <td>61.07</td>
    <td>42.88</td>
  </tr>
  <tr>
    <td>GLA-E8</td>
    <td>8×</td>
    <td>454M</td>
    <td>25.99</td>
    <td>34.74</td>
    <td>32.89</td>
    <td>63.93</td>
    <td>31.52</td>
    <td>51.30</td>
    <td>45.54</td>
    <td>18.77</td>
    <td>36.80</td>
    <td>59.08</td>
    <td>42.48</td>
  </tr>
  <tr>
    <td>Gated DeltaNet</td>
    <td>1×</td>
    <td>400M</td>
    <td>25.25</td>
    <td>36.20</td>
    <td>31.40</td>
    <td>64.31</td>
    <td>31.20</td>
    <td>51.85</td>
    <td>45.03</td>
    <td>18.43</td>
    <td>36.90</td>
    <td>57.61</td>
    <td>42.09</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-L</td>
    <td>1.2×</td>
    <td>455M</td>
    <td>24.17</td>
    <td>33.60</td>
    <td>32.97</td>
    <td>64.80</td>
    <td>32.17</td>
    <td>51.93</td>
    <td>46.84</td>
    <td>19.88</td>
    <td>38.02</td>
    <td>58.62</td>
    <td>43.16</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E2</td>
    <td>2×</td>
    <td>413M</td>
    <td>24.38</td>
    <td>31.93</td>
    <td>34.76</td>
    <td>66.21</td>
    <td>31.73</td>
    <td>50.12</td>
    <td>48.40</td>
    <td>19.03</td>
    <td>37.21</td>
    <td>58.72</td>
    <td>43.27</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E4</td>
    <td>4×</td>
    <td>427M</td>
    <td>23.96</td>
    <td>32.74</td>
    <td>32.74</td>
    <td>65.61</td>
    <td>32.24</td>
    <td>52.49</td>
    <td>47.56</td>
    <td>19.11</td>
    <td>36.64</td>
    <td>58.46</td>
    <td>43.27</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E8</td>
    <td>8×</td>
    <td>454M</td>
    <td>23.61</td>
    <td>29.50</td>
    <td>36.85</td>
    <td>65.29</td>
    <td>32.30</td>
    <td>50.36</td>
    <td>46.46</td>
    <td>19.28</td>
    <td>37.82</td>
    <td>58.53</td>
    <td>43.36</td>
  </tr>
  <tr>
    <td>Gated DeltaNet</td>
    <td>1×</td>
    <td>1.47B</td>
    <td>17.27</td>
    <td>14.57</td>
    <td>44.89</td>
    <td>71.06</td>
    <td>40.28</td>
    <td>56.75</td>
    <td>56.31</td>
    <td>23.81</td>
    <td>40.38</td>
    <td>61.22</td>
    <td>49.34</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-L</td>
    <td>1.2×</td>
    <td>1.68B</td>
    <td>16.42</td>
    <td>10.37</td>
    <td>50.75</td>
    <td>72.63</td>
    <td>41.95</td>
    <td>57.77</td>
    <td>58.63</td>
    <td>26.71</td>
    <td>39.71</td>
    <td>60.49</td>
    <td>51.08</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E2</td>
    <td>2×</td>
    <td>1.52B</td>
    <td>16.45</td>
    <td>10.65</td>
    <td>50.75</td>
    <td>72.91</td>
    <td>41.72</td>
    <td>57.46</td>
    <td>59.24</td>
    <td>26.11</td>
    <td>38.69</td>
    <td>61.62</td>
    <td>51.07</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E4</td>
    <td>4×</td>
    <td>1.57B</td>
    <td>15.99</td>
    <td>10.80</td>
    <td>50.16</td>
    <td>72.20</td>
    <td>41.54</td>
    <td>58.09</td>
    <td>59.47</td>
    <td>24.57</td>
    <td>40.28</td>
    <td>60.52</td>
    <td>50.85</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E8</td>
    <td>8×</td>
    <td>1.68B</td>
    <td>15.66</td>
    <td>10.87</td>
    <td>50.63</td>
    <td>72.52</td>
    <td>42.03</td>
    <td>58.96</td>
    <td>59.64</td>
    <td>25.34</td>
    <td>39.82</td>
    <td>59.05</td>
    <td>51.00</td>
  </tr>
</table>

Table 3: Performance Comparison on Language Modeling and Zero-shot Common-sense Reasoning. “L” denotes scaling the number of layers, and “E” denotes scaling the state size. “State” indicates the state size with respect to corresponding base models. The upward arrow (↑) indicates that higher values are better, while the downward arrow (↓) indicates that lower values are better.

<table>
  <tr>
    <th>Model</th>
    <th>State</th>
    <th>Params</th>
    <th>FDA ↑</th>
    <th>SWDE ↑</th>
    <th>SQuAD ↑</th>
    <th>Avg ↑</th>
  </tr>
  <tr>
    <td>Transformer++</td>
    <td>/</td>
    <td>400M</td>
    <td>72.96</td>
    <td>64.00</td>
    <td>36.96</td>
    <td>57.97</td>
  </tr>
  <tr>
    <td>Transformer++</td>
    <td>/</td>
    <td>1.47B</td>
    <td>77.77</td>
    <td>79.93</td>
    <td>48.26</td>
    <td>68.65</td>
  </tr>
  <tr>
    <td>GLA</td>
    <td>1×</td>
    <td>400M</td>
    <td>11.52</td>
    <td>31.95</td>
    <td>28.32</td>
    <td>23.93</td>
  </tr>
  <tr>
    <td>GLA-L</td>
    <td>1.2×</td>
    <td>455M</td>
    <td>12.34</td>
    <td>31.68</td>
    <td>29.93</td>
    <td>24.65</td>
  </tr>
  <tr>
    <td>GLA-E2</td>
    <td>2×</td>
    <td>412M</td>
    <td>15.52</td>
    <td>28.98</td>
    <td>29.12</td>
    <td>24.54</td>
  </tr>
  <tr>
    <td>GLA-E4</td>
    <td>4×</td>
    <td>427M</td>
    <td>17.06</td>
    <td>32.94</td>
    <td>28.28</td>
    <td>26.09</td>
  </tr>
  <tr>
    <td>GLA-E8</td>
    <td>8×</td>
    <td>450M</td>
    <td>19.06</td>
    <td>34.47</td>
    <td>28.65</td>
    <td>27.39</td>
  </tr>
  <tr>
    <td>Gated DeltaNet</td>
    <td>1×</td>
    <td>400M</td>
    <td>22.50</td>
    <td>32.94</td>
    <td>28.65</td>
    <td>28.03</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-L</td>
    <td>1.2×</td>
    <td>455M</td>
    <td>20.69</td>
    <td>36.27</td>
    <td>28.35</td>
    <td>28.44</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E2</td>
    <td>2×</td>
    <td>413M</td>
    <td>21.14</td>
    <td>38.34</td>
    <td>29.79</td>
    <td>29.76</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E4</td>
    <td>4×</td>
    <td>427M</td>
    <td>34.85</td>
    <td>41.49</td>
    <td>31.23</td>
    <td>33.86</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E8</td>
    <td>8×</td>
    <td>454M</td>
    <td>38.93</td>
    <td>43.02</td>
    <td>30.23</td>
    <td>37.39</td>
  </tr>
  <tr>
    <td>Gated DeltaNet</td>
    <td>1×</td>
    <td>1.47B</td>
    <td>37.66</td>
    <td>53.47</td>
    <td>37.60</td>
    <td>42.91</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-L</td>
    <td>1.2×</td>
    <td>1.69B</td>
    <td>41.74</td>
    <td>60.31</td>
    <td>38.07</td>
    <td>46.71</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E2</td>
    <td>2×</td>
    <td>1.52B</td>
    <td>46.01</td>
    <td>60.40</td>
    <td>39.71</td>
    <td>48.71</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E4</td>
    <td>4×</td>
    <td>1.57B</td>
    <td>52.63</td>
    <td>61.48</td>
    <td>39.51</td>
    <td>51.21</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E8</td>
    <td>8×</td>
    <td>1.68B</td>
    <td>58.53</td>
    <td>66.16</td>
    <td>39.24</td>
    <td>54.64</td>
  </tr>
</table>

Table 4: Performance Comparison on Real-World Retrieval Tasks: FDA, SWDE, and SQuAD.

state size. Furthermore, compared to GLA, Gated DeltaNet exhibits a steeper slope, highlighting its greater potential for state size scaling. We further evaluate the scaling trend on 1.5B-parameter models in Figure 3, which is consistent with our findings on 400M-parameter models.

5.6 Comparison with Transformer under Similar State Sizes

To investigate whether scaling the state size of RNNs can make them competitive with Transformer models, we compare the performance of RNNs and Transformer++ with similar state sizes. Since the state size of a Transformer is proportional to its context length, we control the Transformer’s state size by adjusting the context length in the tasks. For this comparison, we use a more challenging benchmark, M-NIAH (Multi-keys Needle-In-A-Haystack) with numbers as needle values (Hsieh et al., 2024), with context lengths of 512, 1024, and 2048. Specifically, we evaluate the performance of Gated DeltaNet and Transformer++ at the 1.5B parameter scale. Gated DeltaNet, with expansion factors of 2, 4, and 8, has state sizes comparable to those of a Transformer with context lengths of 512, 1024, and 2048, respectively. The results are provided in Table 6.

Finding 4: Gated DeltaNet remains less com-
<table>
  <tr>
    <th rowspan="2">Model</th>
    <th rowspan="2">State</th>
    <th colspan="3">Single-Doc QA ↑</th>
    <th colspan="3">Multi-Doc QA ↑</th>
    <th colspan="3">Summarization ↑</th>
    <th colspan="3">Few-shot ↑</th>
    <th colspan="2">Code ↑</th>
    <th rowspan="2">Avg ↑</th>
  </tr>
  <tr>
    <th>NQA</th><th>QQA</th><th>MFQ</th>
    <th>HQA</th><th>2WM</th><th>Mus</th>
    <th>GvR</th><th>QMS</th><th>MNs</th>
    <th>TRC</th><th>TQA</th><th>SSM</th>
    <th>LCC</th><th>RBP</th>
  </tr>
  <tr>
    <td>Transformer++</td><td>/</td><td>0.35</td><td>6.77</td><td>12.58</td><td>1.36</td><td>5.66</td><td>0.38</td><td>8.14</td><td>3.17</td><td>11.72</td><td>17.50</td><td>12.49</td><td>7.77</td><td>45.00</td><td>19.52</td><td>10.89</td>
  </tr>
  <tr>
    <td>Transformer++-2k</td><td>16x</td><td>5.38</td><td>6.05</td><td>12.74</td><td>7.70</td><td>9.46</td><td>3.21</td><td>9.02</td><td>14.47</td><td>6.71</td><td>35.00</td><td>44.62</td><td>27.69</td><td>52.05</td><td>48.37</td><td>20.18</td>
  </tr>
  <tr>
    <td>Gated DeltaNet</td><td>1x</td><td>2.46</td><td>5.78</td><td>13.67</td><td>5.81</td><td>7.70</td><td>3.71</td><td>7.06</td><td>16.71</td><td>12.16</td><td>28.00</td><td>32.16</td><td>26.97</td><td>40.83</td><td>38.20</td><td>17.23</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-L</td><td>1.2x</td><td>6.18</td><td>6.65</td><td>13.53</td><td>6.92</td><td>9.13</td><td>2.76</td><td>10.27</td><td>16.38</td><td>14.81</td><td>40.50</td><td>36.84</td><td>31.86</td><td>43.48</td><td>37.98</td><td>19.81</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E2</td><td>2x</td><td>7.29</td><td>6.80</td><td>13.66</td><td>6.40</td><td>8.96</td><td>3.88</td><td>11.09</td><td>16.85</td><td>10.20</td><td>48.50</td><td>30.60</td><td>31.11</td><td>41.61</td><td>40.49</td><td>19.82</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E4</td><td>4x</td><td>6.90</td><td>6.45</td><td>13.79</td><td>7.84</td><td>9.97</td><td>4.45</td><td>13.67</td><td>16.40</td><td>10.16</td><td>61.00</td><td>23.69</td><td>30.78</td><td>44.68</td><td>41.42</td><td>20.80</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E8</td><td>8x</td><td>4.11</td><td>6.94</td><td>16.71</td><td>7.98</td><td>10.61</td><td>3.78</td><td>11.07</td><td>16.47</td><td>18.11</td><td>67.50</td><td>41.88</td><td>33.95</td><td>43.09</td><td>39.39</td><td>22.97</td>
  </tr>
</table>

Table 5: Performance Comparison on 14 Tasks from LongBench (Bai et al., 2024). Transformer++-2k is a Transformer++ model using sliding window attention with a window size of 2048.

![Scaling Trend of Maximum Recall Context Length vs. State Size on S-NIAH](page_101_624_627_312.png)

Figure 2: Scaling Trend of Maximum Recall Context Length vs. State Size on S-NIAH. The X-axis represents the state size expansion factor, while the Y-axis denotes the maximum recall context length at which the model maintains recall accuracy above a specified threshold \( t \).

<table>
  <tr>
    <th>Model</th>
    <th>512 ↑</th>
    <th>1024 ↑</th>
    <th>2048 ↑</th>
    <th>AVG ↑</th>
  </tr>
  <tr>
    <td>Transformer++</td>
    <td>100.00</td>
    <td>99.60</td>
    <td>99.20</td>
    <td>99.60</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E2</td>
    <td>68.20</td>
    <td>24.40</td>
    <td>4.00</td>
    <td>32.20</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E4</td>
    <td>65.00</td>
    <td>30.20</td>
    <td>7.20</td>
    <td>34.13</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E8</td>
    <td>94.20</td>
    <td>54.80</td>
    <td>20.80</td>
    <td>56.60</td>
  </tr>
</table>

Table 6: Comparison of Transformer and Gated DeltaNet on M-NIAH: We compare Transformer++ and Gated DeltaNet models with similar state sizes. For the Transformer, the state size is controlled by adjusting the context length ("512", "1024", "2048"). Underlined values indicate that the Gated DeltaNet has the same state size as the Transformer for the corresponding context length.

competitive than Transformer models in terms of recall ability, even with similar state sizes. The results demonstrate that Gated DeltaNet, even when matched in state size, underperforms compared to Transformer models. This highlights the need for further advancements in memory efficiency and the development of new mechanisms to narrow the retrieval ability gap between RNNs and Transformers.

5.7 Ablation Study

We conduct an ablation study to evaluate the contributions of the key components in our proposed methods. 1) We replace the head-wise expansion with other efficient alternatives, including low-rank expansion, DPFP (Schlag et al., 2021), and LFM (Pramanik et al., 2024). 2) We evaluate the impact of removing the subhead mechanism. 3) We assess the effect of replacing the SiLU activation function with ReLU, 1+ELU, or identity.

The results are summarized in Table 7. The results indicate that the head-wise expansion mechanism preserves more information compared to other methods. The subhead division enhances performance by providing finer control over state updates. Additionally, the SiLU activation function improves the model’s performance by introducing greater nonlinear capacity.

Additionally, we compare our method with a naive approach of directly increasing the number of heads, as shown in Table 8. Given the same state size, increasing the number of heads does improve performance, but it also significantly increases the number of parameters, thereby constraining further
<table>
  <tr>
    <th>Model</th>
    <th>Wiki ↓</th>
    <th>Recall ↑</th>
  </tr>
  <tr>
    <td>Gated DeltaNet-E8</td>
    <td>23.61</td>
    <td>37.39</td>
  </tr>
  <tr>
    <td>Head-wise expand -> Low Rank</td>
    <td>23.93</td>
    <td>36.71</td>
  </tr>
  <tr>
    <td>Head-wise expand -> DPFP</td>
    <td>24.41</td>
    <td>33.93</td>
  </tr>
  <tr>
    <td>Head-wise expand -> LFM</td>
    <td>24.59</td>
    <td>34.52</td>
  </tr>
  <tr>
    <td>w/o Subhead</td>
    <td>25.16</td>
    <td>29.84</td>
  </tr>
  <tr>
    <td>SiLU -> ReLU</td>
    <td>23.61</td>
    <td>36.48</td>
  </tr>
  <tr>
    <td>SiLU -> 1+ELU</td>
    <td>24.06</td>
    <td>26.86</td>
  </tr>
  <tr>
    <td>SiLU -> Identity</td>
    <td>23.71</td>
    <td>36.29</td>
  </tr>
</table>

Table 7: Ablation Study on Gated DeltaNet-E8: The table summarizes the impact of replacing the head-wise expansion, subhead division, and SiLU activation function with alternatives. “Recall” represents the average performance across FDA, SWDE, and SQuAD.

<table>
  <tr>
    <th>Model</th>
    <th>State</th>
    <th>Param</th>
    <th>FDA ↑</th>
    <th>SWDE ↑</th>
    <th>SQuAD ↑</th>
    <th>Avg ↑</th>
  </tr>
  <tr>
    <td>Gated DeltaNet-H2</td>
    <td>2×</td>
    <td>526M</td>
    <td>29.76</td>
    <td>39.69</td>
    <td>31.57</td>
    <td>33.67</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E2</td>
    <td>2×</td>
    <td>413M</td>
    <td>21.14</td>
    <td>38.34</td>
    <td>29.79</td>
    <td>29.76</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E8</td>
    <td>8×</td>
    <td>454M</td>
    <td><b>38.93</b></td>
    <td><b>43.02</b></td>
    <td>30.23</td>
    <td><b>37.39</b></td>
  </tr>
</table>

Table 8: Comparison with vanilla increasing the number of heads. “H2” indicates that the number of heads is directly increased by 2 times.

scalability. In contrast, our Gated DeltaNet-E8 achieves better performance with fewer parameters than simply increasing the number of heads. We further compare our models with additional RNN variants, including Mamba, RetNet, and DeltaNet in Appendix A.2. This highlights the superiority of our scaled RNN models in terms of recall ability.

6 Conclusion

To address the limitations of RNNs in handling long-context tasks, we propose an efficient method to scale the state size of RNNs, aiming to match the state size of Transformers with a 2k context length. Our results reveal that the recall ability of RNNs scales nearly linearly with state size, indicating that enlarging the state size of RNNs is a promising approach to improving their long-context understanding ability. We hope this work inspires further research into developing more powerful and efficient RNN models, particularly by scaling their state size effectively.

Limitations

Although we propose subhead division to improve training efficiency, our expanded models still train more slowly than the baseline models. However, as discussed in Appendix A.1, this limitation can be mitigated by developing more efficient RNN kernels, as the FLOPs of our scaled RNN models are comparable to those of self-attention. We leave this exploration for future work.

We do not conduct a precise scaling law analysis in this work, instead focusing on a coarse scaling trend. There are several reasons. First, our coarse scaling trend sufficiently demonstrates the effectiveness of increasing RNN state size for long-context understanding tasks, which meets the needs of this paper. Second, evaluating long-context abilities is sensitive to factors such as evaluation methods, training data, and prompts, making it difficult to obtain a smooth scaling curve comparable to those seen in loss-related scaling laws. Third, training RNN LLMs with large state sizes is slow, limiting the number of experiments we can perform. It’s hard to solve these issues within the scope of this work. Therefore, we leave a precise scaling law analysis for future work.

References

Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. 2023. Gpt-4 technical report. arXiv preprint arXiv:2303.08774.

Simran Arora, Brandon Yang, Sabri Eyuboglu, Avanika Narayan, Andrew Hojel, Immanuel Trummer, and Christopher Ré. 2023. Language Models Enable Simple Systems for Generating Structured Views of Heterogeneous Data Lakes.

Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. 2024. Longbench: A bilingual, multitask benchmark for long context understanding. Preprint, arXiv:2308.14508.

Yonatan Bisk, Rowan Zellers, Ronan LeBras, Jianfeng Gao, and Yejin Choi. 2020. PIQA: reasoning about physical commonsense in natural language. In The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, The Thirty-Second Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, February 7-12, 2020, pages 7432–7439. AAAI Press.

Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. 2019. BoolQ: Exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics:
Human Language Technologies, Volume 1 (Long and Short Papers), pages 2924–2936, Minneapolis, Minnesota. Association for Computational Linguistics.

Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. 2018. Think you have solved question answering? try arc, the ai2 reasoning challenge. ArXiv preprint, abs/1803.05457.

XTuner Contributors. 2023. Xtuner: A toolkit for efficiently fine-tuning llm. https://github.com/InternLM/xtuner.

Tri Dao. 2023. Flashattention-2: Faster attention with better parallelism and work partitioning. Preprint, arXiv:2307.08691.

Tri Dao and Albert Gu. Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. Preprint, arXiv:2405.21060.

Stefan Elfwing, Eiji Uchibe, and Kenji Doya. 2017. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Preprint, arXiv:1702.03118.

Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noach, Haonan Li, Kyle McDonell, Niklas Muenninghoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. 2024. A framework for few-shot language model evaluation.

Alex Graves and Alex Graves. 2012. Long short-term memory. Supervised sequence labelling with recurrent neural networks, pages 37–45.

Albert Gu and Tri Dao. Mamba: Linear-Time Sequence Modeling with Selective State Spaces. Preprint, arXiv:2312.00752.

Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. 2022. Training compute-optimal large language models. Preprint, arXiv:2203.15556.

Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. 2024. Ruler: What’s the real context size of your long-context language models? Preprint, arXiv:2404.06654.

Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. 2020. Scaling laws for neural language models. Preprint, arXiv:2001.08361.

Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. 2020. Transformers are rms: Fast autoregressive transformers with linear attention. Preprint, arXiv:2006.16236.

Xianhang Li, Zeyu Wang, and Cihang Xie. 2023. An inverse scaling law for clip training. Preprint, arXiv:2305.07017.

Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lee. 2023. Visual instruction tuning. Preprint, arXiv:2304.08485.

Colin Lockard, Prashant Shiralkar, and Xin Luna Dong. 2019. OpenCeres: When open information extraction meets the semi-structured web. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 3047–3056, Minneapolis, Minnesota. Association for Computational Linguistics.

Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. 2016. Pointer sentinel mixture models. Preprint, arXiv:1609.07843.

Denis Paperno, Germán Kruszewski, Angeliki Lazarioudou, Ngoc Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. 2016. The LAMBADA dataset: Word prediction requiring a broad discourse context. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1525–1534, Berlin, Germany. Association for Computational Linguistics.

Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Jiaju Lin, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Guangyu Song, Xiangru Tang, Bolun Wang, Johan S. Wind, Stanislaw Wozniak, Ruichong Zhang, Zhenyuan Zhang, Qihang Zhao, Peng Zhou, Qinghua Zhou, Jian Zhu, and Rui-Jie Zhu. RWKV: Reinventing RNNs for the Transformer Era. Preprint, arXiv:2305.13048.

Bo Peng, Daniel Goldstein, Quentin Anthony, Alon Albalak, Eric Alcaide, Stella Biderman, Eugene Cheah, Xingjian Du, Teddy Ferdinand, Haowen Hou, Przemyslaw Kazienko, Kranthi Kiran GV, Jan Kocon, Bartlomiej Koptyra, Satyapriya Krishna, Ronald McClelland Jr., Jiaju Lin, Niklas Muenninghoff, Fares Obeid, Atsushi Saito, Guangyu Song, Haoqiu Tu, Cahya Wirawan, Stanislaw Wozniak, Ruichong Zhang, Bingchen Zhao, Qihang Zhao, Peng Zhou, Jian Zhu, and Rui-Jie Zhu. 2024. Eagle and finch: Rwkv with matrix-valued states and dynamic recurrence. Preprint, arXiv:2404.05892.

Subhojeet Pramanik, Esraa Elelimy, Marlos C. Machado, and Adam White. 2024. Agalite: Approx-
mate gated linear transformers for online reinforcement learning. Preprint, arXiv:2310.15719.

Zhen Qin, XiaoDong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. 2022. The devil in linear transformer. Preprint, arXiv:2210.10340.

Zhen Qin, Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, and Yiran Zhong. HGRN2: Gated Linear RNNs with State Expansion. Preprint, arXiv:2404.07904.

Zhen Qin, Songlin Yang, and Yiran Zhong. 2023. Hierarchically gated recurrent neural network for sequence modeling. Preprint, arXiv:2311.04823.

Pranav Rajpurkar, Robin Jia, and Percy Liang. 2018. Know what you don’t know: Unanswerable questions for SQuAD. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), pages 784–789, Melbourne, Australia. Association for Computational Linguistics.

Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. 2020. Winogrande: An adversarial winograd schema challenge at scale. In The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, The Thirty-Second Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, February 7-12, 2020, pages 8732–8740. AAAI Press.

Maarten Sap, Hannah Rashkin, Derek Chen, Ronan Le Bras, and Yejin Choi. 2019. Social IQa: Commonsense reasoning about social interactions. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pages 4463–4473, Hong Kong, China. Association for Computational Linguistics.

Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear Transformers Are Secretly Fast Weight Programmers. Preprint, arXiv:2102.11174.

Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. 2021. Linear transformers are secretly fast weight programmers. Preprint, arXiv:2102.11174.

Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey. 2023. SlimPajama: A 627B token cleaned and deduplicated version of RedPajama.

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.

Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. 2023a. Retentive network: A successor to transformer for large language models. Preprint, arXiv:2307.08621.

Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. 2023b. Retentive network: A successor to transformer for large language models. Preprint, arXiv:2307.08621.

Philippe Tillet, Hsiang-Tsung Kung, and David Cox. 2019. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, pages 10–19.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. 2023. Llama: Open and efficient foundation language models. Preprint, arXiv:2302.13971.

A Vaswani. 2017. Attention is all you need. Advances in Neural Information Processing Systems.

Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, Garvit Kulshreshtha, Vartika Singh, Jared Casper, Jan Kautz, Mohammad Shoeybi, and Bryan Catanzaro. 2024. An Empirical Study of Mamba-based Language Models. Publisher: arXiv Version Number: 1.

Songlin Yang, Jan Kautz, and Ali Hatamizadeh. 2024. Gated delta networks: Improving mamba2 with delta rule. arXiv preprint arXiv:2412.06464.

Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated Linear Attention Transformers with Hardware-Efficient Training.

Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. 2025. Parallelizing linear transformers with the delta rule over sequence length. Preprint, arXiv:2406.06484.

Songlin Yang and Yu Zhang. 2024. Fla: A triton-based library for hardware-efficient implementations of linear attention mechanism.

Tian Ye, Zicheng Xu, Yuanzhi Li, and Zeyuan Allen-Zhu. 2024. Physics of language models: Part 2.1, grade-school math and the hidden reasoning process. Preprint, arXiv:2407.20311.

Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. 2019. HellaSwag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 4791–4800, Florence, Italy. Association for Computational Linguistics.
Hongxin Zhang, Weihua Du, Jiaming Shan, Qinzhong Zhou, Yilun Du, Joshua B. Tenenbaum, Tianmin Shu, and Chuang Gan. 2024. Building cooperative embodied agents modularly with large language models. Preprint, arXiv:2307.02485.

Jingwei Zuo, Maksim Velikanov, Dhia Eddine Rhaiem, Ilyas Chahed, Younes Belkada, Guillaume Kunsch, and Hakim Hacid. Falcon Mamba: The First Competitive Attention-free 7B Language Model.

A Appendix

A.1 Training Efficiency Discussion

We report the training throughput of our scaled RNN models in Table 9. The results show that the training throughput of Gated DeltaNet-E8 is only 0.21 times that of Gated DeltaNet-E1, indicating a significant decrease in training efficiency as the state size increases. This reduction is primarily attributed to two factors: the RNN kernel and the expansion of queries and keys.

Starting with the RNN kernels, we compare the FLOPs of GLA, Gated DeltaNet, and self-attention in Table 10. As shown, current RNN kernels require significantly fewer FLOPs than self-attention. For example, when \( L = 4096 \) and \( D = 128 \), self-attention consumes 10.7 times more FLOPs than GLA and 8 times more FLOPs than Gated DeltaNet. Consequently, our scaled RNN models have similar FLOPs to self-attention and should exhibit comparable training efficiency. However, we identify the bottleneck within the RNN kernel itself. Despite having a similar number of FLOPs as self-attention, the RNN kernels run significantly slower than FlashAttention-2 (Dao, 2023). Specifically, the kernel TFLOPS of Gated DeltaNet and GLA are approximately one-tenth that of FlashAttention-2, indicating substantial room for improving the kernel efficiency of RNNs. FlashAttention-2 benefits from a highly optimized CUDA implementation, whereas the kernels of GLA and Gated DeltaNet rely on less optimized Triton implementations (Tillet et al., 2019). This highlights the potential to enhance RNN kernel efficiency by adopting CUDA-based optimizations.

The second factor is the expansion of queries and keys, which involves multiple operations with fewer FLOPs, such as head-wise matrix multiplication, SiLU activation, convolutional layers, and normalization. Although these operations require relatively few FLOPs, they incur higher memory access costs. Fortunately, as these operations are head-wise parallel, they can be fused into the RNN kernels, saving most memory operations, which should ideally make their overhead negligible.

As a result, the training efficiency of our scaled RNN models has the potential to match that of self-attention. We leave this as future work.

<table>
  <tr>
    <th></th>
    <th>State</th>
    <th>TGS</th>
    <th>Ratio</th>
  </tr>
  <tr>
    <td>Gated DeltaNet-E1</td>
    <td>1×</td>
    <td>24462</td>
    <td>1.00</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E2</td>
    <td>2×</td>
    <td>13998</td>
    <td>0.57</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E4</td>
    <td>4×</td>
    <td>9736</td>
    <td>0.40</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E8</td>
    <td>8×</td>
    <td>5172</td>
    <td>0.21</td>
  </tr>
</table>

Table 9: Training Throughput Comparison of Scaled Gated DeltaNets. “TGS” refers to training tokens per GPU per second.

<table>
  <tr>
    <th></th>
    <th>FLOPs</th>
    <th>Ratio</th>
    <th>Kernel TFLOPS</th>
  </tr>
  <tr>
    <td>Self-Attention</td>
    <td>2HDL^2</td>
    <td>1×</td>
    <td>213</td>
  </tr>
  <tr>
    <td>GLA</td>
    <td>6HLD^2</td>
    <td>\( \frac{1}{3D} \times \)</td>
    <td>21</td>
  </tr>
  <tr>
    <td>Gated DeltaNet</td>
    <td>8HLD^2</td>
    <td>\( \frac{1}{4D} \times \)</td>
    <td>25</td>
  </tr>
</table>

Table 10: FLOPs and Kernel Comparison of Self-Attention, GLA, and Gated DeltaNet. “FLOPs” refers to floating-point operations, while “TFLOPS” denotes the number of tera floating-point operations per second. \( H, D, \) and \( L \) represent the number of heads, head dimension, and sequence length, respectively. “Kernel TFLOPS” is measured with \( B = 4, H = 8, D = 128, \) and \( L = 4096 \) on an H800 GPU. The self-attention kernel is implemented using FlashAttention-2, whereas the kernels for GLA and Gated DeltaNet are derived from FLA.

A.2 Comparison with Additional RNN Models

We compare our scaled Gated DeltaNet with other RNN models in Table 11. As shown, Gated DeltaNet outperforms other RNN models, including RetNet (Sun et al., 2023b), Mamba (Gu and Dao), and DeltaNet (Yang et al., 2025). Furthermore, our scaled Gated DeltaNet-E8 achieves even greater performance improvements.

A.3 Implementation Details

A.3.1 Architecture

The architectural details of the two basic configurations are summarized in Table 12, serving as the foundation for all our models. Both GLA (Yang et al.) and Gated DeltaNet (Yang et al., 2024) are implemented based on FLA (Yang and Zhang, 2024), incorporating short convolutions and output gates in their RNN layers. Since Transformer++
Figure 3: Scaling Trend of Maximum Recall Context Length vs. State Size on S-NIAH. We use numbers and UUIDs as needle values, respectively, on 1.5B Scale

<table>
  <tr>
    <th>Model</th>
    <th>FDA ↑</th>
    <th>SWDE ↑</th>
    <th>SQuAD ↑</th>
    <th>Avg ↑</th>
  </tr>
  <tr>
    <td>RetNet</td>
    <td>14.3</td>
    <td>42.8</td>
    <td>34.7</td>
    <td>30.6</td>
  </tr>
  <tr>
    <td>Mamba</td>
    <td>6.2</td>
    <td>41.4</td>
    <td>35.2</td>
    <td>27.6</td>
  </tr>
  <tr>
    <td>DeltaNet</td>
    <td>17.2</td>
    <td>49.5</td>
    <td>37.4</td>
    <td>34.7</td>
  </tr>
  <tr>
    <td>Gated DeltaNet</td>
    <td>21.7</td>
    <td>53.6</td>
    <td>37.8</td>
    <td>37.7</td>
  </tr>
  <tr>
    <td>Gated DeltaNet-E8</td>
    <td><b>27.0</b></td>
    <td><b>66.2</b></td>
    <td><b>39.2</b></td>
    <td><b>44.1</b></td>
  </tr>
</table>

Table 11: Comparison with Additional RNN Models on FDA, SWDE, and SQuAD. Following (Yang et al., 2025), we truncate the context length to 2048 for our models. The results of RetNet, Mamba, and DeltaNet are from (Yang et al., 2025).

<table>
  <tr>
    <th>Hyperparameter</th>
    <th>400M</th>
    <th>1.5B</th>
  </tr>
  <tr>
    <td>Model Dim</td>
    <td>1024</td>
    <td>2048</td>
  </tr>
  <tr>
    <td>Num Heads</td>
    <td>8</td>
    <td>8</td>
  </tr>
  <tr>
    <td>Head Dim \( d_{qk} = d_v \)</td>
    <td>128</td>
    <td>256</td>
  </tr>
  <tr>
    <td>Num Layers</td>
    <td>24</td>
    <td>24</td>
  </tr>
  <tr>
    <td>Intermediate Dim</td>
    <td>2816</td>
    <td>5632</td>
  </tr>
</table>

Table 12: Architecture Details

has fewer parameters per layer compared to GLA and Gated DeltaNet, we increase its number of layers to achieve a comparable parameter count.

A.3.2 Training Details
All models are trained on the SlimPajama dataset (Soboleva et al., 2023). The 400M-parameter models are trained on 15B tokens, while the 1.5B-parameter models are trained on 100B tokens. Across all experiments, we employ the AdamW optimizer with a weight decay of 0.1 and gradient clipping of 1.0. We use cosine learning rate scheduling with a peak learning rate of 3e-4. The warm-up steps are set to 1000 and 2000 for the 400M and 1.5B models, respectively. The batch size is set to 0.5M tokens, with a sequence length of 4096 tokens during training. All models use the Mistral tokenizer with a vocabulary size of 32,000.

A.3.3 Evaluation Details
Following (Yang et al., 2024), our common-sense reasoning evaluation includes WikiText (Merity et al., 2016), LAMBADA (Paperno et al., 2016), PIQA (Bisk et al., 2020), HellaSwag (Zellers et al., 2019), WinoGrande (Sakaguchi et al., 2020), ARC-easy (ARC-e) and ARC-challenge (ARC-c) (Clark et al., 2018), SIQA (Sap et al., 2019), and BoolQ (Clark et al., 2019). Real-world retrieval tasks include FDA (Arora et al., 2023), SWDE (Lockard et al., 2019), and SQuAD (Rajpurkar et al., 2018). The above tasks are evaluated using Im-eval (Gao et al., 2024). We also evaluate on 14 English tasks in LongBench (Bai et al., 2024), including Narrative QA, QasperQA, MultiField QA, HotpotQA, 2WikiMulti QA, Musique, GovReport, QMSum, MultiNews, TRec, Trivia QA, SamSum, LCC, and RepoBench-P.

A.3.4 Implementation
Our experiments are primarily conducted using publicly available libraries, including PyTorch, XTuner (Contributors, 2023), and FLA (Yang and Zhang, 2024), with training performed on H100 GPUs. This paper has been refined with the assistance of a chatbot.

A.4 State Size of Transformer++ with Sliding Window Attention
There are two scenarios regarding the state size of Transformer++ with sliding window attention. First, when the context length is smaller than the window size, the state size is the same as that
of the original Transformer++, i.e., \( L \times H \times (d_{qk} + d_v)/2 \). Second, when the window size is smaller than the context length, as is the case in LongBench where the context length can be up to 64k tokens—much larger than the window size of 2048 in Transformer++-2k—the state size of Transformer++-2k with sliding window attention is approximately \( W \times H \times (d_{qk} + d_v) \), where \( W \) denotes the window size.

A.5 Ethical Considerations
Our work focuses on improving the long-context understanding ability of RNNs. The potential applications of our work include enhancing the performance of RNN models in long-context tasks. We believe that our research can contribute to the development of more powerful and efficient RNN models, which can benefit society by improving the performance of AI systems in various applications. Our work is not associated with any ethical concerns.

Metareview score: 4
Metareview: This paper investigates how the recurrent state size impacts the performance of RNN-based LLMs in long-context understanding and retrieval tasks. The authors propose a parameter-efficient approach for scaling up state size, which demonstrates consistent empirical gains across several benchmarks. The core idea—that increasing state size significantly improves recall capacity—is explored in depth, and the work offers a promising direction for improving the performance of RNN-based architectures.
Summary of reasons to publish: - The paper addresses a timely and important problem: long-context modeling remains a key bottleneck for linear RNNs, and understanding the role of state size in such settings is crucial.

- The experimental section is comprehensive. While the study may not fully qualify as a “scaling law” in the strict sense, the experiments are well-executed and span diverse long-context benchmarks, which should be of interest to the RNN community.

- The proposed parameter-efficient state expansion approach is both simple and effective. Importantly, the paper offers ablations that compare against alternative parameter-efficient methods, demonstrating superior expressiveness and performance.