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

The Key to Going Linear: Analysis-Driven Transformer Linearization
Jul 8, 2026Anna Kuzina, Paul N. Whatmough, Babak Ehteshami BejnordiSoftmax Attention LLM Inference Acceleration
Qualcomm AI Research∗ · Qualcomm AI Research
MDN: Parallelizing Stepwise Momentum for Delta Linear Attention
May 7, 2026Yulong Huang, Xiang Liu, Hongxiang Huang +5Momentum Methods Linear Attention
The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China
Blurry Window Attention
May 31, 2026Axel Laborieux, Christos Sourmpis, Juan Gabriel Kostelec +1Long-Context Language Modeling Linear Attention
Huawei, Zurich, Switzerland · Huawei Advanced Computing and Storage Lab, Shenzhen, China

Jul 8, 2026J/K move · Enter open · S save

The Key to Going Linear: Analysis-Driven Transformer Linearization

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

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

The Key to Going Linear: Analysis-Driven Transformer Linearization

MDN: Parallelizing Stepwise Momentum for Delta Linear Attention

Blurry Window Attention