Mamba-2S: Minimal Linear-Attention Model

• Mamba-2S is a simplified variant that retains essential linear-attention components while removing superfluous elements to improve model efficiency.
• It employs a strictly negative softplus-based A-mask, lightweight input convolution, and time discretization to optimize stability and test loss.
• The model serves as a robust baseline for further developments like 2Mamba, offering a linear complexity state-space formulation for practical applications.

Searching arXiv for the cited paper and closely related works mentioned in the supplied data. Mamba-2S is a simplified variant of Mamba-2 that retains only the most fundamental components that empirically drive accuracy while removing architectural elements that add complexity without clear benefit. Introduced in "2Mamba2Furious: Linear in Complexity, Competitive in Accuracy" (Mongaras et al., 19 Feb 2026), it is positioned as a minimal linear-attention construction built around a strictly negative softplus-based A-mask, a lightweight input convolution on $Q$, $K$, and $V$, time discretization applied to values, output RMSNorm, and a causal selective state mechanism. Within the same work, Mamba-2S serves both as an ablated endpoint in its own right and as the base from which 2Mamba is derived.

1. Origin, scope, and defining objective

Mamba-2S is defined as a simplified variant of Mamba-2 that retains only the most fundamental and important components while discarding elements whose empirical contribution is negligible or negative. The stated objective is to identify which specific choices make Mamba-2 most accurate, then preserve those choices in a simpler model. In that sense, Mamba-2S is not merely a smaller implementation of Mamba-2; it is an ablation-derived re-specification of the architecture (Mongaras et al., 19 Feb 2026).

The simplification has a specific technical motivation. Linear attention transformers are described as a strong alternative to softmax attention because of efficiency, but linear attention tends to be less expressive and results in reduced accuracy compared to softmax attention. Mamba-2S is therefore designed as a more accurate linear-attention variant by isolating the components that most improve test loss. This suggests that the model should be understood less as a generic state-space simplification and more as a targeted answer to the accuracy deficit of linear attention.

Within the paper’s overall progression, Mamba-2S is the first-order formulation. It preserves the linear-in-$N$ recurrence and selective forgetting mechanism, then provides the structural basis for 2Mamba, which increases the order of the hidden state from $p=1$ to $p=2$.

2. Architectural simplification of Mamba-2

The defining simplifications from Mamba-2 to Mamba-2S are explicit and component-wise. The most consequential modification is the replacement of the original A-mask parametrization

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

with the decoupled, strictly negative parameterization

$A=-\operatorname{softplus}(h\cdot W_A).$

The rationale given is that this enforces negativity for stability and decouples forgetting from $dt$ for greater flexibility; it is also described as the empirically most impactful ablation.

A lightweight input convolution on $Q$, $K$0, and $K$1 is retained, with kernel size $K$2 and no activation. Its function is to improve inductive bias and accuracy with minimal overhead, and it is reported as the empirically second-most impactful component. Mamba-2S also keeps time discretization

$K$3

applied to values through

$K$4

The paper states that this gives a small additional gain at small scale, although it is later removed for larger models in 2Mamba for numeric stability.

Output normalization is handled by RMSNorm after attention rather than by softmax-like normalization on $K$5. The stated rationale is that this allows inner-product values to be signed, enabling a broader ablation space and avoiding constraints on the $K$6 image being strictly positive. By contrast, the multiplicative $K$7 gate, the additive $K$8 residual, and the SiLU activations in the QKV convolution are removed. Their removal is justified by ablations showing negligible or negative gains and by the resulting reduction in architectural complexity (Mongaras et al., 19 Feb 2026).

A compact summary of the retained and removed components is as follows.

Category	Components
Retained	Softplus A-mask, input convolution with window size 2 on $K$9, $V$0, $V$1, time discretization applied to $V$2, output RMSNorm, causal mask
Removed or altered	Original A-log plus exp plus dt coupling for $V$3, $V$4 gate, $V$5 residual, SiLU in conv, specialized dt_bias initialization

The resulting model is explicitly described as using a minimal conv-only setup and learning $V$6 and $V$7 directly.

3. State-space formulation and recurrent realization

Mamba-2S is presented from the state-space model perspective. The continuous-time SSM is written as

$V$8

In the Mamba family, $V$9 is interpreted as a per-head forget or decay rate, while $N$0 and $N$1 play roles analogous to key, value, and query projections and their combinations. Under discretization with step $N$2, the dynamics become

$N$3

with

$N$4

Mamba-2S introduces per-head scalar decay $N$5 at each position $N$6, constrained to be strictly negative through

$N$7

The softplus constraint is stated to enforce negativity and stability, while time variation across positions and heads enables selectivity. The discretization variable is

$N$8

which is strictly positive and is applied to values through

$N$9

The first-order attention form is

$p=1$0

where output RMSNorm is applied before projection to $p=1$1. The semiseparable structure induced by $p=1$2 yields an efficient $p=1$3 recurrent realization. Defining a recurrent state $p=1$4,

$p=1$5

and the output is obtained by per-head contraction

$p=1$6

The paper states that this recurrence is algebraically equivalent to the masked attention

$p=1$7

This formulation is central to the identity of Mamba-2S. It is simultaneously an attention-like model, because it is written in masked $p=1$8 form, and an SSM-like model, because it admits a scan recurrence with a bounded decay factor.

4. The A-mask and selective forgetting

The A-mask is the key mechanism of Mamba-2S. It is defined through

$p=1$9

followed by a cumulative sum over time,

$p=2$0

and then the mask itself,

$p=2$1

with the causal mask $p=2$2. The effective attention weight is therefore

$p=2$3

The constraint $p=2$4 implies

$p=2$5

which is used to guarantee stability via exponential forgetting. In the paper’s formulation, sparsity comes from the causal mask, while semiseparability comes from exponentials of differences of cumulative sums. This gives Mamba-2S a selective state mechanism in which past information is retained or attenuated through learned, head-wise, time-varying decay.

The work contrasts this A-mask with the original Mamba-2 choice, which couples decay to $p=2$6. The improved form is described as decoupled from $p=2$7, smoother to learn because of softplus curvature, and empirically better in test loss. On the FineWeb small-model ablation, the isolated A-mask comparison is reported as test loss $p=2$8 versus $p=2$9, and with convolution of width $A=-\exp(A_{\log})\odot dt$0 the comparison is $A=-\exp(A_{\log})\odot dt$1 versus $A=-\exp(A_{\log})\odot dt$2 (Mongaras et al., 19 Feb 2026).

A plausible implication is that the A-mask in Mamba-2S is doing more than imposing stable forgetting. By decoupling forgetting from value scaling, it separates temporal attenuation from discretization, which broadens the effective design space without changing the linear-in-sequence recurrence.

5. Complexity, hidden-state order, and relation to 2Mamba

As a first-order model, Mamba-2S corresponds to hidden-state order $A=-\exp(A_{\log})\odot dt$3. In the paper’s terminology, the order of the hidden state is the degree of polynomial features used to approximate softmax attention through a Taylor or Maclaurin expansion of $A=-\exp(A_{\log})\odot dt$4. Linear attention is $A=-\exp(A_{\log})\odot dt$5, whereas 2Mamba sets $A=-\exp(A_{\log})\odot dt$6 by squaring the $A=-\exp(A_{\log})\odot dt$7 inner product. Mamba-2S therefore occupies the first-order endpoint of this hierarchy.

Its per-layer, per-head recurrence has time complexity $A=-\exp(A_{\log})\odot dt$8, linear in sequence length $A=-\exp(A_{\log})\odot dt$9. The hidden-state memory is $A=-\operatorname{softplus}(h\cdot W_A).$0 per head. By contrast, 2Mamba has time complexity

$A=-\operatorname{softplus}(h\cdot W_A).$1

and per-head memory

$A=-\operatorname{softplus}(h\cdot W_A).$2

which is constant in $A=-\operatorname{softplus}(h\cdot W_A).$3 but larger in the hidden feature dimension. Softmax attention is described as requiring $A=-\operatorname{softplus}(h\cdot W_A).$4 memory per head because of the KV cache, specifically $A=-\operatorname{softplus}(h\cdot W_A).$5 elements (Mongaras et al., 19 Feb 2026).

The paper gives a threshold comparison for $A=-\operatorname{softplus}(h\cdot W_A).$6 per head:

$A=-\operatorname{softplus}(h\cdot W_A).$7

so that for $A=-\operatorname{softplus}(h\cdot W_A).$8, $A=-\operatorname{softplus}(h\cdot W_A).$9. Beyond approximately $dt$0K context, 2Mamba is reported to be more memory efficient than softmax with comparable accuracy. Mamba-2S is not described as matching softmax as closely as 2Mamba, but it is described as outperforming vanilla linear attention by a large margin and approaching Mamba-2 and softmax.

This positioning is important for interpretation. Mamba-2S is the low-order, simpler, robust baseline; 2Mamba is the higher-order construction that uses the same A-mask foundation to close the remaining accuracy gap to softmax attention.

6. Empirical behavior, implementation, and limitations

The experimental setup replaces attention blocks in a Llama-2-like base architecture with Mamba-2 variants and evaluates next-token prediction on HuggingFace FineWeb (CC-MAIN-2024-51), with $dt$1 held out and context lengths $dt$2, $dt$3, and $dt$4. Additional datasets are The Pile and SlimPajama. For the small model, the reported architecture is hidden size $dt$5, MLP $dt$6 with SiLU, $dt$7 heads with $dt$8, $dt$9 layers, and vocab $Q$0K; for the medium model, hidden size $Q$1, MLP $Q$2 with SiLU, $Q$3 heads with $Q$4, $Q$5 layers, and vocab $Q$6K. Optimization uses AdamW with $Q$7, $Q$8, learning rate $Q$9, weight decay $K$00, no grad clipping, warmup $K$01K steps, total steps reported as $K$02K/$K$03K, and batch size $K$04 (Mongaras et al., 19 Feb 2026).

The key FineWeb ablation results for the small model are:

Variant	Test loss
Plain linear attention (output norm)	3.13
+ Conv (w=2)	2.93
+ Conv + original A-mask	2.80
+ Conv + softplus A-mask	2.77
+ Conv + softplus A + discretization	2.76

The final row is identified as the Mamba-2S baseline. The paper further states that the $K$05 gate or $K$06 residual do not improve further, and that the $K$07 gate slightly hurts. The softplus A-mask and input convolution are identified as the most impactful components, while value discretization gives a modest gain at small scale.

Implementation emphasizes fused Triton kernels for attention with the A-mask, associative scan or cumulative sum on $K$08 to build $K$09 efficiently, and chunking for long sequences. The provided pseudocode keeps the structure minimal:

• input projections via conv1d with kernel size $K$10;
• decay and discretization through $K$11 and $K$12;
• recurrent update $K$13;
• output $K$14, followed by RMSNorm and projection.

The principal limitation attached to Mamba-2S is the role of time discretization. The paper states that $K$15 improves small-scale accuracy but can cause numerical instability for larger models under TF32; it stabilizes under FP32 at substantial cost, and is therefore removed in some 2Mamba runs. More generally, higher hidden-state order improves expressivity but increases per-token compute and hidden-state memory, so additional kernel engineering is required to realize speed gains.

In practical guidance, Mamba-2S is recommended when a minimal, robust improvement over vanilla linear attention is desired, specifically with strictly negative $K$16, convolution, $K$17, and output RMSNorm, while preserving linear complexity. The same guidance describes it as suitable for moderate sequence lengths and small or medium models. This suggests that Mamba-2S is best viewed as the simplified and stable first-order member of the paper’s family: more accurate than naive linear attention, architecturally leaner than the original Mamba-2, and foundational for the higher-order 2Mamba construction.