DART: Decoded Attention over Recurrent States for Efficient Long-Context Sequence Modeling

Summary (Overview)

• Core contribution: DART (Decoded Attention over Recurrent sTates) is a novel architecture that augments Mamba-2 with state-memory attention (SMA) over compact recurrent states, enabling both recurrent compression and attention-style retrieval from the same memory representation.
• Key insight: Mamba-2's recurrent state can be interpreted as a compressed KV cache, but it only reads values (via $C_t H_t$); DART additionally decodes token-conditioned keys from chunk state memories, enabling explicit attention-style search over compressed history.
• Efficiency gains: DART achieves 75% cache memory savings compared to token-level attention with chunk size $S=256$ and state size $N=128$, while maintaining competitive training efficiency through FlashAttention-style computation.
• Empirical results: DART substantially improves associative recall (MQAR), real-world retrieval tasks (SWDE, FDA), and synthetic NIAH benchmarks over Mamba-2, while preserving general language-modeling quality on Pile pretraining and zero-shot downstream evaluations.
• Design principle: The SMA branch reuses the native Mamba-2 read vector $C_t$ for value extraction (sharing is critical for optimization), and integrates its output via a gated residual connection initialized to zero.

Introduction and Theoretical Foundation

Background: Recurrent Models vs. Attention

• Transformers provide token-level attention with quadratic training compute and linear KV-cache memory during inference.
• Recurrent models (SSMs like Mamba, Mamba-2) use compact states with linear compute but suffer from memory bottlenecks for precise long-range recall—the model cannot know which future query will need a past token when processing it.

State Space Duality (SSD)

Mamba-2 establishes an equivalence between selective SSMs and structured causal attention. The recurrent dynamics are:

$$H_t = A_t H_{t-1} + B_t^\top X_t, \qquad Y_t = C_t H_t$$

The SSD view writes Mamba-2 as structured causal attention with:

• $Q = C$ (read vector), $K = B$ (write vector), $V = X$ (input)
• $H_t$ serves as a compressed KV cache
• The readout $C_t H_t$ extracts values from this cache

Key Observation

Mamba-2 reads values from the compressed state but does not extract keys from it. DART realizes that a chunk state contribution $\Delta H_{[c]} \in \mathbb{R}^{N \times P}$ supports both:

• Value-side readout: $C_t \Delta H_{[c]} \in \mathbb{R}^{1 \times P}$
• Key-side readout: $\Delta H_{[c]} E_t \in \mathbb{R}^{N \times 1}$ (token-conditioned key decoding)

Methodology

Chunk State Memories

For sequence length $L$ and chunk size $S$, DART divides the sequence into chunks. Each chunk's state contribution is:

$$\Delta H_{[c]} = \sum_{s \in [c]} A_{s+1:e(c)} B_s^\top X_s \in \mathbb{R}^{N \times P}$$

where $A_{s+1:e(c)}$ is the decay factor, and $e(c)$ is the last index in chunk $c$.

State-Memory Attention (SMA)

For each token $t$, DART forms three vectors:

• Query: $Q_t = U_t W_Q \in \mathbb{R}^{1 \times N}$ (for chunk routing)
• SSM read vector: $C_t = U_t W_C \in \mathbb{R}^{1 \times N}$ (shared with SSM branch)
• Value-side verifier: $E_t = (U_t W_E)^\top \in \mathbb{R}^{P \times 1}$

For each historical chunk $c$, DART computes:

• Value: $V_{t,c} = C_t \Delta H_{[c]} \in \mathbb{R}^{1 \times P}$
• Key: $K_{t,c} = \Delta H_{[c]} E_t \in \mathbb{R}^{N \times 1}$
• Chunk logit: $\ell_{t,c} = Q_t K_{t,c} / \sqrt{N}$

Gated Residual Integration

SMA logits are normalized over historical chunks via softmax, producing readout $R_t$. The final output is:

$$Y_t = C_t H_t + G_t R_t, \quad G_t = \text{SiLU}(U_t W_G)$$

$W_G$ is initialized to zero, so SMA starts as a zero residual and is learned during training.

Efficient Training and Complexity

FlashAttention-Style Kernel

DART implements SMA with a custom kernel (Algorithm 1) that:

1. Fixes a query tile and streams over historical chunks
2. Constructs $K_{t,c}$ and $V_{t,c}$ from $\Delta H_{[c]}$ on-chip
3. Uses online softmax for numerical stability
4. Recomputes intermediates in the backward pass (no HBM storage of keys/values/logits)

Complexity Analysis

• Training compute: $O(L^2 NP/S)$ for SMA branch (vs. $O(L^2 P)$ for token-level attention)
• Compute ratio vs. attention: $N/S$ (smaller when $N < S$)
• Inference cache: $O(LNP/S)$ per head (vs. $O(2LP)$ for token-level KV cache)
• Cache ratio vs. attention: $N/(2S)$ — with $S=256, N=128$, this gives 25% of token-level cache (75% savings)

Experiments

Synthetic MQAR (Associative Recall)

• Setup: Sweep sequence length $L \in \{256, 512, 1024\}$, model width $d_{model} \in \{32, 64, 128\}$, SSD state size $N \in \{16, 64\}$
• Results: DART consistently outperforms Mamba-2; with $N=64$, DART reaches near-perfect accuracy across all settings. Even with $N=16$, DART exceeds Mamba-2 with $N=64$ in matched-width configurations.

Language Modeling (Pile, 100B tokens)

• Models: 130M, 370M, 780M parameter backbones; DART adds <3% parameters via SMA branch
• Pile validation: DART preserves or slightly improves perplexity and accuracy vs. Mamba-2
• Zero-shot downstream: DART achieves higher average accuracy (e.g., 40.06 vs. 39.07 for 130M scale), with gains especially clear at smaller scales

Retrieval Capacity

• Real-world tasks: DART substantially improves SWDE (e.g., 56.5 vs. 39.9 for 780M) and FDA (28.9 vs. 17.2) over Mamba-2
• NIAH benchmarks: DART recovers large fractions of targets where Mamba-2 has near-zero recall (e.g., NIAH-Single-2 at 1024 tokens: 60.6 vs. 40.6 for 780M)
• Chunk-size ablations: Smaller evaluation chunk sizes ($S=64, 32$) further improve extraction-heavy metrics

Ablations

1. SMA removal: Removing SMA readout at evaluation time collapses retrieval performance (e.g., SWDE drops from 34.9 to 4.2 for 133M), confirming SMA is responsible for gains
2. $C_t$ sharing: Sharing the native Mamba-2 read vector is critical—an independent SMA read vector fails completely (0% accuracy on MQAR)

Efficiency

• SMA kernel has lower forward time than FlashAttention-2 in tested settings
• Inference cache is substantially smaller than token-level attention (consistent with theoretical 75% savings)