2Mamba2Furious: Linear in Complexity, Competitive in Accuracy

Gabriel Mongaras; Eric C. Larson

cs.LGFeb 19, 2026

2Mamba2Furious: Linear in Complexity, Competitive in Accuracy

Authors: Gabriel Mongaras, Eric C. Larson

Organizations: Lyle School of Engineering Southern Methodist University Dallas, TX 75205

Abstract

Linear attention transformers have become a strong alternative to softmax attention due to their efficiency. However, linear attention tends to be less expressive and results in reduced accuracy compared to softmax attention. To bridge the accuracy gap between softmax attention and linear attention, we manipulate Mamba-2, a very strong linear attention variant. We first simplify Mamba-2 down to its most fundamental and important components, evaluating which specific choices make it most accurate. From this simplified Mamba variant (Mamba-2S), we improve the A-mask and increase the order of the hidden state, resulting in a method, which we call 2Mamba, that is nearly as accurate as softmax attention, yet much more memory efficient for long context lengths. We also investigate elements to Mamba-2 that help surpass softmax attention accuracy. Code is provided for all our experiments.

Figures & tables

$\displaystyle QKV$	$\displaystyle=XW_{QKV}$	$\displaystyle\in\mathbb{R}^{N,3\cdot(H\cdot d_{h})}$
$\displaystyle Q,K,V$	$\displaystyle=\text{split}(QKV)$	$\displaystyle\in\mathbb{R}^{H,N,d_{h}}$
$\displaystyle y$	$\displaystyle=((\phi(Q)\phi(K)^{T}\odot M)V)$	$\displaystyle\in\mathbb{R}^{H,N,d_{h}}$
$\displaystyle N$	$\displaystyle=\sum_{j}(\phi(Q)\phi(K)^{T}\odot M)$	$\displaystyle\in\mathbb{R}^{H,N}$
$\displaystyle y_{n}$	$\displaystyle=y/N$	$\displaystyle\in\mathbb{R}^{H,N,d_{h}}$
$\displaystyle y^{\prime}_{n}$	$\displaystyle=\text{vec}(y_{n})_{(H,N,d_{h})\longrightarrow(N,H\cdot d_{h})}$	$\displaystyle\in\mathbb{R}^{N,H\cdot d_{h}}$
$\displaystyle o$	$\displaystyle=y^{\prime}_{n}W_{out}$	$\displaystyle\in\mathbb{R}^{N,d}$

Algorithm 1 Linear attention

	$\displaystyle H$	$\displaystyle=\text{num heads},\quad d_{h}=\text{head (inner) dim},\quad N=\text{seq. length,}$
	$\displaystyle\mathcal{U}[\cdot]$	$\displaystyle=\text{Uniform Dist.},\qquad d_{ssm}=H*d_{h}$
	$\displaystyle d_{conv}$	$\displaystyle=3\cdot d_{ssm},\quad dt_{min}=0.001,\quad dt_{max}=0.1,\quad l_{min}=\log(dt_{min}),\quad l_{max}=\log(dt_{max})$
	$\displaystyle A_{log}$	$\displaystyle=\log(\mathcal{U}[1,16])\in\mathbb{R}^{H},\quad D=\text{ones}({H*d_{h}})\in\mathbb{R}^{H\cdot d_{h}}$
	$\displaystyle dt_{init}$	$\displaystyle=\exp\left(x\right)\in\mathbb{R}^{H}\quad\text{where }\quad x\sim\mathcal{U}\left[l_{min},l_{max}\right]$
	$\displaystyle dt_{bias}$	$\displaystyle=\text{softplus}^{-1}(dt_{init})\quad=dt_{init}+\log(1-\exp(-dt_{init}))\in\mathbb{R}^{H}$

Algorithm 2 Mamba-2 Setup

$\displaystyle Q,K,V$	$\displaystyle=\sigma(\text{conv\_1d}(h\cdot W_{QKV}))$	$\displaystyle\in\mathbb{R}^{H,N,d_{h}}$
$\displaystyle dt_{P}$	$\displaystyle=h\cdot W_{dt}$	$\displaystyle\in\mathbb{R}^{H,N}$
$\displaystyle Z$	$\displaystyle=h\cdot W_{z}$	$\displaystyle\in\mathbb{R}^{H,N,d_{h}}$
$\displaystyle A$	$\displaystyle=-\exp(A_{log})$	$\displaystyle\in\mathbb{R}^{H}$
$\displaystyle dt$	$\displaystyle=\text{softplus}(dt_{P}+dt_{bias})$	$\displaystyle\in\mathbb{R}^{H,N}$
$\displaystyle D_{res}$	$\displaystyle=V\odot D,\qquad V_{dt}=V\odot dt$	$\displaystyle\in\mathbb{R}^{H,N,d_{h}}$
$\displaystyle A^{CS}$	$\displaystyle=\text{cumsum}(A\odot dt)$	$\displaystyle\in\mathbb{R}^{H,N}$

Algorithm 3 Mamba-2 Forward Pass

Figure 1: Accuracy of linear attention, Mamba, and softmax attention, keeping everything but the attention mechanism constant across experiments.

Figure 2: Accuracy of various norm types. Softmax normalization requires a positive inner-product space image, as such we use ReLU.

Model	Train loss	Test loss
ReLU	3.12	3.06
Conv (w=2)	2.87	2.93
Conv (w=3)	2.94	2.91
Conv (w=4)	2.84	2.9
Conv (w=2) + SiLU	2.95	2.9
Conv (w=3) + SiLU	2.95	2.89

Table 1: Adding a convolution to normal linear attention increases accuracy.

Model	Train loss	Test loss
SM Norm + ReLU	3.0	3.03
Out Norm	3.19	3.13
Out Norm + ReLU	3.12	3.06
Out Norm + Conv (w=2)	2.87	2.93
Out Norm + Conv (w=3)	2.94	2.91
Out Norm + Conv (w=2) + SiLU	2.95	2.9

Table 2: Table of isolated component additions added to linear attention.

Panel of Figure 4 — Figure 4: Investigations of major and minor build ups in constructing the simplified Mamba-2 architecture.

Model	Train loss	Test loss
Plain Linear	3.19	3.13
Conv (w=2)	2.87	2.93
Conv (w=2) + $A$ -mask original	2.81	2.8
Conv (w=2) + $A$ -mask softplus	2.67	2.77
Conv (w=2) + $A$ -mask softplus + value discretize	2.7	2.76
Conv (w=2) + $A$ -mask softplus + value discretize + SiLU	2.73	2.76

Table 3: Table of all additions (buildups) leading up to our final algorithm. Italics convey the components in Mamba-2S.

$\displaystyle Q,K,V$	$\displaystyle=\text{conv\_1d}(h\cdot W_{QKV})$	$\displaystyle\in\mathbb{R}^{H,N,d_{h}}$
$\displaystyle dt$	$\displaystyle=\text{softplus}(h\cdot W_{dt})$	$\displaystyle\in\mathbb{R}^{H,N}$
$\displaystyle A$	$\displaystyle={\color[rgb]{0,0,1}-\text{softplus}(h\cdot W_{A})}$	$\displaystyle\in\mathbb{R}^{H,N}$
$\displaystyle V_{dt}$	$\displaystyle=V\odot dt$	$\displaystyle\in\mathbb{R}^{H,N,d_{h}}$
$\displaystyle A^{CS}$	$\displaystyle=\text{cumsum}({\color[rgb]{0,0,1}A})$	$\displaystyle\in\mathbb{R}^{H,N}$
$\displaystyle A^{M}$	$\displaystyle=\exp(A^{CS}-(A^{CS})^{T})\quad\text{where }A^{M}_{ij}=\exp(A^{CS}_{i}-A^{CS}_{j})$	$\displaystyle\in\mathbb{R}^{H,N,N}$
$\displaystyle M_{ij}$	$\displaystyle=\begin{cases}1,&\text{if }i\geq j\\ 0,&\text{if }i<j\end{cases}$	$\displaystyle M\in\mathbb{R}^{H,N,N}$

Algorithm 4 Proposed Mamba-2 Simplified Forward (Mamba-2S): no setup besides weights

Panel of Figure 5 — Figure 5: Experimental results comparing accuracy and training stability.

$\displaystyle Q,K,V$	$\displaystyle=\text{conv\_1d}(h\cdot W_{QKV})$	$\displaystyle\in\mathbb{R}^{H,N,d_{h}}$
$\displaystyle A$	$\displaystyle=-\text{softplus}(h\cdot W_{A})$	$\displaystyle\in\mathbb{R}^{H,N}$
$\displaystyle A^{CS}$	$\displaystyle=\text{cumsum}(A)$	$\displaystyle\in\mathbb{R}^{H,N}$
$\displaystyle A^{M}$	$\displaystyle=\exp(A^{CS}-(A^{CS})^{T})\quad\text{where}\quad A^{M}_{ij}=\exp(A^{CS}_{i}-A^{CS}_{j})$	$\displaystyle\in\mathbb{R}^{H,N,N},$
$\displaystyle M_{ij}$	$\displaystyle=\begin{cases}1,&\text{if }i\geq j\\ 0,&\text{if }i<j\end{cases}$	$\displaystyle M\in\mathbb{R}^{H,N,N}$
$\displaystyle y$	$\displaystyle=((QK^{T})^{{\color[rgb]{0,0,1}2}}\odot A^{M}\odot M)\cdot V$	$\displaystyle\in\mathbb{R}^{H,N,d_{h}}$
$\displaystyle{\color[rgb]{0,0,1}N}$	$\displaystyle{\color[rgb]{0,0,1}=\sum_{j}((QK^{T})^{2}\odot A^{M}\odot M)}$	$\displaystyle{\color[rgb]{0,0,1}\hskip 2.5pt\in\mathbb{R}^{H,N}}$

Algorithm 5 2Mamba Algorithm

Panel of Figure 6 — Figure 6: First row : Small model ( $\sim$ 300M params) test loss for Mamba-2, 2Mamba, and softmax attention trained on 2048, 4096, and 8192 sequence lengths. Second row : Medium model ( $\sim$ 700M params) test loss for Mamba-2, our proposed algorithm, and softmax attention trained on 2048, 4096, and 8192 sequence lengths

Figure 7: Memory usage of a single head of softmax attention, 2Mamba, and linear attention, utilizing softmax-like normalization and a convolution with kernel size 2.

Figure 8: One-shot NIAH benchmark results using maximum likelihood sampling over 1000 sequences. Green indicates the predicted needle was correct more often while red indicates the predicted needle was predicted less often. Scores are indicative of the proportion of sequences in which the model predicted the needle correctly.

Figure 9: Long training run of various models with a batch size of 64 on $8,192$ max sequence length.

Panel of Figure 10 — Figure 10: Medium model ( $\sim$ 700M params) test loss for Mamba-2, 2Mamba-E, and softmax attention trained on 2048, 4096, and 8192 sequence lengths.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

$\displaystyle q_{t},k_{t},v_{t}$	$\displaystyle=h_{t}W_{QKV}$	$\displaystyle\in\mathbb{R}^{H,d_{h}}$
$\displaystyle q^{c}_{t},k^{c}_{t},v^{c}_{t}$	$\displaystyle=\text{conv\_1d}(q_{t},k_{t},v_{t},q_{t-1},k_{t-1},v_{t-1},\text{window\_size}=2)$	$\displaystyle\in\mathbb{R}^{H,d_{h}}$
$\displaystyle a_{t}$	$\displaystyle=exp(-softplus(h_{t}\cdot W_{A}))$	$\displaystyle\in\mathbb{R}^{H}$
$\displaystyle(q^{c}_{t})^{\otimes 2}$	$\displaystyle=q^{c}_{t}\otimes q^{c}_{t}$	$\displaystyle\in\mathbb{R}^{H,D}$
$\displaystyle(k^{c}_{t})^{\otimes 2}$	$\displaystyle=k^{c}_{t}\otimes k^{c}_{t}$	$\displaystyle\in\mathbb{R}^{H,D}$
$\displaystyle H^{\uparrow}_{t}$	$\displaystyle=H^{\uparrow}_{t-1}\odot a_{t}+\left[(k^{c}_{t})^{\otimes 2}\right]^{T}\cdot(v^{c}_{t})$	$\displaystyle\in\mathbb{R}^{H,D,d_{h}}$
$\displaystyle H^{\downarrow}_{t}$	$\displaystyle=H^{\downarrow}_{t-1}\odot a_{t}+\left[(k^{c}_{t})^{\otimes 2}\right]^{T}$	$\displaystyle\in\mathbb{R}^{H,D}$

Appendix

Algorithm 6 2Mamba Inference

Explore similar work

Jul 8, 2026

The Key to Going Linear: Analysis-Driven Transformer Linearization

The quadratic cost of causal self-attention severely bottlenecks long-context transformer inference. While numerous post hoc linearization pipelines exist, it is difficult to identify which components preserve model quality. This work isolates the effect of state update design in a strict frozen-backbone regime. We show that softmax relies on key-dependent, rank-1 orthogonal projections, elucidating why delta-style networks outperform purely gated accumulation. We identify a potential source of approximation errors and introduce structural interventions, specifically sink tokens, short convolutions, and fixed-budget cache routing, which reduces the remaining gap. We scale this linearization approach across LLaMA and Qwen models up to 32B parameters, outperforming prior post hoc baselines on MMLU and matching the long-context retrieval of complex adaptive-caching frameworks.

Anna Kuzina, Paul N. Whatmough, Babak Ehteshami Bejnordi

Qualcomm AI Research∗ · Qualcomm AI Research

Softmax Attention LLM Inference Acceleration Delta-Rule Linear Attention Long-Context Inference

Ask PDF

May 7, 2026

MDN: Parallelizing Stepwise Momentum for Delta Linear Attention

Linear Attention (LA) offers a promising paradigm for scaling large language models (LLMs) to long sequences by avoiding the quadratic complexity of self-attention. Recent LA models such as Mamba2 and GDN interpret linear recurrences as closed-form online stochastic gradient descent (SGD), but naive SGD updates suffer from rapid information decay and suboptimal convergence in optimization. While momentum-based optimizers provide a natural remedy, they pose challenges in simultaneously achieving training efficiency and effectiveness. To address this, we develop a chunkwise parallel algorithm for LA with a stepwise momentum rule by geometrically reordering the update coefficients. Further, from a dynamical systems perspective, we analyze the momentum-based recurrence as a second-order system that introduces complex conjugate eigenvalues. This analysis guides the design of stable gating constraints. The resulting model, Momentum DeltaNet (MDN), leverages Triton kernels to achieve comparable training throughput with competitive linear models such as Mamba2 and KDA. Extensive experiments on the 400M and 1.3B parameter models demonstrate consistent performance improvements over strong baselines, including Transformers, Mamba2 and GDN, across diverse downstream evaluation benchmarks. Code: https://github.com/HuuYuLong/MomentumDeltaNet .

Yulong Huang, Xiang Liu, Hongxiang Huang +5

The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China

Momentum Methods Linear Attention Delta-Rule Linear Attention Linear RNNs

Ask PDF

May 31, 2026

Blurry Window Attention

The Softmax Attention operation in Transformer language models has a quadratic complexity in the sequence length and a growing state size in the form of KV cache, which becomes a bottleneck in long context scenarios. To overcome this limitation, alternative architectures with linear complexity and finite state size have been introduced, such as State-Space Models (SSMs), Linear Attention (LA), and Attention with Bounded-memory Control (ABC). Though linear models achieve similar language perplexity as Transformers, they are still behind in tasks which require retrieval or recall of specific information. In this work, we introduce Blurry Window Attention (BLA) a novel ABC method inspired by SSMs. BLA stores a frequency window from which a blurry KV history is reconstructed via interpolation using Dirichlet kernels. BLA can be understood as a generalization of Sliding Window Attention (SWA) depending on the Dirichlet kernels resolution or as a special case of the Gated Slot Attention (GSA), where the decay factor is implemented with Dirichlet kernels. We describe in details the theory and efficient implementation of BLA. On the Multi-Query Associate Recall (MQAR) synthetic task, we show that the state efficiency of BLA is 8

\times

better than SWA and is competitive with popular linear attention models, and in the RegBench synthetic task, only BLA and SWA improve their performance as the state size grows among the linear models we tested.

Axel Laborieux, Christos Sourmpis, Juan Gabriel Kostelec +1

Huawei, Zurich, Switzerland · Huawei Advanced Computing and Storage Lab, Shenzhen, China

Long-Context Language Modeling Linear Attention Attention Mechanisms Associative Memory

Ask PDF