alphaFeedback

WorksComputer Science

Multi-State Linear Attention with Generalized Second-Order Routing

AI contributions
  • generative AI tools
View details

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.

Introduction

The self-attention mechanism is a core operation in long-sequence modeling [31]. However, its quadratic computational complexity and linearly growing key-value (KV) cache have become a key bottleneck in scaling compute and data [19]. Linear attention [18, 29, 24] addresses these limitations by reformulating attention to operate in sub-quadratic complexity. For example, DeltaNet [36] and Gated DeltaNet [37] remove the softmax in standard attention and replace the KV cache with a single fixed-size memory state. However, a single fixed-size memory state is inherently limited in memory capacity and may lead to information loss when compressing longer contexts [2, 17, 20]. As a result, models based on single-state linear attention still underperform standard attention in practice.

Recently, Multi-State Linear Attention (MSLA) mechanisms [4, 14, 32] have been introduced to improve the long-context capabilities of single-state linear attention. MSLA partitions the context into segments and compresses each segment into its own memory state. Each token extracts segment-level outputs from preceding memory states and then aggregates them to produce its final output representation (Fig. 1a). Operating in the multi-state regime therefore introduces two novel challenges: how to divide the context into segments (separation), and how to combine the segment-level outputs (aggregation).

We study aggregation, which is relatively underexplored compared to separation. Existing methods aggregate segment-level outputs via a weighted sum [4, 14, 32], where the weights represent each segment’s relevance to the query: closely aligned segments contribute more to the final output than others. However, prior works derive the weights either from position (Guo et al., [14]; Wang et al., [32]), assuming recency as the source of relevance, or from the mean (first-order) alignment between each segment’s keys and the query (Behrouz et al., [4]). They overlook the variance (second-order) of the alignment, which captures sharply aligned keys that the mean alone cannot distinguish from a segment of uniformly weak ones.

To analyze why the variance matters, we examine how an MSLA model assigns aggregation weights on a long-sequence reasoning task. Figure 1b shows that the segment containing the ground truth, which is the segment most relevant to the response, receives the highest aggregation weight when the weights are derived from both the mean and variance of the alignment, compared to the mean or position alone. This motivates an aggregation method that incorporates the variance of the alignment in deriving the weights.

Multi-State Linear Attention and aggregation weights

Figure 1. (a) A multi-state linear attention (MSLA) model aggregates segment-level outputs via a weighted sum. Prior methods compute the weights λt,i\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., [38]), 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.

We propose MERLIN (MomEnt Routing for multi-state LINear attention), a method that leverages the mean and variance of key-query alignment to improve aggregation in multi-state linear attention. MERLIN stores the first and second moments of each segment’s keys. Using the moments, a query recovers both the mean (first-order) and variance (second-order) of its alignment with each segment’s keys to derive the weights for the weighted sum of segment-level outputs. Our experiments show that MERLIN improves the long-context capabilities of multi-state linear attention models.

Our contributions are:

  1. We identify segment-level output aggregation as a key bottleneck in multi-state linear attention: existing methods weight each segment’s output using only the mean alignment between the query and the segment’s keys (§ 2).

  1. We propose MERLIN, a method that keeps the first and second moments of each segment’s keys. A query uses these moments to recover the mean and variance of its alignment with the segment (§ 3.1) and to derive the aggregation weights (§ 3.2).

  1. We evaluate multi-state variants of Gated DeltaNet-1.3B and Mamba-2-780M (§ 4.1). On LongBench (Bai et al., [3]), MERLIN improves the average score over the strongest baseline by 13.5% on Gated DeltaNet and 12.0% on Mamba-2 (§ 5).

Background

Given an input sequence of TT tokens, the query, key, and value matrices QQ, KK, V∈RT×dV \in\mathbb{R}^{T \times d} are projected from the tokens’ hidden states x1,…,xTx_1,\ldots,x_T. Standard full attention computes:

O=softmax⁡(QK⊤+M)V,O = \operatorname{softmax}(QK^{\top} + M)V,

where M∈{−∞,0}T×TM \in\{-\infty,0\}^{T \times T} is a causal masking matrix. Softmax attention costs quadratic time, and requires a KV cache that grows linearly with the context. This motivates more efficient sub-quadratic attention formulations.

Linear Attention. Linear attention removes the softmax normalization and instead computes:

O=(QKT⊙M)V,Mi,j=1 if i≥j, and 0 otherwise.O=(QK^{\mathsf{T}}\mathbin{\odot}M)V,\qquad M_{i,j}=1\ \text{if }i\ge j,\ \text{and }0\text{ otherwise.}

This formulation can be reorganized into a recurrence over a memory state St∈Rd×dS_t\in\mathbb{R}^{d\times d}.

St=St−1+vtktT=∑t′vt′kt′T,ot=Stqt,(1)S_t=S_{t-1}+v_tk_t^{\mathsf{T}}=\sum_{t'}v_{t'}k_{t'}^{\mathsf{T}},\qquad o_t=S_tq_t, \tag*{(1)}

where ktk_t, vtv_t, qtq_t are the key, value, and query vectors of a token at position tt. The memory state StS_t accumulates key-value outer products, acting as an associative memory that binds each key to its value. A query reads from the memory state via a matrix-vector product to get the output representation oto_t. The recurrence now runs in linear time and constant memory. Recent works add gating and delta-rule correction to improve the memory capacity of the fixed-size state. For example, Gated DeltaNet (Yang et al. [37]) (GDN) uses the recurrence as follows:

St=St−1(αt(I−βtktktT))+βtvtktT,(2)S_t=S_{t-1}\left(\alpha_t(I-\beta_tk_tk_t^{\mathsf{T}})\right)+\beta_tv_tk_t^{\mathsf{T}}, \tag*{(2)}

where αt∈(0,1)\alpha_t\in(0,1) decays the state and βt\beta_t sets the strength of the update.

Multi-State Linear Attention. Linear attention can be limiting when TT grows because it needs to compress the entire context into a single fixed-size state. Multi-state linear attention relaxes this bottleneck by partitioning the sequence into LL contiguous segments C1,…,CLC_1,\ldots,C_L and compressing each segment in its own state:

SCi=∑j∈CivjkjT.(3)S_{C_i}=\sum_{j\in C_i}v_jk_j^{\mathsf{T}}. \tag*{(3)}

Intuitively, Eq. (3) can be viewed as splitting the first equation of Eq. (1) into LL parts, which we call separation. At step tt, a query reads segment-level outputs from the last ℓ≤L\ell\le L memory states visible at step tt. The intermediate outputs are then aggregated via a weighted sum to derive the final output representation:

ot,i=SCiqt,ot=∑i=1ℓλt,iot,i,(4)o_{t,i}=S_{C_i}q_t,\qquad o_t=\sum_{i=1}^{\ell}\lambda_{t,i}o_{t,i}, \tag*{(4)}

where λt,i\lambda_{t,i} is the aggregation weight that determines how much the output from segment ii contributes to the final output at position tt. Eq. (4) can be viewed as the second equation of Eq. (1) applied to ℓ\ell partitioned states, which we refer to as aggregation.

Unlike single-state linear attention, in multi-state linear attention, the number of memory states increases as the content length increases. Multi-state linear attention therefore interpolates between single-state linear attention (constant memory, linear time) and full attention (linear memory, quadratic time) ([4]). Table 1 summarizes the time and memory complexities of attention variants.

MechanismCompute (Prefill)Memory
Full Attention ([31])O(T2d)O(T^{2}d)O(Td)O(Td)
Linear Attention ([18])O(Td2)O(Td^{2})O(d2)O(d^{2})
Multi-State Linear Attention
Log-Linear Attention ([14])O(Td2log⁡T)O(Td^{2}\log T)O(d2log⁡T)O(d^{2}\log T)
Memory Caching ([4])O(TLd2)O(TLd^{2})O(Ld2)O(Ld^{2})
MERLIN (ours)O(TL(d2+dr2))O(TL(d^{2}+d_{r}^{2}))O(L(d2+dr2))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. drd_r is the dimension of the moments of keys (§ 3.1). Memory Caching ([4]) and MERLIN assume uniform segmentation (L=T/mL=T/m) where mm is the number of tokens per segment.

Aggregation in Multi-State Linear Attention. The aggregation of segment-level outputs in Eq. (4) requires weights λt,i\lambda_{t,i}. These weights function as routers that increase the contribution of outputs from query-relevant segments and decrease the contribution of outputs from others. The simplest methods learn a position-based weight vector λt\lambda_t (with entries λt,i\lambda_{t,i}) directly from the target token, projected from its query or hidden state ([14, 32]):

λt=Wλxtorλt=Wλqt.(5)\lambda_t=W_\lambda x_t\quad\text{or}\quad\lambda_t=W_\lambda q_t. \tag*{(5)}
MERLIN overview showing memory states, context segments, key moments, and aggregation steps

Figure 2. MERLIN Overview. The context is split into segments CiC_i (middle row). Each segment stores a memory state (top row) and the first and second moments of its keys (bottom row). At position tt, the query qtq_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).

These weights do not consider what each segment stores and therefore cannot function as a content-aware router. Behrouz et al. [4] introduces content-awareness in aggregation weights as follows:

Behrouz et al. [4]:

λt,i=⟨kˉi,ut⟩,ut=Wuxt,kˉi=1∣Ci∣∑j∈Cikj(6)\lambda_{t,i} = \langle\bar{k}_i, u_t \rangle,\qquad u_t = W_u x_t,\qquad\bar{k}_i = \frac{1}{\lvert C_i\rvert}\sum_{j\in C_i} k_j \tag*{(6)}

where utu_t is a learned vector projected from hidden state xtx_t and kˉi\bar{k}_i is the first moment of keys in segment ii (i.e., the mean). However, this only captures the mean alignment between the query and a segment’s keys and does not capture the variance.

We revisit the analysis in Figure 1b and demonstrate that the variance of the alignment is necessary for more content-aware aggregation of segment-level readouts. We train multi-state models with different aggregation methods and evaluate on 30 randomly selected TriviaQA samples in HELMET (Yen et al. [38]) at 4K context length. This is a retrieval-augmented generation task where each sample consists of multiple passages and a prompt that requires retrieving information from one of them. Thus, the model should locate the relevant passage and extract the correct answer to the prompt.

We find that different aggregation methods assign varying amounts of aggregation weight to the segment containing the ground truth. Ideally, the weights assigned to the ground truth segment should be higher than others in order to increase the contribution of the information extracted from the segment in the final response. Figure 1b shows that mean- and variance-based aggregation assigns the highest weight to the ground truth segment compared to position-only and mean-only aggregation. This motivates our proposed method MERLIN, which leverages both orders of alignment to compute the weights for segment-level output aggregation.

MERLIN

Our goal is to capture the variance of the alignment between a query and each segment’s keys to derive more informative aggregation weights. We first introduce how the first and second moments of the keys are constructed (§3.1). Then, we discuss how the mean and variance of the alignment can be derived from the moments (§3.2). Finally, we elaborate on the relationship between multi-state linear attention equipped with MERLIN and softmax attention (§3.3).

First and Second Moments of Keys

For each token, we project new keys and queries kt′,qt′∈Rdrk'_t, q'_t \in\mathbb{R}^{d_r} from the hidden states. We use kt′,qt′k'_t, q'_t instead of kt,qtk_t, q_t to construct and read from the first and second moments of keys, primarily to keep moments compact by choosing drd_r that is smaller than dd.

The first moment of keys is the mean-pooled vector of all keys (kj′\mathbf{k}'_j) in the segment. The second moment is the average of the outer products kj′kj′T\mathbf{k}'_j\mathbf{k}'^{\mathsf{T}}_j over the segment. The first and second moments of keys can thus be formulated as:

μi=1∣Ci∣∑j∈Cikj′,Σi=1∣Ci∣∑j∈Cikj′kj′T.(7)\boldsymbol{\mu}_i = \frac{1}{\lvert C_i\rvert}\sum_{j\in C_i}\mathbf{k}'_j,\qquad\boldsymbol{\Sigma}_i = \frac{1}{\lvert C_i\rvert}\sum_{j\in C_i}\mathbf{k}'_j\mathbf{k}'^{\mathsf{T}}_j. \tag*{(7)}

The first and second moments are fixed-size, and their dimensions depend only on drd_r. Together, they summarize segment ii, where μi\boldsymbol{\mu}_i captures the average direction of the keys and Σi\boldsymbol{\Sigma}_i captures how the keys are spread within the segment.

Aggregation Weight Derivation

Given a token at position tt, we compute the following scalar scores for each visible segment CiC_i (the segments leading up to the current one):

rt,iμ=μiTqt′,rt,iΣ=qt′TΣiqt′−(rt,iμ)2.(8)r^\mu_{t,i} = \boldsymbol{\mu}_i^{\mathsf{T}}\mathbf{q}'_t,\qquad r^\Sigma_{t,i} = \mathbf{q}'^{\mathsf{T}}_t\boldsymbol{\Sigma}_i\mathbf{q}'_t - \left(r^\mu_{t,i}\right)^2. \tag*{(8)}

Here, rt,iμr^\mu_{t,i} captures the mean key-query alignment, measuring how well the segment’s keys broadly align with the query. Meanwhile, rt,iΣr^\Sigma_{t,i} captures the variance, measuring the spread of the per-token alignment between the query and the keys in the segment. We mix the softmax-normalized rt,iμr^\mu_{t,i} and rt,iΣr^\Sigma_{t,i} with a learned gate to get the final weight λt,i\lambda_{t,i}:

MERLIN (ours): λt,i=γtsoftmax⁡i(τμrt,iμ)+(1−γt)softmax⁡i(τΣrt,iΣ)\lambda_{t,i} = \gamma_t\operatorname{softmax}_i(\tau_\mu r^\mu_{t,i}) + (1-\gamma_t)\operatorname{softmax}_i(\tau_\Sigma r^\Sigma_{t,i}),

where softmax⁡i\operatorname{softmax}_i denotes the softmax operation over the segment index ii, γt\gamma_t is a scalar gate, and τμ\tau_\mu and τΣ\tau_\Sigma are learnable temperatures that scale the two scores before normalization. The gate γt\gamma_t is computed from the token’s hidden state xtx_t through a learned projection. The temperatures τμ,τΣ\tau_\mu,\tau_\Sigma calibrate the sharpness of each term. This shows how the first and second moments are used to derive the weights for Eq. (4).

Connection to Softmax Attention

Suppose we ran token-level softmax attention over the keys and then grouped its mass by segment. For a target token at position tt and a segment CiC_i, the key-query similarity can be expressed as:

st,i,j=qt′Tkj′,j∈Ci.s_{t,i,j} = \mathbf{q}'^{\mathsf{T}}_t\mathbf{k}'_j,\qquad j\in C_i.

We can define the length-normalized evidence of the segment as:

Zt,i(δ)=1∣Ci∣∑j∈Cieδst,i,j.Z_{t,i}(\delta) = \frac{1}{\lvert C_i\rvert}\sum_{j\in C_i}e^{\delta s_{t,i,j}}.

where δ\delta is a temperature hyperparameter. Grouping the token-level softmax by segment, the total attention mass a full softmax layer would place on segment ii can be expressed as:

At,i=∑j∈Ciexp⁡(δst,i,j)∑ℓ∑j∈Cℓexp⁡(δst,ℓ,j)=softmax⁡i(log⁡∣Ci∣+log⁡Zt,i(δ)).(9)A_{t,i} = \frac{\sum_{j\in C_i}\exp(\delta s_{t,i,j})}{\sum_{\ell}\sum_{j\in C_\ell}\exp(\delta s_{t,\ell,j})} = \operatorname{softmax}_i\left(\log\lvert C_i\rvert+ \log Z_{t,i}(\delta)\right). \tag*{(9)}

Because our segments are equal-sized, log⁡∣Ci∣\log\lvert C_i\rvert is the same for every segment and cancels inside the softmax. Expanding the remaining log-evidence in cumulants around δ=0\delta=0, we obtain:

log⁡Zt,i(δ)=δηt,i+δ22σt,i2+O(δ3),ηt,i=1m∑j∈Cist,i,j,σt,i2=1m∑j∈Cist,i,j2−ηt,i2(10)\log Z_{t,i}(\delta) = \delta\eta_{t,i} + \frac{\delta^2}{2}\sigma^2_{t,i} + O(\delta^3),\qquad\eta_{t,i} = \frac{1}{m}\sum_{j\in C_i}s_{t,i,j},\qquad\sigma^2_{t,i} = \frac{1}{m}\sum_{j\in C_i}s^2_{t,i,j}-\eta^2_{t,i} \tag*{(10)}

Therefore, the relevance of a segment is governed by the first two terms: the mean ηt,i\eta_{t,i} and the variance σt,i2\sigma^2_{t,i}. These are precisely the scores rt,iμr^\mu_{t,i} and rt,iΣr^\Sigma_{t,i}, which are computed from the first and second moments via matrix-vector products:

ηt,i=μiTqt′=rt,iμ,σt,i2=qt′TΣiqt′−(μiTqt′)2=rt,iΣ.(11)\eta_{t,i} = \boldsymbol{\mu}_i^{\mathsf{T}}\mathbf{q}'_t = r^\mu_{t,i},\qquad\sigma^2_{t,i} = \mathbf{q}'^{\mathsf{T}}_t\boldsymbol{\Sigma}_i\mathbf{q}'_t - \left(\boldsymbol{\mu}_i^{\mathsf{T}}\mathbf{q}'_t\right)^2 = r^\Sigma_{t,i}. \tag*{(11)}

Rather than the fixed additive coefficients of Eq. (10) (δ\delta on the mean, δ22\frac{\delta^2}{2} on the variance) inside a single softmax, MERLIN combines the two scores with a learned gated mixture of separately-normalized softmax terms (Eq. (9)), with learnable temperatures τμ,τΣ\tau_\mu,\tau_\Sigma and gate γt\gamma_t. The gated mixture is also flexible to implement, since the two terms are normalized and scaled independently.

Experiments

We conduct experiments to quantify the effect of MERLIN on the long-context capabilities of multi-state linear attention models. We first explain the models, datasets, and baselines used for training and evaluation (§ 4.1). Then, we show the main evaluation results on the core benchmarks (§ 5). Third, we conduct analysis on positional biases and impact of segment length (§ 6). Finally, we conduct ablations on the core design choices in MERLIN (§ 7).

Setup

Models. We use the Gated DeltaNet ([37]) and Mamba-2 ([7]) architectures for the following experiments. Gated DeltaNet is a linear attention model that updates the memory state following the recurrence in Eq. (2). Mamba-2 is a state-space model that uses recurrence of structured semiseparable matrices. We show that MERLIN improves upon the multi-state variants of both architectures with distinct features to demonstrate its robustness and effectiveness. A more detailed pseudocode for each model can be found in Section C.5.

Training. We first pre-train the Gated DeltaNet-1.3B ([37]) model on 50B tokens in FineWeb-edu ([23]) following the official training configurations. For Mamba-2-780M, we use the officially released checkpoint1 pre-trained on 300B tokens from Pile ([10]). These pre-trained models are still single-state: they require mid-training (or continued pre-training) to function as multi-state variants, i.e., training the pre-trained model with the same objective function as pre-training on the mid-training dataset.

In order to convert single-state models to multi-state variants, we introduce separation (Eq. (3)) and aggregation (Eq. (4)) mechanisms during mid-training. We fix the separation mechanism to uniform segmentation, where each segment has the same length of 256 tokens. At training context length of 16K tokens, the model learns to aggregate as many as 64 (16K / 256) segment-level outputs. We specify the full list of training hyperparameters and model configurations in Section C.1.

We mid-train the pre-trained models in two stages. Each stage is defined by the type of training data that was used. Stage 1 uses 2B tokens from Long-Data-Collections (LDC) ([30]), with 10% of the training samples (0.2B tokens) augmented with RULER ([15]) needle-in-a-haystack (NIAH)-like tasks. The augmented data contains artificial keys and values scattered throughout regular LDC sequences. The task query and ground truth answer are attached at the end of the context. This dataset is designed to train basic language modeling capability of a multi-state variant by activating its aggregation mechanism.

In stage 2, we train on a smaller dataset of 60M tokens. It consists of 50% regular LDC sequences, and 50% synthetically generated samples from RULER HotpotQA and SQuADQA. This dataset is designed to further strengthen the model’s aggregation mechanism under reasoning contexts. After stage 1 and stage 2, the pre-trained single-state models now function as multi-state models with long-context capabilities. The detailed mid-training hyperparameters can be found in Section C.2.

Benchmarks. We evaluate the models on various long-context benchmarks. First, we evaluate on the NIAH and QA tasks in RULER (Hsieh et al., [15]) at 16K context length. The single needle variants (S-NIAH-1/2/3) evaluate the model’s ability to retrieve a single key-value pair in different types of natural and synthetic text haystacks. The multi-needle variants (MK-NIAH-1/MQ-NIAH/MV-NIAH) evaluate the model’s ability to retrieve from multiple key-value pairs in different formats. The HotpotQA and SQuADQA tasks in RULER evaluate multi-hop and single-hop reasoning skills.

We also evaluate on the following subtasks in LongBench (Bai et al., [3]): NarrativeQA (NQ), Qasper (QR), MultiFieldQA-en (MF), HotpotQA (HP), 2WikiMultihopQA (2W), GovReport (GR), QMSum (QM), MultiNews (MN), SAMSum (SS), TREC (TC), TriviaQA (TQ). We filter for samples under 16K context length. These tasks largely evaluate single-document QA, multi-document QA, summarization, and few-shot learning capabilities.

Finally, we evaluate on the TriviaQA retrieval-augmented-generation task in HELMET (Yen et al., [38]) at 16K context length. We use this benchmark to analyze how the position of the ground truth in the sequence impacts the model’s performance. Section C.4 provides a detailed explanation of the evaluation protocol used in our experiments.

Baselines. We compare MERLIN with three other baselines: Vanilla, Position, and Meanpool. Vanilla is the unmodified single-state backbone, with neither segmentation nor aggregation mechanisms. The other two are alternative aggregation methods for multi-state variants. Position (Eq. (5)) derives aggregation weights from the hidden states alone, scaling segment-level outputs by a learned recency scalar (Guo et al., [14]). Meanpool (Eq. (6)) derives aggregation weights from first-order alignment between a query and each segment’s keys (Behrouz et al., [4]).

Main Results

We report results on RULER (Hsieh et al., [15]) NIAH and QA tasks. Table 2 shows that MERLIN outperforms baselines in most tasks for both GDN and Mamba-2 architectures. MERLIN applied to GDN leads to an average score of 62.7, while Meanpool has an average score of 57.9, followed by Position’s 54.6. In the case of Mamba-2, MERLIN shows a 7.0% relative improvement over Meanpool (76.2 vs. 71.2), and a 9.5% relative improvement over Position (vs. 69.6). This suggests that mid-training with MERLIN improves long-context retrieval.

MidTrS-NIAH-1S-NIAH-2S-NIAH-3MK-NIAH-1MQ-NIAHMV-NIAHHotpotQASQuADQAAvg
Gated DeltaNet 1.3B–
–18.625.24.216.85.115.810.08.313.0
Vanilla89.099.210.823.480.870.215.08.349.6
Position97.296.816.042.282.578.416.07.354.6
Meanpool99.699.421.248.486.685.012.010.757.9
MERLIN99.2100.047.045.091.088.817.013.762.7
Mamba-2 780M
–0.00.00.00.00.00.02.00.00.3
Vanilla1001008.424.692.286.4164.754.0
Position100.0100.054.291.094.589.220.08.069.6
Meanpool100.0100.066.289.493.791.820.08.171.2
MERLIN99.899.896.094.895.992.222.09.476.2

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 <u>underlined</u>.

We also report results on LongBench (Bai et al., [3]) tasks filtered to 16K. Overall, MERLIN demonstrates strongest average performance, followed by Meanpool, Position, and finally Vanilla. Table 3 shows that MERLIN leads to robust gains particularly on single-doc and multi-doc QA. For GDN, MERLIN leads to a 13.5% improvement in average score compared to Meanpool (21.9 vs 19.3). For Mamba-2, MERLIN leads to a 12.0% (20.6 vs 18.4) improvement compared to Meanpool. This suggests that MERLIN strengthens general long-context capabilities.

Single-doc QAMulti-doc QASummarizationFew-shot Learning
MidTrNQQRMFHP2WGRQMMNSSTCTQAvg
Gated DeltaNet 1.3B-2.45.013.94.57.99.116.19.233.022.319.513.0
Vanilla6.06.112.26.88.59.916.59.440.542.521.416.3
Position11.59.614.914.311.16.114.012.339.048.523.118.6
Meanpool8.815.218.014.69.88.715.310.134.554.822.419.3
MERLIN13.919.420.316.314.47.815.912.840.060.420.121.9
Mamba-2 780M-0.91.63.60.92.23.28.712.68.08.11.44.7
Vanilla1.03.27.64.25.210.212.314.736.554.55.714.1
Position1.25.110.55.76.815.515.017.442.065.112.517.9
Meanpool1.74.211.56.05.616.017.718.544.561.115.118.4
MERLIN1.96.712.67.89.120.416.119.246.072.714.220.6

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

Analysis

Positional Bias Analysis. We evaluate models on the TriviaQA task in HELMET (Yen et al., [38]), at 16K context length. This is a retrieval-augmented generation task, where the context consists of multiple passages and only one of them, which we refer to as the gold passage, is related to the final query of the prompt. The location of the gold passage is shuffled randomly within the context for each sample. Therefore, we can track the accuracy of the model’s responses depending on the location of the gold passage, and detect any positional biases that might exist. Figure 3 shows that the Vanilla baseline demonstrates a strong recency bias: it struggles unless the gold passage is located towards the end of the context. The Position baseline also demonstrates a weak recency bias, where the accuracy gradually increases as the gold passage moves closer to the end. Both the Meanpool baseline and MERLIN perform more evenly across gold passage positions, while MERLIN is better overall. This suggests that MERLIN enhances the model’s ability to extract information from important segments in the context by improving aggregation of segment-level readouts.

Position-controlled retrieval accuracy across gold-passage positions

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.

Segment Length Analysis. We examine whether MERLIN is robust to different segment lengths. Segment length, or the number of tokens per segment, determines the number of segments per context given that the entire length of the sequence is fixed (L=T/mL = T/m). The shorter the segment length, the more states per sequence, and therefore the higher the memory consumption of the model. We experiment with m∈{64,128,256,512}m \in\{64, 128, 256, 512\} for each aggregation method. At m=64m = 64, the model consumes approximately 8 times more memory compared to m=512m = 512. Figure 4 shows that for both GDN and Mamba-2, MERLIN outperforms baselines on LongBench [3] tasks across mm. In particular, MERLIN at m=256m = 256 (default) outperforms Meanpool and Position baselines at m=64m = 64, which consume 4 times more memory. This suggests that MERLIN is more memory efficient compared to its baselines.

Average score across segment lengths

Figure 4. MERLIN Outperforms Baselines Across Segment Length. Average score as the segment size that defines each state varies over {64,128,256,512}\{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.

Memory Cost and Parameter Overhead. In Section D, we compare the memory cost of different multi-state linear attention strategies with MERLIN. We observe that the extra cost of our additional second moment features is limited to at most 12.6% for Gated DeltaNet and 4.2% for Mamba-2 at dr=128d_r = 128, while the memory cost decreases quadratically with drd_r. We also compare the number of parameters that MERLIN adds to each backbone compared to baselines. We find that the MERLIN’s parameter overhead is on par with the Meanpool baseline at 4.6% for GDN and 2.2% for Mamba-2.

Qualitative Examples. In Appendix § 17, we provide examples of aggregation weight (λt,i)(\lambda_{t,i}) assignment on four HELMET (Yen et al., [38]) retrieval-augmented generation samples using MERLIN and baseline methods. As observed in the quantitative evaluations, MERLIN assigns a higher weight to the gold segment, and this supports that the variance of the alignment is critical for aggregation.

Ablations

Ablations on Dimension of Moments of Keys. We report on the impact of drd_r, the dimension of the first and second moments of keys (§ 3.2), on the performance of MERLIN. We vary dr∈{16,32,64,128}d_r \in\{16, 32, 64, 128\} during training and evaluate on LongBench (Bai et al., [3]). Figure 5 shows that the average score on LongBench tasks gradually increases as drd_r increases. This suggests that the width of the moments of keys is critical in deriving informative aggregation weights. However, there is a trade-off in deciding drd_r. The moments of keys can be more expressive when drd_r is larger. However, MERLIN consumes less memory compared to the Meanpool and Position baselines as drd_r gets smaller. More details on the memory consumption can be found in Table 9 of Section D.

Dimension of Moments line chart

Figure 5. Dimension of Moments. Average LongBench scores as drd_r varies over {16,32,64,128}\{16, 32, 64, 128\} (log axis). Performance saturates around dr=32d_r = 32 for GDN and dr=64d_r = 64 for Mamba-2.

(a) Gated DeltaNet-1.3B (dk×dv=256×512d_k \times d_v=256 \times512)(b) Mamba-2-780M (P×N=64×128P \times N=64 \times128)
Method512 (32)256 (64)128 (128)64 (256)Method512 (32)256 (64)128 (128)64 (256)
Vanilla1×1\times1×1\times1×1\times1×1\timesVanilla1×1\times1×1\times1×1\times1×1\times
Position32.0×32.0\times64.0×64.0\times128.0×128.0\times256.0×256.0\timesPosition32.0×32.0\times64.0×64.0\times128.0×128.0\times256.0×256.0\times
Meanpool32.0×32.0\times64.0×64.0\times128.0×128.0\times256.0×256.0\timesMeanpool32.0×32.0\times64.0×64.0\times128.0×128.0\times256.0×256.0\times
MERLIN (dr=16d_r=16)32.1×32.1\times64.1×64.1\times128.3×128.3\times256.5×256.5\timesMERLIN (dr=16d_r=16)32.0×32.0\times64.0×64.0\times128.1×128.1\times256.2×256.2\times
MERLIN (dr=32d_r=32)32.3×32.3\times64.5×64.5\times129.0×129.0\times258.1×258.1\timesMERLIN (dr=32d_r=32)32.1×32.1\times64.2×64.2\times128.3×128.3\times256.7×256.7\times
MERLIN (dr=64d_r=64)33.0×33.0\times66.0×66.0\times132.1×132.1\times264.1×264.1\timesMERLIN (dr=64d_r=64)32.3×32.3\times64.7×64.7\times129.4×129.4\times258.7×258.7\times
MERLIN (dr=128d_r=128)36.0×36.0\times72.1×72.1\times144.1×144.1\times288.3×288.3\timesMERLIN (dr=128d_r=128)33.3×33.3\times66.7×66.7\times133.4×133.4\times266.7×266.7\times

Table 9. Relative Decode Memory Cache (×\times the single fixed-state Vanilla cache) at a 16K context, by segment size mm (number of states L=T/mL=T/m) and moment dimension drd_r. All multi-state methods cache the LL segment states. Position caches no moments, Meanpool caches a first moment (≤0.1%\le0.1\%), and MERLIN caches a dr×drd_r \times d_r second moment. MERLIN’s overhead is larger for Gated DeltaNet (per-head router) than Mamba-2 (per-group router).

Ablations on Parameter Sharing. For each token, we project new keys and queries (kt′,qt′)(k'_t, q'_t) to construct the first and second moments, separate from the original keys and queries for the memory state (kt,qt)(k_t, q_t). We isolate the impact of decoupling the keys and queries by experimenting with a variant of MERLIN where kt′=ktk'_t = k_t and qt′=qtq'_t = q_t (shared-KQ rows in Table 4). We further experiment with a variant of MERLIN where we control the decay and update strength of the first and second moments of keys per token, similar to the decay and write-strength features of the memory state explained in Eq. (2) (write-forget rows in Table 4). We find that both ablations substantially degrade the performance of MERLIN on LongBench (Bai et al., [3]). This suggests that using a dedicated set of keys and queries from moments of keys is particularly important in deriving informative aggregation weights.

SingleMultiSummFew-shotOverall
GDN + MERLIN17.915.412.240.221.9
+ write-forget14.99.211.635.318.5
+ shared-KQ10.79.210.929.815.7
Mamba + MERLIN7.18.518.644.320.6
+ write-forget6.76.916.643.819.6
+ shared-KQ5.16.115.138.317.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.

Conclusion

We investigate the aggregation problem in multi-state linear attention, where segment-level outputs are combined via a weighted sum to produce the final output representation of each token. We propose MERLIN, an aggregation method that assigns weights to each segment based on the mean and variance of the alignment between a query and the keys of each segment. In order to compute the mean and variance of the alignment, MERLIN keeps compact first and second moments of the segment’s keys. We find that MERLIN is superior to existing aggregation methods in terms of performance and memory efficiency. MERLIN presents a generalizable and effective routing mechanism for the aggregation problem in multi-state linear attention models.

Acknowledgments

All experiments, data collection, and processing activities were conducted at Princeton University, and no Google DeepMind computing resources were used for this work. Sanghyuk Chun is supported by the Francis Robbins Upton Fellowship at Princeton University. Xindi Wu is supported by the Apple Fellowship.

AI Use Statement

In this work, we used generative AI tools to polish the writing, generate synthetic training datasets (§ 4.1), and write triton kernels for method implementation. We also used AI tools to plot figures and create a name for our method. We have not used generative AI tools to formulate and develop the conceptual framework of the method, or to identify relevant literature. We have reviewed all AI-assisted work. We take responsibility for the final content of this work, including text, claims or artifacts produced with the aid of generative AI.

Ethics Statement

Our work presents a method that improves general long-context capabilities of multi-state linear attention models. Although we do not see a clear ethical concern, we advise readers to be cautious with the datasets and tools they may use to extend our work beyond the scope of this paper.

Reproducibility Statement

We use publicly available model architectures and datasets for training and evaluation. We discuss the method details in § 3, and specifics of the training procedure in § 4.1. We also provide a full summary of model configurations and training hyperparameters in Appendix Section C. The evaluation details can be found in Section C.4.

Appendices

Notation

SymbolDescription
Sequences and softmax attention
TTSequence / context length (number of tokens)
ttPosition of the target token
xtx_tHidden state of token tt
ddKey/value dimension of one head’s memory state (written dd for both, whereas the backbones use dk×dvd_k \times d_v, see Table 9)
Q,K,V∈RT×dQ, K, V \in\mathbb{R}^{T \times d}Query, key, and value matrices
qt,kt,vtq_t, k_t, v_tQuery, key, and value of token tt
Wq,Wk,WvW_q, W_k, W_vBackbone projections producing qt,kt,vtq_t, k_t, v_t from xtx_t
MMCausal attention mask. In softmax attention it is additive with entries in {−∞,0}\{-\infty, 0\}. In linear attention it is binary, with Mij=1M_{ij} = 1 if i≥ji \ge j and 00 otherwise
O,otO, o_tAttention output, as a matrix and for token tt
softmaxRow-wise softmax normalization
⊙\odotElementwise (Hadamard) product
IIIdentity matrix
Linear attention
St∈Rd×dS_t \in\mathbb{R}^{d \times d}Recurrent memory state at step tt
αt∈(0,1)\alpha_t \in(0, 1)State decay gate in Gated DeltaNet (Eq. (2))
SymbolDescription
βt\beta_tWrite strength of the delta term (Eq. (2))
Multi-state recurrent memory
LLNumber of segments (memory states)
CiC_iThe ii-th context segment
i, ℓi,\ \ellSegment indices (1≤i,ℓ≤L1 \le i,\ell\le L)
jjIndex of a token within a segment (j∈Cij \in C_i)
m=∣Ci∣m = \lvert C_i\rvertNumber of tokens per segment
SCi∈Rd×dS_{C_i} \in\mathbb{R}^{d \times d}Memory state of segment ii. For the plain linear recurrence, SCi=∑j∈CivjkjTS_{C_i} = \sum_{j \in C_i} v_j k_j^{\mathsf{T}}
ot,io_{t,i}Segment-level output at position tt: readout of segment ii’s state, ot,i=SCiqto_{t,i} = S_{C_i}q_t
λt,i\lambda_{t,i}Aggregation weight of segment ii at position tt
λt, Wλ\lambda_t,\ W_{\lambda}Position-based weight vector (entries λt,i\lambda_{t,i}) and its projection
Wu, utW_u,\ u_tTarget-token projection and vector used by mean-key routing
kˉi\bar{k}_iMeanpooled key of segment ii
MERLIN: key moments and aggregation weights
drd_rMoment dimension (hyperparameter)
qt′, kt′∈Rdrq'_t,\ k'_t \in\mathbb{R}^{d_r}New query and key of token tt, projected from its hidden state
Wq′, Wk′W_{q'},\ W_{k'}Projections producing qt′q'_t and kt′k'_t from xtx_t
μi∈Rdr\mu_i \in\mathbb{R}^{d_r}First moment of segment ii’s keys (mean of kt′k'_t over CiC_i)
Σi∈Rdr×dr\Sigma_i \in\mathbb{R}^{d_r \times d_r}Uncentered second moment of segment ii’s keys (mean of kt′kt′Tk'_t {k'_t}^{\mathsf{T}} over CiC_i)
rt,iμr^\mu_{t,i}Mean key–query alignment for segment ii, rt,iμ=μiTqt′r^\mu_{t,i} = \mu_i^{\mathsf{T}}q'_t
rt,iΣr^\Sigma_{t,i}Variance of the key–query alignment over segment ii, rt,iΣ=qt′TΣiqt′−(rt,iμ)2r^\Sigma_{t,i} = {q'_t}^{\mathsf{T}}\Sigma_i q'_t - (r^\mu_{t,i})^2
γt\gamma_tLearned gate mixing the softmax-normalized mean and variance terms (Eq. (9)), distinct from the decay gate above
τμ, τΣ\tau_\mu,\ \tau_\SigmaLearnable temperatures for the mean term (τμ\tau_\mu) and the variance term (τΣ\tau_\Sigma)
Connection to softmax attention
st,i,js_{t,i,j}Query-key similarity, st,i,j=qt′Tkj′s_{t,i,j} = {q'_t}^{\mathsf{T}}k'_j for j∈Cij \in C_i
δ\deltaTemperature in the grouped-softmax analysis (§ 3.3)
Zt,i(δ)Z_{t,i}(\delta)Length-normalized softmax evidence of segment ii
At,iA_{t,i}Grouped softmax attention: total softmax mass on segment ii
ηt,i\eta_{t,i}Mean of st,i,js_{t,i,j} over CiC_i (first cumulant)
σt,i2\sigma^2_{t,i}Variance of st,i,js_{t,i,j} over CiC_i (second cumulant)
Mamba-2 backbone (pseudocode)
P, NP,\ NHead dimension and SSM state size, with SCi∈RP×NS_{C_i} \in\mathbb{R}^{P \times N}
Δt, A\Delta_t,\ AStep size at position tt and per-head scalar decay rate, so the state decays by exp⁡(ΔtA)\exp(\Delta_t A)
Bt, DtB_t,\ D_tInput matrix (the write key) and output matrix (the read query, written CC in Mamba-2)
WB, WD, WΔW_B,\ W_D,\ W_\DeltaProjections producing BtB_t, DtD_t, and Δt\Delta_t from xtx_t

Table 5. Glossary and Notation.

Related Work

Linear Attention and Recurrent Memory Models. Linear attention replaces the softmax normalization in standard attention with a feature map, turning attention into a recurrence with a fixed-size memory state [18]. This state can be viewed as a fast-weight associative memory written by outer products [26, 16]. Recent works improve this recurrence with data-dependent gating [35], delta-rule updates that overwrite rather than accumulate [36], and their combination [37]. Selective state-space models [13, 7], and recurrent memory models such as RWKV [24] and RetNet [29] that are closely related to linear attention share the same fixed-size state form. These models compress the context into a single state, which may be limited in memory capacity due to information loss at longer contexts. Our work builds on these backbones to train more capable multi-state linear attention and recurrent memory models.

Multi-State and Hierarchical Memory. To improve the long-context capabilities of single-state linear attention and recurrent memory models, several methods partition the context into segments and compress each segment in its own memory state. As a result, these models maintain a bank of states over the context. Memory Caching [4] segments the context uniformly at fixed token length per segment. Log-linear attention organizes the recurrent states hierarchically by growing a logarithmic number of states along a Fenwick-tree partition [14] of the sequence. Dynamic linear attention determines state boundaries based on token-level information variation while bounding the cache by merging low-information states [32]. Segment-recurrent and compressive Transformers cache or compress past activations across chunks [6, 25, 5, 33]. These works primarily focus on how states are formed via content segmentation (i.e. separation). However, we identify how the information extracted from each state is merged (i.e. aggregation) as a key bottleneck in multi-state linear attention and propose a method that improves that derivation of aggregation weights.

Routing over Memories via Aggregation. Multi-state linear attention and recurrent memory models extract information from each state and aggregate them to generate the final output representation of a token. Routing occurs during aggregation by adjusting the contribution of each segment to the final output. Mixture-of-Memories routes each token to one of several linear-attention states [8]. Sparse State Expansion routes tokens into state partitions with top-kk hard selection [22]. Gated Slot Attention reads a bounded set of slots through a softmax [39]. Gated Residual Memory scores cached memory modules with a context-aware gate, with a sparse variant that selects the most relevant ones [4]. Such routers descend from mixture-of-experts gating [27, 9] and from content-based addressing of an external memory [12, 28, 34]. These aggregation methods assign weights based on position, the magnitude of its readout, or a single pooled key. However, MERLIN keeps first and second moments of each segment’s keys to derive more informative aggregation weights.

Associative Recall and Long-Context Retrieval. Retrieving specific information from a long sequence is a core capability of models that handle long-context. In particular, such tasks evaluate the memory capacity of linear attention and recurrent memory models that compress the entire context into a single memory state. Multi-query associative recall and the Zoology analysis attribute much of the quality gap between recurrent memory models and standard transformers to recall [1], and the recall-throughput tradeoff frames recall as bounded by recurrent state size [2]. Copying and factual-recall analyses make the capacity limit of the memory state more precise [17, 21]. Furthermore, studies show that long-context models demonstrate positional biases where they prefer information located in certain depth within the context, usually close to the end [15, 38]. We use these benchmarks to demonstrate that MERLIN mitigates such biases by improving aggregation of segment-level outputs.

Experiment Details

We describe our experimental setup in detail. Section C.1 describes the model and method configurations, Section C.2 the mid-training hyperparameters, Section C.3 the datasets and benchmarks, and Section C.4 the evaluation protocol. In Section C.5 we provide the pseudocode for MERLIN.

Model Configurations

Table 6 lists the two backbones and the MERLIN hyperparameters. Gated DeltaNet-1.3B (21 layers, hidden size 2048) was pre-trained on 50B tokens of FineWeb-edu, and Mamba-2-780M (48 layers, hidden size 1536) on 300B tokens of the Pile. For both, MERLIN uses a segment size of m=256m=256 and a moment dimension of dr=128d_r=128.

Gated DeltaNet-1.3BMamba-2-780M
Parameters≈1.3B\approx1.3\mathrm{B}≈780M\approx780\mathrm{M}
Pretraining corpusFineWeb-eduPile
Pretraining tokens50B300B
Hidden size20481536
Layers2148
Heads648
Head dimension25664
SSM state size–128
Expansion factor22
Short-conv kernel44
MERLIN
Segment size mm256256
Moment dimension drd_r128128
Key/query normalizationℓ2\ell_2-normalizedℓ2\ell_2-normalized
Temperatures τμ,τΣ\tau_\mu,\tau_\Sigma, gate γt\gamma_tlearnedlearned

Table 6. Model Configurations.

Mid-Training Hyperparameters

Table 7 shows the mid-training recipe, which is shared by both backbones and by all mm. We use AdamW with a weight decay of 0.1 and a peak learning rate of 1×10−41 \times10^{-4}, which follows a cosine schedule after 305 warmup steps. Each run uses 16,384-token sequences and a global batch of 131,072 tokens for 15,259 steps, which amounts to 2.0B tokens in bfloat16 mixed precision.

HyperparameterValue
OptimizerAdamW (ϵ=10−15\epsilon= 10^{-15})
Weight decay0.1
Peak learning rate1×10−41 \times10^{-4}
Learning-rate schedulecosine (min ratio 0.1)
Warmup steps305
Gradient clipping (max norm)1.0
Sequence length16,384
Global batch size131,072 tokens
Training steps15,259
Total tokens2.0B
Precisionbfloat16 (mixed)
Random seed42
DataLong-Data-Collections, augmented with RULER NIAH/QA

Table 7. Mid-Training Hyperparameters.

Datasets and Benchmarks

Table 8 summarizes the mid-training data and the evaluation benchmarks, together with the metric, context length, and number of samples for each. We evaluate long-context retrieval on RULER ([15]) and long-context understanding on LongBench ([3]), both at up to a 16K context.

PhaseDataset / TaskMetricContextSize
Mid-trainingLong-Data-Collections (+ RULER NIAH/QA)–16K2.0B tokens
RULER (retrieval)S-NIAH-1/2/3recall16K500 each
RULER (retrieval)MK-NIAH-1recall16K500
RULER (retrieval)MQ-NIAHrecall16K500
RULER (retrieval)MV-NIAHrecall16K500
RULER (retrieval)HotpotQArecall16K500
RULER (retrieval)SQuADQArecall16K500
LongBenchNarrativeQA (NQ)F1≤16K61
LongBenchQasper (QR)F1≤16K197
LongBenchMultiFieldQA (MF)F1≤16K150
LongBenchHotpotQA (HP)F1≤16K161
LongBench2WikiMQA (2W)F1≤16K198
LongBenchGovReport (GR)ROUGE≤16K174
LongBenchQMSum (QM)ROUGE≤16K116
LongBenchMultiNews (MN)ROUGE≤16K200
LongBenchSAMSum (SS)ROUGE≤16K172
LongBenchTREC (TC)accuracy≤16K200
LongBenchTriviaQA (TQ)F1≤16K142

Table 8. Summary of Datasets and Evaluation Benchmarks. We mid-train on Long-Data-Collections augmented with RULER-style NIAH/QA, and evaluate retrieval on RULER and long-context understanding on LongBench (filtered to ≤16K tokens). Counts vary slightly by backbone tokenizer.

Evaluation Protocol

All evaluations use a 16K context length. We evaluate on RULER ([15]) and LongBench ([3]) following the official metrics and configurations from Language Model Evaluation Harness ([11]). For RULER, we run the eight NIAH and QA tasks listed in Table 8 and report the score on each task together with the average across tasks. For LongBench, we run the eleven subtasks in Table 8, keeping only the samples of at most 16K tokens, and score each task with its standard metric (F1, ROUGE, or accuracy).

For HELMET ([38]), we use the TriviaQA retrieval-augmented generation task and follow the official evaluation setup. Each prompt contains several passages, and only one of them, the gold passage, contains the answer to the question. We randomly shuffle the position of the gold passage within the context and report accuracy as a function of that position, from the first sixth (1/6) to the last sixth (6/6) of the context.

Pseudocode

Algorithm 1 shows MERLIN on a Gated DeltaNet backbone, assuming equal-size segments of mm tokens each. The outer loop iterates over the LL segments, and the inner loop iterates over the mm tokens within each segment. For each token, the algorithm first updates the memory state of its segment together with the first and second moments of the segment’s keys, so both accumulate token by token. It then reads out from every visible segment, meaning all closed segments plus the current, partially built one. Finally, it aggregates these readouts using the moment-derived weights from Eq. (9).

Algorithm 2 shows MERLIN on a Mamba-2 backbone. The structure is identical to Algorithm 1, but differs only on the memory state recurrence. Each token’s hidden state is projected to the SSM

Algorithm 1 MERLIN (Gated DeltaNet backbone) with equal-size segments C1,…,CLC_1,\ldots,C_L of mm tokens, learnable temperatures τμ,τΣ\tau_\mu,\tau_\Sigma, and a learnable gate γt\gamma_t.

  1. for i=1,…,Li=1,\ldots,L do

  2. SCi,μi,Σi←0,0,0S_{C_i},\mu_i,\Sigma_i\gets0,0,0

  3. for t∈Cit\in C_i do

  4. qt,kt,vt←Wqxt,Wkxt,Wvxtq_t,k_t,v_t\gets W_qx_t,W_kx_t,W_vx_t

  5. qt′,kt′←Wq′xt,Wk′xtq'_t,k'_t\gets W_{q'}x_t,W_{k'}x_t

  6. SCi←SCi(αt(I−βtktkt⊤))+βtvtkt⊤S_{C_i}\gets S_{C_i}\left(\alpha_t(I-\beta_tk_tk_t^\top)\right)+\beta_tv_tk_t^\top ▷ GDN memory state, Eq. (2)

  7. μi←μi+1mkt′,Σi←Σi+1mkt′kt′⊤\mu_i\gets\mu_i+\frac{1}{m}k'_t,\quad\Sigma_i\gets\Sigma_i+\frac{1}{m}k'_tk_t^{\prime\top} ▷ moments of keys

  8. for ℓ=1,…,i\ell=1,\ldots,i do

  9. ot,ℓ←SCℓqt,rt,ℓμ←μℓ⊤qt′,rt,ℓΣ←qt′⊤Σℓqt′−(rt,ℓμ)2o_{t,\ell}\gets S_{C_\ell}q_t,\quad r^\mu_{t,\ell}\gets\mu_\ell^\top q'_t,\quad r^\Sigma_{t,\ell}\gets q_t^{\prime\top}\Sigma_\ell q'_t-(r^\mu_{t,\ell})^2

  10. end for

  11. λt,ℓ←γtsoftmax⁡ℓ(τμrt,ℓμ)+(1−γt)softmax⁡ℓ(τΣrt,ℓΣ)\lambda_{t,\ell}\gets\gamma_t\operatorname{softmax}_\ell(\tau_\mu r^\mu_{t,\ell})+(1-\gamma_t)\operatorname{softmax}_\ell(\tau_\Sigma r^\Sigma_{t,\ell}) ▷ Eq. (9)

  12. ot←∑ℓ≤iλt,ℓot,ℓo_t\gets\sum_{\ell\leq i}\lambda_{t,\ell}o_{t,\ell} ▷ aggregate readouts, Eq. (4)

  13. end for

  14. end for

  15. return o1,…,oTo_1,\ldots,o_T

Algorithm 2 MERLIN (Mamba-2 backbone) with equal-size segments C1,…,CLC_1,\ldots,C_L of mm tokens, learnable temperatures τμ,τΣ\tau_\mu,\tau_\Sigma, and a learnable gate γt\gamma_t. The state SCi∈RP×NS_{C_i}\in\mathbb{R}^{P\times N}, and the cached moments are shared per group.

  1. for i=1,…,Li=1,\ldots,L do

  2. SCi,μi,Σi←0,0,0S_{C_i},\mu_i,\Sigma_i\gets0,0,0

  3. for t∈Cit\in C_i do

  4. vt,Bt,Dt,Δt←Wvxt,WBxt,WDxt,softplus⁡(WΔxt)v_t,B_t,D_t,\Delta_t\gets W_vx_t,W_Bx_t,W_Dx_t,\operatorname{softplus}(W_\Delta x_t)

  5. qt′,kt′←Wq′xt,Wk′xtq'_t,k'_t\gets W_{q'}x_t,W_{k'}x_t

  6. SCi←exp⁡(ΔtA)SCi+ΔtvtBt⊤S_{C_i}\gets\exp(\Delta_tA)S_{C_i}+\Delta_tv_tB_t^\top ▷ Mamba-2 SSD state

  7. μi←μi+1mkt′,Σi←Σi+1mkt′kt′⊤\mu_i\gets\mu_i+\frac{1}{m}k'_t,\quad\Sigma_i\gets\Sigma_i+\frac{1}{m}k'_tk_t^{\prime\top} ▷ moments of keys

  8. for ℓ=1,…,i\ell=1,\ldots,i do

  9. ot,ℓ←SCℓDt,rt,ℓμ←μℓ⊤qt′,rt,ℓΣ←qt′⊤Σℓqt′−(rt,ℓμ)2o_{t,\ell}\gets S_{C_\ell}D_t,\quad r^\mu_{t,\ell}\gets\mu_\ell^\top q'_t,\quad r^\Sigma_{t,\ell}\gets q_t^{\prime\top}\Sigma_\ell q'_t-(r^\mu_{t,\ell})^2

  10. end for

  11. λt,ℓ←γtsoftmax⁡ℓ(τμrt,ℓμ)+(1−γt)softmax⁡ℓ(τΣrt,ℓΣ)\lambda_{t,\ell}\gets\gamma_t\operatorname{softmax}_\ell(\tau_\mu r^\mu_{t,\ell})+(1-\gamma_t)\operatorname{softmax}_\ell(\tau_\Sigma r^\Sigma_{t,\ell}) ▷ Eq. (9)

  12. ot←∑ℓ≤iλt,ℓot,ℓo_t\gets\sum_{\ell\leq i}\lambda_{t,\ell}o_{t,\ell} ▷ aggregate readouts, Eq. (4)

  13. end for

  14. end for

  15. return o1,…,oTo_1,\ldots,o_T

value vtv_t, input matrix BtB_t (the write key), output matrix DtD_t (the read query), and step size Δt\Delta_t. The segment state SCi∈RP×NS_{C_i}\in\mathbb{R}^{P\times N} (PP the head dimension, NN the SSM state size) follows the Mamba-2 SSD recurrence with per-head scalar decay exp⁡(ΔtA)\exp(\Delta_tA). In Mamba-2, the cached moments are shared per group. The aggregation step is identical to Eq. (9).

Additional Analysis

Memory Cost. Table 9 reports the recurrent memory retained to decode each token, relative to the single fixed-state (Vanilla) cache. All multi-state methods cache the L=T/mL=T/m segment states, which dominate the cost. Position adds no cached moments, and Meanpool adds a negligible first moment (≤0.1%\leq0.1\%), while MERLIN additionally caches a dr×drd_r\times d_r second moment. This extra cost is controlled by drd_r and is at most 12.6% for Gated DeltaNet and 4.2% for Mamba-2 (whose moments are shared per group) at dr=128d_r=128, shrinking quadratically as drd_r decreases.

Parameter Overhead. We show the number of parameters that each aggregation method adds to the backbone under default model configurations (m=256m=256, dr=128d_r=128). The parameter overhead of MERLIN roughly equals the Meanpool baseline (4.5% for Gated DeltaNet-1.3B and 2.2% for

Mamba-2-780M). The Position baseline is relatively lightweight, with only 1.2% and 0.6% overhead for Gated DeltaNet-1.3B and Mamba-2-780M respectively.

BackboneAggregationAdded%
GatedDeltaNet-1.3BVanilla00.0%
GatedDeltaNet-1.3BPOSITION16.8M1.2%
GatedDeltaNet-1.3BMEANPOOL66.1M4.5%
GatedDeltaNet-1.3BMERLIN66.3M4.6%
Mamba-2-780MVanilla00.0%
Mamba-2-780MPOSITION4.8M0.6%
Mamba-2-780MMEANPOOL18.9M2.2%
Mamba-2-780MMERLIN18.9M2.2%

Table 10. Parameter Overhead of the Aggregation Mechanisms. Parameters added in the aggregation pathway on top of each shared backbone. We assume default model configurations where m=256m=256, dr=128d_r=128.

Qualitative Examples. Table 11 illustrates the aggregation weight assignment on individual segments given retrieval-augmented generation samples in HELMET Yen et al. [38] (two Natural Questions, two TriviaQA, all at 16K context). For each query we show the gold segment that contains the ground truth as well as four randomly sampled distractor (irrelevant) segments, and the aggregation weight λt,i\lambda_{t,i} each receives in the last layer. MERLIN assigns highest weights on the gold segment, well above the recency-driven POSITION router and above MEANPOOL.

IndexTypeCandidate segmentPOSITIONMEANPOOLMERLIN
Question: “what is another name for the water cycle?” Answer: hydrological cycle
9Distractor. . . of the world’s water supply, about 321,000,000 cubic miles of it. . .0.0060.0180.017
22Distractor. . . water vapour or aqueous vapor is the gaseous phase of water. . .0.0050.0190.016
33Ground Truth. . . also known as the hydrologic cycle or the hydrological cycle, the continuous movement of water. . .0.0050.0300.055
35Distractor. . . cultural eutrophication of lakes is primarily due to phosphorus. . .0.0050.0080.005
54Distractor. . . loch or fen; each has a special meaning, though modern usage. . .0.0050.0100.007
Question: “where are the Washington Redskins based out of?” Answer: the Washington metropolitan area
8Distractor. . . Franchise history: the Daniel Snyder ownership era (1999–present). . .0.0040.0100.012
28Distractor. . . Daniel Snyder . . . W. Smith, the founder of Federal Express. . .0.0040.0160.013
37Ground Truth. . . a professional football team based in the Washington metropolitan area; they compete in the NFL. . .0.0040.0340.050
39Distractor. . . History. The Redskins relocated from Boston, Massachusetts. . .0.0040.0180.031
50Distractor. . . a home for that team; its highest-ranking teams as of 2015. . .0.0040.0300.027
Question: “in which Charles Dickens novel does the character Mr M Choakumchild appear?” Answer: Hard Times
16Distractor. . . the illiterate poor chipped in ha’pennies to have each new monthly episode. . .0.0040.0100.024
26Ground Truth. . . where he wrote “Bleak House” (1852–53) and “Hard Times” (1854). . .0.0040.0210.060
42Distractor. . . Mrs Nickleby is based on his mother, though she did not recognize herself. . .0.0040.0120.010
45Distractor. . . the novel “Nicholas Nickleby”, and Mrs Micawber in “David Copperfield”. . .0.0040.0230.010
53Distractor. . . Dickens in films and TV: Lennox Pawle (1935). . .0.0040.0190.007
Question: “which chain of European retailers began near a crossroads in Annecy, France in 1958?” Answer: Carrefour
6Distractor. . . the family also holds a stake in 3 Susines, Norauto and others. . .0.0070.0210.016
25Ground Truth. . . Carrefour S.A. is a French multinational retailer. . .0.0070.0190.049
31Distractor. . . within its metropolitan boundaries; the Cagliari–Sarroch port system. . .0.0070.0130.016
37Distractor. . . 17 September 2012 . . . Jacqueline de Jong; an internment camp. . .0.0070.0080.006
42Distractor. . . some trades are also present along the avenue de Verdun. . .0.0070.0110.007

Table 11. Qualitative Examples. We evaluate the GDN-based models on 4 random RAG samples (2 from Natural Questions and 2 from TriviaQA) in HELMET ([38]) 16K. We show the last layer’s aggregation weights assigned to the ground truth and distractor segments for each aggregation method. MERLIN assigns highest weights to the ground truth segment compared to the Meanpool and Position baselines.

Limitations

Due to compute constraints, our experiments cover two backbones (Gated DeltaNet-1.3B and Mamba-2-780M) at a 16K context length, and we mid-train multi-state models from single-state checkpoints. We compare MERLIN against Vanilla, Position, and Meanpool aggregation with uniform segmentation. Studying larger models, longer contexts, and a broader set of aggregation baselines is beyond our current compute budget, and we leave it to future work.

References

  1. [1]Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Y Zou, Atri Rudra, and Christopher Ré. Zoology: Measuring and improving recall in efficient language models. In International conference on learning representations, volume 2024, pp. 15664–15730, 2024a.arxiv.org/abs/2312.04927
  2. [2]Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. Simple linear attention language models balance the recall-throughput tradeoff. arXiv preprint arXiv:2402.18668, 2024b.arxiv.org/abs/2402.18668
  3. [3]Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al. Longbench: A bilingual, multitask benchmark for long context understanding. In Proceedings of the 62nd annual meeting of the association for computational linguistics (volume 1: Long papers), pp. 3119–3137, 2024.arxiv.org/abs/2308.14508
  4. [4]Ali Behrouz, Zeman Li, Yuan Deng, Peilin Zhong, Meisam Razaviyayn, and Vahab Mirrokni. Memory caching: Rnns with growing memory. arXiv preprint arXiv:2602.24281, 2026.arxiv.org/abs/2602.24281
  5. [5]Aydar Bulatov, Yury Kuratov, and Mikhail Burtsev. Recurrent memory transformer. Advances in neural information processing systems, 35:11079–11091, 2022.arxiv.org/abs/2207.06881
  6. [6]Zihang Dai, Zhilin Yang, Yiming Yang, Jaime G Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. In Proceedings of the 57th annual meeting of the association for computational linguistics, pp. 2978–2988, 2019.arxiv.org/abs/1901.02860
  7. [7]Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.21060, 2024.arxiv.org/abs/2405.21060
  8. [8]Jusen Du, Weigao Sun, Disen Lan, Jiaxi Hu, Zhang Tao, and Yu Cheng. Mom: Linear sequence modeling with mixture-of-memories. In International Conference on Learning Representations, volume 2026, pp. 106613–106631, 2026.arxiv.org/abs/2502.13685
  9. [9]William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. arXiv 2021. arXiv preprint arXiv:2101.03961, 2021.arxiv.org/abs/2101.03961
  10. [10]Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020.arxiv.org/abs/2101.00027
  11. [11]Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac’h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. The language model evaluation harness, 07 2024.
  12. [12]Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv preprint arXiv:1410.5401, 2014.arxiv.org/abs/1410.5401
  13. [13]Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023.arxiv.org/abs/2312.00752
  14. [14]Guo Guo, Songlin Yang, Tarushii Goel, Eric P Xing, Tri Dao, and Yoon Kim. Log-linear attention. In International Conference on Learning Representations, volume 2026, pp. 91296–91317, 2026.arxiv.org/abs/2506.04761
  15. [15]Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What’s the real context size of your long-context language models? arXiv preprint arXiv:2404.06654, 2024.arxiv.org/abs/2404.06654
  16. [16]Hee Seung Hwang, Xindi Wu, Sanghyuk Chun, and Olga Russakovsky. Reinforced fast weights with next-sequence prediction. arXiv preprint arXiv:2602.16704, 2026.arxiv.org/abs/2602.16704
  17. [17]Samy Jelassi, David Brandfonbrener, Sham M Kakade, and Eran Malach. Repeat after me: Transformers are better than state space models at copying. arXiv preprint arXiv:2402.01032, 2024.arxiv.org/abs/2402.01032
  18. [18]Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pp. 5156–5165. PMLR, 2020.arxiv.org/abs/2006.16236
  19. [19]Feyza Duman Keles, Pruthuvi Mahesakya Wijewardena, and Chinmay Hegde. On the computational complexity of self-attention. In International conference on algorithmic learning theory, pp. 597–619. PMLR, 2023.arxiv.org/abs/2209.04881
  20. [20]Yuri Kuratov, Aydar Bulatov, Petr Anokhin, Ivan Rodkin, Dmitry Sorokin, Artyom Sorokin, and Mikhail Burtsev. Babilong: Testing the limits of llms with long context reasoning-in-a-haystack. Advances in Neural Information Processing Systems, 37:106519–106554, 2024.arxiv.org/abs/2406.10149
  21. [21]Eshaan Nichani, Jason Lee, and Alberto Bietti. Understanding factual recall in transformers via associative memories. In International Conference on Learning Representations, volume 2025, pp. 4166–4207, 2025.arxiv.org/abs/2412.06538
  22. [22]Yuqi Pan, Yongqi An, Zheng Li, Yuhong Chou, Ruijie Zhu, Xiaohui Wang, Mingxuan Wang, Jinqiao Wang, and Guoqi Li. Scaling linear attention with sparse state expansion. arXiv preprint arXiv:2507.16577, 2025.arxiv.org/abs/2507.16577
  23. [23]Guilherme Penedo, Hynek Kydlíček, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, Thomas Wolf, et al. The fineweb datasets: Decanting the web for the finest text data at scale. Advances in Neural Information Processing Systems, 37:30811–30849, 2024.arxiv.org/abs/2406.17557
  24. [24]Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Leon Derczynski, et al. Rwkv: Reinventing rnns for the transformer era. In Findings of the association for computational linguistics: EMNLP 2023, pp. 14048–14077, 2023.arxiv.org/abs/2305.13048
  25. [25]Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. arXiv preprint arXiv:1911.05507, 2019.arxiv.org/abs/1911.05507
  26. [26]Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International conference on machine learning, pp. 9355–9366. PMLR, 2021.arxiv.org/abs/2102.11174
  27. [27]Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06538, 2017.arxiv.org/abs/1701.06538
  28. [28]Sainbayar Sukhbaatar, Jason Weston, Rob Fergus, et al. End-to-end memory networks. Advances in neural information processing systems, 28, 2015.arxiv.org/abs/1503.08895
  29. [29]Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307.08621, 2023.arxiv.org/abs/2307.08621
  30. [30]TogetherAI. Long data collections database. https://huggingface.co/datasets/togethercomputer/Long-Data-Collections, 2024.
  31. [31]Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.arxiv.org/abs/1706.03762
  32. [32]Xin Wang, Hui Shen, Boyuan Zheng, Xueshen Liu, Minkyoung Cho, Zhongwei Wan, Zesen Zhao, Zhuoqing Mao, Shen Yan, and Mi Zhang. Dynamic linear attention. arXiv preprint arXiv:2606.10650, 2026.arxiv.org/abs/2606.10650
  33. [33]Xindi Wu, Sven Elflein, James Lucas, Olga Russakovsky, Laura Leal-Taixé, Despoina Pachalidou, Jonathan Lorraine, and Aljoša Ošep. Addressable memory for video world models. arXiv preprint arXiv:2608.07408, 2026.arxiv.org/abs/2608.07408
  34. [34]Yuhuai Wu, Markus N Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing transformers. arXiv preprint arXiv:2203.08913, 2022.arxiv.org/abs/2203.08913
  35. [35]Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training. arXiv preprint arXiv:2312.06635, 2023.
  36. [36]Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transformers with the delta rule over sequence length. Advances in neural information processing systems, 37:115491–115522, 2024.arxiv.org/abs/2406.06484
  37. [37]Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule. In International Conference on Learning Representations, volume 2025, pp. 29687–29707, 2025.arxiv.org/abs/2412.06464
  38. [38]Howard Yen, Tianyu Gao, Minmin Hou, Ke Ding, Daniel Fleischer, Peter Iszak, Moshe Wasserblat, and Danqi Chen. Helmet: How to evaluate long-context models effectively and thoroughly. In International Conference on Learning Representations, volume 2025, pp. 98914–98965, 2025.arxiv.org/abs/2410.02694
  39. [39]Yu Zhang, Songlin Yang, Ruijie Zhu, Yue Zhang, Leyang Cui, Yiqiao Wang, Bolun Wang, Freda Shi, Bailin Wang, Wei Bi, et al. Gated slot attention for efficient linear-time sequence modeling. Advances in Neural Information Processing Systems, 37:116870–116898, 2024.arxiv.org/abs/2409.07146

Paper details

vlnc:W_zydnf4v8
v1 · latestHistory (1)

published · Fresh — added to the Hub within the last 30 days

added ·

Usage

views · 0

downloads · 0

A view counts once per browser per day, a download each time the PDF is requested.

CC BY 4.0

Original source

  • Computer Science

generative AI tools · View statement →

Looking for related papers…
Contents