Sparse Delta Memory: Scaling the State of Linear RNNs through Sparsity

Unknown

Conceptio › Archive › arXiv CS

arXiv CSopen access

Sparse Delta Memory: Scaling the State of Linear RNNs through Sparsity

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

machine learning, deep learning, neural networks

Sparse Delta Memory: Scaling the State of Linear RNNs through Sparsity Loïc Cabannes1,2 , Pierre-Emmanuel Mazaré1 , Gergely Szilvasy1 , Matthijs Douze1 , Maria Lomeli1 , Ilze Amanda Auzina1,3 , Justin Carpentier2 , Gabriel Synnaeve1 , Hervé Jégou1 Meta FAIR, 2 Inria Paris & ENS-PSL University, 3 University of Tübingen

Linear attention models allow a fixed state size and a fixed amount of compute per token. However, due to their limited state size, linear attention models fall behind in long-context recall compared to softmax-attention-based transformer architectures. Increasing the state size of linear attention improves recall performance but at the cost of higher FLOPs. In this work, we introduce Sparse Delta Memory (SDM), an architecture that scales the hidden state of gated linear RNNs to orders of magnitude higher capacity using a sparse addressing scheme. SDM extends the Gated DeltaNet architecture by replacing the dense key-value outer product with sparse reads and writes to a large explicit memory. We show that, under an isoFLOP constraint and with an identical number of parameters, a higher state memory capacity significantly improves performance on in-context learning and long-context retrieval tasks. Moreover, by learning the initial state of the SDM memory and therefore using it as a parametric memory, we show that the model further improves on a wide range of common-knowledge and reasoning tasks. Date: July 9, 2026 Code: https://github.com/facebookresearch/sparse-delta-memory

Introduction

As frontier models continue to progress, they are leveraged in increasingly more complex tasks. In particular, the emergence of agentic settings involve sustained reasoning over long contexts, including software engineering, research assistance, and personal assistants. These applications demand memory mechanisms that preserve long-range dependencies across extended interactions. In the standard transformer architectures equipped with a vanilla softmax attention, the interaction memory is stored in a Key-Value (KV) cache. It effective on long-context tasks, yet the KV cache and therefore the compute and memory per token all grow linearly with the sequence length, (Fig. 1 green line). This unbounded growth limits in-context learning over very long sequences, such as entire codebases or extended reasoning traces, which are increasingly central to autonomous agents, but also videos which play a central role notably in world models and robotics.

1GB

State size

arXiv:2607.07386v1 [cs.LG] 8 Jul 2026

128k large state 100MB SDM (Ours) 32k 8k (KV Cache) 10MB efficient 1MB 100KB

GDN 108 109 FLOPs / token

Figure 1 State size vs FLOPs per token (global layer, 1.4B). KV cache scales linearly with the sequence length. Our SDM approach offers a large state with constant FLOPs.

A possible way to avoid this growing KV cache is to replace the explicit storage of all past tokens with a compressed recurrent state. Recurrent Neural Networks (RNNs), including State Space Models (Gu & Dao, 2024) and Linear Attention variants (Katharopoulos et al., 2020; Beck et al., 2024), rely on this strategy: they compress information into a fixed-size hidden state, maintaining constant memory and compute per token

regardless of sequence length. This enables the processing of arbitrarily long contexts without an explicit token limit. However, their extremely small state sizes limit recall capability compared to transformers (Fu et al., 2023). Indeed, Arora et al. (2025) show that long-context performance is fundamentally bounded by the hidden state size. Simply increasing the RNN memory size would improve recall, but modern linear RNNs like Mamba2 (Dao & Gu, 2024) and Gated DeltaNet (GDN) (Yang et al., 2025) are bottlenecked by dense state updates that become prohibitively expensive as the state size grows. To address the limitations of dense state updates, we introduce Sparse Delta Memory (SDM), a novel architecture based on the observation that the GDN update rule can be sparsified. This enables a three-orderof-magnitude increase in memory state size while maintaining the same compute budget, as shown in Fig. 1. Thanks to its larger state size, SDM significantly outperforms GDN on long-context recall tasks from the RULER (Hsieh et al., 2024) benchmark and also shows better in-context learning capabilities on sequences of up to 1 million tokens. Moreover, we show that contrary to GDN, learning the initial state of SDM allows the model to store meaningful pretraining knowledge that significantly improves performance on a wide range of metrics. Indeed, in isoFLOP comparisons, an SDM with a learned initial state consistently achieves a lower training loss than GDN across all parameter scales. We validate our findings by training 8B activated parameter models on more than 1 trillion tokens. At this scale, SDM reaches an even lower loss and a slightly better short-context accuracy than a model trained with full attention. Overall, we show that SDM, thanks to its large state, can keep the constant-space and memory advantages of Linear RNNs while significantly improving on long-context tasks which have been the main limitation of Linear RNNs like GDN and Mamba so far. Given the constant compute and memory footprint of SDM and its strong long-context performance, we believe that SDM opens up new possibilities in developing agents with improved long-term memory and in-context understanding over extended sequences and also has the potential to address long-term memory issues found in other modalities such as long-video processing.

Background

Linear Attention as an Associative Memory. A fundamental perspective introduced by Katharopoulos et al. (2020) is that attention writes outer products of keys kt ∈ Rdqk and values vt ∈ Rdv into a memory. Then, to read from the memory, a query qt ∈ Rdqk is compared to the previous keys using a similarity metric, usually the inner product ⟨q, k⟩. Moreover, to improve the accuracy of information retrieval, a pre-processing feature mapping ϕ can be applied to the keys and queries. We can then define the memory tensor Mt and normalization factor zt : Mt =

t X

ϕ(ki ) vi⊤ ∈ Rdqk ×dv ,

zt =

i=1

A normalized read is then:

M⊤ ϕ(qt ) yt = ⊤t = zt ϕ(qt )

t X

ϕ(ki ) ∈ Rdqk .

i=1

i≤t P

⟨ϕ(qt ), ϕ(ki )⟩ vi

i≤t ⟨ϕ(qt ), ϕ(ki )⟩

(1)

Note that there exists an infinite-dimensional feature mapping ϕ such that ⟨ϕ(qt ), ϕ(ki )⟩ = e⟨qt ,ki ⟩ , in which case Equation (1) computes exactly softmax attention. On the contrary, if ϕ is a finite-dimensional mapping, then the feature mapped keys and queries as well as the memory tensor Mt are finite-dimensional and can thus be materialized and cached for future retrievals. That is, (Mt , zt ) is stored in constant memory, no matter the sequence length. Gated DeltaNet (GDN). DeltaNet (Schlag et al., 2021) improves upon vanilla linear attention by introducing the delta rule: before writing a new association into the memory, the model first retrieves and subtracts the existing value associated with the key, thereby preventing interference and keeping the memory norm bounded. Gated DeltaNet (GDN) further adds a decay (forget) gate αt ∈ (0, 1) to control how quickly old associations are forgotten. This gives the following state update formula: Mt ← αt Mt−1 + βt kt vt − αt M⊤ t−1 kt 2

⊤

(2)

Mt-1 N ×d

M̃t

+ βt ·kt ⊗∆vt

N ×d

∆vt M̃⊤ t kt

αt ∈ R

kt ∈ R N

ṽt

vt − ṽt

βt ∈ R

vt ∈ Rd

top-W (k1′ ⊗k2′ ) k′ → [k1′ , k2′ ]

qt ∈ R N

gt ∈ Rd

ot ∈ R d

√

gt ⊙LN(yt )

top-R(q1′ ⊗q2′ ) q ′ → [q1′ , q2′ ] q ′ ∈ R2 N

k ′ ∈ R2 N

Wα

yt = M⊤ t qt

Wβ

xt ∈ Rd Figure 2 SDM layer. Gray: operations present in GDN. Purple: operations modified in SDM. Dashed borders indicate sparse operations (W or R out of N slots).

where βt ∈ [0, 1] is a learned input gate modulating the memory update strength. In GDN, kt and qt are dense vectors in Rdqk , and the state Mt ∈ Rdqk ×dv is a dense matrix. The per-token computational cost is O(dqk × dv ). Thus, increasing the size of the state represents a linear increase in FLOPs. Product Key Memory. Product-Key Memories (PKM) (Lample et al., 2019) propose an indexing scheme that allows indexing to k arbitrary indices among N possible memory slots, while scaling the number of computations sublinearly with respect to N . Indeed, PKM can scale to memories of size N × d with N ≈ 106 (Berges√et al., 2024), while previous dense approaches are limited to N ≈ 103 or 104 . Given two sets of scores √ s1 ∈ R N and s2 ∈ R N , one can get N scores, one for each possible memory index, by doing the outer √ √ sum between the two score vectors: s ∈ R N × N = s1 ⊕ s2 . We simply get the topk of the flattened scores s. Since topk (s1 ⊕ s2 ) = topk (topk (s1 ) ⊕ topk (s2 )), we do not need to materialize the entire N scores s but 2 only the scores s1 and s2 requires √ k scores, making the indexing operation very efficient. Since obtaining 2 O( N × d) operations and obtaining the final topk scores requires O(k ) operations, the indexing scheme of √ PKM has a O( N × d + k 2 ) time complexity. It is important to note that in PKM, the memory state is not updated by the context and is learned only through training, analogous to the FFN weights in the transformer architecture. Other Related Works. The work most closely related to ours is the recent Fast-Weight Product Key Memory (Zhao & Jones, 2026), which aims at sparsifying the existing Test-Time-Training methods (Sun et al., 2025; Zhang et al., 2025) using a sparse PKM memory. Their design diverges from ours in their update rule, their hybridization design and ablation choices and other experimental settings. Moreover, their work does not provide results in an iso-FLOPs setting. Finally, their analysis remains limited to small-scale experiments involving models with at most 100M non-embedding parameters. Nonetheless, their design of a sparse online Memory shows promising results, and we concur with their motivation that sparsity enables significant improvements in RNN long-context capabilities.

Method

3.1

Sparse Delta Memory

Our key insight is that the GDN update rule in Eq. (2) can be sparsified: rather than applying the decay and delta update to the entire dense state Mt ∈ Rdqk ×dv , we maintain an explicit memory table Mt ∈ RN ×dv with N slots, and apply gated delta updates only to the W slots selected by sparse keys. Concretely, at each timestep t and for each SDM head, we propose: 1. Sparse Key Selection. Pre-PKM write keys k′ t√and read queries q′ t are projected from the input xt via learned linear projections Wk , W ∈ Rd×2 N . Each projected vector is split into two halves √ √ q √ √ ′ ′ N ′ ′ N (k1,t , k2,t ∈ R for keys; q1,t , q2,t ∈ R for queries), and their outer sum k′1,t ⊕ k′2,t ∈ R N × N yields N scores: one per memory slot. Applying top-W to the write scores and√top-R to the read scores selects the W write indices Itw and R read indices Itr from N possible slots with O( N ×d+W 2 +R2 ) compute, as this can be done without materializing the full score matrix (since topk (s1 ⊕ s2 ) = topk (topk (s1 ) ⊕ topk (s2 ))). 2. Gated Delta Write. For each selected write slot i ∈ Itw : M̃t [i] ←

αt ·Mt−1 [i] |{z}

(3)

(i) βt ·kt · vt − M̃t [i] |{z}

(4)

forget gate

Mt [i] ← M̃t [i] +

input gate

(i)

where αt = exp(−A · softplus(Wa xt + bdt )) is a per-head forget gate, βt = σ(Wb xt ) is the input gate, kt is the sparse key value (writing weight) for slot i, vt = Wv xt is the value vector, and A is a learnable decay parameter. Unselected slots (i ∈ / Itw ) remain unchanged: Mt [i] = Mt−1 [i]. 3. Sparse Read. The memory is then read by a weighted sum of the R selected read slots: X (i) yt = M⊤ qt · Mt [i] t qt =

(5)

i∈Itr

4. Norm, Gating, and Head Mixing. The retrieved memory yt is normalized through RMS-Norm, elementwise gated with g ∈ Rdv , and finally a projection Wo mixes the outputs from all SDM heads to produce the final layer output ot ∈ Rd . (i)

Connection to GDN. When N = dqk , W = R = dqk (all slots selected), and the sparse key values kt form a dense vector, Eq. (4) recovers exactly the GDN update. In this case, the only difference is the lack of 1D convolutions on the qkv vectors, present in GDN but not in SDM. Learned Initial State M0 . Compared to GDN, which has a very small state size, SDM has a much larger state. This property might not serve only as a storage mechanism for in-context knowledge, but also, if one considers M0 as a learnable parameter of the model, the SDM memory can learn knowledge during pretraining and reuse it at test time. Since having a learned M0 does not add any FLOPs at inference time compared to a null-initialized M0 , we chose the learned M0 variant as the default setting for SDM. We ablate the impact of learning M0 in Section 6. Efficient Training. We detail how SDM is trained efficiently using chunk-wise parallelism (via the WY representation from GDN/FLA) and a memory-efficient backward pass in Appendix A.

3.2

IsoFLOP Design: Matching GDN Parameters and FLOPs

We ensure that SDM uses the same number of parameters and FLOPs as the dense GDN baseline. Hence, any improvement stems from the larger memory capacity alone. total

Parameters. Both GDN and SDM share identically-sized linear projections: Wq , Wk ∈ Rd×dqk and Wv ∈ total GDN Rd×dv , where dtotal = d/2 and dtotal = H × dGDN = d. v v qk = H × dqk 4

FLOPs. GDN’s per-token cost is O(H × dGDN × dGDN ) = O(dGDN × dtotal v v ) because every head accesses its qk qk SDM SDM full dqk × dv state. SDM’s cost is instead O(H × (W +R) × dv ) = O((W +R) × dtotal v ), independent of GDN the memory size N . Setting W = R = dqk thus yields matching FLOPs. To be precise, PKM’s top-k on the outersum of scores is an additional computation, but it amounts for less than 1% of the layer’s FLOPs. We set dGDN = 64, dGDN = 128, giving Wq , Wk ∈ Rd×d/2 and Wv ∈ Rd×d for both GDN and SDM. v qk

3.3

Limiting the State Size Expansion

With SDM (and unlike GDN), fewer heads do not incur increased FLOPs but still increasing memory size, which improves performance (Arora et al., 2025). Having a single head under a parameter constraint maximizes the size of the memory. Indeed, the Memory size of SDM scales as such: # " total 2 total (dtotal dtotal dqk dtotal qk ) · dv qk v = . (6) × × Msize = H × 2H 2H H 4H 2 However, the total memory size per layer of SDM would scale as O(d3 ), faster than both GDN’s state growth (O(d) with GDN heads and O(d2 without heads) and even faster than the model’s parameter count (O(d2 )). This would make the sparse memory impractically large at scale. Therefore, H serves as a hyper-parameter with no impact on FLOPs but allowing one to control the state size in SDM.

Experimental Setup

Architecture. All models use a hybrid architecture with Multi-Head Attention (MHA) layers using Sliding Window Attention (SWA) and interleaved global receptive field layers in a 3:1 short:long ratio. To avoid harmful competition between the short and long layers (Cabannes et al., 2025), we choose a small window size of 128 tokens which has been adopted for similar reasons in Team et al. (2026) and in OpenAI et al. (2025). All full-attention layers use grouped-query attention (GQA) with group size 2 (nkv = nh /2) and gated attention output (Qiu et al., 2025). Each block uses a gated MLP (Liu et al., 2021) with SiLU activation (Elfwing et al., 2017) and a hidden dimension made to match the parameters of an equivalent non-gated MLP with a hidden dimension of 4d. Position encoding uses RoPE (Su et al., 2023) with θ = 500,000. SDM Configuration. The SDM layers use W = R = 64 reads and writes, and a number of slots N = (d/4H)2 memory slots (Section 3.2). Read and write activations use softmax-normalization. The forget gate parameter A is initialized uniformly in [0, 16] and the time-step bias bdt is initialized from inv_softplus(U(0.001, 0.1)), both matching GDN/FLA conventions. The number of SDM heads H is set per scale to keep the state-toparameter ratio around 1:1 (see Section 3.3 and Table 1). GDN and Mamba2 Baselines. The GDN and Mamba2 baseline uses dqk = 64, a value dimension vdim = 128, and nh heads, matching the number of attention heads. Training. All models were pretrained on 8192-token sequences of diverse text data. Training follows a Warmup–Stable–Decay (WSD) LR schedule with gradient clipping at 1.0 using the AdamW optimizer (Loshchilov & Hutter, 2019) with β1 = 0.9 and β2 = 0.95. Learning rates were tuned for the transformer baseline and confirmed optimal for GDN via small grid searches. We thus use a uniform LR across all architectures, varying only by model scale (Table 5.1). To maximize recall performance, 1.4B and 8B models undergo a long-context fine-tuning stage on 128k-token sequences using 4B and 16B tokens, respectively. Scaling Ladder. To evaluate the capability of our approach to scale at larger scale, we train a ladder of model sizes, where we adopt a 160 tokens-per-parameter (160TPP) compute budget. The architectures considered for the ladder are FullAttn, GDN and SDM. We do not rely on optimal token budget following the Chinchilla scaling laws (Hoffmann et al., 2022) because most architectures and models currently deployed at scale are trained way beyond their training-optimal compute budget. We thus choose this 160 tokens per non-embedding parameter budget to compare the architectures in a rather inference-optimal setting. Except for the head number H, GDN, SDM and FullAttn share identical hyperparameters at each level. Evaluation. We evaluate models using validation NLL on held-out natural text and coding data. We also evaluate them on a diverse set of reasoning and commonsense tasks as listed in Table 2. 5

Table 1 Scaling ladder configurations. Params include embedding weights but excludes SDM memory state. The State columns correspond to the total state size across all global layers, with St:Param reporting this size as a fraction of the non-embedding parameters.

Level

Layers

Params

Tokens

LR (×10−3 )

State

GDN St:Param

SDM State St:Param

1 2 3 4 5 6 8

768 768 1024 1024 1280 1536 1920

9 11 11 14 14 15 21

257M 280M 407M 465M 658M 847M 1.48B

10.7B 14.9B 24.3B 34.0B 52.8B 81.0B 168.4B

1.78 1.50 1.35 1.24 1.20 0.99 0.87

98k 98k 131k 197k 246k 295k 614k

0.16% 0.12% 0.09% 0.10% 0.07% 0.07% 0.06%

1 1 1 1 1 1 2

57M 57M 134M 201M 393M 679M 553M

94% 68% 93% 99% 119% 150% 56%

3840

8.14B

1.141T

0.50

2.2M

0.03%

7.963B

111%

Results

In the following sections, we report SDM performance compared to GDN, Mamba2 and FullAttn as global layers. First, we examine scaling laws across the scaling ladder after pre-training (Section 5.1). Second, we report the performance on short and long-context tasks of the long-context finetuned models (Sections 5.2, 5.3,5.4). Last, we perform extensive ablations verifying SDM architecture’s utility and memory use (Section 6).

5.1

Power-Law Scaling: SDM Outperforms GDN at All Compute Levels

We start by evaluating SDM’s compute efficiency by measuring how training loss decreases as total training compute (FLOPs) increases. Specifically, we compare SDM against GDN at each level of the scaling ladder (Table 1), keeping both architectures matched in FLOPs and parameters (excluding SDM’s sparse embedding memory). We used the final losses of the levels from 1 to 8 to compute scaling laws and predict the loss at 8B (level13) scale. As shown in Figure 3, SDM outperforms GDN at every level of the scaling ladder. Moreover, SDM provides predictable scaling with a correlation coefficient R2 = 0.999. As predicted by our scaling law, when trained at level-13 (8B parameters) scale, SDM reaches a significantly lower loss than GDN (which slightly underperforms its predicted loss but stays within the 95% confidence interval computed using bootstrapping) and even outperforms the 8B model with FullAttn. These results demonstrate that SDM consistently outperforms an iso-FLOP GDN across all compute levels while remaining competitive or even surpassing FullAttn at larger scales.

5.2

Sparse Memory Improves Short-Context Performance

We evaluate how SDM’s large sparse memory impacts performance on reasoning and knowledge tasks that do not specifically involve large context capabilities, referred to as “short context”. At both 1.4B and 8B scales, SDM achieves lower DCLM NLL and higher average accuracy than GDN, see Table 2. Notably, SDM obtains the lowest DCLM NLL among all models at both scales, outperforming even FullAttn. at 1.4B scale SDM improves over GDN on 13/15 tasks and achieves an average accuracy much higher than GDN and closer to FullAttn. At 8B scale, SDM again improves on most benchmarks and reaches an average accuracy not only higher than GDN but even higher than the model with FullAttn. This underscores the advantage of having a large state containing not only in-context information but also pretraining knowledge in the form of a learned initial state compared to the states of GDN or FullAttn which only store in-context knowledge and not pre-training knowledge.

5.3

Scaling Memory Improves Long-Context Retrieval

RULER results demonstrate that scaling the hidden memory state significantly improves long-context recall. SDM achieves the highest overall scores among fixed-state models at both 1.4B and 8B scales (31.2 and 50.2 respectively), outperforming GDN (20.0 and 34.2) by a wide margin (Table 2). Indeed, at both scale SDM 6

3.0

SWA:GDN =0.0652, R²=0.997 SWA:SDM =0.0628, R²=0.999 SWA:FullAttn =0.0605, R²=1.000

Training Loss (avg last 10)

2.8 2.6 2.4 2.2 2.0 1.8 1.6

1.80 1.78 1.76 1.74 1.72 1.70 1.68 1023

1019

1020

1021

Total Compute (FLOPs)

1022

1023

Figure 3 Scaling laws: loss vs FLOPs. Average of last 10 log entries vs total compute. SDM consistently outperforms GDN across all scales. As predicted by the scaling law, SDM outperforms full attention at 8B scale.

improves (or matches in the case of the single1 where GDN is already at 100% accuracy) on 6 out of the 6 RULER tasks we evaluated the models on. At 8B, FullAttn achieves 76.2 accuracy overall thanks to its unbounded KV cache. However, SDM actually matches or exceeds FullAttn on 4 of 6 tasks at 1.4B and 3 of 6 at 8B, despite using a fixed memory representation. On multikey 2 however, FullAttn maintains a large advantage over both SDM and GDN. Detailed results on the obtained performance gap at every sequence length scale are reported in Figures 9 and 10. Overall, the RULER results provide strong evidence that the larger state size enabled by the SDM architecture significantly improves performance on long-context recall.

5.4

SDM Reaches Lower Perplexity with More Context

SDM achieves consistently lower perplexity on code data compared to Mamba2 and GDN (Fig. 4). Even at short sequences (512 tokens), SDM outperforms both baselines, a benefit we attribute to the learned initial state M0 . The advantage grows substantially at long contexts (32k–512k tokens), where SDM’s perplexity decreases to near 2.0 while Mamba2 and GDN remain around 2.2–2.3. On this evaluation data, the validation perplexity of all models seems to increase for token positions beyond 256k. However, this is an artifact of the local, token-level perplexity being higher at those positions. When we measure only the perplexity gain contributed by the long-context layers, SDM continues to improve beyond 256k tokens, up to 1 million tokens. This highlights SDM’s key strength: its large sparse memory retains information across extremely long sequences, whereas fixed-size state methods suffer from capacity constraints.

Ablations: What Makes SDM Work?

Disentangling Memory Capacity from Learned Initialization. To verify that SDM’s gains stem from increased memory capacity rather than the learned initial state M0 , we ablate both components. Figure 5 shows that SDM without a learned M0 substantially outperforms GDN, confirming that state size is the primary driver of performance. Adding a learned M0 to a vanilla GDN does not measurably improve performance, which is not surprising considering the limited state size. For detailed results, see Table 3. Overall, the ablation confirms that the in-context learning gains come from the increased state size more than from the learned initial state. Impact of Memory State Size. To confirm that memory size drives SDM’s long-context advantage, we ablate the number of memory slots N by varying the PKM key dimension dqk (which controls N = (dqk /2)2 ) 7

Table 2 Performance of different global layer choices at 1.4B (L08) and 8B (L13) scale. All models use SWA as the local layer. Delta column shows delta in performance between SDM and the isoFLOP GDN. Top: Validation NLL, reasoning and common-knowledge tasks. Bottom: exact-match accuracy averaged across lengths 4k–131k on all 13 RULER long-context tasks.

Local layer → Global layer →

FullAttn

1.4B (L08) SWA Mamba2 GDN

SDM

FullAttn

8B (L13) SWA GDN

SDM

Validation text NLL ↓

2.660

2.686

2.683

2.654 -0.029

2.285

2.298

2.253 -0.040

HellaSWAG (2019) ↑ WinoGrande (2019) ↑ ARC easy (2018) ↑ ARC challenge (2018) ↑ PIQA (2019) ↑ OpenBookQA (2018) ↑ RACE.mid (2017) ↑ RACE.high (2017) ↑ CommonsenseQA (2019) ↑ BoolQ (2019) ↑ TQA (2017) ↑ HumanEval+/pass@1 (2023) ↑ NaturalQuestions (2019) ↑ MMLU (2021) ↑ GSM8K (2021) ↑

61.47 60.54 61.56 33.56 73.83 35.80 53.55 39.62 19.57 62.42 24.24 8.54 7.70 24.54 2.81

61.31 62.67 62.88 34.08 73.67 36.40 50.28 36.25 20.31 62.75 23.48 9.76 8.09 25.03 3.03

61.27 61.56 62.49 34.59 73.83 36.80 50.42 36.68 19.33 63.03 23.68 8.54 7.76 26.02 2.96

62.04 +0.77 62.12 +0.55 63.00 +0.51 34.76 +0.17 74.27 +0.44 36.00 -0.80 49.65 -0.77 37.76 +1.09 20.72 +1.39 62.78 -0.24 26.61 +2.93 9.76 +1.22 8.23 +0.47 27.24 +1.22 2.73 -0.23

79.33 73.64 78.10 50.82 79.98 43.80 64.83 47.94 68.88 71.83 55.23 24.39 22.94 58.73 29.34

79.10 73.24 79.15 52.27 81.01 43.00 60.17 43.40 66.83 74.98 55.36 18.29 22.60 57.24 28.81

80.02 +0.92 75.30 +2.06 78.52 -0.63 53.05 +0.78 80.47 -0.54 44.60 +1.60 62.05 +1.88 44.97 +1.57 70.60 +3.77 67.13 -7.85 59.11 +3.75 24.39 +6.10 25.93 +3.33 57.81 +0.57 28.66 -0.15

Average accuracy ↑

37.98

38.00

37.93

38.51 +0.58

56.65

55.70

56.84 +1.14

single_1 avg ↑ single_2 avg ↑ single_3 avg ↑ multikey_1 avg ↑ multikey_2 avg ↑ multikey_3 avg ↑ multivalue avg ↑ multiquery avg ↑ vt avg ↑ cwe avg ↑ fwe avg ↑ qa_1 avg ↑ qa_2 avg ↑

64.2 53.1 41.1 44.0 25.3 8.6 39.8 39.0 32.1 5.8 30.9 19.3 20.9

53.8 19.2 8.0 14.7 1.0 0.2 17.5 10.5 8.6 3.6 37.5 11.5 18.5

99.9 20.7 12.1 13.6 0.7 0.1 11.7 11.8 16.6 6.6 35.6 13.3 18.3

100.0 +0.1 70.8 +50.1 46.3 +34.2 35.0 +21.4 0.8 +0.1 0.2 +0.1 37.4 +25.7 41.1 +29.3 23.7 +7.1 8.7 +2.1 12.3 -23.2 13.3 +0.0 18.0 -0.3

99.3 89.4 70.1 75.3 58.0 29.8 74.3 73.0 67.2 5.5 77.8 39.0 38.4

100.0 45.1 32.6 32.9 1.0 0.1 31.1 31.2 46.2 11.4 62.7 23.6 28.2

100.0 +0.0 71.5 +26.4 74.9 +42.3 59.3 +26.4 4.7 +3.7 1.2 +1.1 66.1 +35.0 68.6 +37.4 72.3 +26.1 12.8 +1.5 65.0 +2.4 27.2 +3.6 30.4 +2.3

Average accuracy on RULER ↑

32.5

17.9

20.0

31.2 +11.2

61.2

34.2

50.2 +16.0

while keeping W , R, dv , and all non-SDM layers fixed. Table 3 (bottom) shows monotonic NLL degradation as memory shrinks from 432 MB to 27 MB (0.914 → 0.947), confirming that larger memory state improves modeling. All SDM variants outperform GDN on long-context recall and show monotonic improvement with larger state sizes. Training Efficiency. A common limitation of sparse approaches is that despite matching dense models in FLOPs, their larger memory footprints incur slower memory accesses in practice. GDN’s compact state fits in fast GPU SRAM, while SDM’s state must remain in HBM, which has a 10× lower bandwidth than SRAM. This is why the current SDM kernel MFU (Model FLOPs Utilization) is around an order of magnitude lower than the highly optimized GDN kernel from the FLA library (Yang & Zhang, 2024). We expect that many gains remain to be made in developing more efficient SDM kernels that reach an MFU closer to that of GDN. Despite this, thanks to the hybrid attention architecture design, SDM’s end-to-end training overhead is modest: at 8B scale, the training throughput of SDM was 1.49x slower than the GDN model. This is remarkable given SDM’s state is ∼4,000× larger at that scale. For the 1.4B SDM, inference decode is around

Mamba2 GDN SDM

3.0

2.8

2.6

Perplexity

2.8

2.4

2.6 2.4

2.2

2.0 512

32k

128k 512k

Token Position

State/layer Code NLL ↓ Acc↑ RULER↑

rned : lea .4B) n o i 1 t abla l state ( a i t i in

GDN GDN SDM SDM

null learned null learned

0.2 MB 0.5 MB 211 MB 211 MB

0.849 0.850 0.845 0.822

37.9 38.0 37.3 38.5

20.0 20.6 28.0 31.2

ory

GDN SDM SDM SDM

null learned learned learned

0.2 MB 27 MB 108 MB 432 MB

0.963 0.947 0.937 0.914

33.8 33.7 33.6 34.6

16.0 20.2 20.7 21.5

m : me tion abla 0.8B) ( size

32k

128k 512k

Token Position

Figure 5 Learned initial state ablation on code data (1M token documents, 1.4B model size, posttrained). PPL by token position. Learning M0 benefits SDM but not GDN.

Figure 4 Perplexity by token position on code data (1M token documents, 1.4B model size). Solid: 128k long-context finetuned, dashed: pre-trained. Local layers are SWA for all models.

Model M0

GDN (null M0) GDN (learned M0) SDM (null M0) SDM (learned M0)

3.0

Table 3 SDM ablations. Top: learned initial state ablation (L08, 1.4B). Making the memory M0 a learnable parameter improves code-data NLL, accuracy (%), and RULER recall (%). Bottom: memory size ablation (L06). SDM degrades gracefully when memory is reduced; even at 27 MB it outperforms GDN on RULER.

10% slower than GDN but 6 times faster than FullAttn. Adaptive Memory Access Patterns. The memory slot utilization is controlled by the softmax over read and write keys. We analyze the peakiness of these distributions to assess whether all keys are actively used. We pre-train SWA:SDM 1.4B models with varying read/write configurations (W ∈ {32, 64}, R ∈ {32, 64, 128}) and extract softmax values per token across all layers and Pmheads (for details see Appendix D.1). First, we compute cumulative mass in the top-m keys: C(m) = i=1 p(i) where p(1) ≥ p(2) ≥ · · · ≥ p(k) . Figure 6 shows that write distributions are moderately peaked: with k = 64 writes, the top 32 keys capture ∼85% of probability mass. Read distributions are more uniform: with k = 128 reads, the top 64 keys hold only ∼77% of mass, indicating broader access patterns. Second, we compute the effective number of keys per token as u(p) = exp(H(p))/k ∈ (0, 1], which measures utilization independent of k. Figure 7 reveals distinct read/write behaviors: reads are consistently more uniform than writes across all configurations. Increasing reads to 128 makes the distribution less uniform (more selective), while decreasing writes to 32 makes writes more uniform. Notably, read and write distributions adapt to each other: with limited reads (R = 32), writes become more uniform to compensate; with abundant reads (R = 128), writes become more peaked. This demonstrates SDM’s adaptive memory access, the model dynamically balances read selectivity against write dispersion based on available capacity. For a more detailed performance comparison across runs, we refer to reader to Table 4 in Appendix D .

W64_R32 (k=64) W64_R64 (k=64) W64_R128 (k=64) W32_R32 (k=32) W32_R64 (k=32) W32_R128 (k=32)

0.4 0.2 0.0 100

101

m (rank in sorted softmax)

102

0.8 0.6

W64_R32 (k=32) W64_R64 (k=64) W64_R128 (k=128) W32_R32 (k=32) W32_R64 (k=64) W32_R128 (k=128)

0.4 0.2 0.0 100

(a) Write keys

m=128

0.95 mass

m=64

1.0

m=32

cum mass in top-m read keys

0.6

m=128

0.8

m=64

0.95 mass

m=32

cum mass in top-m write keys

1.0

101

m (rank in sorted softmax)

102

(b) Read keys

Figure 6 Coverage of the top-m selected memory keys. In each panel, the solid line is the layer-averaged mean cumulative mass curve across all tokens; the shaded band shows the corresponding [p10, p90] range. Panel (a) shows write keys and panel (b) shows read keys. Write keys exhibit less uniform distribution than read. uniform-over-k

0.8 0.6 0.4 0.2 0.0

uniform-over-k

1.0

eff_keys / k (read)

eff_keys / k (write)

1.0

0.8 0.6 0.4 0.2 0.0

W64_R32 W64_R64 W64_R128 W32_R32 W32_R64 W32_R128 (k=64) (k=64) (k=64) (k=32) (k=32) (k=32) (a) Write keys

W64_R32 W64_R64 W64_R128 W32_R32 W32_R64 W32_R128 (k=32) (k=64) (k=128) (k=32) (k=64) (k=128) (b) Read keys

Figure 7 Per-token effective-key utilization: analysis of eff _keys/k. For each ablation, each violin is built from all pooled token-level values for each run; the overlaid markers indicate the mean and median. Panel (a) shows write keys and panel (b) shows read keys. Read keys have a more uniform distribution than write keys, while the write key distribution varies more with the assigned read and write key budget, indicating that SDM memory access adapts to the given key constraints.

Conclusion

In this work, we introduced Sparse Delta Memory, a sparse extension of the Gated DeltaNet architecture based on a Product-Key Memory sparse design. This sparsity offers state sizes thousands of times larger than GDN while keeping the FLOPs identical. Thanks to this much larger state size, SDM reaches much better training loss and NLL than GDN and even than Full Attention. It also demonstrates much better long-context performance as measured across a wide range of long-context tasks from the RULER benchmark. Finally, we provide ablations showing which gains are attributed to the larger state size and which are obtained from learning the initial state of the large memory. Limitations. Although our implementation of SDM gives a training speed allowing us to scale to 8B models, more research is needed to design more efficient kernels to further scale the SDM models. The main limitation of SDM is also its strength: SDM memory requirements are not negligible, as the memory footprint may be as large as the model parameters, which is not adapted to certain resource-constrained contexts. However, the KV cache memory usage of FullAttn models is significant when processing sequence lengths of hundreds of thousands of tokens. More precisely, the SDM state for the 8B model occupies as much memory as 203400 tokens in the KV cache of the 8B FullAttn model.

References Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. Simple linear attention language models balance the recall-throughput tradeoff, 2025. URL https://arxiv.org/abs/2402.18668. Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory, 2024. URL https://arxiv.org/abs/2405.04517. Vincent-Pierre Berges, Barlas Oğuz, Daniel Haziza, Wen tau Yih, Luke Zettlemoyer, and Gargi Ghosh. Memory layers at scale, 2024. URL https://arxiv.org/abs/2412.09764. Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language, 2019. URL https://arxiv.org/abs/1911.11641. Loïc Cabannes, Maximilian Beck, Gergely Szilvasy, Matthijs Douze, Maria Lomeli, Jade Copet, Pierre-Emmanuel Mazaré, Gabriel Synnaeve, and Hervé Jégou. Short window attention enables long-term memorization, 2025. URL https://arxiv.org/abs/2509.24552. Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions, 2019. URL https://arxiv.org/abs/1905.10044. Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge, 2018. URL https://arxiv.org/abs/1803.05457. Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168, 2021. Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. In International Conference on Machine Learning (ICML), 2024. Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning, 2017. URL https://arxiv.org/abs/1702.03118. 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, 2023. URL https://arxiv.org/abs/2212.14052. Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, 2024. URL https://arxiv.org/ abs/2312.00752. Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. Proceedings of the International Conference on Learning Representations (ICLR), 2021. Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training compute-optimal large language models, 2022. URL https://arxiv.org/abs/2203.15556. Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What’s the real context size of your long-context language models?, 2024. URL https://arxiv.org/ abs/2404.06654. Mandar Joshi, Eunsol Choi, Daniel Weld, and Luke Zettlemoyer. TriviaQA: A large scale distantly supervised challenge dataset for reading comprehension. In Regina Barzilay and Min-Yen Kan (eds.), Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1601–1611, Vancouver, Canada, July 2017. Association for Computational Linguistics. doi: 10.18653/v1/P17-1147. URL https://aclanthology.org/P17-1147/. Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast autoregressive transformers with linear attention, 2020. URL https://arxiv.org/abs/2006.16236. Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, Kristina Toutanova, Llion Jones, Matthew Kelcey, Ming-Wei Chang, Andrew M. Dai, Jakob Uszkoreit, Quoc Le, and Slav Petrov. Natural questions: A benchmark for question answering

research. Transactions of the Association for Computational Linguistics, 7:452–466, 2019. doi: 10.1162/tacl_a_00276. URL https://aclanthology.org/Q19-1026/. Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard Hovy. Race: Large-scale reading comprehension dataset from examinations, 2017. URL https://arxiv.org/abs/1704.04683. Guillaume Lample, Alexandre Sablayrolles, Marc’Aurelio Ranzato, Ludovic Denoyer, and Hervé Jégou. Large memory layers with product keys, 2019. URL https://arxiv.org/abs/1907.05242. Hanxiao Liu, Zihang Dai, David R. So, and Quoc V. Le. Pay attention to MLPs, 2021. URL https://arxiv.org/abs/2105. 08050. Jiawei Liu, Chunqiu Steven Xia, Yuyao Wang, and Lingming Zhang. Is your code generated by chatGPT really correct? rigorous evaluation of large language models for code generation. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https://openreview.net/forum?id=1qvx610Cu7. Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, 2019. URL https://arxiv.org/abs/1711.05101. Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. In EMNLP, 2018. OpenAI, Sandhini Agarwal, Lama Ahmad, Jason Ai, Sam Altman, Andy Applebaum, Edwin Arbus, Rahul K. Arora, Yu Bai, Bowen Baker, Haiming Bao, Boaz Barak, Ally Bennett, Tyler Bertao, Nivedita Brett, Eugene Brevdo, Greg Brockman, Sebastien Bubeck, Che Chang, Kai Chen, Mark Chen, Enoch Cheung, Aidan Clark, Dan Cook, Marat Dukhan, Casey Dvorak, Kevin Fives, Vlad Fomenko, Timur Garipov, Kristian Georgiev, Mia Glaese, Tarun Gogineni, Adam Goucher, Lukas Gross, Katia Gil Guzman, John Hallman, Jackie Hehir, Johannes Heidecke, Alec Helyar, Haitang Hu, Romain Huet, Jacob Huh, Saachi Jain, Zach Johnson, Chris Koch, Irina Kofman, Dominik Kundel, Jason Kwon, Volodymyr Kyrylov, Elaine Ya Le, Guillaume Leclerc, James Park Lennon, Scott Lessans, Mario LezcanoCasado, Yuanzhi Li, Zhuohan Li, Ji Lin, Jordan Liss, Lily, Liu, Jiancheng Liu, Kevin Lu, Chris Lu, Zoran Martinovic, Lindsay McCallum, Josh McGrath, Scott McKinney, Aidan McLaughlin, Song Mei, Steve Mostovoy, Tong Mu, Gideon Myles, Alexander Neitz, Alex Nichol, Jakub Pachocki, Alex Paino, Dana Palmie, Ashley Pantuliano, Giambattista Parascandolo, Jongsoo Park, Leher Pathak, Carolina Paz, Ludovic Peran, Dmitry Pimenov, Michelle Pokrass, Elizabeth Proehl, Huida Qiu, Gaby Raila, Filippo Raso, Hongyu Ren, Kimmy Richardson, David Robinson, Bob Rotsted, Hadi Salman, Suvansh Sanjeev, Max Schwarzer, D. Sculley, Harshit Sikchi, Kendal Simon, Karan Singhal, Yang Song, Dane Stuckey, Zhiqing Sun, Philippe Tillet, Sam Toizer, Foivos Tsimpourlas, Nikhil Vyas, Eric Wallace, Xin Wang, Miles Wang, Olivia Watkins, Kevin Weil, Amy Wendling, Kevin Whinnery, Cedric Whitney, Hannah Wong, Lin Yang, Yu Yang, Michihiro Yasunaga, Kristen Ying, Wojciech Zaremba, Wenting Zhan, Cyril Zhang, Brian Zhang, Eddie Zhang, and Shengjia Zhao. gpt-oss-120b & gpt-oss-20b model card, 2025. URL https://arxiv.org/abs/2508.10925. Zihan Qiu, Zekun Wang, Bo Zheng, Zeyu Huang, Kaiyue Wen, Songlin Yang, Rui Men, Le Yu, Fei Huang, Suozhi Huang, Dayiheng Liu, Jingren Zhou, and Junyang Lin. Gated attention for large language models: Non-linearity, sparsity, and attention-sink-free, 2025. URL https://arxiv.org/abs/2505.06708. Jack W Rae, Jonathan J Hunt, Tim Harley, Ivo Danihelka, Andrew Senior, Greg Wayne, Alex Graves, and Timothy P Lillicrap. Scaling memory-augmented neural networks with sparse reads and writes, 2016. URL https://arxiv.org/ abs/1610.09027. Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale, 2019. URL https://arxiv.org/abs/1907.10641. Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers, 2021. URL https://arxiv.org/abs/2102.11174. Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, 2023. URL https://arxiv.org/abs/2104.09864. Yu Sun, Xinhao Li, Karan Dalal, Jiarui Xu, Arjun Vikram, Genghan Zhang, Yann Dubois, Xinlei Chen, Xiaolong Wang, Sanmi Koyejo, Tatsunori Hashimoto, and Carlos Guestrin. Learning to (learn at test time): RNNs with expressive hidden states, 2025. URL https://arxiv.org/abs/2407.04620. Alon Talmor, Jonathan Herzig, Nicholas Lourie, and Jonathan Berant. Commonsenseqa: A question answering challenge targeting commonsense knowledge, 2019. URL https://arxiv.org/abs/1811.00937. Core Team, Bangjun Xiao, Bingquan Xia, Bo Yang, Bofei Gao, Bowen Shen, Chen Zhang, Chenhong He, Chiheng Lou, Fuli Luo, Gang Wang, Gang Xie, Hailin Zhang, Hanglong Lv, Hanyu Li, Heyu Chen, Hongshen Xu, Houbin Zhang, Huaqiu Liu, Jiangshan Duo, Jianyu Wei, Jiebao Xiao, Jinhao Dong, Jun Shi, Junhao Hu, Kainan Bao, Kang Zhou, Lei Li,

Liang Zhao, Linghao Zhang, Peidian Li, Qianli Chen, Shaohui Liu, Shihua Yu, Shijie Cao, Shimao Chen, Shouqiu Yu, Shuo Liu, Tianling Zhou, Weijiang Su, Weikun Wang, Wenhan Ma, Xiangwei Deng, Bohan Mao, Bowen Ye, Can Cai, Chenghua Wang, Chengxuan Zhu, Chong Ma, Chun Chen, Chunan Li, Dawei Zhu, Deshan Xiao, Dong Zhang, Duo Zhang, Fangyue Liu, Feiyu Yang, Fengyuan Shi, Guoan Wang, Hao Tian, Hao Wu, Heng Qu, Hongfei Yi, Hongxu An, Hongyi Guan, Xing Zhang, Yifan Song, Yihan Yan, Yihao Zhao, Yingchun Lai, Yizhao Gao, Yu Cheng, Yuanyuan Tian, Yudong Wang, Zhen Tang, Zhengju Tang, Zhengtao Wen, Zhichao Song, Zhixian Zheng, Zihan Jiang, Jian Wen, Jiarui Sun, Jiawei Li, Jinlong Xue, Jun Xia, Kai Fang, Menghang Zhu, Nuo Chen, Qian Tu, Qihao Zhang, Qiying Wang, Rang Li, Rui Ma, Shaolei Zhang, Shengfan Wang, Shicheng Li, Shuhao Gu, Shuhuai Ren, Sirui Deng, Tao Guo, Tianyang Lu, Weiji Zhuang, Weikang Zhang, Weimin Xiong, Wenshan Huang, Wenyu Yang, Xin Zhang, Xing Yong, Xu Wang, Xueyang Xie, Yilin Jiang, Yixin Yang, Yongzhe He, Yu Tu, Yuanliang Dong, Yuchen Liu, Yue Ma, Yue Yu, Yuxing Xiang, Zhaojun Huang, Zhenru Lin, Zhipeng Xu, Zhiyang Chen, Zhonghua Deng, Zihan Zhang, and Zihao Yue. Mimo-v2-flash technical report, 2026. URL https://arxiv.org/abs/2601.02780. Songlin Yang and Yu Zhang. Fla: A triton-based library for hardware-efficient implementations of linear attention mechanism, January 2024. URL https://github.com/fla-org/flash-linear-attention. Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transformers with the delta rule over sequence length, 2024. URL https://arxiv.org/abs/2406.06484. Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule, 2025. URL https://arxiv.org/abs/2412.06464. Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence?, 2019. URL https://arxiv.org/abs/1905.07830. 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, 2025. URL https://arxiv.org/abs/2505.23884. Tianyu Zhao and Llion Jones. Fast-weight product key memory, 2026. URL https://arxiv.org/abs/2601.00671.

Appendices A

Efficient Training of SDM

Training SDM requires computing exact outputs and gradients for the gated delta rule (Eq. 4–5) across long sequences. We decompose this into intra-chunk parallel computation (batched across all chunks) and chunkwise recurrent computation (sequential across chunks), following the WY representation approach from GDN (Yang et al., 2024).

A.1

Intra-Chunk Parallel vs. Chunkwise Recurrent

We split the sequence into chunks of size C. Within each chunk, the gated delta rule creates causal dependencies between tokens—token t’s memory read depends on all prior writes within the same chunk. The WY representation (Yang et al., 2024) resolves these dependencies analytically via a triangular solve, enabling parallel computation within each chunk. Phase 1: Intra-chunk parallel (batched). All chunks are processed simultaneously. Let V ∈ RC×d denote the stacked value vectors for a chunk, and β ∈ RC the per-token input gates. For each chunk, we compute: • Segmented cumulative decay. For each slot i accessed in the chunk, accumulate the log-decay log αt across tokens that write to slot i. This is a segmented prefix sum over slot indices, implemented via scatter-add in log space: for each entry (t, w) with slot index i,P we atomically add log αt to an accumulator indexed by i, yielding the cumulative log-decay λt,w = t′ ≤t, I w′ ∋i log αt′ . t

• Sparse interaction matrices A and QK. We define two C × C lower-triangular matrices via the gated sparse inner product (Eq. 7): (i) A[i, j], the write–write self-interaction using write keys on both sides, capturing how token j’s write to shared slots affects token i’s delta-rule subtraction; and (ii) QK[i, j], the read–write cross-interaction using read queries for token i and write keys for token j, capturing how token j’s write affects token i’s read output. Both are extremely sparse since two tokens interact only if their top-W slot indices overlap (W ≪ N ). We compute them via a sparse inner product kernel (Section A.2). • Triangular solve. Construct the lower-triangular system Msys [i, j] = δij + βi · A[i, j] for j < i, and solve Msys ∆Vconst = diag(β) · V via solve_triangular. This yields the intra-chunk delta-v values assuming reads from the chunk-initial memory state. We also derive the correction matrix B = M−1 sys · diag(−β), which accounts for intra-chunk memory modifications. All operations above are batched over chunks (one bmm/solve call for all chunks simultaneously), achieving high GPU utilization. Phase 2: Chunkwise recurrent (sequential). Chunks are processed one at a time, since each chunk’s memory writes depend on the previous chunk’s memory state. Per chunk: P (n) 1. Read: Gather memory (from HBM) at write indices to obtain retrieved[t] = n keff,t · M[Itw [n]] (where keff = kval · eλ incorporates cumulative decay), and at read indices to obtain the inter-chunk output P (n) inter[t] = n qeff,t · M[Itr [n]]. 2. Correct: Compute the full delta-v incorporating the current memory state: δv = ∆Vconst +B·retrieved. The correction term B · retrieved accounts for the fact that earlier tokens in the chunk have already modified the memory that later tokens read from. 3. Write: Apply per-slot decay and scatter the delta-v contributions back to memory. 4. Output: Compute the final chunk output combining inter-chunk reads with intra-chunk corrections: y = inter + tril(QK) · δv.

The sequential Phase 2 is memory-bound: it is dominated by random gathers and scatters to the N × d memory table. Phase 1 is compute-bound and runs once for all chunks. This decomposition keeps the total compute at O(T · (W 2 + W · d)) per layer, while the sequential bottleneck scales as O((T /C) · W · d) random memory accesses.

A.2

Sparse Inner Product via Two-Pointer Merge

The interaction matrices A[i, j] and QK[i, j] measure how much tokens i and j interact through shared memory slots. For GDN with dense keys, these are standard matrix products (A = KK⊤ ). For SDM with sparse keys, most entries are zero—two tokens interact only if their top-W slot indices overlap. We compute the gated sparse inner product: X A[i, j] =

(n)

(m)

· kj

· exp(λi,n − λj,m ) ,

j<i

(7)

n,m : Iiw [n]=Ijw [m]

where the sum ranges only over matching slot indices between tokens i and j, and λ are the cumulative log-decays. Two-pointer algorithm. The PKM top-W selection includes sort-by-index step (O(W log W ), √an additional 2 which does not change the overall PKM complexity of O( N · d + W )). The resulting indices are thus already sorted by slot index, enabling the classical two-pointer merge to find matching slots between two tokens in O(W ) time rather than O(W 2 ): 1: Input: Sorted indices Ii [0..W −1], Ij [0..W −1]; values ki , kj ; log-decays λi , λj 2: a, b ← 0, 0; result ← 0 3: while a < W and b < W do 4: if Ii [a] = Ij [b] then (a) (b) (a) (b) 5: result += ki · kj · exp(λi − λj ) 6: a += 1; b += 1 7: else if Ii [a] < Ij [b] then 8: a += 1 9: else 10: b += 1 11: return result This is implemented as a CUDA kernel with one thread per (i, j) pair in the C × C interaction matrix, yielding O(C 2 · W ) total work. Only the causal entries (j < i for A, j ≤ i for QK) are computed; non-causal entries are zero. Compared to a dense inner product (O(C 2 · dqk ) for GDN), the sparse version scales as O(C 2 · W ) where W = 64 ≪ N , making it significantly cheaper than addressing the full memory and enabling SDM to scale to large memory tables without increasing the interaction cost.

A.3

Memory-Efficient Backward

Following Sparse Access Memory (Rae et al., 2016), we apply sparse updates to the memory in-place during the forward pass and undo these sparse updates during the backward pass to recover the correct memory state for each timestep. This avoids storing a full copy of the N × d memory at each chunk boundary, reducing peak memory from O((T /C) × N × d) to O(N × d + T × W × d). for example at 16384 token training sequence length, and for an 8B training, chunk size C = 128 W = 64, this reduces memory consumption for the memory state checkpoints of an SDM layer from 226 GB to 8GB. Nonetheless, for longer sequence length such as 128k long-context finetuning, this memory snapshot can become prohibitively large. We therefore investigated quantization of the snapshot to fp8 and int4 and found that both quantization levels seem to have no detrimental on training loss or end performance and therefore recommend it for long-context finetuning.

SDM Memory Utilization During Training

Figure 8 shows the evolution of key memory access statistics during training of the SDM model at 1.4B scale (L08). All metrics are averaged across the 5 SDM layers.

Forget Gate

1.0

0.8

0.6

sigmoid(b)

exp(g)

0.8 mean min max

0.4 0.2

mean min max

0.4

0.0 0

100

Tokens (B)

125

150

175

Slot Access Entropy

1.0 0.8

0.6

0.4 0.2

100

Tokens (B)

125

150

175

Memory Slot Utilization Read unique % Write unique %

40 20

Read entropy Write entropy

100

% of slots

Normalized Entropy

0.6

0.2

0.0

Input Gate

1.0

100

Tokens (B)

125

150

175

100

Tokens (B)

125

150

175

Figure 8 SDM memory utilization during training (1.4B, L08). Statistics are accumulated over 10 training steps (∼21M tokens) per data point and averaged across 5 SDM layers. Top: forget and input gate statistics (min/mean/max). Bottom left: normalized entropy of slot access distributions. Bottom right: percentage of unique memory slots accessed. The forget gate (top left) converges early to exp(g) ≈ 0.95. The minimum forget gate stays near zero, indicating that some slots are fully decayed when overwritten. The input gate (top right) stabilizes at σ(b) ≈ 0.50, with the maximum and minimum reaching ∼1 and ∼0 respectively. Overall, the model learns a wide dynamic range for both forget and input strength, showing the importance of both mechanisms in the model. The read utilization (bottom right) quickly reaches 100% of the memory, showing that no memory slot is left unutilized. The write utilization (bottom right) grows steadily from 28% to 42% during training, suggesting that the model learns to diversify its write patterns over time.

RULER Per-Task Accuracy by Sequence Length

Figure 9 and Figure 10 show the per-task RULER accuracy at each sequence length for different architectures at 1.4B scale and 8B scale respectively. All 13 RULER tasks are shown. Solid lines denote post-trained models (128k finetuned); dashed lines denote pre-trained models.

RULER per task niah_single_1

niah_single_2

1.4B (solid=post, dashed=pre) niah_single_3

niah_multikey_1

niah_multikey_2

100

FullAttn (post) FullAttn (pre) GDN (post)

GDN (pre) SDM (post) SDM (pre)

16k

32k

niah_multikey_3

64k

128k

16k

32k

niah_multivalue

64k

128k

16k

32k

niah_multiquery

64k

128k

16k

32k

64k

128k

100

16k

fwe

32k

64k

128k

16k

qa_1

32k

64k

128k

100

16k

32k

64k

128k

16k

32k

64k

128k

16k

qa_2

32k

64k

128k

32k

64k

128k

16k

32k

64k

128k

16k

cwe

32k

64k

128k

32k

64k

128k

Figure 9 RULER per-task accuracy (%) by sequence length at 1.4B scale. Solid lines: after 128k post-training. Dashed lines: pre-trained only. SDM maintains strong recall on NIAH single-needle tasks across all lengths, while Full Attention degrades rapidly beyond its training length. GDN struggles on multi-step retrieval (single_2, single_3, multiquery). For some reason on the variable tracking task, GDN performs better and better on longer sequence lengths. We believe this to be due to the design on the vt task which clusters distractor variables early in the sequence, unintentionally making their influence weaker and weaker the longer the sequence is. RULER per task niah_single_1

niah_single_2

8B (solid=post, dashed=pre) niah_single_3

niah_multikey_1

niah_multikey_2

100

FullAttn (post) FullAttn (pre) GDN (post)

GDN (pre) SDM (post) SDM (pre)

16k

32k

niah_multikey_3

64k

128k

16k

32k

niah_multivalue

64k

128k

16k

32k

niah_multiquery

64k

128k

16k

32k

64k

128k

100

16k

fwe

32k

64k

128k

16k

qa_1

32k

64k

128k

100

16k

32k

64k

128k

16k

32k

64k

128k

16k

qa_2

32k

64k

128k

32k

64k

128k

16k

32k

64k

128k

16k

cwe

32k

64k

128k

32k

64k

128k

Figure 10 RULER per-task accuracy (%) by sequence length at 8B scale. Solid lines: after 128k posttraining. Dashed lines: pre-trained only. At 8B, post-training dramatically improves FullAttn (green) which achieves near-perfect recall on single-needle tasks. SDM maintains strong recall and outperforms FullAttn on single_3 and vt tasks. The multikey_2 task remains challenging for all fixed-state models.

Ablations

D.1

Memory Utilization and Access

The values are collected by an offline probe run on each ablation model’s final checkpoint using a fixed held-out batch from coding data. Concretely, for every run, the script reads the first B = 8 documents whose tokenized length is at least T = 8192, truncates each document to 2048 tokens, and performs a single deterministic forward pass on this 8 × 8192 batch. During this pass, forward hooks are attached to every memory-controller layer to capture the post-softmax weights over the selected top-k memory slots for each token and attention head, with k = R on the read side and k = W on the write side. All reported statistics are then computed from these per-token top-k distributions and aggregated across tokens, heads, and layers. Cumulative Mass in top-k Keys For the cumulative mass in the top-m keys, let p ∈ Rk denote the postsoftmax distribution over the selected keys for a single token, and let p(1) ≥ p(2) ≥ · · · ≥ p(k) be the same values sorted in descending order. We then define the cumulative mass curve as C(m) =

m X

p(i) ,

m = 1, . . . , k.

i=1

This quantity measures how quickly the probability mass concentrates in the highest-weighted keys: if C(m) rises rapidly toward 1, then only a few keys account for most of the mass and the distribution is highly peaked; if it rises more slowly, the controller is making broader use of its available top-k budget. In the notebook, we summarize this curve by reporting its mean and token-level quantiles across all captured (batch × head × time) positions. Effective Number of Keys per Token per-token top-k distribution,

For the effective-keys calculation, we compute the entropy of the H(p) = −

k X

pi log pi ,

i=1

and convert it into an entropy-based effective number of keys, eff_keys(p) = exp(H(p)). This is equal to 1 for a delta-like distribution concentrated on a single key and approaches k when the distribution is close to uniform over the selected budget. To make the quantity directly comparable across runs with different values of k, we normalize it as u(p) =

eff_keys(p) ∈ (0, 1]. k

Thus, u(p) ≈ 1 indicates near-uniform use of the available top-k slots, whereas u(p) ≪ 1 indicates that the controller effectively relies on only a small fraction of them. Performance across different read and write key pairs for SDM All values are reported on 1.4B model scale after pre-training on 8k sequence length.

Run W64_R128 W64_R64 W32_R64 SWA:FullAttn W32_R128 W64_R32 W32_R32

DCLM NLL ↓

∆

Avg. RULER ↑

∆

Avg. Reasoning ↑

∆

2.6703 2.6719 2.6726 2.6729 2.6737 2.6746 2.6766

-0.0016 0.0000 +0.0007 +0.0010 +0.0018 +0.0027 +0.0047

0.4011 0.4176 0.3501 0.3474 0.4039 0.3764 0.3612

-0.0164 0.0000 -0.0675 -0.0702 -0.0136 -0.0412 -0.0564

0.3842 0.3868 0.3890 0.3805 0.3821 0.3887 0.3871

-0.0026 0.0000 +0.0022 -0.0063 -0.0047 +0.0019 +0.0003

Table 4 Comparison across varying reads (R) and writes (W). Deltas are computed w.r.t. W64_R64. Lower is better for DCLM NLL; higher is better for Avg. RULER and Avg. Reasoning. Bold indicates the best value in each non-delta metric column. RULER is averaged across 6 sub-tasks for computational reasons (niah_single 1/2/3, multikey 2, multiquery, vt).

Mixed FullAttn + SDM Hybrid

We explore a three-way hybrid architecture combining SWA (local), FullAttn (global), and SDM (global) layers within the same model. In this configuration, global layer positions alternate between FullAttn and SDM: for 5 global layers, the pattern is FA, SDM, FA, SDM, FA. All other layers remain SWA with a window of 128 tokens. We denote this architecture SWA:(FA/SDM).

3.0

SWA:GDN =0.0627, R²=0.997 SWA:SDM =0.0603, R²=0.999 SWA:FullAttn =0.0581, R²=1.000 SWA:(FA/SDM) =0.0606, R²=0.999

Training Loss (avg last 10)

2.8 2.6 2.4 2.2 2.0 1.8 1.6

1.80 1.78 1.76 1.74 1.72 1.70 1.68 1023

1019

1020

1021

Total Compute (FLOPs)

1022

1023

Figure 11 Training loss scaling law for four architectures at levels 1–8. As shown in Figure 11, the 3-way hybrid architecture outperforms both 2-way hybrids at larger scale thanks to it better scaling coefficient. This hints at the fact that SDM and FullAttn might provide complementary capabilities and hints at the possibility of augmenting existing local-global architectures with SDM layers.

Compute Resources

All experiments in this work were conducted using NVIDIA H100 Hopper 80GB GPUs using the latest stable PyTorch release available at run time. We estimate the GPU-hour usage of this work for pre-training, post-training, ablations and evaluation to be around 200k gpu hours in total.

Validation text data NLL ↓ Validation code data NLL ↓ Avg. Accuracy (%) ↑ RULER (6-task, 4k–8k) ↑

SWA:FA

SWA:(FA/SDM)

SWA:SDM

SWA:GDN

2.658 0.798 38.1 80.1

2.645 0.797 38.6 64.3

2.654 0.821 38.6 56.9

2.684 0.849 37.9 32.9

Table 5 Performance of the mixed SWA:(FA/SDM) hybrid at 1.4B scale (L08, pre-trained). RULER is averaged across 6 tasks (niah_single 1/2/3, multikey_2, multiquery, vt) at sequence lengths 4k–8k (within training length).

Record · ID 349646 · SHA-256 e3fca2e9260999f5
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.