# Multi-State Linear Attention with Generalized Second-Order Routing

> MERLIN improves multi-state linear attention by using both mean and variance of segment keys for aggregation, yielding consistent gains of up to 13.5% on long-context benchmarks with 4x less memory.

- **Source:** [arXiv](https://arxiv.org/abs/2610.37645)
- **Published:** 2026-10-31
- **Permalink:** https://picx.dev/p/9DJtcP
- **Whiteboard:** https://picx.dev/p/9DJtcP/image

## Summary

## Summary (Overview)

- **MERLIN (Moment Routing for multi-state LINear attention)** is proposed as a novel aggregation method for Multi-State Linear Attention (MSLA) mechanisms, leveraging both the **first and second moments** of segment keys to compute aggregation weights.
- The method addresses a key limitation of existing MSLA aggregation approaches: they only use **mean alignment** (first-order) or position-based weights, ignoring the **variance** of the alignment distribution within segments.
- MERLIN stores compact first moments ($\mu_i$) and second moments ($\Sigma_i$) of each segment's keys, allowing a query to recover both the mean and variance of its alignment with each segment.
- Experiments on multi-state variants of Gated DeltaNet-1.3B and Mamba-2-780M show consistent improvements: **+8.3%** (GDN) and **+7.0%** (Mamba-2) on RULER, and **+13.5%** and **+12.0%** on LongBench over the strongest baseline.
- MERLIN outperforms baselines that consume **4× more memory**, demonstrating both performance and memory efficiency gains.

---

## Introduction and Theoretical Foundation

### Background and Motivation

Standard self-attention (Vaswani et al., 2017) is the core operation in long-sequence modeling but suffers from **quadratic computational complexity** and a **linearly growing KV cache**. Linear attention mechanisms (Katharopoulos et al., 2020; Sun et al., 2023; Peng et al., 2023) address these limitations by reformulating attention to operate in sub-quadratic complexity, replacing the KV cache with a single fixed-size memory state.

However, single-state linear attention models (e.g., DeltaNet, Gated DeltaNet) are **inherently limited in memory capacity**, leading to information loss when compressing longer contexts. This motivates **Multi-State Linear Attention (MSLA)** mechanisms, which partition the context into segments and compress each segment into its own memory state.

### Two Novel Challenges in MSLA

1. **Separation**: How to divide the context into segments
2. **Aggregation**: How to combine segment-level outputs

The paper focuses on **aggregation**, which is relatively underexplored. Existing methods aggregate segment-level outputs via weighted sums where weights represent segment relevance to the query, but they derive weights from either:
- **Position** (Guo et al., 2026; Wang et al., 2026): assuming recency as the source of relevance
- **Mean alignment** (Behrouz et al., 2026): first-order alignment between query and segment keys

### Key Theoretical Insight

The paper demonstrates that **variance matters** for aggregation. Figure 1b shows that on retrieval-augmented generation tasks, the segment containing the ground truth receives the highest aggregation weight when weights are derived from **both mean and variance** of the alignment, compared to mean or position alone.

The theoretical foundation connects to softmax attention: when grouping token-level softmax attention mass by segment, the log-evidence expands in cumulants as:

$$\log Z_{t,i}(\delta) = \delta \eta_{t,i} + \frac{\delta^2}{2}\sigma^2_{t,i} + O(\delta^3)$$

where $\eta_{t,i}$ is the mean alignment and $\sigma^2_{t,i}$ is the variance—precisely the scores MERLIN computes.

---

## Methodology

### 3.1 First and Second Moments of Keys

For each token, MERLIN projects new keys and queries $k'_t, q'_t \in \mathbb{R}^{d_r}$ from hidden states (with $d_r < d$ to keep moments compact). The moments for segment $C_i$ are:

$$\mu_i = \frac{1}{|C_i|} \sum_{j \in C_i} k'_j, \quad \Sigma_i = \frac{1}{|C_i|} \sum_{j \in C_i} k'_j k'^{\top}_j$$

- $\mu_i$ captures the **average direction** of keys in segment $i$
- $\Sigma_i$ captures how keys are **spread** within the segment

### 3.2 Aggregation Weight Derivation

For a token at position $t$ and each visible segment $C_i$, MERLIN computes:

$$r^{\mu}_{t,i} = \mu^{\top}_i q'_t, \quad r^{\Sigma}_{t,i} = q'^{\top}_t \Sigma_i q'_t - (r^{\mu}_{t,i})^2$$

where:
- $r^{\mu}_{t,i}$: **mean key-query alignment** (first-order)
- $r^{\Sigma}_{t,i}$: **variance** of the alignment (second-order)

The final aggregation weight is computed via a **gated mixture** of softmax-normalized scores:

$$\lambda_{t,i} = \gamma_t \text{softmax}_i(\tau_\mu r^{\mu}_{t,i}) + (1 - \gamma_t)\text{softmax}_i(\tau_\Sigma r^{\Sigma}_{t,i})$$

where:
- $\gamma_t$: scalar gate computed from hidden state $x_t$ via learned projection
- $\tau_\mu, \tau_\Sigma$: learnable temperatures calibrating the sharpness of each term

### 3.3 Connection to Softmax Attention

The paper shows that MERLIN's scores correspond exactly to the first two cumulants of the segment-level attention mass in softmax attention:

$$\eta_{t,i} = \mu^{\top}_i q'_t = r^{\mu}_{t,i}, \quad \sigma^2_{t,i} = q'^{\top}_t \Sigma_i q'_t - \left(\mu^{\top}_i q'_t\right)^2 = r^{\Sigma}_{t,i}$$

Rather than using fixed additive coefficients from the cumulant expansion, MERLIN uses a learned gated mixture with independently normalized softmax terms, providing more flexibility.

### Complexity Comparison

| Mechanism | Compute (Prefill) | Memory |
|---|---|---|
| Full Attention | $O(T^2 d)$ | $O(Td)$ |
| Linear Attention | $O(Td^2)$ | $O(d^2)$ |
| Log-Linear Attention | $O(Td^2 \log T)$ | $O(d^2 \log T)$ |
| Memory Caching | $O(TLd^2)$ | $O(Ld^2)$ |
| **MERLIN (ours)** | $O(TL(d^2 + d_r^2))$ | $O(L(d^2 + d_r^2))$ |

---

## Empirical Validation / Results

### Setup

- **Models**: Gated DeltaNet-1.3B (pre-trained on 50B tokens from FineWeb-edu) and Mamba-2-780M (official checkpoint pre-trained on 300B tokens from Pile)
- **Mid-training**: Two-stage process to convert single-state models to multi-state variants
  - Stage 1: 2B tokens from Long-Data-Collections (LDC), 10% augmented with RULER NIAH-like tasks
  - Stage 2: 60M tokens with 50% LDC and 50% synthetic RULER HotpotQA/SQuADQA samples
- **Segmentation**: Uniform, 256 tokens per segment
- **Baselines**: Vanilla (single-state), Position (Eq. 5), Meanpool (Eq. 6)

### RULER Results (16K context)

| Model | Method | S-NIAH-1 | S-NIAH-2 | S-NIAH-3 | MK-NIAH-1 | MQ-NIAH | MV-NIAH | HotpotQA | SQuADQA | **Avg** |
|---|---|---|---|---|---|---|---|---|---|---|
| **GDN-1.3B** | Vanilla | 89.0 | 99.2 | 10.8 | 23.4 | 80.8 | 70.2 | 15.0 | 8.3 | 49.6 |
| | Position | 97.2 | 96.8 | 16.0 | 42.2 | 82.5 | 78.4 | 16.0 | 7.3 | 54.6 |
| | Meanpool | 99.6 | 99.4 | 21.2 | 48.4 | 86.6 | 85.0 | 12.0 | 10.7 | 57.9 |
| | **MERLIN** | **99.2** | **100.0** | **47.0** | 45.0 | **91.0** | **88.8** | **17.0** | **13.7** | **62.7** |
| **Mamba-2-780M** | Vanilla | 100 | 100 | 8.4 | 24.6 | 92.2 | 86.4 | 16 | 4.7 | 54.0 |
| | Position | 100.0 | 100.0 | 54.2 | 91.0 | 94.5 | 89.2 | 20.0 | 8.0 | 69.6 |
| | Meanpool | 100.0 | 100.0 | 66.2 | 89.4 | 93.7 | 91.8 | 20.0 | 8.1 | 71.2 |
| | **MERLIN** | 99.8 | 99.8 | **96.0** | **94.8** | **95.9** | **92.2** | **22.0** | **9.4** | **76.2** |

### LongBench Results (16K filtered)

| Model | Method | NQ | QR | MF | HP | 2W | GR | QM | MN | SS | TC | TQ | **Avg** |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| **GDN-1.3B** | Vanilla | 6.0 | 6.1 | 12.2 | 6.8 | 8.5 | 9.9 | 16.5 | 9.4 | 40.5 | 42.5 | 21.4 | 16.3 |
| | Position | 11.5 | 9.6 | 14.9 | 14.3 | 11.1 | 6.1 | 14.0 | 12.3 | 39.0 | 48.5 | 23.1 | 18.6 |
| | Meanpool | 8.8 | 15.2 | 18.0 | 14.6 | 9.8 | 8.7 | 15.3 | 10.1 | 34.5 | 54.8 | 22.4 | 19.3 |
| | **MERLIN** | **13.9** | **19.4** | **20.3** | **16.3** | **14.4** | 7.8 | **15.9** | **12.8** | **40.0** | **60.4** | 20.1 | **21.9** |
| **Mamba-2-780M** | Vanilla | 1.0 | 3.2 | 7.6 | 4.2 | 5.2 | 10.2 | 12.3 | 14.7 | 36.5 | 54.5 | 5.7 | 14.1 |
| | Position | 1.2 | 5.1 | 10.5 | 5.7 | 6.8 | 15.5 | 15.0 | 17.4 | 42.0 | 65.1 | 12.5 | 17.9 |
| | Meanpool | 1.7 | 4.2 | 11.5 | 6.0 | 5.6 | 16.0 | 17.7 | 18.5 | 44.5 | 61.1 | 15.1 | 18.4 |
| | **MERLIN** | **1.9** | **6.7** | **12.6** | **7.8** | **9.1** | **20.4** | 16.1 | **19.2** | **46.0** | **72.7** | 14.2 | **20.6** |

### Positional Bias Analysis (HELMET TriviaQA)

- **Vanilla**: Strong recency bias—struggles unless gold passage is near the end
- **Position**: Weak recency bias with gradually increasing accuracy toward the end
- **Meanpool**: Flatter but weaker overall
- **MERLIN**: Stays high and roughly uniform across all gold passage positions

### Segment Length Analysis

MERLIN outperforms baselines across segment lengths $m \in \{64, 128, 256, 512\}$. Notably, **MERLIN at $m=256$ outperforms Meanpool and Position at $m=64$**, which consume 4× more memory.

### Ablations

**Dimension of Moments ($d_r$)**:
- Performance increases with $d_r \in \{16, 32, 64, 128\}$, saturating around $d_r=32$ (GDN) and $d_r=64$ (Mamba-2)
- Trade-off: larger $d_r$ improves expressiveness but increases memory

**Parameter Sharing Ablations**:

| Variant | Single | Multi | Summ | Few-shot | **Overall** |
|---|---|---|---|---|---|
| GDN + MERLIN | 17.9 | 15.4 | 12.2 | 40.2 | **21.9** |
| + write-forget | 14.9 | 9.2 | 11.6 | 35.3 | 18.5 |
| + shared-KQ | 10.7 | 9.2 | 10.9 | 29.8 | 15.7 |
| Mamba + MERLIN | 7.1 | 8.5 | 18.6 | 44.3 | **20.6** |
| + write-forget | 6.7 | 6.9 | 16.6 | 43.8 | 19.6 |
| + shared-KQ | 5.1 | 6.1 | 15.1 | 38.3 | 17.1 |

Both ablations (sharing keys/queries between memory state and moments, and using write-forget gates) substantially degrade performance, confirming that **dedicated keys and queries for moments** are crucial.

---

## Theoretical and Practical Implications

### Theoretical Contributions

1. **Identifies aggregation as a key bottleneck** in multi-state linear attention, which was relatively underexplored compared to separation.
2. **Establishes a theoretical connection** between MSLA aggregation and softmax attention via cumulant expansion, showing that mean and variance of key-query alignment are the first two cumulants of the segment-level attention mass.
3. **Demonstrates the necessity of second-order information** for content-aware routing: variance captures sharply aligned keys that mean alone cannot distinguish from uniformly weak alignments.

### Practical Implications

1. **Memory efficiency**: MERLIN at $m=256$ outperforms baselines at $m=64$ (4× more memory), showing that better routing can substitute for more memory states.
2. **General architecture-agnostic**: Works on both Gated DeltaNet (linear attention with delta rule) and Mamba-2 (state-space model), suggesting broad applicability across MSLA architectures.
3. **Modest overhead**: Extra memory cost limited to at most 12.6% (GDN) and 4.2% (Mamba-2) at $d_r=128$, decreasing quadratically with $d_r$. Parameter overhead is on par with Meanpool baseline (4.6% for GDN, 2.2% for Mamba-2).

---

## Conclusion

### Main Takeaways

- MERLIN provides a **simple and general routing mechanism** for improving long-context capabilities of MSLA models by leveraging both first and second moments of segment keys.
- The method consistently outperforms existing aggregation methods (position-based and mean-based) across multiple benchmarks, architectures, and segment lengths.
- The variance of key-query alignment is critical for content-aware aggregation, enabling the model to identify sharply aligned keys within segments.

### Future Directions

- Exploring **adaptive segmentation strategies** (separation) in conjunction with MERLIN's aggregation, potentially optimizing both jointly.
- Scaling to larger model sizes and longer contexts (beyond 16K tokens).
- Investigating the application of moment-based routing to other sequence modeling paradigms beyond linear attention.
- Further reducing the memory footprint of the second-moment features through compression or low-rank approximations.

---

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