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

---

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)

LongBench Results (16K filtered)

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:

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.