RAT: Bridging RNN Efficiency and Attention Accuracy via Chunk-based Sequence Modeling
Authors: Xiuying Wei
Anunay Yadav
Razvan Pascanu
Caglar Gulcehre
Venue: NeurIPS 2025 poster ()
Keywords: language modeling
attention
recurrence
Primary area: deep_learning

Abstract: Transformers have become the cornerstone of modern large-scale language models, but their reliance on softmax attention poses a computational bottleneck at both training and inference. Recurrent models offer high efficiency, but compressing the full sequence into a fixed-size and holistic representation can suffer from memory degradation in long contexts and limit fine-grained retrieval. To address this, we propose RAT, an intermediate design that bridges the efficiency of RNNs and capacity of attention. RAT partitions the input into chunks, applies recurrence within each chunk for local dependencies, and softmax-based attention across chunks for long-range interactions. This design mitigates memory degradation and enables direct access to distant tokens, while retaining computational efficiency. Empirically, with a chunk size of 16, the RAT block achieves a $7\times$ improvement in training speed for 100K sequence length and $9\times$ in generation at the 4K position, while maintaining similar performance compared to standard attention. We demonstrate this by training 1.3B parameter models from scratch and performing large-scale evaluations, including short- and long-context benchmarks, as well as supervised fine-tuning (SFT). We further propose a hybrid architecture that interleaves RAT with local attention. By combining efficient long-range modeling with strong local interactions, this hybrid design not only improves inference speed and reduces cache memory usage, but also consistently enhances performance and shows the overall best results. Code is available at \url{https://github.com/CLAIRE-Labo/RAT}.

RAT
Bridging RNN Efficiency and Attention Accuracy via Chunk-based Sequence Modeling
Xiuying Wei'; Anunay Yadav', Razvan Pascanu', Caglar Gulcehre' 'CLAIRE, EPFL 'Google DeepMind
Abstract
Transformers have become the cornerstone of modern large-scall seage-scalella language models, but their reliance on softmax attention poses a computational bottleneck at both training and inference. Recurrent models offer high efficiency, compressing the full sequence into fixed-size and holistic representation can suffer from memory degradation long contexts and limit fine-grained retrieval. To address this, we propose RAT, an intermediate design that bridges the efficiency of RNNs and capac- ity attention. RAT partitions the input into chunks, applies recurrence within each chunk for local dependencies, and softmax- attention across chunks for long- range interactions. This design mitigates memory degradation and enables direct access to distant tokens, while retaining computational efficiency. Empirically, with a chunk size of 16, the RAT block achieves a improvement in training speed for 100K sequence length and 9x in generation at the 4K position, while maintain- ing similar performance compared to standard attention. We demonstrate this by training 1.3B parameter models from scratch and performing large-scale evalua- tions, including short- and long-context benchmarks, as well as supervised fine- tuning (SFT). We further propose hybrid architecture that interleaves RAT with local attention. By combining efficient long-range modeling with strong local in- teractions, this hybrid design only improves inference speed and reduces cache memory usage, but also consistently enhances performance and shows the overall best results. Code is available at https: // github........ cOm/CLAIRE Labo/RAT. /RAT.
Introduction
Language modeling has long been dominated by Transformer- based architectures due to their strong performance across a wide range of tasks. However, their reliance on full self-attention [1] results in quadratic time and memory complexity with respect to sequence length, which limits scalability in long-context processing. This limitation has motivated a wave of recent efforts to revisit recurrent mod- els or propose novel linear recurrent models such state space models, linear attention methods [2-8].
By comparing these architectures, we observe a key difference between recurrent models and self-attention. Recurrent approaches compress the full sequence history into fixed-size and holistic representations, which can lead to degraded memory when modeling long sequences and limit precise information retrieval. In contrast, self-attention retains full-token access and thus do not suffer from the two problems, but at the cost of heavy computation. This motivates us to explore an intermediate design that partially compresses the sequence while still maintaining global access.
We propose a RAT layer, a simple yet effective temporal mixing method with a chunk- based design. It divides long sequences into chunks, applying recurrence within each chunk for local modeling and softmax-based based attention across chunks for direct access distant information (see Fig. 1a). Recur- rence efficiently captures short-range dependencies while avoiding the memory degradation common
"Correspondence to xiuying, wei@epfl.ch
39th Conference on Neural Information Processing Systems (NeurIPS 2025).

---

ouiput projection
inter-chunk Attention
Attention
RAT
((F)
Intra-chunk RNN
(a) Comparison of three structures: Circles denote token representations and arrows show the information flow for the 5th token. (1) Attention: Each token keeps its own representation and has full access tion is progressively compressed along the sequence, and the current token accesses only the last hidden state. (3) RAT (L=2): Intra-chunk recurrence compresses local information, while inter-chunk attention enables direct access to previous chunks. and respectively.
norm and input projection
previous tokens. (2) RNN: Informa-
(b) RAT structure. The sym- bol bol (-) further applies sig- moid function the gates.
(1) and (2) can be viewed as RAT
in long sequences, whereas attention chunk-level representations enables long-range retrieval. By adjusting the chunk size L, RAT interpolates between attention (when L= and RNN (when L=T)..
To make RAT scalable and efficient, we further explore its positional encoding, parameter allocation, and efficient implementations that specifically address the causal masking problem in this chunk-based design. Moreover, we investigate a hybrid architecture that interleaves RAT with sliding-windon attention [9-12], which focuses computation on local windows. As RAT suggests that overusing attention in local contexts underutilizes its strengths, and such dependencies can be handled more efficiently via lightweight recurrence, the two approaches can complementary.
We demonstrate that RAT is both efficient and performant. For instance, with a chunk size of L= 16. the latency of a single temporal mixing block at position 4096 is up 9x lower than that of attention. As shown in Table 1, the full model achieves up to 10x higher maximum throughput, and interleaving with local attention yields about 4x improvement. We pretrain models at 1.3B scale and compare their results on 1) shot commonsense reasoning with short contexts, (2) long-com congccttext understanding on LongBench [13], and (3) supervised fine-tunin tne-tuning on long-context tasks. RAT with L 16 performs on par with full attention in most benchmarks and even outperforms it on several I LongBench tasks. When interleaved with local attention, achieves the best overall results across all variants while maintaining high efficiency. Our main contributions are as follows:
We propose the RAT layer, novel intermediate architecture that bridges the efficiency of recurrence and the capacity of attention. It compresses only local context while preserving global access, enabling direct retrieval and avoiding the memory degradation caused by full sequence compression in long contexts.
RAT is simple, scalable, and efficient. It requires no custom CUDA kernels and naturally compatible with existing multi-dimension parallelism schemes. We also introduce hybrid variant by interleaving RAT with aiding-windoo attention, enabling efficient long-range modeling with strong local interactions.
We validate RAT through extensive experiments at the 1.3B scale across diverse tasks, including 7 short-context reasoning benchmarks, 11 long-context tasks, 4 supervised fine-tuning objectives, and 9 retrieval-heavy synthetic evaluations. RAT with L=16 shows comparable performance to full attention with 9x faster single-layer decoding. Its hybrid variant with local attention yields the best overall results, along with 3 - 4x maximum throughput boost-for instance, +1 accuracy commonsense reasoning, +4 code completion, +4 challenging QA task, and +1 on a difficult summarization task.
Table I: Representative results for 1.3B models across pretraining, direct evaluation, and SFT. -SWA denotes interleaving with sliding-window window attention (SWA) (window size 1024). Maximum throughput is measured by generating 1024 tokens given a prompt of 3072 tokens on a H100 GPU in GH200 system. See Sec. 4 for details.
Model
Throughput
Pretrain
Direct Evaluation
SFT
Val.
CSR Avg.
SQA Avg. F1
Summ Avg. Rouge-L
Code
NQA! F1
QMSum Rouge-L
token/sec
PPL
Avg. EditSum
Attention
3052
7.61
$6.9
18.2
19.5
23.9
61.3
23.4
RAT (L=16)
31170
7.67
56.7
19.6
20.2
17.4
60.8
23.3
Attention-SWW
4605
7.61
$7.1
17.4
19.4
21.7
63.3
23.4
RAT (I-16) SMA
13582
7.57
58.0
19.5
28.2
63.2
24.6

---

2 Overall architecture
To motivate our design. we first examine how attention and recurrence compress contextual information at the token level, along with the equations that will be reused later. We then introduce RAT an intermediate mechanism that inherits the respective advantages of both.
2.1 Attention: Full-Token Access
In softmax attention mechanisms, each token has access to all preceding tokens through a learned token-specific weighted aggregation:
y f= VV.
(1)
f()) denotes causal masking and the softmax function, and f(q. K. represents the attention weights used to aggregate V. We adopt Python-style indexing in K, V. emphasize that each query attends to all keys and values. This full-token access makes attention dominant architecture in sequence modeling, but suffers from high computational cost in both training and inference.
2.2 Recurrence: Full-Sequence Compression
Recurrent models maintain a summary of past information in a fixed-size representation, allowing each step depend only on the previous state and the current input [14, 15]. To initiate the design of RAT, we adopt a simple and fast linear recurrence [16-18], which uses diagonal matrices and no nonlinearities to simplify the classic recurrence into an EMA-like gating mechanism:
ú = 9: 3 02:00,11 1 (1 (1lgg))...
y/ 2, 11,
where Ie and 2, Rd denote the per-dimension forget and output gates, respectively, computed via linear projections of the input followed by a sigmoid activation. Here, D refers to the dimension of the model. This can be efficiently implemented using the parallel scan algorithm [18]. We take this minimal choice to highlight the core idea, but the recurrence in RAT is not limited to this certain form and can be extended more advanced variants such as 2D recurrence nonlinear RNNs. We leave this as future work.
The state-based formulation of recurrence yields low computational cost and efficient inference, and performs well on short sequences. However, compressing the entire sequence history into a fixed-size and holistic state, even with more expressive designs, can still suffer from memory degradation when the sequence length grows [19], and limits precise information retrieval, particularly in noisy contexts.
2.3 RAT: Chunk-Based Intermediate Design
To bridge the strong performance of attention enabled by full-token access with the efficiency of RNNs derived from dall-sequence compression, we propose an intermediate design by reinterpreting the input as sequence of shorter chunks. A recurrent module is applied within each chunk to model local dependencies, followed by cross-chunk attention to enable global interactions, as illustrated Fig. la. This design mitigates the fixed-size representation limitation of RNNs and the inefficiency of attention.
Technically, we divide sequence of length T into C chunks of length L,, such that 7 C- L. A token originally at position t is re-indexed as (c,l), where e denotes the chunk index and the position within the chunk. Within each chunk, forget gate 9c.1 is used recurrently aggregate the value ve,1 and key kz,, vectors, yielding updated representations T.,1 and kg.- For each query qe.t, we compute attention on the chunk-lever key and value vectors, including K for all preceding chunks, and k for the current chunk. The causal masking function f(-) restricts attention to the chunks before it, followed by a softmax operation. Finally, an output gate is applied to produce the output:
ر., = 00. (1-
(Intra-chunk RNN)
k. 9e.1 Oo.,kka.ttt 9c.1) Oka
(Intra-chunk RNN)
(3)
ye.1 = [(q.,) K.T V.
(Inter-chunk attention)
yo. = Zc,1 3yo.
3

---

The overall computation flow is illustrated in Fig. lb. Thanks to the short length of each chunk, the recurrent module efficiently captures local dependencies without significant information loss. Cross-chunk attention then enables global access without compressing distant information. This design preserves the accuracy of softmax-base attention while reducing computation, as inter-chunk attention operates over shorter sequences with I. as the FLOPs reduction ratio. We also experimented with a reversed variant by applying attention within chunks and recurrence across them, but found that standard RAT achieves better FLOPs utilization and overall performance (see Appendix C).
Benefits of the chunk design We position this as an intermediate architecture between RNNs and attention, and refer to it as RAT. By adjusting the chunk size T., RAT interpolates between their behav- iors: when 7. it reduces to a pure RNN; when 1.11, 1, it resembles full attention. This allows us to pursue the trade-off between attention and RNN with a single-layer design, and offers greater flexibil- ity for hybrid modeling by using different chunk sizes to emulate varying behaviors. From a mechanis- tic perspective, unlike either classic or advanced recurrence (state space linear attention models) that rely on fixed-size and holistic representations, our chunk-based structure enables memory capacity to scale with sequence length while maintaining fixed FLOPs reduction ratio. Its partial compression with direct access to prior chunks ensures its superior performance on retrieval -heayy tasks.
3 Scalable and Efficient Modeling with RAT
3.1 Design details
Parameter allocation We aim to control the parameters of RAT at 4D², given D for the output projection, 3D2 for the query, key, and value projections in attention, and an additional 2D for the two gates in RNN. We explore lightweight alternatives by using low-rank projections for the gates g and z, or by sharing query and key projections across heads. Empirical results in Sec. 4.2 show that sharing the query and key l slightly outperforms using low-rank gates. Notably, this design does not collapse into single-head attention, since the forget gate operates the per-dimension level and produces distinct gated keys l for the inter-chunk attention.
Positional encoding We examine how to encode positional information in RAT, due to the presence of cross-chunk attention. Motivated by the fact that RNN captures positional information through its sequential structure, we find that applying positional encoding at the chunk level, rather than relying on original token positions, yields slightly better fidelity. This strategy also improves length generalization, as the number of positions requiring encoding (i.e., the number of chunks) is much smaller than the full sequence length. In our main experiments, indices for inter-chunk attention. In the length generalization study (see Subsec. C.2), we further explore NoPE [21], which yields the best overall generalization performance.
RoPE [20] based on chunk
Hybrid design with Local Attention RAT is a hierarchical architecture, but not hybrid model, as it applies the same strategy across all tokens, layers, and heads. It is compatible with various hybrid strategies and offers more flexibilities in hybrid modeling by varying the chunk size in different layers or heads. In particular, we explore interleaving RAT with hiding-window attention (SWA) [9-11], a widely adopted technique in recent models [22, 23, 5, 7J. We find that the two are highly complementary: while local attention methods allocate most computation within fixed windows, RAT reserves attention for global access and handles local modeling more efficiently. Interleaving them enables the model to efficiently and effectively capture both short-range and long-range dependencies.
3.2 Efficiency
We discuss efficiency-related aspects of RAT, and show that our current implementation does not rely on custom CUDA or Triton kernels, yet achieves significant speed-ups in both training and generation in experiments. The pseudocode for training is provided in Listing 1, and the decoding algorithm is shown in Listing 2.
To begin, the FLOPs per token of RAT are O D), compared to O(D) for RNN and O(T D) for full attention, where C' is the number of chunks, D the model dimension, and T the sequence length. Most components of RAT are simple and easy to implement, except that the chunk-based design introduces non-trivial causal masking challenge during training, which we address below.

---

Causal masking problem In training, where tokens are processed in parallel, special care must be taken to apply causal masking in inter-chunk attention. First, block-wise causal mask is required. Second, each token must also attend to its own chunk's key and value, which should be gated up to its position. They vary across tokens due to causal masking, preventing efficient parallel computation. To address this, we adopt an online softmax formulation [24]: we separately compute [q - and ) and then combine the results by adjusting the softmax denominator. The first term can be implemented in parallel using existing attention frameworks, while the second is handled via a simple einsum.
Practical implementation For training, we implement intra-chunk recurrence in Eq. (3) using PyTorch's associative scan, enabling forward and backward passes with (TT) FLOPs. Compared to full-sequence recurrence, chunking reduces scan depth from (logT) to (log and thus improves parallelism. For inter-chunk attention, we use PyTorch's ' flex attention above. It supports flash attention with flexible features such as custom masks and returning the softmax denominator, and thus aligns well with our needs. For decoding, tokens are generated sequentially. Intra-chunk recurrence only requires single-step updates and can be implemented directly. For inter-chunk attention, standard implementations like flash attention [25] can be used without modification, as no complex causal masking is required at inference time.
implement the first term
Parallelism We think RAT is compatible with both tensor parallelism and context parallelism, which are commonly used for large model dimensions and long sequence lengths, respectively. Both the recurrence and cross-chunk attention RAT are head-independent, making it easy to apply standard tensor parallelism by assigning different heads to different GPUs. For context parallelism, the intra-chunk recurrence is chunk-independent, allowing chunks to be distributed across GPUs. Since RAT stores much fewer number of chunk-level key/value vectors (e.g.. 16x fewer than full attention), the cross-chunk attention may even avoid ring-style communication.
Experiments
In this section, we present the efficiency and large-scale evaluations of RAT, along with comparisons to other models. Additional discussions are provided in Appendix C.
4.1 Setup
For brevity, we summarize the setup for the 1.3B model with a 4K context window, which is used in most of our experiments. Full implementation details are available in Appendix A.
Model We adopt a Transformer architecture that interleaves token mixing block with hidden-state mixing block (FFN), each wrapped with residual connections and LayerNorm. We compare variants that use different token mixing modules, including RNN, attention, and RAT, as well as recent state space and linear attention models. In addition, we provide hybrid models that interleave RAT with sliding-wing window attention (window size 1024), such as Attention-SWy SWA and RAT-SWA
Efficiency We benchmark the latency of a single token mixing block, including input and output projections, on a single H100 GPU (GH200 system, 120GB), and also report the maximum throughput of the full model. We measure the time required train on a full sequence or generate a batch of tokens at specific positions. For fair comparison, we use flash attention for the attention baseline and associative scan for RNN. All models are compiled using compile and evaluated in bfloat 6 with torch. cmp..
Accuracy We pretrain all 1.3B models on 100B tokens from FineWeb-Edu [26] using the same setup: a learning rate of 8.0e-4 decayed to 1.0e-6 (cosine schedule) and a global batch size of 2M to- kens, following DeepSeek's 's hyperparameter guidelines [27]. We directly evaluate the models on both short-context tasks- several classical commonsense reasoning benchmarks [28]-and long-context tasks, including 11 tasks from LongBench [13] covering QA, summarization, and code completion. Since LongBencc includes instruction-heavy prompts that pretrained models often struggle with, we also include SFT-based evaluations. Specifically, we use NarrativeQA [29] (two modes), QMSum [30], and WikiSum [31] to test long-context understanding with SFT. To assess retrieval capabilities,
5

---

include the Ruler benchmark [32] and test nine synthetic needle-in sosele-in-haystarn haystack tasks with varying con- figurations. A single round of lightweight fine-tuning applied to adapt models to specific prompts.
4.2 Analyses of RAT
Figure 2: Latency of the temporal mixing block (including linear projections) with a model dimension of 2048. (a): full-sequence latency with 200K tokens; (b): generation of 512 tokens at specified positions. We flash atfention for Attention.
Table Maximum throughput full models (tokens/sec). measured by gen- erating 1024 tokens from a 3072-token prompt. By reducing the KV cache memory and boosting speed, we achieve 10x maximum throughput compared to flash attention. and even more on 13B models, as attention suffers from poor GPU utilization larger scale.
1750
RATIL-1129
1500
COM!
RAT(L-10)
1250
Albestion
1000
9
Model
1.3B
7B
RAT(1=16)
31170
10103
5749
Attention
3152
983
534
5-
Ratio
10.2x
10.3x
10.8x
250
0
619
3192
143544
33766
48036
19877
3192
19384
33748
68536
130772
(8) Train
(b) Generation
Efficiency study Fig. 2 and Table 2 present the efficiency comparison, with additional results provided in Appendix B. For training, on the strong H100 GPU, when the sequence length is short (e.g., 4096), RAT is slightly slower than attention due underutilized GPU parallelism in the flex attention (with only 256 chunks) and the overhead introduced by associative scan. We expect these can be improved through further kereel-level optimizations. As the sequence length increases, RAT becomes increasingly efficient: for 16 we observe approximately 2 speedup at 16K, 3x at 32K, 4x 64K, and 7 x at 100K tokens. For generation, RAT (L=16) achieves 9X speedup at position 4K and around 10x for longer sequences. It also reduces KV cache usage, making it significantly less prone out-of-memory (OOM) errors and enabling much higher maximum throughput.
Ablation study We conduct the ablation study on a 200M-parameter model trained on the book dataset. Allocating more parameters the intra-chunk RNN gates yields significantly better performance than assigning them to the inter-chunk attention query and key projections, improving the perplexity by 0.4-0.5. Further replacing the original RoPE with cross-chunk variant, where positions are indexed by chunk index, brings additional improvement, especially at long sequence lengths, with 0.3 drop in PPL observed.
15.6
Skared qli croe-chusk RaPE, (RAT
Attentinn
Laterank gatss
14:75
RAT BONN TwItfs
Shared and ton-naik eutpol
Shared
14,50
11.-
14.25
14.00
3.5
13.25
13.5
13.50
7.53
R192
14/944
32766
1-1
1/464
1 128
256
L-T
1-1
1-16
1-44
L 128
L-T
Sey fms (T)
Chrsak vize (1.)
Cherc (l.)
(n) Parameter allocation and PE design
(b) 200M model
15B tokeas (beok)
(4) 1.38 model
1088 tokens (web)
Figure 3: (a) Ablation study on (L-64) (b) and (c) show pretraining results on 200M and 1.3B models, respectively, RAT lies between RNN and attention in terms of pretraining perplexity.
Pretraining In Fig. 3(b) and (c), we begin our study by examining pretraining perplexity foundational measure of architectural performance, demonstrating that RAT falls between attention and RNN. In (b), using 200M model, we observe that increasing the chunk size leads to higher

---

perplexity while reducing FLOPs. Viewing attention as L = 1 and RNN as L= we observe perplexities of 13.26 for attention, 13.42 for RAT (L=64), 13.72 for (LT and 14.31 for RNN. Interestingly, increasing the training context from 8 to = 327668 leads to sharp rise in perplexity for RNN, while both attention and RAT remain stable, demonstrating the pitfalls of full-sequence compression. We then scale up to 1.3B model trained on web data, which primarily consists of short sequences (often under IK tokens) concatenated with separator tokens to construct 4K-length contexts. We find that RAT (L=16) gives the comparable performance to attention, and again RAT with different chunk sizes exhibits behavior intermediate between attention and RNN.
4.3 Large-scale evaluation and comparisons
In this section, we present direct evaluation and SFT results of our trained 1.3B models. We mainly compare (L=16) with its two extremes: RNN and attention. Additionally, we include results from recent state space and linear attention models [5], which are trained under the same settings (model size, sequence length, and dataset) as ours. We report the value/slot sizes, which indicates the size of memory slots and FLOPs, based on sequence length T', model dimension D, and number of layers N. Note that state space and linear attention models use fixed-size memory and expand the classic recurrence state, while RAT scales its memory capacity with the sequence length.
Commonsense reasoning In Table 3, the performance gap (within two points on average) between RNN and attention is small on these short commonsense reasoning tasks, indicating that RNN be surprisingly competitive in short-context settings. RAT (L-16) consistently outperforms RNN on most tasks and surpasses other recent models on 4 out of 7 benchmarks, with only at most 20N D memory slots. We also observe that some binary-choiee tasks (c.g.. BoolQ. WinoGrande) show inconsistent results under small-scale pretraining. For instance, attention struggles on BoolQ, while RAT under- performs on WinoGrande. Interestingly, the hybrid variant RAT (L=16) SWA, which interleaves RAT with SWA, achieves the best overall accuracy across almost all tasks. This suggests a complementary architecture: using SWA for strong local attention and RAT for long-range dependencies can bring us a model that is not only more efficient but also more accurate than its dense counterpart.
Table 3: We evaluate models with model dimension
2048 and number of layers N 24 seven [28]. Since all tasks (except BoolQ. which only 13 sequences > 300 tokens) have sequences I 300 tokens, we set T = 300. Results marked with are from a recent study [5] trained on the same setting ours. For SWA, slot sizes remain 150N D due to the short context.
commonsense reasoning tasks using an--evaltation-hanee
Model
Value/Slot size
ARC-C
ARC-E
HellaSwag acc/acc_n
LAMBADA
PIQA
WinoGrande atc
BoolQ
T=300
acc/acc_n
accface_n
3EC
acolacc_
acc
RNN
ND
35. 15/35578.5
70.66666.46
43.5555
71.49/2.69
$3.99
62.54
Mamba*
.433.88
61.13
Mamba2*
256ND
/37.88
72,47/-
JS5.67
J45,66
71.87/-
55.24
60.13
DeltaNet*
128NDD
-/35.66
68.177--
-50.93
-11..66
70.72/-
$3.35
$5,29
DeltaDe
288ND
.33..39
71.21/-
J55.76
J46.65
72.25/-
$7.45
60.24
Attention
300ND
35.67/37.71
25/666667
44.16/57.44
47.84
73.29/72.91
57.70
58.23
(L-16)
2AND
36.35/39.33
7264467.80
43.27/66.44
72.03/72.69
53.20
GoneDDelaNetSSWN
288N/2D 300
./40.1 10
7L.75/-
.56.53
47.73
72.57/-
58.4
63.21
Attontion: SWA.
M0N/2DD-00N/20
35.24/37.20
7L.97/6.200
44.36/57.09
47.97
72.25/72.74
$6.91
61,25
(L-16) SSA
~ 20N/2D 300N/2D
37.20/40.78
72.56/68.22
44.33/57.85
49.33
73.29/73.94
$6.91
63.21
LongBench In Table 4, we present results 11 long-context tasks from LongBench. Since I LongBench includes instruction-heavy prompts across diverse domains, it's natural that no single model dominates all tasks. Evaluating purely pretrained models on such instruction- based datasets can be very challenging. Therefore, we focus on general performance trends, especially average scores by task category. As shown in the table, RAT (L=16) and its hybrid variant RAT (L=16) -SW achieve top performance in many groups, such as question answering and summarization. This is noteworthy. given that RAT(L=16) slightly lags behind attention pretraining perplexity and underperforms in short-context commonsense reasoning. We hypothesize that its advantage on LongBench stems from structural adaptation to very long sequences-fir instance, achieving 24.9 on GovReport (3-1) (vs. 18.6 for the attention baseline) and 16.7 on HotpotQA (2-1) (vs. 13.4). Second, unlike the small gap (>>2 points) between RNN and attention on short-context tasks, the difference becomes much larger on LongBench, further demonstrating the problem of full-sequence compression. We also notice that
7

---

Table 4: Evaluation results on LongBench with 7 = 1096, D 2048 and N = 24. NQA: NarrativeQA, MQA: MulriFicldQA-en. HQA: HotpotQA, WQA: 2WikiMulthopOAA, MSQ: Musique, GR: GovReport, MN: MultiNews, TQA: TA. RBP: RepoBench-P. Metrics: for QA tasks, Rouge-L for summarization, and EditSum for code tasks. Results marked with are from a recent study [5] trained under the same settings as ours. Note that they use local attention with a window size of 2048, while we use 1024.
Single-Decument QA
QA
Summortzation
Completion
Model
T U196
NQA
MQA
HUA
WQA
MSQ
GR
(MSuum
LCC
RBP
RNN
ND
11.7
5.5
19.5
13.5
9.1
15.1
5.4
9.9
16.9
16.8
16.4
16.7
14.0
17.6
15.4
Manbs2*
N
11.1
11.3
18.6
13.7
11.8
15.1
6.7
11.2
6.7
14.5
7.4
9.5
17.9
20.6
19.3
D
128
12.9
10.5
21.3
15.1
10.9
13.2
5.1
9.3
6.5
13.5
7.2
9.1
17.6
20.3
19.0
ConeDDehuNee
14.1
14.0
23.3
13.1
13.7
14.4
5.8
11.3
7.5
16.4
7.9
10.6
18.7
22.1
20.4
Amention
12.3
14.0
28.2
18.2
13.4
12.3
499
12.0
18.5
18.9
21.0
19.5
21.3
26.5
23.9
644
6END
1399
130
27.0
15.0
14.9
16.3
4.3
11.8
18.9
17.7
19.9
18.8
16.2
19.7
18.0
BATCC110
145
28,3
19.6
16.7
18.9
6.3
14.0
17.5
18.3
20.2
20.6
17.4
GosellellNNsSSA
145
15.8
12.6
13.6
6.1
9.1
16.1
12.8
12.7
17.4
Amentinn SWA
4PN6N/2D
1024 /2
24.5
6.1
12.9
19.4
17.3
19.4
BATCL-16)-00 -SA
250 N,N0200-10 1/AN/20
127
186
28.2
18.8
14.2
18.1
7.4
132
20.1
18.6
19.8
19.5
26.3
30.1
28.2
(AT performs less favorably on code completion. However, its hybrid variant significantly improves performance in this setting, outperforming the second-best model by about four points.
Supervised Fine-tuning Since LongBench tasks are instruction- and pretrained-only models often struggle to follow specific prompts, we further assess the performance of SFT on some long- context QA and summarization tasks, as shown in Table For NarrativeQA, where passages are extremely long and may not contain the answer within the truncated 4K-token context, we evaluate two settings: (1) summary-only, which uses only the passage summary, and (2) summary plus passage, which concatenates the summary with the full passage. In both cases, RAT (AT11)))-SM delivers strong performance while remaining efficient. Although both the summary and passage contain useful information, we observe that models, particularly the attention- based ones, perform worse in the second setting, with Fl dropping to 33. We hypothesize that in this setup, attention may be more easily distracted due the absence of structural bias, which can lead flattened attention score distributions [33]. In contrast, models based on RAT appear more robust, likely due to their chunk- wwie structure. For summarization tasks, RNN performs better than in QA, and RAT (L=16) -SW achieves the best result on QMSum, surpassing the next-best model by approximately Rouge-L point.
Table 5: Left: Average and 95th percentile of context lengths of datasets with the LLaMA2 tokenizer. Right: SFT performance on long-context datascts with T 4096, D 2048, and N 24. F1 is used for QA tasks, and Rouge- L used for summarization tasks. NarrativeQA : only the summary of passage is provided. NarrativeQA 2: the summary plus the full passage is provided; the passage typically contains much irrelevant information.
Task
Avg.
95th petl.
Model
Value/Siot
NarrativeQa
NurrativeQAA
QMSam
RNN
ND
30,5
27.4
23.2
34.3
N
838
1296
BAT(L-16)
43.5
2
97k
254k
Amcntion
61.1
33.0
23.4
QMSum
16k
29k
SWention-SWA
4/86//20 1024
63,3
39.5
WikiSum
1828
3473
RAT(L=16)
256
1024//20
36.7
Table 6: Retrieval ability: Accuracy performance with exact match scoring the Needle-in-Hayee Haystack tasks with different configurations from the RULER benchmark [32]. We use T= 1096, 2048, and N 24.
Model
ValueiSutttie size
single_1
aingle_2
aingle__
mult:kkyy_
multikey__ 2
multikkyy_
multiquary
yulcivvalue
RNN
ND
99.4
99.8
0.0
2.2
5.8
5.7
RAT(-66))
GIND
100.0
99.8
55.6
96.4
62.2
1.0
86.3
89.1
RAT(1-16))
256/V D
99.6
10010
94.6
99.6
99.6
82.6
91.2
9488
Atscntion
100.0
100.0
100.0
99.6
99.6
95.2
98.6
99.1
Retrieval ability In Table 6, we evaluate the three models on challenging synthetic needle-in- haystack tasks that test retrieval ability (see Subsec. C.3 for task descriptions). It can be seen that attention consistently performs well due to full-token access. RNN handles simpler tasks like single_1 1 and single_2 reasonably well, but its performance drops on harder settings. RAT (L=16) matches attention closely on most tasks but struggles with UUID ones (single_3 3, multikey_5 due to their difficulty and evaluation metric exact match scoring. Mcanwhile, RAT (L=64) falls
8

---

between RNNs and RAT (L6 as expected given its partial access to long-range context. These results align with the model structures: attention sees all tokens directly, recurrence compresses all history into single state, and RAT blends both, preserving chunk-level access while reducing F1.OPs.
5 Related work
State space and linear attention models Recurrent neural networks (RNNs) have long been used for sequential modeling [34-36]. Recent state space models [37, 18, 2, 7, 8] extend recurrence by ex- panding hidden states from vectors to matrices for higher accuracy, removing non- non-linearity and using structured operations (e.g., diagonal matrices) to enable parallel computation [18]. Per-head gating connects these models [3] to linear attention [4, 38, 39, 8, 40, 5] and shows the duality between linear recurrence and linear attention. Unlike this duality perspective, we take an intermediate view between per-dimension-gee -gate RNNs and softmax-based attention, without enforeing strict linearity. Linear re- currence in RAT is intended primarily for training efficiency, rather than imposed as design constraint.
Despite state expansion, these models inherit the fixed-size representation limitation of classic RNNs, leading to degraded memory on long sequences. A concurrent work [41] proposes addressing this by growing the memory slots logarithmically. We similarly target the fixed-size memory problem, but use softmax- based attention to access distant tokens. Hybrid strategies [23, 42] combining state space models and attention are also popular. While RAT is not hybrid with the same computation graph for all tokens and all layers, it is orthogonal to such designs and provides more flexibility, such using different chunk sizes for different layers and heads, which we leave for future work.
Softmax-based attention Attention mechanisms suffer from slow computation due to their full- token access. Earlier works have explored local attention with chunking, combined with recurrence or sparse mechanisms to access earlier context [9, 10, 43-45) For example, Ma et al. [43] gates query/key vectors before local attention, while Hua et al. [10] adds the outputs of local attention to cross-chunk linear attention. In work, we instead leverage attention's strength in directly accessing distant tokens with lightweight recurrence for local contexts, and hierarchically organize their outputs. As shown in Sec. C. RAT utilizes FLOPs more effectively than local attention. We also regard RAT and local attention as structurally complementary. Although we experimented with interleaving them in this work, we leave the combination with more advanced local attention variants in future research.
RAT can be viewed as an attention mechanism with resolution, akin to how humans process informa- tion: short-term inputs are integrated into coherent representations, while long-term events are stored as memory anchors for selective retrieval. Compared to dilated attention [9], our simple recurrence en- ables global perception, and the chunk-level attention has high training parallelism and reduced cache memory (see Sec. C.I for more). Finally, a recent work [46] also aggregates the key and value vectors and applies softmax-based attention over them. However, it compresses the full-sequence vectors into a 2D memory matrix, unlike our chunk- based design, which enables flexible memory slots.
6 Conclusion
This paper proposes the RAT structure, an intermediate architecture between RNN and attention. It segments sequences into chunks, applies RNN within each chunk to capture local dependencies, and employs attention across chunks to model long-range dependencies. We detail its architectural design and investigate its combination with the sliding-window window attention mechanism. Experimental results across various settings demonstrate the efficiency and accuracy benefits of RAT, highlighting its potential for future language model development.
Limitations. Due to resource constraints and the exploratory nature of this work, our experiments are limited to models with up 1.3B parameters. We have not yet scaled RAT to larger language models such as 7B or 14B to confirm our results hold at that scale. Additionally, we did not adopt supervised fine-tuning techniques commonly used in current industry practices. which typically involve substantial data resources and many engineering details. Instead, we followed a more classical approach based on train-test splits to evaluate performance, as our focus lies primarily on architectural design during pretraining. Lastly, RAT may still face length generalization issues from positional encoding, similar to attention. But as shown Subsec. C.2, this can be significantly reduced by shorter inter-chunk attention spans, and we also explore the use of NoPE
9

---

Acknowledgments and Disclosure of Funding
We sincerely thank the RCP and the SCITAS team EPFL for GPU support. We also thank the Swiss Al Initiative and the Swiss National Supercomputing Centre (CSCS) for supporting this work through grants under project IDs a06 and al0. We thank Skander Moalla, Anja Surina, and Lars Quaedvlieg for helpful discussion SFT tasks. We extend our appreciation to Karin Getaz for administrative support.
References
[1] A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
[2] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.007522 2023.
[3] Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405. 2024.
[4] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5156-5165, PMLR, 2020.
[5] Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule. arXiv preprint arXiv:2412.06664 2024.
[6] Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory. arXiv preprint arXiv:2405.0451/ 2024.
[7] Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, al. Griffin: Mixing gated linear recurrences with local attention for efficient language models. arXiv preprint 2rr.ii:2222291422,,, 2024.
[8] Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences. In International Conference on Machine Learning, pages 16670-26698 PMLR, 2023.
[9] Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.051/111 2020.
[10] Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc Le. Transformer quality in linear time. In International conference on machine learning, pages 9099-99999 9119. PMLR, 2022.
[11] Zihang Dai, Zhilin Yang, Yiming Yang. Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. arXiv preprint arXiv:1901.02860 2019.
[12] DeLesley Hutchins, Imanol Schlag, Yuhuai Wu, Ethan Dyer, and Behnam Neyshabur. Block-recurrent transformers. Advances in neural information processing systems, 35: 33248-33261, 2022.
[13] Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng. Lei Hou, al. Longbench: bilingual, multitask benchmark for long context understanding. arXiv preprint ar0iv:2308. 14508, 2023.
[14] Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint arXiv: 2014.
[15] S Hochreiter. Long short-term memory. Neural Computation MIT-Press, 1997.
[16] Leo Feng, Frederick Tung, Mohamed Osama Ahmed, Yoshua Bengio, and Hossein Hajimirsadeghi. Were rnns all we needed? arXiv preprint arXiv:2410.01221, 2024.
10

---

[17] Eric Martin and Chris Cundy. Parallelizing linear recurrent neural nets over sequence length. arXiv preprint vrXiv:779.04005/ 2017.
[18] Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling, arXiv preprint N.XKi:2200.0493333 2022.
[19] Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. International conference on machine learning, pages 1310-1318. Pmlr, 2013.
[20] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputingg 568: 127063, 2024,
[21] Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact of positional encoding on length generalization in transformers. Advances in Neural Information Processing Systems, 36:2489222444 24928, 2023.
[22] Team Cohere, Arash Ahmadian, Marwan Ahmed, Jay Alammar, Milad Alizadeh, Yazeed Alnumay, Sophia Althammer, Arkady Arkhangorodsky, Viraat Aryabumi, Dennis Aumiller, et al. Command An enceprise-ready. large language model. arXiv preprint arXiv:2504. 006600 2025.
[23] Liliang Ren, Yang Liu, Yadong Lu, Yelong Shen, Chen Liang, and Weizhu Chen. Samba: Simple hybrid state space models for efficient unlimited context language modeling, arXiv preprint arXiv:2406.07755555 075522 2024.
[24] Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax. arXiv preprint 1805. 2018.
[25] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in neural information processing systems, 35:16344-16335 2022.
[26] Guilherme Penedo, Hynek Kydlíček, Anton Lozhkov, Margaret Mitchell, Colin A Raffel, Leandro Von Werra, Thomas Wolf, et al. The fineweb datasets: Decanting the web for the finest text data scale. Advances in Neural Information Processing Systems, 37:3081 1-30849, 2024.
[27] Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du, Zhe Fu, et al. Deepseek Ilm: Scaling open-source language models with longtermism. arXiv preprint arXiv:2401.02 1.000..0944..... 2024.
[28] Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac"h, h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. The lan- guage model evaluation harness, 07 2024. URL https:t////// / //tp///zz/ddo///// org/ http://coorrssee// seccords/12608602. 12608602.
[29] Tomáš Kočiský, Jonathan Schwarz, Phil Blunsom, Chris Dyer, Karl Moritz Hermann, Gábor Melis, and Edward Grefenstette. The narrativeqa reading comprehension challenge. Transactions of the Association for Computational Linguistics, 6:317-328, 2018
[30] Ming Zhong, Da Yin, Tao Yu, Ahmad Zaidi, Mutethia Mutuma, Rahul Jha, Ahmed Hassan Awadallah, Asli Celikyilmaz, Yang Liu, Xipeng Qiu, et al. Qmsum: A new benchmark for query-based multi-domain meeting summarization. arXiv preprini 1...i::2144009332 2021.
[31] Nachshon Cohen, Oren Kalinsky, Yftah Ziser, and Alessandro Moschitti. Wikisum: Coherent summarization dataset for efficient human-evaluuttinn. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing ( (Volume 2: Short Papers), pages 212-219, 2021.
[32] Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What's 's the real context size of your long-context language models? arXiv preprint NXV::2404.06654. 2024.
[33] Ken M Nakanishi. Scalable-softmax superior for attention. arXiv preprint arXiv:2501. 2025.
11

---

[34] Jeffrey L Elman. Finding structure in time. Cognitive science, 14(2): 221, 1990.
[35] Hava T Siegelmann and Eduardo D Sontag. On the computational power of neural nets. In Pro- ceedings of fifth annual workshop on Computational learning theory, pages 440-449, 1992.
[36] Yoshua Bengio, Patrice Simard, and Paolo Frasconi. Learning long-term dependencies with gradient descent is difficult. IEEE transactions on neural networks, 5(2):157-166, 1994.
[37] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces, arXiv preprint arXiv:211 2021.
[38] Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307.0866. 2023.
[39] Bo Peng. Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, et al. Rwkv: Reinventing rens for the transformer era. arXiv preprint arXiv:2305. /3048, 2023.
[40] Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transformers with the delta rule over sequence length. arXiv preprint arXiv:2406 06484, 2024.
[41] Han Guo, Songlin Yang, Tarushii Goel, Eric P Xing, Tri Dao, and Yoon Kim. Log-linear attention, arXiv preprint 250.::25000047761, 2025.
[42] Xin Dong, Yonggan Fu, Shizhe Diao, Wonmin Byeon, Zijia Chen, Ameya Sunil Mahabalesh- warkar, Shih- Yang Liu, Matthijs Van Keirsbilck, Min-Hung Chen, Yoshi Suhara, et al. Hymba: hybrid-head architecture for small language models. arXiv preprint ar
2024.
[43] Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig. Jonathan May, and Luke Zettlemoyer. Mega: moving average equipped gated attention. arXiv preprint arXiv:2209. 10655, 2022.
[44] Xuezhe Ma, Xiaomeng Yang, Wenhan Xiong, Beidi Chen, Lili Yu. Hao Zhang, Jonathan May. Luke Zettlemoyer, Omer Levy, and Chunting Zhou. Megalodon: Efficient I1m pretraining and inference with unlimited context length. Advances in Neural Information Processing Systems, 37:71831-71854, 2024.
[45] Amirkeivan Mohtashami and Martin Jaggi. Landmark attention: Random-access infinite context length for transformers. arXiv preprint arXiv:2305. 16300, 2023.
[46] Yu Zhang, Songlin Yang. Rui-Jie Zhu, Yue Zhang, Leyang Cui, Yiqiao Wang. Bolun Wang. Freda Shi, Bailin Wang, Wei Bi, et al. Gated slot attention for efficient linear-time sequence modeling. Advances in Neural Information Prucessing Systems, 1117666176888, 2024.
[47] Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, Chloe Hillier, and Timothy Lillicrap. Compressive transformers for long-range sequence modelling. arXiv preprint, 2019. URL https://arxiv. owgo/og//////.eooo abs/1911 05507.
[48] Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 7432-7439, 2020.
[49] Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv:1803.05457v1 2018.
[50] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? arXiv preprint arXivv1905.07830, 2019.
[51] Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
[52] Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. arXiv preprint 000000200000001, 2023.
12

---

[53] bloc97. Aware Scaled RoPE allows LLaMA models to have extended (8k+) context size without any fine-tuning and minimal perplexity degradation, 2023.
[54] Jie Wang, Tao Ji, Yuanbin Wu, Hang Yan, Tao Gui, Qi Zhang, Xuanjing Huang, and Xiaoling Wang. Length generalization of causal transformers without position encoding. arXiv preprint arXiv:2404.12224 2024.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We ensure that the abstract and introduction clearly summarize the paper's core contributions: the RAT architecture, its intra- and inter-chunk mechanisms, the efficiency, and empirical evaluations on both pretraining and downstream tasks. All major claims are supported by experiments and analyses throughout the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss limitations in Sec.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory assumptions and proof's
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
14

---

Answer: [NA]
Justification:
The paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We describe our experimental setup and implementation details in Sec. 4 and Appendix A.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access to a hosted model (e.g.. in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate the research performed.
While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature the contribution. For example
(a) the contribution primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution should either be a way to access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how construct the dataset).
new model (e.g., large language model), then there
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path reproducing verifying the results.
5. Open access to data and code
15

---

Question: Does the paper provide open access to the data and code, with sufficient instructions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We open-source our code with link put in Appendix. The dataset we use publicly available.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https:
for more details. While we encourage the release of code and data, understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source source benchmark).
hntpps.cc/publiceguideee SeaidddeeCooeSSumssssiiiiiti
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https://nips http://app.ccppplic//gg/iieeeeee
for more details. The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments should state which ones are omitted from the script and why.
reproducible, they
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyperpa- rameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification:
We describe our experimental setup Sec. and give details in Appendix A.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental
material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: We do not report error bars statistical significance due computational cost and limited variance observed across repeated runs in preliminary testing. We follow the common practice in similar-scale LLM experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
16

---

The method for calculating the error bars should be explained (closed form formula,
call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out of range (c.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Computing resources are discussed in Sec. 4 and Appendix A.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g.. preliminary failed experiments that didn't t make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: attps://neett http///eerrppp.eeeeeeeee
Answer: [Yes]
Justification:
The research conducted in the paper conforms with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consideration due laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Societal impacts are discussed Appendix D.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
17

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: We do not include safeguards in the current version, model and data release is not the focus of this work, and the model is only 1B scale. Should we release models in the future, will follow community standards mitigate misuse risk.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
All the assets used have been cited.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, possible, include a URI..
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperewithcode. thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
18

---

For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer:
[Yes]
Justification:
We provide documented code.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
19

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: LLMs were not used in any component of the core research methodology. Any usage was limited to writing assistance and did not affect the technical contributions of the work.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy httpe//errrrppcccooooeeeeeeeeee ccc//oonferrrrccc/222////// for what should or should not be described.
20

---

A Implementation details
A.1 Algorithm
We provide the pseudocode for the training and prefilling modes of RAT in Listing 1, and the pseudocode for the generation mode in Listing
def gagto__lss_ttke___eeee
(inter_out, intra_out, inter_lse, intra_lse):
# Nerge previous chunk's output and current chunk 's output safely # by adjusting the base of the esponential softmar max_lse maximum (inter_lse intra_lse)
inter_lse_exp exp(inter_1se max_lse)
intra_lse_exp exp(intra_1se max_lse)
denom iatra_lse_exp
injua_ (intra.lse_exp denom).... unsqueeze (-1)
injer_adjust (inter_lse_ex, / denom) unsqueeze (-1) (-1)
10
return inter_out inter_adjur + intra_ intra_adjust
13
def Srain_or_prefit
(z, 8, 9, k, v, num_chunk, chunk_size, num num_head,
softmax_scale).
14
Args:
16
2: (B,
D), output gate after sigmosd function
17
g: (B, I, D), forget gate after sigmoid function
18
v: (B, I, D), value vectors in attention
9. b: (B, I, D), shared query and key vectors in attention
Returns:
21
out: (B, 7. D)
22
23
bs, len d_model - shape (2)
k, v
rearrange (m, (c 1) a b d", chnum_chunk) for in (k, v)
25
26
# mntra-chund RиИ: apply associative scan along the "I" dimension intra_v = ascan(g, (1.0 g) * v)
intra_k ascan(g, (1.0 g) k)
29
30
# Rearrange for multi head attention intra_v
31
rearrange (intra_v, "b 1 (a p) -> b (c 1) p", a=num_head)
32
intra_k
Eeiraane(intta.. "b c (a p) - p", a=num_head) p) -> batp", a=num_head)
33
= rearrange "b t p)
34
# Apply seter-chunt RoPE
36
9, intra_k apply_rope (q, intra_k)
37
38
# Estract chunk's end representation for Unter-chund attention
Younk_intran = intra_k[.... chunk size 1: chunkk chunk_intra_v intra_v[.... chunk_size c::nkkk_aiize
41
42
# Inter chunk attention using flez_attentie attention
43
# with chunk- causal mask and returning log sum eap scores blaate_block_maat
block_mase lambda b, h, q_idx, kv_idx: q_idx // chunk_size kv_idx,
45
46
1, 1, shapee [2], num_chunk
47
inter_out, inter_las = flex_attention(
9. chunk chunk scale=softmax.sca scale, block_mes was.-mbocc..kk.kkk return_lse-True
50
51
52
21

---

# Compute logits with
the current chunk's h, required separately due causal masking intra_lse einsum einsum("batp) ("batp, batp -> bat", 9, intra_k) softmax_scale
56
Nerge outputs
from previous chunks and current chunk (due to causal masking) out = merge_ late__lstttttte__eeeee_
39
(inter_out intra_v, inter intra_lse) rearrange (out, "b P -> b (a p)", a=num_head) return out
61
Listing Pseudo onde training or prefilling modes of RAT. use Pytorch's flex attention and associative scan for implementation.
def gen (z, 8, k, Chunk_start num head, cache):
A A
2: (B, 1, D), output gate after sigmoid
g: (B, 1, D), forget gate after sigmoid
v: (B, 1, D), volue vector
g, b: (B, 1, D), shared query and key vectors across heads Shunk_starts int, indea of current chunk cache: tuple of four tensors lastkcache
10
(B, 1, D), last token in in current chunk (for RND style update) lastucache
: (B, D), last token v in current chunk (for RMB style update) keache: (B, A, C. P) or
12
(B, num head, num chunk, head_dim): cached k at each chunk 's end veache: (B, A, C, P)
(B, num head, num_ chunk, head dim): cached v each chunk end # Note: KVCache is stored per chunk, not per token.
14
15
Returns:
out: (B, 1, D)
18
lastkcache, lastvcache, keache, veache cache
19
20
# Intra- chunk AНЛ:: Recurrent update within the current chunk lastkcache (1.0 g) * k
intra_k
22
intra_v =g lastvcache (1.0 g)
25
lastkcache... copy_ (intra_k)
24
lastvcache. (intra_v)
# Rearrange input for attention
27
intra_k rearrange (intra_k, t (a p) babtti p", a-num_head) intra_v rearrange (intra_v, "b t (a p) >> b p" a-num_head)
28
29
= rearrange(q, "b t (a -> bat p", a=num_heas head)_hedd)
# Apply ROPE
32
9. intra_k apply_rope intra_k)
33
# Update the current chunk's kucache
kcache [: i, stukksstartcchnkk start + 1] intra_k
36
veache :, Chunksstart.chuk..coo
intra_v
37
38
# Inter-chunk
attention: standard attention over cached chunk representations
attn_out scaled_dov daaled_d
40
q, keache [i, :,
chunk_stare i], veachel;, i, :chunk + 1], is is_causal=Fale
41
rearrange (out, "b b (a ")", a=num_
22

---

43 return * out
Listing 2:
Pscudo code for the generation mode of RAT. We simply use the flash attention (Pytorch's 's onc). Note that KVCache of RAT is reduced from T to C' compared to the attention module.
A.2 Experiments
Dataset In our preliminary studies, we use the PG19 dataset [47], a long-form English book corpus with inherently long contexts. For the 1.3B model experiments, we adopt the FineWeb-Edu dataset [26], using its 100B-token randomly sampled version downloaded from the HuggingFace repository. To match the pretraining context length, we concatenate documents using separator token. Note that web samples are usually very short, compared to the book dataset. For downstream evaluation, consider suite of classical commonsense reasoning benchmarks from the Eleuther Al evaluation harness [28], including PIQA [48], ARC-C [49], and HellaSwag [50]. In the LLaMA2 tokenizer, we observe that the inputs of these tasks typically contain fewer than 300 tokens. For the I LongBench evaluation, in Table 7 we provide input length for each task, offering a rough indication of task difficulty with respect input length. For SFT-based tasks, we have elaborated sequence lengths in the main text. As LongBench and SFT tasks often involve very long inputs, we apply truncation in the middle to preserve information at both the beginning and the end.
Table We report the average input length and the 95th percentile input length of each task, measured in tokens using the LLaMA2 tokenizer. NQA (NarrativeAA), MQA (MultiFieldQA enn), HQA (HolpotQA), WQA (2Wiki- MultihopQA), MSQ (Musique), GR (GovReport), MN (MultiNews), TQA (TriviaQA), and RBP (RepoBench-P).
Task
Deecum-D QA
Multi- QA
Summarization
Code Completion
Name NQA
Qasper
MQA
HQA
WQA
MSQ
GR
QMSum
MN
1.CC
RBP
Avg.
36037
5780
8115
15329
8483
18555
12280
15980
3156
4307
14818
95th pctl.
77966
10164
14994
19755
16939
20066
25721
29069
7135
10401
31937
M00M model We start with 200M-parameter model in our preliminary study, with model dimension of 1024, 12 transformer layers, and head dimension of 64. The rotary position embedding (RoPE) base is set to 10,000. We use the GPT2 tokenizer. Following Mohtashami and Jaggi [45], we repeat the PG19 training split five times to reach a total of 15B training tokens. The learning rate is scheduled using cosine annealing, starting at 6.0 10 and decaying to 1.0 10-6, with warm-up ratio of 10%. We use the Adam W optimizer with a weight decay of 0.1 and B = (0.9,..98). Gradient clipping is applied with threshold of 1.0, and the global batch size is set to 1M tokens. We explore training with three different context lengths: 8K, 16K, and 32K. Training these models on 4 H100 GPUs takes approximately 14 hours when the sequence length is T= 32768, whereas the RAT (L=1) model completes training in about hours under the same setting.
14 hours. In particular, the attention model requires up to
1.3B model The 1.3B parameter model uses a model dimension of 2048, 24 transformer layers, and a head dimension of 128, equipped with RMSNorm [51] and without bias. The RoPE base is also set to 10,000. The model parameters are initialized using a Gaussian distribution with standard deviation of 0.02. We adopt the ILaMA2 tokenizer in the following studies. For pretraining, we use cosine-annealed learning rate schedule starting at 8.0 10-1 and decaying to 1.0x 10-6, with 5% warmup. The global batch size is set to 2M tokens, and the context window is set to 4096. Each model is trained on 16 H100 GPUs, requiring approximately 2 to 3 days to complete.
For LongBench evaluation, we follow the default prompts with greedy decoding for all tasks except summarization. For summarization, we apply a repetition penalty of 1.2 to address the common issue in pretrained-only models of generating repetitive outputs under instructional prompts. For SFT tasks, we train the models on the official training splits with an answer-only loss and evaluate them on the corresponding test sets. Although we explored different hyperparameters during the early experimentation, we observed that the relative trends remained largely stable. Thus, we fix the learning rate and batch size to (1.0x 10-5, 5,2) for large datasets, and (1.0 x 10-5.32) for the smaller QMSum [30] task, following common practice for 1B-scale models. The weight decay is
23

---

set as 0,01, and all other hyperparameters follow the pretraining setup. We sweep over the number of epochs and report the best result for each dataset and architecture. In practice, QA tasks typically converge in I 1-2 epochs, while summarization benefits from slightly longer training.
B
Efficiency
B.1 Latency
To supplement Fig. 2 in main text, we put the concrete latency number of single layer in Table 8, Table 9, and Table 10.
Table 8: Single temporal-min mimporal-mxxing layer (including input and output projections) training time across different sequence lengths. The latency (ms) tested 200K tokens.
Model
4096
8192
16384
32768
65536
131072
2621144
RNN
77.00
78.42
77.50
83.50
115.83
130.62
195.30
RAT (L-128)
150.19
150.76
156.65
154.52
159.94
165.07
206.46
RAT (L-64)
146.25
148.46
151.10
155.83
165.58
177.78
227.76
RAT (L-16)
150.08
153.63
158.98
173.89
202.00
249.79
378.11
Attention
122.06
176.44
277.61
474.90
877.14
1713.82
3417.48
(ttention/RAT(L=16))
0.81x
1.15x
1.75x
2.73x
4.34x
6.86x
9.04x
Table 9: Single temporal-mixing layer (including input and output projections) prefilling time across different sequences lengths. The latency (ms) tested on 200K tokens.
Model
4096
8192
16384
32768
65536
131072
262144
RNN
24.93
24.75
25.32
27.09
30.54
39.40
56.16
RAT (L-128)
44.74
44.78
45.09
45.50
46.67
48.62
51.75
RAT (L=64)
43.46
44.27
44.27
44.83
47.08
50.83
57.38
RAT (L-16)
44.03
44.90
45.53
51.08
57.31
71.18
99.53
Attention
36.93
52.71
80.21
135.68
245.14
494.62
997.50
Attention/RAT (L=16)
0.84x
1.17x
1.76x
2.66x
4.28x
6.95x
10.02x
10.2x
Attention
30000
Alteatinn
RATUL-I8),5WA
20000
12.8s
15000
4.5%
8x
15.6x
13.9%
1.5v
1.0x
1.0x 1.7x
1.0s 2.0%
Figure 4: We measure the maximum throughput of the full 1.3B model for generating 1024 tokens under different prefilling lengths. For each total sequence length 7', the pretilling length is set to 7' 1024. For example, ''= 4096 corresponds to prefilling of 3072 tokens, while 7'= 8192 and 7 16384 correspond to and 15360 tokens, respectively. When sequence length increases, the maximum throughput ratio between RAT and attention rises from x to 15.6x, highlighting the strong efficiency advantage of RAT in long-context generation.
B.2 Maximum throughput
As RAT also reduces cache memory usage, we report the maximum throughput of the full 1.3B model across different sequence lengths in Fig. 4. For example, RAT RAT(L=16) achieves an approximately

---

Table 10: Single temporal-mixing layer (including input and output projections) generation time at the specified position. The latency (ms) is tested generating batches of tokens with 64, = 512, and B 1024.
Model
4096
8192
16384
32768
65536
131072
262144
B 64
RNN
0.36
0.33
0.34
0.33
0.33
0.33
0.33
RAT (L=128)
0.62
0.60
0.61
0.63
0.66
0.66
0.68
RAT (L=64)
0.67
0.64
0.66
0.70
0.73
0.74
1.03
RAT (L=16)
0.62
0.65
0.64
0.68
1.04
1.82
3.42
Attention
1.15
1.73
3.33
6.66
12.94
26.8
OOM
B 512
RNN
0.35
0.34
0.33
0.33
0.34
0.34
0.33
RAT (L=1128)
0.63
0.64
0.66
0.75
1.17
2.44
3.46
RAT (L=64)
0.70
0.68
0.74
1.16
2.41
3.44
6.69
RAT (L=16)
0.75
1.22
2.45
3.48
6.73
13.09
25.76
Attention
6.56
12.89
25.84
OOM
OOM
OOM
OOM
B 1024
RNN
0.36
0.36
0.35
0.35
0.35
0.36
0.35
RAT (L=128)
0.73
0.75
0.95
1.35
2.54
4.64
6.98
RAT (L=64)
0.74
0.94
1.35
2.48
4.67
7.02
13.45
RAT (L=16)
1.38
2.56
4.63
7.04
13.51
26.43
OOM
Attention 13.04
25.70
OOM
OOM
OOM
OOM
OOM
10.2x higher throughput than the baseline attention model at T 4096, and improves further to 15.6x T== 16384. Similarly, RAT (L=16) SWA reaches 4.5x, 8.0x, and 13.9x higher throughput over attention at T 4096, 8192, and 16384, respectively.
C Accuracy
C.1 Supplementary ablation study
All experiments here are conducted the PG19 Book dataset with 200M-parameter models, and FLOPs are controlled for comparison.
Reversed design of RAT Instead of employing recurrence within each chunk, we also explored a reversed design that applies attention inside the chunk, followed by recurrence to capture long- range dependencies. For this reversed design, standard RoPE is sufficient. Regarding the parameter allocation, we found crucial to retain the key vector in the attention module as full tensor, as sharing together with the query vector leads to collapse into single-head attention. Thus, the final reversed design shares the query vector across attention heads, employs low-rank matrices for the projection used for output gate, and preserves full projections for both forget gate and the key vector.
Results of the reversed design are put in Table 11. It can be observed that, under identical FLOP constraints, this design significantly underperforms RAT. We think this is because short-context dependencies are easier to capture and do not suffer from memory degradation, so using lightweight recurrence locally is more efficient. Long-context dependencies are harder and may require direct retrieval of distant information, where attention is more suitable. For example, to match the FLOPs budget of (TT /64), the reversed design requires chunk sizes of 128, 256, and 512 for sequence lengths of 81 8192, 16384, and 32768, respectively. With this small local window size, the long-range dependencies are attributed to the recurrence, thus leading to high perplexity. Only when we reduce the number of chunks to 16 does the perplexity drop below 14.00. Therefore, while the reversed design also allows for interpolation between attention and RNN (reducing to attention when T.=T, and to RNN when == 1), we choose to focus RAT due to its more efficient utilization of FLOPs.
25

---

Comparison to FLASH [10] As pointed out in Sec. 5, Hua et al. [10] also uses the concept of chunking; however, we adopt a fundamentally different framework, both in the components involved and in how the two outputs are combined. Specifically, Hua et al. [10] simply adds the outputs of inter- chunk and intra-chunk computations, whereas RAT organizes them hierarchícally. applying recurrence over the key and value before inter-chunk attention. We argue that these differences lead to the weaker performance of Hua et al. [10], as shown in Table 13. First, as discussed above, RAT utilizes FLOPs more effectively than methods based on local attention, since its long-range dependencies must rely on memory-degrading components such as recurrence and linear attention. Second, we find that directly adding the outputs of softmax-bases attention and linear attention introduces differences in output scale and potential representational conflicts, which may lower down the performance.
Comparison to dilated attention We discuss the difference between RAT and the design that interleaves recurrence and dilated attention [9] with the sliding pattern. First, from an efficiency perspective, RAT has advantages over dilated attention in both training and inference. During training. dilated attention reduces parallelism by the dilation rate, since tokens within the same dilation span attend to different KV vectors. In contrast, RAT allows tokens to share the same chunk-level representation, except within their own chunk. During inference, the sliding dilation pattern causes a memory issue, since the KV cache of all tokens must be stored, whereas RAT reduces the cache size proportionally to the chunk size. Second, in terms of accuracy, to match the FLOPs, we set the dilation rate to 64 for RNN-Dilated attention, since dilated attention occurs in only half of the layers. It can be seen from Table 14 that dilated attention performs poorly. We observed significantly slower convergence at the beginning of training, likely due to its lack of global perception.
Table 11:
Perplexity results of the reversed design, where attention is applied within chunks and RNN is used across chunks. Experiments are conducted on the 200M model trained on the PG19 book dataset. Under the same FLOPs, it can be observed that RAT significantly outperforms its reversed counterpart.
Table 12: Pretraining and downstream evaluation re- sults of 1.3B models using either RoPE or NoPE posi- tional encodings. RoPE is used in the main text, while NoPE is trained for investigating length extrapolation. Notably, NoPE achieves reasonable performance even at the 1B model scale with only 100B training tokens.
Method
FLOPs
T-8192
'T-
VY-16344
T 32768
Method
PE
Pretrain 191.
HollaSwag an: nem
LAMBADA
PIQA
BA7(L-128)
O11/122)
13.57
13.44
13.66
acc
acr nor
(L-(4)
((T/64)
13.42
13.34
13.30
Atteniinn SWA
RaPE
7.61
$7.1
48.0
22.7
RAI (C-
OCT/ 128)
14.53
14.35
14.37
(LT(1-16)-
RaPF
7.57
$7.9
49.3
73.9
BAT (C-6)))
((1/64)
14.21
14.06
14,02
Ancmion-SWa SWA
NoPE
7.69
56.4
47.3
72.7
RAT-
O(T/16)
13,00
13,50
RAT (1T(1-16)554
NoPP
7.69
56.8
47.8
73.7
Attention
O(2)
13.26
13.10
13.19
Table 13: Performance of different chunk based designs with T T 16384 and FLOPs (64)
Table 14:
Performance compared to dilated attention with sliding pat- terns under the same FLOPPs.
Method
Intra-chunk
Inter-chunk
Organization
PPL.
RAT (L-64)
Recurrence
Attention
Hierarchically
13.34
Method
PPL.
RAT-Reversed (C+64)
Attention
Recurrence
Hierarchically
14.06
RNN-Dilated attention
15.22
FLASH (C=
Attention
Linear attention
Add together
14.4
RAT (L-128)
13.44
C.2 Length generalization
Because of the use of softmax- -based attention and RoPE at the inter-chunk level, is reasonable to expect that RAT may also face challenges in length generalization, where the model is trained on short sequences but evaluated on much longer ones. To study this, we consider SWA variants using either RoPE or NoPE for attention and RAT, following recent practices that interleave attention with local attention modules and apply positional encodings only in the local attention. This design has been adopted in practice and has shown promising results [22].
As shown in Fig. 5, the RNN performs very steadily as the test sequence length increases. However, within the training context, it has the highest loss among all models. With RoPE, both Attention SWA and RAT (L-()) -SWA experience a sharp increase in loss when the sequence length goes beyond 6000. The increase is more severe in the attention model than in the RAT variant. This is likely because RAT reduces the effective attention span by attending only to inter-chunk positions, which may offer better robustness to extrapolation.
26

---

3.8
Attention-SWA (RoPE)
RAT(L-16)-SWA (RoPE)
Attention-SWA (NoPE)
3.4
RAT(L-16)-SWA (NoPE)
3.2
Test
3.0
2.8
2.6
2000
4000
6000
8000
10000
12000
14000
16000
Test Length
Evaluation at different test lengths for pretrained models trained with a 4K context window. RAT (L-16) -SWA with NoPE achieves the best overall performance, exhibiting strong generalization up to T= 16384 while maintaining low loss within the training context.
Figure 5:
We also pretrain NoPE versions, where no positional encodings are used in attention or RAT layers. Interestingly, these NoPE models show reasonable performance, reported in Table 12. While their pretraining perplexities are higher than those of the RoPE variants, the gap remains small, especially considering the 1.3B model size and 100B-token training budget. To further verify this, we evaluate the models on several commonsense reasoning tasks and obtain promising results. We expect that NoPE will perform length extrapolation, using NoPE significantly improves the robustness for both Attention-SW. and RAT (L=16) SWA. Among them, (LT with NoPE shows the most stable performance, maintaining low loss even at sequence length of 16,384.
better with larger models and longer training. In terms of
In conclusion, find that RAT outperform the attention module in length generalization, both with RoPE and with NoPE. Still, the problem is not fully solved, as none of the models reach the stability level of the simple RNN. While many techniques have been proposed to improve length extrapolation in attention models, they should also be considered for RAT, as RAT essentially redirects attention from every position inter-chunk locations. For example, RoPE extrapolation in the attention module can be improved using methods from Peng et al. [52], bloc97 [53], and NoPE extrapolation has been that NoPE still has context length limit, though it
studied in Wang et al. [54], which points performs better than RoPE. Their work links the failure to shifts in attention distributions and proposes tuning the temperature of attention heads improve extrapolation. We believe such methods could also be applied to RAT, and leave further exploration of these techniques to future work.
C.3 Needle-in- Haystack (retrieval ability)
To evaluate retrieval capabilitics an area where attention crenntion-based architectures are known to excel, and where RNNs typically fall short-we conduct experiments the Needle-in-Haystace tasks. In these tasks, a "needle" (c.g.. a magic number or UUID) is embedded within irrelevant passages, noisy text, mixed with multiple key-value pairs. Each key-value pair is presented short identifier (the key) followed by its associated value, and the model is prompted to retrieve the correct value given specific key.
We adopt the RULER benchmark [32] and evaluate a range Needle-in-Haastaal Haystack task configurations. Specifically, task single_1 involves retrieving specific number from background filled with repeated noise. In single_2 2, the background replaced with natural stories. single_3 increases the difficulty by requiring retrieval of long and complex UUID. The multikey tasks more
27

---

challenging than the single-key ones, they introduce multiple such key-value pairs into the context. In particular, multikey_2 and multikey_S consist almost entirely of densely packed key-value pairs, among which only a single key is queried at the end. This makes the task especially challenging. as the model must retrieve the correct item from highly cluttered input full of distractors. And multikey: further incorporates the complex UUID format. In the multi muetiquery and multivaluu settings, the model is required to resolve multiple retrieval targets, either by answering several distinct queries or by retrieving all values associated with single key. For detailed definitions of each task, we refer the reader to the benchmark documentation.
During evaluation, we observed that models may fail to interpret certain prompts correctly. For instance, prompts like "A special magic number is hidden within the following text. Make sure to memorize it. I will quiz you about the number afterwards. can cause failures even in attention- based models, especially on the harder multikey tasks. To mitigate this, we apply a light, one-round supervised fine-tuning stage before evaluation to adapt the models to the instruction patterns. We generate 1000 synthetic training samples for each of the 8 tasks, resulting a total of 8000 examples disjoint from the validation sets. The models are trained on this dataset for one round and then are evaluated directly on all 8 tasks. This procedure yields a fairer and more stable comparison. As shown in Table 6, RNNs perform reasonably well on the simpler tasks (single_1. single_ e_2), but their performance drops to near zero on the harder ones. Attention based models perform consistently well, benefiting from their ability to access all previous tokens directly. RAT (L=16) achieves results close to full attention, especially in numeric tasks. However, it struggles on UUID tasks due to their complexity and the use of exact match scoring (single_3 3, multikey_53 Mcanwhile, RAT (L=64) falls between RNNs and RAT (L=16), as expected given its partial access to long-range context.
These results are consistent with the underlying architecture designs: attention provides full direct access to past tokens, RNNs compress all past information into a hidden state, while RAT compresses part of the history but also retains direct access at the chunk level. As result, RAT naturally exhibits retrieval capabilities that lie between those of attention and RNNs.
D Broader Impacts
Enhancing the efficiency of Large Language Models (LLMs) significantly reduce computational resources and energy consumption, benefiting the environment and democratizing access to advanced Al technologies. However, increased efficiency could also lead to greater dissemination of disinformation and the creation of deepfakes, posing risks to public trust and security and potentially reinforcing existing biases that impact specific groups unfairly. This research aims promote the responsible development and deployment of LLMs, maximizing societal benefits while acknowledging potential harms.
E
License information
FineWeb-Edu (dataset); Open Data Commons License Attribution family.
Link: https://hugginnfaae htt.a/cc.ccatttttttttttuaaaaaaaa
LongBench (dataset and code): MIT License.
Link: httppp///giiiub...oo/o/////////p
NarrativeQA (datastt)) Apache License 2.0.
Link:
QMSum (dataset): MIT License.
Link: http:////hhhhaaaooco/aaaa111////
WikiSum (dataset): Custom license (unspecified).
Link: https://hugginngface http:p/:/waggggggg..eeeeaaaattts co/fattasets//dd:j//kkk..oo
Im-evaluatiion-hamess (code)) MIT License.
Link: htt.://wo...oootaaeeeeeoo/oooooa
RULER benchmark (code): Apache 2.0 License.
Link:
28


Reviews:
[{"review_id": "E33CclQhp1", "reviewer": "Reviewer_CiSa", "summary": "This paper introduces RAT (Recurrence-Attention Transition), a novel architecture that interpolates between the efficiency of RNNs and the expressiveness of attention by chunking sequences. \nWithin each chunk (size L,16 in the paper), a lightweight gated linear RNN captures local dependencies, while softmax attention spans chunks to model long-range interactions. \nThis design reduce the burdens of the quadratic complexity of full self-attention (O(T²)) by reducing inter-chunk attention to O(C·D) flops, where C = T/L is the number of chunks. RAT flexibly trades off efficiency and accuracy via L: L=1 mimics attention, while L=T collapses to an RNN.\n\nExperiments on 1.3B-parameter models demonstrate 7× faster training (1M-token sequences) and 9× faster generation (4K-length) with L=16, while maintaining perplexity within 0.06 of full attention. \nOn short-context benchmarks (e.g., HellaSwag, PIQA), RAT w/ L 16 trails attention by <2 points but improves with hybrid RAT-SWA (sliding-window attention + RAT), gaining +1 point on commonsense tasks. On long-context LongBench, RAT-SWA outperforms attention on 12/14 tasks, including +4 points on code completion.", "questions": "* I'm curious if the author try other recent RNN variants in RAT, for example, HGRN and RWKV4?\n* In terms of param allocation, if I understand correctly, the final layer would have 6d^2 params with shared q/k? Did the authors tied key and forget gate param allocation as in HGDN2?\n* It would make the paper more comprehensive if the authors add some discussions with GAU in terms of local/global chunk design, and Log Linear Attention in terms of O(log N) caching. Some recent works, e.g., GSA, also makes some efforts in combing attention with RNNs\n* I'm especially curious what if combining RAT with some state expansion techniques like in Mamba or Hedgehog, which intuitively could make the 1d RNNs more powerful once expanded to 16x or larger state size\n\nHGRN: https://arxiv.org/abs/2311.04823\n\nGAU: https://arxiv.org/pdf/2202.10447\n\nGSA: https://arxiv.org/abs/2409.07146\n\nHedgehog: https://arxiv.org/abs/2402.04347", "limitations": "yes", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "I find the RAT (Recurrence-Attention Transition) architecture to be an elegant and natural solution for balancing expressiveness and efficiency in long-context modeling. Unlike recent works (e.g., Log-Linear Attention), which impose hierarchical inductive biases with O(log N) memory growth, RAT’s chunking strategy feels more intuitively grounded:\n* At 1M context, RAT with L=16 retains a fixed 64K KV cache, a constant-sized compression that scales gracefully without the logarithmic overhead of alternatives (log₂(1M)/log₂(4K) ≈ 2× growth in Log-Linear Attention).\n* Hybridization with SWA further reinforces its practicality: local attention handles short-range dependencies (where RNNs might underperform), while RAT’s inter-chunk attention preserves global coherence. This synergy outperforms pure attention or SSMs on LongBench, suggesting the architecture generalizes well beyond synthetic benchmarks.\n\nAnd I am particularly eager to see the results of RAT in terms of param and context scaling (if possible):\n* will RAT’s efficiency gains persist under larger model regimes like 7B/14B?\n* how does RAT fare against other RNNs like Mamba-2/HGRN2/DeltaNet/GDN on retrieval-heavy tasks?", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their constructive feedback. Following is our responses.\n> Q1: will RAT’s efficiency gains persist under larger model regimes like 7B/14B?\n\nA1: In response to the reviewer’s question we provide a detailed efficiency study which we will add to the appendix of our paper:\n\n  * First, as the model dimension D grows, the latency of the temporal mixing block is increasingly dominated by the input and output linear projections, so the efficiency gains during compute-intensive training diminish. For instance, when increasing D from 2048 to 4096, the speed-up drops from 7× to 5.5× for 100K tokens (Attention block: 1776.51s vs. RAT(L=16) block: 322.43s). However, training a 7B model usually involves tensor parallelism, which effectively reduces the dimension per GPU (e.g., D=2048 with 2 GPUs), allowing the expected speed-up to still be achieved.\n\n  * Second, the scaled dot-product operation of attention remains a major bottleneck in memory-intensive decoding, even for large models and especially in the long–context settings. We report the latency of a single temporal mixing block with D = 4096 for generating a batch of tokens at a specific position and the maximum throughput of the full 7B and 13B models for generating 1024 tokens with a 3072-token prompt. RAT shows a significant speed-up over attention. In particular, on the 13B model, the maximum batch size of attention is very limited, leading to poor GPU utilization and allowing RAT to achieve even greater throughput gains.\n\n| Latency of Temporal mixing block (D=4096) | Position=4096, Batch=1024 | Position=8192, B=512 | Position=16384, B=256 | Position=32768, B=128 | Position=65536, B=64 |\n|-|-|-|-|-|-|\n| RAT(L=16)   | 3.01s| 2.7s | 2.25s | 2.21s  | 2.15s  |\n| Attention   | 26.16s  | 25.82s | 25.87s  | 25.8s   | 26.22s |\n| Ratio | 8.7x  | 9.6x | 11.5x| 11.7x   | 12x |\n\n| Maximum throughput of a whole model | 1.3B (results in our paper) | 7B              | 13B            |\n|-|-|-|-|\n| RAT(L=16) | 31170 token/sec | 10103 token/sec | 5749 token/sec |\n| Attention | 3152 token/sec  | 983 token/sec   | 534 token/sec  |\n| Ratio  | 10.2x     | 10.3x           | 10.8x          |\n\n> Q2: how does RAT fare against other RNNs like Mamba-2/HGRN2/DeltaNet/GDN on retrieval-heavy tasks?\n\nA2: Thanks for this suggestion. We present comparisons on the Ruler Benchmark, which contains challenging synthetic tasks for retrieval ability (NIAH tasks), as well as QA tasks in LongBench. We compare our model against the results from the recent paper [1], which trained several linear attention models on the same dataset, with the same model size and sequence length as ours. As shown in the two tables below, we achieve better performance with the same or even smaller value/state size. We think the improvement comes from RAT’s direct access to the past, which can be beneficial in noisy backgrounds compared to other models that compress all history together.\n\n|Model  | value/state size (T=4096, D=2048, N=24) | NIAH-1 | NIAH-2 | NIAH-3 |\n|-|-|--------|--------|--------|\n| DeltaNet       | 128ND  | **99**     | 18.6   | 22.4   |\n| Mamba2         | 256ND  | 65.4   | 56.2   | 4.6    |\n| Gated DeltaNet | 288ND                                      | 91.4   | 92.2   | 27.6   |\n| RAT(L=16)      | 256ND                                      | 97     | **97.6**   | **56.4**   |\n\n| Model         | value/state size (T=4096, D=2048, N=24) | NQA  | Qasper | MQA  | HQA  | WQA  | MSQ |\n|---------------|-----------------------------------------|------|--------|------|------|------|-----|\n| Mamba2        | 256ND   | 11.1 | 11.3   | 18.6 | 11.8 | 15.1 | **6.7** |\n| DeltaNet      | 128ND    | 12.9 | 10.8   | 21.5 | 10.9 | 13.2 | 5.1 |\n| GatedDeltaNet | 288ND    | 14.1 | 14     | 23.3 | 13.7 | 14.4 | 5.8 |\n| RAT(L=16)     | 256ND    | **14.5** | **16.1**   | **28.3** | **16.7** | **18.9** | 6.3 |\n\nSpecifically, the value/state size calculation is based on sequence length T, model dimension D, and number of layers N. RAT (L=16) has 4096/16ND = 256ND value vectors. For other models, we multiply the value dimension with the state expansion factor. Note that Mamba2 and GatedDeltaNet expand the value dimension compared to standard attention. Finally, for the Ruler Benchmark (NIAH tasks), we would like to mention that in Appendix C.4 and Table 12, we compared RAT with attention in terms of retrieval ability and obtained comparable performance after light fine-tuning for prompt adaptation. Here, we directly evaluate RAT’s performance in the same way as theirs to ensure a fair comparison.\n\n[1]. Gated Delta Networks: Improving Mamba2 with Delta Rule. ICLR’25\n> Q3: I'm curious if the author try other recent RNN variants in RAT, for example, HGRN and RWKV4? I'm especially curious what if combining RAT with some state expansion techniques like in Mamba or Hedgehog, which intuitively could make the 1d RNNs more powerful once expanded to 16x or larger state size.\n\nA3: We tried this in our early experiments with a design inspired from  Mamba2 architecture, but found that using state expansion and reduction in the intra-chunk RNN, while enhancing expressiveness, introduced additional optimization difficulties. The experiments were conducted on a 4-layer model on WikiText-103 with model dimension 256, sequence length 256, and chunk size 64. These issues might disappear in larger models and datasets. And the added expressivity could become more important for very long chunks (e.g., 2k or 4k), but as an initial step, we primarily focus on shorter chunks and leverage softmax attention more, which currently dominates modern architectures. In detail, we considered two cases:\n\n  * State expansion → cross-chunk attention → state reduction: We applied state expansion on the “key” and “value” vectors and then performed softmax-based attention on the expanded states across chunks, followed by state reduction. We observed that the loss curve flattened early and remained at a higher loss level than its 2-D RNN counterparts. We think this is because involving many matrix multiplications without intermediate non-linearity can make optimization more challenging.\n  * State expansion → reduction → cross-chunk attention: Here, we applied expansion and reduction within each chunk and then performed inter-chunk attention. This approach only performs well when the state size is relatively small (e.g., 4 or 16). Given 64 as the chunk size, it is possible that the extra expressiveness of even larger state size is unnecessary and instead makes optimization harder.\n\nUnder the scope of RAT, 2-D linear recurrence could be explored  in the future for large chunk-length cases.  We also mentioned in the paper that non-linear RNNs could be explored for short chunk-length cases because the computational efficiency is no longer an issue there. For now, to introduce the idea, we use the simple and fast 1-D linear recurrence in the paper.\n\n> Q4: In terms of param allocation, if I understand correctly, the final layer would have 6d^2 params with shared q/k? Did the authors tied key and forget gate param allocation as in HGDN2?\n\nA4: The final layer is still 4d^2 parameters, and we did not tie the key and forget gates. We share “query” and “key” across different heads. And due to the existence of the per-dimension forget gate over key, such a design won’t collapse into a single-head attention. Thus, there’s the d^2 parameters for forget gate, output gate, value vector, and output projection, respectively. Query and key vectors here only require a small portion of parameters. It is D*256 in our 1.3B model with D=2048.\n\n> Q5: It would make the paper more comprehensive if the authors add some discussions with GAU in terms of local/global chunk design, and Log Linear Attention in terms of O(log N) caching. Some recent works, e.g., GSA, also makes some efforts in combining attention with RNNs.\n\nA5: Thanks for the suggestions. Here are some discussions and we’ll add them in the revision. \n  * Yes, we cited GAU in the paper and described its use of local attention. In the revision, we will further clarify that GAU applies softmax-based attention within local chunks and uses linear attention across chunks to cache history. The two outputs are simply added together, whereas we use recurrence within chunks and softmax-based attention across chunks and organize their outputs in a hierarchical manner.\n\n  * The concurrent work Log-Linear attention focuses on the fixed state size limitation of linear attention by increasing the memory cache size in a logarithmic manner. Our work also targets the fixed memory size issue in recurrent models, which we address by applying softmax-based attention sparsely to enable direct access to distant tokens.\n\n  * As for the interesting paper GSA, sorry for missing it. GSA aggregates and caches all “key” and “value” vectors into a 2D memory matrix, where the “query” attends to the compressed key and value representations. Compared to standard attention, this approach can be implemented as a two-pass linear attention method and is thus more efficient. A key difference from our method is that we do not aggregate the entire sequence’s keys and values; instead, we store each chunk’s information through a simple recurrence and use cross-chunk softmax-based attention to directly retrieve history. This design enables a flexible state/memory size depending on sequence length and allows for efficient implementation via FlexAttention."}, {"review_id": "bCqR0XhJSw", "reviewer": "Reviewer_aCcd", "summary": "The paper proposes RAT (Recurrent-Attention Transformer), an architectural design that bridges RNN efficiency and Transformer attention accuracy for language modeling. RAT partitions the input sequence into fixed-size chunks, applies a lightweight gated RNN within each chunk for capturing local dependencies, and then performs softmax attention across chunks to model long-range interactions. This design interpolates between standard RNNs and full attention depending on the chunk size, allowing a trade-off between efficiency and expressiveness.\n\nRAT models, especially with a chunk size of 16, achieve substantial speedups—up to 9× faster generation at 4K length—while matching or surpassing attention baselines on both pretraining perplexity and downstream tasks. Moreover, a hybrid variant interleaving RAT with sliding-window attention improves both efficiency and accuracy across commonsense reasoning, long-context understanding (LongBench), and supervised fine-tuning tasks, outperforming full-attention models in several benchmarks.", "questions": "See weaknesses.", "limitations": "N/A", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n\n1. Proposes a new intermediate architecture (RAT) that blends the efficiency of RNNs with the accuracy of attention mechanisms.\n\n2. The hybrid variant (RAT interleaved with sliding-window attention) further boosts performance, showing strong results on reasoning, summarization, and code tasks.\n\nWeaknesses:\n\n1. Missing strong baselines: While the paper discusses the landscape of linear and local attention methods, it only benchmarks against vanilla RNNs and standard attention. This makes the advantages of RAT less convincing, as prior work—such as hybrid architectures mixing linear and softmax attention (e.g., Hymba [1])—has shown similar trade-offs. Including such comparisons is necessary to establish RAT’s utility over existing efficient attention methods.\n\n2. Sensitivity to chunk size: RAT’s performance declines with larger chunk sizes, requiring careful tuning to balance accuracy and efficiency, particularly for long-context or noisy inputs.\n\n3. Lack of large-scale evaluation: The study is limited to 1.3B parameter models, leaving open questions about how well RAT scales to larger models used in practical deployment.\n\n[1] Hymba: A Hybrid-head Architecture for Small Language Models", "was_revised": true, "final_justification": "Overall, the rebuttal strengthens the work, and I am increasing the score to 4. I still consider this paper borderline. The chunking concept appears similar to the approach in the cited Transformer Quality in Linear Time (see Fig. 4). The authors should clearly articulate the differences, analyze and justify the superiority of RAT’s design, and include that work as a baseline for comparison. Another round of revision incorporating the rebuttal results and additional analysis would be needed to further strengthen the paper.", "rebuttal": "We thank the reviewer for the feedback and suggestions. Here are the responses.\n> Q1: Missing strong baselines: While the paper discusses the landscape of linear and local attention methods, it only benchmarks against vanilla RNNs and standard attention. This makes the advantages of RAT less convincing, as prior work—such as hybrid architectures mixing linear and softmax attention (e.g., Hymba [1])—has shown similar trade-offs. Including such comparisons is necessary to establish RAT’s utility over existing efficient attention methods.\n\nA1: To address the reviewer’s concerns, we clarify RAT's relationship to hybrid models and provide comparisons to efficient attention methods. RAT uses a simple recurrence within each chunk and applies softmax-based attention across chunks.\n\n  * Thus, RAT is not a hybrid model. It applies the same mechanism across all heads, tokens, and layers, and over time, while hybrid models like Hymba alternate state-space modules and attention in different heads and Griffin [1] alternates them in different layers. The hybrid strategy is orthogonal to RAT. In this paper, we explored interleaving RAT with sliding-window attention (RAT-SWA). Future work could adopt Hymba-like strategies, such as using RAT in some heads and attention in others, or applying RAT with different chunk sizes in different heads.\n\n  * Second, we mainly compared RAT to vanilla RNNs and standard attention because RAT directly interpolates between these two models and even achieves comparable or better performance than the strong baseline attention. Although RAT is also not a linear attention model (the simple recurrence uses per-dimension gating rather than per-head gating, the softmax-based attention is used across chunks), we noted in the paper (lines 245–246 with footnotes) that our best model outperforms strong baselines reported in a recent linear attention work [2], where the models below were trained on the same datasets, model sizes, and sequence lengths as ours. For completeness, we post detailed results below and will add them in the revision. \n| Model                   | value/state size (T=300, D=2048, N=24) | arc_challenge | arc_easy     | hellaswag    | lambada | piqa         | winogrande | boolq |\n|-|-|-|-|-|-|-|-|-|\n|    |     | acc/acc_norm  | acc/acc_norm | acc/acc_norm | acc     | acc/acc_norm | acc        | acc   |\n| Mamba1   | 64ND        | -/35.40       | 69.52/-      | -/52.91      | 43.98 | 71.32/-      | 52.95      | 61.13 |\n| Mamba2  | 256ND    | -/37.88       | 72.47/-      | -/55.67      | 45.66 | 71.87/-      | 55.24      | 60.13 |\n| DeltaNet  | 128ND    | -/35.66       | 68.47/-      | -/50.93      | 42.46 | 70.72/-      | 53.35      | 55.29 |\n| GatedDeltaNet    | 288ND  | -/38.39       | 71.21/-      | -/55.76      |46.65 | **72.25**/-      | 57.45      | 60.24 |\n| RAT(L=16)  | ~20ND   | 36.35/**39.33**   | **72.64**/67.80  | 43.27/**56.44**  | 44.4    | 72.03/72.69  | 53.2       | **63.3**  |\n| **Hybrid**     |          |         |    |      |         |     |      |       |\n| GatedDeltaNet-SWA(2048) | 288N/2D-300N/2D      | -/40.10       | 71.75/-      | -/56.53      | 47.73   | 72.57/-      | **58.4**       | 63.21 |\n| RAT(L=16)-SWA(1024)     | ~20N/2D-300N/2D     | 37.20/**40.78**   | **72.56**/68.22  | 44.33/**57.85**  | **49.33**   | **73.29**/73.94  | 56.91      | **63.21** |\n\n\n    These commonsense reasoning tasks have very short contexts (≤300 tokens), so RAT with L = 16 at most has 20ND value size. Even so, it surpasses other models in 4 of 7 tasks, and RAT-SWA outperforms GatedDeltaNet-SWA in 6 of 7 tasks. Note that the current RAT has a clean structure without using extra short convolutions or normalization layers, which other works often adopt to boost performance.\n\nHope these can address the reviewer’s concern.\n\n[1]. Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models\n\n[2]. Gated Delta Networks: Improving Mamba2 with Delta Rule. ICLR’25\n> Q2: Sensitivity to chunk size: RAT’s performance declines with larger chunk sizes, requiring careful tuning to balance accuracy and efficiency, particularly for long-context or noisy inputs.\n\nA2: Thanks for the question. In this paper, we showed that using a short chunk size with the simple recurrence is both performant and highly efficient. The chunk design in RAT can also help with long-context and noisy-input cases compared to recent state-space and linear attention models. We’ll elaborate on each point below.\n\n  * First, it is natural that performance declines with larger chunk sizes, as FLOPs also decrease. Other models have similar trade-offs (e.g., the state expansion factor in state-space models). In the paper, we showed that L = 16 yields 7×/9× speed-ups while achieving comparable results with the strong attention baseline on challenging tasks, including LongBench (Table 3), SFT (Table 4), and the Ruler benchmark (Table 12 in Appendix).\n  * Second, when the context becomes very long, the number of chunks (memory slots) in RAT grows with sequence length. This helps maintain accuracy while preserving the same FLOPs reduction ratio, while state-space or linear attention models, which has fixed memory slots, can suffer from memory degradation on very long inputs. For noisy inputs, the cross-chunk softmax attention in RAT directly retrieves past information rather than compressing all history together (as in Mamba2), which we believe is advantageous. We report results on the retrieval-heavy tasks Needle-in-Haystack tasks and QA tasks in LongBench, which contain very noisy backgrounds. It can be seen that RAT outperforms others in most tasks.\n\n| Model          | Value/state size(T=4096, D=2048, N=24 (layer)) | NIAH-1 | NIAH-2 | NIAH-3 |\n|-|--|-|-|--|\n| DeltaNet       | 128ND      | **99**     | 18.6   | 22.4   |\n| Mamba2         | 256ND| 65.4   | 56.2   | 4.6    |\n| Gated DeltaNet | 288ND  | 91.4   | 92.2   | 27.6   |\n| RAT(L=16)| 256ND   | 97     | **97.6**   | **56.4**   |\n\n| Model  | Value/state size (T=4096, D=2048, N=24) | NQA  | Qasper | MQA  | HQA  | WQA  | MSQ |\n|-|-|-|-|-|-|-|-|\n| Mamba2  | 256ND| 11.1 | 11.3   | 18.6 | 11.8 | 15.1 | **6.7** |\n| DeltaNet | 128ND  | 12.9 | 10.8   | 21.5 | 10.9 | 13.2 | 5.1 |\n| GatedDeltaNet| 288ND | 14.1 | 14     | 23.3 | 13.7 | 14.4 | 5.8 |\n| RAT(L=16)  | 256ND  |**14.5**| **16.1**   | **28.3** | **16.7** | **18.9** | 6.3 |\n\nTherefore, we think that the chunk design proposed in RAT is meaningful and important. Currently, with the simple recurrence, performance can drop at very large chunk sizes as the model approaches vanilla RNN behavior. Future work may explore more advanced RNNs for different chunk sizes, such as using state-expanded recurrence for large chunks and nonlinear RNNs for short chunks when computational efficiency is no longer an issue.\n\n> Q3: Lack of large-scale evaluation: The study is limited to 1.3B parameter models, leaving open questions about how well RAT scales to larger models used in practical deployment.\n\nA3: Thanks for pointing this out, and we have highlighted it as one of the limitations in our paper. Since we are introducing a novel architecture in this work, we prioritized model and architecture ablations (including the core idea, parameter allocation, positional encoding, and efficient implementation) and benchmarked RAT across multiple tasks. We believe these provide more insight into how the new architecture works, especially given the limited compute resource available to us.\n\nAlthough we did not train 7B or 13B models in this paper, we provide speed-related results on larger models below to demonstrate its efficiency potential at scale. We also noted in the paper that RAT can be easily combined with tensor parallelism (as the core part is head-independent) and context parallelism, making the method scalable to larger settings. With reduced number of key and value, it might even avoid the need for ring attention in context parallelism and thus further saves the communication cost.\n\nThe first table below tested the generation latency with 7B model dimension 4096 at a specific position, yielding 8-12x speed-up. The second table shows the maximum throughput of the whole model by generating 1024 tokens given the prompt of 3072. In particular, on the 13B model, attention utilizes GPU resources poorly due to the very limited maximum batch size, allowing our method to achieve even greater improvements in terms of maximum throughput.\n| Latency of Temporal mixing block (D=4096) | Position=4096, Batch=1024 | Position=8192, B=512 | Position=16384, B=256 | Position=32768, B=128 | Position=65536, B=64 |\n|-|-|-|-|-|-|\n| RAT(L=16)   | 3.01s| 2.7s | 2.25s | 2.21s  | 2.15s  |\n| Attention   | 26.16s  | 25.82s | 25.87s  | 25.8s   | 26.22s |\n| Ratio | 8.7x  | 9.6x | 11.5x| 11.7x   | 12x |\n\n| Maximum throughput of a whole model | 1.3B (results in our paper) | 7B              | 13B            |\n|-------------------------------------|-----------------------------|-----------------|----------------|\n| RAT(L=16)                           | 31170 token/sec             | 10103 token/sec | 5749 token/sec |\n| Attention                           | 3152 token/sec              | 983 token/sec   | 534 token/sec  |\n| Ratio                               | 10.2x                       | 10.3x           | 10.8x          |"}, {"review_id": "K3WbW2GNeK", "reviewer": "Reviewer_pHG2", "summary": "Attempts to bridge the efficiency of RNNs (actually gated linear recurrence) and the accuracy of attention mechanisms for language modeling.\n - RAT divides input sequences into chunks, applying lightweight gated linear recurrence within each chunk to capture local dependencies and softmax attention across chunks for long-range interactions.\n- RAT interpolates between RNN and attention behaviors, balancing computational cost and performance.\n- RAT achieves significant efficiency gains: up to 7× training speedup for 1M-token sequences and 9× generation speedup at 4K tokens, with reduced KV cache usage\n- Empirically, RAT matches or surpasses standard attention on benchmarks like LongBench (e.g., 24.9 vs. 18.6 on GovReport) and commonsense reasoning tasks.\n- A hybrid RAT-SWA variant, interleaving RAT with sliding-window attention, further boosts performance (e.g., +4 points on code tasks, +1 Rouge-L on summarization).", "questions": "- I am surprised at the 7x training speedup and 9x generation speedup, and I am curious what the baseline comparisons are. Was flash attention used?", "limitations": "None.", "rating": 2, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 1, "originality": 2, "strengths_and_weaknesses": "Strengths:\n- Achieves up to 7× training speedup (1M tokens) and 9× generation speedup (4K tokens), with reduced KV cache usage, minimizing OOM risks. However, these efficiency gains may fade at a larger scale when MLP computation dominates.\n\nWeaknesses:\n- One big issue I have is that Eqn. 2. is a specific case of Mamba, and that this proposal for a new sequence mixing method is a mix of attention and this EMA construction. \n  - A more meaningful comparison then should also be between Mamba-Transformer Hybrid architectures, both in terms of performance on downstream tasks, and on efficiency metrics. \n- Because of the EMA-style construction of the \"RNN\" part, this also does not afford the model the capability to do state-tracking: https://arxiv.org/pdf/2404.08819.\n   Since I suspect the speedups gained in the paper may fade with 1) different implementation comparisons, or 2) scale of model, I think there is little gain in capability if we consider the gains from using a RAT-Attention hybrid vs a Mamba-Attention hybrid.", "was_revised": true, "final_justification": "I've made minor modifications to the rating based on the author response, but ultimately there are some fundamental limitations of the paper.", "rebuttal": "We thank the reviewer for the feedback. Here are the responses.\n\n> Q1:  One big issue I have is that Eqn. 2. is a specific case of Mamba, and that this proposal for a new sequence mixing method is a mix of attention and this EMA construction. A more meaningful comparison then should also be between Mamba-Transformer Hybrid architectures, both in terms of performance on downstream tasks, and on efficiency metrics.\n\nA1: We would like to clarify an important point which might have been misunderstood for RAT. RAT first chunks the inputs, applies recurrence within each chunk to aggregate “key” and “value” vectors, and then applies cross-chunk nonlinear attention to query chunk-level “key” and “value”. This same mechanism is used across all layers, heads, and tokens, making RAT not a hybrid model, whereas Mamba-Transformer alternates state-space modules and attention in different layers. And we're not putting the EMA construction in one layer and placing the attention in the second layer. In detail, \n\n  * The way recurrence used in RAT is different from Mamba and can be extended to different kinds of RNN in the future. Eqn. (2) in RAT is applied only within chunks and over both “key” and “value”. This chunking design allows the simple recurrence to work well on short chunks, and enables future exploration of more advanced RNNs, such as 2D recurrence for large chunks or nonlinear RNNs for small chunks. By contrast, Mamba-1 aggregates the “value” over the entire sequence (which can suffer from memory degradation issues on very long inputs) and is limited to linear RNNs. Mamba-2 cannot be reduced to Eqn. (2), as it is the linear attention with per-head gating.\n  * RAT’s cross-chunk attention is applied sparsely over previous chunks, reducing FLOPs compared to standard attention. We also designed an efficient implementation using FlexAttention to handle causal masking during training. \n\nSince RAT interpolates between attention and RNN layers, we mainly compared against these baselines in the paper and noted (lines 245–246, footnotes) that RAT outperforms recent linear attention works. Any hybrid strategy (e.g., mixing mechanisms in different layers or heads) are indeed orthogonal to RAT. In this paper, we explored RAT-SWA, which interleaves RAT with sliding-window attention. We now post comparisons with other state-space models below.\n| Hybrid models | value/state size (T=300, D=2048, N=24 (number of layers)) | arc_challenge | arc_easy     | hellaswag    | lambada | piqa  | winogrande | boolq |\n|-|-|-|-|-|-|-|-|-|\n|||acc/acc_n|acc/acc_n|acc/acc_n|acc|acc/acc_n|acc|acc|\n|Mamba1|64ND|-/35.40|69.52/-|-/52.91|43.98|71.32/-|52.95|61.13|\n|Mamba2|256ND|-/37.88| 72.47/-|-/55.67|45.66|71.87/-|55.24|60.13|\n|DeltaNet|128ND|-/35.66| 68.47/-|-/50.93|42.46|70.72/-|53.35|55.29|\n|GatedDeltaNet| 288ND| -/38.39|71.21/- |-/55.76|**46.65**|**72.25**/-|**57.45**|60.24|\n|RAT(L=16)|~20ND| 36.35/**39.33**|**72.64**/67.80| 43.27/**56.44**|44.4|72.03/72.69|53.2|**63.3**|\n|Hybrid|\n|Samba (Mamba1-SWA(2048) | 64N/2D-300N/2D| -/36.17| 68.81/- | -/53.42 | 44.94   | 70.94/- | 55.56 |62.11|\n|GatedDeltaNet-SWA(2048) | 288N/2D-300N/2D| -/40.10 | 71.75/-   | -/56.53 | 47.73   | 72.57/- | **58.4** |63.21|\n|RAT(L=16)-SWA(1024)| ~20N/2D-300N/2D| 37.20/**40.78**| **72.56**/68.22  | 44.33/**57.85** | **49.33**|**73.29**/73.94 | 56.91| **63.21**|\n  \nResults for other models are taken from a recent paper [1], where they were trained with the same model size, datasets, and sequence lengths as ours. These commonsense reasoning tasks have short inputs (≤300 tokens), so RAT with a chunk size of 16 has at most \\~20 chunks (\\~20ND value vectors), yet it achieves better performance with much smaller “value” size.  The RAT-SWA also surpasses other Mamba-SWA or GatedDeltaNet-SWA. Also note that the current RAT has a clean structure without using extra short convolutions or normalization layers, which other works often adopt to boost performance.\n\nFinally, the novel chunk-based design allows partial compression of history while preserving direct access to past tokens, which helps retrieve distant information, especially in noisy contexts (see results in Reviewer 1's A2) . Unlike recent state space models that has fixed memory slots and may suffer from memory degradation on long inputs, RAT’s memory slots grows with sequence length while maintaining a fixed FLOPs reduction ratio. RAT can complements the popular local attention well by effectively and efficiently preserving long-term information. Hope these could address the reviewers’ concern.\n\n[1]. Gated Delta Networks: Improving Mamba2 with Delta Rule. ICLR’25\n\n> Q2: Because of the EMA-style construction of the \"RNN\" part, this also does not afford the model the capability to do state-tracking: https://arxiv.org/pdf/2404.08819. \n\nA2: The simple Eqn. (2) with per-dimension gating operates only within its own chunk (e.g., 16 or 64 tokens). Because the chunk length is so short, the simple recurrence can effectively encode the entire chunk without information loss and is therefore fully expressive for summarizing it. The cross-chunk softmax attention then handles long dependencies by providing direct access to distant tokens. This design also allows for a flexible value/state size, whereas state space models typically have fixed memory slots size even when the sequence length becomes very long.\n\n> Q3: I am surprised at the 7x training speedup and 9x generation speedup, and I am curious what the baseline comparisons are. Was flash attention used? Since I suspect the speedups gained in the paper may fade with 1) different implementation comparisons, or 2) scale of model, I think there is little gain in capability if we consider the gains from using a RAT-Attention hybrid vs a Mamba-Attention hybrid. \n\nA3: As stated in the paper (lines 185–187, bold “efficiency” in the experimental setup), we ensured fair comparisons by using FlashAttention for the attention module, associative scan for the RNN, bf16 and torch.compile for all models, and by measuring the latency of a single token mixing block (including input and output projections) on the same GH200 GPU.\n\nWe analyzed the efficiency in Method Section 2.4 and showed that the observed 7×–9× speed-up is reasonable. Compared to standard attention, RAT reduces theoretical FLOPs by the chunk length (e.g., 16), and has a clean structure with no extra short convolutions or normalizations, and can be implemented very efficiently with our optimized implementation. \n  * During training, the recurrence does not involve state expansion or reduction, and different chunks can be stacked together to pass through the parallel scan. The causal masking problem in cross-chunk attention is addressed using FlexAttention and the online softmax trick. \n  * At inference time, a single-step update is used for the RNN, and the same FlashAttention can be used for cross-chunk attention. \n\nThus, when benchmarking the latency of the temporal mixing block (including the linear projections), the observed 7×–9× speed-up is expected. We have also provided the anonymous code link in line 808 of the Appendix. \n\nFor hybrid models, we interleave RAT with sliding-window attention rather than full attention for efficiency. Detailed differences and comparisons with Mamba are provided in the question above. Finally, we emphasize that RAT is a new architecture, not a hybrid or state-space model, and not an incremental variant of Mamba. We proposed the idea, explored its positional encodings, parameter allocations, and efficient implementations, and validated it across multiple tasks.\n\n> Q4: However, these efficiency gains may fade at a larger scale when MLP computation dominates.  The speedups gained in the paper may fade with 2) scale of model.\n\nA4: First, this is not specific in RAT but a general property of all temporal mixing methods: temporal mixing and channel mixing compete for runtime dominance, and which dominates depends on model dimension and sequence length. Then we analyzed the speed-up on larger models for training and generation cases, respectively.\n  * For computation-intensive training, we benchmark the latency of the temporal mixing block (including input and output projections) and find that efficiency gains decrease when the model dimension D increases because the linear projections contribute a larger fraction of FLOPs. For example, the speed-up drops from 7× to 5.5× at 100K tokens when D increases from 2048 to 4096 (Attention block: 1776.51s, RAT(L=16) block: 322.43s). However, training a 7B model typically uses tensor parallelism, which reduces the model dimension on each GPU (e.g., D=2048 with 2 GPUs), so the expected 7× speed-up can still be achieved.\n  * For memory-intensive decoding, the scaled-dot product operation in attention remains a major bottleneck, even for large models. We report latency for a single temporal mixing block with D=4096 when generating one batch of token at a specific position, and the maximum throughput of the full 7B and 13B models when generating 1024 tokens with a 3072-token prompt. RAT consistently achieves significant speed-ups over attention. On the 13B model in particular, attention utilizes GPU resources poorly due to its very limited maximum batch size, which allows RAT to achieve even greater throughput improvements.\n\n| Latency of Temporal mixing block (D=4096) | Position=4096, Batch=1024 | Position=8192, B=512 | Position=16384, B=256 | Position=32768, B=128 | Position=65536, B=64 |\n|-|-|-|-|-|-|\n| RAT(L=16)  | 3.01s| 2.7s | 2.25s | 2.21s | 2.15s  |\n| Attention  | 26.16s  | 25.82s | 25.87s | 25.8s   | 26.22s |\n| Ratio | 8.7x  | 9.6x | 11.5x| 11.7x | 12x |\n\n|Maximum throughput of a whole model | 1.3B (results in our paper) | 7B | 13B|\n|-|-|-|-|\n| RAT(L=16)| 31170 token/sec | 10103 token/sec | 5749 token/sec |\n| Attention| 3152 token/sec| 983 token/sec| 534 token/sec  |\n| Ratio  | 10.2x  | 10.3x| 10.8x |"}, {"review_id": "KCDJ7SIyFw", "reviewer": "Reviewer_FAfY", "summary": "The paper introduces RAT (Recurrent-Attention Tradeoff), a novel architecture designed to bridge the efficiency of RNNs with the accuracy of attention mechanisms in language modeling. RAT segments input sequences into chunks, applying a lightweight linear recurrence within each chunk for local modeling and using softmax attention across chunks to capture long-range dependencies. This hybrid structure enables flexible trade-offs via chunk size, interpolating between RNN-like and attention-like behavior. The authors demonstrate that RAT significantly improves training and inference speed, up to 10×, without substantial loss in accuracy. Through extensive pretraining and evaluation at the 1.3B parameter scale, RAT matches or exceeds standard attention on various tasks, especially in long-context settings. A further enhancement interleaves RAT with sliding-window attention to boost local precision, showing consistent performance gains across commonsense reasoning, code, and summarization tasks. The study highlights RAT's efficiency, architectural simplicity, and strong empirical results, positioning it as a promising alternative for scalable language models.", "questions": "1. How well does RAT scale to larger models like 7B or 13B? Please comment on any challenges with parallelism, memory, or throughput at scale. Concrete results or estimates would help assess broader relevance.\n\n2. RAT underperforms on tasks like WinoGrande and few-shot QA. Can you explain whether this is due to chunking, attention granularity, or training setup? A deeper analysis would clarify limits of generalization.\n\n3. The reversed design (attention within chunk, recurrence across) is said to perform worse. Can you provide more insight or experiments to show why? Understanding this failure would help validate your design choices.\n\n4. How does RAT compare directly to recent efficient models like Mamba or DeltaNet under similar compute or FLOP budgets? A head-to-head baseline is needed to contextualize the claimed efficiency and accuracy trade-offs.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Quality: The paper presents a technically solid and well-executed study, with claims supported by strong empirical results across pretraining, evaluation, and fine-tuning settings. The proposed RAT architecture is carefully implemented and benchmarked against attention and RNN baselines. However, theoretical analysis is limited, and some performance gaps (e.g., in tasks like WinoGrande or few-shot QA) are acknowledged but not deeply investigated. The work is complete and honestly evaluated, though larger model scaling and ablations at higher capacity would strengthen the evidence.\n\nClarity: The paper is mostly clear and well-organized, with helpful diagrams and consistent terminology. Key ideas such as intra-chunk recurrence and inter-chunk attention are explained effectively. However, some descriptions (e.g., implementation efficiency or positional encoding strategies) are overly condensed and require more unpacking. Presentation could be improved with clearer separation between core contributions and peripheral optimizations.\n\nSignificance: RAT targets a critical challenge in scaling language models efficiently, and the results suggest meaningful practical benefits in speed and memory while maintaining accuracy. The ability to interpolate between attention and recurrence is valuable, and the hybrid design with sliding-window attention shows promise. While the impact at larger scales remains untested, the demonstrated improvements at 1.3B suggest the approach is relevant and could influence future model design.\n\nOriginality: The idea of chunked recurrence combined with sparse cross-chunk attention is a novel architectural middle ground, distinct from both state-space models and purely linear attention. While the components are not new individually, their integration is original and motivated by practical trade-offs. The paper provides a new way to think about temporal mixing that departs from both strict attention and RNN paradigms, though more conceptual framing and discussion of prior hybrid efforts would make the novelty clearer.", "was_revised": true, "final_justification": "I am satisfied with the author's response and have updated my evaluation slightly upward to reflect the clarifications provided.", "rebuttal": "We would like to thank the reviewer for the constructive feedback. Here are our responses.\n> Q1: How well does RAT scale to larger models like 7B or 13B? Please comment on any challenges with parallelism, memory, or throughput at scale. Concrete results or estimates help assess broader relevance.\n\nA1: This is an important question. We did not train RAT on 7B or 13B models because we prioritized model and architecture designs in this paper, especially given the compute resources available to us. However, we provide efficiency analyses and results at larger model scales to demonstrate its potential.\n  * For parallelism, we mentioned in the paper that RAT is compatible with both tensor parallelism and context parallelism. Both the recurrence and cross-chunk attention are head-independent, enabling standard tensor parallelism by assigning different heads to different GPUs. For context parallelism, the within-chunk recurrence is chunk-independent and works well with small chunk sizes, allowing chunks to be distributed across GPUs. Since RAT stores only a small number of key/value vectors (e.g., 16× fewer), cross-chunk attention may even avoid ring-style communication. \n  * For compute-intensive training, we benchmarked the temporal mixing block (including input/output projections) and found that speed-up decreases with larger model dimension D, as linear projections dominate FLOPs more. For instance, when D increases from 2048 to 4096, the speed-up drops from 7× to 5.5× (Attention block: 1776.51s vs. RAT(L=16) block: 322.43s for 100K tokens). However, training a 7B model typically involves tensor parallelism, which reduces the per-GPU dimension (e.g., D = 2048 with 2 GPUs), allowing the 7× speed-up to be retained.\n  * For memory-intensive decoding, scaled-dot product attention remains a major bottleneck, even in large models. We benchmark latency on a single temporal mixing block when generating tokens at a specific position, and measure the maximum throughput of full models when generating 1024 tokens with a 3072-token prompt. RAT consistently outperforms attention, and on the 13B model, where attention suffers from limited batch size and poor GPU utilization, RAT achieves even higher improvements.\n\n|Latency of Temporal mixing block (D=4096)|Position=4096, Batch=1024|P=8192, B=512|P=16384, B=256|P=32768, B=128|P=65536, B=64|\n|-|-|-|-|-|-|\n| RAT(L=16)| 3.01s| 2.7s|2.25s| 2.21s| 2.15s  |\n| Attention|26.16s| 25.82s|25.87s| 25.8s| 26.22s |\n| Ratio|8.7x|9.6x|11.5x|11.7x| 12x |\n\n|Maximum throughput of a whole model (token/sec)|1.3B (results in our paper)|7B|13B|\n|-|-|-|-|\n|RAT(L=16)|31170|10103|5749|\n|Attention|3152|983|534|\n| Ratio|10.2x|10.3x|10.8x |\n\n> Q2: How does RAT compare directly to recent efficient models like Mamba or DeltaNet under similar compute or FLOP budgets? A head-to-head baseline is needed to contextualize the claimed efficiency and accuracy trade-offs.\t\n\nA2: Thanks for the suggestion. Since RAT is a novel idea with a simple structure that interpolates between RNN and attention, we primarily compared it with strong attention and RNN baselines for clarity and showed comparable performance to attention. We also noted in lines 245–246 and the footnotes that our best model outperforms the best model in a recent linear attention paper [1]. Below we provide detailed comparisons with these models and will include them in the revision.\n\n[1] trained these models with the same model size (1.3B), sequence length (4096), and datasets (fineweb_edu 100B) as ours. We calculate the value/state size based on task sequence length T, dimension D, and number of layers N. For state space models, we multiply the value dimension by the state expansion factor and note that Mamba2 and GatedDeltaNet expand the value dimension. Since these commonsense reasoning tasks contain very short contexts (≤300 tokens), RAT(L=16) only has at most 20ND value vectors, yet it surpasses others in 4 of 7 tasks, and RAT-SWA outperforms GatedDeltaNet-SWA in 6 of 7 tasks. On retrieval-heavy tasks (Needle-in-Haystack and LongBench QA), RAT also consistently outperforms them.\n\n[1]. Gated Delta Networks: Improving Mamba2 with Delta Rule. ICLR’25\n|Model|value/state size (T=300, D=2048, N=24) |arc_challenge|arc_easy|hellaswag|lambada|piqa|winogrande|boolq|\n|-|-|-|-|-|-|-|-|-|\n|||acc/acc_n|acc/acc_n|acc/acc_n|acc|acc/acc_n|acc|acc|\n|Mamba2|256ND|-/37.88| 72.47/-|-/55.67|45.66|71.87/-|55.24|60.13|\n|DeltaNet|128ND|-/35.66| 68.47/-|-/50.93|42.46|70.72/-|53.35|55.29|\n|GatedDeltaNet| 288ND| -/38.39|71.21/- |-/55.76|**46.65**|**72.25**/-|**57.45**|60.24|\n|RAT(L=16)|~20ND| 36.35/**39.33**|**72.64**/67.80| 43.27/**56.44**|44.4|72.03/72.69|53.2|**63.3**|\n|Hybrid|\n|GatedDeltaNet-SWA(2048)|288N/2D-300N/2D|-/40.10|71.75/-|-/56.53|47.73|72.57/-|**58.4**|63.21|\n|RAT(L=16)-SWA(1024)|~20N/2D-300N/2D| 37.20/**40.78**|**72.56**/68.22| 44.33/**57.85**|**49.33**|**73.29**/73.94| 56.91|**63.21**|\n\n|Model|value/state size (T=4096, D=2048, N=24)|NIAH-1|NIAH-2|NIAH-3|\n|-|-|-|-|-|\n|DeltaNet|128ND|**99**|18.6|22.4|\n|Mamba2|256ND|65.4|56.2|4.6|\n|Gated DeltaNet|288ND|91.4|92.2|27.6|\n|RAT(L=16)|256ND|97|**97.6**|**56.4**|\n\n|Model|value/state size (T=4096, D=2048, N=24) |NQA|Qasper|MQA|HQA|WQA|MSQ|\n|-|-|-|-|-|-|-|-|\n|Mamba2|256ND|11.1|11.3|18.6|11.8|15.1|**6.7**|\n|DeltaNet|128ND|12.9|10.8|21.5|10.9|13.2|5.1|\n|GatedDeltaNet|288ND|14.1|14|23.3|13.7|14.1|5.8|\n|RAT(L=16)|256ND|**14.5**|**16.1**|**28.3**|**16.7**|**18.9**|6.3|\n\n> Q3: The reversed design (attention within chunk, recurrence across) is said to perform worse. Can you provide more insight or experiments to show why? \n\nA3: We provided relevant discussion and results in Appendix C.1 and Table 9. We found that the reversed design uses FLOPs less effectively. Short-context dependencies are easier to capture and do not suffer from memory degradation, so using lightweight recurrence locally is more efficient.  Long‑context dependencies are harder and may require direct retrieval of distant information, where attention is more suitable.\n\nIn the reversed design, maintaining good performance requires a large chunk size; otherwise, the recurrence must handle long-term history alone and may suffer from memory degradation. But with large chunk sizes, most FLOPs are spent computing local information. For example, Table 9 shows that on the Book dataset (sequence length 8192), the reversed design needs a chunk size of 512 (with 16× FLOPs reduction) to reach low perplexity. In contrast, RAT achieves even lower perplexity with a chunk size of 64 and 64× FLOPs reduction (128 chunks for inter-chunk attention).\n\nOverall, RAT uses FLOPs more effectively. That said, the two designs may be complementary with RAT being better at preserving long-term information, and the other emphasizing local modeling. Due to resource constraints, we only explored interleaving RAT with local attention (without recurrence across chunks) in our large-scale experiments and leave combining both designs for future work.\n> Q4: RAT underperforms on tasks like WinoGrande and few-shot QA. Can you explain whether this is due to chunking, attention granularity, or training setup? \n\nA4: We think the weak performance of RAT (L=16) on WinoGrande is due to task and model instability in small‑scale pre‑training. The evaluation set contains only ~1k items and it is a binary choice task, so biases in the small-scale pretraining on this task can easily push accuracy to 50–60.  Another task we did not include but found a similar problem in the paper is BoolQ, where attention performs poorly (62.54 for RNN, 63.30 for RAT, and 58.23 for attention).\n\nAs for the weak performance of RAT (L=16)-SWA on TREC (results on other few-shot QA tasks are comparable, so we focus on TREC here), we think the main reason is that LongBench frames this classification task with 50 mixed coarse‑ and fine‑grained labels in a few‑shot task. Without post‑training or instruction tuning, models may generate semantically correct but non‑exact answers and be uncertain about the required specificity. RAT (L=16)-SWA often produces more descriptive outputs for some labels (e.g., “Date of holiday” for “Date”), and simply correcting these would bring its performance close to the attention and RAT. For location‑related questions, which form a large portion of the test set, all models except attention‑SWA tend to predict the generic “Other location” instead of fine‑grained labels like “City,” explaining attention‑SWA’s higher accuracy on TREC.\n> Q5: The paper provides a new way to think about temporal mixing that departs from both strict attention and RNN paradigms, though more conceptual framing and discussion of prior hybrid efforts would make the novelty clearer.\n\nA5:  We’d like to first clarify that RAT is not a hybrid model. It applies the same mechanism to all tokens, heads, and layers. But hybrid models use different modules like attention or state-space in different places. Hybrid strategies are orthogonal to RAT. In our paper, we explored interleaving RAT with local attention and leave other directions, such as varying chunk sizes across layers or heads, for future work.\n\nRegarding novelty, one key contribution is the chunk-based design. It allows partial compression of history while preserving direct access to past tokens, which helps retrieve distant information, especially in noisy contexts. Unlike recent state space models that has fixed memory slots and may suffer from memory degradation on long inputs, RAT’s memory slots grows with sequence length while maintaining a fixed FLOPs reduction ratio. Finally, while local attention has been widely adopted in industry, RAT can complements it well by effectively and efficiently preserving long-term information.\n\nWe will make these points clearer in the revision.\n>Q6: Presentation could be improved with clearer separation between core contributions and peripheral optimizations.\n\nA6: Thanks for pointing this out. We'll update the revision based on the suggestion."}]

Author rebuttal:

Meta-review: The paper proposes RAT, a chunked recurrence-plus-attention mechanism that achieves strong efficiency gains (up to 9–10× speedup) while maintaining or improving accuracy on both short- and long-context benchmarks. The approach is simple and empirically validated at the scale of1.3B, with further evidence provided in the rebuttal for 7B/13B settings. The main concerns raised were incomplete benchmarking against other efficient or hybrid attention methods and the need for clearer distinctions from related chunked/state-space approaches. While one reviewer remained unconvinced, the additional results and analysis alleviated many concerns, and the overall consensus is that this work makes a meaningful and novel contribution. I recommend acceptance, with the expectation that the camera-ready version will resolve the rest of the concerns.
Decision: Accept (poster)