Rodimus*: Breaking the Accuracy-Efficiency Trade-Off with Efficient Attentions
Authors: Zhihao He
Hang Yu
Zi Gong
Shizhan Liu
Jianguo Li
Weiyao Lin
Venue: ICLR 2025 Poster (2025)
Keywords: Large Language Models
Fixed-Size Hidden States
Data-Dependent Tempered Selection
Sliding Window Shared-Key Attention
Primary area: foundation or frontier models, including LLMs

Abstract: Recent advancements in Transformer-based large language models (LLMs) have set new standards in natural language processing. However, the classical softmax attention incurs significant computational costs, leading to a $O(T)$ complexity for per-token generation, where $T$ represents the context length. This work explores reducing LLMs' complexity while maintaining performance by introducing Rodimus and its enhanced version, Rodimus$+$. Rodimus employs an innovative data-dependent tempered selection (DDTS) mechanism within a linear attention-based, purely recurrent framework, achieving significant accuracy while drastically reducing the memory usage typically associated with recurrent models. This method exemplifies semantic compression by maintaining essential input information with fixed-size hidden states. Building on this, Rodimus$+$ combines Rodimus with the innovative Sliding Window Shared-Key Attention (SW-SKA) in a hybrid approach, effectively leveraging the complementary semantic, token, and head compression techniques. Our experiments demonstrate that Rodimus$+$-1.6B, trained on 1 trillion tokens, achieves superior downstream performance against models trained on more tokens, including Qwen2-1.5B and RWKV6-1.6B, underscoring its potential to redefine the accuracy-efficiency balance in LLMs. Model code and pre-trained checkpoints are open-sourced at https://github.com/codefuse-ai/rodimus.

Published as a conference paper at ICLR 2025
RODIMUS* BREAKING THE ACCURACY-EFFICCIEEA TRADE-OFF WITH EFFICIENT ATTENTIONS
Zhihao He! Hang Yu² Zi Gong", Shizhan Liu?, Jianguo Li?; Weiyao Lin
Shanghai Jiao Tong University, Ant Group
ABSTRACT
Recent advancements in Transformer-based large language models (LLMs) have set standards in natural language processing. However, the classical softmax atten- tion incurs significant computational costs, leading to a (TT) complexity for per- token generation, where 7 represents the context length. This work explores reduc- ing LLMs' complexity while maintaining performance by introducing Rodimus and its enhanced version, Rodimus+. Rodimus employs an innovative data-dependent tempered selection (DDTS) mechanism within a linear attention-basel purely recurrent framework, achieving significant accuracy while drastically reducing the memory usage typically associated with recurrent models. This method exemplifies semantic compression by maintaining essential input information with fixed-size hidden states. Building on this, Rodimus+ + combines Rodimus with the innovative Sliding Window Shared-Key Attention (SW-SKA) a hybrid approach, effectively leveraging the complementury semantic, token, and head compression techniques. Our experiments demonstrate that Rodimus+-116B trained on trillion tokens, achieves superior downstream performance against models trained more tokens, including Qwen2-1.55 1.5B and RWKV6-1 16666, underscoring its potential to redefine the accuracy-efficiency balance in LLMs. Model code and pre-trained checkpoints are open-sourced at httph::t///tgt /tittp, com/ codefuse ai/rodimmn
INTRODUCTION
Recent advancements have positioned Transformer- bassed large language models (LLMs) at the fore- front of natural language processing, establishing them as state-of-the-art. Their strong capabilities stem from the softmax attention mechanism, which selectively focuses on relevant tokens stored in the kcy-value (KV) cache when predicting the next token. By maintaining a historical record of tokens, the KV cache allows all pertinent data to remain accessible during inference. However, the demand to store this extensive historical information incurs notable computational costs, leading to (T) complexity for per-token generation, where T represents the length of the context preceding the generated token. Indeed, a 7B Llama (Touvron al.) 2023h), without inference optimization, takes over a minute to generate 2K-length sequences on an A10-24G (Ilyas Moutawwakil 2023).
This has sparked research into next-generation foundation models aimed at efficient alternatives to attention. Among them, three main categories stand out that aim to compress the KV cache in the original softmax attention from distinct perspectives: semantic, token, and head compression.
Semantic Compression:
The first category, also known as linear attention (Katharopoulos et al. 2020), or linear state space models (SSMs) (Gu & Dao 2023; Dao & Gu 2024), substitutes the exponential kernel in softmax attention with simplified inner product of (transformed) query and key vectors. This approach results in recurrently updated hidden state of fixed size that retains historical semantic information like linear RNNs (Qin et al., 2024b) Peng et al. (023), successfully reducing the per-token generation complexity from O(T') to (1). However, compressing the softmax caariiinn-ccaaaeeriiieede by its unlimited capacity into linear attention with fixed capacity inevitably leads to some information loss. To address this, current methods strive to (i) increase the recurrent state size to enhance memory capacity (Qin et al. 2024a) and (ii) utilize the fixed-sized states more effectively (Yang et al.) 2024b). Despite these advancements, there remains a trade-off between complexity and performance (refer to Figure [ Blindly expanding the hidden states compromises
"Equal contribution. This work was done when Zhihao He was a research intern at Ant Group. Corresponning authors.

---

Published as a conference paper at ICLR 2025
Language Modeling (Sequence Longth 512)
MQAR (Sequence Length 1024 Pairs 256)
GLA
30
HGRN
Mamba
0.8
Mamba-2
Validaton PPL
StreamingLLM
Accuracy
-0
Rodimus
Rodimus
04
Mamba
Mamba-2
0.3
Rodimus
Rodimus-
50
0.0
0.5
1.0
1.5
2.0
9.5
3.0
Memory Foosprint (M Byces)
Expansion Factor
(a) LM Task on the WikiText- 103 Dataset.
(b) Multi Saali--uery Associative Recall (MQAR) Task.
Figure I: Memory Footprint vs. Performance: (a) This experiment is conducted on the WikiText-100 dataset (Details in Appendix D.3), and focuses on the recorded best perplexity (PPL). The memory footprint is adjusted by modifying the expansion factor. (b) The model's 's recall capability is assessed using the MQAR Task (See Appendix (D.7), as described in Arora et al. (2024a). Among all models evaluated, Rodimus* achieves the optimal balance between space complexity and performance.
time and space efficiency, while incorporating data-dependent decay gating mechanisms to filter out irrelevant past information may hinder parallel training efficiency, as noted in (Yang et al. 2024c).
Token Compression:
This second category introduces sparsity into the attention mask, allowing it to follow predefined patterns and focus on strategically chosen tokens, such as those at the beginning of sequence or close to the answer (Xiao al. 2024 Han et al., 2024). Consequently, this method can also achieve the (1) complexity for per-token generation. However, the elimination of masked tokens can lead complete loss of information and potentially degrade the overall performance.
The third category reshapes attention by modifying the design of attention heads. This involves grouping heads and sharing keys and values within each group, as seen multi -query attention (MQA) (Shazeer (019) and grouped-query attention (GQA) (Ainslie et al. (023)). While these methods reduce the cache size by constant factor, they are inherently lossy compared to multi-head attention (MHA) (Vaswani 2017) as the same values are used within each group.
Head Compression:
Ultimately, none of these methods can fully supplant the original softmax attention, prompting an essential question: Is it feasible to reduce the complexity of LLMs while preserving performance? Our research affirms this possibility, showing that a linear attention model enhanced with well- designed data-dependent tempered selection (DDTS) mechanism- an improvement within the first category-and its hybridization with our proposed sliding window shared-key attention (SW-SKA), which integrates all three types of attention compressions, offers viable solutions. We designate the former model as Rodimus and the latter as Rodimus-. Remarkably, Rodimus maintains a much smaller hidden state size than most current SOTA recurrent models (c.g.. half of the size in Mamha2 (Dao & Gu 2024)), but outperforms softmax attention- Transformers. Moreover, Rodimus+--1.6 (trained
trillion tokens) achieves an average performance that is 0.31% higher than Qwen2-1.5B (trained on 7 trillion tokens) and 2.3% higher than RWKV6-1.66 (trained on 1.4 trillion tokens) in downstream task evaluations. Our contributions can be summarized as follows:
We introduce a linear attention-based, purely recurrent model, Rodimus, effectively overcoming the accuracy-efficiency trade-off found in existing recurrent models. By incorporating the innovative DDTS, Rodimus can autonomously filter out irrelevant information during semantic compression, resulting improved validation perplexity (PPL) with much smaller memory footprints compared to other SOTA methods, as illustrated by the green curves in Figure
We present a hybrid model, Rodimus+, which combines Rodimus with the innovative SW-SKA. While the recurrent hidden state in Rodimus offers a comprehensive semantic view of the historical context, the sliding window attention highlights the nearby tokens that are most influential in predicting the next token, and SKA further compresses cache size in a lossless manner. As shown by the red curves in Figure I Rodimus+ continues to extend the accuracy-efficienge frontier of existing methods, achieving superior performance with reduced complexity.
We validate the effectiveness of Rodimus* (encompassing both Rodimus and Rodimus+) through thorough experimentation. In downstream evaluations with model sizes from 130M to 1.3B, * demonstrates enhanced language modeling performance relative to current SOTA models of similar sizes. Furthermore, Rodimus* achieves a performance improvement of up to 7.21% % over Mamba2, and outperforms even softmax attention- based Pythia on NeedleBench- tasks where recurrent models typically underperform (Waleffe et al. 2024).
suite of
2

---

Published as a conference paper at ICLR 2025
2 PRELIMINARIES
In this section, provide a brief introduction to softmax attention in Vaswani (2017) and derive its recurrent form. We will then derive the aforementioned three attention compression methods, including linear attention for semantic compression, sparse attention for token compression, and sharing-based attention for head compression. The first two methods simplify attention by reducing context length, while the last one makes modifications along the dimension of attention heads.
Softmax Attenion and its Recurrent Form: Suppose that X = ( 2r E RR denotes the input sequence with length T and dimension d. Standard autoregressive Transformers (Vaswani 2017 utilize a softmax attention mechanism to generate the output O E R'x' that is,
(I) where Wi E Rddx WK Rdxn Wv E Rd are learnable weight matrices associated with the queries Q. the keys K, and values V. respectively. The matrix M E (-00, 1} E RxxT denotes the upper-triangular attention mask, ensuring that the model does not attend to future tokens. Multi-Head Attention (MHA) further splits the dimensionality d into hi heads, computing attention for each head in dh individually with distinct weights. This approach enhances the model's ' capacity to capture diverse sequential information while reducing the computational costs for matrix operations.
Q, Q.K.V = XWo, XWK: XWv; O softmax ((QK o GM)M
The above parallel form allows for the computation of O = for....., Or} } in parallel given the full input X, facilitating efficient training. contrast, during inference, Transformers rely on the following recurrent formulation (Katharopoulos et al. 2020):
E exp(g:k: )o:
qr, kr, 02= = x. WQ, I. WK, n. Wv, 02 EL (sp(
(2)
Here, the query q: E I key k2 E R Ixn, and value v. R lxm vectors are computed based on the representation of current token o, E Bixd, Attention is subsequently performed over the evolving collection of keys (ki,..., k,} and values {Vi,..., v. (i.e., the KV cache). Thus, the time and space complexity for generating the next token at time stamp is (tt.
Linear Attention for Semantic Compression: To optimize efficiency, linear attention mechanisms replace the exponential kernel exp(q:k; in Eq. (2) by a kernel k(q,, k;) paired with an associated feature map é, i.e., ki k(g;, k.) = ((g:) (Ak;) As a result, the calculation of o, can be simplified as:
O Σt (( d (k.) v; ((q:) ΣS
(kk:) To;
E p(q (k:)))
((q:) E (k;,)
Letting S. = Et= o(k;) T 0, E mx and 2, = Ei=1 (k:)) E Rax be the KV state and the state respectively, we can rewrite the previous equation as linear state-space model (SSM) or RNN:
(gq.)
S, S + ((k)) Vt: 2, = = - (kk) O2 ((q.)2,
We emphasize that S, provides a semantic compression of the historical context up to f. The denominator (z E IR may introduce numerical instabilities and hinder the optimization (Qin
ct al. (022a), prompting many recent studies to replace it with a normalization (Qin ct al.) 2022a
Sun et 2023a). Moreover, it is common to use the identity mapping for á (Sun et al. 2023a 3a Yang et al. 2024b), and so Eq. (4) amounts to
S + kg Ur o: = S..
(5)
More generally, when framed within the context of linear SSMs (Gu Dao 2023 Dao Gu or RNNs (De et al. 2024), these equations can be reformulated as:
2024
S. A, 3 S.1 B, O u, o, C,S..
(6)
where A, 1, B: 1, and C. = q: represent the state transition, input, and output matrix respectively, while u, k v, signifies the input control. In stark contrast to softmax attention, the per-token generation complexity at time stamp t is (1), enabling efficient inference. Meanwhile, linear attention can be expressed in a parallel form similar to Eq. I allowing for training parallelism.
However, this additive formulation of updating the hidden states S, with new key-value pairs at each time step does not possess the capability to forget irrelevant information, which leads to the phenomenon known as attention dilution (Qin et al., 2022b a). To address this concern, recent works (i) increase the state size (either n or d₂) retain more information and (ii) incorporate decay
3

---

Published as a conference paper at ICLR 2025
factors gating mechanisms that enable the state St to discard irrelevant past information. The first approach typically sacrifices efficiency for performance, whereas the second method focuses on designing A, and B, to manage memory retention and forgetting behaviors within the hidden states.
We summarize the functional forms of A, and B, used in the SOTA methods in Table 4 Three major trends can be gleaned from this table. Firstly, it is beneficial for A, and B, to be negatively correlated, as shown in Mamba (Gu & Dao 2023). ). Mamba2 (Dao & Gu (024). and HGRN2 QQin et al.] (0244)), allowing them to collaboratively regulate the deletion and addition of information in the hidden state S,. This contrasts with methods like RetNet (Sun et al. 2023a), gRetNet (Sun et al. 2024). and GLA (Yang et al. 2024b), which use A, as decay factor while keeping B, constant at 1. Secondly. allowing A, and B, to be functions of the input enables dynamic adjustments over time in a data-dependent manner. This approach has been validated by the superior performance of GLA (Yang et al.) (024b), gRetNet (Sun et al.) 2024) compared to RetNet(Sun et al. (223a)). Lastly, designs for gating mechanisms must be compatible with GPU acceleration, ensuring that the recurrent expression 6 aligns with parallel format, similarly to
In Section 3.1 we further analyze these functional forms in relation to their equivalence with linear attention and propose novel data-dependent tempered selection (DDTS) mechanism. This mechanism is capable of succinctly compressing historical information within a recurrent hidden state of fixed capacity, all while facilitating parallel training.
This group of methods aim to sparsify the attention mask M in Eq. [ The goal is to compute only the (gt, k;) pairs associated with nonzero elements of Mi. Various types of attention masks have been proposed in the literature, including window attention (Child al.) 2019), dilated attention (Beltagy al. 2020), and bridge attention (Guo et al., 2023), among others. Research shows that the softmax attention maps in vanilla Transformer models often exhibit localized behavior (Qin et al. 2022a Xiao et al. 2024). Thus, we exploit a sliding window (SW) attention to enhance the local context understanding in Rodimus+.
Sparse A Attention for Token Compression:
Sharing- based Attention for Head Compression: In the original MHA, each head has distinct linear transformations Wo, WK, Wy for the input X, allowing for diverse representations and attention maps across heads. Sharing- and value heads to be shared among multi-query heads. Specifically, in MQA (Shazeer 2019). all heads use a single set of key and value weights, which reduces parameters and memory usage but risks diminishing the attention mechanism's s expressiveness. As better alternative, GQA assigns one key and value head for each group of query heads (Ainslie et al., 2023). J. However, it still limits the expressiveness of the original MHA by overly constraining learned relationships. As remedy, we propose the Shared-Key Attention (SKA) in Section 3.2.1 which compresses the MHA while preserving its expressiveness.
attention seek to compress the MHA by allowing key
3 METHODOLOGY
In this section, we introduce two models: Rodimus and Rodimus+ whose architecture is depicted in Figure Rodimus is a purely recurrent model that iteratively compresses historical context into a fixed-size hidden state (i.e., semantic compression) and further exploits this state for the next token prediction. With the equipment of the newly proposed DDTS mechanism, Rodimus effectively filters out irrelevant information, thereby enhancing performance while reducing the size of the hidden state. On the other hand, Rodimus- builds upon Rodimus by integrating the proposed SW-SKA technique. This enhancement improves performance without sacrificing efficiency, allowing for a seamless combination of semantic, token, and head compression methods.
3.1 RODIMUS: OVERCOMING PERFORMANCE BOTTLENECK IN SEMANTIC COMPRESSION
3. ANALYSIS OF EXISTING LINEAR ATTENTION MODELS
As can be observed from Table 4 the state transition equation in all existing linear attention models can be expressed recurrently as:
S, = (a, (1) @ S + (á, 3r) @ (kTTe,) v.) = (a, 3,) o S,-1 + (ā, @ k2) T о v.), where a, E glxn and B: E g gl denotes the gating mechanism along the dimension of n and
(7)
m respectively, and &, and B, are negatively correlated with a and B,, allowing them to select between the current input u, k v. and the previous state S Substituting Eq. (7) into Eq. (6)
4

---

Published as a conference paper at ICLR 2025
Ltter
FFN
Emrbedding
SW-SKA
LM Head
RMM
Shamd kry
RMSNom
3M6NNom
SW-SKA
Ix
Rodimus
SharkCoep
MulE qUEY
NEWN va2a8
Embodding
Liver
Libser
Liber
Liner
Licaer
Rodimus Black
Token Embeeding
Rodimus" Model
(b) Rodimus
(c) The SW-SKA part of Rodimus+ Block
Figure 2: Overview of the Proposed Models. The Rodimus* Model serves as a template for both Rodimus and Rodimus+-- When modules within the gray dashed box included, it becomes the Rodimus+ Model; otherwise, it is the Rodimus Model. The architecture comprises L layers of stacked Rodimus Blocks along with essential modules for language modeling (c.g., Embedding, RMSNorm, LM Head). Rodimus Flow (TI) denotes our proposed recurrent computation method. The purple arrows depict information flow between layers (ignoring RMSNorm). In the Rodimus- Model, the Rodimus block compresses the Token Embedding into a global semantic embedding, enhancing SW-SKA's ability to perceive long contexts.
and unfolding the recurrent relation yields:
q
II ag ROoaOk J=++1
II B, o B,0Bor)) j=it1
8
Here, I1 Oj captures the relative positional information hetween time stamp f. and i, specifically between the query and the key k; Indeed, this component acts as a more flexible version of relative positional embeddings for softmax altention (e.g., xPos (Sun al. 2023b) and ALiBi (Press et al.) 2022)). In addition, the elements within the vector Or vary from one another, enabling the capture of high- and low-frequency information akin to sinusoidal positional encoding (Vaswani (017) and RoPE (Su et al. 2021). Notably, Iīi=+t1 a also incorporates absolute positional information since is also function of the absolute positions j i 1,..., Consequently, linear attention models expressed in the form of Eq. (7) capture first-order dependencies within sequences in a position-aware manner. In contrast, the original softmax attention (I) captures only second-order dependencies in position-agnosse framework, necessitating the addition of positional embeddings (Ren et al.) (023)..
On the other hand, If 33 also regulates the relative positional information between the time stamp / and i, but between the query q and the value v,, This aspect has not been witnessed in softmax attention. Moreover, earlier studies (Yang et al. 2024b) and ablation study (see Table T6) indicate that incorporating 1[; B3 cannot result in improvements compared setting B3 1m probably because such positional information has already been described by II; -ii C',. Moreover, including By can impede training efficiency, as the dimensionality of B3 E R xm is significantly larger than that of 04 Rlxn This discrepancy lead redundant I/O operations when B are frequently transferred between high bandwidth memory (HBM) and on-chip SRAM during chunkwise parallelism (cf. appendix(B.3]
As consequence, our focus is on optimizing the design of ar while setting Br = 1m for all t. We retain 3, to allow for flexible selection among the various elements in the value vector u,
3.1.2 DESIGN OF at, â AND B. We propose the following formulations for a and Qy:
ax exp(-xp gt ① Tt), â; 9."
(9)
5

---

Published as a conference paper at ICLR 2025
where 9. ((x; W و + bby) and T, (2,,W, + b,). In this design, It serves selection gate, determining whether to retain the previous state S,-1 or incorporate the current input kt. Note that exp(-9:)
oxx W , + b,), thus ensuring that a; E [0, 1] and preventing complete oblivion of the previous state. In contrast, Og can assume values greater than due to the softplus function. This asymmetry between discarding the previous state and retaining the current introduces greater flexibility into the state transition equation (Gu & Dao (023).
Furthermore, different from all existing works, we introduce a temperature gate t. that governs the sharpness or sensitivity of the selection gate gt- As visualized in Figure 7 in the appendix, as T decreases, both a and d changes more slowly with gr. Note that Tı is function of the input zt- as opposed to being fixed constant as is the case in gRetNet (Sum et al.) 2024) and GLA (Yang al. 2024b). This provides an additional degree of freedom in the selection process. Indeed, recent findings indicate that tempered losses can improve robustness against noise during training (Papernot al., 2021; Wang et al. 2023). our case, we find that т, helps sharpen the original selection gate gr. (see Figure 8 in the appendix), thereby facilitating more aggressive filtering of irrelevant information.
On the other hand, 3r can be expressed as:
(10) where W mx and W Rx are two low-rank matrices (i.c., E < m). This low-rank B formulation helps mitigate noise in the input while keeping the overall increase in model parameters manageable, thus providing parameter efficiency (Hu et al. (002). This comprehensive design is referred to as the data- detendependent tempered selection (DDTS) mechanism and is proven to be a selection mechanism as demonstrated below:
3. - = W + bz).
Proposition 1. Given the specifications for Ar and B, in Eqs. O realize the selection between the previous state S and the current imput u,
and 100. DDTS can
Proof. See Appendix C.2
3.1.3 THE OVERALL RODIMUS BLOCK
The SSM formulation for the overall Rodimus block can be articulated follows:
S: (exp(-ge @ T:) 1mm @ S,-1 ((g.') T 8.) @ (k. vt), O: = qu + d, o x/.
(11)
In this formulation, we define q: = x; Wq. q. ke = xz WR. 2 = x₁. We follow the Mamba series (Dao & Gu 2024) to set 02 = Zt, and the computation of the control u, = k, 0g can be interpreted as the state expansion operation within SSMs. Additionally, we employ ShortConv to process the original input ot into x/, from which we derive gi and T.. Note that Short Conv is widely used in recent recurrent models (Gu & Dao 2023 Dao & Gul 2024; Yang et al., (024bb;; enhances local context aggregation and introduces nonlinearity for and TE. The learnable weight de functions as the feedthrough matrix in SSMs. To ensure stability during back- propaaappa,,,,,n we implement post -nomallzzatinn after activation. This entails normalizing kn and dividing qe by vn (Dao & Gu 2024).
The final configuration of the Rodimus block is depicted in Figure Zp. The afforementioned SSM merges into Gated Linear Unit (GLU) (Dauphin et al. (017). allowing simultaneous token and channel mixing compactly, akin the GAU (Hua et al. (022) and Mamba series (Gu Dao 2023] Dao & Gu 2024). The initial linear layers within the GLU increase the size of the hidden states, thereby enhancing the performance of Rodimus. The last linear layer finally reduces the hidden dimension back to its original size d. In contrast to the original Transformer block, which mixes tokens and channels separately through softmax attention and feedd-Fowwar Network (FFN), this compact architecture demonstrates greater parameter efficiency. Additionally, the output gate of the GLU enriches the gating mechanisms of the Rodimus block.
During inference, the Rodimus block retains only a fixed-size hidden state, resulting in ((1) time and space complexity. In contrast, during training, the chunkwise parallelization of the Rodimus block ensures sub-quadratic computational complexity, which is explained in the Appendix B.3
3.2 RODIMUS+: INTEGRATION WITH TOKEN COMPRESSION AND HEAD COMPRESSION
Although the hidden states S, in Rodimus provide comprehensive semantic overview of historical context, their fixed capacity may still limit performance. On the other hand, it has been observed in the literature (Qin et al. 2022a b Xiao et al. 2024) that the attention map in vanilla Transformer

---

Published as a conference paper at ICLR 2025
MHA
MQA
GQA
SKA
Volues
Keys
Quetles
Figure 3: Comparison of Different Attention Mechanisms. In MQA and GQA, values are shared among the same group of query heads, resulting in lossy compression compared to the individual head-specific values used in MHA. SKA maintains the multi-value setting of MHA but uses shared key across all heads. This approach produces a separate attention map for each head, preserving the expressiveness of MHA while reducing the memory footprint.
tends to be localized, indicating that nearby tokens are significantly more influential in generating responses. To leverage the strengths of both approaches, we introduce Rodimus+, which integrates the Rodimus block with sliding window attention (i.e., token compression). This integration enables the hidden states to act as robust global representation of the context, incorporating information even from distant tokens relative to the answer, while the sliding window attention sharpens focus on the most influential nearby tokens. Notably, this enhancement occurs without increasing computational complexity. Furthermore, we optimize the KV cache of the sliding window attention by a constant factor by substituting the traditional MHA with our proposed SKA (i.e., head compression).
3.2.1
SHARED- SEAR ATTENTION FOR LOSSLESS HEAD COMPRESSION
Recall from Section
2 that the MHA for head h can be expressed as:
O softmax
(Q**"* T o M) V = softmax (XX X )
M ( xW
= softmax (XWWWX (WT)(WXTT) O M) XW
(12)
where Wx is a learnable weight matrix that is invariant with regard to head h, and W K denotes the pseudo inverse of the transposed WK- Thus, we can redefine the query as Q X X Wx and the key as K X XWK. which is invariant across heads. This means for all heads without diminishing the expressiveness of MHA. We refer to this mechanism as Shared- Key Attention (SKA), which achieves lossless compression of the KV cache in MHA by utilizing single key for all heads. In contrast, approaches like MQA (Shazeer) 2019) and GQA (Ainslie et al.) 2023) also share the values V among the same group of query heads, which typically differ for each head in the original MHA formulation. Consequently, MQA and GQA can only be seen as lossy compressions of MHA. TigureBillustree Billutraates the distinctions between MHA, MQA, GQA, and SKA.
employ single key
We further combine SKA with sliding window attention to create SW-SKA, which introduces a local structure where each token focuses solely on its immediate contextual window, thereby maintaining a constant memory footprint during both training and inference.
3.2.2 THE OVERALL RODIMUS+ BLOCK
Recall that the Rodimus block builds upon the foundations established by the Mamba series (Ciu & Dao 2023 Dao & Gu 2024) and GAU (Hua al. 2022), enabling the simultaneous mixing of tokens and channels while offering a unified residual connection for both mixers (see Figure Zb). Specifically, integrates the SSM block (token mixer) into the FFN layers (channel mixer). This structured arrangement allows for an increase in state size without necessitating additional parameters, as the FFN typically amplifies the dimension before subsequently reducing In contrast, a conventional Transformer block alternates between self-attention (token mixer) and FFN (channel mixer), with each component linked by its own residual connection. Notably, self-attention is generally not placed in between the FFN layers due to the linear growth of the KV cache with input sequence length, which can lead to significant memory consumption when key and value dimensions are increased.
When integrating the Rodimus block with the SW-SKA, it is beneficial to have the attention mech- anism and FFN closely interdependent, as these components collaboratively deliver a localized understanding of context. However, they should maintain relative independence from the Rodimus block, which is responsible for providing a global perspective. To achieve this balance, we employ a two-hop residual approach, adapted from Ma et al. (2024), effectively bind the attention and FFN
7

---

Published as a conference paper at ICLR 2025
Scaling Curve (Sequence Length 2048)
Table I: Results on WikiText-103.
Transkmmart
Models
Test PPL.
Params(M)
(log scale)
7=10"
Mamba-2
Rodimus
Rodimus
21.90
46.68
Rodimus*
Rodimus+
21.56
48.45
Mamha
22.58
46.08
6x10
Mamba2
22.78
46.38
Aesssno pp
Transformer+
22.02
46.24
Llama3
22.12
46.48
GLA
22.
44.71
5.10"
IIGRN2
22.00
46.20
The Jower the PPL, the better the language modeling effect. Bold indicates the best result, and undertine indicutes the second best result.
12"
Training FLOPs (log scale)
Figure Scaling Curve based on Scaling Laws.
layers. This ensures that the information derived from the token mixer is cohesively integrated into the channel mixer, preserving the overall integrity of the system. The resulting Rodimus+ block is:
X State Rodimus(Norm( X)) + X,
Y = SW SWW-SK(NNrml staate) + X X
(13)
Y = FFN(Norm(Y)) (Y)) +X
where FFN represents the GLU, Norm denotes RMSNorm (Zhang & Sennrich 2019). and Rodimus signifies the Rodimus block, generating state information for the subsequent SW-SKA Taken together, the Rodimus block, SW-SKA, and FFN constitute the final Rodimus+ block. In Appendix B3.4| we
further discuss the relationship between Rodimus+ and existing hybrid models.
EXPERIMENTS
In this section, we compare Rodimus* with other SOTA methods across various benchmarks, in- cluding language modeling and recall benchmarks. brief introduction to the SOTA methods is presented in Appendix B.5
We conclude this section with a series of ablation studies.
4.1 LANGUAGE MODELING
Wiki 103: The WikiTexx-1003 language modeling dataset (Merity et 2022
consists of over
100 million tokens derived from verified good and featured articles on Wikipedia. In this study, we compare Rodimus* to six other models. Among these, Mamba (Gu & Dao 2023) and Mamba2 (Dao
& Gu 2024) utilize purely recurrent architectures, similar to Rodimus. GL.A Yang et al.] 2024b)
and HGRN2 (Qin al., 2024a) serve as linear attention counterparts to the Mamba series. We also include TransTormer- (Touvron et al. 2023) which employs multihead softmax attention, and Llama3 (Dubey et al.] 2024), the grouped query attention version of +ransformer++. All models trained from scratch the training dataset, each with approximately 44 million parameters. To ensure a fair comparison, we set Rodimus* expansion factor n to 32, aligning its parameter count with that of the other models. Other experimental configurations are also the same across all models, including batch size, learning rate, and training iterations (see Appendix D.T.
The results, summarized in T present the perplexity (PPL.) of the testing data. Notably, Rodimus, a purely recurrent model, surpasses both Transformer++ and Llama3 in performance, despite having an inference complexity one power less than that of the latter two models. When compared to Mamba2, which has a state expansion factor of 128, Rodimus employs a smaller expansion factor of 32, leading a significantly reduced hidden state size while achieving lower PPL. Furthermore, the performance of Rodimus+ even exceeds that of Rodimus. These findings indicate that Rodimus* effectively challenges the traditional accuracy- -eeuracy-efficienyy trade-off in language modeling.
Scaling Laws (Hoffmaan et al. )022): We further expand the number of parameters in our models from approximately 125 million 1.3 billion, adhering to the architectural guidelines GPT-3 (Brown et al. 2020). For comparison, we used Mamba2 (Dao & Gu 2024) and Transformer (Touvron et al. 2023) as baselines. Mamba2 represents the most advanced recurrent models, while Transformer+ + exemplifies the leading softmax attention models. All models are trained on subsets of Pile Gao et al. 2020). Detailed experimental setups are provided in Appendix D.4 Figure illustrates that both Rodimus and Rodimus+ outperform Mamba2 and Transformer+ across all model sizes, ranging from 125 million to 1.3 billion parameters. Interestingly, Rodimus+, which integrates Rodimus with the SW-SKA, achieves even greater improvements over Rodimus with the
8

---

Published as a conference paper at ICLR 2025
Table 2: Results of Downstream Tasks via Zero-Shot Evaluation.
Params
Tokens
ARC-e
ARC-e
HS
LMB
LMB
OBQA
PIQA
WG
AVG
Model
Tokenizer
(B)
10/1
mmm1
29.
somt
Mamba'
0.13
100
NeoX
23.72
46,84
23.37
19.72
40.33
29.80
65.07
52.17
41.61
Mambaze
0.13
100
NeoX
23.81
45.50
34.13
20.42
39.55
29.20
64.47
$1.85
41.22
Rodimus
0.13
100
NeoX
23.38
46.84
34.22
17.95
42.44
30,00
65.40
$1,46
41.96
GPT-Neo
0.13
300
GPT2
23.21
43.69
30.42
30.26
37.38
26.40
62.89
50.51
39.21
OPT
0.13
300
OPT
22.78
43.52
31.35
26.02
37.90
28.00
63.00
50.36
39.56
Pythia
0.16
300
NeoX
23.72
43.60
30.19
38.31
32.80
26.00
61.59
50.51
38.34
Mamba
0.13
300
NeoX
24.32
47.90
35.20
16.05
44.21
28.60
64.69
52.33
42.46
Mamhu2
0.13
300
NeoX
24.23
47.35
35.32
16.79
43.78
30.40
64.85
$2.64
42.65
RWKV4
0.17
300
NeoX
23.55
47.56
32.27
31.73
32.64
27 80
64.25
$0.99
39.87
Rodimus
0.13
300
NeoX
23.72
49.07
35.45
15.75
45.00
29.20
64.36
$3.35
42.88
Rodimus-
0.13
300
NeoX
24.23
47.22
35.49
12,94
48.09
29.20
64.36
53.43
43.15
OPT
0.35
300
OPT
23.98
44.19
36.66
16.40
45.10
28.20
64.47
52.41
42.14
Pythia
0.41
300
NeoX
24.49
52.10
40.59
10.83
$1,45
29.40
66.97
53.59
45.51
BLOOM
0.56
300
BLOOM
23.89
47.47
36.89
28.83
34,10
28.80
64.20
52.01
41.05
Mamba
0.37
300
NeoX
27.82
54.92
46.49
8.14
55.60
30.80
69.53
55.17
48.62
Mambu2
300
NeoX
26.62
54.67
46.96
7.98
32,60
70.46
55.64
48.97
RWKV4
0.43
300
NeoX
25.26
47.18
40.77
13.38
45.39
30.60
68.12
53.28
44.37
Qwen2
0.5
7000
Owen2
28.84
54.97
49.03
11.66
50.05
33,20
57.62
49.02
Rodimus
0.46
150
Redimus
27.65
55.77
48.78
10.17
50.65
32.60
70.73
55.41
48.80
Rodimus
0.47
150
Rodimus
28.58
57.28
52.13
10.22
51.14
35.00
72.91
53.59
50.09
GPT-Neo
1.3
300
GPT2
25.77
56.14
48
7.50
57.21
33
71.22
55.01
49.69
OPT
1.3
300
OPT
29.52
56.94
$3.75
6.65
57.89
33.20
71.60
59.59
51.78
Pythia
1.0
300
NeoX
27.05
56.94
47.15
7.92
56,22
31,40
70.67
$3.51
48.99
Pythia
NeoX
51.98
6.08
33.20
57.30
51.97
BLOOM
1.1
300
BLOOM
25.51
$1.47
42.97
17.28
42.64
29.40
67.19
55.01
44.88
GLA
1.3
100
Mistral
27.82
55.13
48
15.37
46.09
33.00
69.80
S3.20
47.72
RetNet
1.3
100
Mistral
26.45
57.32
48.04
16.44
43.37
32.00
69.42
53.43
47.15
HGRN
1.3
100
Mistral
25.51
55.01
48.82
20.24
37.71
31.80
69.80
50.28
45.56
HGRN2
Mistral
49.97
32.80
71.27
52.17
49.18
Mamba
1.3
300
NeoX
32.85
65.49
$9.08
5.04
64.87
36.40
74.10
61.40
56.31
Mamha2
300
NeoX
33.36
59.92
5.02
37.80
73.29
56.43
RWKV4
1.5
300
NeoX
29.01
60.94
$2.90
7.08
57.19
33.40
72.09
S5.33
$1.55
RWKV6
1.6
1420
RW
33.36
60.69
61.42
4.61
67.26
37.40
73.67
60.38
56.31
Rec-Gemma
2.0
2000
Gemma
29.69
48.91
61.82
9.27
53.83
29 40
67.52
$7.06
49.75
Qwen2
1.5
7000
Qwen2
35.75
65.95
65.51
5.52
63.61
36.80
75.19
65.27
58.30
Rodimus
1.4
500
Rodimus
36.09
68.01
62 44
6.33
59.62
35.60
74.32
59.75
56.55
Rodimus-
1.6
1000
Rodimus
36.35
68.35
64.51
5.38
63.59
38,80
76.17
62.51
58.61
Here, bold, underting, wayy upslering indicates the best, secund, third result, respectively.
volumm "Tokens", underling
indicates models using lewer tokens. The superscrign demotes models trained from scratch using the same configuration for fair comparison.
increase of the model size. As the number of layers and parameters grows, the benefits of the two-hop residual connections may become more pronounced (Ma et al.) 2024). Consequently, Rodimus- demonstrates superior scaling potential, especially when approaching 1.3 billion parameters.
Downstream Evaluation: We evaluate Rodimus* across various downstream tasks, as outlined
in Gu & Dao (2023). These tasks include content analysis (LAMBADA (Paperno et al. 2016)).
commonsense reasoning (PiQA (Bisk al.. (020) and HellaSwag (Zellers et al. 2 coreference
resolution (WinoGrande (Sakaguchi et al. 2019)), reading comprehension (ookn (Mihaylov)
al. 2018)), and professional examinations (ARC-Easy and ARC-Challenge (Clark et al.] 201))..
The evaluation results are obtained using the 1m-eva luat oo-ooon-aaeee (Gao et al. 2023). For
comparison, we choose well-known open source models as baselines, indicating the number of tokens and the tokenizers used, similar approaches commonly adopted in prior works (Gu & Dao 2023
Dao & Gu 2024). Our experimental structure is as follows: (i) We first train both Rodimus and Mamba series with 130M parameters from scratch using the same settings
100B tokens from the Pile dataset (denoted as Rodimus Mamba and Mamba2t) for a fair comparison. (ii) We then train - 130M 300B tokens from Pile and compare it other similarly sized models also trained on 300B tokens from Pile, noting that the training corpora for these models may still differ. (iii) We evaluate R with about 460M parameters against other open-source models like BLOOM- 560M and Qwen2-500M. Since these models are trained on self-cleaned high gugh-quality data, we also utilize curated datasets, including Weh (Penedo et al. 22024), Pile (Gao et al., 2020). etc., to train our Rodimus* on 150B tokens. (iv) To fully explore the model's 's capabilities, we train R.6imus+-1.6B on IT tokens from the curated dataset, while Rodimus- 1.4B is trained on 500B tokens due to resource constraints. Both models are compared against other open-source models of comparable sizes, such as BLOOM, RWKV6, and Qwen2, which also use self-cleaned datasets for training. Note that the tokenizers used in the last two experiments are developed based on their respective curated datasets, The specific results are presented in Table 2 More details can be found in Appendices D.5 D.5 and D.6

---

Published as a conference paper at ICLR 2025
Mambaa21138 Overall Score-6.51
Pythia-1.48 Overall Score:3.8
Rodimus-1.48
Rodimus+-1
Overall Spore:2.4
Overall Score
102
212
50
65
21
04
Figure 5: Results NeedleBench. The average accuracy of each subtask of NeedleBench in contexts of different lengths. Overall Score represents the global average score.
In experiment (i), where Rodimus Mamba*, and Mamba2 are trained under identical conditions, Rodimus outperforms the Mamba series on downstream tasks, validating the effectiveness of the DDTS design. In experiment (ii), where all models are trained with the same number of tokens from Pile, Rodimus outperforms all existing models across the benchmarks. Moreover, Rodimus+ achieves even better results, indicating that the proposed hybrid structure is more effective for downstream tasks. Moving experiment (iii), note that the number of training tokens can impact performance. Due to a relatively small token count, Rodimus-460M underperforms compared to Qwen2-500M and Mamba2-370M but still surpasses Mamba-370M, which trained on 300B tokens. In contrast, Rodimus-4-4000 delivers the best average result, even though it is trained on only 150B tokens, further validating the proposed hybrid structure's effectiveness. Lastly, in experiment (iv), training larger models on a larger corpus reinforces Rodimus** potential as practical LLM. Specifically, Rodimus- 1.44 outshines all existing purely recurrent models of similar sizes, including those from the Mamba and RWKV series. Furthermore, Rodimus+- --666 surpasses the performance of both Qwen2-1.5B and Saaurrent-Gemma- 2B, despite being trained on only 1T tokens, in contrast to the 7T and 2T tokens used for the latter two models, respectively. In addition, we refer the readers to where we continue-pretre prettrii Rodimus+ as a practical LLM for math and code.
Appendix E.3
4.2 RECALL BENCHMARKS
In this section, we assess the recall or retrieval capabilities of Rodimus* in the context of the MQAR task and NeedleBench. The MQAR task involves training a small-paramess using a specific dataset, while NeedleBench evaluates models trained on large-scale corpus. Due to space limitations, we offer brief overview here and direct readers to Appendices D.7 and D.8 for more details. For MQAR, we primarily compare Rodimus* * with other recurrent models and find that Rodimus consistently outperforms existing models the state expansion factor n (see Figure lb) or the model dimension d (see Figure () increases, thanks to the proposed DDTS mechanism. For NeedleBench, both Rodimus and Rodimus+ even exceed the performance of Pythia, a softmax attention- based model (see Figure (). is noteworthy that many studies suggest pre-trained recurrent models often struggle with challenging recall-intensive tasks like NeedleBench when compared to models utilizing full attention mechanisms (Waleffe et al., 2024). The superior performance of Rodimus and Rodimus+ indicates that the Rodimus family effectively balances the accuracy- efficiency trade-off, particularly in enbancing recall ability. More results on other recall benchmarks, such as FDA, SWDE, NQ, SQuAD, TriviaQA, and DROP, can be found in Appendix D.9
model from scratch
4.3 ABLATION STUDIES
In this section, we summarize the key conclusions from ablation studies. A detailed analysis is provided in Appendix DD10 (i) Our assessment of the components in Rodimus reveals that the inclusion of 9t. T:, and B, is crucial for the DDTS framework. However, the incorporation of B, tends to complicate the training process, leading to suboptimal results. Additionally, the two-hop residual connection is beneficial as the model depth increases. (ii) In terms of head compression methods, we find that SKA offers a balanced compromise berween GQA and MHA, delivering superior performance compared to GQA while utilizing fewer parameters relative to MHA, irrespective of the backbone models used. (iii) We have determined that the optimal hyperparameter values are = 16 and = 64. Notably, the Rodimus configuration with = 64 operates with a state size that is only half that of Mamba2, yet it outperforms Mamba2 on various benchmarks.
CONCLUSION
We introduce here Rodimus and Rodimus+,, two innovative models that utilize semantic, token, and head compression. We refer the readers to Appendices G for more discussion and future work.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
The work was supported by the National Natural Science Foundation of China (No. 62325109, U21B2013) and Ant Group.
REFERENCES
AbdeIrahman Abouelenin, Atabak Ashfaq, Adam Atkinson, Hany Awadalla, Nguyen Bach, Jianmin Bao, Alon Benhaim, Martin Cai, Vishrav Chaudhary, Congcong Chen, et al. Phi-4-mini technical report: Compact yet powerful multimodal language models via mixture-of-loree arXiv preprint arXiv:2503.01743, 2025.
Joshua Ainslie, James Lee- Leo-Th Michiel de Jong, Yury Zemlyanskiy, Federico Lebron, and Sumit Sanghai. GQA: Training generalized multi-query transformer models from multi-head checkpoints. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 4895-4999999 4900,, Singapore, December 2023. Association for Computational Linguisties. doi: 10. (006633/1/2223emmm//pp. https::/acllantho..ooo www.yyy.rgg222223..
URL
1p ma in. 298
Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Re. Zoology: Measuring and improving recall in efficient language models. In The Twelfth International Conference on Learning Representations, 2024a. URL https: http///pweeweeeeeeeeeeeeee wwt.:/w..eeeeeeeeeeeeeeeeeeeeee wtt.://www...eoooooooeeeeuuuaaa.
Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, James Zou, Atri Rudra, and Christopher Re. Simple linear attention language models balance the recall- throughput tradeoff. In ICLR 2024 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2024b. URL BtKLhhttss://oeeeeee
ew. net./ ftt:/formm? bimmmm?ddopg d= coPhEoD
Ben Athiwaratkun, Sanjay Krishna Gouda, Xiaopeng Wang. Zijian an Li, Yuchen Tian, Ming Tan, Wasi Uddin Ahmad, Shiqi Wang, et al. Multi-lingual evaluation of code generation models. arXiv preprint arXiv:2210.14800 2022.
Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, et al. Program synthesis with large language models. arXiv preprint 177732, 2021.
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. 055::2004400150, 2020.
Stella Biderman, Hailey Schoelkopf. Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pp. 2397-2433. PMLR, 2023.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language. Thirty AAAI Conference on Artificial Intelligence, 2020.
Sid Black, Leo Gao, Phil Wang. Connor Leahy, and Stella Biderman. GPT-Neo: Large Scale Autore- gressive Language Modeling with Mesh-Tensorflow, March 2021. URL httlh:tts:/ddd org/ 00.5281/zenodo. 5297715 If you use this software, please cite it using these metadata.
Aleksandar Botev, Soham De, Samuel L Smith, Anushan Fernando, George-Cristian Muraru, Ruba Haroun, Leonard Berrada, Razvan Pascanu, Pier Giuseppe Sessa, Robert Dadashi, et al. Re- currentgemma: Moving past transformers for efficient open language models. arXiv preprint arXiv:2404.07839, 2024.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 13:1877-1901, 2020.
11

---

Published as a conference paper at ICLR 2025
Federico Cassano, John Gouwar, Daniel Nguyen, Sydney Nguyen, Luna Phipps-Costin, Costin, Donald Pinck- ney, Ming-Ho Yee, Yangtian Zi, Carolyn Jane Anderson, Molly Q Feldman, Arjun Guha, Michael Greenberg, and Abhinav Jangda. Multipl-e: A scalable and extensible approach to benchmarking neural code generation, 2022.
org/ / /208. 08227
Mark Chen. Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained code. arXiv preprint vrXiv22107.03374 2021.
Rewon Child, Scott Gray, Alee Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019.
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, abs/1 haa//800305455, 2018.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 10. /4/68, 2021.
Codefuse and Ling Team. Every sample matters: Leveraging -ixturr- and qual-qy data for efficient and accurate code Ilm, 2025. URL /tl.htttps//axxxi.rrr///a/s/
7793
OpenCompass Contributors. Opencompass: A universal evaluation platform for foundation models. hhttps://githue cwm/opem///ocooooccee wttpass//pppcccccopaaaaaaaa 2023.
Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality, In International Conference on Machine Learning (ICML), 2024.
Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In International conference machine learning, pp. 933-941. PMLR. 2017.
Soham De, Samuel I. Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al. Griffin: Mix- ing gated linear recurrences with local attention for efficient language models. arXiv preprint arXiv:2402. 19427 2024.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, et al. The Ilama herd of models. arXiv preprint 2xXiv:240..21788 2024.
Daniel Y. Fu, Tri Dao, Khaled K. Saab, Armin W. Thomas, Atri Rudra, and Christopher Ré. Hungry Hungry Hippos: Towards language modeling with state space models. In International Conference on Learning Representations, 2023.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint 1eXiv22101000022 2020.
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding. Jeffrey Hsu, Alain Le Noac' 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. A framework for few-shot language model evaluation, 12 2023. URL httlstt::: /twpooo/////oooooo o org/ http://ooorooo/o0o6/66
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00075,, 2023.
Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. Advances in neural information processing systems, 33: 1474-1487, 2020.
12

---

Published as a conference paper at ICLR 2025
Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initialization of diagonal state space models. Advances in Neural Information Processing Systems, 35:35971- 35983, 2022a.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In The International Conference on Learning Representations (ICLR), 2022b.
Alex Gu, Baptiste Rozière, Hugh Leather, Armando Solar-Lezama, Gabriel Synnaeve, and Sida I. Wang. Cruxeval: A benchmark for code reasoning, understanding and execution. arXiv preprint arXiv:2401.03065 2024.
Daya Guo, Canwen Xu, Nan Duan, Jian Yin, and Julian McAuley. Longcoder: A long-range pre- trained language model for code completion. In International Conference on Machine Learning. pp. 12098-12107. PMLR, 2023.
Chi Han, Qifan Wang. Hao Peng, Wenhan Xiong, Yu Chen, Heng Ji. and Sinong Wang. Lm-infinite: Zero-shot extreme length generalization for large language models. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguisties:
Human Language Technologies (Volume 1: Long Papers), pp. 3991-4008, 2024.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. Proceedings of the International Conference on Learning Representations (ICLR), 2021a.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. NeurIPS, 2021b.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. An empirical analysis of compute-optimal large language model training. Advances in Neural Information Processing Systems, 11:30016-30000, 2022.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. International Conference on Learning Representations, 2022. URL
id=nzevKeeFyE>
Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc Le. Transformer quality in linear time. In International conference on machine learning, pp. 9099-9117. PMLR, 2022.
Yuzhen Huang, Yuzhuo Bai, Zhihao Zhu, Junlei Zhang, Jinghan Zhang, Tangiun Su, Junteng Liu, Chuancheng Lx, Yikai Zhang, Jiayi Lei, Yao Fu, Maosong Sun, and Junxian He. C-eval: A multi-level multi-discipline chinese evaluation suite for foundation models. In Advances in Neural Information Processing Systems, 2023.
Binyuan Hui, Jian Yang, Zeyu Cui, Jiaxi Yang, Dayiheng Liu, Lei Zhang, Tianyu Liu, Jiajun Zhang, Bowen Yu, Keming Lu, et al. Qwen2. 5-coder technical report. arXiv preprint arXiv:2409.1211 2024.
Régis Pierrard Ilyas Moutawwakil. LIm-perf leaderboard. spaces/
https://huggingface co/
2023.
Naman Jain, King Han, Alex Gu, Wen-Ding Li, Fanjia Yan, et al. Livecodebench: Holistic and contamination free evaluation of large language models for code. arXiv preprint arXiv:2403.07974,
2024.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mnn:: Fast autoregressive transformers with linear attention. In International conference on machine learning, pp. 5156-5165. PMLR, 2020.
Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Iliè. Daniel Hesslow, Roman Castagné, Alexandra Sasha Luccioni, François Yvon, Matthias Gallé, et al. Bloom: A 176b- parameter open-access multilingual language model. 2023.
13

---

Published as a conference paper at ICLR 2025
Haonan Li, Yixuan Zhang, Fajri Koto, Yifei Yang, Hai Zhao, Yeyun Gong, Nan Duan, and Timothy Baldwin, Cmmlu: Measuring massive multitask language understanding in chinese, 2023.
Mo Li, Songyang Zhang, Yunxin Liu, and Kai Chen. Needlebench: Can Ilms do retrieval and reasoning in million context window?, 2024. URL
11963
Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, et al. Jamba: A hybrid transformer- mamba language model. arXiv preprint arXiv:2403. 19887, 2024.
Jiawei Liu, Chunqiu Steven Xia, Yuyao Wang, and Lingming Zhang. Is your code generated by chatGPT really correct? rigorous evaluation of large language models for code generation. In NeulPS, 2023. URL
net forum? wttpmmmmddiiiixx111
Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Zettlemoyer Luke. Mega: Moving average equipped gated attention. arXiv preprint PrXiv:2209.10655. 2022.
Xuezhe Ma, Xiaomeng Yang, Wenhan Xiong, Beidi Chen, Lili Yu, Hao Zhang, Jonathan May. Luke Zettlemoyer, Omer Levy, and Chunting Zhou. Megalodon: Efficient Ilm pretraining and inference with unlimited context length, 2024.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In International Conference on Learning Representations, 2022.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? new dataset for open book question answering. In EMNLP, 2018.
Niklas Muennighoff, Qian Liu, Armel Zebaze, Qinkai Zheng, Binyuan Hui, Terry Yue Zhuo, Swayam Singh, Xiangru Tang, Leandro von Werra, and Shayne Longpre. Octopack: Instruction tuning code large language models, 2023.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The LAMBADA dataset: Word prediction requiring broad discourse context. In Katrin Erk and Noah A. Smith (eds.), Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Pupers), pp. 1525-1534, Berlin, Germany, August 2016. Association for Computational Linguistics. doi: 10.18653/v1/P16 144. DBLL:tttee////aaattooooooo..g/gp
Nicolas Papernot, Abhradeep Thakurta, Shuang Song, Steve Chien, and Ülfar Erlingsson. Tempered sigmoid activations for deep learning with differential privacy. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 9312-9321, 2021.
Guilherme Penedo, Hynek Kydlíček, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, Thomas Wolf, et al. The fineweb datasets: Decanting the web for the finest text data at scale. arXiv preprint arXiv:2406. 7557, 2024.
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Leon Derczynski, Xingjian Du, Matteo Grella, Kranthi Gv, Xuzheng He, Haowen Hou, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Jiaju Lin, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito. Guangyu Song, Xiangru Tang, Johan Wind, Stanistaw Wožniak, Zhenyuan Zhang, Qinghua Zhou, Jian Zhu, and Rui-Jie Zhu. RWKV: Reinventing RNNs for the transformer era. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 14048-14077, Singapore, December 2023. Association for Computational Linguistics. doi: 10.18653/v1/2023.1 fnndii findings-emmip 936
findingg-emmlll.955. URL https: bttpss://aclanth... ogy www.eoy.....oo.oo..... org/2023
Bo Peng, Daniel Goldstein, Quentin Anthony, Alon Albalak, Eric Alcaide, Stella Biderman. Eugene Cheah, Teddy Ferdinan, Haowen Hou, Przemystaw Kazienko, et al. Eagle and finch: Rwky with matrix-valued states and dynamic recurrence. arXiv preprint a40iv:2404 05892, 2024.
14

---

Published as a conference paper at ICLR 2025
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations, 2022. URL
lttps://<penreviee:
Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. The devil in linear transformer. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 7025-7041, 2022a.
Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan. Lingpeng Kong, and Yiran Zhong. cosformer: Rethinking softmax in attention. In International Conference on Learning Representations, 2022b.
Zhen Qin, Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, and Yiran Zhong.
HGRN2: Gated linear RNNs with state expansion. First Conference on Language Modeling. 2024a. htLhhtts::// httcttps://openeevieee net
mete/oourme forum? 1--rrmmmdddy
Zhen Qin, Songlin Yang, and Yiran Zhong. Hierarchically gated recurrent neural network for sequence modeling. Advances in Neural Information Processing Systems, 36, 2024b.
Liliang Ren, Yang Liu, Shuohang Wang, Yichong Xu, Chenguang Zhu, and Cheng Xiang Zhai Sparse modular activation for efficient sequence modeling. Advances in Neural Information Processing Systems, (6:19999--1982) 2023.
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 preprins 24X6::2406.07522, 2024.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53-68, 2021.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. arXiv preprint 19Xivv/907.100 2019.
Noam Shazeer. Fast transformer decoding: One write-head is all you need, 2019. URI. https
7/arxiv. wtwxv..org/abs//911.1
Jimmy T.H. Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for sequence modeling. In The Eleventh International Conference on Learning Representations, 2023.
(tRl:tttss:// httttpps:://ppeeeeeeeee iew. n fenet/forune forum? id=Ai 8Hw3AXqks
Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, 2021.
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 preprins vrXii:2307.08621, 2023a.
Yutao Sun, Li Dong, Barun Patra, Shuming Ma, Shaohan Huang, Alon Benhaim, Vishrav Chaudhary, transformer. In Anna Rogers, Jordan Boyd- Xia Song, and Furu Wei. A Pruceedings of the 61 Annual Meeting of the Association Graber, and Naoaki Okazaki (eds.),
for Computational Linguistics (Volume 1: Long Pupers), pp. 19.11455 14590-14604, Toronto, Canada, July 2023b. Association for Computational Linguistics. doi: 10. wt..1865/w1/2223..o lolng long.816. URL pttpp:://cllathhooggy
8116
Yutao Sun, Li Dong, Yi Zhu, Shaohan Huang, Wenhui Wang, Shuming Ma. Quanlu Zhang, Jianyong Wang, and Furu Wei. You only cache once: Decoder-decoder architectures for language models, 2024. URL
hbtp///..//4/e.e// 05254
Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoe Le, Ed Chi, Denny Zhou, et al. Challenging big- bench tasks and whether ththugh can solve them. In Findings of the Association for Computational Linguistics: ACI. 2023, pp. 13003-13051, 2023.
15

---

Published as a conference paper at ICLR 2025
Gemma Team, Morgane Riviere, Shreya Pathak, Pier Giuseppe Sessa, Cassidy Hardin, Surya Bhupatiraju. Léonard Hussenot, Thomas Mesnard, Bobak Shahriari, Alexandre Ramé, et al. Gemma 2: Improving open language models at practical size. arXiv preprint arXfr:2408.00118, 2024.
Gemma Team, Aishwarya Kamath, Johan Ferret, Shreya Pathak, Nino Vieillard, Ramona Merhej, Sarah Perrin, Tatiana Matejovicova. Alexandre Ramé, Morgane Rivière, al. Gemma technical report. arXiv preprint 1xXi::2503.19786 2025.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXIv:2302.1397,, 2020.
Jos Van Der Westhuizen and Joan Lasenby. The unreasonable effectiveness of the forget gate. arXiv preprint wrXiv:1804004849, 2018
A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, et al. An empirical study of mamba- based language models. arXiv preprint arXiv:2406.078877 2024.
Fengjun Wang, Sarai Mizrachi, Moran Beladev, Guy Nadav, Gil Amsalem, Karen Lastmann Assaraf, and Hadas Harush Boker. Mumic-multimoda embedding for multi-label image classification with tempered sigmoid. In Proceedings of the AAAI Conference Artificial Intelligence, volume 37, pp. 15603-15611, 2023.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Representations, 2024.
An Yang, Baosong Yang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li, Dayiheng Liu, Fei Huang, et al. Qwen2 technical report. arXiv preprint arXiv:2407. 10671, 2024a.
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with efffweee-efficceen training. In Forty-first International Conference on Machine Learning, 2024b. URLLAtttss//opeeeeee
net/ forum? wiw.unndddiiixaiiiiaa.aaaaa UJT
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.0644 06484, 2024c.
Tao Yu, Rui Zhang, Kai Yang. Michihiro Yasunaga, et al. Spider: A large-scale human-labeled dataset for complex and cross-domain semantic parsing and text-to-sql task, 2019.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.0n 00066 2022.
Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri,, Myle Ott, Sam Shleifer, et al. Pytorch fsdp: experiences on scaling fully sharded data parallel. arXiv preprint arXiv: /177, 2023.
Terry Yue Zhuo, Minh Chien Vu, Jenny Chim, Han Hu, Wenhao Yu, Ratnadira Widyasari, al. Bigcodebench: Benchmarking code generation with diverse function calls and complex instructions. arXiv preprint arXiv:2406. 2024.
16

---

Published as a conference paper at ICLR 2025
A
MAIN NOTATION INTRODUCTION
Table 3: Meanings of Notations
Notation
Size
Mcaning
T
Constant
The length of the context.
d
Constant
The model dimension.
Constant
dimensionality the hidden corresoonding token mixer.
Constant
The expansion factor of matrix hidden state, w when is set to 1 in vector hidden state. head dimension
da
Constant
multi-head mechanism.
1
I X
The ull-ones vector of dimension
Tx
input sequence module. The output sequence of module.
O
T d
Q
T xx.
The query.
x
The key.
V
Txm
The value.
W
dxn
weight matrix the query.
Wx
dxn
The weight matrix of the key. weight matrix the value. The hidden state stamp transition matrix at stamp f.
Wv
S.
axm m
axm
The input matrix stamp The output matrix at stamp f.
C.
x
1 1xn
The gating mechanism along dimension
the trunsition.
6
xxa
The gating mechanism along the dimension acts on the input, which is negatively correlated with gating mechanism along dimension The guting mechunism along the dimension acts on the input, which is negatively correlated with 8,
transition.
B.
1
4
x m
The input term of recurrent calculations.
The output of ShortConv.
g:
Variable
The gating mechanism of selection or decay. In Rodimus, it's a selection gate with the size x
W
The weight matrix of gr-
b,
1 1xn
The bias of gt-
temperature gate Rodimus.
W.
2A X
The weight matrix of 7, in Rodimus.
b,
1
bias Rodimus,
Constant
The low-rank dimensionality of projections.
W!
m x
The down-projection weight matrix Rodimus.
W
Exm
The up-projection weight matrix Rodimus.
ba
1x m
The bias of Rodimus.
d.
x
The feedthrough matrix in SSMs.
B
METHOD DETAILS
B. SIMILAR FUNCTIONAL FORMS OTHER METHODS
We summarize the functional forms of A, and B, used in state-of-the-art methods in Table 4
As illustrated in the above table, the specific formulation of DDTS provides several notable advan- tages:
Dynamic Control of Hidden State Updates: DDTS allows for data-dependent selection through A, and B, when writing and deleting information, enabling dynamic control of hidden state updates. This feature is discussed in Section 3.1.2] with a formal proof provided in Appendix C22 Efficient Parallelization: As DDTS performs data-dependent selection solely along the dimension n, supports the implementation of chunkwise parallel algorithms. This capability enables efficient parallel training while maintaining flexibility, elaborated in Section 3.1.3 and Appendix B.3
17

---

Published as a conference paper at ICLR 2025
Table 4: Comparison of A, and B₁ across various models. For models employing the multi-head mechanism, only the gates for single head are presented. The term "Gate Dim" represents the dimensionality of the gates gt A; and B₁, and "L.A." denotes linear attention without scal- ing (Katharopoulos et al. 2020 Qin et al. 2022a). Elements with the subscript & are data-dependent gates, with most elements being learnable. Notably, the non-learnable constants include ? in RetNet, in GLA and gRetNet, and in Hawk. Although the "State Size" of RWKV6 is smaller than that of
Rodimus, the use DDTS in Rodimus yields superior performance across all metrics (sce Tables 2
State Size x
Parallelization like L.A.
Recurrent
Model
Gate Dim
Type
A
B
L.A.
Kutharopoulos al.
2020
No Gates
256
Gates
1,, 1m
1 1m
RetNet Sun al 20230
1x] 1
256
Decay
1m
1, 1m
gRetNet
(Sun
2024
1x1
256
Decay
(o(9))4At)
- 1m
1 1m
GLA
Yang et
2024b
1 x
256
Decay
(a(g:))
!r] 1m
1 1m
RWKV6
Peng et al.
2024
1 xx
64 x
Decay
exp
- exp(a)) 1m
1 1m
HGRN2
Qin al.
20244
128
Selection
x(9)) 1m
(o(()))
1
Hawk
et al 2024
1xd
1xd
X
Selection
vI (as
S4D
2022a
1xd
16
X
Selection
exp( 1.19)
Mamha
Gu & Dan
2023
1xd
16 x 2d
X
Selection
exp( - @ 1.
Mamba2
Dao Cu.
2024
1x1
128 2d
Selection
exp(-o
1m
Clai) 11.1
Rodimus
18x
64 x 2d
DDTS
exp( @ T2) 1m
(s 8.
Enhanced Flexibility through the Temperature Gate: The temperature gate тę introduces
additional flexibility the selection mechanism, allowing for more precise control over information compression and representation. This aspect sets our approach apart from existing works, as noted in Section 3.1.2 and detailed in Appendix E.2
B.2 RELATED WORK
Linear Attention with Gates: Linear attention replaces the traditional exponential kerne,, exp(g, kz), with the dot-product feature map <(g:) ((k,) (Katharopoulos et al. 2020). This approach allows inference to be reformulated as a linear RNN with matrix hidden state. However, without gating mechanisms, it often struggles to effectively capture imput and forget historical infor- mation, resulting in irrelevant information within the hidden state (Van Der Westhuizen & Lasenby 2018). To address this, gating mechanisms such as purely decay (e.g., GLA (Yang al. (024b)) and selection mechanisms for associating input and decay (e.g., HGRN2 (Qin et al. 2024a have been introduced.
These improvements, however, often lead to the issue of sacrificing hidden state size to enhance performance (Qin al. 2024a), lacking thorough analysis and design of gates, which hinders performance even with larger hidden states. In Section 3.1.1] we analyze effective gate design and propose a solution for a more efficient gating mechanism named DDTS.
State-Space Models: Recent advancements in state-space models based on HiPPO (Gu et al. 2020 have led notable results in long sequence modeling, inspiring further work like S4 (Gu et al. 2022b), S5 (Smith et al.) 2023). and H3 (Fu et al. 2023). Mamba (Gu & Dao 2023 builds on S4D GGu al. (202aa), the parameterized diagonal version of S4, making state-space model (SSM) parameters data-dependent. Its input-based selection mechanism arises from ZOH discretization. However, Mamha's diagonal state space structure limits calculations with larger hidden states expanded by the expansion factor nt restricting memory capacity improvement.
Mamba2 (Dao & Gu 2024) addresses this by using a scalar state space structure instead of the original diagonal state space structure in Mamba. This change allows Mamba2 to be trained with chunkwise parallelization, similar to linear attention (Sun et al. 2023a Ma et al. (022). permitting an expansion factor of 128. However, this approach mirrors linear attention' challenge of over-relying on hidden state size while neglecting gating design. The oversimplified gating mechanism struggles approximate the original softmax attention performance (see Section 3.1.1
18

---

Published as a conference paper at ICLR 2025
Our method features a similar architecture to Mamba2, enabling chunk-wise parallel calculations, while our gating mechanism, DDTS, better approximates softmax attention.
Softmax Attention with Custom Attention Mask: Several studies have enhanced the attention mechanism while retaining the softmax function. One approach involves modifying the attention scope by integrating prior knowledge create diverse custom attention masks. During inference, these masks function similarly to token-level sequence compression, storing unmasked tokens in the KV cache (Xiao et al. 2024). Many studies observe local behavior in attention, leading to the use of local attention instead of full attention, as seen in models like LongFormer (Beltagy et al. (220). Additionally, some methods employ learnable patterns, such as the Routing Transformer (Roy et al.. 2021). to construct attention masks in data-driven manner. However, custom attention masks only support token-level compression, which can result in information loss for masked tokens.
In Rodimus++ the Rodimus block can compress long contexts and transform token embeddings into context embeddings by dynamically fusing multiple tokens to create semantic compression, as shown in Figure 2p. This enables the local attention in SW-SKA to capture token information beyond the window.
B.3 CHUNKWISE PARALLEL REPRESENTATION OF THE RODIMUS BLOCK
Let A. = I[ aj denote the cumulative product of gates aj, and define A = (AI,...,,T)) RTx", then A, - G denote the cumulative product of gates By, and define A
(A AT} E Tm After the following derivation, we can obtain the parallel form of the Rodimus block:
S. (a, (.) @ S + (a, 3.) ② (k] ve)
= (a, B.) @ S (ā. @ kt) (Be o vz),
(14)
0, S.S
from the derivation in Yang et al.
2024b) we can have:
o = q: Σ
II
(a ( (a, o k.) (
i=1
j=i+1
T II k))(((0k)T(
q:
j=it1
II 0,o00A1))
(15)
j=++
II O100000)) s=1+1
Σ (g O (404))(ak/44))
/ (0.00/AA))
T A.
Q QGA, K =K V-V/AOB,
0 (QK @ MM,,
(16)
Ô Oc A.
It follows from Eq. 16) that the parallel formula can be written as:
[Q& C (K/AA @ T @ M) o 3) A.
(17)
Since we set B, in Section 3.1.1 the above equation can be simplified as:
0
(Q ()(K/A T c
(V о (
(18)
which the parallel form for the proposed Rodimus block. However, the complexity of computing
O via Eq. (18) is quadratic. To further reduce the training complexity, we exploit chunkwise
integration of the recurrent and parallel forms, in analogy to other linear attention approaches (Yang
19

---

Published as a conference paper at ICLR 2025
et al. 2024b).
Specifically, first divide the entire sequence into T/B consecutive chunks, each with length B. For each chunk, the parallel computation as described in Eq. (18) is employed to obtain the intra-chunk output Ointra, Subsequently, information between chunks is integrated using recurrence approach akin to Eq.
111
S = S 1| O (CI (((1)))+ 1)
(K o Eqg O II aj) (T)))(V ① STa),
O11 ginter
Q @ (=(1))) Oz S|]=
(+-1)++b
(19)
where [i] represents the i-th chunk. The parallel computations within each chunk can further leverage tensor cores speed up matrix operations (Yang et al.j 2202bb Dao & Gul (024). Additionally, this sequential form for inter-chunk computation reduces overall computational complexity from
quadratic to subquadratic and optimizes memory utilization during training. To further optimize training efficiency, we incorporate Fl reatttentt
B.4 RODIMUS+'S RELATIONSHIP TO EXISTING HYBRID MODELS
Here we focus on models that integrate self-attention with recurrent architectures. One notable approach is Griffin (De et al. 2024 Botev al. 2024), which combines sliding window attention with a linear RNN. However, unlike Rodimus+, which employs matrix hidden states, the hidden state in Griffin's linear RNN is a vector, limiting its capacity to retain historical information. Another hybrid model, Based (Arora et al.) (024b), also integrates linear attention and sliding window attention. Unfortunately, lacks decay or gating mechanism within linear attention framework, making it susceptible to the attention dilution problem. On different note, Samba (Ren et al.) (024) explores various methods to combine Mamba with self-attention. The original version of Samba inserts MLP (FFN) layers between Mamba and self-attention, which neglects Mamba's channel mixing capabilities and introduces unnecessary parameters. Another variant of Samba, Mamba- SWA-MLP, then removes the MLP between Mamba and self-attention, but it does not implement the two-hop residual connection used in Rodimus + Different from the aforementioned works, Mamba-2- Attention (Dao & Gu 2024) seeks enhance retrieval (i.e., recall) capabilities by replacing some of the SSM layers with self-attention. This approach contrasts with the model architecture of Rodimus+ which regularly incorporates SWA between SSM layers. Moreover, determining the optimal number of layers to replace with attention requires further experimentation, and the introduction of full attention layers with their quadratic complexity complicates model scaling. Lastly, Jamba (Lieber 2024) combines the Transformer block with Mamba and MoE layers, yet it also adds extra MLP
et al. (024) layers after Mamba, similar Samba. Additionally, establishing the optimal number of Transformer blocks remains an unresolved challenge.
B.5 MODELS IN EXPERIMENTS
Transformer Brown et al. 2020 The foundational architecture with causal attention mask used in GPT-3.
Transformer+ Touvron et al. 2023): An enhancement of the Transformer, Transformer
incorporates RoPE, GLU, and RMSNorm for improved performance. is currently the most robust and widely used Transformer configuration (Gu & Dao 2023).
Mamba Gu & Dao 2023): Mamba makes the original SSM parameters data-dependent and introduces an 1/O-aware associative scan.
Mamba2 Dao & Gul 2024):
Mamba2 simplifies the diagonal state space structure in Mamba to the scalar state space structure, enhancing computational efficiency through chunkwise parallelism algorithms like linear attention.
GPT-Neo (Black et al. 2021): Similar to the Transformer, GPT-Neo employs full attention and local attention in alternating layers and uses ALiBi for positional encoding.
OPT (Zhang et al. )022): OPT mirrors the Transformer architecture but utilizes learnable positional encoding.
https///gtthh.ccc////rttteeg//na inarr- attentittetiiii
20

---

Published as a conference paper at ICLR 2025
BLOOM (Le Scao al.
2023):
BLOOM is akin to the Transformer, except it uses ALiBi for
positional encoding.
An improvement upon the Transformer, Pythia uses RoPE for Pythia Biderman et al.
2023):
positional encoding.
RWKV4 (Peng al. states with Transformer-sty -tyle attention mechanisms.
2023
RWKV4 integrates techniques such token-shift to merge RNN-like
RWKV6 (Peng et al.
2024): : An enhanced version of RWKV4, RWKV6 features more flexible
gating mechanism and introduces multi-headed matrix-valued states to boost the memory capacity of the recurrent model.
RetNet (Sun et al.
y RetNet applies RoPE with a fixed scalar decay y to linear attention.
GLA Yang al.
2024b): GLA omits positional encoding on linear attention, opting instead for
data-dependent decay.
HGRN (Qin al. 2024b): HGRN introduces decay lower bound based on traditional linear RNNs, employs recurrent calculation in the complex domain, and uses a selection mechanism the real domain.
HGRN2 (Qin et al., 2024a):
HGRN2 removes complex domain calculations from HGRN and introduces expansion similar to Mamba2 and linear attention.
Qwen2 Yang et al.
2024a): Built on +ransformer++ Qwen2 uses GQA to reduce inference costs.
Llama3 (Dbeyy al.
2024 Llama3's architecture is similar that of Qwen2.
RecurrentGemma (Botev et al. 2024): RecurrentGemma a series of open language models built
upon Google's innovative Griffin architecture (De al. 2024).
C
PROOF
C.I LEMMA FOR SELECTION MECHANISM
Lemma 1. For the recurrent model defined below,
S: = ( o St + @b(1)) o us
(20)
where St E Raxm, 2t E IR 1x us = dx(xx) Ruxm. .T7 = (x1))) E Raxm If the following condition met, is a model with the selection mechanism.
mm]. d dob (Nts))) 0.
Vi (0, .,.....) Vj (0,11 1....,m,,,
(21)
dm dma..
Proof. Since the product of the gradient of (M(.5))) and ΦB (T).)) is negative, )a(7)))) is negatively correlated with (n(.)))) satisfying the selection mechanism.
C.2 PROOF OF PROPOSITION
To prove Proposition we first verify that Rodimus satisfies Lemma 1 In Rodimus, = k n and we define 7): a W, + by)" 1 Consequently, g softplus( m.) For simplicity, we omit temperature T/ in this proof. We define ( = exp(-g) and da (7:) g/. which corresponds to Eq. IT). Finally, we find that:
d don(m)ss)) -a (T...)) o (m.,))) < 0, dmu.s dm,is
(22)
where is the sigmoid function. Thus, Rodimus satisfies the selection mechanism.
21

---

Published as a conference paper at ICLR 2025
D EXPERIMENTAL DETAILS
D.1 LANGUAGE MODELING ON WIKITEXT-103
The training settings are listed in Table 5 For the models trained on WikiText- 103, the number of layers and dimensions are set to and 512, respectively. The main difference from Qin et al. (2024b) that we increase the total batch size, which reduces the training steps to obtain results more quickly.
Table 5: Settings details for WikiText-103 and the Pile subset.
Dataset
WikiText-103
Pile subset
Tokenizer method
BPE
GPT-Ncox Tokenizer
Vocab size
50265
50277
Sequence length
512
2048
Total Batch size
256
256
Training steps
25000
13351
Warmup steps
2000
572
Peak learing rate
5e-4
5e-4
Min learning rate
le-5
1e-5
Optimizer
AdamW
AdamW
Adam $1
0.9
0.9
Adam B2
0.95
0.95
Weight decay
0.1
0.1
Gradient clipping
1.0
1.0
D.2 LANGUAGE MODELING ON PILE
In addition to Tiki we also train various model architectures, including Transformer (full attention), Transformers+ with sliding window attention (sparse attention), HGRN2 (linear attention), Mamba/Mamba2 (state-space models), Rodimus, and Rodimus+, using 7B tokens in the Pile dataset. This dataset is considerably larger than WikiText- 103. This training aims investigate the language modeling performance across different model architectures, with the training settings adhering to the 350M parameter configuration outlined in Table g The results are summarized in T
Consistent with the findings in Table T Rodimus demonstrates superior language modeling perfor- mance compared to other models, including Transforner+ and Mamba2. However, unlike the results from Table [ we find that Rodimus+ and Rodimus exhibit similar performance levels under this scenario. We attribute this similarity to the limited scale of the dataset. Note that Rodimus- needs to distribute the tasks of global and local dependency modeling respectively to the Rodimus block and the SW-SKA, making it more challenging train than the original Rodimus and thus requiring more data. Despite this, Rodimus+ shows greater potential for enhanced language modeling performance compared to Rodimus when deployed at larger parameter scale, illustrated in Figure 4
D.3 ACCURACY-EFFICIENG TRADE-OFF
We first present the accuracy-efficiency trade-off among Rodimus* recurrent models (GLA, HGRN- 2, Mamba, and Mamba-2), and sparse attention models (StreamingL using the WikiText- dataset. The experimental settings are detailed in the WikiText- 103 column of Table while the results are illustrated Figure Ia To manipulate the memory footprint (focusing solely on the cache), we adjust the scaling factor for the recurrent models. In contrast, for StreamingLLMM, modify the number of recent tokens or the size of the sliding window. is important to note that the input history sequence length is set at 512, which is also the recent token size corresponding to the rightmost endpoint of the StreamingLLM curve (i.c., the blue curve). This configuration aligns with that of the standard Transformer, which attends to the entire historical token sequence. Consequently,
22

---

Published as a conference paper at ICLR 2025
Table Language modeling on the Pile dataset.
Model
Valid PPL
Transformer+
6.04
Transformer+ SWA
6.13
HGRN2
6.08
Mamba
6.19
Mamba2
6.02
Rodimus
5.88
Rodimus+
5.88
the performance of the standard Transformer can be referenced at the rightmost endpoint of the StreamingLLMM curve.
As demonstrated in Figure Ta Rodimus* exhibits superior modeling performance compared to recur- rent models, even when operating under the same fixed memory footprint. Furthermore, Rodimus* also surpasses both the sparse attention model (StreamingLLMM and the standard Transformer base- line. These findings underscore Rodimus capacity to effectively balance accuracy and memory efficiency, achieving substantial trade-off between the two.
To further validate the memory efficiency of Rodimus, Rodimus, and StreamingLLMM on larger dataset, as shown in Table 7 The foundational experimental settings are aligned with the 255M configurations described in Table Once again, we adjust the memory footprint by modifying the expansion factor n. The results of PPI. summarized in Table indicate that Rodimus continues to outperform the accuracy-efficiency trade-off when trained at larger scale compared to WikiText-103.
we conduct experiments using Mamba2,
Table 7: Modifying the memory footprint to test the LM task of PPL on the Pile dataset.
Model
1.2M Bytes
2.3M Bytes
4.7M Bytes
9.4M Bytes
StreamingLLLM
24.76
17.54
12.57
9.39
Mamba2
8.16
7.91
7.61 7.81
Rodimus
7.60
7.47
7.35
7.33
D.4 SCALING LAWS
We train all models on subset of the Pile, using the parameter settings from GPT-3 (Brown et al. 2020). The context length is 2048 tokens. The training steps and total token counts follow the Mamba settings (Gu & Dao 2023). Detailed configurations are provided in Table 8 For methods using purely recurrent models that combine GLU and token mixers, such as Mamba series and Rodimus, the number of layers should be doubled match the GPT-3 specifications.
is worth noting that Rodimus+, we replace the second Rodimus block in each layer with SW-SKA and FFN, but the overall layer count remains unchanged. To clarify the parameter count: the Rodimus Block consists of 6d parameters, where d the model dimension. Since each layer contains two Rodimus blocks, the total parameter count for Rodimus sums 12f In contrast, the Rodimus+ Block consists of 13d2 parameters, comprising one Rodimus block (6d2), alongside multi-head projection, multi-query projection, and output projection in the attention mechanism (3d?). and the FFN (4d2).
D.5 DOWNSTREAM EVALCATION
Our default training settings align with those in Table & of the scaling law experiment. Unless otherwise specified, these settings remain unchanged. All our models use mixed precision and
23

---

Published as a conference paper at ICLR 2025
Table Model parameters and training settings used in our scaling experiments. "Steps" indicates the number of training steps, and "Batch Size" refers to the number of tokens used for each update.
Params
llayer
d mmodel
n heads/d head
Steps
Learing Rate
Batch Size
Tokens
125M
12
768
12/64
4800
6e-4
0.5M
2.5B
350M
24
1024
16/64
13500
3e-4
0.5M
7B
760M
24
1536
16/96
29000
2.5e-4
0.5M
15B
1.3B
24
2048
32/64
50000
2e-4
0.5M
26B
FSDP (Zhao et al. 2023) to improve training speed. To maximize GPU memory usage, we adjust the batch size according to Table 1 and modify the maximum and minimum learning rates.
The Rodimus Tokenizer uses Hugging Face's s byte-level BPE algorithm, featuring a vocabulary of 126,340 tokens, including special tokens, trained 60 billion token subset of the curated dataset.
Notably, we do not apply weight decay to the biases of linear projections and LayerNorm, nor to the weights of LayerNorm or RMSNorm. Additionally, set the head size of SW-SKA in Rodimus to 128 and the window size to half the training context, as described in Samba Renn et al. (024).
Table 9: The extra training settings of Rodimus* in downstream task.
Params (B)
Tokenizer
Training context
Batch size
Max learning rate
Min learning rate
0.13
GPT-NeoX
2048
1.5M
3e-3
le-5
0.46
Rodimus
2048
0.8M
1 1.5c-3
le-5
1.4
Rodimus
2048
6.3M
Ic-3
le-5
D.6
FAIR COMPARISON OF LARGER SCALE MODELS ON DOWNSTREAM TASKS
In the first block of Table 2 noted that the model size relatively small, consisting of only 130M parameters. To further examine the performance of Rodimus in comparison to Mamba2, we conducted experiment using larger models. Specifically, with around 1.4B parameters from scratch under identical conditions, utilizing 100B tokens from the Pile dataset. We follow the configurations outlined in the last row of Table g The results presented in Table I0 where it is evident that Rodimus, trained using the same dataset and tokenizer as Mamba2, consistently outperforms Mamba2. This performance enhancement can be attributed to the innovative design of the DDTS.
trained both Rodimus and Mamba2
Table 10: More Results of Downstream Tasks via Zero-Shot Evaluation.
Model
Params
Tokens
Tokenizer
ARC-e ax.aom
ARC-e
HS
LMB F:1
LMB =:1
OBQA
PIQA
WG
AVG
db
sct
scc
Mamba
0.13
100
NeoX
23.72
46.84
33.37
19.72
40.33
29.80
65.07
52.17
41.61
Mamba2
0.13
100
NeoX
23.81
45,50
34.13
20.42
39.55
29.20
64.47
51.85
41.22
R
0.13
100
NeoX
23.38
46.84
34.22
17.95
42.44
30.00
65.40
51.46
41.96
Mamba2
1.4
100
NeoX NeoX
26.54
56.06
48.74
8.22
55.25
33.40
70.57
54.14
49.24
R
100
28.75
56.61
49,51
7.78
56.05
34.00
70.78
52.41
49.73
Here, bold, underting, wavy undertline indicates best, second, indicates mouels using fewer tokens. The superserign demotes models trained from scratch using the same configuration for fair comparisom.
result, respectively. column "Tokens", underling
D.7 MQAR (MULTI-QUERY ASSOCIATIVE RECALL)
The MQAR task (Arora et al. 2024a) is a benchmark widely used to assess the associative recall capabilities of recurrent models. Note that quadratic softmax attention based models, such as the vanilla Transformer, achieves perfect scores across all configurations, and so these models are excluded from detailed analysis (Yang al. 2024b Arora et al. 2024a). In this task, input sequences
24

---

Published as a conference paper at ICLR 2025
consist of several key-value pairs followed by queries. When presented with query, the model must retrieve the corresponding key-value pair from earlier in the sequence in order to accurately predict the next token.
For example, consider the input "A 3 B 2 C I". The key-value pairs are "A 2B and "C 1". Later, when "A ?", "B ?", and "C ?" appear queries, the correct outputs are "3", "2", and "1", respectively.
Notably, models utilizing full softmax attention tend to achieve perfect scores on this benchmark. Consequently, our study concentrates exclusively recurrent models, specifically Mamba (Gu & Dao ).23), Mamba2 (Dao & Gu. 2024), RWKV (Peng et al. 2023). and RWKV6 (Peng et al. 2024).
The experiment in Figure 6 follows the settings from MQAR (Arora al.. 2024ap. We use sequence lengths of 128, 256, and 512, with corresponding key-value pairs of 8, 16, and 32. The total vocabulary size is set to 8192. the number of model layers to 2, and the model dimension d varies as 64, 128, 256, and 512. Learning rates of 10 10-3.3 10020-1.- and 10-2 are applied to train for 64 epochs using the cosine annealing learning rate schedule, with the best accuracy recorded as the reported accuracy for each model dimension. In Figure [b] the sequence length is increased to 1024, the key-value pairs are set to 256, and the model dimension is fixed at 256. Only the expansion factor n. is varied. Learning rates of 10-2, 10-2.5 and 10-3.5 are used to train for 32 epochs with the cosine annealing learning rate schedule, and the best accuracy is recorded as the reported accuracy corresponding to the expansion factor.
We first depict the accuracy as a function of the expansion factor n in Figure lb In this analysis, we focus on the Mamba series for comparison, as their architecture closely resembles that of Rodimus* allowing for consistent settings across all models. Our results indicate that Rodimus" exhibits superior recall ability at the same expansion factor This enhancement can be attributed to the model's DDTS mechanism, which filters out irrelevant past information and current inputs more effectively.
Next, we present the accuracy as function of model dimension d in Figure 6 where the expansion factor n is maintained at its original value across different models. Once again, Rodimus achieves the highest performance. In comparison, Mamba and RWKV underperform, particularly when the model dimension is small. This suggests that the fixed-capacity state of these models can only deliver reliable recall results when the state size is sufficiently large to retain relevant past information.
Sequence Length: 128 KV Pairs:
Sequence Length: 256
Sequence Length: 512
Pairs: 16
KV Pairs: 32
1.00
Mamba
Accuracy
0.75
Mamba-2
0.50
RWKV-6
0.25
RWKV
Rodimus
0.00
Rodimus
64
128
256
512
64
120
256
512
64
128
256
512
Model dimension
Model dimension
Model dimension
Figure 6: Recall accuracy of MQAR Task at different sequence lengths.
D.8 NEEDLEBENCH
The NeedleBench (Li et al. 2024 is retrieval evaluation method that randomly inserts key infor- mation into long texts to create prompts for LLMs. The primary objective of this test is to assess the ability of LLMs to effectively extract crucial information from extensive documents, thereby evaluating their proficiency in processing and comprehending long-form content. The NeedleBench comprises several tasks: the Single-Needle Retrieval Task (S-RT), the Multi-Needle Retrieval Task (M-RT), the Multi-Needle Reasoning Task (M-RS), and the Ancestry Trace Challenge (ATC). The S-RT focuses on the precise extraction of specific details, while the M-RT evaluates the retrieval of multiple related pieces of information, simulating complex queries. The M-RS task involves synthesizing and understanding interconnected information, and the ATC examines logical reasoning, requiring comprehensive engagement with all provided content to solve intricate problems.
In this study, we evaluate Mamba-1.44 1.4B (Gu Dao 2023), Mamba2-1.3B (Dao & Gu 2024), Pythia-1.4B (Biderman et al. 2023). Rodimus-1 1 .4B, and 4-limus-+11.63 on NeedleBench using OpenCompasa Contributors 2023) since the sequence length of all these models is 2048 tokens.
25

---

Published as a conference paper at ICLR 2025
The brief evaluation results, featuring Mamba2-1.33 1.33B, Pythia-1.4B, Rodimus-1.49 and Rodimus--- 1.6B, are presented in Figure Full evaluation results, including subtasks, are detailed in Figures M.4)ba-144B), TO (Mamba2-T.3B), [T (Rodimus-144B) [ (Pythia-1.4B), and [3 (Rodimus+-1.68)
The "Overall Score" represents the average result of these subtasks. Additionally, we present the NeedleBench results for Rodimus++ during training in Appendix [3.3] where the model has trained on the 4K-length context with 2.5T tokens, and demonstrates superior long-distance retrieval capabilities, as shown in Figure
Many studies have suggested that pre-trained recurrent models tend to underperform on challenging recall-nnennsive tasks like NeedleBench when compared to models utilizing full attention mechanisms (Fu et al., [023] Waleffe et al.) 2024). This trend is illustrated by comparisons between Pythia and Mamba2. Nevertheless, Rodimus demonstrates superior recall ability relative to Pythia. While Pythia shines in scenarios involving shorter contexts (Ok-4k), Rodimus exhibits better extrapolation capabilities for longer contexts, leading to overall enhanced performance. Furthermore, Rodimus hybrid model, further improves recall ability by integrating global semantic-level and local token- level representations. The results for Rodimus and Rodimus+ indicate that the Rodimus family effectively navigates the eecuracy-efficieney trade-off, particularly in terms of enhancing recall ability. addition, training with longer contexts can effectively enhance the model's recall ability in Needlebench. Rodimus+ trained with the 4K-length context performs better Needlebench than when trained with the standard 2K-length context.
D.9 MORE RECALL -INTENIVE BENCHMARKS OF PRE-TRAINED LANGUAGE MODELS
The NeedleBench benchmark may be excessively demanding for small pre- pre-trained models with context length limitation of 2K tokens. To mitigate this issue, we introduce several additional, simpler recall benchmarks that are outlined below:
FDA: This benchmark involves extracting specific details, such as device codes, classifications, and intended uses, from medical device reports submitted to the FDA.
SWDE: This task focuses retrieving particular information, such movie directors and univer- sity tuition, from HTML pages across 14 websites in domains like Movies and Universities.
SQuADv2: A question-answering task where models identify answer spans within Wikipedia articles based on given queries.
TriviaQA: This benchmark challenges models to answer questions using information from Wikipedia and other broad web sources, featuring diverse question formats.
NQ: Utilizing real Google search queries, this task requires models to find answers within Wikipedia, necessitating the extraction of relevant text snippets.
DROP: A massoning-intensin question-anse answering task where models must carry out opera- tions incluinc counting, sorting, or performing arithmetic- information derived from Wikipedia passages.
The same checkpoints used in the NeedleBench section employed for evaluation, with the accuracy results presented in Table II] The average scores for both Rodimus and Rodimus outperform those of Mamba and Mamba2. However, contrary the trends observed in NeedleBench, Rodimus demonstrates a slight performance lag behind Pythia. This finding supports our earlier analysis in Appendix [D.8 which indicates that the softmax attention-based model Pythia surpasses recurrent models in recall tasks within the con-cxtrappllatin interval. Notably. the hybrid model Rodimus which integrates SW-SKA, successfully exceeds Pythia's performance, highlighting that the incorporation of this hybrid model effectively addresses the performance limitations of recurrent models.
D.10 ABLATION STUDIES DETAILS
In this section, we present ablation studies designed to achieve four primary objectives: (i) assessing the contributions of various components-suce as the gates and the two-hop residual-in R (ii) evaluating performance with different head compression methods, and (iii) determining the optimal hyperparameter values (L and n), (iv) analyzing the ordering of the Rodimus block and SW-SKA plus FFN in Rodimus+. We begin with experiments conducted on the WikiText- 103 dataset, chosen for its relatively small size, which allows for quicker results and makes the impact
26

---

Published as a conference paper at ICLR 2025
Table 11: Recall results of models via zero-shot evaluation measured by accuracy. The input context length is set to 1024, and the maximum generation length is 512.
Model
SWDE
FDA
SQUAD
Drop
NQ
TriviaQA
AVG
Pythia
53.05
81.65
41.89
21.85
34.49
58.18
48.52
Mamba
45.36
59.85
41.28
20.84
34.43
60.43
43.70
Mamha2
52.67
71.12
42.19
22.86
35.79
63.57
48.03
Rodimus
49.77
79.47
43.10
22.76
32.56
61.85
48.25
Rodimus+
58.76
80.84
48.74
24.58
36.27
64.63
52.30
of removing specific components easily observable. The second of experiments utilizes a 7B subset of Pile, providing a larger dataset to discern subtle differences arising from hyperparameter adjustments. Unless otherwise specified, the models trained on WikiText- and Pile follow the training settings detailed in the second and third columns of Table respectively. Regarding model configurations, for the models trained on WikiText- 103, the number of layers and dimensions are set to 6 and 512, respectively. For the models trained on Pile, the various model configurations provided in Table 3 For the WikiText- 103 dataset, we report both the test PPL and the best test PPL. The test PPL refers to the evaluation of the test set after training the training dataset for fixed number of epochs, while the best test PPL reflects the lowest value achieved during early stopping of training. Although the former approach is commonly employed in the literature, it does not always accurately capture model performance.
The necessity of gt, Tr, and Be: As presented in Table 15 the removal of these gates results in decreased performance, indicating their essential role in DDTS.
The necessity of 3; Table 15 shows that incorporating 3e does not reduce the PPL further; in fact, it tends to increase it slightly. As shown in Table 16 this finding suggests that excluding B, is preferable, as it does not provide performance improvements but may complicate the training process and hinder the training efficiency (see detailed explanations Section 3.1.1.
The necessity of the two-hop residual: According to Table 17 the two-hop residual in Rodimus can slightly enhance performance when the model deepens with more layers. This mechanism aids in stabilizing training, as noted by Ma et al. (002)). Given that modern Large Language Models (LLMs) typically consist of many layers and the two-hop residual incurs extra computational cost, implement it in Rodimus+.
Comparison between SKA, GQA, and MHA:
In this experiment, the model settings refer to 125M item in the Table 8 Table inndicatsss that SKA consistently outperforms GQA in terms of Transformer demonstrating the effectiveness of SKA. However,
PPL across both Rodimus+ SKA performs less optimally than MHA, likely due to increased training difficulty, requiring more iterations data for comparable performance. Nevertheless, since SKA is more parameter-effiiien than MHA and surpasses GQA in performance, we opt to utilize SKA in Rodimus+. Additionally, we provide analysis of the memory footprint of the cache, excluding the model weights. SKA demonstrates a substantial reduction in memory footprint compared to MHA, although its cache size slightly larger than that of GQA.
Sensitivity to l and n: Table reveals that increasing the rank é of the weight matrices in Eq. 10 generally improves performance but also raises the parameter count. We find that setting f 16 effectively balances performance and parameter quantity, and thus we adopt this value for our experiments. Regarding the state expansion factor n, shown in Table [3] we observe that valid PPI. decreases increases; however, the rate of decrease slows significantly for 64. Consequently, we establish n 64 for our experiments. Note that n = 128 in Mamba2. Thus, Rodimus outperforms Mamba2 in most tests with only half the state size.
The ordering of the Rodimus block and SW-SKA plus FFN: We investigate the impact of the ordering of the Rodimus Block, SW-SKA, and FFN model performance. The results detailed in Table Specifically, for the Pile subset, the model settings correspond to 350 million items, as referenced Table Since SW-SKA and FFN integrated within a two-residual hop structure,
27

---

Published as a conference paper at ICLR 2025
treat them single unit and focus the relative positioning of the Rodimus Block with respect to SW-SKA and FFN. The term "Rodimus+X" denotes the original order presented Eq. (13). while "X+Rodimus" indicates the rearrangement of Rodimus's position relative to SW-SKA and FFN in the same equation.
The findings in Table 18 reveal that for shallow models tested the WikiText-103 dataset, it is essential to place the Rodimus Block before SW-SKA and FFN. In contrast, the results from the Pile dataset suggest that the number of layers increases, the significance of the Rodimus Block's positioning diminishes. We hypothesize that this occurs because an increase in layers lessens the impact of module types that are closer the word embedding and output projection on overall model performance.
Table 12: Results of the ablation experiments on the low-rank term.
Table 13: Results of ablation experiments on ex-
pansion factor
Valid PPL
Params (M)
Valid PPL
Params (M)
Valid PPL
Params (M)
Valid PPL
Params (M)
7.41
126.3
5.84
360.8
7.32
126.4
5.81
361.0
16
5.75
363.7
7.26
127.4
7.36
126.5
5.77
361.4
5.66
370.0
7.13
129.8
7.30
126.8
5.76
362.2
7.26
127.4
5.75
363.7
382.6
7.01
134.5
32
7.31
128.6
5.74
366.9
128
5.58
407.8
6.95
143.9
7.24
130.9
5.71
373.2
128
7.17
135.7
5,70
385.8
256
5.57
458.1
6.9
162.8
Table 14: Results of ablation experiment of MHA, GQA and SKA. For the Memory Footprint, we log only the cache usage, excluding the contribution of model weights.
Table 15: Results of ablation experi-
ments on gates.
B
Tr
Test PPL
Best Test PPL
Model
Head Type
Valid PPL.
Memory Footprint (M Bytes)
X
X
23.23
22.66
V
X
X
23.24
22.20
Transformer
Transfomme:
GQA
MIA
7.38
7.25
75.5
25.2
X
V
X
22.34
21.33
Transformer
X
22.14
21.15
Rodimus+
MHA
7.09
40.1
V
X
22.24
21.17
Radimus +
Rodimus+
GQA
SKA
7.15
7.19
24.4
14.9
V
21.90
20.84
Table 16: Results of ablation experiment on type of B.
Table 17: Results of ablation experiment models with or with- out two-hop residual.
8,
=?
Test PPL
Best Test PPL
THR
Test PPL
Best Test PPL
21.56
20.46
[T; 1-1111
3
24.36
23.32
21.77
20.37
1m
23.24
22.20
24
21.50
20.85
24
21.52
20.95
two hop residual, and represents number of layers.
E MORE EXPERIMENTAL RESULTS
E. THROUGHPUTS, FLOPS, AND MEMORY
In this section, we present the FLOPs and throughputs (words per second, wps) of Rodimus*, Mamba2, and the Transformer model, both during training and inference. Additionally, we provide an analysis of memory usage during inference, excluding model weights to ensure fair evaluation of space complexity. The specific configurations examined correspond to the 1.3B settings outlined in Table For the Transformer model, evaluate its attention mechanism implemented in both PyTorch (denoted as "torch") and Flash Attention 2 (denoted as "F.A."). The PyTorce implementation adheres to the theoretical time and space complexity of softmax attention, while Flash Attention significantly optimizes computation speed and reduces memory usage. Our analysis FLOPs and throughputs focuses scenarios with batch size of one, with results detailed in Tables 19 and 20
28

---

Published as a conference paper at ICLR 2025
Table 18: Results of ablation experiment on the ordering of Rodimus, SW-SKA, and FFN on WikiText- and Pile.
WikiText-105 PPL
Pile PPL la
Configuration
n_layenw6
Rodimus + X X Rodimus
21.56
5.9
21.95
5.9
The results reveal that the FLOPs and memory footprint of Rodimus are the smallest among all models tested, indicating the time and space efficiency of the proposed architecture. Particularly, Rodimus consumes only half the memory of Mamba2, making it attractive option for resource-constrained environments, such as edge devices or low-power systems. Conversely, Rodimus+ exhibits a slight increase in FLOPs but significant rise memory consumption due the integration of SW-SKA and FFN components, Overall, both Rodimus and Rodimus+ demonstrate greater time and space efficiency than the Transformer model, with Rodimus being more efficient than Mamba2 as well. Notably, Rodimus, Rodimus+, and Mamba2 can manage long sequences (e.g., pre-filling 24K tokens) more effectively than Transformers.
However, it is noteworthy that the throughputs of Mamba2 and the Transformer (F.A.) is significantly higher than that of Rodimus*. This discrepancy arises because Rodimus* does not incorporate the 1/O-aware optimizations that implemented in these baselines. In future work, we aim to address these throughputs limitations by integrating 1/O-aware optimizations to enhance Rodimus efficiency during both training and inference.
Table 19: Throughputs and FLOOSs during training.
4K-length contexts
Model
FLOPs x
Throughputs (wps)
Transformer (torch)
4.46
2726.64
Transformer (F.A.)
8547.29
Mamba2
3.65
10147.10
Rodimus
3.60
8592.59
Rodimus+
4.29
9410.38
Table 20: Throughputs, FLOPs, and Memory Footprint during inference.
4K-length contexts
24K 24K-knnghh contests
Model
Throughputs (wps)
Memory Foctpcint Bytes)
x110
Throughputs (wpo)
Memory Footprint Bytes)
Transformer (lurch)
3.6
2.4
805
7.7
00M
4831,8
Transformes (FA.)
49.4
11.0
Mamha2
3,0
59,3
0.8
Rodimus
29
21.4
0.4
2.9
21.4
0.4
E.2 GATES WITH TEMPERATURE Tz.
We further visualize the effect of gates influenced by the temperature gate T₁. Figure 7 we present the curves of the decay term exp(gi Tτ) and the input term g,¹¹ as the temperature varies. As 7i decreases, the lower bounds of both terms gradually rise, enhancing information retention.
In Figure compare the gates of layer in models from Table 15 with and without the temperature gate T, By comparing Figure Sa with Figure 8b, and Figure & with Figure 3d we observe that the inclusion of temperature increases the flexibility of the gates in selecting information.
29

---

Published as a conference paper at ICLR 2025
1.0
1
7:310
-2
(a) Decay term.
(b) Input term.
Figure 7: The curves of the gates with temperature gate Tç.
25
25
25
25
(a) Input w/o T.
(b) Input W T.
32
(c) Decay wlo
(d) Decay W
Figure 8: The visualization results of the gates are shown. The y-axis represents the n-th dimension of the gates, and the x-axis displays the sub-words after tokenization. (a) and (b) show the heatmap visualizations the input gate and g, (c) and (d) present the corresponding results of exp(g:) and exp(g: @ Tr).
E.3 MORE PRACTICAL LLM: MULTI-STAGE TRAINING RODIMUS+
To evaluate the practical capabilities of Rodimus+ and develop it into functional LLM, we develop -Coder, lightweight code generation model available in 1.6B and 4B parameter versions. Both variants have achieved SOTA performance in their respective size. The training process for both the 1.6B and 4B models begins from seratch and follows a three-stage pre-training approach:
Initial: Training on 6T tokens of primarily natural language data (20% code).
Coding: Continued training 1.5T tokens of code-focused data (60% code).
Annealing: Annealing with 500B tokens (more than 50% synthetic data).
30

---

Published as a conference paper at ICLR 2025
Table 21: Comprehensive evaluation results on the base models.
Qwen2.5- Coder-1.5B
Rodimus+- Coder-1.6B
Gemma2- 2B-PT
Qwen2.5- Coder-3B
Rodimus+- Coder-4B
Gemma3- 4B-PT
Qwen2.5-
Datasets
Coder-7B
Coding Tasks
Humanlival
41.5
51.2
19.5
51.8
60.4
36.0
60.4
HumanFval+
34.8
45.1
40.9
52.4
50.6
MВPP
57.2
51.2
31.0
46.0
70.0
MBPP+
66.1
62.2
65.9
71.4
70.1
BCBCOMPLETION
21.6
17.9
26.2
30.8
30.4
MultiPL-E
46.1
52.5
49.4
60.7
56.9
CRUXEval
38.5
45.1
44.6
56.4
56.8
Coding Avg.
43.7
46.5
56.7
General Täsks
C-EVAL
55.2
56.7
70.2
69.1
CMMLU
54.5
52.3
65.4
68.3
72.7
MMLU
55.5
51.1
52.2
63.3
62.6
59.6
70.5
BBH
21.8
46.8
42.4
32.5
61.9
50.9
67.3
General Avg.
46.8
51.7
56.6
65.8
69.9
Mathematics Tasks
GSM8K
60.4
68.7
25.0
78.5
38.4
83.4
23.7
29.0
31.9
37.0
24.2
42.2
Mathematics Avg.
41.9
48.9
20.7
52.0
57.8
31,3
62.8
Overall
44.4
48.4
51.7
59.6
61.6
The total pre-training dataset comprises about 8T tokens. The post training pipeline includes Supervised Tine-Tuning (SFT) and Direct Preference Optimization (DPO), utilizing the same dataset as for Ling- Coder-lite (Codefuse & Team 2025). The model weights have been open-sourced
We conduct a comprehensive evaluation of Rodimus---Coder models on 19 diverse benchmarks, including 13 code-related tasks (HumanEval (Chen al., 2021. HumanEval+ (Liu et al. (023)).
MBPP (Austin et al. 2021). MBPP+ (Liu et al. 2023). LiveCodeBench Jain et al. (2024). Big- CodeBench (Zhuo et al. 2024), MultiPL-E (Cassano et al.. 2022). MBXP-PLUS (Athiwaratkun
et al.
2022). CRUXEval (Gu et al.) (024), HumanEvalFix (Muennighoff et al. (223)) ,023),Spider (Yu
al. 2019)), ), general knowledge assessments (C-Eval (5-shot) (Huang et al. 2023), CMMLU
(5-shot) (Li et al. 2023). .MLUU (5-shot) (Hendrycks et al. 2021a). and BBH (3-shot) (Suzgun et al.
2023)), and mathematical reasoming challenges (GSM8K (4-shot) (Cobbe et al. (021) and MATH
(4-shot) (Hendrycks et al. 2021b)).
We systematically compare performance against various scales of contemporary models, including Qwen2.5-Coder (Hui al. 2024), Gemma (Team et al. 2024 2025). and Phi-4-Mini (Abouelenin al. 2025), to establish a comprehensive analysis of coding proficiency, general knowledge com- prehension, and mathematical problem-solving capabilities. To ensure direct comparability, re-evaluate Qwen series models using the same methodology as Rodimus+-Code Metrics for other models (Gemma, Phi-4-Mini) are taken from their original publications. All evaluation scripts are repository for transparency and reproducibility.
publicly available in the
E.3.1 PERFORMANCE ON BASE MODEI.
Overall assessment demonstrates that Rodimus+-Coode base models outperform similarly sized Qwen2.5-Coder and Gemma models across benchmarks. In particular, Rodimus+-Coder--t achieves a comparable average coding performance to Qwen2.5-Coder-77 although it has fewer parameters. Furthermore, even after extensive code-specific training. Rodimus+-Code maintains competitive performance in general natural language understanding and mathematical reasoning tasks, indicating effective preservation of its general capabilities alongside cooainn-speciio enhancement.
31

---

Published as a conference paper at ICLR 2025
Table 22: Comprehensive evaluation results on the Chat model.
Qwen2.5-Codes 1. 1.5.-mmtruct
Rodimus-- Coder-1.6B-CCl
GenamaZ- 2B-TT
Qwen25-Cadrr- 38-Instucct
Phi-4- Mini-38B
Rodiannc-- Cuder-BB-hal
Gemma3- 4B-IT
Q-ene--Cooler-
7B-Intruct
Coding Tasks
Humaskvvll
64.6
76.8
74.4
Foal
63.4
73.8
68.
MBPP
51.0
59.0
36.5
59.2
65.3
68.0
63.2
75.8
MBPP
33.0
66.4
4,0
1099
BC BCBNsmeETT
10.8
21.5
21.7
33.8
26.6
30.6
Humanival-Mau
50.8
57.3
67.4
76.1
52.4
MBXP-EN
55.8
75.5
76.0
87.3
87.7
MBXP CN
18.8
75.0
68.7
$3.5
28.6
55.0
EvalFis
18.9
52.6
55.5
68.8
69.3
Spider
61.2
71.
71.8
42.
73.5
$2.0
Coding Arg.
57.5
General Tooks
C-EVAL
51.5
50.8
CMMIUU
50.5
MMLU
$2.0
49.3
56.1
61.7
67.3
$7.5
58.1
66.1
BUH
24.2
58.7
57.3
70.4
Avg.
43.2
52.3
Muthommtticz Tucka
GSMSK
54.4
68.5
62.6
13.5
88.6
79.2
89.2
MAIH
38.1
33.5
27.2
75.6
Mathematics Avg.
46.2
$1.0
44.9
58.8
68.8
61.7
82.4
70.1
Overail
44.2
558
58.9
64.3
68.4
E.3.2 PERFORMANCE ON CIAT OR INSTRUCT MODELS.
Overall assessment reveals that A-Cider---CCor---e models maintain superior performance com- pared to models of similar size. Remarkably, Rodimust-Coder- Gooder-l.6B-. 1.6B-Chat and Codimus--Coddr-- Chat achieve coding performance comparable to their larger counterparts, -wwen..--Coder-3BB-B- and Swee22.5-Coder. -a-eer--BBnnaaae IIIItrutet respectively, demonstrating efficient parameter utilization while delivering competitive performance with larger instruction-tumi models.
F DISCUSSION ON WIDESPREAD USAGE OF RECURRENT AND HYBRID MODELS
The technical challenges associated with our proposed models, particularly in comparison to the widespread adoption of softmax attention-based large language models (LLMs), are well-recogmized. However, the evolving landscape of machine learning research suggests a shift in this dynamic. The growing interest in recurrent and hybrid models within both academic and industrial domains is beginning to address these barriers, paving the way for broader exploration and application.
Recent developments illustrate this trend, as numerous new recurrent and hybrid models have emerged. Notable examples include Mamba, Mamba2, RWKV, RWKV6, RetNet, HGRN, HGRN2, GLA, Hawk, and Griffin. Moreover, Mistral has successfully scaled Mamba to 7B parameters and has made it open- open-source. Similarly, Google has scaled Griffin to an impressive scale of 9B and open- sourced Ceturrent-Gemman These advancements indicate growing confidence in the performance of recurrent and hybrid models, often achieving comparable results compared to traditional softmax attention-based L.LMs while maintaining lower complexity.
From an inference perspective, it is indeed worth highlighting that recurrent and hybrid models are inherently designed for efficiency, The per-token generation complexity for these models is theoret- ically O(1), in stark contrast to the O(T') complexity of softmax-based models, where T denotes the context length. As a result, even without optimization, recurrent and hybrid models demonstrate rapid inference capabilities. Furthermore, state-of- of-the-art inference optimization frameworks such as vLLM, TeMS,rRT-LLM, and OLlama have already begun to support the inference of models like Mamba and Mamba2, further paving the way for their more widespread usage.
32

---

Published as a conference paper at ICLR 2025
G
FUTURE WORK
Rodimus* R has delivered excellent results across benchmarks; however, there are areas for improve- ment. Due limited computing resources, we have not expanded its parameters match open-source models like SWKV6-14B and Qwen2-72B. Additionally, Rodimus* lacks the highly 1/O-aware opti- mization found in models such as Mamba and Mamba2. There's 's potential enhance its performance by designing aware multi-head scalar decay and integrating it with Rodimus's 's DDTS to broaden gating mechanisms without significantly impacting training efficiency. Furthermore, for Rodimus+, the memory usage of SW-SKA can be further reduced while achieving better practical application performance. We aim address these issues in future work.
33

---

Published as a conference paper at ICLR 2025
Sam
Figure 9: Mamba-ll 1.4b
Figure 10: Mamba2-1.35
34

---

Published as a conference paper at ICLR 2025
Redimus-3.40
Aedircs-3.
Srgh: Nitin SA
Langl
Figure 11: Rodimus-1.4B
Pythis
Pythee3448
Seat Dept
Aestige Sep
Figure 12: Pythia- 1.4B
35

---

Published as a conference paper at ICLR 2025
Aodimaa--1.. 60
Rodimas 33.00
Trgh: Niria Cotr 20
Dusrall Scu 31.75
Longh
Langl
Figure 13: Rodimus Rodimus+-LL6B- --B.B-2K
Redimasn 1-1 68
Seotig Dept
Man Langt
BMC
Figure 14: Rodimus--I.6B-44
36


Reviews:
[{"review_id": "1lDZOHOWhe", "reviewer": "Reviewer_Cz2E", "summary": "This paper introduces two efficient Transformer architectures that incorporate semantic, token, and head compression: Rodimus, which is purely based on linear recurrence and, Rodimus+, which combines linear recurrence with sliding window attention to improve the quality and efficiency tradeoff in LLMs. The key components proposed include a data-dependent tempered selection mechanism (DDTS) and sliding window shared-key attention (SW-SKA). Experiments with models between 130M to 1.6B on language modeling, various downstream tasks, and retrieval tasks show that Rodimus+ achieves better performance than models based on recurrence such as Mamba and based on standard attention such as Pythia.", "questions": "1. Was there any specific reason for not using the same number of parameters for Rodimus and Mamba models? For instance, in the third block of Table 2 Mamba has 0.37B parameters while Rodimus has 0.46+B parameters and in the fourth block Mamba has 1.3/1.4 while Rodimus 1.4/1.6. Such differences make it difficult to understand if the improvement comes from the proposed gating design or from pure increase in capacity. To strengthen the claims made in the paper, I'd suggest to the authors to provide a head-to-head comparison with that model. \n2. There is lack of comparison with hybrid models, even though Rodimus+ is in that category of models. Have you considered comparing with such models? A few examples include: Jamba, Griffin, Recurrent-Gemma. \n3. What is the memory consumption and training/inference times for the models compared in the experiments? Even though the efficiency aspect is emphasized a lot in the beginning of the paper, it is under explored in the later sections. \n4. Can the authors report language modeling scores with the larger model sizes? The ones reported on Wikitext-103 are only with very small models and the conclusions that we can draw from it are limited.\n5. In Figures 1a and 1b, what is the performance of a standard Transformer? \n6. What is the memory benefit of SKA compared to alternative attention mechanisms such as MQA and GQA? It would be useful to include an ablation study that highlights their differences.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper introduces two new techniques, data-dependent tempered selection (DDTS) and shared-key attention (SKA). DDTS is a new gating-based recurrent type for state-space models which enhances performance and increases parameter efficiency through reducing hidden state size. SKA shares a single key representation to reduce memory footprint while maintaining the multi-value setting of the original multi-head attention. \n2. These techniques are evaluated extensively in general tasks (content analysis, commonsense reasoning, reading comprehension, etc), long-context retrieval tasks (NeedleBench) and show improved performance compared to various efficient and standard transformer models.\n3. Overall, the paper is satisfactorily written, the proposed techniques are clearly described, and the prior work moderately well covered. The experiment section is thorough and has an organization that is easy to follow.", "weaknesses": "1. While the main focus of the paper is to improve the trade-off in terms of accuracy and efficiency, the experiments do not discuss the efficiency aspects such as memory cost and latency or times for  training/inference. \n2. Even though the specific gates proposed haven't been used in prior work, the novelty in the design can be considered rather incremental as the formulation is very similar to previous ones (Mamba, S4D).\n3. Direct comparisons with alternative gating functions in SSMs using the exact same architecture and model size are lacking. This prevents the comprehensive understanding of the benefits brought by the specific gating techniques.\n4. The evidence provided in this paper to support the claim of \"breaking\" the accuracy-efficiency tradeoff is not comprehensive. \n   - Several dimensions are lacking in the experimental setup: comparison with a wide range of efficient attention models (sparse, linear, hybrid), head-to-head comparisons (equal number of params, HPO budget), and rigorous hyperparameter optimization. \n   - In addition, the performance of Rodimus compared to Mamba when using the same number of parameters is actually the same or worse. \n   - Only when combining with SWA and increasing number of parameters performance improves further with Rodimus+ but this is not a valuable comparison since Mamba2 is not a hybrid method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "h3MtxIwXrI", "reviewer": "Reviewer_kK5a", "summary": "Attention is the most memory-costly operation in transformer architecture. It scales quadratically in sequence length, as it can correlates each position to each other. To overcome this bottleneck, state space models were introduced, however usually degenerated in downstream performance.\nThe authors show that their particular gating mechanics of state space models \"Rodimus/DDTS\" improve downstream performance as they may model language more naturally. It in particular comprises a low-rank auto-encoder approximation and a softmax based temperature parameter.\nFurthermore the authors introduce a shared key attention \"SKA\", that reduces some memory footprint and can be shown to have the same expressiveness to standard attention. Both methods are \"Rodimus+\".\nLots of efforts was done with evaluations, albeit, i'm afraid, not in a comparable fashion..", "questions": "1) line 120 : \"matrices\"\n2) Tab 1 only ~2M parameters increase to Rodimus+ is the SKA + FFM? other hyperparameters are the same? the changes appear to be pretty random in tab 2", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "the paper is generally well written and presented. ideas appear thought-through and are (somewhat) mathematically reasoned.\nmore reasonable interplay of gating mechanics / ssm's in general and attention is pretty important to iterate towards more robust language models.", "weaknesses": "1) i understood the change to Rodimus+ is the two norms, SW-SKA and FFN per layer. how can this sum up to those negligible changes in parameters in Tab 2 - or do you use entirely different architecture configs?\n2) i do not find anything conclusive about runtime estimates (only a vague big-O estimation and some chunking approach for training). This should not be optional!\n3) my major concern is that only Fig 1 / Tab 1 / Fig 4 appear to be \"somewhat fair\" comparisons but in a ridiculous low scale. i cannot follow steps in section 4 most of the time. in particular i could not reconstruct the experiment design due to missing descriptions partially found in appendix.\n4) table 2 appears pretty delusive. variances in the dataset and tokenizer alone are able to cause significant variations already, which is not discussed at all and not negligible! the only comparable lines appear to be the first 3. the only thing i can derive here is that rodimus can  model language. a sentence like line 478 is pretty off standard to me. you pretend to have a better method so do a fair 1:1 training/ scaling analysis (on less models). tab 14 should be sufficient for such a case study, but is also pretty uncomparable?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NlbHNkfR5J", "reviewer": "Reviewer_Az8Q", "summary": "The paper introduces Rodimus and Rodimus+, two efficient alternatives to traditional Transformer models that aim to balance accuracy and computational efficiency. Rodimus uses a Data-Dependent Tempered Selection (DDTS) mechanism which is an improved parameterization of gated linear attention. Rodimus+ extends this with Sliding Window Shared-Key Attention (SW-SKA) for enhanced local token focus, further optimizing efficiency without sacrificing performance. Experiments demonstrate these models outperform comparable models on language modeling and recall tasks, making Rodimus+ a scalable, efficient alternative for large language models.", "questions": "- How significant is $ d_t $ in Eq. (11)?\n\n- Have you investigated the ordering of Rodimus, SW-SKA, and FFN in Eq. (13)?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This work comprehensively examines the design space of gated linear attention (GLA), particularly focusing on an outer-product-based gating structure with input and forget gates. It addresses a gap in systematic studies on GLA's gating mechanisms, positioning itself well to bridge this gap. The proposed Data-Dependent Tempered Selection (DDTS) mechanism is both intuitively appealing and practically useful, as demonstrated by experimental results. Additionally, the definition of the selection mechanism in this context is compelling.\n\n- The ablation study on architectural design is relatively thorough.", "weaknesses": "- The work appears somewhat incremental, as it largely builds upon existing concepts. Technical contribution is limited. \n\n- Regarding language modeling experiments, the choice of the WikiText-103 dataset for language modeling may not be ideal, as it is relatively simplistic and sensitive to hyperparameter tuning. For Figure 1a, larger-scale experiments using more extensive datasets would provide a more robust evaluation. Additionally, the controlled experiment scale in Table 2 is too limited; ideally, models with around 1B parameters trained on 100B tokens should be tested.\n\n- Regarding recall-intensive eval,  the NeedleBench benchmark may be overly challenging for smaller-scale models, as evidenced by Figure 5, where all models perform poorly. Consider testing on simpler recall-intensive benchmarks (e.g., FDA, SWDE, NQ, SQuAD, TriviaQA, DROP as used in [1, Table 1]) for more interpretable results. Additionally, it would be insightful to know if Rodimus can outperform Mamba2 in recall performance with half the state size under equivalent training conditions. Results on these tasks could be easily added using the lm-eval-harness during the rebuttal period.\n\n- Training and inference throughputs would be reported.\n\n- Several minor errors are present throughout. See the list of specific corrections below.\n\n---\n\n### Typos and Minor Corrections:\n\n- **Table 1:** Replace “suboptimal” with “second-best.”\n- **Line 279:** Incorrect citation for GLA.\n- **Table 4:** “State size d” should consistently be represented as $d^2 $.\n- **Table 4:** For Rodimus,  $B_t$ should likely be $ (g_t^{\\tau_t})^T \\hat{\\beta}_t $.\n- **Line 1280:** The symbol $\\beta$ in the first line should be $\\hat{\\beta}_t $.\n\n---\n\n### Reference:\n\n[1] Just Read Twice: Closing the Recall Gap for Recurrent Language Models", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "G9eyS96boz", "reviewer": "Reviewer_ZtZu", "summary": "The paper introduces Rodimus, a linear attention based model, purely recurrent, and and its enhanced variant Rodimus+, a hybrid model, which combines Rodimus with SW-SKA. The paper aims at reducing the computational inefficiencies of traditional softmax attention mechanisms in large language models (LLMs). Rodimus uses a linear attention approach combined with a data dependent tempered selection mechanism to achieve O(T) complexity for per-token generation. This significantly reduces memory usage while maintaining performance (a key limitation of recurrent models). Rodimus+ extends this by integrating a sliding window shared key attention (SW-SKA), which combines the strengths of semantic, token, and head compression to enhance efficiency and accuracy.\n\nThe paper is in the broader context of linear attention mechanisms such as those introduced by Katharopoulos et al. (2020), which reduce computational complexity through kernel-based approximations. Plus, it builds on the concept of sparse attention, similar to the approaches used in Longformer (Beltagy et al., 2020), but enhances it by integrating local and global context understanding. Rodimus+ also leverages ideas from Grouped-Query Attention (GQA) (Ainslie et al., 2023) by implementing shared-key attention, compressing the memory footprint while preserving performance. \n\nThe experimental results show that Rodimus+ achieves better downstream performance even against models trained on more tokens.", "questions": "How does the semantic compression affect the fine-tuning performance on specific downstream tasks that require retaining nuanced context?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The proposed approach seems to achieve interesting results in complexity vs performances: \n- Rodimus+ manages to achieve high performance even when trained on fewer tokens compared to other models. \n- The authors provide extensive validation across multiple benchmarks, showing that Rodimus+ outperforms existing models. \n- The introduction of the data dependent tempered selection mechanism shows an innovative way to manage the flow of information in recurrent architectures. \nThe code will be available after the review process.", "weaknesses": "- While the innovations presented are technically very interesting, they also introduce a level of complexity that could be a barrier for widespread usage. \n- Although Rodimus improves on the traditional recurrent models, it may still inherit some of the limitations associated with recurrent architectures (like the sequential processing bottlenecks).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)