PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Loïc Cabannes, Pierre-Emmanuel Mazaré, Gergely Szilvasy, Matthijs Douze, Maria Lomeli, Ilze Amanda Auzina, Justin Carpentier, Gabriel Synnaeve, Hervé Jégou

arXiv:2607.07386v1cs.LG

TL;DR

Linear RNNs offer fixed per-token computation but limited state capacity constrains long-context recall. SDM sparsifies Gated DeltaNet updates to expand memory, improving long-context and in-context learning performance at matched compute.

  • Problem

    Linear RNNs have limited recall because increasing their dense state size makes computation prohibitively expensive.

  • Method

    SDM replaces dense Gated DeltaNet state updates with sparse reads and gated writes to a large explicit memory table.

  • Results

    Under identical FLOPs and parameters, SDM substantially improves long-context recall and in-context learning, while learned initialization further improves common-knowledge and reasoning tasks.

  • Takeaways & Limitations

    SDM’s gains primarily arise from increased memory capacity, with learned initialization additionally enabling useful parametric memory.

  • Takeaways & Limitations

    SDM requires substantial memory and currently needs more efficient kernels to scale further, limiting suitability for some resource-constrained settings.

Abstract

from arXiv · show

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.

1 Introduction

Long-context applications require memory mechanisms that preserve dependencies across extended interactions, but recurrent models’ fixed hidden states limit recall. Sparse Delta Memory addresses this tradeoff by greatly expanding state capacity without increasing the compute budget.

  • Agentic tasks such as software engineering, research assistance, and personal assistants require memory mechanisms that preserve long-range dependencies across extended interactions.
  • Recurrent models maintain constant memory and compute per token regardless of sequence length, but their small hidden states limit long-context recall compared with transformers.Increasing recurrent memory size would improve recall, but modern linear RNNs use dense state updates.
  • Sparse Delta Memory sparsifies the Gated DeltaNet update rule, enabling a three-order-of-magnitude increase in memory state size under the same compute budget.Figure 1 characterizes SDM as providing a large state with constant FLOPs.
  • SDM significantly outperforms Gated DeltaNet on long-context recall tasks from RULER and shows better in-context learning capability.

2 Background

Linear attention stores key–value outer products in a finite-dimensional, constant-memory associative state, but dense state updates cost O(dqk × dv) per token. Gated DeltaNet improves this memory with delta-rule interference reduction, decay, and input gates, while Product-Key Memory enables sublinear sparse indexing into much larger memories.

  • Linear Attention as an Associative Memory: Linear attention writes key–value outer products to memory and retrieves associations by comparing queries with previous keys, optionally after feature mapping ϕ.With finite-dimensional ϕ, the memory tensor and normalization state can be materialized and cached in constant memory regardless of sequence length.
  • Linear Attention as an Associative Memory: Infinite-dimensional ϕ can make linear attention compute exactly softmax attention, whereas finite-dimensional ϕ yields a materializable finite state.The finite-dimensional state consists of the memory tensor Mt and normalization state zt.
  • Gated DeltaNet: Gated DeltaNet retrieves and subtracts an existing value before writing, while decay and input gates control interference, forgetting, and update strength.The decay gate is αt ∈(0, 1), and the learned input gate satisfies βt ∈[0, 1].
  • Gated DeltaNet: O(dqk × dv) per token is the dense Gated DeltaNet update cost, so enlarging its state produces a linear increase in FLOPs.Gated DeltaNet uses dense kt and qt vectors and a dense state Mt ∈Rdqk×dv.
  • Product Key Memory: Product-Key Memory indexes k arbitrary slots among N using sublinear computation, scaling to N ≈106 memory slots versus N ≈103 or 104 for earlier dense approaches.Its top-k construction avoids materializing all N scores, and the memory state is learned only through training rather than updated by context.

3 Method

SDM sparsifies Gated DeltaNet by addressing a large explicit memory table with selected writes and reads, while retaining gated delta updates and sparse retrieval. The architecture recovers GDN in the dense-selection limit, supports a learned initial memory state, and controls memory capacity through the number of heads without changing FLOPs.

  • Sparse memory updates: SDM maintains an N-slot memory table and applies gated delta updates only to W slots selected by sparse keys.This replaces dense-state decay and delta updates with sparse addressing.
  • Sparse addressing: Top-W write selection and top-R read selection choose memory slots from projected input keys and queries.The sparse key-selection procedure scores memory slots and selects write and read indices.
  • Sparse retrieval: The selected memory slots are read by a weighted sum, then RMS-normalized, gated, and mixed across SDM heads.The retrieved output is normalized, elementwise gated, and projected into the layer output.
  • Connection to GDN: When N = d_qk and W = R = d_qk with dense key values, SDM exactly recovers the GDN update, except for GDN’s 1D qkv convolutions.The equivalence holds when all memory slots are selected.
  • Learned initial state: A learned initial state M0 lets SDM memory acquire knowledge during pretraining and reuse it at test time without adding inference FLOPs.The learned-M0 variant is the default setting.
  • Memory scaling: The number of heads controls SDM state size without affecting FLOPs, avoiding the impractical O(d^3) total-memory growth of unrestricted scaling.A single head maximizes memory size under a parameter constraint, while H remains a state-size hyper-parameter.

4 Experimental Setup

The experiments use hybrid attention architectures, SDM-specific sparse-memory configurations, and matched training procedures across a fixed-compute scaling ladder. Models are evaluated on language modeling, coding, reasoning, commonsense, in-context, and retrieval-oriented settings.

  • Architecture: All models interleave Sliding Window Attention and global receptive-field layers in a 3:1 short:long ratio, using 128-token windows and grouped-query attention.The architecture uses Multi-Head Attention layers with sliding windows alongside global layers.
  • SDM Configuration: SDM uses W = R = 64 reads and writes, with N = (d/4H)^2 memory slots and softmax-normalized read and write activations.The forget gate and time-step bias follow GDN/FLA initialization conventions, while heads are scaled to keep state-to-parameter ratio near 1:1.
  • Baselines: GDN and Mamba2 baselines use d_qk = 64, value dimension 128, and a number of heads matching the attention heads.These settings define the baseline configuration used for comparison.
  • Training: Models are pretrained on 8192-token sequences with AdamW, gradient clipping at 1.0, a Warmup–Stable–Decay schedule, and a uniform learning rate by model scale.Learning rates were tuned for the transformer baseline and confirmed optimal for GDN through small grid searches.
  • Scaling Ladder: The scaling ladder trains FullAttn, GDN, and SDM models at a 160 tokens-per-parameter compute budget rather than Chinchilla-optimal token budgets.The choice reflects that many deployed architectures are trained beyond their training-optimal compute budget.
  • Evaluation: Evaluation uses validation NLL on held-out natural-text and coding data, plus a diverse set of reasoning and commonsense tasks.The task suite is listed in Table 2.

5 Results

SDM consistently improves over dense-state recurrent baselines across compute scaling, short-context reasoning and knowledge tasks, long-context retrieval, and code perplexity. Its sparse memory can match or exceed FullAttn on many retrieval tasks despite using a fixed memory representation.

  • Short-context tasks: At 1.4B and 8B scales, SDM achieves lower DCLM NLL and higher average accuracy than GDN, with the lowest DCLM NLL among all models.At 1.4B, SDM improves over GDN on 13/15 tasks.
  • Long-context retrieval: SDM reaches RULER scores of 31.2 and 50.2 at 1.4B and 8B, respectively, versus GDN’s 20.0 and 34.2.These are the highest overall scores among fixed-state models at both scales.
  • Compute efficiency: SDM outperforms GDN at every scaling-ladder level under matched FLOPs and parameters, with scaling laws predicting an advantage over full attention at 8B scale.SDM’s sparse embedding memory is excluded from the parameter matching.
  • Long-context retrieval: SDM improves or matches GDN on all 6 RULER tasks and matches or exceeds FullAttn on 4 of 6 tasks at 1.4B and 3 of 6 at 8B.FullAttn scores 76.2 overall at 8B, while retaining a large advantage on multikey 2.
  • Code evaluation: SDM has lower code perplexity than Mamba2 and GDN from 512-token sequences through 32k–512k contexts, reaching near 2.0 while baselines remain around 2.2–2.3.The short-sequence advantage is attributed to the learned initial state M0.

6 Ablations: What Makes SDM Work?

The ablations identify memory capacity as SDM’s primary performance driver, while learned initialization provides additional benefits only when paired with SDM’s larger memory. SDM’s adaptive access patterns support broad reads and concentrated writes, but its HBM-resident state currently reduces kernel utilization despite favorable throughput.

  • Disentangling Memory Capacity from Learned Initialization: SDM without learned M0 substantially outperforms GDN, showing that increased memory capacity—not initialization—is the primary performance driver.Adding learned M0 to vanilla GDN does not measurably improve performance because its state remains limited.
  • Impact of Memory State Size: 0.914 →0.947 NLL as memory shrinks from 432 MB to 27 MB confirms that larger memory states improve modeling, with all SDM variants outperforming GDN on long-context recall.Performance degrades monotonically as the memory state is reduced, rather than collapsing abruptly.
  • Training Efficiency: SDM’s HBM-resident state has 10× lower bandwidth than SRAM, leaving its current kernel MFU around an order of magnitude below optimized GDN.GDN’s compact state fits in GPU SRAM, whereas SDM’s larger state must remain in HBM.
  • Training Efficiency: SDM is 10% slower than GDN but 6 times faster than FullAttn despite its current memory-access inefficiency.This comparison reflects observed end-to-end training efficiency under the ablation discussion.
  • Adaptive Memory Access Patterns: With k = 64 writes, the top 32 keys capture ∼85% of probability mass, whereas with k = 128 reads, the top 64 hold only ∼77%.Writes are therefore more concentrated, while reads use broader access patterns.

7 Conclusion

Sparse Delta Memory extends Gated DeltaNet with sparse Product-Key Memory, enabling state sizes thousands of times larger at identical FLOPs. This larger state improves training loss, NLL, and long-context performance, while memory requirements and kernel efficiency remain limitations.

  • Contribution: SDM extends Gated DeltaNet with a sparse Product-Key Memory design.Its sparse addressing replaces the dense memory design used by Gated DeltaNet.
  • Contribution: Thousands of times larger state sizes are achieved than GDN while keeping FLOPs identical.The architecture scales state capacity substantially without increasing the stated compute budget.
  • Results: SDM reaches much better training loss and NLL than both GDN and Full Attention.The passage attributes these gains to SDM’s much larger state size.
  • Results: SDM demonstrates much better long-context performance across a wide range of long-context tasks.The cited passage identifies the RULE benchmark family as part of this evaluation, but the excerpt ends before specifying its full name.
  • Limitations: SDM still requires more efficient kernels for further scaling and may use memory as large as the model parameters.This footprint is not adapted to certain resource-constrained contexts, despite the implementation scaling to 8B models.

Appendices … A.3 Memory-Efficient Backward

The appendices describe SDM’s efficient training procedure: chunk-parallel computation handles intra-chunk dependencies, while sequential memory access handles inter-chunk recurrence. They also exploit sparse interactions and reversible in-place updates to reduce computation and backward-pass memory requirements.

  • A Efficient Training of SDM: SDM training decomposes the gated delta rule into batched intra-chunk parallel computation and sequential chunkwise recurrent computation.The decomposition follows the WY representation approach from GDN and computes exact outputs and gradients across long sequences.
  • A.1 Intra-Chunk Parallel vs. Chunkwise Recurrent: The WY representation resolves causal dependencies within each chunk analytically through a triangular solve, enabling parallel computation.Token reads depend on prior writes within the same chunk, so the sequence is split into chunks of size C.
  • A.2 Sparse Inner Product via Two-Pointer Merge: Sparse interaction matrices are lower-triangular and capture write–write and read–write effects only through shared memory slots.For sparse keys, two tokens interact only when their top-W slot indices overlap.
  • A.1 Intra-Chunk Parallel vs. Chunkwise Recurrent: Phase 2 reads and writes the recurrent memory sequentially, correcting intra-chunk delta values before producing outputs.The correction term accounts for earlier tokens in the chunk having already modified the memory read by later tokens.
  • A.1 Intra-Chunk Parallel vs. Chunkwise Recurrent: The decomposition keeps per-layer compute at O(T · (W^2 + W · d)), while sequential memory access scales as O((T/C) · W · d).Phase 1 is compute-bound and batched across chunks, whereas Phase 2 is memory-bound and dominated by random gathers and scatters.
  • A.2 Sparse Inner Product via Two-Pointer Merge: A sorted-index two-pointer merge computes sparse inner products in O(W) time per token pair instead of O(W^2).The sparse interaction kernel performs O(C^2 · W) total work over causal entries.
  • A.2 Sparse Inner Product via Two-Pointer Merge: Compared with GDN’s dense inner product, the sparse interaction cost is O(C^2 · W), with W = 64 ≪N, allowing large memory tables without increasing interaction cost.The sparse formulation avoids addressing the full memory during interaction computation.
  • A.3 Memory-Efficient Backward: Memory-efficient backward recovery applies sparse updates in-place during the forward pass and undoes them during backward propagation.This reduces peak memory from O((T/C)×N ×d) to O(N ×d+T ×W ×d) by avoiding full memory copies at chunk boundaries.

B SDM Memory Utilization During Training

During training, SDM learns broad dynamic ranges for forgetting and input strength while rapidly using all memory slots and gradually diversifying its writes. These utilization patterns indicate active decay, overwrite, and increasingly distributed memory access.

  • Gate dynamics: The forget gate converges early to exp(g) ≈0.95, while its minimum remains near zero as some overwritten slots are fully decayed.The input gate stabilizes at σ(b) ≈0.50, with its maximum and minimum reaching ∼1 and ∼0.
  • Memory utilization: Read utilization quickly reaches 100% of memory, indicating that no memory slot remains unutilized.Statistics are averaged across 5 SDM layers during training of the 1.4B-scale L08 model.
  • Memory utilization: Write utilization increases steadily from 28% to 42%, suggesting that SDM diversifies its write patterns over training.The statistics are accumulated over 10 training steps, approximately 21M tokens, per data point.

C RULER Per-Task Accuracy by Sequence Length … F Compute Resources

The supplied material reports per-task RULER behavior across sequence lengths and scales, defines SDM memory-utilization diagnostics, and describes hybrid-layer and compute-resource evaluations. Results emphasize SDM’s sustained retrieval performance, complementary hybrid scaling, and the evaluation setup used for ablations.

  • C RULER Per-Task Accuracy by Sequence Length: At 1.4B scale, SDM maintains strong recall on NIAH single-needle tasks across sequence lengths, while Full Attention degrades beyond its training length.GDN struggles on multi-step retrieval tasks, although it improves on variable tracking at longer lengths.
  • C RULER Per-Task Accuracy by Sequence Length: At 8B scale, post-training gives FullAttn near-perfect single-needle recall, while SDM outperforms FullAttn on single_3 and vt.The multikey_2 task remains challenging for all fixed-state models.
  • D.1 Memory Utilization and Access: SDM memory utilization is assessed from post-softmax selected-key weights using cumulative top-k mass and entropy-based effective-key measures.The analysis summarizes mean and token-level quantiles across batch, head, and time positions; normalized utilization near 1 indicates near-uniform use of top-k slots.
  • D Ablations: The ablation compares varying numbers of read and write key pairs against W64_R64 using DCLM NLL, average RULER, and average reasoning metrics.RULER averages six subtasks: niah_single 1/2/3, multikey_2, multiquery, and vt.
  • E Mixed FullAttn + SDM Hybrid: The three-way SWA:(FA/SDM) hybrid alternates FullAttn and SDM among five global layers, with all remaining layers using SWA with a 128-token window.Its global-layer pattern is FA, SDM, FA, SDM, FA.
  • E Mixed FullAttn + SDM Hybrid: At larger scale, the three-way hybrid outperforms both two-way hybrids because of its better scaling coefficient, suggesting complementary FullAttn and SDM capabilities.The passage also suggests augmenting existing local-global architectures with SDM layers.
  • F Compute Resources: Experiments used NVIDIA H100 Hopper 80GB GPUs and the latest stable PyTorch release available at run time.Estimated total usage across pre-training, post-training, ablations, and evaluation was around 200k GPU hours.
  • F Compute Resources: The mixed SWA:(FA/SDM) hybrid was evaluated at 1.4B scale in a pre-trained L08 model, with RULER averaged across six tasks at 4k–8k sequence lengths.These sequence lengths are within the training length.
Loading 2607.07386v1…