2Mamba: Second-Order Linear Attention

• 2Mamba is a second-order linear attention mechanism derived from Mamba-2S that closely approximates softmax accuracy while maintaining linear computational cost.
• It integrates a softplus A-mask and convolution to harness key inductive biases, leading to improved performance and memory efficiency for long-context tasks.
• The design shifts computational load from sequence length to feature dimensions, demonstrating memory benefits especially for sequences exceeding about 1,000 tokens.

2Mamba is a Mamba-2-derived linear-attention architecture introduced to close the accuracy gap between linear attention and softmax attention while preserving linear complexity. In the formulation presented in “2Mamba2Furious: Linear in Complexity, Competitive in Accuracy” (Mongaras et al., 19 Feb 2026), the method emerges through a two-step reduction-and-extension pipeline: Mamba-2 is first simplified into a minimal effective variant called Mamba-2S, and then the hidden state is increased to a second-order form, yielding 2Mamba. The resulting model is described as nearly as accurate as softmax attention, yet much more memory efficient for long context lengths (Mongaras et al., 19 Feb 2026).

1. Position within efficient sequence modeling

2Mamba is motivated by a familiar trade-off in sequence modeling: softmax attention is highly expressive but quadratically expensive in sequence length during training, while linear attention is efficient but typically less accurate (Mongaras et al., 19 Feb 2026). The method therefore belongs to the class of efficient attention alternatives that seek to retain linear scaling without accepting the usual loss in expressivity.

In the paper’s framing, Mamba-2 serves as the starting point because it is treated as a very strong linear attention variant rather than merely as a recurrent state-space model. The authors explicitly interpret Mamba-2 as a structured, decay-masked variant of linear attention with extra inductive biases, including a convolution over inputs, query/key/value projections, a decay-like A-mask, time discretization, value scaling, a residual $D$ term, a $Z$ gate, and normalization (Mongaras et al., 19 Feb 2026). This positioning is central: 2Mamba is not introduced as an unrelated architecture, but as a controlled manipulation of Mamba-2 intended to isolate which components actually matter for accuracy.

The paper states the standard softmax attention form as

$Q = XW_Q,\quad K = XW_K,\quad V = XW_V$

$$O = \text{softmax}(QK^T + M)V = \frac{\exp(QK^T + M)V}{\sum_j \exp(QK^T + M)}.$$

By contrast, the baseline linear attention form is written as

$$O = \frac{\left[\phi(Q)\phi(K)^T\right]V}{\sum_j \phi(Q)\phi(K)^T} = \frac{\phi(Q)\left[\phi(K)^T V\right]}{\sum_j \phi(Q)\phi(K)^T},$$

with $\phi(\cdot)=\text{ReLU}$ in the baseline experiments (Mongaras et al., 19 Feb 2026). The role of 2Mamba is to narrow the expressivity gap between these two regimes without abandoning linear sequence complexity.

2. Simplification from Mamba-2 to Mamba-2S

A defining feature of the work is the claim that Mamba-2 contains many components that may be “arbitrary or underexplored,” so the authors first perform ablations before proposing 2Mamba (Mongaras et al., 19 Feb 2026). The simplification step produces Mamba-2S, which is the intermediate architecture from which 2Mamba is derived.

The full Mamba-2 block is written with the following major components: $Q,K,V = \sigma(\text{conv\_1d}(hW_{QKV}))$

$dt_P = hW_{dt}$

$Z = hW_z$

$A = -\exp(A_{\log})$

$Z$0

$Z$1

$Z$2

$Z$3

with $Z$4, a residual value term

$Z$5

and output gating/normalization

$Z$6

This description matters because the subsequent ablations are precisely about which of these elements should be kept (Mongaras et al., 19 Feb 2026).

The paper reports that the two strongest contributors are the A-mask, especially in its softplus form, and the input convolution. Time discretization gives a smaller gain, while the $Z$7-gate and $Z$8-residual are not helpful enough to justify complexity (Mongaras et al., 19 Feb 2026). From these findings, the authors define Mamba-2S as the minimal but effective Mamba-2 variant with just an input convolution of window size 2, a softplus A-mask, time discretization, and output RMSNorm.

Its forward pass is: $Z$9

$Q = XW_Q,\quad K = XW_K,\quad V = XW_V$0

$Q = XW_Q,\quad K = XW_K,\quad V = XW_V$1

$Q = XW_Q,\quad K = XW_K,\quad V = XW_V$2

$Q = XW_Q,\quad K = XW_K,\quad V = XW_V$3

$Q = XW_Q,\quad K = XW_K,\quad V = XW_V$4

$Q = XW_Q,\quad K = XW_K,\quad V = XW_V$5

$Q = XW_Q,\quad K = XW_K,\quad V = XW_V$6

$Q = XW_Q,\quad K = XW_K,\quad V = XW_V$7

$Q = XW_Q,\quad K = XW_K,\quad V = XW_V$8

$Q = XW_Q,\quad K = XW_K,\quad V = XW_V$9

Mamba-2S is therefore not an end in itself. It is a controlled reduction that retains what the authors judge to be the strongest Mamba-2 inductive biases before introducing the second-order modification that defines 2Mamba (Mongaras et al., 19 Feb 2026).

3. Second-order hidden state and the defining 2Mamba update

The paper argues that even Mamba-2S still falls short of softmax attention accuracy, and that the remaining gap should be addressed by increasing the order of the hidden state (Mongaras et al., 19 Feb 2026). This is the specific technical meaning of the “2” in 2Mamba: the model uses a second-order hidden state.

The motivating intuition is linked to a Taylor/Maclaurin expansion view of softmax attention. In the paper’s summary, softmax attention can be viewed as a sum over higher-order query-key interactions; ordinary linear attention is roughly a first-order approximation; and a second-order representation captures more structure and is more expressive (Mongaras et al., 19 Feb 2026). The authors present second order as a practical sweet spot: higher order improves approximation to softmax, but the hidden state size grows rapidly with order, so $$O = \text{softmax}(QK^T + M)V = \frac{\exp(QK^T + M)V}{\sum_j \exp(QK^T + M)}.$$0 is still practical.

Starting from Mamba-2S, 2Mamba makes three defining changes: it removes RMSNorm, uses softmax normalization, and replaces $$O = \text{softmax}(QK^T + M)V = \frac{\exp(QK^T + M)V}{\sum_j \exp(QK^T + M)}.$$1 with $$O = \text{softmax}(QK^T + M)V = \frac{\exp(QK^T + M)V}{\sum_j \exp(QK^T + M)}.$$2 (Mongaras et al., 19 Feb 2026). The complete forward pass is given as

$$O = \text{softmax}(QK^T + M)V = \frac{\exp(QK^T + M)V}{\sum_j \exp(QK^T + M)}.$$3

$$O = \text{softmax}(QK^T + M)V = \frac{\exp(QK^T + M)V}{\sum_j \exp(QK^T + M)}.$$4

$$O = \text{softmax}(QK^T + M)V = \frac{\exp(QK^T + M)V}{\sum_j \exp(QK^T + M)}.$$5

$$O = \text{softmax}(QK^T + M)V = \frac{\exp(QK^T + M)V}{\sum_j \exp(QK^T + M)}.$$6

$$O = \text{softmax}(QK^T + M)V = \frac{\exp(QK^T + M)V}{\sum_j \exp(QK^T + M)}.$$7

$$O = \text{softmax}(QK^T + M)V = \frac{\exp(QK^T + M)V}{\sum_j \exp(QK^T + M)}.$$8

$$O = \text{softmax}(QK^T + M)V = \frac{\exp(QK^T + M)V}{\sum_j \exp(QK^T + M)}.$$9

$$O = \frac{\left[\phi(Q)\phi(K)^T\right]V}{\sum_j \phi(Q)\phi(K)^T} = \frac{\phi(Q)\left[\phi(K)^T V\right]}{\sum_j \phi(Q)\phi(K)^T},$$0

$$O = \frac{\left[\phi(Q)\phi(K)^T\right]V}{\sum_j \phi(Q)\phi(K)^T} = \frac{\phi(Q)\left[\phi(K)^T V\right]}{\sum_j \phi(Q)\phi(K)^T},$$1

$$O = \frac{\left[\phi(Q)\phi(K)^T\right]V}{\sum_j \phi(Q)\phi(K)^T} = \frac{\phi(Q)\left[\phi(K)^T V\right]}{\sum_j \phi(Q)\phi(K)^T},$$2

The paper ties the square operation to several specific benefits. It provides higher-order expressivity, makes the attention scores nonnegative, allows normalization to shift to the attention weights rather than the output, and improves accuracy at long context (Mongaras et al., 19 Feb 2026). The authors also note that placing the square on the A-mask would square unnecessary terms, so the square is applied directly to the $$O = \frac{\left[\phi(Q)\phi(K)^T\right]V}{\sum_j \phi(Q)\phi(K)^T} = \frac{\phi(Q)\left[\phi(K)^T V\right]}{\sum_j \phi(Q)\phi(K)^T},$$3 inner product instead.

A plausible implication is that 2Mamba should be read less as a generic “second Mamba version” than as a deliberately second-order efficient attention mechanism whose lineage runs through Mamba-2 and Mamba-2S.

4. The A-mask, convolution, and related variants

The A-mask is one of the paper’s central design elements. In the original Mamba-style form it is described as

$$O = \frac{\left[\phi(Q)\phi(K)^T\right]V}{\sum_j \phi(Q)\phi(K)^T} = \frac{\phi(Q)\left[\phi(K)^T V\right]}{\sum_j \phi(Q)\phi(K)^T},$$4

which couples the decay mask directly to the discretization parameter (Mongaras et al., 19 Feb 2026). In Mamba-2S and 2Mamba, the improved form is

$$O = \frac{\left[\phi(Q)\phi(K)^T\right]V}{\sum_j \phi(Q)\phi(K)^T} = \frac{\phi(Q)\left[\phi(K)^T V\right]}{\sum_j \phi(Q)\phi(K)^T},$$5

The authors give three reasons for the importance of this change: it is still strictly negative, preserving decay behavior; it is learned directly rather than tied to discretization; and it empirically improves accuracy more than the original A-mask (Mongaras et al., 19 Feb 2026). The decay matrix is then built as

$$O = \frac{\left[\phi(Q)\phi(K)^T\right]V}{\sum_j \phi(Q)\phi(K)^T} = \frac{\phi(Q)\left[\phi(K)^T V\right]}{\sum_j \phi(Q)\phi(K)^T},$$6

which acts as a multiplicative temporal bias favoring recent information.

The paper explicitly claims that the softplus A-mask is the most impactful isolated improvement among Mamba-2 components, and that convolution is the second strongest (Mongaras et al., 19 Feb 2026). This finding is important because it recasts Mamba-2’s complexity: the empirical gains do not seem to require retaining every architectural feature of the original block.

The work also introduces an exponentiated variant, 2Mamba-E, which uses

$$O = \frac{\left[\phi(Q)\phi(K)^T\right]V}{\sum_j \phi(Q)\phi(K)^T} = \frac{\phi(Q)\left[\phi(K)^T V\right]}{\sum_j \phi(Q)\phi(K)^T},$$7

instead of $$O = \frac{\left[\phi(Q)\phi(K)^T\right]V}{\sum_j \phi(Q)\phi(K)^T} = \frac{\phi(Q)\left[\phi(K)^T V\right]}{\sum_j \phi(Q)\phi(K)^T},$$8, together with the A-mask and convolution (Mongaras et al., 19 Feb 2026). The authors report that exponentiating the inner product yields slightly better accuracy than softmax attention in their tested medium-model setting, but requires a KV cache and therefore loses the pure linear-memory advantage at inference. They also connect this construction to the Forgetting Transformer through

$$O = \frac{\left[\phi(Q)\phi(K)^T\right]V}{\sum_j \phi(Q)\phi(K)^T} = \frac{\phi(Q)\left[\phi(K)^T V\right]}{\sum_j \phi(Q)\phi(K)^T},$$9

arguing that the A-mask and convolution are not merely Mamba-specific tricks but can also benefit softmax-like attention variants (Mongaras et al., 19 Feb 2026).

5. Complexity, hidden-state size, and memory efficiency

Although 2Mamba uses a second-order hidden state, the paper emphasizes that it remains linear in sequence length (Mongaras et al., 19 Feb 2026). The cost of the higher-order state is shifted from sequence length dependence to feature-dimension dependence.

For second-order features, the number of unique quadratic terms is

$\phi(\cdot)=\text{ReLU}$0

rather than $\phi(\cdot)=\text{ReLU}$1, because of symmetry and repeated terms (Mongaras et al., 19 Feb 2026). This reduction is used to argue that the state size, while larger than first-order linear attention, can still be smaller than softmax’s KV cache for sufficiently long sequences.

The paper derives the condition

$\phi(\cdot)=\text{ReLU}$2

and therefore the threshold

$\phi(\cdot)=\text{ReLU}$3

For the head dimension used in the experiments, $\phi(\cdot)=\text{ReLU}$4, this threshold is about

$\phi(\cdot)=\text{ReLU}$5

(Mongaras et al., 19 Feb 2026). In other words, beyond roughly one thousand tokens, the second-order 2Mamba state is described as memory-friendlier than the softmax KV cache.

This memory analysis is one of the paper’s strongest claims. It explains why the authors present 2Mamba not only as an accuracy-improved linear attention mechanism, but as a practical long-context alternative whose competitiveness increases with sequence length (Mongaras et al., 19 Feb 2026).

6. Experimental protocol, empirical findings, and stability

The experiments are performed in a Llama 2-style architecture where the attention block is replaced with candidate linear or Mamba variants (Mongaras et al., 19 Feb 2026). The main setup uses next-token language modeling on FineWeb (CC-MAIN-2024-51), with 100k training steps for most experiments, batch size 32 across devices, AdamW with betas $\phi(\cdot)=\text{ReLU}$6, learning rate $\phi(\cdot)=\text{ReLU}$7, warmup of 10k steps, weight decay 0.01, and no gradient clipping. The reported model scales are approximately 300M and 700M parameters, and the sequence lengths tested are 2048, 4096, and 8192 (Mongaras et al., 19 Feb 2026). Additional evaluations include long training at 400k steps, NIAH retrieval tests, and tests on The Pile and SlimPajama.

The main empirical conclusions are qualitative but specific. Mamba-2 is already much better than naive linear attention; Mamba-2S recovers most of this gain with fewer components; and 2Mamba closes most of the remaining gap to softmax attention (Mongaras et al., 19 Feb 2026). The central summary is that 2Mamba is nearly as accurate as softmax attention while remaining linear in sequence length.

The component ablations are equally important. The softplus A-mask is identified as the strongest single enhancement, convolution as the second strongest, discretization as mildly helpful, and the $\phi(\cdot)=\text{ReLU}$8-gate and $\phi(\cdot)=\text{ReLU}$9-residual as not worth keeping in the simplified model (Mongaras et al., 19 Feb 2026). This is a stronger claim than mere incremental tuning: it implies that the path to competitive linear attention may depend more on a small number of carefully chosen inductive biases than on preserving the entire Mamba-2 recipe.

The paper also reports a stability caveat for the squared model. Adding discretization to the squared version can create numerical instability, especially for the medium model, because larger values can blow up over training steps (Mongaras et al., 19 Feb 2026). To stabilize training, discretization is removed in the medium-scale squared variant, and the final 2Mamba model therefore trades part of the Mamba-2S structure for robustness.

7. Terminological scope and distinction from other “2Mamba” usages

A recurring source of confusion is that the term “2Mamba” can also surface in searches for models that are “2D Mamba,” “dual Mamba,” or “bidirectional Mamba,” even when those methods address different problems. In the paper under discussion, however, 2Mamba refers specifically to the second-order hidden-state architecture derived from Mamba-2S (Mongaras et al., 19 Feb 2026).

This distinction matters because several nearby Mamba-family names encode different meanings. “2DMamba,” for example, is a native 2D selective SSM for image representation and giga-pixel whole slide image classification rather than a second-order linear-attention LLM (Zhang et al., 2024). “V2M” generalizes the state space model itself to the 2-dimensional space so that image tokens are processed directly in 2D rather than through flattening (Wang et al., 2024). “DTMamba,” short for Dual Twin Mamba, is a time-series forecasting model whose relevance comes from using two Mamba modules inside each block plus two TMamba blocks overall (Wu et al., 2024). “LBMamba” is a locally bi-directional SSM block for vision that embeds a lightweight locally backward scan inside the forward selective scan and executes it entirely in per-thread registers (Zhang et al., 19 Jun 2025).

This suggests that “2Mamba” has become a potentially ambiguous search term in the broader Mamba literature. In strict usage, however, the term designates the second-order Mamba-2-derived model introduced in “2Mamba2Furious,” whose defining contribution is to combine a strong decay mask, local convolution, and a second-order hidden state in order to make linear attention competitive with softmax attention (Mongaras et al., 19 Feb 2026).