Multi-State Linear Attention with Generalized Second-Order Routing

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

cs.CLOct 27, 2026

Multi-State Linear Attention with Generalized Second-Order Routing

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

Organizations: Princeton University · Google DeepMind

Abstract

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.

Figures & tables

Panel of Figure 1 — Figure 1: (a) A multi-state linear attention (MSLA) model aggregates segment-level outputs via a weighted sum. Prior methods compute the weights $\lambda_{t,i}$ from position or from the mean of each segment’s keys. Merlin instead reads both the first and second moments of the keys. (b) On retrieval-augmented generation samples in HELMET ( Yen et al., 2025 ) , we measure the aggregation weight each method places on the ground truth (GT) segment at the query token. Position-based aggregation largely misses the GT and mean-based aggregation recovers more, while Merlin places the most weight on it.

Mechanism	Compute (Prefill)	Memory
Full Attention ( Vaswani et al., 2017 )	$O(T^{2}d)$	$O(Td)$
Linear Attention ( Katharopoulos et al., 2020 )	$O(Td^{2})$	$O(d^{2})$
Multi-State Linear Attention
Log-Linear Attention ( Guo et al., 2026 )	$O(Td^{2}\log T)$	$O(d^{2}\log T)$
Memory Caching ( Behrouz et al., 2026 )	$O(TLd^{2})$	$O(Ld^{2})$
Merlin (ours)	$O(TL(d^{2}+d_{r}^{2}))$	$O(L(d^{2}+d_{r}^{2}))$

Table 1: Computational and Memory Complexity. Memory is the size of the state(s) retained to decode each token.

d_{r}

is the dimension of the moments of keys (§ 3.1 ). Memory Caching ( Behrouz et al., 2026 ) and Merlin assume uniform segmentation (

L=T/m

) where

m

is the number of tokens per segment.

Figure 2: Merlin Overview. The context is split into segments $C_{i}$ (middle row). Each segment stores a memory state (top row) and the first and second moments of its keys (bottom row). At position $t$ , the query ${\bm{q}}_{t}$ reads segment-level outputs from memory states (step 1). The query computes the mean and variance of the key-query alignment from the moments (step 2). The final output is a weighted sum of segment-level outputs, where the weights are derived from a gated mix of the alignment mean and variance (step 3).

	MidTr	S-NIAH-1	S-NIAH-2	S-NIAH-3	MK-NIAH-1	MQ-NIAH	MV-NIAH	HotpotQA	SQuADQA	Avg
Gated DeltaNet 1.3B	–	18.6	25.2	4.2	16.8	5.1	15.8	10.0	8.3	13.0
	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	–	0.0	0.0	0.0	0.0	0.0	0.0	2.0	0.0	0.3

Table 2: Performance on Long-Context Retrieval Tasks. We evaluate mid-trained ( MidTr ) models with different aggregation methods on the NIAH and QA tasks in RULER at 16K context length, – is the base model without mid-training. Highest scores in each column are in bold and second-highest are underlined .

		Single-doc QA			Multi-doc QA		Summarization			Few-shot Learning
	MidTr	NQ	QR	MF	HP	2W	GR	QM	MN	SS	TC	TQ	Avg
Gated DeltaNet 1.3B	-	2.4	5.0	13.9	4.5	7.9	9.1	16.1	9.2	33.0	22.3	19.5	13.0
	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

Table 3: Performance on Long-Context Tasks in LongBench. We evaluate on eleven LongBench tasks, filtered for samples with at most 16K tokens.

Panel of Figure 3 — Figure 3: Position-Controlled Retrieval. Retrieval accuracy as the gold passage sweeps from the start ( $1/6$ ) to the end ( $6/6$ ) of the context. Vanilla is recency-biased: it recovers the gold only when placed near the end and is otherwise flat, while Position and mean-based aggregation are flatter but weaker. Merlin stays high and roughly uniform across positions.

Panel of Figure 4 — Figure 4: Merlin Outperforms Baselines Across Segment Length. Average score as the segment size that defines each state varies over $\{64,128,256,512\}$ . Merlin leads at every segment length, and its margin over mean-based aggregation persists as segments grow. Vanilla (dashed) is segment-independent.

Figure 5: Dimension of Moments. Average LongBench scores as

d_{r}

varies over

\{16,32,64,128\}

(log axis). Performance saturates around

d_{r}\,{=}\,32

for GDN and

d_{r}\,{=}\,64

for Mamba-2.

	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

Table 4: Parameter Sharing. We train two Merlin variants: sharing the keys and queries between the memory state and moments of keys (+shared KQ) and using write-forget gates on moments of keys (+write-forget). We evaluate on LongBench and report per-category average scores. The categories are identical to Table 3 .

Long-Context Language Modeling

Long-context language modeling studies how language models can use, retain, and efficiently process information across very long sequences.

See all 189 papers

Dynamic Linear Attention
Jun 9, 2026Xin Wang, Hui Shen, Boyuan Zheng +7Long-Context Language Modeling Linear Attention
The Ohio State University · University of Michigan · ByteDance Seed
MKA: Memory-Keyed Attention for Efficient Long-Context Reasoning
Mar 21, 2026Dong Liu, Yanxuan Yu, Ben Lengerich +1Attention Mechanisms Memory-Efficient Attention
University of California, Los Angeles, Los Angeles, California, USA · Columbia University, New York, New York, USA · University of Wisconsin-Madison, Madison, Wisconsin, USA
Triadic Linear Attention: Three-Dimensional Recurrent States for Long-Context Sequence Modeling
Sep 29, 2026Oliver Sieberling, Bharat Runwal, David Jin +3Linear Attention Long-Range Dependency Modeling
Massachusetts Institute of Technology · MIT-IBM Computing Research Lab

Jun 9, 2026J/K move · Enter open · S save

Dynamic Linear Attention

Xin Wang, Hui Shen, Boyuan Zheng, Xueshen Liu+6

The Ohio State University · University of Michigan · ByteDance Seed

Long-Context Language Modeling Linear Attention

Open paper AskPDF arXiv

The scalability of Large Language Models (LLMs) to long contexts is fundamentally constrained by the quadratic complexity of standard attention, motivating the adoption of linear attention mechanisms with sub-quadratic cost. To improve representation capacity under long contexts, recent approaches organize memory in a multi-state manner. However, existing multi-state linear attention methods rely on fixed state merging policies that cannot adapt to dynamically varying token importance, irreversibly obscuring critical tokens and causing severe error accumulation over long sequences. To address this limitation, we propose DLA, a dynamic memory modeling framework for multi-state linear attention. DLA introduces (i) Information-Aware Dynamic State Merging, which adaptively determines state boundaries based on token-level information variation, preserving high-resolution representations around semantic transitions while aggressively summarizing stable regions, and (ii) Capacity-Bounded Memory Modeling, which maintains a fixed-size, chronologically ordered state cache by selectively merging adjacent low-information states to control memory growth with minimal information loss. We pre-train DLA on two different linear attention models and evaluate on 16 datasets across three categories. Experimental results demonstrate the superiority of DLA over state-of-the-art.

Abstract

Long-Context Language Modeling

Explore similar work

Dynamic Linear Attention

MKA: Memory-Keyed Attention for Efficient Long-Context Reasoning

Triadic Linear Attention: Three-Dimensional Recurrent States for Long-Context Sequence Modeling