Technical Summary: Sparse Delta Memory (SDM)

Problem Statement

Linear Recurrent Neural Networks (RNNs), such as Gated DeltaNet (GDN) and Mamba2, offer the advantage of fixed memory and compute per token, making them theoretically suitable for arbitrarily long contexts. However, their practical performance on long-context recall and in-context learning is limited by their small hidden state sizes. Increasing the state size in standard linear RNNs leads to a linear increase in FLOPs due to dense state updates (matrix-vector multiplications), creating a prohibitive computational bottleneck. Conversely, standard Transformer architectures with softmax attention maintain high recall but suffer from linearly growing Key-Value (KV) cache memory and compute costs as sequence length increases.

The core challenge addressed is how to scale the hidden state capacity of linear RNNs by orders of magnitude to improve long-context performance without increasing the computational cost (FLOPs) per token.

Methodology: Sparse Delta Memory (SDM)

The authors introduce Sparse Delta Memory (SDM), an architecture that extends the Gated DeltaNet by replacing the dense key-value outer product with a sparse addressing scheme based on Product-Key Memory (PKM).

Core Mechanism

1. Explicit Sparse Memory: Instead of maintaining a dense matrix $M_t \in \mathbb{R}^{d_{qk} \times d_v}$, SDM maintains an explicit memory table $M_t \in \mathbb{R}^{N \times d_v}$ with $N$ slots, where $N$ can be orders of magnitude larger than the hidden dimension (e.g., $N \approx 10^6$).
2. Sparse Key Selection: At each timestep, projected keys ($k'$) and queries ($q'$) are split into two halves. Their outer sum generates $N$ scores for memory slots. Using a top-$W$ (write) and top-$R$ (read) selection strategy, the model identifies specific memory slots to update or read from. This selection is performed efficiently without materializing the full $N$-sized score matrix, leveraging the property that $\text{topk}(s_1 \oplus s_2) = \text{topk}(\text{topk}(s_1) \oplus \text{topk}(s_2))$.
3. Gated Delta Updates: For the selected $W$ write slots, the model applies the Gated DeltaNet update rule:
   • Forget Gate: $\tilde{M}_t[i] = \alpha_t \cdot M_{t-1}[i]$
   • Delta Update: $M_t[i] = \tilde{M}_t[i] + \beta_t \cdot k^{(i)}_t \cdot (v_t - \tilde{M}_t[i])$
   • Unselected slots remain unchanged.
4. Sparse Read: The output is a weighted sum of the $R$ selected read slots.

IsoFLOP Design

To ensure a fair comparison, the authors design SDM to match the parameter count and FLOPs of a dense GDN baseline:

• FLOPs: The per-token cost is $O((W+R) \times d_v)$, independent of the memory size $N$. By setting $W$ and $R$ equal to the query dimension of the GDN baseline, the computational cost remains identical.
• Parameters: Linear projections are sized identically to the baseline. The memory size $N$ is controlled via the number of heads $H$ and the projection dimensions, allowing the state size to scale as $O(1/H^2)$ without affecting FLOPs.

Learned Initial State

Unlike GDN, which typically initializes its small state to zero, SDM treats the initial memory state $M_0$ as a learnable parameter. This allows the model to store pretraining knowledge directly in the memory structure, acting as a parametric memory that persists across inference.

Key Contributions

1. Architecture: The proposal of SDM, which scales linear RNN state capacity by three orders of magnitude (e.g., from kilobytes to gigabytes) while maintaining constant FLOPs per token.
2. IsoFLOP Validation: Demonstration that increasing state capacity alone (without increasing compute) significantly improves performance on long-context tasks.
3. Parametric Memory: The introduction of a learned initial state $M_0$ for linear RNNs, showing that this mechanism effectively stores common knowledge and improves reasoning tasks.
4. Scaling Laws: Empirical validation that SDM follows predictable scaling laws, outperforming dense GDN and, at larger scales (8B parameters), surpassing Full Attention models in training loss.

Experimental Results

The authors trained models ranging from 1.4B to 8B parameters on over 1 trillion tokens.

• Scaling Laws: On a scaling ladder, SDM consistently achieved lower training loss than GDN at all compute levels. At the 8B scale, SDM reached a lower loss than both GDN and a Full Attention model.
• Long-Context Recall (RULER Benchmark): SDM significantly outperformed GDN on long-context retrieval tasks. At 1.4B and 8B scales, SDM achieved scores of 31.2 and 50.2 respectively, compared to GDN's 20.0 and 34.2. Notably, SDM matched or exceeded Full Attention performance on several RULER sub-tasks despite using a fixed memory representation.
• Short-Context and Reasoning: SDM achieved lower perplexity on code data and higher average accuracy on reasoning benchmarks (e.g., MMLU, GSM8K, HumanEval) compared to GDN. At the 8B scale, SDM outperformed Full Attention on average accuracy.
• Ablation Studies:
  • State Size: Performance degraded monotonically as memory size was reduced, confirming that the large state is the primary driver of gains.
  • Learned Initialization: Learning $M_0$ improved performance for SDM but provided negligible gains for GDN, suggesting the large capacity is necessary to utilize the learned initialization effectively.
  • Efficiency: While SDM's training throughput was 1.49x slower than GDN at the 8B scale (due to HBM memory access vs. SRAM), it was 6x faster than Full Attention during inference.

Significance and Claims

The paper claims that SDM successfully bridges the gap between the efficiency of linear RNNs and the long-context capabilities of Transformers. By decoupling memory capacity from computational cost through sparsity, SDM enables:

• Constant Space/Compute: Processing arbitrarily long sequences (up to 1 million tokens demonstrated) without the unbounded memory growth of KV caches.
• Superior Long-Context Performance: Achieving state-of-the-art results among linear RNNs and competitive results against Full Attention on long-context retrieval and in-context learning.
• Parametric Knowledge Storage: The ability to learn and store pretraining knowledge within the memory state itself, improving general reasoning and common-knowledge tasks.

The authors acknowledge limitations, specifically that the large memory footprint (comparable to model parameters) requires significant HBM resources, and that current kernel implementations are less efficient than optimized dense GDN kernels. However, they posit that SDM opens new possibilities for developing agents with improved long-term memory and for processing long-video or extended reasoning traces.