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

Jun 9, 2026

Dynamic Linear Attention

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.

Xin Wang, Hui Shen, Boyuan Zheng +7

The Ohio State University · University of Michigan · ByteDance Seed

Long-Context Language Modeling Linear Attention

Ask PDF

Mar 21, 2026

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

As long-context language modeling becomes increasingly important, the cost of maintaining and attending to large Key/Value (KV) caches grows rapidly, becoming a major bottleneck in both training and inference. While prior works such as Multi-Query Attention (MQA) and Multi-Latent Attention (MLA) reduce memory by sharing or compressing KV features, they often trade off representation quality or incur runtime overhead. We propose Memory-Keyed Attention (MKA), a hierarchical attention mechanism that integrates multi-level KV caches (local, session, and long-term) and learns to route attention across them dynamically. We further introduce Route-Fused MKA (FastMKA), a broadcast-routed variant that fuses memory sources before attention computation for improved efficiency. Experiments on different sequence lengths show that FastMKA achieves a favorable accuracy-efficiency trade-off: comparable perplexity to MLA while achieving up to 5x faster training throughput and 1.8x lower evaluation latency. These results highlight MKA as a practical and extensible framework for efficient long-context attention.

Dong Liu, Yanxuan Yu, Ben Lengerich +1

University of California, Los Angeles, Los Angeles, California, USA · Columbia University, New York, New York, USA · University of Wisconsin-Madison, Madison, Wisconsin, USA

Attention Mechanisms Memory-Efficient Attention Long-Context Language Modeling KV-Cache Management LLM Inference Acceleration Memory-Augmented Language Models

Ask PDF

Sep 29, 2026

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

Recurrent neural networks (RNNs) compress the historical context into a memory state of fixed size, thus allowing for constant-time inference. The memory state size is a crucial factor in their performance, as exemplified by the strong performance and resurgence of linear attention, which extends the vector-valued hidden states of ordinary RNNs to matrix-valued hidden states. Crucially, linear attention does so in a parameter-efficient way, in particular by using an outer product of the key and value vectors to write to the matrix-valued hidden state. We generalize this construction and propose triadic linear attention, which writes the triadic outer product of a key, a second key, and a value, into a third-order (i.e., 3D) tensor state, and reads from it by contracting both key axes with two queries. An

E

-dimensional second key thus yields an

E

-fold increase in state size while adding only two projections. Triadic linear attention is compatible with data-dependent forgetting, the delta rule, and chunkwise-parallel training. Applied to Gated DeltaNet and scalar-gated linear attention, triadic linear attention substantially improves long-context language modeling and recall, outperforming alternatives that enlarge the state.

Oliver Sieberling, Bharat Runwal, David Jin +3

Massachusetts Institute of Technology · MIT-IBM Computing Research Lab

Linear Attention Long-Range Dependency Modeling Long-Context Language Modeling

Ask PDF

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