Dynamic Linear Attention

Xin Wang$^{1}$ $^{}$, Hui Shen$^{2}$ $^{}$, Boyuan Zheng$^{2}$, Xueshen Liu$^{2}$, Minkyoung Cho$^{2}$, Zhongwei Wan$^{1}$, Zesen Zhao$^{2}$, Zhuoqing Mao$^{2}$, Shen Yan$^{3}$, Mi Zhang$^{1}$
$^{1}$The Ohio State University, $^{2}$University of Michigan, $^{3}$ByteDance Seed

$^{*}$Equal contribution.

Abstract

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.

Correspondence: Xin Wang at [email protected], Mi Zhang at [email protected]

Executive Summary: Large language models struggle to handle long contexts efficiently because standard attention scales quadratically with sequence length, driving up compute and memory costs. Linear attention offers a faster alternative, yet recent multi-state versions still lose accuracy on long inputs. Their reliance on fixed merging rules mixes important and redundant tokens indiscriminately, producing irreversible compression errors that grow with context length.

This paper introduces Dynamic Linear Attention (DLA) to fix that mismatch. The authors pre-train two linear-attention backbones—Mamba-2 (780 M) and Gated DeltaNet (1.3 B)—using an information-aware merging rule that creates new memory states only when token content drifts significantly, plus a fixed-size cache that merges the least informative adjacent states to keep memory predictable. They evaluate the resulting models on 16 tasks that span commonsense reasoning, in-context retrieval, and long-context benchmarks, using the same 50 B-token pre-training regime as prior work.

DLA improves accuracy over the leading multi-state baseline (Log-Linear Attention) on every task, with relative gains reaching 52 % on commonsense benchmarks and 350 % on harder multi-needle retrieval tests. On Mamba-2, the DLA variant matches or exceeds a full-attention Transformer of similar size. Inference throughput rises and peak memory falls relative to Log-Linear Attention once batch size or context length grows. Ablations confirm that both the dynamic-merging and capacity-bounded components contribute to these gains, and results remain stable across modest changes in cache size or decision thresholds.

These outcomes matter because they narrow the accuracy gap between fast linear models and full attention while preserving the predictable cost needed for production serving. The approach therefore offers a practical route to longer contexts without proportional increases in hardware demand.

Adopters should integrate DLA-style merging into new linear-attention training runs and test it at larger scales. Further work is needed on sequences beyond 16 k tokens and on integration with retrieval-augmented pipelines before broad deployment decisions are made. The main uncertainties are the academic-scale training budget and the fixed cache size of 30 states; both may need re-tuning for production workloads.

1. Introduction

Section Summary: Large language models struggle to handle very long texts efficiently because standard attention mechanisms become too slow and memory-heavy as sequences grow. Existing linear-attention approaches try to fix this by compressing history into a small number of memory states, but they rely on rigid, fixed-size merging rules that cannot adapt to changes in information density and therefore lose important details over long contexts. This paper introduces Dynamic Linear Attention, which builds and merges memory states on the fly according to how much new tokens differ from what is already stored, while strictly limiting the total number of states to keep computation predictable; experiments show the method outperforms prior linear-attention techniques and can even match full-attention transformers on many tasks.

Large Language Models (LLMs) have demonstrated remarkable capabilities across a wide range of natural language understanding and generation tasks. However, scaling LLMs to long-context settings remains a fundamental challenge due to the quadratic computational and memory complexity of standard self-attention ([1, 2, 3]). This limitation has motivated extensive research on efficient attention mechanisms that enable long-sequence modeling without retraining from scratch. Among these approaches, linear attention ([4, 5]) has emerged as a promising direction, as it approximates full attention with sub-quadratic complexity and offers favorable scalability to long contexts.

To further improve the representation capacity of linear attention under long sequences, recent works organize historical context in a multi-state manner, where long token histories are partitioned into chunks and summarized into compact memory states. Representative methods such as Log-Linear Attention ([6]) demonstrate improved efficiency and practicality for long-context inference. By operating on summarized states rather than individual tokens, these approaches significantly reduce memory footprint and computation cost.

Despite their success, existing multi-state linear attention methods still suffer from notable performance degradation as context length increases. This limitation stems from a fundamental mismatch between fixed memory construction policies and the non-uniform, dynamically evolving information structure of long sequences. In particular, current methods typically rely on fixed block sizes or rule-based merging schedules, implicitly assuming uniform information density across the sequence. Such designs fail to adapt to dynamically emerging semantic transitions, forcing critical tokens to be prematurely absorbed into coarse summaries. Moreover, merge decisions made under fixed policies are irreversible: once heterogeneous tokens are compressed into a single state, their individual contributions cannot be recovered, leading to error accumulation.

These observations suggest that effective long-context linear attention requires memory modeling mechanisms that are both information-aware and capacity-controlled. On one hand, state construction should adapt to local representation variation, allocating higher resolution to semantically volatile regions while aggressively summarizing stable spans. On the other hand, the total number of memory states must be explicitly bounded to ensure predictable computation and memory cost during inference.

In this work, we propose Dynamic Linear Attention (DLA), a new framework for multi-state linear attention that addresses these challenges. DLA differs from prior approaches in two key aspects. First, DLA introduces Information-Aware Dynamic State Merging, which determines state boundaries on the fly based on token-level information variation. Instead of relying on fixed merging policies, DLA evaluates the representation change of each incoming token relative to the current memory state, merging low-variation tokens while initiating new states at semantic transition points. Second, DLA incorporates Capacity-Bounded Memory Modeling, which maintains a fixed-size, chronologically ordered state cache. When the cache reaches its capacity, DLA selectively merges adjacent low-information states, preserving temporal order while minimizing information loss.

We pre-train DLA on two linear-attention backbones, Mamba-2-780M and Gated DeltaNet-1.3B, following the design in ([6]). We evaluate DLA on 16 datasets spanning three aspects: eight commonsense reasoning benchmarks, six in-context retrieval datasets, and two long-context modeling datasets. We highlight three main findings. (1) DLA consistently outperforms the state-of-the-art multi-state method, Log-Linear Attention, across all tasks. (2) When applied to Mamba-2, the DLA variant even achieves performance comparable to full-attention Transformers with similar parameter budgets. (3) DLA achieves superior efficiency, delivering higher throughput and lower runtime memory consumption than Log-Linear Attention.

2. Preliminary

Section Summary: This section introduces efficient alternatives to standard self-attention for sequence modeling, which normally requires quadratic time and memory due to its pairwise computations. Linear attention simplifies the process by dropping the softmax operation, allowing past information to be summarized in a single evolving state matrix that supports fast, constant-memory inference, though at the cost of limited representational power over long sequences. Variants such as DeltaNet add a data-dependent forgetting mechanism to this single state, while multi-state approaches like log-linear attention maintain multiple hierarchical states to capture context at different scales, achieving better capacity with near-linear complexity but relying on rigid, non-adaptive partitioning of history.

We consider a sequence modeling task with input length $T$ and hidden dimension $d$. Let $\mathbf{Q}, \mathbf{K}, \mathbf{V} \in \mathbb{R}^{T \times d}$ denote the query, key, and value matrices. Standard self-attention computes the output $\mathbf{O} \in \mathbb{R}^{T \times d}$ as

$ \mathbf{O} = \mathrm{softmax}(\mathbf{Q}\mathbf{K}^\top \odot \mathbf{M})\mathbf{V}, $

where $\mathbf{M}$ is the causal mask. While effective, this operation incurs quadratic time and memory complexity in $T$, motivating the development of sub-quadratic attention mechanisms. In this section, we review linear attention and its multi-state variants that form the foundation of our approach.

2.1 Linear Attention

Linear attention mitigates the quadratic cost of Transformers by removing the softmax normalization, enabling the reordering of computation via associativity. A causal linear attention layer can be written in a parallel form as

$ \mathbf{O} = (\mathbf{Q}\mathbf{K}^\top \odot \mathbf{M})\mathbf{V}, \qquad \mathbf{M}_{ij} = \mathbb{I}(i \ge j).\tag{1} $

This formulation admits an equivalent recurrent implementation. Let $\mathbf{q}_t, \mathbf{k}_t, \mathbf{v}_t \in \mathbb{R}^d$ denote the query, key, and value vectors at time step $t$. Linear attention maintains a single state matrix $\mathbf{S}_t \in \mathbb{R}^{d \times d}$ that summarizes all past tokens:

$ \begin{aligned} \mathbf{S}t &= \mathbf{S}{t-1} + \mathbf{v}_t \mathbf{k}_t^\top, \ \mathbf{o}_t &= \mathbf{S}_t \mathbf{q}_t. \end{aligned} $

This recurrent form enables linear-time inference with constant memory, but compresses the entire history into a single state, which can limit representation capacity under long contexts. We use $\phi(\cdot)$ to denote the feature map used in linear attention. Unless otherwise specified, $\phi:\mathbb{R}^d \rightarrow \mathbb{R}^d$ is implemented as an identity mapping or a learnable linear projection, following prior work.

2.2 Linear Attention with the Delta Rule

To improve state tracking and introduce controlled forgetting, DeltaNet ([4]) extends linear attention with a delta-style update rule:

$ \mathbf{S}_t

\mathbf{S}_{t-1}(\mathbf{I} - \beta_t \mathbf{k}_t \mathbf{k}_t^\top)

• \mathbf{v}_t \mathbf{k}_t^\top,\tag{2} $

where $\beta_t$ is a data-dependent step size. While this formulation improves adaptivity over a pure accumulator, it still relies on a single global state and therefore cannot selectively preserve fine-grained information over long sequences.

2.3 Multi-State Linear Attention

To increase modeling capacity while retaining sub-quadratic complexity, recent work organizes linear attention in a multi-state manner by partitioning the historical context into segments and summarizing each segment into a separate state ([6, 7]). Among them, log-linear attention ([6]) replaces the single recurrent state with a logarithmic number of multi-scale states constructed via a Fenwick-tree decomposition of the causal prefix. Concretely, at time step $t$, the prefix $[0, t]$ is decomposed into at most $L=\lceil \log_2(t+1)\rceil+1$ disjoint buckets ${B_t^{(\ell)}}_{\ell=0}^{L-1}$, with finer resolution near the current position and coarser resolution for distant history. The corresponding linear attention states and the final aggregated output are computed as:

$ \mathbf{S}t^{(\ell)} = \sum{s \in B_t^{(\ell)}} \mathbf{v}_s \mathbf{k}_s^\top \in \mathbb{R}^{d \times d}, \mathbf{o}t = \sum{\ell=0}^{L-1} \lambda_t^{(\ell)} , \mathbf{S}_t^{(\ell)} \mathbf{q}_t\tag{3} $

This design achieves $O(T\log T)$ training complexity and $O(\log T)$ time and memory per decoding step. However, the granularity of its memory states is determined by a fixed hierarchical schedule, independent of token-level representation variation. As a result, semantically salient tokens may be prematurely absorbed into coarse summaries, and disturbances introduced at critical positions can propagate through the fixed multi-scale states. This limitation motivates the need for information-aware and adaptive memory construction, which we address in the next section.

3. Dynamic Linear Attention (DLA)

Section Summary: Dynamic Linear Attention is a method that builds a compact, fixed-size memory of past information while processing long sequences, allowing efficient attention without growing compute costs. Instead of merging tokens on a fixed schedule, it measures how much each new token changes the current summary and either folds it in or starts a fresh state when meaningful shifts occur. When memory fills up, the least informative adjacent states are combined, so the model keeps resolution where the content is changing and compresses stable stretches, producing output by attending over these adaptive states with a simple linear formula.

Figure 1 provides an overview of DLA. DLA is an information-aware linear attention framework that dynamically constructs a compact set of memory states for efficient long-context modeling. Unlike prior approaches that rely on fixed temporal schedules or predefined block boundaries, DLA adaptively determines state granularity based on token-level information variation. Specifically, tokens are processed sequentially. For each new token, DLA computes a lightweight State Information Score measuring its representation change relative to the most recent memory state. Tokens with low information variation are merged into the current state, while tokens exhibiting significant drift initiate a new state. This enables fine-grained modeling around semantic transitions while aggressively summarizing stable token spans. To bound memory and computation, DLA maintains a capacity-limited state cache. When the cache reaches its maximum size, two adjacent states with the lowest information density are merged, preserving temporal order while minimizing information loss. The resulting memory consists of a fixed-size, chronologically ordered set of summary states. At each decoding step, DLA produces the output by attending over the maintained memory states using a linear attention formulation, where each state contributes with a query-dependent weight. Together, information-aware state construction and capacity-bounded memory modeling enable DLA to achieve adaptive resolution, stable inference cost, and efficient long-context representation.

3.1 Information-Aware Dynamic State Merging

Motivation: Existing multi-block linear attention methods typically rely on fixed schedules (e.g., block and merge every $K$ tokens) ([6]) or hard, rule-based boundaries ([7]) to determine the block of historical tokens that should be merged into summary states. While such designs improve memory and compute efficiency, they are largely agnostic to the semantic evolution of the sequence. In practice, information density is highly non-uniform: critical semantic transitions may occur abruptly, whereas long stretches of tokens can be locally redundant. As a result, fixed or hard block policies often suffer from two key limitations. First, they cannot adapt to dynamically emerging semantic changes, forcing important transitions to be prematurely absorbed into coarse summaries simply because a pre-defined boundary is reached. Second, merge decisions made without regard to local semantic continuity are inherently irreversible: once tokens are merged under a fixed policy, their individual contributions cannot be recovered, even if subsequent context reveals their importance. These misalignments between merge decisions and the true semantic structure lead to sub-optimal generation and finally degrades the representation quality.

In the following, we provide a theoretical proof on why the fixed merging policy is sub-optimal.

########## {caption="Theorem 1: State deviation"}

Let ${u_t}_{t=1}^T \subset \mathbb{R}^d$ denote per-token additive contributions to a linear attention state. Consider a blocking policy $\pi$ of token list ${1, \dots, T}$ into $m$ disjoint contiguous blocks ${\mathcal{C}i}{i=1}^m$. For each block $\mathcal{C}_i$, let $\bar{u}i \in \mathbb{R}^d$ be a representative summary vector. Therefore, for any query vector $q \in \mathbb{R}^d$, the exact output $y(q)$ and the summarized output $\tilde{y}{\pi}(q)$ for a query vector are:

$ y(q) \triangleq \sum_{t=1}^T \langle q, u_t\rangle, \quad \tilde{y}\pi(q) \triangleq \sum{i=1}^m \sum_{t\in\mathcal{C}_i} \langle q, \bar{u}_i\rangle\tag{4} $

The deviation induced by summarization $\operatorname{Err}(\pi ; q) \triangleq\left|y(q)-\tilde{y}_\pi(q)\right|$ admits the bound:

$ \big|y(q)-\tilde{y}_\pi(q)\big| ;\le; |q|2 \cdot \sum{i=1}^m \sqrt{|\mathcal{C}i|} ;\sqrt{\sum{t\in\mathcal{C}_i}|u_t-\bar{u}_i|_2^2}\tag{5} $

Proof: The deviation between the exact and summarized outputs $y(q)-\tilde{y}(q)$ can be further rewritten as:

$ \sum_{i=1}^m \sum_{t\in\mathcal{C}_i} \langle q, u_t-\bar{u}_i\rangle

\left\langle q, ; \sum_{i=1}^m \sum_{t\in\mathcal{C}_i} (u_t-\bar{u}_i)\right\rangle $

By Applying Cauchy–Schwarz ([8]) Inequality, we have:

$ \operatorname{Err}(\pi ; q) = \big|y(q)-\tilde{y}\pi(q)\big| \le |q|2 \cdot \left|\sum{i=1}^m \sum{t\in\mathcal{C}_i} (u_t-\bar{u}_i)\right|_2\tag{6} $

We then use Triangle Inequality ([9]) over chunks to further get:

$ \left|\sum_{i=1}^m \sum_{t\in\mathcal{C}_i} (u_t-\bar{u}i)\right|2 \le \sum{i=1}^m \left|\sum{t\in\mathcal{C}_i} (u_t-\bar{u}_i)\right|_2\tag{7} $

Similarly, for each block $\mathcal{C}_i$, we also have:

$ \begin{aligned} \left|\sum_{t\in\mathcal{C}_i} (u_t-\bar{u}_i)\right|2 &\le \sum{t\in\mathcal{C}_i} |u_t-\bar{u}_i|_2 \notag \ &\le \sqrt{|\mathcal{C}i|}\cdot \sqrt{\sum{t\in\mathcal{C}_i}|u_t-\bar{u}_i|_2^2}. \end{aligned}\tag{8} $

By combining Equation 6, Equation 7, and 8, we finally obtain the upper-bound $B(\pi ; q)$ of the deviation $\operatorname{Err}(\pi ; q)$:

$ B(\pi ; q) \triangleq|q|2 \sum{i=1}^m \sqrt{\left|\mathcal{C}i\right|} \sqrt{\sum{t \in \mathcal{C}_i}\left|u_t-\bar{u}_i\right|_2^2}\tag{9} $

This upper bound shows that the deviation induced by block-wise summarization is controlled by the within-block heterogeneity. As a result, content-agnostic fixed blocking policies, which do not adapt to representation variation, can incur a larger bound on non-stationary sequences, especially when tokens with large representation variance are mixed into the same block.

########## {caption="Corollary 2: Fixed blocking is sub-optimal on non-stationary sequences"}

There exists a class of non-stationary token sequences for which any fixed blocking policy $\pi_{\mathrm{fix}}$ yields a strictly larger deviation bound $B(\pi_{\mathrm{fix}};q)$ than an adaptive blocking policy $\pi_{\mathrm{dyn}}$ that aligns block boundaries with semantic change points.

Proof sketch: Consider a non-stationary sequence consisting of two contiguous segments $\mathcal{A}$ and $\mathcal{B}$ with distinct means $\mu_A \neq \mu_B$. For simplicity, assume $u_t=\mu_A$ for $t\in\mathcal{A}$ and $u_t=\mu_B$ for $t\in\mathcal{B}$ (a special case of $u_t\sim\mathcal{D}_A, \mathcal{D}_B$).

Let $\pi_{\mathrm{fix}}$ be any fixed blocking policy that yields at least one block $\mathcal{C}$ overlapping both segments, and denote $n_A = |\mathcal{C}\cap\mathcal{A}|$, $n_B = |\mathcal{C}\cap\mathcal{B}|$. For this block, the choice $\bar{u}$ that minimizes $\sum_{t\in\mathcal{C}}|u_t-\bar{u}|_2^2$ is the block mean $\bar{u} = \frac{n_A\mu_A+n_B\mu_B}{n_A+n_B}$, and the minimum within-block heterogeneity satisfies

$ \sum_{t\in\mathcal{C}}|u_t-\bar{u}|_2^2

\frac{n_A n_B}{n_A+n_B}, |\mu_A-\mu_B|_2^2

0 . $

In contrast, an adaptive policy $\pi_{\mathrm{dyn}}$ that places a boundary at the change point produces blocks contained in $\mathcal{A}$ or $\mathcal{B}$ only, for which the optimal heterogeneity term is $0$ under the same construction. Since the deviation bound $B(\pi;q)$ is a sum of nonnegative per-block terms $|q|2\sqrt{|\mathcal{C}i|}\sqrt{\sum{t\in\mathcal{C}i}|u_t-\bar{u}i|2^2}$, the overlapping block $\mathcal{C}$ alone contributes a strictly positive amount to $B(\pi{\mathrm{fix}};q)$ for any $q\neq 0$, while $B(\pi{\mathrm{dyn}};q)$ does not incur this cross-segment term. Hence, there exists such a sequence for which $B(\pi{\mathrm{fix}};q) > B(\pi{\mathrm{dyn}};q)$, proving the claim.

Key Design: The pseudocode of Information-Aware Dynamic State Merging of DLA is provided in Algorithm 1. We also plot the difference between Vanilla Linear Attention, Log-Linear Attention, and our Dynamic Linear Attention (DLA) in Figure 2. Specifically, to dynamically determine whether a newly generated token $t$ should be merged into an existing memory state or initiate a new one, we first introduce a new metric named State Information Score ($I_t$) to measure the amount of novel information carried by the current token relative to the most recent memory state. Concretely, let $s_{t} \triangleq \phi\left(k_t\right) v_t^{\top}$ denote the state of new token $t$, and let $S_{t-1}$ denote the previous memory state, which summarizes multiple past tokens. We quantify the information variation between $S_t$ and $S_{t-1}$ as follows:

$ I_t = \frac{| S_t - S_{t-1} |F}{| S{t-1} |_F + \epsilon}\tag{10} $

In practice, we apply RMSNorm to both $S_t$ and $S_{t-1}$ prior to score computation to further stabilize the scale across layers and timesteps. During inference, we measure the following boundary indicator

$ b_t \triangleq \mathbf{1}\left[I_t \geq \tau\right]\tag{11} $

Let $S_{t-1}^{\text cur}$ denote the most recent memory state in the cache. The state update rule at inference is then defined as

$ S_t^{\text cur}= \begin{cases}S_{t-1}^{\text cur} + s_t, & b_t=0, \ S_t, & b_t=1, \end{cases}\tag{12} $

where $b_t=1$ indicates that the current token initiates a new memory state, while $b_t=0$ continues to accumulate information into the existing state. When $b_t=1$, the newly created state $S_t$ is appended to the memory cache, preserving the chronological order of states.

We apply soft gating to decide the boundary in a differentiable manner during pre-training and then switch to a hard segmentation strategy during inference to ensure that inference produces a discrete set of memory states aligned with semantic boundaries, while retaining the same information-aware criterion learned during training.

Discussion.

Theorem 1 shows the summarization deviation is dominated by the within-block heterogeneity term in Equation 9. Fixed blocking policies are content-agnostic and therefore may mix tokens from distinct semantic regimes into the same block, which yields a strictly larger deviation bound on non-stationary sequences (Corollary 2). In contrast, DLA monitors token-level representation drift and only merges a new token when the induced increase of heterogeneity is small, using the State Information Score $I_t$ in Equation 10. Therefore, DLA can be viewed as a greedy online strategy that approximately minimizes the dominant term in the deviation bound, while fixed policies ignore it, making it less competitive than DLA.

3.2 Capacity-Bounded Memory Modeling

Motivation: While the previous design enables flexible and information-aware memory state construction, maintaining an unbounded number of states is impractical for efficient inference, especially in long-context or high-throughput serving scenarios, as dynamic memory growth leads to irregular memory layouts, variable attention costs, and reduced batching efficiency. To address these challenges, it is essential to explicitly limit the number of memory states while preserving the most informative summaries.

Key Design: The pseudocode of Capacity-Bounded Memory Modeling of DLA is provided in Algorithm 2. Specifically, DLA maintains a state cache $\mathcal{M} = {(S_i, n_i, \bar{I}i)}{i=1}^{m}$ with $m \le K$, where $S_i \in \mathbb{R}^d$ denotes the $i$-th memory state in chronological order, $n_i$ is the number of tokens summarized by $S_i$, and $\bar{I}_i$ is an aggregated information score of all tokens in this state. We maintain $\bar{I}_i$ as the sum of per-token information scores within each state, such that $\bar{I}_i / n_i$ measures information density. Newly generated tokens are first converted to per-token representations $S_t$, and a tentative state is produced following Section 3.1. The resulting state is appended to the cache, preserving temporal order. When the cache is not full ($m < K$), we simply insert the new state. When the cache reaches capacity ($m = K$), we trigger a compression step that merges two adjacent states to free one slot. Restricting merges to adjacent states preserves the temporal order and avoids distorting positional semantics. Concretely, among all consecutive pairs $(i, i!+!1)$, we select the pair with the lowest information density:

$ (i^\star, i^\star!+!1) = \arg\min_{i \in {1, \dots, K-1}} \frac{\bar{I}i + \bar{I}{i+1}}{n_i + n_{i+1}} $

We then merge them using a summarization operator as in Section 3.1,

$ \begin{aligned} S_{i^\star} &\leftarrow S_{i^\star} + S_{i^\star+1}, \notag \ n_{i^\star} &\leftarrow n_{i^\star} + n_{i^\star+1}, \quad \bar{I}{i^\star} \leftarrow \bar{I}{i^\star} + \bar{I}_{i^\star+1} \end{aligned} $

and shift the remaining states accordingly to keep $m=K-1$ before inserting the incoming state.

Given the capacity-bounded cache $\mathcal{M}$, we compute the output at time step $t$ by attending over the stored memory states. Let $q_t$ denote the query vector of the current token. The final output is then computed as

$ o_t = \sum_{i=1}^{m} \lambda_{t, i} , q_t^\top \left(\sum_{s \in \mathcal{C}i} v_s k_s^\top \right) = \sum{i=1}^{m} \lambda_{t, i} , q_t^\top S_i,\tag{13} $

where $\lambda_{t, i}$ is the weight learned during pre-training with the same shape as the memory capacity. Following ([6]), $\lambda_{t, i}$ is produced by a learned linear layer over the query representation and reused at inference. In this way, DLA provides a unified way to read from a temporally ordered, capacity-bounded memory, enabling stable inference cost while retaining the ability to emphasize informative states during the inference.

4. Experiments

Section Summary: In the experiments, researchers trained DLA variants of Mamba-2 and Gated DeltaNet from scratch on 50 billion tokens and compared them to standard linear-attention models, log-linear attention versions, and a full-attention Transformer of similar size. They tested all models on 16 benchmarks covering commonsense reasoning, in-context retrieval, and long-context understanding. DLA consistently delivered higher accuracy than the baselines across these tasks, with particularly large gains on challenging long-context retrieval problems, while maintaining practical inference speed and memory use.

::: {caption="Table 1: Performance comparison of DLA and baseline methods on zero-shot commonsense reasoning tasks on Mamba-2 (780M) and Gated DeltaNet (1.3B). Commonsense reasoning datasets (LAMBADA, PIQA, HellaSwag, WinoGrande, ARC-e, ARC-c, OpenBookQA, and CommonsenseQA) are measured by accuracy (↑). The best performance is marked in bold. The relative performance gain compared to the best-performing baseline is marked in green inside bracket."}

:::

4.1 Experimental Setups

Baselines. We compare DLA against two groups of models: (1) Vanilla linear attention models, including Mamba-2-780M ([10]) and Gated DeltaNet-1.3B ([5]). (2) Multi-state linear attention models, including Mamba-2 with Log-linear Attention, and Gated DeltaNet with Log-linear Attention ([6]). Following the design in previous work ([6]), we also compare the DLA version of Mamba-2-780M with full attention Transformers with 24 layers and 778M parameters.

Datasets. To demonstrate the generalizability of DLA, we evaluate the performance of DLA on 16 datasets covering three categories, including eight commonsense reasoning datasets (LAMBADA ([11]), PIQA ([12]), HellaSwag ([13]), WinoGrande ([14]), OpenBookQA ([15]), CommonsenseQA ([16]), ARC-e, and ARC-c ([17])), six in-context retrieval datasets (SWDE ([18]), SQuAD ([19]), FDA ([20]), TriviaQA ([21]), Drop ([22]), NQ ([23])), and two long-context datasets (RULER ([24]) and LongBench ([25])). All of the evaluations are conducted using the LM-Evaluation-Harness framework ([26]).

Implementation Details. To ensure a fair comparison, we followed the same configuration used in Log-Linear Attention ([6]) to train the full attention Transformer-778M, Mamba-2-780M, Gated DeltaNet-1.3B, and their variants in Log-Linear and DLA forms. Specifically, we perform academic-scale language modeling pretraining from scratch using 50B tokens on the Long-Data-Collections dataset, using a sequence length of 16K. We set the capacity of the state cache in DLA to 30, which is the same as the maximum state number in Log-linear attention. All of our experiments are conducted on 4 NVIDIA A100 GPUs.

4.2 Overall Comparison

We evaluate the overall performance of DLA from three main aspects: (1) performance on commonsense reasoning tasks, (2) performance on in-context retrieval tasks, and (3) performance on long-context modeling tasks.

Performance on Commonsense Reasoning. Following prior work ([10]), we evaluate all models on eight commonsense reasoning benchmarks. Results are summarized in Table 1. We make two key observations. First, DLA consistently outperforms both the vanilla and log-linear variants of linear-attention–based models across all tasks. In particular, compared to the log-linear variant, DLA achieves up to 52% and 22% relative accuracy improvement on Mamba-2-780M and Gated DeltaNet-1.3B, respectively. Second, when applied to Mamba-2-780M, DLA also consistently outperforms a full-attention Transformer with a comparable parameter size, demonstrating that DLA can close and even surpass the accuracy gap between linear attention and full attention.

::: {caption="Table 2: Performance on in-context retrieval benchmarks measured by accuracy (↑). The best performance is marked in bold. The relative performance gain compared to the best-performing baseline is marked in green inside bracket."}

:::

Performance on In-Context Retrieval Tasks. Then, we evaluate the models on six in-context retrieval tasks following prior work ([27]). As shown in Table 2, DLA consistently outperforms the baseline methods with at most 49 $%$ performance improvement.

::: {caption="Table 3: Evaluation results of single-needle tasks (S-NIAH-1–3) and multi-needle tasks (MK-1, MQ, MV) on RULER (4K context)."}

:::

Performance on Long-Context Modeling Tasks. We next evaluated the models on long-context tasks, including long-context retrieval on RULER with 4k, 8k, 16k length and long-context understanding on LongBench. As shown in Table 3 and Table 4, we make two main observations.

First, DLA substantially improves long-context retrieval performance on RULER across both single-needle and multi-needle settings. Compared to the log-linear variant, DLA achieves consistent and often large gains on Mamba-2 and Gated DeltaNet, with particularly pronounced improvements on harder multi-needle tasks (e.g., up to 350% relative improvement on S-NIAH-3 and 67% on MQ-NIAH). These results indicate that DLA more effectively preserves and aggregates long-range information under extended contexts.

::: caption="Table 4: Performance on LongBench datasets ([25]) with different types of tasks."

:::

Second, on LongBench, DLA consistently outperforms both vanilla and log-linear variants across diverse long-context understanding tasks, including narrative QA, multi-field QA, summarization, and few-shot learning. Notably, DLA delivers strong and uniform gains across different task categories, suggesting that the benefits of DLA extend beyond retrieval and generalize to complex reasoning and generation under long contexts.

4.3 Inference Efficiency of DLA

We next evaluate the efficiency of DLA from two aspects: (1) efficiency under varying batch sizes and (2) efficiency under varying input context lengths.

Efficiency Under Various Batch Sizes. Figure 3(a) and (c) report the throughput and runtime memory footprint under varying batch sizes with a fixed context length of 128 and decode length of 1. As batch size increases, both the log-linear and DLA variants exhibit smaller throughput gains and higher memory usage than the vanilla Mamba-2, due to caching multiple summary states. Nevertheless, compared to log-linear attention, DLA consistently achieves higher throughput with lower memory consumption, indicating better compute and memory efficiency.

Efficiency Under Various Context Lengths. Figure 3(b) and (d) show the throughput and KV memory footprint under varying context lengths with a fixed batch size of 1 and decode length of 1. As context length increases, both log-linear and DLA variants incur higher memory usage and limited throughput improvement relative to the vanilla model, again due to maintaining multiple summary states. In contrast, DLA consistently outperforms the log-linear variant in throughput while maintaining lower and more stable memory consumption, demonstrating superior efficiency under long-context settings.

4.4 Ablation Studies

::: {caption="Table 5: Ablation study of Mamba-2 and Gated DeltaNet with different variants. DLA(I) denotes the version of DLA with information-aware dynamic state merging only."}

:::

Module Sensitivity Study. We conduct ablation studies to evaluate the separate contribution of the two components of DLA. Let DLA(I) denote the version of DLA with information-aware dynamic state merging only. As shown in Table 5, we have two observations. (1) DLA(I) and DLA variants consistently outperform Log-Linear variants across all benchmarks. (2) DLA consistently outperforms DLA(I) across all benchmarks. This result demonstrates the unique contributions of the two components in DLA.

::: caption="Table 6: Ablation study of Mamba-2 DLA variant with different memory budget k and merge boundary tau."

:::

Impact of capacity $k$ and boundary $\tau$. To study the impact of memory budget $k$ and the merge boundary $\tau$ on performance, we adjust the default budget in DLA variant of Mamba-2 from 30 to 20 and 40 and adjust the default boundary from 0.6 to 0.5 and 0.7. We then compare the changes in performance. As shown in Table 6, changes in the memory budget and merge boundary have only a marginal effect on the final performance of DLA, indicating that the proposed memory modeling is robust to these two hyperparameters.

5. Related Work

Section Summary: Researchers have created linear attention methods and state space models to process long sequences efficiently by avoiding the slowdown of standard attention calculations. Models such as DeltaNet and Mamba compress all prior information into one ongoing summary state, with gating added to selectively drop outdated details. Newer variants maintain multiple memory states merged on a fixed schedule, yet they still rely on rigid rules that do not adapt to the varying importance of individual tokens.

To overcome the quadratic bottleneck of softmax attention on long sequences, linear attention and state space models (SSMs) reformulate attention computation to achieve $O(T)$ complexity. Representative methods such as DeltaNet ([4]) and Mamba ([28]) compress the entire history into a single recurrent state, continuously merging incoming tokens into a fixed-size summary for inference. To alleviate the resulting over-compression, gating mechanisms ([5]) introduce data-dependent modulation to selectively attenuate obsolete information. More recent approaches extend linear attention to multi-state memory. In particular, Log-Linear Attention ([6]) maintains a logarithmic number of hierarchical states, where tokens are deterministically merged according to a fixed temporal schedule. Despite their effectiveness, these methods rely on fixed merging policies that ignore token-level information variation, leaving open the question of how to adaptively control state construction to preserve fine-grained information under long contexts.

6. Conclusion

Section Summary: The paper introduces DLA, a framework for multi-state linear attention that improves how models process information by dynamically building memory states based on the actual content of each token rather than using fixed rules. It also applies capacity limits to keep memory use and inference costs stable and predictable. The authors pre-trained the approach on two existing backbones and showed consistent gains over leading methods across 16 datasets.

In this paper, we presented DLA, a framework for multi-state linear attention. DLA replaces fixed merging with information-aware dynamic state construction and uses capacity-bounded memory modeling to keep inference cost predictable. By allocating memory resolution based on token-level information variation, DLA improves representation quality while preserving efficiency. We pre-train DLA on two linear-attention backbones and evaluate it on 16 datasets across three aspects, where it consistently outperforms state-of-the-art baselines.

Acknowledgement

This work is supported in part by NSF Award NeTS-2312675.

References

Section Summary: This section compiles citations to recent research papers on efficient large language models, including new architectures that aim to reduce computational demands while handling long sequences of text. It also references established benchmarks and datasets used to test model performance on tasks like commonsense reasoning, reading comprehension, and question answering. A few older works on mathematical inequalities and network measurements are included alongside the dominant focus on AI and machine learning advancements from 2023 onward.

[1] Wan et al. (2023). Efficient Large Language Models: A Survey. arXiv preprint arXiv:2312.03863.

[2] Wang et al. (2024). IoT in the Era of Generative AI: Vision and Challenges. arXiv preprint arXiv:2401.01923.

[3] Zhongwei Wan et al. (2025). D2O: Dynamic Discriminative Operations for Efficient Long-Context Inference of Large Language Models. In ICLR.

[4] Songlin Yang et al. (2024). Parallelizing Linear Transformers with the Delta Rule over Sequence Length. In NeurIPS.

[5] Songlin Yang et al. (2025). Gated Delta Networks: Improving Mamba2 with Delta Rule. In ICLR.

[6] Han Guo et al. (2025). Log-Linear Attention. CoRR. abs/2506.04761.

[7] Xiuying Wei et al. (2025). RAT: Bridging RNN Efficiency and Attention Accuracy in Language Modeling. CoRR. abs/2507.04416.

[8] Johnston et al. (2025). Generalizing the Cauchy-Schwarz inequality: Hadamard powers and tensor products. arXiv preprint arXiv:2507.10327.

[9] Lumezanu et al. (2009). Triangle inequality variations in the internet. In Proceedings of the 9th ACM SIGCOMM Conference on Internet Measurement. pp. 177–183. doi:10.1145/1644893.1644914. https://doi.org/10.1145/1644893.1644914.

[10] Tri Dao and Albert Gu (2024). Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. In ICML.

[11] Denis Paperno et al. (2016). The LAMBADA dataset: Word prediction requiring a broad discourse context. In ACL (1).

[12] Yonatan Bisk et al. (2020). PIQA: Reasoning about Physical Commonsense in Natural Language. In AAAI. pp. 7432–7439.

[13] Rowan Zellers et al. (2019). HellaSwag: Can a Machine Really Finish Your Sentence?. In ACL (1). pp. 4791–4800.

[14] Keisuke Sakaguchi et al. (2021). WinoGrande: an adversarial winograd schema challenge at scale. Commun. ACM. 64(9). pp. 99–106.

[15] Todor Mihaylov et al. (2018). Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering. In EMNLP. pp. 2381–2391.

[16] Alon Talmor et al. (2019). CommonsenseQA: A Question Answering Challenge Targeting Commonsense Knowledge. In NAACL-HLT (1). pp. 4149–4158.

[17] Sumithra Bhakthavatsalam et al. (2021). Think you have Solved Direct-Answer Question Answering? Try ARC-DA, the Direct-Answer AI2 Reasoning Challenge. CoRR. abs/2102.03315.

[18] Lockard et al. (2019). Openceres: When open information extraction meets the semi-structured web. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers). pp. 3047–3056.

[19] Rajpurkar et al. (2018). Know What You Don’t Know: Unanswerable Questions for SQuAD. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers). pp. 784–789.

[20] Arora et al. (2023). Language Models Enable Simple Systems for Generating Structured Views of Heterogeneous Data Lakes. Proceedings of the VLDB Endowment. 17(2). pp. 92–105.

[21] Joshi et al. (2017). TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). pp. 1601–1611.

[22] Dua et al. (2019). DROP: A Reading Comprehension Benchmark Requiring Discrete Reasoning Over Paragraphs. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers). pp. 2368–2378.

[23] Kwiatkowski et al. (2019). Natural questions: a benchmark for question answering research. Transactions of the Association for Computational Linguistics. 7. pp. 453–466.

[24] Cheng-Ping Hsieh et al. (2024). RULER: What's the Real Context Size of Your Long-Context Language Models?. CoRR. abs/2404.06654.

[25] Yushi Bai et al. (2024). LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding. In ACL (1). pp. 3119–3137.

[26] Gao et al. (2024). The Language Model Evaluation Harness. doi:10.5281/zenodo.12608602. https://zenodo.org/records/12608602.

[27] Arora et al. (2024). Simple linear attention language models balance the recall-throughput tradeoff. In Proceedings of the 41st International Conference on Machine Learning. pp. 1763–1840.

[28] Albert Gu and Tri Dao (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. CoRR. abs/2312.00752.