MERLIN: Enhancing Long-Context Capabilities of Linear Attention with Second-Order Routing

The Efficiency Dilemma of Linear Attention and the Breakthrough of MSLA

Long-sequence modeling has always been a core challenge for large language models. Although standard self-attention is powerful, its quadratic time complexity fundamentally limits the scalability of models when processing ultra-long contexts. In recent years, Multi-State Linear Attention (MSLA) has attracted attention as an efficient alternative: it partitions the context into multiple segments, compresses each segment into an independent memory state, and thus models long sequences while maintaining linear complexity.

However, existing MSLA mechanisms have a neglected flaw — when aggregating information from each memory state, they rely solely on the average alignment between the query and each segment's keys, completely ignoring the distribution of alignment within the segment. In other words, the model only knows "how relevant this segment is to the query overall," but not "whether this relevance is uniformly distributed or highly concentrated within the segment." This information loss directly affects aggregation accuracy.

The quadratic complexity of standard self-attention stems from its computation: for a sequence of length n, each token must compute attention weights with all other tokens, forming an n×n attention matrix, so both time and memory costs grow quadratically with sequence length. Linear attention approximates or reformulates softmax attention into a kernel function form, rearranging the computation to "first aggregate keys and values, then multiply with the query," reducing complexity to O(n). MSLA further partitions long sequences into multiple segments, each maintaining a fixed-size memory state, and queries interact only with these compressed states, thus avoiding the capacity bottleneck of a single state while maintaining linear complexity. This "segmented compression + state aggregation" design is the basis for understanding MSLA's efficiency advantage, and also the source of its aggregation accuracy problem — when information within a segment is compressed into a state, distribution details are discarded.

MERLIN's Core Idea: From First to Second Moments

To address the above problem, researchers proposed MERLIN (MomEnt Routing for multi-state LINear attention), a moment routing method for multi-state linear attention. Its core innovation is: instead of tracking only the first moment (mean) of each segment's keys, it simultaneously tracks both first and second moments.

With the moment information of keys, a query can simultaneously recover the mean and variance of its alignment with each segment. This additional variance dimension is crucial — it enables the aggregation process to perceive the dispersion of alignment within a segment, thereby achieving more content-aware per-state readout aggregation. Intuitively: if the query's alignment with a segment has a high mean but large variance, the relevance is concentrated in a few positions; if the mean is high and variance is small, the relevance is more uniform. MERLIN allows the model to distinguish these two drastically different cases and make more refined routing decisions.

The paper describes this mechanism as a "simple and general routing mechanism," meaning it does not depend on a specific architecture and can be flexibly embedded into different MSLA variants.

First and second moments are basic concepts in probability and statistics. The first moment is the mean, describing the central location of a random variable; the second moment (or more precisely, the second central moment, i.e., variance) describes the dispersion of the random variable around the mean. In the context of MERLIN, the alignment between each segment's keys and the query is treated as a distribution: the first moment tells the model how relevant the segment is to the query overall, and the second moment tells the model whether this relevance is concentrated or dispersed within the segment. With only the mean, two segments with the same mean but drastically different distributions cannot be distinguished; with variance, the model can determine whether a high mean is driven by a few strongly relevant positions or contributed by many moderately relevant positions. This discriminative ability is the core of "distribution-aware" aggregation and the extra information dimension MERLIN provides compared to traditional MSLA. Notably, moment tracking only requires maintaining a few statistics, with computational overhead far lower than exact per-position computation of alignment distributions.

Experimental Validation: Consistent Improvements Across Architectures

MERLIN's effectiveness has been validated on multi-state variants of two mainstream architectures: Gated DeltaNet-1.3B and Mamba-2-780M. Experiments cover various long-context tasks, and results show high consistency.

On RULER needle-in-a-haystack and long-context question answering tasks, MERLIN achieves significant improvements over baseline methods: 8.3% on Gated DeltaNet-1.3B and 7.0% on Mamba-2-780M. On multiple LongBench subtasks, the improvements are even more considerable, reaching 13.5% and 12.0%, respectively.

One particularly noteworthy detail: MERLIN even surpasses baseline methods that consume 4× the memory. This indicates that the gains from second-order moment routing do not come from more lavish computational resources, but from the improved information efficiency of the aggregation mechanism itself — achieving smarter aggregation with fewer resources.

RULER is a benchmark for evaluating long-context language model capabilities, in which the needle-in-a-haystack task requires locating and extracting specific information from a large amount of irrelevant text, a classic test of long-range retrieval ability. LongBench is a multi-task benchmark for long-text understanding, covering question answering, summarization, code understanding, and other subtasks, used to comprehensively evaluate models' practical capabilities in long-context scenarios. These two benchmarks measure long-context capabilities from different angles: RULER focuses on precise retrieval and positional robustness, while LongBench focuses on multi-task generalization. MERLIN achieves consistent improvements on both benchmarks, indicating that the benefits of its second-order moment routing are not limited to a single task type but represent a general improvement in long-context information utilization.

Impact on Long-Context Technical Roadmaps

MERLIN provides a lightweight and general plug-in solution for enhancing the long-context capabilities of MSLA models. Its significance can be understood on two levels.

At the methodological level, it reveals a long-neglected information dimension in linear attention aggregation: the distributional characteristics of alignment. Introducing second-order statistics into routing decisions essentially evolves the model from "average relevance" to "distribution-aware relevance," which has reference value for any architecture involving segmented memory compression.

At the engineering level, MERLIN does not change the basic framework of MSLA, only adding the computational overhead of moment tracking, yet it yields stable gains across architectures and tasks, with better memory efficiency than heavier baselines. For deployment scenarios pursuing long-context capabilities but constrained by GPU memory and inference costs, this "small change, big gain" approach has practical appeal.

Of course, current experiments are limited to models with 1.3B and 780M parameters; performance on larger-scale models remains to be seen. But as a new direction for long-context optimization of linear attention, the second-order routing idea proposed by MERLIN has already demonstrated clear value.

Gated DeltaNet and Mamba-2 represent two different classes of linear attention/state space architectures. The DeltaNet series updates memory states based on the delta rule, a modern implementation of linear attention; Mamba-2 is a representative of state space models (SSMs), achieving linear-complexity sequence modeling through structured state space recurrence. Although their mathematical forms differ, both adopt the paradigm of "compressing history into states," and thus can naturally be extended to multi-state (MSLA) variants. MERLIN's effectiveness on both architectures validates its "architecture-agnostic" design claim — second-order moment routing acts on the state aggregation layer and does not depend on the underlying state update rule. This cross-architecture generality means the method has broad applicability: any model that compresses long sequences into multiple memory states and performs query aggregation could potentially benefit.

One-line Summary

MERLIN uses second-order moment routing to make multi-state linear attention aware of alignment distributions, consistently improving long-context tasks.

Paragraph Summary

The article introduces the efficiency advantages of Multi-State Linear Attention (MSLA) in long-sequence modeling and its neglected flaw: when aggregating memory states of each segment, it relies only on the average alignment between the query and segment keys, ignoring the distribution of alignment within the segment, causing information loss. To address this, researchers propose MERLIN, which simultaneously tracks the first and second moments of each segment's keys, enabling queries to recover the mean and variance of alignment, thereby achieving distribution-aware per-state readout aggregation. On Gated DeltaNet-1.3B and Mamba-2-780M, MERLIN improves RULER long-context QA by 8.3% and 7.0%, and LongBench by 13.5% and 12.0%, even surpassing baselines consuming 4× memory. The article argues that methodologically it reveals the neglected information dimension of alignment distribution, and engineering-wise it provides a lightweight, general plug-in enhancement for long-context capabilities.