PaperlayerView exact revision on arXiv

Source-linked AI summary

2Mamba2Furious: Linear in Complexity, Competitive in Accuracy

Gabriel Mongaras, Eric C. Larson

arXiv:2602.17363v3cs.LG

TL;DR

Softmax attention is accurate but has quadratic training complexity, while linear attention is more efficient yet less expressive. This paper simplifies Mamba-2 and adds a higher-order hidden state, producing 2Mamba with softmax-level accuracy and constant memory, while 2Mamba-E surpasses softmax attention accuracy.

  • Problem

    Softmax attention has quadratic training complexity, whereas linear attention offers efficiency but tends to be less expressive and less accurate.

  • Method

    The paper isolates Mamba-2's important components to build Mamba-2S, then uses a softplus A-mask and second-order hidden state to create 2Mamba and an exponentiated variant, 2Mamba-E.

  • Results

    2Mamba achieves accuracy comparable to softmax attention with constant memory, while 2Mamba-E is better than softmax attention.

  • Takeaways & Limitations

    The results support bridging linear and softmax attention by increasing hidden-state order while retaining lower memory requirements for sufficiently long contexts.

  • Takeaways & Limitations

    Time discretization can cause numerical instability and divergence in medium models, while stabilization with IEEE FP32 input precision is about 8x slower than TF32.

Abstract

from arXiv · show

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.

1 Introduction

Linear attention offers better training and inference complexity than softmax attention, but its lower expressiveness leaves an accuracy gap. This work modifies Mamba-2 to approach softmax-level accuracy while retaining linear complexity.

  • Motivation: Softmax attention routes information across tokens accurately but has quadratic training complexity in sequence length.Flash Attention improves implementation efficiency without changing the underlying quadratic FLOP count.
  • Motivation: Linear attention replaces softmax’s exponential nonlinearity with a decomposable kernel, enabling linear training complexity and constant inference complexity.It can also be implemented as an RNN for efficient long-sequence inference.
  • Motivation: Mamba and DeltaNet improve vanilla linear attention’s accuracy through decay or delta-rule mechanisms, but remain less accurate than full softmax attention.The paper targets this remaining accuracy gap.
  • Contributions: The paper simplifies Mamba-2, increases hidden-state order, and uses the resulting designs to improve accuracy while keeping model complexity linear.It also studies an exponentiated query-key inner product that can outperform softmax attention but requires a KV cache.

2 Background

The background contrasts softmax and linear attention, then motivates Mamba-2’s component analysis and higher-order hidden states as routes toward greater expressiveness. Higher-order states can approach softmax accuracy, but their memory cost grows rapidly with order.

  • Softmax Attention: Softmax attention’s exponential term causes quadratic training FLOPs and linear inference FLOPs with respect to sequence length.Flash Attention improves GPU utilization but leaves the FLOP count quadratic.
  • Linear Attention: Linear attention replaces softmax with a kernel feature map, allowing computation to be reordered to avoid sequence-length quadratic complexity.The key trade-off is reduced computational complexity alongside lower expressiveness than softmax attention.
  • Mamba-2: Mamba-2 improves linear attention through an associative scan and decay mask, while its implementation also includes additional inductive biases.The paper questions which components are necessary for its accuracy.
  • Softmax as a Recurrent Neural Network: Increasing linear-attention order increases hidden-state size on the order of dh^p, making expressiveness gains more expensive in memory.First-order linear attention has hidden-state size (dh, dh), while higher orders scale approximately as (dh^p, dh).
  • Softmax as a Recurrent Neural Network: Softmax attention can be represented as a recurrent computation whose hidden state is the KV cache, requiring 2 × N × dh elements per head.This memory grows linearly with sequence length.
  • Softmax as a Recurrent Neural Network: Higher-order hidden states add powers of the query-key inner product and improve downstream accuracy toward softmax accuracy.A second-order state has dimension roughly (dh^2, dh), offering a memory-feasible compromise for sufficiently long contexts.

3 Isolating Mamba-2 Accuracy Gains

The study isolates Mamba-2 components through controlled Llama 2 ablations, identifying the A-mask and convolution as the largest accuracy contributors relative to vanilla linear attention. It then uses these findings to construct a simpler architecture while documenting normalization choices and evaluation constraints.

  • Experimental setup: Ablations replace Llama 2 softmax attention with Mamba-2 blocks in approximately 300-million-parameter models, with larger models reserved for scalability tests.Experiments use FineWeb next-token language modeling and evaluate base, 4096, and 8192 sequence lengths.
  • Isolated components: The study varies QK activation, A-mask type, convolution window, additive residual, multiplicative gate, normalization, and value discretization.The A-mask comparison includes original masking and a softplus variant that removes the discretization parameter.
  • Normalization and A-mask: Output normalization is used for Mamba-2 ablations because QK normalization requires a strictly positive inner-product image, whereas Mamba-2 queries and keys are not necessarily positive.The authors note that softmax normalization cannot be tested by itself without an activation function.
  • Evaluation caveat: The training code evaluates one step after each 10K-step interval rather than the intended checkpoints, and the authors do not rerun all experiments because retraining would take months.Reported comparisons therefore use test loss through 90K steps rather than a correctly evaluated 100K-step endpoint.
  • Convolution: Adding a convolution with kernel size 2 substantially improves normal linear attention, while larger windows add only small gains and require more memory.The authors therefore select a window size of 2 without an activation function for a minimal implementation.
  • Isolated components: The softplus A-mask is the most impactful isolated component, followed by convolution; other components provide only minor accuracy gains over vanilla linear attention.The value discretization method is the strongest among the minor improvements but remains less influential than the A-mask and convolution.

4 Building Up to the Mamba-2S Base Model

The authors build a minimal Mamba-2 variant by selecting components according to isolated ablation results, while avoiding exhaustive component-interaction testing. The resulting Mamba-2S retains the softplus A-mask, size-2 input convolution, and time discretization.

  • Buildup strategy: Because all component interactions are too numerous to test exhaustively, the buildup uses insights from isolated component ablations.The progression is shown in Figure 4a and summarized in Table 3.
  • Buildup strategy: The buildup starts from an output-norm variant of linear attention and incrementally adds components.The authors use this base to construct a minimal model comparable in accuracy to original Mamba-2.
  • Mamba-2S: Mamba-2S uses a softplus A-mask, an input convolution of size 2, and time discretization.These are the components retained in the simplified algorithm.

5 Mamba-2 with a Squared Hidden State

2Mamba squares the QK inner product in a simplified Mamba-2 variant, achieving nearly softmax-level accuracy while reducing memory usage beyond sufficiently long contexts.

  • 5 Mamba-2 with a Squared Hidden State: Squaring the QK inner product produces a significant accuracy gain and reaches softmax-level accuracy in the simplified Mamba-2 variant.The method removes RMSNorm, uses softmax normalization, and replaces the QK inner product with a squared inner product.
  • 5 Mamba-2 with a Squared Hidden State: Removing time discretization for the medium model preserves numerical stability in low-precision kernels because discretization can cause value growth and training divergence.Discretization offers minimal accuracy gains relative to adding convolution and the softplus A-mask, but can destabilize medium-model training.
  • 5 Mamba-2 with a Squared Hidden State: 2Mamba is nearly as accurate as softmax attention across model sizes and sequence lengths.The comparison covers small and medium models trained with sequence lengths of 2048, 4096, and 8192.
  • 5 Mamba-2 with a Squared Hidden State: The square is applied to the QK inner product rather than the A-mask because squaring the A-mask provides no accuracy gain and introduces unnecessary squared mask terms.The A-mask is applied directly to attention scores derived from the squared QK inner product.
  • 5.1 2Mamba Algorithm Efficiency: For head dimension 64, sequence lengths above approximately 1058 require less memory with the squared hidden state than with softmax attention.The paper verifies this crossover by accumulating raw hidden-state memory for both methods.
  • 5 Mamba-2 with a Squared Hidden State: The NIAH evaluation tests retrieval of facts across long prompts after training for 400,000 steps with a maximum context of 8192 tokens.The reported result provides evidence that 2Mamba can use its context beyond next-token prediction tasks.

6 2Mamba With an Exponentiated Hidden State

The paper exponentiates the QK inner product to form 2Mamba-E and reports slightly better accuracy than softmax attention, connecting the design to softmax attention with a softplus A-mask.

  • 6 2Mamba With an Exponentiated Hidden State: Higher-order hidden states can improve accuracy with diminishing gains, while the limiting exponential form yields softmax attention.This motivates testing exponentiation after the squared hidden-state construction.
  • 6 2Mamba With an Exponentiated Hidden State: 2Mamba-E exponentiates the QK inner product and achieves slightly better accuracy than softmax attention.The method can also be interpreted as softmax attention using a softplus A-mask and input convolution.
  • 6 2Mamba With an Exponentiated Hidden State: The softplus A-mask used here is nearly equivalent to the forgetting transformer's log(sigmoid(x)) construction, differing only by sign.The paper presents this as an intuitive connection between Mamba-2-style state-space models and transformer variants.

7 Conclusion and Future Work

The conclusion presents 2Mamba as a more efficient linear-attention model built from a simplified Mamba-2, while identifying extensions and hidden-state choices for future work.

  • 7 Conclusion and Future Work: 2Mamba builds on a minimal Mamba-2S implementation and uses a second-order hidden state.The paper examines individual Mamba-2 components before increasing the hidden-state order.
  • 7 Conclusion and Future Work: 2Mamba is reported to match softmax-attention accuracy while remaining constant in memory.The conclusion frames this as progress toward bridging linear- and softmax-attention architectures more efficiently in practice.
  • 7 Conclusion and Future Work: The exponentiated variant 2Mamba-E is reported to outperform softmax attention and resemble the forgetting transformer.This variant extends the implementation beyond the base 2Mamba model.
  • 7 Conclusion and Future Work: Future work could combine DeltaNet with 2Mamba to further improve the algorithm.The paper notes that Mamba-2 and DeltaNet have already been combined in Gated DeltaNet.
  • 7 Conclusion and Future Work: Future work could vary hidden-state size and compare squared query-key inner products with equal-sized query and key projections.Optimizing hidden-state size could improve 2Mamba’s memory efficiency.

A Inference Algorithm

The inference procedure caches only the previous window_size−1 queries, keys, and values when using an input convolution; with window size 2, one of each is cached.

  • A Inference Algorithm: Input convolution requires caching window_size−1 previous queries, keys, and values.The passage gives the window-size dependency for the inference cache.
  • A Inference Algorithm: With window size 2, inference caches one previous query, key, and value.The example illustrates the general caching rule concretely.

B Model Ablation Details

The ablation studies replace attention throughout Llama 2 while keeping other architectural components fixed, using multiple sequence lengths, model scales, and a specified training setup.

  • B Model Ablation Details: Ablated attention mechanisms replace the attention mechanism in all Llama 2 layers, with other architecture components held constant.This isolates changes to the attention mechanism.
  • B Model Ablation Details: Experiments use a base maximum sequence length of 2048, with 4096 and 8192 tested for longer contexts.Most experiments use approximately 300 million parameters, with a 700-million-parameter model used for scale testing.
  • B Model Ablation Details: Training uses FineWeb, a held-out 0.1 percent test split, AdamW, batch size 32, learning rate 1e-4, 100,000 steps, and weight decay 0.01.The setup also specifies 10,000 warmup steps and no gradient clipping.
  • B Model Ablation Details: The 300-million-parameter model uses hidden size 1024, intermediate MLP size 2048, 16 attention heads, 20 layers, and vocabulary size 32,000.Its hidden MLP activation is SiLU and each head has dimension 64.
  • B Model Ablation Details: The 700-million-parameter model uses hidden size 1536, intermediate MLP size 3072, and 24 attention heads with 64-dimensional heads.The passage identifies this as the medium model used for scale testing.

C Gradients

The gradients section derives gradients for the required attention kernels, treating query, key, and value matrices as inputs and post-attention outputs as outputs.

  • C Gradients: Each gradient is derived for an attention kernel whose inputs are query, key, and value matrices.The kernels perform the attention operation without preprocessing or postprocessing.
  • C Gradients: The kernel outputs are the post-attention outputs and exclude preprocessing and postprocessing.This defines the scope of the gradient equations.

C.1 Linear

This section distinguishes row/key-value and column/query axes to derive gradients for the A-mask, then describes generalization runs on The Pile and SlimPajama.

  • C.1 Linear: M denotes the row, key, and value axis, while N denotes the column, query axis, and causal mask notation also uses M.The notation separates axis roles even though M also denotes the causal mask.
  • C.1 Linear: P_j sums along the M axis, whereas P_i sums along the N axis.This distinction is needed for the A-mask gradient derivation.
  • C.1 Linear: The axis distinction is retained for gradient derivation even in self-attention, where N = M.The equality of the axes in self-attention does not remove the notational distinction.
  • C.1 Linear: 2Mamba generalization is evaluated with medium-sized models trained on The Pile and SlimPajama.The accompanying plots report test loss on each dataset.
Loading 2602.17363v3…