Multi-State Linear Attention with Generalized Second-Order Routing
Hee Seung Hwang, Sanghyuk Chun, Xindi Wu, Olga Russakovsky, Zhiwei Deng
arXiv:2610.37645v1This paper tackles a practical bottleneck in long-context language modeling: how to keep models efficient without losing the ability to retrieve and combine information from far back in the sequence. Standard self-attention is powerful, but its cost grows quadratically with sequence length, and the key-value cache grows linearly. That becomes a major obstacle when you want large language models to handle very long documents, long conversations, or retrieval-heavy reasoning tasks. A lot of recent work has tried to fix this with linear attention and related recurrent memory models. The basic idea is to compress past information into a fixed-size memory state, so the model can process tokens in roughly linear time. But there is a tradeoff: if you squeeze an entire long context into one state, you can lose information. That is one reason single-state linear attention methods, including gated variants like Gated DeltaNet, still tend to fall short of standard attention on difficult long-context tasks. Multi-state linear attention is a newer direction that tries to soften that bottleneck. Instead of forcing the whole context into one memory, it splits the sequence into segments and gives each segment its own memory state. Then, when the model produces an output for the current token, it reads from several prior segment states and combines those segment-level outputs. This gives the model more capacity than a single state, while still avoiding the full quadratic cost of standard attention. The key question in this paper is not how to split the sequence, but how to combine the segment outputs once you have them. The authors argue that this aggregation step is underexplored, and that it is a major bottleneck in multi-state linear attention. Existing methods often decide how much to trust each segment using only position, or using only the mean alignment between the query and each segment’s keys. In other words, they ask: how well does the query match the average content of this segment? The paper argues that this is not enough. Why not? Because two segments can have the same mean alignment but very different internal structure. One segment might contain a few strongly relevant tokens and many irrelevant ones. Another might contain uniformly weak matches. If you only look at the mean, those cases can look similar. The authors point out that the variance of the query-key alignment within a segment carries important information. High variance can indicate that some tokens in the segment are sharply aligned with the query, even if the average alone does not make that obvious. So the core idea is to route information using both first-order and second-order statistics: the mean and the variance. That is what MERLIN stands for: MomEnt Routing for multi-state LINear attention. The method stores the first and second moments of each segment’s keys. With those moments, a query can recover not only the mean alignment to each segment, but also the variance of that alignment. Then it uses both signals to compute aggregation weights over the segment-level outputs. To make the background concrete, the paper first reviews standard attention and linear attention. In normal attention, each token attends over all previous tokens with a softmax over query-key similarities, which is expressive but expensive. Linear attention removes the softmax and reorganizes the computation into a recurrence over a memory state. That memory state accumulates key-value outer products, acting like an associative memory. A query reads from that state with a matrix-vector product. This makes the computation linear in sequence length and constant in memory, but the fixed-size state limits capacity. Gated DeltaNet improves on this by adding gating and a delta-rule style correction to the memory update. Multi-state linear attention then goes further by maintaining multiple states, one per segment, which increases capacity. But once you have multiple states, you need a routing mechanism to decide how much each state should contribute to the current token’s output. That routing is the aggregation problem this paper focuses on. The authors contrast three broad aggregation styles. The simplest is position-based weighting, where the model learns a weight vector from the current token and uses that to mix the segment outputs. This can encode recency bias, but it does not look at the segment contents. A more content-aware prior approach uses only the mean key of each segment, which gives a first-order measure of how the query matches the segment. MERLIN extends this by also tracking the second moment, so it can estimate alignment variance as well. Mechanically, MERLIN uses a lower-dimensional projection for the keys and queries used in the routing calculation, mainly to keep the stored moments compact. For each segment, it stores the mean of the projected keys, which is the first moment, and the average outer product of those projected keys, which is the second moment. Together, these summarize the segment’s average direction and how spread out the keys are inside that segment. At inference time, for a given query, MERLIN computes two scalar scores for each visible segment. One score measures mean alignment between the query and the segment keys. The other measures the variance of that alignment. Then it normalizes each score with a softmax over segments, and combines the two normalized distributions with a learned gate. The gate depends on the current token’s hidden state, so the model can adapt how much it trusts the mean-based signal versus the variance-based signal. The paper also introduces learnable temperatures to control the sharpness of the two terms. There is also a nice conceptual connection to full softmax attention. If you imagine doing standard token-level softmax attention and then summing the attention mass within each segment, the relevance of a segment can be approximated by a cumulant expansion. In that view, the segment’s contribution depends on the first two terms: the mean and the variance of query-key similarity. So MERLIN is motivated not just heuristically, but by a connection to how full attention would distribute mass across segments. The difference is that MERLIN uses a gated mixture of separately normalized mean and variance terms, rather than a single fixed expansion inside one softmax. On the experimental side, the authors evaluate MERLIN on multi-state variants of two architectures: Gated DeltaNet-1.3B and Mamba-2-780M. The focus is on long-context capability, so the comparison is against different aggregation baselines within the multi-state setting. They test on long-context benchmarks including RULER, which contains needle-in-a-haystack style retrieval and long-context question answering, and LongBench, which covers a diverse set of long-context tasks. They also analyze positional bias, segment length effects, and perform ablations on design choices. The headline result is that MERLIN consistently improves over existing aggregation methods. On RULER needle-in-a-haystack and long-context QA, it improves by 8.3% for Gated DeltaNet-1.3B and 7.0% for Mamba-2-780M. On LongBench, it improves average performance by 13.5% and 12.0%, respectively. The paper also notes that MERLIN can outperform baselines that use four times more memory, which is important because it suggests the gains are not just from giving the model more storage, but from routing that storage more intelligently. What do these results actually show? First, they support the paper’s central claim that aggregation is a real bottleneck in multi-state linear attention. If the model already has multiple memory states, the way it mixes them matters a lot. Second, they show that incorporating second-order information gives a consistent advantage over using only position or only mean alignment. That is especially meaningful on retrieval-heavy tasks, where the model needs to identify which segment actually contains the answer and then amplify that segment in the final output. The paper’s analysis on a TriviaQA-based setup is also telling. When they inspect how different aggregation methods distribute weight across segments, the mean-plus-variance approach gives the highest weight to the segment containing the ground-truth answer more often than position-only or mean-only methods. That aligns with the intuition that the best router should not just prefer recent segments or broadly similar segments, but should identify the segment with the most informative internal alignment structure. As for limitations, the paper is fairly clear that MERLIN is specifically a routing mechanism for the aggregation stage in multi-state linear attention. It does not solve the separate challenge of how best to partition the context into segments, which the paper explicitly identifies as another major problem in MSLA. So MERLIN improves one half of the multi-state story, but not the whole system. Another thing to watch is that the method relies on storing and using first and second moments of projected keys. That is still compact, but it introduces extra machinery compared with the simplest aggregation schemes. The paper argues this is flexible and efficient, yet in practice there will always be a balance between richer routing signals and implementation simplicity. The broader takeaway is that long-context efficiency is not only about compressing more aggressively. It is also about preserving the right statistics of the past so the model can make better decisions about what to read. MERLIN’s main contribution is to show that second-order information, not just mean similarity, can make multi-state linear attention much better at routing information from the right segment at the right time.
Also spotted that day
Previous daily papers