2Mamba2Furious: Linear in Complexity, Competitive in Accuracy

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

💡 Research Summary

The paper tackles the well‑known trade‑off between the computational efficiency of linear attention and the superior accuracy of softmax‑based attention. While softmax attention incurs quadratic time and memory costs with respect to sequence length, linear attention replaces the softmax kernel with a decomposable feature map, achieving O(N) complexity but typically lagging in performance. The authors focus on Mamba‑2, a recent linear‑attention architecture that incorporates state‑space modeling, an associative‑scan algorithm, and a decay mask (A‑mask) to improve both speed and expressiveness.

Their first contribution is a systematic dissection of Mamba‑2’s implementation. By isolating each architectural choice—query/key/value activation (SiLU, ReLU, none), A‑mask formulation (original negative‑exp vs. softplus), input convolution window size, additive D‑residual, multiplicative Z‑gate, normalization strategy (output RMSNorm vs. softmax‑style), and value discretization—they conduct an extensive ablation study on a 300‑million‑parameter Llama‑2‑derived model trained on the massive CC‑MAIN‑2024‑51 dataset. The experiments reveal that (1) adding a simple 1‑D convolution (kernel size 2) to the QKV projection yields the largest baseline gain, (2) replacing the original A‑mask (A = −exp(A_log)·dt) with a smooth softplus version (A = −softplus(A)) improves stability and reduces memory pressure, and (3) ReLU activation combined with softmax‑style normalization on the query‑key inner product is nearly as effective as the more complex RMSNorm used in the original code. Other components such as the D‑residual and Z‑gate provide marginal benefits, while value discretization can even harm performance.

Building on these insights, the authors propose two key modifications that together form the “2Mamba” model. First, they introduce a second‑order hidden state. Drawing on the Maclaurin expansion of the exponential function, they note that softmax attention can be viewed as an infinite sum of non‑negative powers of the QK inner product. Linear attention corresponds to the first‑order term; by extending the hidden state to include the square of the inner product (order p = 2), the model captures richer interactions while keeping the hidden state size at O(d_H²) per head—still smaller than the KV‑cache of softmax attention (2 × N × d_H). This higher‑order representation yields a strictly positive inner‑product space, allowing the use of softmax normalization without instability.

Second, they refine the A‑mask. The softplus‑based mask removes the need to couple the mask with the discretization parameter, resulting in a smoother decay profile that is easier to train and more memory‑efficient. Combined with the second‑order hidden state, 2Mamba achieves test loss reductions of roughly 5‑7 % over the original Mamba‑2 and narrows the gap to softmax attention to within 0.1‑0.2 loss units on the validation set. Importantly, for a context length of 2048 tokens, 2Mamba’s memory footprint is under 30 % of a comparable softmax‑based transformer, confirming its linear‑complexity advantage.

The paper’s contributions can be summarized as follows:

1. Component‑wise analysis of Mamba‑2, clarifying which architectural choices truly matter for accuracy.
2. A‑mask redesign using a softplus function, improving both stability and memory usage.
3. Higher‑order hidden state (second‑order RNN) that bridges the expressivity gap between linear and softmax attention without incurring quadratic costs.
4. Empirical validation on large‑scale language modeling, demonstrating that 2Mamba matches or exceeds softmax performance while retaining linear time and space complexity.

The authors acknowledge that the second‑order hidden state’s d_H² scaling may become prohibitive for very large head dimensions, suggesting future work on low‑rank approximations or adaptive order selection. They also propose exploring KV‑cache integration for inference speedups and extending the evaluation to downstream tasks such as translation, summarization, and code generation. Overall, the work presents a compelling pathway to make linear attention not only fast but also competitively accurate, positioning 2Mamba as a strong candidate for long‑context language models.