Scaling Linear Attention Capacity with Sparse State Expansion
Authors: Yuqi Pan
Yongqi An
Zheng Li
Yuhong Chou
Rui-Jie Zhu
Xiaohui Wang
Mingxuan Wang
Jinqiao Wang
Guoqi Li
Venue: ICLR 2026 Poster (2026)
Keywords: Linear Attention
Language Model
Primary area: foundation or frontier models, including LLMs
TL;DR: SSE is a partitioned state expansion method under a row-sparse update framework for linear attention, improving retrieval and reasoning.

Abstract: The Transformer architecture, despite its widespread success, struggles with long-context scenarios due to quadratic computation and linear memory growth. While various linear attention variants mitigate these efficiency constraints by compressing context into fixed-size states, they often degrade performance in tasks such as in-context retrieval and reasoning. To address this limitation and achieve more effective context compression, we propose two key innovations. First, we introduce a row-sparse update formulation for linear attention by conceptualizing state updating as information categorization. This enables sparse state updates via softmax-based top-$k$ row selection, thereby extending receptive fields and reducing information interference. Second, we present Sparse State Expansion (SSE) within the sparse framework, which expands the contextual state into multiple partitions, effectively decoupling parameter size from state capacity while maintaining the sparse row-selection paradigm. Supported by efficient parallelized implementations, our design achieves highly discriminative state representations. We extensively validate SSE in both pure linear and hybrid (SSE-H) architectures across language modeling, in-context retrieval, and mathematical reasoning benchmarks. SSE demonstrates strong retrieval performance and scales favorably with state size. Moreover, after reinforcement learning (RL) training, our 2B SSE-H model achieves state-of-the-art mathematical reasoning performance among small reasoning models, scoring 64.5 on AIME24 and 50.2 on AIME25, significantly outperforming similarly sized open-source Transformers. These results highlight SSE as a promising and efficient architecture for long-context modeling.

Published as a conference paper at ICLR 2026
SCALING LINEAR ATTENTION CAPACITY WITH SPARSE STATE EXPANSION
1.21 Zheng Li Yuhong Chou², Ruijie Zhu", Xiaohui Wang', Mingxuan Wang', Jinqiao Wang", Guoqi Li 'ByteDance Seed Institute of Automation, Chinese Academy of Sciences The Hong Kong Polytechnic University 'UC Santa Cruz lizheng m@byt edance. com.
Yuqi Pan Pan!-2' 1.21 Yongqi An!.?'
guoqi 110 aee.aaacccom
ABSTRACT
The Transformer architecture, despite its widespread success, struggles with long- context scenarios due to quadratic computation and linear memory growth. While various linear attention variants mitigate these efficiency constraints by compress- ing context into fixed-size states, they often degrade performance in tasks such in-context retrieval and reasoning. To address this limitation and achieve more effective context compression, we propose two key innovations. First, we intro- duce row-sparse update formulation for linear attention by conceptualizing state updating as information categorization. This enables sparse state updates via softmax-bate -baed top-k row selection, thereby extending receptive fields and reduc- ing information interference. Second, we present Sparse State Expansion (SSE) within the sparse framework, which expands the contextual state into multiple partitions, effectively decoupling parameter size from state capacity while main- taining the sparse row-selection paradigm. Supported by efficient parallelized implementations, our design achieves highly discriminative state representations. We extensively validate SSE both pure linear and hybrid (SSE-H) architec- tures across language modeling, in-context retrieval, and mathematical reasoning benchmarks. SSE demonstrates strong retrieval performance and scales favorably with state size. Moreover, after reinforcement learning (RL) training, our 2B SSE- H model achieves state-of-the-art mathematical reasoning performance among small reasoning models, scoring 64.5 AIME24 and 50.2 AIME25, signif- icantly outperforming similarly sized open-source Transformers. These results highlight SSE as a promising and efficient architecture for long-context modeling.
INTRODUCTION
The Transformer architecture Vaswann al., 2017) has achieved remarkable success in various sequence modeling tasks, leveraging its attention mechanism for expressive modeling and high par- allelism. However, it inherently struggles with processing long sequences due quadratic compu- tational complexity and linearly growing memory usage of key-value (KV) caches during inference.
To address these limitations, numerous linear attention variants have been proposed (Katharopoulos) ett al. 2020 Yang et al., 2024h [ao & Gu 2024 Yang et al. (022a)). Such approaches typically achieve sub-quadratic complexity and enable partial parallelism in training, along with RNN-like constant memory usage during decoding. However, most methods compress contextual information into fixed-size state matrices (e.g., 128 d), which often compromises performance on tasks such as in-context retrieval and reasoning (Arora et al., 22044a b Jelassi et al. 2024 Akyürek et al. 2024). This raises a key question: how can context compression be designed to balance modeling fidelity and computational efficiency in long-context scenarios?
Motivated by the pursuit of more effective context compression, we propose two key innovations. First, we introduce a row-sparse update formulation for linear attention by conceptualizing state up- dating as an information categorization process. Specifically, we treat the rows of contextual states as distinct latent categories, interpret the key vector mapping k, f(x:, Wx) category as- signment function, and employ a softmax based top-k mechanism to select rows. At each step, this formulation selectively updates only state rows corresponding to predicted categories. Both theo- retical analysis and synthetic experiments validate that our approach achieves larger receptive fields and reduces inter-category interference. Second, to mitigate the memory capacity constraints caused

---

Published as a conference paper at ICLR 2026
by limited state sizes, extend the contextual state within the row-sparse update framework and propose Sparse State Expansion (SSE). SSE expands the state into partitions with shared atten- tion parameters and utilizes write-read gate for partition selection, followed by softmax-based selection of state rows within partitions. This design effectively decouples parameter size from state capacity while maintaining the sparse row-selection paradigm, and empirical results indicate that the states of SSE achieve lower inter-row similarity and higher singular value entropy, which enables more discriminative state representations. We also propose efficient parallelized implementations of SSE, employing masking and the varlen technique optimized for various training contexts.
We extensively validate SSE under both linear and hybrid architectures (denoted SSE-H), across three core capabilities: language modeling. in-context retrieval, and mathematical reasoning. SSE demonstrates strong language modeling performance among advanced linear models (Yang et al. 2024b|a a Du et al. 2025), while SSE-H outperforms both hybrid and Transformer baselines. On long-context retrieval tasks, SSE consistently improves upon other linear attention models, and its hybrid variant significantly narrows the gap with softmax attention Transformers. For reasoning, 2B SSE-H model achieves scores of 64.5 on AIME24 and 50.2 on AIME25 with advanced RL training (Shao al. 2024 Yu et al., 2025)- exceeding the best reported results from similarly sized open-source Transformers. Moreover, SSE exhibits strong scalability with respect to state capacity, and supports efficient conversion from pretrained Transformers different scales, demonstrating its flexibility across model architectures and training regimes.
Our main contributions are follows:
We introduce a row-sparse state update framework that conceptualizes state updating as information categorization. This framework enables sparse state updates through softmax- based top-k row selection, supported by both theoretical and empirical analysis.
Within the proposed framework, we introduce Sparse State Expansion (SSE), an efficient state expansion mechanism designed to effectively manage parameter count and preserve the sparse row-selection paradigm. Furthermore, we develop parallelized implementations of SSE, suitable for various training contexts.
We extensively validate SSE and the hybrid SSE-H across diverse training stages and bench- marks. Our 2B hybrid model achieves state-of-the-art reasoning performance among small reasoning models, scoring 64.5 on AIME24 and 50.2 on AIME25 exceeding previous open- source softmax Transformers of comparable size.
2 PRELIMINARY
Autoregressive language models predict the next token by modeling the conditional probability Px xxx). At inference, historical tokens Xxr are encoded into a compact contextual state S, through different token-mixing mechanisms, facilitating sampling via (x, 2 Px(x)(xc1).
Softmax Attention. Softmax attention (Vaswani et al., 2017) stores the key-value vectors derived from historical tokens in a KV cache (K, E Rixa V, E Rtxd). As a new token X, arrives, the corresponding vectors k, V, are appended to this cache. Attention is then computed by evaluating the interaction between the current token's query q and the KV cache:
S. (K,, Vih.
K,-1
V
K,
V
k
O, softmax(q.KK) V
This append-based mechanism for state representation naturally enables high parallelism. However. as historical information is neither compressed nor pruned. incurs efficiency bottlenecks for long sequences: memory and computation grow linearly with sequence length per step, and quadratically when all queries are processed in parallel during training.
Vanilla Linear Attention. Vanilla linear attention (Katharopoulos et al. 2020) compresses histor- ical context into a fixed-size state matrix (S, E Roxa) using an outer proouet update, where is predefined constant (e.g., 128) independent of the sequence length:
= S Vt, o, S..
(4)
2

---

Published as a conference paper at ICLR 2026
For long contexts, this compression achieves significantly higher inference efficiency, requiring only constant memory and computational cost per time step. However, when с < t, the lossy compres- sion causes performance gaps compared to softmax attention, particularly in tasks such in-context retrieval and reasoning. By reformulating the linear recurrence into chunk-wise parallel form, linear attention enables hardware-efficient and sub-quadratic training (Sun et al. 2023 Qin al. 2023b
Yang et al. 2024b; Dao & Gu 2024).
Thus, effective contextual compression has become essential for long-sequence processing, yet ex- isting linear attention mechanisms have not fully exploited this potential. To address this, we in- troduce two improvements: a novel row-sparse update formulation for linear attention by concep- tualizing state updating as information categorization, and Sparse State Expansion (SSE), which efficiently augments state capacity within the row-sparse update framework.
3 ROW-SPARSE STATE UPDATE: AN INFORMATION CATEGORIZATION PERSPECTIVE
Understanding the operational mechanism of contextual states is crucial for developing effective compression methods. To this end, we adopt recurrent outer-product update formulation that provides a unified view of linear attention modeling (Chou et al. 2024):
S: = A,S,-1 +((k)
(5)
Here, A, the state transition matrix, and Φ the feature map embedding k, into finite- dimensional space. This framework highlights two key design directions:
(1) The role of A, in historical information management. Many modern linear attention vari- ants (Sun et al. 2023 Qin et al., 2022; Peng et al. 2023) introduce recency bias via exponential decay or input-dependent gating YYang et al. 2024b Gu Dao 2023; Dao Gul 2024). while more recent works introduce structured updates such diagonal plus-low low--oo- (DPLR) matrices
derived from the delta-rule Yangg et al. 2024c Siems et al. 2025).
(2) The role of o(k)) in new information processing. Early designs sought to approximate the exponential kernel of softmax attention (Katharopoulos et al. 2020 Choromanski et al. 2020) Peng et al. 2021 Zhang al., 2024a), whereas recent variants often rely simple mappings such as identity, ReL.U, or SILU (Snn et al. 2023] Yang et al. 2024b) Chou et Yang el. 2024c a)
While structured state transitions have seen steady progress, the role of the key feature map in ef- fective state updating remains less understood. In this section, we focus on this second direction, proposing a row-sparse update formulation that conceptualizes state updating information cate- gorization. This is implemented via a top-k-then-solee soohen--ofmaxx strategy for row-selection.
3.1 TREATING STATE ROWS LATENT CATEGORIES
In Equation 5 we treat state rows distinct latent categories and conceptualize k; f(x,, Wx) as a category assignment functior where larger values assign more information to the corresponding state rows. Consequently, each state row represents distinct feature subspace, storing information that shares similar categorization outcomes. For a linear categorization function (k, = x, W&). we theoretically demonstrate that information assigned to the same state row exhibits similar features, as reflected by large inner products under this linear mapping:
Proposition 1 For imputs X, satisfying ||x.122 {x, x, > k where ki, > O is category-specific threshold. For any inputs x,,, E C'i, the
and given the categorization rule O
s, E {l,..., t}, and following inequality holds: x d cos 20', where 0² = arccos( Vallall
vice versa.
See Proof in Appendix D for the detailed proof. To validate this, we analyze the information as- signments of linear attention using top-1 row-selection (Figure 99. and the information composition exhibits a clustering pattern, supporting our proposed categorization perspective. However, vanilla linear attention often suffers from inter-category mixing without this explicit top-l: hard assignment.
For clarity, k, denotes the vector after the mapping throughout the remainder of this paper.
3

---

Published as a conference paper at ICLR 2026
Linear Function
SILU Function
(c) Softmax Function
(4) SSE Function
Figure I: Row-wise cosine similarity of contextual states in linear attention models with dif- ferent categorization functions. Cosine similarity is computed over the final states of multiple long-context inputs. Each plot shows a 128 x 128 similarity matrix across state rows; darker blue indicates lower similarity, reflecting more effective latent categorization.
This perspective allows us to reinterpret various designs of f(x, Wx). Early models (Sun et al, 2023 Yang al. 2024b) used single linear transformation (f(x; W&) W&). Later works incorporated more expressive categorization functions, such as nonlinear activations (Yang et al.,
2024c a) gating mechanisms (Dao Gu 2024: Beck al. (024), and MLPs (Zhang et al. 2024a Kasai et al. 2021). Beyond these examples, viewing state construction as information categorization also motivates novel architectural designs, including the softmax based hard row oow-selection presented in the next section.
We evaluate different categorization strategies by analyzing row-wise cosine similarity of contex- tual states across models. As Figure (b)(b)) shows, results highlight both the benefits of more ex- pressive categorization functions and the limitations of existing linear attention- state rows exhibit substantial homogenization. We posit that this limitation arise from implicitly involving the notion of information categorization without fully exploiting category assignments for compressed storage.
3.2 ROW-SPARSE STATE UPDATE VIA TOP-K-THEN-SOTMAX
In contrast to existing linear attention mechanisms, we explicitly leverage category assignments derive a row-sparse update formulation for linear attention. From the perspective of latent categorization, we combine the top-k hard assignment strategy with softmax categorization function:
k. = softmax((op- &(x, Wt)), S. = ;==AAS,,- + v..
(6)
We apply the softmax only to the top-k entries of x, WW. setting all remaining values to zero Moreover, in gated linear attention variants such as GL.A, state transition matrix A, incorporates gating factor ar E (0, 1) (see Appendix C for details). In these cases, eliminate the gates of non-selected rows using the same top-k index set, yielding the modified transition matrix A,
Using the linear categorization function as an illustration, we theoretically identify two fundamental limitations of existing linear attention designs together with the corresponding advantages of hard top-k row selection:
Updating all state rows introduces cross-category noise into memory, resulting in state row homogenization and reduced querying discrimination. In contrast, row-sparse updates encourage each state row store similar information, yielding higher inter-row discriminability and more precise information organization.
(1) Information interference (Propositions 2-3
Proposition 2
Under mild regularity conditions, the row-wise cosine similarity in vanilla linear St S... is greater than that obtained under row-sparse updates, Cos S S.).
allention, Cos
Proposition 3 For arbitrary queries Pr O, and states St with strictly bounded row norms, define the lower bound of the state-distiguuisaabitty measure as minp..q p. Cos(qiS,, p.SS) Row- sparse updates reduce the lower bound compared vanilla linear attention.
(2) Receptive fields (Propasition 4): In gated variants of linear attention, uniform decay driven by strong recency bias restricts the effective receptive field. Row-sparse updates mitigate unnecessary decays, extending the receptive field over longer contexts and improving retrieval performance.
Proposition 4 Let q: and k, be all-ones vectors with 2 s >>>0 decay factors az (0, 1), the receptive field is upper-bounded: the attention score pts falls below
O. For gated linear attention with
4

---

Published as a conference paper at ICLR 2026
new information
Vanilla Linear Attention
S EReaa
attention
contextual state
W
projection weights
(x, W.
02
selected rows
selected rows
Sparse State Expansion (SSE)
R(Nc)ool
softmax
selected
selection
top-k selection
fx, Wa)
(for partition))
N partitons with shared parametors
Classification function: assign information to specific state rows.
write-read gate
write-read gate
Figure 2: Comparison between vanilla linear attention and SSE. SSE expands the state into N partitions within the row-sparse update framework, where a categorization function assigns infor- mation to specific state rows. All partitions share attention parameters. Sparse selection follows two steps: (I) top-k partition selection based on a write-read gate (blue indicates selected partitions; green marks always-selected partition for training stability), and (2) selection within the cho- sen partitions via softmax over key vectors.
given threshold Puh- In contrast, row-sparse update formulation can retain important information across arbilrary spans, ensuring the attention score remains above Pth-
Detailed proofs provided in Appendix D (Proofs
Synthetic MQAR (Arora al. 2024a)
experiments highlight the benefits of top-k row-sparse updates for information recall (see Figure 7 Moreover, Figure (e) illustrates, the low row-wise similarity exhibited by the states produced by the softmax-based row selector provides empirical support for the effectiveness of our framework.
To improve stability and mitigate category imbalance, encourage a more uniform distribution of samples across state rows via an auxiliary loss (Fedus et al. [2022] Lepikhin et al. (020): Loatance af Ei= f kj. where a is a manually specified coefficient, the number of state rows, and f³ the selection frequency of the i-th row. This loss complements the sparsity induced by top-k row- selection and is optimized jointly with the main objective.
SSE: STATE EXPANSION UNDER THE ROW-SPARSE FRAMEWORK
The most critical gap between linear and softmax attention lies in long-contex. retrieval and rea- soning, central focus of recent research. This primarily stems from the limited memory capacity caused by small state size. For instance, with c = 128, the memory budget of linear attention equals that of softmax attention with context window of only 64 tokens. Therefore, increasing the contex- tual state size linear models key to addressing their current shortcomings. Recent works explore this issue from various perspectives (Arora et 2024b Du al. 2025 Guo al., 2025b Zhang et al. 2025). Building on this foundation, we pose the following question: How can state expansion be designed from the perspectives of information categorization and sparse row-selection?
4.1 SPARSE STATE EXPANSION FOR LINEAR ATTENTION
Within the row-sparse update framework, we expand the state rows by factor of N, yielding N x rows grouped into N partitions. For each token X,, a top-k-then-softmal mechanism selects rows from the chosen partitions. Furthermore, isolate the effect of state expansion, we fix the parameter size. This follows from the observation that softmax attention, with 4d parameters, can effectively manage unbounded states while maintaining strong in-context learning. Thus, in linear attention, the primary bottleneck lies in state size rather than parameter count.
Our design therefore aims scale state capacity without increasing parameters while preserving sparse row selection. For efficiency, we avoid top-k over all rows (Equation (). instead applying hard top-k selection partitions followed by soft row selection within the chosen partitions using
5

---

Published as a conference paper at ICLR 2026
softmax. With parameter sharing across partitions, this procedure is equivalent to applying softmax independently within each partition. Our core design consists of the following components:
Shared attention parameters across partitions. We use single set of projection weights (QKV) across all partitions. This design choice, together with sparse partition selection, is equivalent to per- forming re-segmentation of input information along the sequence dimension, where each segment corresponds to a different partition, thereby enabling sparse token-level interactions. Additionally, parameter sharing mitigates training instability, such as overly sparse parameter updates.
Write-read gate for hard partition selection, followed by soft row selection via softmax. A gate vector e, R.N is first used select the top-k partitions, and within these partitions softmax(k; further selects the state rows. Only the selected partitions are updated, and ablation results show that the softmax step yields over 40% row sparsity. Furthermore, the gate e, applied to both state input (KV) and output (Q), thereby governing both information writing and reading.
We refer to the linear attention variant that integrates the two core components described above as Sparse State Expansion (SSE), illustrated in Figure 2
The complete procedure is as follows:
e, softmax X, W.),
{i [ E top- k(e:)),
q: X₁ W. 9) Vi = x w., k. = softmax(x: Wx),
(8)
S: AnSi t klv,, for ET T for i & T
(9)
S
0. e q.S;
(10)
1ET
(I) We compute gate vector e, E R.V via linear projection and softmax, and use it for top-k par- tition selection (Equation 77 (2) Next, we apply softmax across all rows of the selected partitions.
Since partitions share the projection W., this is equivalent to applying softmax within each parti- tion, with the scaling factor absorbed into parameters (Equation 8) (3) To ensure the gate remains trainable, e, is applied to both state input (KV) and output (Q) (Equations IT)), thereby controlling both writing and reading and enabling more targeted optimization. (4) For partition i, the gated key
ki- = ei softmax (x, Wx) (Equation 9), which ensures normalization Σ, Ea())a 1. Only
the selected partitions Si are updated. (5) Within the write-read gate framework, we restrict state reading to the top-k partitions using the same gate e, (Equation I0).
Local interactions between consecutive tokens serve an effective prior for language modeling (Fu et al. 2022 Peng et al. 2023: Arora et al. 2024b), which we leverage to stabilize the sparse com- ponent's training. To this end, we incorporate an always-selected partition and apply LoRA QK projections, keeping the parameter count nearly constant. We further adopt partition-based aux- iliary loss in place of a row-based one. Leveraging parameter sharing and the sparsity induced by hard top-k selection, SSE enables scaling up the state size of linear attention while keeping both computational and parameter overhead nearly constant.
Empirical analysis reveals that SSE's state partitions exhibit low inter-partition similarity; further- more, state rows within each partition remain diverse, as illustrated in FigureTTd), Overall, the state exhibits diagonal pattern, indicating discriminative state representations. Furthermore, Figure shows, the singular value entropy of the SSE states (definition and discussion in Appendix F.T is higher than that of GLA (Yang et al. 2024b). This indicates more diverse and less compress- ible state composition, suggesting a more efficient utilization of state capacity in SSE. In addition, analyze the receptive field of SSE and observe that it is significantly larger than that of GLA, indicating improved long-range information access (see Appendix F.I for details).
4.2 EFFICIENT IMPLEMENTATIONS OF SSE
SSE expands the state into multiple partitions. During operator execution, we prioritize maintaining parallelism across partitions, avoiding sequential computation, which requires invoking the lincar- attention kernel separately for each partition in for-loop. Concurrently, we leverage sparsity to minimize unnecessary computational overhead. This section briefly outlines three practical imple- mentation strategies, with further details and pseudocode provided in Appendix [E]
Naive implementation via masking, For short and variable-length sequences, partitions are han- dled by replication and masking, enabling the use of larger chunk sizes. Specifically, activations

---

Published as a conference paper at ICLR 2026
Short Context
Long Contest
Sinqular Value Entropy
0.4
Nm4, sequertial
GLA SSE
Fwd Biwd Time (log
Nmm. sequential
0.3
3.0
N-4. masking
N6 sequentell
0.2
2.5
K- k- masking
N- masking
2.0
N-4, varlen
Nmm, varten
1.5
N-26. varien
16k 64k Sequence Length
Layer Index
Sequence Length
Figure 3: Singular value en- tropy of states. SSE with di- agonal gating exhibits higher entropy than GLA.
Figure 4: Speed comparison of different SSE implementations. The number of selected partitions is fixed at K = 1, while the total number of partitions N is varied. We set attention heads to 8, head dimension to 128, and cu.seqlens to [0, L/2, L].
repeated N times, and the top-k indices utilized mask QKV vectors zero The additional dimension introduced by replication can be merged into the head dimension for execution.
Efficient implementation via varlen technique. For long contexts, redundant computation is avoided by first reordering QKV vectors with the top-k indices, grouping tokens by partition. A new eu seglens parameter then computed from the reordered sequences for operator execution. See Figure 8 for an illustration.
Fusing always-selected partitions. (wways-selected partitions are fused through concatenation along appropriate dimensions, eliminating multiple sequential calls to the linear attention operator.
We evaluate the runtime performance of different SSE implementations, including sequential kernel invocation, masking, and varlen, across input lengths from 8 to 256k tokens. As Figure 4 illustratess masking competitive for short sequences ($1k), whereas varlen more efficient for longer se- quences. Furthermore, the varlen technique enables favorable scalability with respect to state size N, maintaining nearly constant runtime when is fixed.
5 EXPERIMENTS
We evaluate SSE across benchmarks spanning language modeling, in-context retrieval, and rea- soning. All models are built on the MHA-SwiGLU backbone, with attention as the only varying componen and trained at two scales (600M and 2B parameters) under a unified pipeline cover- ing pretraining, long-context extension, reasoning-focusee distillation, reinforcement learning, and Transformer conversion. In Section 5.1 we show that SSE consistently outperforms existing lin- ear attention models, with particularly strong gains in retrieval settings. When integrated into hybrid architecture (-H), SSE achieves performance on par with Transformers. Notably, hybrid SSE-H demonstrates reasoning capabilities comparable to state-of-the-art Transformers under iden- tical training regimes (Section 5.2). We primarily focus on SSE with GLA-style diagonal gating and additionally investigate integrating the delta-rule mechanism into SSE, denoted as SSE-GDN (Section (.3). We further benchmark the efficiency of SSE in Section 5.4 and ablate several key design choices in Section 5.5 Additional details architectures, baselines, and training setups
provided in Appendix
5.1 LANGUAGE MODELING AND RETRIEVAL
Language Modeling. We conduct small-scale pretraining to compare different attention mecha- nisms, using 600M and 2B models trained on 15B and 100B tokens, respectively. We then evaluate these models zero-shot commonsense reasoning performance. As shown in Table [1] SSE-n4 achieves strong results among advanced linear attention models, outperforming Transformers and matching MoM despite using fewer parameters (600M vs. 800M). In the hybrid setting, SSE-H also surpasses both GLA-H and Transformers at the 2B scale, reaching state-of th-the-art performance.
In-context Retrieval SSE enhances retrieval through sparse state expansion and row-sparse up- dates. As Table [ illustrates, SSE consistently outperforms other linear attention models on real- world retrieval-intentie tasks, narrowing the gap with Transformer. At both 600M and 2B scales,
To ensure fair comparison, all of models vary only in the attention mixer and do not include convo- lutional excers for the Mamba family, which inherently 1D convolutions and unified blocks.
7

---

Published as a conference paper at ICLR 2026
Table 1: Zero-shot language modeling and recall performance. For SSE and MoM, and
(e.g., 14k1) denote expansion and sparsity ratios, respectively. Both adopt diagonal gating, as in
GLA. MoM is implemented without token truncation for stronger baseline. SSE-GDN uses DPLR gating, following GDN. The best and second-best results are bolded and underlined, respectively,
CommonSense Reasoning
Real-world Recall
Model
Wiki.p
PIQA
Hella.
Wino.
ARC-e
ARC-c
SIQA
Avg.
FDA
SWDE
SQuAD
Avg.
600M params. 15B tokens
Transformer
33.47
61.92
31.60
51.38
48.02
23.12
37.31
42.22
74.50
60.67
32.67
55.95
GLA
43.97
61.21
29.61
51.46
47.01
23.29
36.59
41.53
9.44
23.40
23.06
18.63
Mamba
36.62
63.82
32.20
SU.43
49.62
23.38
38.59
43.01
6.35
21.60
24.93
17.63
Mamba2
34.89
62.89
32.84
48.93
49.20
23.81
37.92
42.60
20.51
30.42
27.21
26.05
GDN
35.26
62.79
32.16
51.85
50.46
24.23
36.80
43.05
14.79
33.48
26.24
24.84
KDA
34.76
62.19
31.81
$1.22
49.96
24.83
37.41
42.90
24.50
35.55
26.71
28.92
SSE-n4k1
35.52
61.59
31.71
51.62
49.45
24.74
38.33
42.91
33.67
33.30
26.51
31.16
SSE-n4k2
35.33
62.95
31.64
50.43
49 66
25.34
38.28
43.05
30.85
36.90
27.82
31.86
SSE-GDN-n4kI
34.15
63.11
32.22
50.36
50.59
23.98
37.46
42.95
41.56
41.13
30.83
37.84
800M (600M active) params,
15B tokens
MoM-n4k1
36.05
62.57
31.64
$1.70
48.53
23.46
37.92
42.6 64
18.97
36.36
27.75
27.6
MoM-n4k2
35.11
63.87
32.47
49.64
49.54
23.81
37.36
42.78
21.05
37.08
26.68
28.27
params, 100B tokens
Translormer
16.46
71.82
$2.70
58.80
63.76
31.31
42,89
53.55
86.48
82.99
49.53
73.00
GLA
21.30
68.44
42.82
54.38
59.43
28.24
41.45
49.13
51.27
59.86
36.73
49.29
Mamba2
17.27
71.87
53.34
57.06
64.52
32.59
42.17
53.59
$6.62
63.28
43.06
54.32
GDN
17.08
72.25
53.20
57.70
65.61
32,68
42.89
54.05
54.63
66.88
39.88
53.79
SSE-n4k1
16.92
71.98
53.62
60.14
66.67
32.94
42.07
54.57
71.23
69.94
43.20
61.46
GLA-H
17.27
70.95
51.58
57.30
65.03
31.31
42.73
53.15
78.22
78. 44
45.84
67.52
SSE-H-n4k1
16.47
71.93
53.65
60.38
65.28
31.91
43.71
54.48
84.48
80.65
47.49
70.87
SSE-n4 significantly improves recall over GLA and GDN on all tasks, despite GDN's use of an explicit erasure mechanism. SSE-GDN further improves performance (see Section 5.3). While pure linear variants still lag behind Transformer, the hybrid SSE-H further boosts retrieval performance, surpassing GLA-H and approaching Transformer-lewe levell accuracy. We also benchmark our 2B models on synthetic S-NIAH tasks within the RULER (Hsich et al.) (2024) suite. As shown in Table 2 SSE demonstrates superior retrieval capabilities, further validating the benefits of its architectural design.
Table 2: Performance comparison on Single-NIAH tasks in RULER. All models have 2B param- eters and are trained on 100B tokens using 8k variable-length sequences.
S-NIAH--
S-NIAH-2
S-NIAH-3
Model
2K
4K
8K
2K
4K
8K
2K
4K
8K
Transformer
100.0 100.0
100.0
100.0
100.0
100.0
100.0
100.0
100.0
GLA
100.0
100.0
87.6
100.0
81.8
23.2
88.0
62.2
16.2
GDN
100.0
100.0
100.0
98.8
62.4
8.2
96.2
63.8
9.0
SSE-n4k1
100.0
100.0
100.0
100.0
99.6
85.2
100.0
88.2
62.2
GLA-H
100.0
100.0 100.0
100.0
100.0
100.0
98.8
96.4
70,8
SSE-H-n4k1
100.0
100.0 100.0 100.0
100.0
100.0
100.0
99.4
97.4
Data Scaling and Length Extension. We scale both 2B SSB-H and Transformer baseline to 2T pre- training tokens, followed by a 32k context extension with 250B tokens. These models are then eval- uated on benchmarks spanning knowledge, reasoning, and retrieval. As presented Table SSE-H remains competitive with the Transformer baseline across all task categories, while achieving higher average accuracy. Notably, outperforms the Transformer on several challenging benchmarks, in- cluding MMLU, MMLU-Pro, and C-Eval, which are widely used to assess general knowledge and multilingual reasoning ability. We also report results on six single- and 32k from RULER, where SSE-H shows strong long-context retrieval performance (Appendix H2 Table 10p. These results highlight the robustness and scalability of our hybrid architecture under large-scale training and long-context settings.
tasks up to
Converting Pre-trained Transformers. Following prior work (Kasai et al. 2021). we convert a pretrained Transformer by replacing part of its softmax attention layers with linear attention, fol-
8

---

Published as a conference paper at ICLR 2026
Table Benchmarking 2B SSE-H and Transformer baseline after long-context extension.
Model
MMLU
MMLU-Pro
C-Eval
AGlEval
TrviaQA
BBII
SWDE
SQuAD
Drop
GSM8K
Avg.
Transformer
24.2
55.9
21.8
85.5
52.0
30,9
50.6
45.1
SSE-H-n4k1
54.5
26.1
59.7
39.6
22.7
39.2
84.0
51.1
30.1
49.2
45.6
Table 4: Reasoning ability of SSE-H. Average accuracy (avg @ 32) for AIME24/25 reported.
Model
AIME24
AIME25
MATH500
OlympiadRench
AMC23
DeepScaleR-I.F B-eallRR-----Peeee 3BB-Preview (Luo et al
2025
43.1
87.8
50.0
73.6
Qwen3-1.7B (Thinking) Yang ct al.
2025
48.3
36.8
93.4
DeepSeek-B RI -iit-Dll-- Distill-Qwen- I.SB Guo et
2025a
28.9
23.5
83,9
43.3
62.9
DeepSeek-BB R1 Distill-- Qoall--eeennBBB Guo et al. 2025a
55.5
39.2
92.8
M1-3B Wang al. 2025b)
23.0
22.0
81.7
43.6
$6.0
PromptCol (Zhao al. 2025
35.2
24.6
84.6
50.7
Transfommc-22 (Ours)
64.1
52.8
93.0
83.3
92.0
SSE-H-n4k1-2B (Ours)
64.5
50.2
92.1
91.4
lowed by continued training. Applied to 2B model during the 128k long-context training stage and evaluated on three 32k retrieval tasks, the converted SSE-H outperforms GLA-H H and further narrows the gap with the original Transformer (Table TT Appendix F.2)
5.2 REASONING ABILITY
Reasoning is particularly challenging for efficient architectures, as it demands multi-step computa- tion, symbolic manipulation, and precise memory access. We evaluate whether hybrid linear models like SSE-H effective for reasoning under standard pipelines with supervised distillation and rein- forcement learning. We use math-focused benchmarks including MATH500, AIME24/25, AMC23, and OlympiadBench. For AIME24/25, we repeat the evaluation set 32 times and report avg 32 with temperature 1.0 and top-p 0.7, comparing against strong open-source small-scale reasoning models covering both softmax and hybrid designs.
As shown in Table 4 our 2B source Transformers improviing AIME24 from 48.3 to 64.5 (+16.2) and AIME25 from 36.8 to 50.2 (+13.4) and approaches much larger models such
achieves substantial gains over similarly sized open-
BaaBSeek-RR--issiil---eenn cct al. 2025a). Trained under the same pipeline, SSE-H matches the Transformer baseline, and 12B SSE-H variant converted from Transformer also maintains parity with its softmax counterpart (Appendix F22 Table D. These results show that SSE-H competitive with state-of-the-an small reasoning models and scales effectively for reasoning tasks.
(Guo
5.3 INTEGRATING SSE WITH GDN
Beyond GLA-style transitions, additionally provide comprehensive evaluation of integrating the delta-rule transition into SSE, denoted as SSE-GDN. Experimental details are provided in Ap- pendix F3 Overall, SSE-GDN achieves state-of-the-ar recall performance (Tables [ and 14), and its
hybrid variant, SSE-GDN-H, consistently outperforms GDN-H in the 15B MoE setting (Table 15).
5.4 EFFICIENCY ANALYSIS
We benchmark the efficiency of SSE through attention runtime and end-to-end training throughput comparisons. Experimental details are provided in Appendix F4 Overall, SSE surpasses Trans- former at sequence lengths of 32k and beyond, while remaining slower than GLA and GDN due to the inherent overhead introduced by state expansion, reflecting clear speed reeed--rrroaeeeeee tradeoff.
5.5 ABLATION STUDY
Effectiveness of State Size Scaling Recall. We ablate models (20B tokens) by varying the number of partitions (π)) and the top-k selection size (k), As shown in Figure 5 SSE scales recall effectively without increasing parameters: (1) with fixed sparsity ratio k/n, recall grows nearly lin- early with n; (2) with fixed larger k/n improves recall up moderate levels (e.g., 50%), but saturates, reducing to vanilla linear attention when k/n=1. Since higher n or k/n increases compu- tation and latency, adopt the n441 setting for most experiments, leaving room for improvement.
Effectiveness of boltmax-based Row Selection. We ablate 600M models (15B tokens) to test softmax as the feature map for row-sparse updates. As shown in Table softmax yields much higher

---

Published as a conference paper at ICLR 2026
Average Recall Accuracy (%)
0.7
k/h=25%
GLA
k/%=50%
SSELLEWAYSSSELECTE
0.5
SSE
38
0.4
0.3
36
0.1
34
0.0
8
16
10 12
16
Number of State Partitions (n)
Layer
Figure 5: Scalability of SSE with respect to state capacity. denotes the number of parti- tions; is the top-k hard selection size. Average recall accuracy reported over FDA, SWDE, and SQuAD.
Figure 6: Row sparsity ratio state updating. Measured as the fraction of k, channels with values below 1e-5. For SSE, the sparsity evaluated within the selected partitions.
Table Ablation on shared parameters.
Table 5: Ablation softmax selection.
Params. (non-embed)
Models
Recall-Avg.
Models
Recall-Avg.
SSE-n4kI-k.shl
24.23 31.16 (+6.93)
SSE-n4k1
600M (300M)
31.16
SSE-n4kl 1441-k4.sooomaa
w/o shared pareams-
890M (580M)
(-5.21)
recall than the common SiLU map, consistent with information categorization perspective. We also analyze k, sparsity in 2B GLA and SSE-n4k1 within the selected partitions (threshold 1e-5). Figure shows that GLA with SiLU produces sparsity, whereas SSE with softmax induces 5% and 42% sparsity, writing densely to the always-selected partition and sparsely to others.
Effectiveness of Shared Attention Parameters. SSE shares attention parameters (QKV and gate projection) across partitions, with LoRA applied only to QK of the shared one. We compare against variant with separate parameters for all N-1 partitions ("w/o shared-params"), As shown in Table 6, despite nearly doubling non-embedding parameters, this variant yields worse recall and only minor CSR gains (see Table 13}. This supports our core assumption: linear attention limited by state capacity, not parameter count, and parameter sharing provides both efficiency and stability.
Effectiveness Write-Read Gating. We ablate 2B models (100B tokens) test four variants: no gate, write-only, read-only, and our proposed write-read gate (applied to both input KV and output Q) shown Table 7 the write-read gate consistently yields the best recall and commonsense reasoning (CSR) performance, validating its role in effective partition selection.
Table 7: Ablation write-read gating. Comparison of four variants: gate, write-only, read- only, and write-read (SSE). / indicates applied; X indicates not applied.
Models
Write
Read
Params.
Tokens
Wiki
FDA
SWDE
SQuAD
Recull-Avg,
CSR-Avg.
SSE-n4k1
2B
100B
18.44
64.07
65.17
40.65
56.63
53.63
gate
x
2B
100B
18.63
51.91
63.91
39.81
51.87 (-4.76)
53.01 (-0.62)
write gate
x
2B
100B
18.65
54.54
61.12
37.10
50.92 (-5.71)
52.82 (-0.81)
read gate
2B
100B
18.62
57.26
64.45
40.18
53.96 (-2.67)
52.91 (-0.72)
CONCLUSION
We enhance the long-context modeling capacity of linear attention through two innovations: row- sparse state updates and Sparse State Expansion (SSE). Our framework conceptualizes state updates as information categorization and enables efficient expansion contextual states into multiple parti- tions with controlled parameters, yielding scalable capacity and more discriminative representations. We demonstrate that SSE and its hybrid variant (SSE-H) achieve superior performance across lan- guage modeling, retrieval, and reasoning tasks. Notably, our 2B SSE-H achieves sfate-of-the-an the-art mathematical reasoning among small-scale models, surpassing similarly sized Transformers. These results highlight SSE scalable and efficient architecture for high-fidelity long-context modeling.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Ekin Akyürek, Bailin Wang, Yoon Kim, and Jacob Andreas. In-context language learning: architec- tures and algorithms. In Proceedings of the 41 International Conference on Machine Learning.
pp. 787-812, 2024.
Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré. Zoology: Measuring and improving recall in efficient language mod- els. In Proceedings of 12th International Conference on Learning Representations (ICLR). ICLR, 2024a.
Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, James Zou, Atri Rudra, and Christopher Re. Simple linear attention language models balance the tecall-throughput tradeoff. In International Conference on Machine Learning, pp. 7763-1840 PMLR, 2024b.
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xistm: Ex- tended long short-term memory. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Ali Behrouz, Peilin Zhong, and Vahab Mirrokni. Titans: Learning memorize at test time. arXiv preprint M.XXv:2501.00666 2024.
Ali Behrouz, Meisam Razaviyayn, Peilin Zhong, and Vahab Mirrokni. It's all connected: A jour- ney through test-time memorization, attentional bias, retention, and online optimization. arXiv preprint arXiv:2504.13111 2025.
Aili Chen, Aonian Li, Bangwei Gong, Binyang Jiang, Bo Fei, Bo Yang, Boji Shan, Changqing Yu, Chao Wang, Cheng Zhu, al. Minimax-ml: Scaling test-time compute efficiently with lightning attention. arXiv preprint arXiv:2506.1330 2025.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. In International Conference on Learning Representations, 2020.
Krzysztof Marcin Choromanski, Han Lin, Haoxian Chen, Arijit Sehanobish, Yuanzhe Ma, Deepali Jain, Jake Varley, Andy Zeng, Michael S Ryoo, Valerii Likhosherstov, et al. Hybrid random features. In International Conference on Learning Representations, 2021.
Yuhong Chou, Man Yao, Kexin Wang, Yuqi Pan, Rui-Jie Zhu, Jibin Wu, Yiran Zhong, Yu Qiao, Bo Xu, and Guoqi Li. Metala: Unified optimal linear approximation to softmax attention map. Advances in Neural Information Processing Systems, 37:7 37:71034-71067, 2024.
Róbert Csordás, Piotr Piekos, Kazuki Irie, and Jürgen Schmidhuber. Switchhead: Accelerating transformers with mixture-of-experts attention. Advances in Neural Information Processing Sys- tems, 37:74411-744388 2024.
Tri Dao and Albert Gu. Transformers are ssms: generalized models and efficient algorithms through structured state space duality. In Proceedings of the 41st International Conference on Machine Learning. pp. 10041-10071, 2024.
Xin Dong, Yonggan Fu, Shizhe Diao, Wonmin Byeon, Zijia Chen, Ameya Sunil Mahabaleshwarkar. Shih-Yas Yang Liu, Matthijs Van Keirsbilck, Min-Hung Chen, Yoshi Suhara, al. Hymba: A hybrid- head architecture for small language models. arXiv preprint arXiv:2411.136. 2024.
Jusen Du, Weigao Sun, Disen Lan, Jiaxi Hu, and Yu Cheng. Mom: Linear sequence modeling with mixture-of-memoriaa arXiv preprint arXiv:2502.13685. 2025.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-39, 2022.
11

---

Published as a conference paper at ICLR 2026
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. arXiv preprint DrXiv:2212.14052. 2022.
Tianyu Fu. Haofeng Huang, Xuefei Ning. Genghan Zhang, Boju Chen, Tianqi Wu, Hongyi Wang, Zixiao Huang, Shiyao Li, Shengen Yan, al. Moa: Mixture of sparse attention for automatic large language model compression. arXiv preprint arXiv:2406. /4909, 2024.
Paolo Glorioso, Quentin Anthony, Yury Tokpanov, James Whittington, Jonathan Pilault, Adam Ibrahim, and Beren Millidge. Zamba: A compact 7b ssm hybrid model. arXiv preprint arXiv:2405.16712 16712, 2024.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In First Conference on Language Modeling, 2023.
Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces, In International Conference on Learning Representations, 2021.
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, 2022.
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arXiv:2501. 12948, 2025a.
Han Guo, Songlin Yang. Tarushii Goel, Eric Xing, Tri Dao, and Yoon Kim. Log-linear attention. arXiv preprint 250vv25506.0471,, 2025b.
Cheng-Ping Hsich, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. Ruler: What's 's the real context size of your long-context language models? In First Conference on Language Modeling, 2024.
Samy Jelassi, David Brandfonbrener, Sham M Kakade, and Eran Malach. Repeat after trans- formers are better than state space models at copying. In Proceedings of the 41st International Conference on Machine Learning, pp. 21502-21521, 2024.
Peng Jin, Bo Zhu, Li Yuan, and Shuicheng Yan. Moh: Multi-head attention as mixture-of-head attention. arXiv preprint arXiv:2410.1 1842, 2024.
Jungo Kasai, Hao Peng, Yizhe Zhang, Dani Yogatama, Gabriel Ilharco, Nikolaos Pappas. Yi Mao, Weizhu Chen, and Noah A Smith. Finetuning pretrained transformers into mns. In 2021 Confer- ence Empirical Methods in Natural Language Processing, EMNLP 2021, pp. 10630-10643. Association for Computational Linguistics (ACL), 2021.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In International conference on ma- chine learning, pp. 5156-5165. PMLR, 2020.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu. Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. In International Conference Learning Representations, 2020.
Aonian Li, Bangwei Gong, Bo Yang, Boji Shan, Chang Liu, Cheng Zhu, Chunhao Zhang, Congchao Guo, Da Chen, Dong Li, et al. Minimax-O1: Scaling foundation models with lightning attention. arXiv preprint www.iv:2501.88/111 2025a.
Yixing Li, Ruobing Xie, Zhen Yang. Xingwu Sun, Shuaipeng Li, Weidong Han, Zhanhui Kang, Cheng, Chengzhong Xu, Di Wang, et al. Transmambu: Flexibly switching between trans- former and mamba. arXiv preprint arXiv:2503.24063 2025b.
12

---

Published as a conference paper at ICLR 2026
Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, et al. Jamba: A hybrid transformer-
Shaked Meirom, Yonatan Belinkov, Shai mamba language model. arXiv preprint arXi::2403 2024.
Ao Liu, Botong Zhou, Can Xu, Chayse Zhou, ChenChen Zhang, Chengcheng Xu, Chenhao Wang, Decheng Wu, Dengpeng Wu, Dian Jiao, et al. Hunyuan-turbos: Advancing large lan- guage models through mamba-transformer synergy and adaptive chain-of-thought arXiv preprini arXiv:2505.15431, 2025.
Michael Luo, Sijun Tan, Justin Wong, Xiaoxiang Shi, William Y. Tang, Manan Roongta, Colin Cai, Jeffrey Luo, Li Erran Li, Raluca Ada Popa, and Ion Stoica. Deepscaler: Surpassing ol-peeveew with a 1.5b model by scaling rl, 2025. URL htt.l: ptttt::://petttttooiiiioooooeeee -ot.ettyyoaiio--b...ooo.........
notion.
te/Deep e
Og-RR-11968 902c1 4600000e0 with-a 1--5-------- 5B- Mode poooeel-b-- -by-bby-lll
Notion Blog.
Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pas- canu, and Soham De. Resurrecting recurrent neural networks for long sequences. In International Conference Machine Learning, pp. 16670-26699) PMLR, 2023.
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Leon Derczynski, et al. Rwkv: Reinventing mns for the transformer era. In Findings of the Association for Computational Linguistics: pp. 14048-14077, 2023.
Linguittti EMNLP 2023,
Bo Peng, Ruichong Zhang, Daniel Goldstein, Eric Alcaide, Xingjian Du, Haowen Hou, Jiaju Lin, Jiaxing Liu, Janna Lu, William Merrill, al. Rwkv-7" goose" with expressive dynamic state evolution. arXiv preprint arXiv:2503. 14456, 2025.
H Peng, N Pappas, D Yogatama, R Schwartz, N Smith, and L Kong. Random feature attention. In International Conference on Learning Representations (ICLR 2021), 2021.
Piotr Piekos, Róbert Csordás, and Jürgen Schmidhuber. Mixture of sparse attention: Content-based learnable sparse attention via expert-choice routing. arXiv preprint 2rXiv:2505.003355 2025.
Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. The devil in linear transformer. Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 7025-7041, 2022.
Zhen Qin, Dong Li, Weigao Sun, Weixuan Sun, Xuyang Shen, Xiaodong Han, Yunshen Wei, Bao- hong Lv, Xiao Luo, Yu Qiao, et al. Transnormerllm: A faster and better large language model with improved transnormer. arXiv preprint arXiv: 2023a.
Zhen Qin, Dong Li, Weigao Sun, Weixuan Sun, Xuyang Shen, Xiaodong Han, Yunshen Wei, Bao- hong Lv, Fei Yuan, Xiao Luo, et al. Scaling transnormer to 175 billion parameters. arXiv preprint arXiv:2307. 209999....99...... 4, 2023b.
Zhen Qin, Songlin Yang. and Yiran Zhong. Hicrarchically gated recurrent neural network for se- quence modeling. Advances in Neural Information Processing Systems, 36:33202-33221, 2023c
Olivier Roy and Martin Vetterli. The effective rank: A measure of effective dimensionality. In 2007
15th European signal processing conference. pp. 606-610. IEFF, 2007.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International conference machine learning, pp. 9355-9366. PMLR, 2021.
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Y Wu, et al. Deepseekmath: Pushing the limits mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300 2024.
Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.0522 2020.
Julien Siems, Timur Carstensen, Arber Zela, Frank Hutter, Massimiliano Pontil, and Riccardo Grazzi. Deltaproduct: Improving state-tracking tracking in linear mns via houscholder products. In ICLR 2025 Workshop on Foundation Models in the Wild, 2025.
13

---

Published as a conference paper at ICLR 2026
Yu Sun, Xinhao Li, Karan Dalal, Jiarui Xu, Arjun Vikram, Genghan Zhang, Yann Dubois, Xinlei Chen, Xiaolong Wang, Sanmi Koyejo, et al. Learning to (lear test time): Rnns with expressive hidden states. arXiv preprint arXiv:2407.04620, 2024.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue. Jianyong Wang, and Furu Wei. Retentive network: A successor transformer for large language models. arXiv preprint arXiv:2307.08622, 2023.
Kimi Team, Yu Zhang, Zongyu Lin, Xingcheng Yao, Jiaxi Hu, Fanqing Meng, Chengyin Liu, Xin Men, Songlin Yang. Zhiyuan Li, et al. Kimi linear: An expressive, efficient attention architecture. arXiv preprint (..iv:55/00266990 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 13911, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- fion processing systems, 30, 2017.
Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, al. An empirical study of mamba- based language models. arXiv preprint arXiv:2406.07887, 2024.
Dustin Wang, Rui-Jie Zhu, Steven Abreu, Yong Shan, Taylor Kergan, Yuqi Pan, Yuhong Chou, Zheng Li. Ge Zhang, Wenhao Huang. et al. A systematic analysis of hybrid linear attention. arXiv preprint arXiv:2507.06457, 2025a.
Junxiong Wang, Wen-Ding Li, Daniele Paliotta, Daniel Ritter, Alexander M Rush, and Tri Dao. M1: Towards scalable test-time compute with mamba reasoning models. arXiv preprint 1rX11:2504.10449, 2025b.
Ke Alexander Wang, Jiaxin Shi, and Emily B Fox. Test-time regression: unifying framework for designing sequence models with associative memory. arXiv preprint arXiv:2501.12333 2025c.
An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint 0988888 2025.
Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule. arXiv preprint arXiv:2412.006444 2024a.
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training. In International Conference on Machine Learning, pp. 56501-56523.. PMLR, 2024b.
Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transform- ers with the delta rule over sequence length. The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024c.
Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yu Yue, Weinan Dai, Tiantian Fan, Gaohong Liu, Lingjun Liu, et al. Dapo: An open-source Ilm reinforcement learning system at scale. arXiv preprint Xrxiv:2003. 14476, 2025.
Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Re. The hedgehog & the poreu- pine: Expressive linear attentions with softmax mimicry. In The Twelfth International Conference on Learning Representations, 2024a.
Tianyuan Zhang, Sai Bi, Yicong Hong, Kai Zhang, Fujun Luan, Songlin Yang. Kalyan Sunkavalli, William T Freeman, and Hao Tan. Test-time training done right. arXiv preprint arXiv:2505.2388 238844 2025.
14

---

Published as a conference paper at ICLR 2026
Yu Zhang, Songlin Yang, Rui-Jic Zhu, Yue Zhang, Leyang Cui, Yiqiao Wang, Bolun Wang, Freda Shi, Bailin Wang, Wei Bi, et al. Gated slot attention for efficient linear-time sequence modeling. Advances in Neural Information Processing Systems, 37: 1187110887-11888,, 2024b.
Xueliang Zhao, Wei Wu, and Lingpeng Kong. Scaling reasoning without attention. arXiv preprint arXiv:2505.22425, 2025.
15

---

Published as a conference paper at ICLR 2026
A LIMITATIONS AND FUTURE WORK
While our experiments demonstrate the promising scalability of SSE with GLA-style transition mechanisms, we have not yet conducted comprehensive evaluations with delta-rule baselines (e.g., GDN), Moreover, more sophisticated categorization functions, such as state-aware mechanisms, could be effectively combined with SSE-GDN, offering a promising direction for future work.
Concurrently, we observe that non-hybrid SSE still exhibits notable retrieval gap compared to well- trained Transformer models. Inspired by SSE's effective state size scalability, we consider further investigation into SSE's state scaling law potential solution. Given the high search cost of partition count as a hyperparameter, we plan to utilize an up-training scheme to expand state size based on existing checkpoints. Thanks our parameter sharing strategy, SSE can smoothly transition be- tween different N, K configurations, quickly recovering general ability without requiring extensive retraining or distillation. We consider SSE's state up-training key direction for future work.
We acknowledge that SSE's current context-only hard-selection mechanism may not be optimal and warrants broader exploration. The computation of e, provides flexible interface for incorporating inductive biases into the model. A compelling future direction involves incorporating dependencies on positional and historical information, leading to more general form such e₁ gt, x,, S;-1)-
B RELATED WORK
Linear Attention. Linear attention (Katharopoulos et al. 2020) replaces the softmax kernel of Transformer attention with dot-product of feature maps, thereby achieving linear complexity and more efficient long-context inference. Early research in this area (Choromanski et al., 2020 Peng et al.) 2021 Choromanski et al.) 2021: Zhang et al., 2024a) primarily focused designing spe- cialized non non-negative feature maps to approximate the behavior of the softmax kernel. Subsequent works (Sun et al., 2023 Qin et al. 2023a) further simplified linear attention, removing unneces- sary normalizer terms and other specialized feature maps. Concurrent with the evolution lin- attention, other generalized linear recurrent mechanisms emerged, such State Space Models (SSMs) Gu et al. [2021] 2022) Fu et al. ; Gu & Dao (2023); Dao & Gu (2024) and linearized RNNs (Qin al. 2023c Peng al., 2023 Orvieto et al. 2023 Beck et al. 2024). Through con- tinuous development, these mechanisms have converged towards a unified decay and outer-pros outer-product
update formulation (Chou et al. 2024), similar the unnormalized linear attention discussed in our work.
highlight the benefit of introducing input- dependency into state decay, enabling more adaptive memory control through gating mechanisms. More recent trends involve densifying the state transition matrix, exemplified by delta-rule-basee at- Siems et al. (2225). While sophisticated structures, 2024 Behrouz et al.] 2024) Wang et al.) 2025c; Behrouz et al.) 2025) designed based on online learning. offer advanced state updates, their nonlinear de- pendencies often introduce trade-offs in implementation efficiency. In contrast to these works that primarily focus on the state transition matrix, our work re-examines how new information is more effectively incorporated into the state. We achieve this by innovatively connecting the state update mechanism to an information categorization process.
Recent findings (Yang et al.
2024b Gu
Dao
2023)
tentions (Yang et al. 2024c at Peng et al.
2025
such as various TTT variants (Sun et al.
State Expansion. Recent literature (Arora et al., 2024a Jelassi et al. 2024 Akyürek et al. 2024
Waleffe al. 2024) consistently highlights the challenges posed by limited state size in linear tention models, particularly evident in retrieval tasks. Consequently, state expansion has emerged as a seemingly indispensable direction. Beyond directly expanding the state. Log-Line Linear Atten- tion (Guo et al., (0255) adopts a Fenwick tree structure to enable logarithmically growing state sizes. Our work bears the closest resemblance to Lfxtuure-ffMMmorrr which models the state using ffxture-of-Fxxeene with SSE. However, SSE designs a sparse state expansion scheme specifically from an information categorization perspective, explicitly decouples state capacity from parameter count, and introduces improved feature mapping and gating mechanisms. Furthermore, we provide a comprehensive the- oretical framework demonstrating that SSE's state achieves more effective information storage. Our extensive experimental pipeline, spanning pretraining to RL, robustly confirms the proposed model's stability.
(MoM) (Du et al. (025), (MoE) approach and shares a similar sparse spirit
16

---

Published as a conference paper at ICLR 2026
Hybrid Models and Reasoning. Beyond solely expanding the state of linear attention, common compromise involves hybridizing some softmax attention layers. This approach sacrifices a degree of computational efficiency in exchange for significant improvements in retrieval performance. Re-
cent works have explored both inter-layer (Lieber et al. 2024, Glorioso et al.) 2024 Yang et al.
2024c] [i al. 2025aa Wang al., (025a) and intra-layer (Dong et al. [024] Li et al. 2025b
Behrouz et al.) 2024) hybrid paradigms. Notably, such hybrid architectures have recently demon-
strated promising overall performance, even large scales (Li et al. 2025a Chen al. 2025
Liu et al., 2025). Concurrently, the reasoning capabiliiiees of pure linear attention models remann
relatively underexplored. largely relying on hybrid architectures (Wang et al., 2025b Liu et al.) Chen et al., 2025) Zhao et al. 2025 Our work is the first to demonstrate that, under iden-
2025
tical pre-training paradigms, hybrid linear models can achieve reasoning performance on par with
Transformer models, thereby enjoying the advantages of test-time scaling. Through fine-tuning and advanced RL training, we report that our 2B SSE-hybrid model exhibits exceptional mathematical reasoning ability, achieving state-of-the-ar performance at its scale.
Sparse MoE for Quadratic Attention. In line with the Mixture-of-Exxee Expeets (MoE) paradigm widely
used in feed-forward networks (FFNs), dynamic computation
also be enabled in quadratic soft-
max attention through routing mechanisms. Existing approaches largely fall into two categories: (1)
applying MoE attention projection weights (e.g.. W, g: WE, k. W.., W., W.) (Fedus et al. 2022 Csordás iet al. 2024), which closely parallels FFN-style MoF; and (2) performing MoR at the attention-head
level by selecting token-dependent subset of heads for computation (Jin et 2024 Pickos al. 2025 Fu et al.j 2024). For the latter, attention heads and their associated KVs can be viewed as
coarse form of state allocation. In contrast, the write-read gating used for top-k partition selection in SSE extends this idea to the linear-attention setting, leveraging parameter sharing to achieve both computational and parameter efficiency. Moreover, softmax- based row selection can be interpreted as finer-grained MoE mechanism operating at the level of individual state rows.
C FORMULATION OF LINEAR ATTENTION VARIANTS
We first present a recurrent outer- peceerproduttt update formulation that provides a unified view of various linear attention modeling (Chou al. 2024):
S. = A,S,-1 + @(ki) vv. Vv.
(1I)
Here, A, denotes the state transition matrix, and & is the feature map embedding k, into a finite- dimensional space.
Gated linear attention variants (Sun et al., 2023 Qin et al. 2022 Peng et al. 2023
introduce recency
bias via fixed exponential decay or input-dependent gating (Yang et al. 2024b Gu & Dao 2023 Dao & Gu 2024). Specifically, gating mechanisms often employ diagonal (or scalar) matrices, A, diag aa)), where ax is input-dependent with values in (0, 1):
S = diag( ox S + o(k:) vv..
(12)
Another line of work introduces structured updates such as plas-oo--plus-oow-ree (DPLR) matrices
derived from the delta-rule (Schlag et al.) 2021 Yang et al. 2024c a Siems et al. 2025). General
DPLR matrices are typically structured as A, diag at) + a, b,:
S, (diag(o:) (x,) a. b.).:))._111 + ((k,) vv.
(13)
A notable special case is DeltaNet, where A, = I - I-3ck, k, Its gated variant GDN adopts A₁ = ar(I l.k,,,).
D THEORETICAL ANALYSIS OF ROW-SPARSE STATE UPDATES
In this section, we present theoretical analyses and key insights related to the information catego- rization perspective and the row-sparse update framework. Our discussion primarily focuses on the case of a linear categorization function, namely:
k, = f(x,. W+) = x, WE.
(14)
S, = A,S,11 1 + k, Vn.
(15)
17

---

Published as a conference paper at ICLR 2026
The contextual state updates in linear attention act an information categorization process. By treating state rows as distinct latent categories and conceptualizing k; = x, W as a categoriza- exhibits similar
tion function, can demonstrate that information assigned to the same state features. In the linear function setting, information similarity is measured by the inner product. The assignment result for category is defined as C {x,|xx X Wi kig ). where ki 0 is category- specific threshold. Proposition I shows that information assigned to the same state row exhibits large mutual inner products.
Proposition 1 For imputs x, satisfying ||x, = and given the categorization rule C (x, X W >> k. inputs belonging to the same category (row) Cs satisfy x,x d coos 28', where 0* arccos(
kjg Vallwal ), 8. (I,...,,}, and vice versa.
Proof 1 The categorization rule implies inputs with high acceptance intensity k; Wi belong to the i-th category (row), and the dategory-specippe threshold kig > 0.
If Xg E C., then
x, W = Will cos kin.
(16)
kigh
cos 0
cos
0',
(17)
Vd Will
which means all imputs from category have an angle with Wi less than 0', where kig E (0, WallWill and 6¹ E (0, 3). This is equivalent to saying that the angle between inputs of the same category is less than 20' Thus obtain:
x,x > ||x,|| ||x cos 20 cos 28
(18)
where x,, X2 are arbitrary inputs from category
Furthermore, if we set the threshold large enough to ensure that H where r min.,, ((Wi, W ). we can ensure that for inputs from different categories, x,X, S cos 20
Current linear attention suffers from noise and forgetting due to incomplete use of categoriza- tion assignments. As shown in Proposition [ vanilla linear attention implicitly incorporating the notion of information categorization, but does not fully leverage the resulting category assignments for compressed storage:
S
w Xs, j=1 k=1EEE
Σ
X3, W..
(19)
This mixing of information from different categories leads to major issues: (1) Noise: Propo- sitions [ and show that state rows in vanilla linear attention exhibit high similarity due to inter- category interference, resulting in state row homogenization and reduced querying effectiveness. (2) Forgetting:
Proposition 4 demonstrates that applying uniform decay to all state rows adversely limits the receptive field.
In contrast, by explicitly leveraging category assignments, we introduce row-sparse update for- mulation of linear attention:
Si
Σ W x. W..
(20)
E.EC,
Through this hard assignments, each row of the contextual state stores similar information belonging to the same category, resulting in higher inter-row discriminability and more precise information organization. Simultaneously, each element undergoes fewer decay operations, effectively extending the receptive field over longer contexts. The top-k operation introduced in this work is one of the methods achieve row-sparse updates.
18

---

Published as a conference paper at ICLR 2026
Definition The row-wise cosine similarity in vanilla linear attention is defined as Cos (Si.
(S |8:|| ||S Similarly, the row-wise cosine similarity under the row-sparse update formulation is denoted as Cos (S), St) for differentiation.
Proposition 2 Assume that W. is orthogonal and w'y 2 [i,j E I,....c) in Equation 19 the row-wise cosine similarity in vanilla linear attention is greater than that of the corresponding row- sparse version, iee, Cos S) St > Cos (S).S.)
Proof 2 We use binary-calegories as an illustration, which can be generalized to multi-calegories scenario. Let =1 = 2, S - wh, Xsy EmpeC, w Xxx Because Wa is an orthogonal matrix, we have:
S S.W.
S.
(21)
Since multiplying by an orthogonal matrix does not change the norm of a vector, therefore:
Cos (5 (S..st)
Cos S
(22)
We have demonstrated that all inputs from category j have an angle with W less than @ (Propo- sition TV. For Xs, E Cj and w', 2 O, the positive weighted sum Eayec, Wa,Xa, X5, still have an angle with W less than 8², due to the parallelogram law. Thus we can let Exjec, w2,x,, xx, = azx}, where x E Cg and a² > 0 implies module changed (1xx)|1==0 d).
Then, for vanilla linear attention, have S = a x + x'x} and S a²x} + 82x3. For nota- tional simplicity, we use a², B representing aj a²y respectively. Let cos min ( 20' cos 20 }, we have:
x > cos 28' 2 doos A,
(23)
(x), > cos 20 2 doss O,
(24)
xx) S dcos
(25)
(x2, x} < doos O,
(26)
| (a')² d + (B + 2a 8 (x), & d(a + B
(27)
Because the row index symmetric, we suppose (x², xf) 2 x2) without loss of generality. Then obtain:
(a'a² + (C0s) e)d (a (x),
Cos S..
(28)
d(a² + (a²
where refers to a sufficiently small constant. For the corresponding row-sparse linear atlention, each state row corresponds to latent category, so the contextual state stores information from different categories separately. That is, S ax} and S = BxZ. The similarity of the row-sparse attention is computed by:
Cos S) S). = [x] (x|,x3) x3)
(29)
d
The proof is given by:
dcos0 + de > xx), (holds according to Equation 25
(30)
(a' aa² + 8 (cos + +ed
(a'a²
B x3) (a 0, 0)
(31)
Cos s S Cos
S! S!
(32)
(holds according to Equations
28
29 and
This means the state cosine similarity of vanilla linear attention is larger than the corresponding row-sparse version.
19

---

of S. are strictly Published as a conference paper at ICLR 2026
For the multi-categoriee scenario, still utilize triangle inequality of norm || and Propo- sition D to reach the same conclusion. Note that in this case, we can solely focus on the two most different categories, namely (x), x3) m x to complete the derivation, which is enough
prove Proposition
Combining Equations
and
32
we derive:
Cos [8), (8),s:) > Cos S S!)
(33)
Finally, in Proposition 3 we demonstrate that more precise contextual state (established in Propo- sition enables queries to extract diverse information more effectively. For well-designed cate- gorization function f(x,, W&), expect the reading operation q,S, allow different queries to extract distinet information, a property we refer to as state distinguishability. Higher distinguisha- bility, achieved via row-sparse updates, reflects more expressive and structured state. In contrast, lower distinguishability. typically observed in vanilla linear attention, leads to state homogeniza- tion and inter-category interference, which in turn hampers the model's ability generate accurate output representations.
Definition The measure minp.q, Cos(q:S,, P.S.), referred to as state distinguishability, quanti- fies the minimum cosine similarity between states outputs, where qr Pt are arbitrary queries.
Proposition 3 For arbitrary queries 4:, 0, assume that the row bounded. The lower bound of the state distinguishability measure, minp. Cos(q. PtS,). is given by (S S.)) Moreover, row-sparse updates reduces this lower bound compared to vanilla linear atlention.
Proof Considering the cosine similarity, we have:
(q,S, PtS
Co
|q
(q.Se. P.
(34)
([qq
S |.
(()
1
where q.
[[q:1] 9.
P..
Then the lower bound of the state distinguishability measure is given
[pll
by:
min St), 2
min (q.S
(35)
||
p...
The right side of Equation
35 is:
min Sq P.S.)
min
J.S. p/S
p..q
min
girl (S S
(36)
||0
|p
1.
min (S. S/)
43
We assume the norm of S, is bounded, which means:
3c s.t. min |S 2
(37)
20

---

Published as a conference paper at ICLR 2026
The assumption implies that the contextual state is numerically bounded, thus the inner-product between state rows is governed by the cosine similarity:
min (:S( 2 P..9
||
min Cos %
S St
||S ||
(holds according to Defination
|s
2²
min Cos
S S (holds according to Equation 37
|S
(38)
which shows that the lower bound of the state distinguishability measure is related to the cosine similarity between the two most different categories. Because row-sparse update formulation has reduced the right-side similarity Cos of the measure.
it can be seen way reduce the lower bound (i St
In Proposition 4 we demonstrate that the exponential cumulative decay inherent to vanilla gated lin- ear attention adversely affects receptive fields. In contrast, row-sparse update formulation mitigates this issue and can theoretically achieve infinitely long receptive fields, as sparse decay ensures the retention of important information.
Definition The attention score of x, with respect to x, defined as Pts q:((II)=++ a,) o k.)T.
Definition Let Pts denote the attention scores and Ptk. be a given threshold, where 2 8 > 0.
The receptive field at time t is defined as: M,- max(t Pts 2 Pa.).
Let qe and kg be all-ones vectors, and let I. attention with decay factors 0z, E 10, the receptive field upper-hounded by M, max:=1..... a log( Paa)/ 1og ) such that the attention score pts is less than given threshold P.k. Conversely, row-sparse update formulation can retain important information across arbitrary spans, ensuring the attention score remains above the threshold Pet. i.e. M, t.
Proposition 4
O. For vanilla gated lin-
Proof 4 Assume that qr = K. = =[,,... 1] E Rd,, then the attention score:
Pts qi(l II (a))kk))))
II II aj
(39)
j=st1
k=1j=st1
which can characterize the utility of only decay. Effective receptive field requires the attention score Pts 2 Pih. The corresponding necessary condi- tion is:
Pe d
3 ki s.l.
max stisist
(40)
That is,
Pet
3 k, s.t. t-s logl d
log( max s+1
at))
(41)
Therefore, M, max log( // log(max at) represents the upper bound of recep- tive fields. For earlier tokens at time s', the receptive field exceeds this upper bound (t s' M2), the information is forgotten as Pts < Pih-
In contrast, for row-sparse update variant of linear attention whose decay values be scattered to actually equal 1, we can ensure the retention of important information across arbitrary spans by simply considering:
3 kk s.t. = 1 for Vs 1sjst t.
(42)
This implies that row-sparse update linear atlention ensure Mt t.
21

---

Published as a conference paper at ICLR 2026
GLA: Setting
GSA: Setting 1
MetaLA: Setting
1.00
1.00
1.00
0.75
0.75
0.75
Accuracy
Accuracy
0.50
GLA-topk
0.50
65A topk
0.50
MetalA-topk
MetalA
0.25
0.25
0.25
0.00
0.00
0.00
64
128
256
512
64
128
756
512
64
128
256
512
Model dimension
Model dlmmnsion
Model dimension
GLA: Setting 2
GSA: Setting 2
MetaLAA Setting 2
Accuracy
Accuracy
Accuracy
0.50
GLA GLA-topk
0.50
GS0itopk
MetalA MetauN-topK
64
128 dimension
256
512
128
256
512
64
128
256
512
Model dimension
Model dimension
Figure 7: MQAR results of linear attention models and their top-k row-sparse variants. Set- tings and 2 introduce zero tokens and random noise tokens, respectively. Top-k row-sparse updates lead to longer effective context and more accurate information storage within the state. Moreover, with row-sparse updates, recall performance increases monotonically with respect to state size (i.e., model dimension).
Synthetic experiment validation. We validate our theoretical analysis by applying top-k sparsity during the state update process of three linear attention models, and evaluate recall performance on the synthetic MQAR task (Arora al., 2024a). Specifically, only the state rows corresponding to the k largest values in kç are updated at each step. In MQAR (Multi- Bulli--Qery Associative Recall), the model must retrieve previously stored key-value pairs in response to multiple queries. We evaluate recall under two settings: (1) In the first setting, zero tokens are inserted between key-value pairs and queries to explicitly evaluate the model's effective receptive field. (2) In the second setting, random noise tokens are added, testing the model's robustness to interference during memory retention.
We conduct experiments on GLA (Yang et al. 2024b). MetaLA (Chou et al., 2024), and GSA (Zhang eet al. 2024b) three representative linear attention variants. These models respectively incorporate gated updates, decay-key coupling, and two-pass recurrence. We set = d/4, and train 2-layer models with alternating token and channel mixers, using 2 attention heads. As shown in Figure 7 the results demonstrate that row-sparse updates offer a simple and general approach for significantly enhancing the recall capabilities of vanilla linear attention, leading to longer effective context and more accurate information storage within the state. Moreover, with row-sparse updates, recall per- formance increases monotonically with respect to state size (i.c., model dimension).
E DETAILS OF EFFICIENT SSE IMPLEMENTATIONS
SSE expands the contextual state into multiple partitions, each consisting of different subsets of tokens. During operator execution, our goal is preserve parallelism across partitions, rather than relying on sequential invocation. To minimize unnecessary computational overhead, we introduce two implementations tailored to different sequence length regimes: masking- based version for short sequences, and a parlen-based version for long sequences. In addition, we describe a simple fusion strategy for managing always-selected partitions.
Naive implementation via masking for short contexts. In the chunk-wise operator of linear atten- tion, each kemne instance is responsible for processing single data chunk. To improve execution
22

---

Published as a conference paper at ICLR 2026
efficiency, the chunk size is typically set power of two (e.g.. 32, 64, or 128), which facilitates memory aligmment and GPU-level parallelism. The optimal chunk size depends on both the tar- get GPU architecture and the specific characteristics of the operator, and is typically determined through empirical tuning. However, variable-length training often includes a large proportion of short sequences, which can lead per-partition lengths falling below the optimal chunk size, or even below the minimum threshold (c.g.. 16). In such cases, the varlen implementation incurs ad- ditional operational overhead, and grouping tokens by their assigned partitions and processing them with cu sulleel can become inefficient. To address this, we adopt a masking-based strategy that increases parallelism through replication while enabling the use of larger chunk sizes. Specifically, each activation is replicated across partitions, and masking is applied based on the top-k selection to ensure that each partition only attends to its assigned tokens. This approach improves GPU utiliza- tion without requiring token reordering. The corresponding pseudocode is presented in Algorithm
Algorithm 1: SSE Implementation via Top-k Masking
Input: Q K. V E RR EEEEE ERRNN, cu_seqlens E R Number of partitions N, Top-k
partitions
Output: O R
// Compute top-k selection mask
Mapk = =TopkAndMask(E, K, dim = 1) {0, 1)
// Repeat input across partitions and apply masking
Q, K, V Repeat(Q, V,
N, dim 1) R
Q Q O Mopk
K K Mook:
V = V @ Mupk
// Rearrange for linear attention
Q, K, V = Rearrange(Q. K, did = (1,2)) E RLX(NH)xD
// Linear attention computation
O = LinearAttention) Q, K, V, cusseqlens);
return O:
Efficient implementation via varlen technique for long contexts. During the long-context contin- ual training phase, the sequence lengths are generally longer and more evenly distributed. allowing chunk-wise computation to operate only on the relevant tokens within each partition. This avoids the redundant computations over masked tokens inherent in the naive masking- based implementation. Specifically, we first derive the top-k partition indices and use them reorder the QKV vectors, grouping tokens sequentially by partition (from to N) within each sample. Next, new cu cu_seqlens is computed based on reordered sequences and their corresponding partition assignments. stage, each resulting subsequence corresponds to specific partition within specific sample. Given this reordering and the updated cu seglens, all partitions can be processed in parallel using chunk- wise linear attention, without introducing additional computational overhead. This implementation exhibits favorable scalability with respect to state size N, maintaining nearly constant overhead long as K (the number of selected partitions) remains fixed. The corresponding pseudocode is presented in Algorithm
this
Algorithm 2: SSE Implementation via Varlen Technique
Input: Q K, V E LRLHDD E E BLXN, cu seglees E R Number of partitions N, Top-k partitions K
Output: O E 4xDD
// Obtain reorder index IE KL and updated sequence offsets I, new cu seglens = GethndexAncdOffsettsst. cu seqlen))))
// Group tokens by partition order
Q, K, Reorder(Q, V. dim = 0) E ((K))IIIDD
V
// Linear attent i on computati on with varlen partitioning
0 = Linear Attention Q, K, V, new_cu_seqlens):
return O:
23

---

Published as a conference paper at ICLR 2026
Reorder inputs compute new cu seglens
Linear attention computation
3. Reorder outputs
Query/Key/Valle Value
partition!
partition4
Linear
Output o
Attention
Reorder
Kernel
Original order
Grouping by partition
Original order
( (cu_ (u__sellen==00.7)
new (u.segleen==(0224666711
example). We first use top-k partition indices to reorder the QKV vectors, grouping tokens by partition (from I to N) within each sample. A cu_seqlens is then constructed from the reordered sequences, where each subsequence corresponds to specific partition. All partitions are subsequently processed in parallel using linear- attention kernels. Finally, the outputs are reordered back to the original token order.
Figure 8: Varlen implementation of SSE (n4k1
Sequential kernel invocation via for-loop. For comparative analysis, we present the vanilla se- quential implementation. Initially, additional dimension representing partitions is introduced, and tokens are gathered according to their assigned partitions. Subsequently, the chunk linear at- tention kernel is sequentially invoked over all partitions within for-loop, bypassing varlen control. While this approach circumvents the overhead potentially introduced by varlen, sacrifices inherent parallelism across partitions. Consequently, this method exhibits poor scalability: runtime increases significantly with larger values of N due to its sequential computation.
Fusing always-selected partitions. Always-selected partitions can be fused through concatena- tion to avoid multiple sequential calls to the linear attention operator. In the naive implementation, always-selected partitions are concatenated along the head dimension, and no masking is applied to this portion. In the varlen implementation, always-selected partitions are concatenated with the reordered sequence along the sequence dimension, while segment- wise computation remains con- trolled via the cu.seqlens parameter. These modifications allow the linear attention operator to be invoked only once.
In the inference setting, the chunk-wise implementation described above applies directly to the pre- filling stage. During decoding, we leverage sparse indices to perform recurrent computation solely on the selected partitions, significantly reducing computational overhead.
F EXPERIMENT DETAILS
Model Configurations. Model configurations for our experiments are summarized in Table 8 We build models on the MHA-SwiGLU architecture (Touvron et al. 2023 Shazeer 2020), varying only the attention mechanism to ensure fair comparisons. Accordingly. all models exclude convolu- tional layers, except for the Mamba family, which inherently employs ID convolutions and unified blocks. We compare SSE against standard Transformer (using softmax attention) and several rep- resentative linear attention baselines, such as GLA (Yang et al. 2024b), GDN (Yang et al. (024a), Mamba (Gu Dao 2023), Mamba2 (Dao Gu 2024). KDA (Team et 2025) and MoM (Du et al. 2025). Both SSE and MoM adopt GLA-style state transition mechanisms, with n denoting the state expansion factor (i.e., the number of state partitions) and & representing the top-k hard selection size per token. The sparsity ratio is defined as k/n. We further investigate integrating the delta-rule mechanism into SSE, denoted as SSE- GDN, For the SSE family, we utilize a single always-selected selected partition and set the low-rank dimension to 64 for QK projections. The coefficient for the auxiliary loss 0.01 across both SSE(-GDN) and MoM. Additionally, we also evaluate layer- wise hybrid ar- chitectures, constructed by adding one sofumax attention layer after every five linear attention layers. In our 2B setting, this results in 3 softmax attention layers out of 18 total layers. These models are denoted by the "-H" suffix (e.g., SSE-H).
Training Setup. We conducted our primary experiments with two model scales: 600M and 2B pa- rameters. The non-embedding parameters are 300M and 1.3B. respectively. During pretraining, all models are trained with Adam W, using a maximum sequence length 8192 and a total batch size of 4M tokens. The learning rate schedule includes a linear warm-up phase followed by a constant
24

---

Published as a conference paper at ICLR 2026
Table 8: Model Architectures.
Model Size
Non-Emmbedding Params
Layers
Hidden Dimension
Heads
600M
300M
24
1024
8
2B
1.3B
18
2304
18
Layer Head 6
Layer 1 Head 12
Layer 5 Head
Layer 5 Head
Figure 9: Clustering of information within linear attention state rows. We observe that learned state representations reveal clear clustering patterns. Specifically, assign each token's value vector (represented as a point) to a specific state row (indicated by color) by taking the maximum activation its corresponding key vector. This assignment demonstrates that information within the same row tends to share similar feature representations.
learning rate 6e-4, We apply a weight decay of 0.1 and a gradient clipping of 1.0. Following pre-training, perform context length extension up 32k tokens, using cosine learning rate de- cay. Subsequently, we perform a oriinnte data distillation stage using a fine-tuning set of approximately 80k examples, trained for 5 epochs. Finally, we apply reinforcement learning with GRPO (Shao et al. 2024) for 230 steps. During RL training, we sample responses per prompt with generation limit of 32k tokens.
F.1 ANALYSIS OF CONTEXTUAL STATES
Singular Value Entropy. The singular value entropy (Roy & Vetterli
2007) is computed by:
log of 2ド=1
of
H logn
(43)
13=103
where o, (for {1....,.1 n2)) are the singular values of contextual state matrix. We average H across all attention heads for a specific layer.
The singular value distribution directly determines compressibility. Given a singular value decom- position (SVD), low-rank compression removes small singular values, since each singular value indicates the importance of its corresponding directional component. When the singular value tropy is low (a concentrated spectrum), most information lies in a few dominant singular values, and truncation preserves nearly all content. When entropy high (a flatter spectrum), removing any singular value entails noticeable information loss, making compression significantly harder.
Receptive Field Analysis. To further validate the effectiveness of SSE in long-context modeling, we analyze the receptive field of pretrained 2B SSE model and compare it with that of a GLA baseline. As visualized Figure I0 we compute the receptive field across all layers by examining the -epededepnndent gating matrices after pretraining (with maximum sequence length of 8k). Specifically, we extract the last 128 tokens and measure the effective receptive field width across different channels in each layer using threshold of 0.001. The results show that SSE consistently exhibits larger receptive fields than GL.A across all layers, confirming SSE's enhanced capacity for long-range information integration.
F.2 EXTENDED EXPERIMENTAL RESULTS
Conversion and Reasoning Results. To further evaluate the scalability of SSE, we convert pretrained 12B Transformer into SSE-H variant. The conversion involves layer-wise
12B SSE-H:
25

---

Published as a conference paper at ICLR 2026
Comparison Effective Receptive Fields between GLA SSE
Mean Channel Receptive Field
4000
2000
GLA avg.crf
GLA = Std
SSE (Sparse) avg.crf
SSE Std
14
15
36
Layer
Figure 10: SSE exhibits larger receptive fields than GLA. We visualize the receptive field width of pretrained 2B SSE and GL.A models after 8k-seqlen training. Values are computed over the last 128 tokens using threshold of 0.001. SSE shows consistently broader receptive fields across all layers.
stacking pattern of SSE-SWA-SSE-SSA-SSoE training stage, followed by supervised distillation. As shown in Table SSE-H matches the math- ematical reasoning accuracy of its softmax-attentinn counterpart, demonstrating the model's scala- bility to larger model sizes.
layers, applied during the 128k long-context
Table 9: Reasoning ability of the 12B SSE-H variant under the conversion paradigm.
Model
AIME24
AIME25
MATH500
OlympiadBench
AMC23
Transformer-I 12B
58.7
95.8
82.7
SSE-H-n4k1-12B
74.3
57.3
96.0
85.3
96.3
2B SSE-H: Results 32k RULER Benchmarks. As discussed in Section 5.1 scale the 2B SSE-H and Transformer baselines to 2T pretraining tokens, followed by 32k extension with 250B additional tokens. We evaluate six single- and multi-needle tasks from the RULER benchmark at lengths up to 32k. As shown in Table TO, SSE-H demonstrates strong long-context retrieval, matching surpassing Transformer several 32k tasks. These results underscore the robustness and effectiveness of our hybrid design under large-scale and long-context training.
Table 10: Performance comparison on NIAH tasks in RULER after long-context extension. All models have 2B parameters with a context length of 32k. Results are reported in zero-shot setting.
S-NIAH-1
S-NIAH-2
S-NIAH-3
Model
8K
16K
32K
8K
16K
32K
8K
16K
32K
Transformer
100.0
100.0
99.4
100.0
100.0
70.8
97.8
98.8
67.4
SSE-H-n4k1
100.0
100.0
93.6
100.0
100.0
84.4
100.0
98.2
90.2
MK-NIAH-1
MQ-NIAH
MV-NIAH
Model
8K
16K
32K
8K
16K
32K
8K
16K
32K
Transformer
95.8
90.0
55.4
86.1
65.5
28.6
92.4
53.0
18.1
SSE-H-n4k1
91.0
84.8
64.6
89.1
67.8
44.6
87.3
66.8
40.6
2B SSE-H: 32k Retrieval after Conversion. As discussed in Section 5.] strategy that replaces subset of softmax attention layers in pretrained Transformer with linear
we explore a conversion
26

---

Published as a conference paper at ICLR 2026
attention layers, followed by continued training limited dataset. Specifically, apply this method to a 2B Transformer during its 128k long-context training stage, using approximately 100B tokens. We evaluate the resulting models on three 32k-context retrieval tasks. As shown in Table IT] the converted SSE-H model achieves substantially better retrieval than GLA-H and further narrows the performance gap with the original Transformer.
Table 11: Evaluation long-range retrieval (32k) under the conversion paradigm.
Model
Params.
Quest
Qampari
Table Query
Avg.
Transformer
2B
49.1
71.5
60.6
60.4
GLA-H
25.6
56.6
48.2
43.5
SSE-H-n4k1
2B
35.6
70.6
50.0
52.1
Full Ablation Row Selection and Parameter Sharing. As discussed in Section 5.5 we ablate softmax row selection and partition parameter sharing on 600M models with 15B tokens. Full results confirming the effectiveness of these design choices.
provided in Tables 12 and 13
Table 12: Full ablation on the effect softmax- based row selection.
Models
Params.
Tokens
FDA
SWDE
SQuAD
Avg.
SSE-n4k1-k.silu
600M
15B
20.87
26.91
24.90
24.23 31.16 (+6.93)
SSE-n441-k.softmn
600M
15B
33.67
33.30
26.51
Table 13: Full ablation on shared parameters. SSE shares attention parameters across partitions, while "w/o shared- shared-params" assigns separate parameters to each partition.
Models
Params. (non-embed)
Tokens
Wiki.]
FDA
SWDE
SQuAD
Recall-Avg.
CSR-Avg.
SSE-n4k1
600M (300M)
15B
35.52
33,67
33.30
26.51
31.16
42.91 42.99 (+0.08)
w/o shared--pand params
890M (580M)
15B
34.94
16.52
36.45
24.87
25.95 (-5.21)
F.3 EXPERIMENTAL RESULTS OF SSE-GDN
We primarily focus SSE integrated with GLA-style transitions in the main experiments. Here, provide a more comprehensive evaluation of integrating the delta-rule mechanism into SSE, denoted SSE-GDN. Overall, SSE-GDN achieves state-of-the-art recall performance, and its hybrid variant (SSE-GDN-H) consistently outperforms GDN-H hybrid architectures.
600M SSE-GDN: Recall Performance. Following same training setting as in Table [ SSE- GDN establishes new state-of-Ue of-tt---ttt-tttt recall benchmarks (Table [4). In addition, write-raad gating design remains effective compared with read- sadd-atc-only variants adopted prior work such MoM.
15B MoE (700M Active) SSE-GDN-H: Scaling Performance. Building upon the SSE-H results based on GLA, we further evaluate SSE-GDN-H. We use 15B MoE model with 700M active pa- rameters, trained on 450B tokens with 1:6 hybrid ratio. The results are shown in Table 15 We observe that: (1) SSE-GDN-H achieves significant gains over both GDN-H and the Transformer baseline MMLU and other general benchmarks: (2) it delivers substantial improvements on RULER, surpassing GDN-H and matching Transformer performance under the n4k2 setting, while maintaining strong extrapolation up 16k tokens (with training context lengths 8k).
These results demonstrate that SSE extends effectively beyond diagonal linear attention and remains well-suited for delta-rule-based architectures.
F.4 EFFICIENCY ANALYSIS
In Figure we compare the speed (forward backward time) of different SSE implementations. The masking implementation competitive for short sequences (5Ik), whereas the varlen imple- mentation becomes more efficient sequence length increases. Here, further expand ef-
27

---

Published as a conference paper at ICLR 2026
Table 14: Recall performance comparison on 600M models trained with 15B tokens.
Models
Params.
Tokens
FDA
SWDE
SQuAD
Avg.
Transformer
600M
15B
74.50
60.67
32.67
55.95
GLA
600M
15B
9.44
23.40
23.06
18.63
Mamba
600M
15B
6.35
21.60
24.93
17.63
Mamba2
600M
15B
20.51
30.42
27.21
26.05
KDA
600M
15B
24.50
35.55
26.71
28.92
GDN
600M
15B
14.79
33.48
26.24
24.84
MoM-n441
800M (A600M)
18.97
36.36
27.75
27.69
MoM-n442
800M (A600M)
15B
21.05
37.08
26.68
28.27
SSE-n4k1
600M
15B
33.67
33.30
26.51
31.16
SSF-GDN-n441 (read gate)
600M
15B
25.77
38.25
28.02
30.68
SSF-GDN-n4k1
600M
15B
41.56
41.13
30.83
37.84
Table 15: SSE-GDN-H performance under the 15B MoE setting.
Model
MMLU
MMLU-Pio
C-Eval
AGiEval
BBH
DROP
GSMSK
TrivieQA
Ruler (-k)
Ruler (8k)
Ruler (16k)
Transfomner
61.5
60.6
64.6
62.3
80.0
67.5
73.0
89.5
GDN-H (1:6)
62.1
62.4
64.6
60.9
77.1
68.5
86.3
72.4
81.0
63.6
38.9
SSE-GDN-H-n4k1 (1:6)
62.1
618
65.3
61.5
79.4
71.1
86.9
71.3
82.0
72.1
46.6
SSE-GDN-H-odk2
ficiency study through two complementary evaluations: (1) attention runtime and (2) end-to-end training throughput.
Attention Runtime (Forward + Backward, ms). We benchmark SSE (with GLA transition).
GLA, and full attention across sequence lengths ranging from 8k to 500k. As shown in Table SSE demonstrates linear scaling behavior and becomes substantially faster than full attention beyond 32k tokens. Compared with GLA, SSE ineurs moderate overhead due router index sorting, Q/K/V/O reordering, increased varlen segmentation (resulting in more kernel-level chunk boundaries), and computing attention over k+1 state partitions per step.
Training Throughput (K tokens/s). We measure training throughput on single A100 GPU. For sequence lengths from 2k to 16k, we use a 1.3B model; for 32k, we use a 340M model due to memory constraints. We additionally evaluate SSE-GDN, which integrates SSE with the GDN architecture. Aggregated results are summarized in Table I7
Overall, SSE surpasses Transformer in throughput at 32k sequence length, while remaining slower than GLA, GDN, and Mamba2 due to the inherent overhead introduced by state expansion, reflect- ing a clear speed- pereddderformance tradeoff. SSE-GDN achieves approximately 70% of GDN throughput (comparable to Mamba), and SSE(-GLA) reaches approximately 60% of GLA throughput. Despite this gap, SSE maintains sufficient efficiency for large-scale long-context pretraining. We consider targeted low-level kernel optimization promising direction for further reducing the remaining over- head.
28

---

Published as a conference paper at ICLR 2026
Table 16: Attention runtime comparison across sequence lengths. Runtime measured as the
time (ms) required for forward and backward pass.
Model / Seq Len
8k
16k
32k
64k
128k
256k
500k
Full Attention (Quadratic)
9
24
83
315
1250
4376
GLA (Linear)
10
18
36
71
131
ASE-n4k1-varles (Linear)
14
97
191
358
Table 17: Training throughput comparison. Throughput is measured in K tokens/s on a single 100 GPU. Models with 1.3B parameters are used for sequence lengths from 2k 16k, while 340M model is used for 32k due to memory constraints.
Model / Seq Len & Batch Size
2k & 8
4k &
8k 2
16k
32k
Transformer
24.24
22.35
19.41
15.25
25.14
GLA
19.55
19.47
19.44
19.23
53.01
GDN
21.37
21.07
20.87
20.48
55.71
Mamba
12.81
12.79
12.48
14.72
35.87
Mamba2
18.32
18.37
18.13
18.29
45.81
MoM-GDN-n4kl
11.12
11.23
11.16
10,20
23.71
SSE(-GLA)-44KI
12.40
12.33
12.29
12.29
28.79
SSF-GDN-n4k1
14.88
14.86
14.73
14.45
35.08
29


Reviews:
[{"review_id": "ZyCXTzc2XM", "reviewer": "Reviewer_dHEC", "summary": "The paper proposes row-selector to bottleneck the update of state and use multiple partitions to expand state size for improved expressiveness. The results show their methods can get comparable performance with Transformers.", "questions": "1. can you provide some efficiency analysis, especially empirical evidence, when compared to baseline linear attention models?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "1. results are strong: SOTA in 2B reasoning model.\n2. preliminaries are well-organized, proofs are completed.\n3. enable multiple efficient parallelized implementations.", "weaknesses": "1. phrasing can be simplified and changed for better delivery: information classification, row-sparse -> row-selector. (TBH information classification is very confusing).\n2. miss an important baseline Mamba/Mamba2", "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": "YFw52mwlxX", "reviewer": "Reviewer_Jg4D", "summary": "The paper proposes Sparse State Expansion (SSE) for linear attention. Two ideas drive the method: (i) row‑sparse updates that treat state updates as a classification problem and write only to top‑k rows via a softmax head, and (ii) state expansion into N shared‑parameter partitions chosen by a write–read gate, so capacity (state size) scales without growing parameter count. Efficient masked/varlen implementations are provided. Empirically, SSE and its hybrid variant are tested on language modeling, retrieval, and math reasoning.", "questions": "See weaknesses.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 2, "strengths": "1. Modeling state updates as information classification is well argued and operationalized \n2. Decoupling the model’s parameters and state's parameters is important direction for improving linear transformers.", "weaknesses": "1. There is no convincing explanation on why SSE is only effective for linear attention but not deltanet.\n2. Baselines seem to be cherry picked since they are neither the most powerful nor fundamentally relevant sequence models. \n3. The paper uses hard partition selection (top‑k) and softmax row selection, but the gradient treatment for the discrete top‑k isn’t described. Can you please clarify how gradients flow through Eqs. 7–9 and whether any implementation tricks are needed for stability.", "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": "2EE5YHbpQy", "reviewer": "Reviewer_heMc", "summary": "This paper improves context compression in linear attention models with row-sparse update and sparse state expansion (SSE). The row-sparse update learns a sparse mask with top-k and softmax to update only a few rows in the contextual states, which utilizes the state space more effectively than dense update. Since different tokens are associated with different rows, the row-sparse update eliminates the need of gating at each step, thereby avoiding the limited receptive field caused by gating. The authors then propose SSE to extend row-sparse update to larger state space. SSE first divides a large state space into N partitions, and then perform dense updates on k partitions. By sharing the attention parameters across partitions, SSE decouples the parameters size and the memory capacity, thereby solving the capacity bottleneck of linear attention models. Experiments on language modeling, needle-in-a-haystack and reasoning benchmarks show that SSE and the hybrid SSE-H achieves state-of-the-art performance among linear attention models.", "questions": "1. Is the cosine similarity computed for a single step then averaged over a sequence or the whole dataset?\n2. The logic in Line 174-178 is hard to understand. It’s hard to understand what Propositions 2-4 are without looking into the Appendix. Propositions 2-4 use the row-sparse update, which is introduced only in later sections. Besides, it’s not very clear how the conclusion of decay in gated variants is derived before looking into Proposition 4. As row-sparse update is complementary to gating and they may co-exist, it’s also hard to understand why SSE solves the issue caused by gating. You may discuss eliminating gating as a benefit of SSE after introducing the method.\n3. Line 184-186: By theoretical analysis, do you mean Proposition 2? It’s a little bit hard to understand this without looking into the appendix. You need to add more details in the main paper.\n4. The functions softmax and top-k produce a $k$ dimension vector by their definition, which is not correct. Do you mean softmax on the $k$ non-zero elements and keep the rest as 0? Then you need to re-define your softmax function.\n5. Line 263-265: This key insight should be brought to early paragraphs of Sec 4.1. Otherwise, it’s hard to understand how SSE is connected with row-sparse update. You’re essentially factorizing a row-sparse update for N*c rows into two small parameter matrices.\n6. Font size in Figure 3 & 4 should be increased.\n7. Line 284 & 286: Why can the always-selected partition capture local interactions? The inputs are accumulated by addition, which doesn’t model interactions.\n8. Line 294-296: Please explain why singular value entropy reflects the difficulty of compression. It’s not a commonsense to audience.\n9. Line 301: Sequential computation is never explained in the main paper.\n10. Line 308-310: I would recommend to draw a figure for grouping and varlen technique.\n11. In Table 1, linear attention models are worse than Transformer. Is it because language modeling requires more pairwise interactions? Then why SSE becomes on par with Transformer in Table 3 & 4?\n12. Figure 5: Which task is it?\n13. Figure 6: Does SSE-Shared refer to the always-selected partition? Please be consistent in the terms.\n14. Line 457: Is the sparsity measured within the selected partitions? For n4k1, I understand the upper bound of sparsity for all partitions should be 25%, right?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 2, "contribution": 4, "strengths": "1. This paper tackles the memory capacity problem, a key issue in linear attention models with clear motivations. The row-wise sparse update aims to improve the utilization of a fixed size state space, while the SSE extends the memory capacity without increasing the number of parameters.\n2. Experiments are thorough and solid. The authors visualizes the cosine similarity of the state space and show that row-sparse update significantly improves state space utilization compared to existing designs. The final SSE model is evaluated against Transformer and linear attention baselines on a wide range of benchmarks, with a rigorous setup of fixed number of parameters.\n3. The efficiency of SSE is not confined to theory. The authors implemented SSE by grouping tokens into subsequences according to their partitions and executing them with a linear attention kernel call. This makes SSE practical for real-world long context use.", "weaknesses": "1. The writing of this paper may be largely improved. The authors used quite a few terminologies or preliminaries without enough explanations, which make this paper hard to follow. For example, the authors didn’t explain how they computed the cosine similarity in Figure 1. Line 176 compares SSE against gated linear attention, but the form of gated linear attention is never mentioned in the paper. There aren’t ground truth classes nor a classification task, but the authors keep using the term classification to refer to state space utilization. In Line 245, the term segmented clustering is used, but this isn’t a commonsense for audience. Captions of figure 1 and 3 need to be extended with their implications. See more in questions.\n2. The title doesn’t exactly reflect the contribution of this paper. “Scaling linear attention\" sounds like this paper studies the model performance under different linear attention sizes. I would suggest modifying it to be “Extending linear attention capacity with sparse state expansion”.\n3. While SSE focuses on improving linear attention models, there isn’t any comparison of wall time for SSE and baselines. Could you please report the performance-time trade-off curve for SSE, Transformer and other linear attention models? That will help audience know which model to use given a specific context length.", "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": "aTSAYPRgRM", "reviewer": "Reviewer_KhC1", "summary": "This paper proposes a new framework to enhance the efficiency and scalability of linear attention models for long-context processing. It introduces two innovations: (1) row-sparse state updates, which treat information storage as a classification task using top-k softmax selection to reduce interference and extend context range, and (2) Sparse State Expansion (SSE), which partitions the state into multiple sparsely updated segments to increase memory capacity without adding parameters. Experiments show that SSE and its hybrid variant (SSE-H) outperform earlier linear attentions across language modeling, retrieval, and reasoning tasks.", "questions": "On which datasets was the model trained? Will these data be open-sourced?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Clear motivation, sound implementation, and good performance with model at scale.", "weaknesses": "1. From the Fig.4 we can see that the wall time is not linear with respect to the sequence length. Why? Hope the authors could provide some more experiment results comparing the wall time of their architecture against earlier linear attentions, GLA, GDN, Mamba2, and full attention.\n2. Hope the authors could also provide scaling experiments compare their architecture against earlier linear attentions, GLA, GDN, Mamba2, and full attention, considering training wall time vs perplexity.\n3. Hope the authors could also discuss MoE for quadratic attention in their related work.", "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)