PaperlayerView exact revision on arXiv

Source-linked AI summary

Multi-State Linear Attention with Generalized Second-Order Routing

Hee Seung Hwang, Sanghyuk Chun, Xindi Wu, Olga Russakovsky, Zhiwei Deng

arXiv:2610.37645v1cs.CL

TL;DR

Self-attention’s quadratic compute and growing KV cache constrain long-sequence modeling, while MSLA aggregation overlooks alignment variance. MERLIN uses key moments to route among segment outputs, and outperforms existing aggregation methods across long-context evaluations, including LongBench.

  • Problem

    Prior MSLA aggregation derives segment weights from position or mean alignment, overlooking alignment variance that can distinguish sharply aligned keys from uniformly weak ones.

    Source: §1 · pp. 1–2

  • Method

    MERLIN stores first and second moments of segment keys so queries can recover alignment mean and variance and use both to weight segment-level outputs.

    Source: §1 · p. 2

  • Results

    13.5% and 12.0% average-score gains over Meanpool on LongBench for Gated DeltaNet and Mamba-2, respectively, alongside gains on RULER tasks.

    Source: §4.2 · p. 7

  • Takeaways & Limitations

    MERLIN provides a generalizable routing mechanism for aggregation in multi-state linear attention models.

    Source: §5 · p. 9

    Experiments cover two backbones at 16K context length; larger models, longer contexts, and broader aggregation baselines remain beyond the compute budget.

    Source: E LIMITATIONS · p. 19

Abstract

from arXiv · show

Self-attention plays a key role in long-sequence modeling, but its quadratic time complexity fundamentally limits the scaling of Large Language Models. Recently, Multi-State Linear Attention (MSLA) mechanisms have been proposed to make sequence modeling more efficient by splitting the context into segments and compressing each in its own memory state. MSLA aggregates the information extracted from each memory state using only the mean alignment between a query and each segment's keys, ignoring how the alignment is distributed within the segment. In order to model the relevance of a segment to a query more comprehensively and incorporate the variance of the alignment in the aggregation, we propose MERLIN (MomEnt Routing for multi-state LINear attention), a method that tracks the first and second moments of each segment's keys. Using the moments of keys, a query can recover both the mean and the variance of its alignment with each segment to perform a more content-aware aggregation of per-state readouts. Experiments on multi-state variants of Gated DeltaNet-1.3B and Mamba-2-780M show that MERLIN consistently outperforms existing aggregation methods on various long-context tasks. In particular, MERLIN outperforms baselines on RULER needle-in-a-haystack and long-context QA (+8.3% for Gated DeltaNet-1.3B and +7.0% for Mamba-2-780M) and diverse subtasks in LongBench (+13.5% and +12.0%). It even outperforms baselines that consume 4x more memory. MERLIN provides a simple and general routing mechanism for improving long-context capabilities of MSLA models.

1 INTRODUCTION

Multi-state linear attention addresses the capacity limits of single-state linear attention by storing segments in separate memory states, but segment aggregation remains underexplored. MERLIN uses the mean and variance of query-key alignment to route among segment outputs, and the paper reports gains over existing methods on LongBench.

1 INTRODUCTION

  • Single-state linear attention compresses the whole context into a fixed-size state, which may lose information on longer contexts and underperform standard attention.Source: §1 · p. 1
  • MSLA partitions context into segments, compresses each into its own memory state, and aggregates segment-level outputs for each token [Behrouz et al., 2026; Wang et al., 2026].Source: §1 · p. 1
  • On HELMET retrieval-augmented generation samples, MERLIN places the most aggregation weight on the ground-truth segment, compared with mean-based and position-based methods.Source: Figure 1 · p. 2, §1 · p. 2
  • Prior aggregation weights use position or mean alignment, overlooking variance that can distinguish sharply aligned keys from uniformly weak ones [Wang et al., 2026; Behrouz et al., 2026].Source: §1 · p. 2
  • MERLIN stores the first and second moments of each segment’s keys so a query can recover alignment mean and variance and derive weights for segment-level outputs.Source: §1 · p. 2
  • On LongBench, MERLIN improves average score over the strongest baseline by 13.5% for Gated DeltaNet-1.3B and 12.0% for Mamba-2-780M.Source: §1 · p. 2

2 BACKGROUND

Full softmax attention has quadratic compute and a linearly growing KV cache, while linear attention uses a recurrent fixed-size memory state. MSLA divides context among segment states and aggregates their outputs, but existing content-aware routing uses only mean alignment.

2 BACKGROUND

  • Softmax attention costs quadratic time and requires a KV cache that grows linearly with context length.Source: §2 · p. 2
  • Linear attention removes softmax normalization and reorganizes computation as a recurrence over a memory state that accumulates key-value outer products.A query reads the state via a matrix-vector product; this recurrence runs in linear time and constant memory.Source: §2 · p. 3
  • Multi-state linear attention partitions context into contiguous segments, stores each in its own state, and aggregates outputs read from visible segment states.Unlike single-state attention, its number of memory states increases with content length, interpolating between single-state and full attention [Behrouz et al., 2026].Source: §2 · p. 3
  • Aggregation weights route each segment’s output into the final token representation, with larger weights intended for query-relevant segments.Source: §2 · p. 3
  • Position-based methods derive weights from the target token’s hidden state or query, without considering what each segment stores [Wang et al., 2026].Source: §2 · pp. 3–4
  • Mean-based routing uses the first moment of a segment’s keys to estimate alignment, but does not capture alignment variance [Behrouz et al., 2026].Source: §2 · p. 4
  • In 30 TriviaQA samples at 4K context length, mean-and-variance aggregation assigns the highest weight to the ground-truth segment compared with position-only and mean-only methods.The task uses retrieval-augmented generation prompts requiring the model to find information in one of multiple passages and answer correctly.Source: §2 · p. 4

3 MERLIN

MERLIN augments multi-state linear attention with segment key moments so queries can use both mean and variance of alignment to derive aggregation weights. The section develops these summaries, their use in routing, and their connection to grouped softmax attention.

3 MERLIN

  • MERLIN stores each segment’s first and second key moments, letting a query recover mean and variance of its alignment to derive weights for aggregating segment outputs.The section frames this as the method’s goal and outlines the moment construction, alignment derivation, and softmax connection.Source: §3 · p. 4

3.1 FIRST AND SECOND MOMENTS OF KEYS

  • Fixed-size first and second key moments summarize a segment: the first captures average key direction, while the second captures how keys are spread.The moments’ dimensions depend only on the projected key dimension d_r.Source: §3.1 · pp. 4–5

3.2 AGGREGATION WEIGHT DERIVATION

  • For each visible segment, MERLIN derives mean-alignment and variance scores, then mixes their softmax-normalized forms with a learned gate.A token’s hidden state determines the gate, while learnable temperatures calibrate each score’s sharpness.Source: §3.2 · p. 5

3.3 CONNECTION TO SOFTMAX ATTENTION

  • Grouped token-level softmax assigns segment mass through length-normalized evidence; its log-evidence expansion identifies mean and variance as the first two relevance terms.MERLIN uses a learned gated mixture of separately normalized mean and variance scores rather than fixed additive coefficients in one softmax.Source: §3.3 · p. 5

4 EXPERIMENTS

The experiments evaluate multi-state Gated DeltaNet and Mamba-2 on long-context benchmarks, then analyze aggregation across passage positions and segment lengths and ablate MERLIN’s design choices. MERLIN improves benchmark performance, while the analyses examine memory use and the contribution of its moment features.

4.1 SETUP

  • The study converts pretrained single-state Gated DeltaNet and Mamba-2 into multi-state models through two-stage mid-training with uniform 256-token segmentation, then evaluates long-context performance on RULER and LongBench.Stage 1 uses 2B tokens, including 0.2B RULER NIAH-augmented tokens; Stage 2 uses 60M tokens, half regular LDC and half synthetic RULER HotpotQA and SQuADQA samples.Source: §4.1 · pp. 6–7

4.2 MAIN RESULTS

  • On RULER, GDN scores 62.7 with MERLIN versus 57.9 with Meanpool and 54.6 with Position; Mamba-2 improves 7.0% over Meanpool (76.2 vs. 71.2).MERLIN also improves 9.5% over Position on Mamba-2 (76.2 vs. 69.6).Source: §4.2 · p. 7
  • MERLIN improves LongBench average scores by 13.5% for GDN (21.9 vs 19.3) and 12.0% for Mamba-2 (20.6 vs 18.4), with strongest average performance among the compared methods.The evaluation filters tasks to samples of at most 16K tokens.Source: §4.2 · p. 7, Table 3 · p. 7

4.3 ANALYSIS

  • In HELMET retrieval, MERLIN performs better overall and, like Meanpool, remains more even across gold-passage positions than the recency-biased Vanilla baseline.The evaluation shuffles the gold passage among six positions in a 16K context.Source: §4.3 · pp. 7–8, Figure 3 · p. 8
  • Across segment sizes 64, 128, 256, and 512, MERLIN outperforms baselines; at the default size 256, it beats Meanpool and Position at 64 using four times less memory.At size 64, memory consumption is approximately eight times that at size 512.Source: §4.3 · p. 8, Figure 4 · p. 8
  • MERLIN’s extra second-moment features cost at most 12.6% for Gated DeltaNet and 4.2% for Mamba-2 at dr = 128, while memory cost decreases quadratically with dr.Its parameter overhead matches Meanpool’s: 4.6% for GDN and 2.2% for Mamba-2.Source: §4.3 · p. 9

4.4 ABLATIONS

  • LongBench scores rise as moment dimension dr increases, but smaller dimensions use less memory than Meanpool and Position; performance saturates around dr = 32 for GDN and dr = 64 for Mamba-2.The dimension sweep tests dr ∈ {16, 32, 64, 128}.Source: §4.4 · p. 9, Figure 5 · p. 9
  • Sharing the memory-state keys and queries or controlling moment decay and write strength like the memory state substantially degrades MERLIN’s LongBench performance.The ablations point to dedicated keys and queries for the moments as important for informative aggregation weights.Source: §4.4 · p. 9

5 CONCLUSION

The conclusion presents MERLIN as a mean-and-variance-based method for routing among segment-level outputs in multi-state linear attention. It reports superior performance and memory efficiency relative to existing aggregation methods.

5 CONCLUSION

  • The parameter-sharing study compares variants with shared keys and queries and with write-forget gates on key moments, reporting LongBench per-category averages.Source: Table 4 · p. 9
  • MERLIN weights segment-level outputs using the mean and variance of query–key alignment, computed from compact first and second moments of each segment’s keys.The conclusion describes this as a generalizable and effective routing mechanism for multi-state linear attention models.Source: §5 · p. 9
  • MERLIN is reported to be superior to existing aggregation methods in both performance and memory efficiency.Source: §5 · p. 9

C Experiment Details

The appendix supplies notation and pseudocode context, reviews related linear-attention, memory, and routing work, and describes the experiment details and model configurations.

A NOTATION

Multi-state recurrent memory

  • MERLIN’s routing notation includes projected query and key vectors, segment key moments, and a learned gate and temperatures for combining mean and variance terms.Source: Multi-state recurrent memory · p. 14

Connection to softmax attention

Mamba-2 backbone (pseudocode)

B RELATED WORK

  • Linear attention and related recurrent models use fixed-size memory states; the appendix situates MERLIN among methods that extend these models with multiple states and alternative routing approaches.Prior approaches include fixed-length, hierarchical, and dynamic segmentation, while routing methods include state selection, slot attention, and context-aware gates.Source: B RELATED WORK · p. 15

C EXPERIMENT DETAILS

  • The experiment-details appendix covers configurations, mid-training, datasets and benchmarks, evaluation, and MERLIN pseudocode; the listed models use a segment size of 256 and moment dimension of 128.The two backbones are Gated DeltaNet-1.3B and Mamba-2-780M.Source: C EXPERIMENT DETAILS · p. 16, §C.1 · p. 16, Table 6 · p. 16

C.2 MID-TRAINING HYPERPARAMETERS

The mid-training recipe is shared across both backbones and all segment sizes, using AdamW and a cosine learning-rate schedule.

C.2 MID-TRAINING HYPERPARAMETERS

  • AdamW uses weight decay 0.1 and a peak learning rate of 1 × 10−4, followed by a cosine schedule after 305 warmup steps.Source: §C.2 · p. 16, Table 7 · p. 16
  • Each run trains for 15,259 steps on 16,384-token sequences with a global batch of 131,072 tokens, totaling 2.0B tokens.Source: §C.2 · p. 16, Table 7 · p. 16
  • The shared mid-training runs use bfloat16 mixed precision for both backbones and all segment sizes.Source: §C.2 · p. 16, Table 7 · p. 16

C.3 DATASETS AND BENCHMARKS

The paper evaluates long-context retrieval on RULER and long-context understanding on LongBench, with evaluation contexts up to 16K tokens.

C.3 DATASETS AND BENCHMARKS

  • LongBench measures long-context understanding, while RULER is used for long-context retrieval.Source: §C.3 · p. 17, Table 8 · p. 17
  • Both benchmarks are evaluated at context lengths up to 16K tokens.Source: §C.3 · p. 17, Table 8 · p. 17
  • The benchmark summary lists each dataset’s metric, context length, and sample count; mid-training data includes Long-Data-Collections augmented with RULER-style NIAH/QA.Counts vary slightly by backbone tokenizer.Source: Table 8 · p. 17, §C.3 · p. 17

C.4 EVALUATION PROTOCOL

The appendix specifies long-context evaluation protocols and describes how MERLIN processes segments on Gated DeltaNet and Mamba-2 backbones.

C.4 EVALUATION PROTOCOL

  • RULER evaluation covers eight NIAH and QA tasks, while LongBench covers eleven subtasks, all at a 16K context length.The evaluations follow official metrics and configurations; LongBench retains samples of at most 16K tokens.Source: §C.4 · p. 17
  • HELMET evaluation uses TriviaQA retrieval-augmented generation, shuffles the gold passage, and reports accuracy by its position across six context divisions.The reported positions range from the first sixth to the last sixth of the context.Source: §C.4 · p. 17

C.5 PSEUDOCODE

  • For each token, MERLIN updates its segment state and key moments, reads all visible segments, then aggregates their readouts with moment-derived weights.The Mamba-2 algorithm follows the same structure as Gated DeltaNet but uses a different memory-state recurrence.Source: §C.5 · pp. 17–18

D ADDITIONAL ANALYSIS

Additional analyses report MERLIN’s decode-memory and parameter overheads and compare its segment weights on retrieval examples.

Memory Cost

  • At dr=128, MERLIN’s decode-memory overhead is at most 12.6% for Gated DeltaNet and 4.2% for Mamba-2.The second-moment cache grows quadratically with dr, and Mamba-2 shares moments per group.Source: D ADDITIONAL ANALYSIS · p. 18, Table 9 · p. 19

Parameter Overhead

  • MERLIN’s parameter overhead roughly matches Meanpool: 4.5% for Gated DeltaNet-1.3B and 2.2% for Mamba-2-780M.Position adds 1.2% and 0.6% overhead for those respective backbones.Source: D ADDITIONAL ANALYSIS · pp. 18–19

Qualitative Examples

  • On four HELMET retrieval examples, MERLIN assigns the gold segment higher weights than both Meanpool and Position.The examples comprise two Natural Questions and two TriviaQA samples at 16K context.Source: D ADDITIONAL ANALYSIS · p. 19

E LIMITATIONS

The experiments are limited in model, context-length, training, and aggregation-method coverage.

E LIMITATIONS

  • Compute constraints limit experiments to two backbones at 16K, mid-trained from single-state checkpoints, with uniform segmentation and three comparison methods.Larger models, longer contexts, and broader aggregation baselines are left to future work.Source: E LIMITATIONS · p. 19

Related papers

Loading 2610.37645v1…