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

> DART augments Mamba-2 with attention-style key-value retrieval over compressed recurrent states, achieving 75% cache memory savings while substantially improving long-range recall and retrieval tasks.

- **Source:** [arXiv](https://arxiv.org/abs/2608.02032)
- **Published:** 2026-08-29
- **Permalink:** https://picx.dev/p/1RIQCk
- **Whiteboard:** https://picx.dev/p/1RIQCk/image

## Summary

# 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)

---

_Markdown view of https://picx.dev/p/1RIQCk, served by PicX — AI-generated visual whiteboard summaries of research papers._
