PaperlayerView exact revision on arXiv

Source-linked AI summary

Mamba-3: Improved Sequence Modeling using State Space Principles

Aakash Lahoti, Kevin Y. Li, Berlin Chen, Caitlin Wang, Aviv Bick, J. Zico Kolter, Tri Dao, Albert Gu

arXiv:2603.15569v1cs.LG

TL;DR

Mamba-3 addresses the quality, capability, and hardware-efficiency limitations of linear sequence models by combining SSM-inspired recurrence improvements, complex-valued state updates, and MIMO processing. It improves downstream accuracy, state tracking, and the performance-efficiency Pareto frontier while maintaining competitive inference efficiency.

  • Problem

    Linear models offer efficient inference but can sacrifice quality and state-tracking capability, while their theoretically linear decoding can remain hardware-inefficient.

  • Method

    Mamba-3 combines exponential-trapezoidal discretization, complex-valued state transitions, and a MIMO formulation with architectural refinements.

  • Results

    At 1.5B scale, Mamba-3 (SISO) improves downstream accuracy by +0.6 points over GDN, while Mamba-3 (MIMO) improves by +1.8 points over GDN and matches Mamba-2 perplexity with half its state size.

  • Takeaways & Limitations

    Mamba-3 advances the performance-efficiency Pareto frontier, with MIMO providing stronger modeling power while maintaining competitive inference efficiency compared to Mamba-2.

  • Takeaways & Limitations

    The discretization result is stated under standard stability assumptions, including a trapezoidal parameter satisfying λ_t = 1.

Abstract

from arXiv · show

Scaling inference-time compute has emerged as an important driver of LLM performance, making inference efficiency a central focus of model design alongside model quality. While the current Transformer-based models deliver strong model quality, their quadratic compute and linear memory make inference expensive. This has spurred the development of sub-quadratic models with reduced linear compute and constant memory requirements. However, many recent linear models trade off model quality and capability for algorithmic efficiency, failing on tasks such as state tracking. Moreover, their theoretically linear inference remains hardware-inefficient in practice. Guided by an inference-first perspective, we introduce three core methodological improvements inspired by the state space model (SSM) viewpoint of linear models. We combine: (1) a more expressive recurrence derived from SSM discretization, (2) a complex-valued state update rule that enables richer state tracking, and (3) a multi-input, multi-output (MIMO) formulation for better model performance without increasing decode latency. Together with architectural refinements, our Mamba-3 model achieves significant gains across retrieval, state-tracking, and downstream language modeling tasks. At the 1.5B scale, Mamba-3 improves average downstream accuracy by 0.6 percentage points compared to the next best model (Gated DeltaNet), with Mamba-3's MIMO variant further improving accuracy by another 1.2 points for a total 1.8 point gain. Across state-size experiments, Mamba-3 achieves comparable perplexity to Mamba-2 despite using half of its predecessor's state size. Our evaluations demonstrate Mamba-3's ability to advance the performance-efficiency Pareto frontier.

1 Introduction

Mamba-3 targets the quality, capability, and hardware-efficiency limitations of prior linear sequence models through three SSM-inspired changes. Its variants improve downstream accuracy, state tracking, and the performance-latency tradeoff.

  • Motivation: Prior linear models often sacrifice expressivity and state-tracking capability for efficiency, while their theoretically efficient inference can remain hardware-inefficient.Mamba-2 traded some expressivity for training speed and simplicity, and sub-quadratic inference can be limited by hardware utilization.
  • Method: Mamba-3 introduces exponential-trapezoidal discretization, a more expressive recurrence that can replace the short causal convolution with explicit B,C bias terms.The recurrence exposes an implicit convolution over the SSM input.
  • Method: Complex-valued state transitions provide a richer update rule that addresses state-tracking limitations and are equivalent to an efficiently computed data-dependent rotary embedding.The paper evaluates this mechanism on synthetic tasks beyond prior linear models’ capabilities.
  • Method: MIMO replaces an outer-product state update with matrix multiplication, adding computation during memory-bound decoding without increasing state size or compromising speed.This formulation corresponds to moving from SISO to MIMO sequence dynamics.
  • Results: +1.8 points is Mamba-3 (MIMO)’s downstream accuracy gain over Gated DeltaNet at 1.5B scale, while SISO gains +0.6 points.MIMO also matches Mamba-2’s perplexity at half its state size, using state size 64 versus 128.
  • Results: Up to 4× more decoding FLOPs are achieved by Mamba-3 (MIMO) than Mamba-2 at fixed state size while maintaining similar wall-clock decode latency and improving perplexity and downstream performance.The authors also report that both Mamba-3 variants advance the performance-latency Pareto frontier.

2 Preliminaries

SSMs model sequence processing through continuous-time linear dynamics and discretized recurrences, while SSD provides an equivalent parallel form. Mamba-3 extends this framework with a more expressive structured mask than Mamba-2.

  • SSM formulation: SSMs describe continuous-time linear dynamics using a hidden state, input, state-transition matrix, and input/output projections.The hidden state evolves through the transition and state-input terms, while the output is obtained from the state through C.
  • Mamba-2 parameterization: Mamba-2 parameterizes the transition as a scalar multiple of the identity, enabling efficient recurrence computation with GPU tensor cores.Its data-dependent step size controls both memory retention and the strength of the current-token contribution.
  • Mamba-2 parameterization: The data-dependent transition α_t controls each SSM’s memory horizon: larger Δ_t forgets faster, whereas smaller Δ_t retains more of the hidden state.The same step size also up-weights the current token for larger values and reduces its contribution for smaller values.
  • State space duality: SSD establishes a duality between recurrent SSMs and parallel matrix-multiplication forms using a structured element-wise mask L.The mask represents state-transition decay, and different structures on L produce different SSD instantiations.
  • State space duality: SSD also connects recurrence to attention by interpreting L as a data-dependent mask, with causal linear attention as its simplest case.This connection uses C as Q, B as K, and X as V.
  • Mamba-3 extension: Mamba-3 generalizes Mamba-2 within SSD by using a more expressive structured mask, induced by exponential-trapezoidal discretization.Figure 1 characterizes this mask as the product of decay and a two-band convolutional mask.

3 Methodology

Mamba-3 combines a more expressive exponential-trapezoidal recurrence, complex-valued state updates, and MIMO computation with architectural refinements to improve model quality, state tracking, and inference efficiency.

  • Mamba-3 introduces exponential-trapezoidal discretization, complex-valued state spaces, and MIMO to improve dynamics, state tracking, modeling power, and hardware utilization.These innovations are combined into an updated Mamba architecture block.
  • 3.1 Exponential-Trapezoidal Discretization: The generalized discretization framework converts continuous-time SSMs into recurrences that combine prior-state decay with state-input terms from adjacent time steps.It formalizes the prior Mamba discretization as exponential-Euler and extends it with the more expressive exponential-trapezoidal method.
  • 3.1 Exponential-Trapezoidal Discretization: Exponential-trapezoidal discretization induces a convolution on the state-input within the core recurrence, unlike standard short convolutions applied outside the recurrence.The resulting state-input operation can be viewed as a width-2 convolution before the linear recurrence.
  • 3.2 Complex-Valued State Space Model: Complex-valued SSM updates provide richer state dynamics and are implemented as data-dependent rotary embeddings on the input and output projections.The complex transition can be represented with real 2 × 2 rotation blocks and computed efficiently using the RoPE trick.
  • 3.2 Complex-Valued State Space Model: The complex SSM solves state-tracking tasks that real-valued SSMs with and without standard RoPE cannot.This empirical result supports the theoretical connection between complex SSMs and data-dependent RoPE embeddings.
  • 3.3 Multi-Input, Multi-Output: MIMO replaces the outer-product state update with a matrix-multiplication update, increasing arithmetic intensity and expressivity while practically preserving decoding speed.The recurrence computes R× more FLOPs, with the resulting matrix-matrix product using fast tensor cores at marginal latency cost.
  • 3.4 Mamba-3 Architecture: Mamba-3 adds RMS normalization after the B and C projections, and the resulting BCNorm can stabilize large-scale pure Mamba-3 runs.In hybrid models, however, the removed post-gate RMSNorm remains crucial for long-context extrapolation.
  • 3.4 Mamba-3 Architecture: Data-independent biases combined with exponential-trapezoidal discretization can remove the short causal convolution and its activation function.The paper reports this result empirically for Mamba-3.

4 Empirical Validation

Mamba-3 is evaluated across language modeling, retrieval, state tracking, ablations, and inference-efficiency tradeoffs. Results show stronger modeling with compact states, MIMO gains over SISO, and competitive practical latency.

  • Downstream evaluation: Mamba-3 outperforms popular baselines across downstream tasks at all four model scales.
  • MIMO evaluation: 0.11 perplexity gain and 1.2 percentage-point average gain distinguish MIMO Mamba-3 from its SISO variant at 1.5B scale.
  • Retrieval: Mamba-3 is competitive on associative recall, question answering, and synthetic NIAH retrieval, but struggles with semi-structured and unstructured extraction.
  • State tracking: Data-dependent RoPE enables Mamba-3 to solve Parity and unbracketed Modular Arithmetic while nearly closing the bracketed-task gap; comparable variants fail to learn these tasks.
  • Ablations: BC bias and exponential-trapezoidal discretization synergize, making the short convolution optional in Mamba-3.
  • Efficiency frontier: Mamba-3 achieves better perplexity than Mamba-2 at half the state size, while MIMO shifts the frontier further without increasing state size.
  • Inference efficiency: Mamba-3 SISO is fastest among recurrent models in end-to-end decoding, while MIMO is on par with Mamba-2 and incurs moderate prefill overhead.

5 Related Work

Related work frames Mamba-3 within linear attention, test-time training, and structured state-space models. The paper distinguishes its complex transitions and MIMO design from earlier efficiency- or simplification-driven approaches.

  • Frameworks: Linear-time sequence models span linear attention, test-time training, and state-space-model frameworks.
  • Expressivity: Prior recurrent models simplified state transitions for efficiency, but analyses identified limitations in expressivity and state tracking.
  • Complex SSMs: Earlier complex language models generally underperformed Transformers and typically lacked Mamba-3’s data-dependent RoPE treatment.
  • Complex SSMs: Mamba-3 reintroduces complex-valued state transitions specifically to increase expressivity and state-tracking ability in a modern recurrent model.
  • MIMO: Mamba-3’s MIMO formulation targets increased modeling power while preserving inference efficiency, unlike prior MIMO work focused on simplifying training.
  • Design viewpoints: The paper presents exponential-trapezoidal discretization as SSM-specific, while noting MIMO can extend beyond SSMs and other frameworks remain convergent but distinct.

6 Conclusion And Future Work

The paper concludes that Mamba-3 combines improved discretization, complex state transitions, and MIMO to improve efficient sequence modeling. Its SISO and MIMO variants advance the performance-efficiency tradeoff with different training and inference characteristics.

  • Contributions: Mamba-3 combines exponential-trapezoidal discretization, complex-valued transitions, and MIMO to improve recurrence expressivity and inference efficiency.
  • Conclusions: The SISO model delivers strong language modeling results in standalone and interleaved hybrid architectures while advancing the Pareto frontier.
  • MIMO: MIMO trades slower training for stronger modeling power while maintaining competitive inference efficiency relative to Mamba-2.
  • Future directions: The methodology suggests state-space principles can provide theoretically motivated design directions for efficient sequence models.

A.2 Exponential-Trapezoidal Discretization Error Rate

This appendix analyzes the error rate of the exponential-trapezoidal discretization under smoothness, stability, and boundedness assumptions. It also reports the default gate parameterization used in experiments.

  • Assumptions: The analysis assumes bounded, C3 coefficients and inputs, Lipschitz state dynamics, and bounded λ_t to ensure zero-stability.
  • Error analysis: The state-input integral is approximated using Taylor expansion and the trapezoidal rule within the exponential-trapezoidal derivation.
  • Error rate: O(Δ_t^3) error follows when 1/2 − λ_t = O(Δ_t), under the stated parameterization condition.
  • Experimental setting: The error experiments use the 440M SISO model with shared architectural and optimization settings.
  • Gate parameterization: The default data-dependent gate λ_t = σ(u_t) performs best among tested parameterizations and is therefore selected.

B.1 Proof of Proposition 2

This section develops equivalences between complex-valued SSMs and real-valued systems with block-diagonal rotations, then describes the MIMO Mamba parameterization and output computation.

  • B.1 Proof of Proposition 2: Complex-valued SSM states can be represented by real block-diagonal systems with rotational transitions.The equivalence follows by pairing coordinates and expanding the complex exponential under discretization.
  • B.1 Proof of Proposition 2: The RoPE trick absorbs data-dependent rotary embeddings into the SSM input and output components through cumulative rotation products.This yields a vanilla scalar-transition SSM equivalent to the complex formulation.
  • B.1 Proof of Proposition 2: Exponential-trapezoidal discretization produces a real N-dimensional recurrence whose rotations arise from the complex phase term.The scalar decay commutes with the block-diagonal rotation matrix.
  • C MIMO for Mamba-3: MIMO Mamba obtains SSM parameters from projections of each position’s input using head dimension P and MIMO rank R.The projections define B, C, X′, and X while preserving a lightweight parameterization.
  • C MIMO for Mamba-3: The formulation transforms independent SISO SSMs within each head into MIMO SSMs, followed by gated activation and down projections.The residual gate uses Y_t ⊙ SiLU(Z_t), and the resulting layer output has dimension R^D.
  • C MIMO for Mamba-3: MIMO parameter matching reduces MLP hidden dimensions so the MIMO and SISO variants have comparable parameter counts.The supplied model-size row lists 180M, 440M, 880M, and 1.5B scales.

D Experimental Details

The experiments pretrain and evaluate Mamba-family models and baselines on language modeling, retrieval, synthetic retrieval, and state-tracking tasks using standardized procedures.

  • Language Modeling: All pretrained models follow shared bfloat16 procedures, with Mamba models using expand factor 2, state size 128, and head dimension 64.Transformer and Gated DeltaNet baselines follow their respective reference configurations.
  • Language Modeling: Zero-shot language modeling is evaluated on LAMBADA, HellaSwag, PIQA, Arc-Easy, Arc-Challenge, WinoGrande, and OpenBookQA.Evaluation uses the LM Evaluation Harness.
  • Real-World and Synthetic Retrieval: Real-world retrieval evaluation covers SWDE, SQuAD, FDA, TriviaQA, NQ, and DROP, while synthetic NIAH tasks are also evaluated.The retrieval tasks use cloze-formatted versions, with inputs truncated to the training context length.
  • Architecture Evaluation: The pretrained hybrid Mamba-3 SISO norm ablation finds no additional norm strongest for average in-context retrieval.Pre-gate grouped RMS performs best on synthetic retrieval, particularly beyond the training context length.
  • State-Tracking Synthetics: State-tracking models use a curriculum from maximum length 40 to 160 and are evaluated at length 256.Parity uses one layer, modular arithmetic uses three layers, and validation accuracy is selected across model dimensions and learning rates.

E Additional Experimental Results

Additional experiments show that Mamba-3 extrapolates better to longer contexts, improves pretraining performance, and remains on the quality–state-size Pareto frontier.

  • Context-Length Extrapolation: Mamba-3 exhibits strong length extrapolation on held-out FineWeb-Edu contexts, whereas Mamba-2 falters at longer lengths.The comparison uses pretrained 1.5B models evaluated across varying context lengths.
  • Pretraining Performance: Mamba-3 demonstrates better FineWeb-Edu validation perplexity than strong baselines including Mamba-2 and Gated DeltaNet.The comparison uses fully pretrained 1.5B models.
  • State-Size Trade-offs: State-size experiments compare inference speed and pretraining perplexity, with Mamba-3 and Mamba-3 MIMO continuing to set the Pareto frontier.Gated DeltaNet comparisons require caution because its multihead structure differs from Mamba’s multi-value structure.

F Architecture Ablations

Architecture ablations at the 440M scale examine normalization and B/C bias choices, showing that positive bias initialization is robust while B and C biases interact synergistically.

  • Experimental Setup: Ablation models are trained at 440M scale using Chinchilla-optimal tokens and the pretrained-model experimental procedures.This establishes a shared training setting for the architecture comparisons.
  • B, C Bias Parameterization: Mamba-3 uses head-specific, channel-wise B and C biases added after QK-Norm and initialized to all ones.The all-ones choice is adopted for simplicity.
  • B, C Bias Parameterization: Only B bias slightly hurts performance, while using both B and C biases produces synergistic benefits.This differs from the cited prior finding that B bias alone improved Mamba-1 performance.
  • B, C Bias Parameterization: Performance is not sensitive to bias initialization provided that the biases remain positive.The selected initialization gives the best performance while retaining simple parameterization.
  • B, C Bias Parameterization: Table 10 reports ablations of B/C bias initialization and bias presence for Mamba-3.The table separates initialization effects from whether each bias is included.

G.1 Kernel Implementations and Fusion Structure

The latency analysis documents the kernel DSLs and fusion structures used for Mamba-2, Gated DeltaNet, and Mamba-3 across forward and decode implementations.

  • Kernel implementations: Mamba-2 and Gated DeltaNet use their authors’ publicly released Triton kernels, while Mamba-3 uses newly implemented kernels with comparable fusion structure.Mamba-3’s forward SISO kernel is fused with rotary position embeddings; its forward MIMO kernel uses TileLang at the same fusion level.
  • Kernel implementations: Table 6 details the DSLs and fusion levels used in the latency analysis.The listed DSLs are Triton, TileLang, CuTe, and PyTorch.
  • Kernel implementations: Tables 11 and 12 abbreviate input projection, one-dimensional convolution, gating, and output projection, with colors marking implementation backends.The backends are Torch, Triton, TileLang, and CuTe.

G.2 Extended Prefill and Prefill+Decode Latency Measurements

The experiments compare recurrent models and a Transformer baseline across extended prefill and decode settings. Recurrent mixers show gentler context-length scaling, while Mamba-3 adds minimal forward-pass cost and maintains competitive recurrent-model decode latency.

  • Experimental setup: The benchmark includes Mamba-3 SISO, Mamba-2, Gated DeltaNet, and a vLLM Llama-3.2 1B Transformer baseline.The recurrent models are trained at the 1.5B scale.
  • Experimental setup: Prefill sequence lengths span L∈{512, 1024, 2048, 4096, 16384}, with an equal number of decoded tokens and batch size 128.At sequence length 16384, the vLLM result uses batch size 16 and is scaled by 8 because direct measurement at batch size 128 exceeds GPU memory.
  • Latency results: Mamba-3 adds minimal forward-pass cost while retaining lightweight exponential-trapezoidal, complex-state, and MIMO components.The observation attributes the low cost to the lightweight nature of these three design components.
  • Latency results: Recurrent mixers scale more gently with context length than vLLM Llama-3.2-1B, whose latency grows much faster because of KV-cache overhead.This pattern is reported across the swept sequence lengths.
  • Latency results: Decode latency is competitive across the recurrent models.The comparison includes Mamba-3, Mamba-2, and Gated DeltaNet at the 1.5B scale.
Loading 2603.15569v1…