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

cs.LGMar 16, 2026

Mamba-3: Improved Sequence Modeling using State Space Principles

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

Organizations: Carnegie Mellon University · Princeton University · Together AI · Cartesia AI

Abstract

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.

Figures & tables

Figure 1: Left: The structured mask induced by the exponential-trapezoidal rule ( Section 3.1 ) is a product of the decay and two-band convolutional mask. Right: Euler (hold endpoint) versus Trapezoidal (average endpoints) integral approximation.

Discretization Method	$\alpha_{t}$	$\beta_{t}$	$\gamma_{t}$	Appearance
Forward Euler	$I+\Delta A$	—	$\Delta$	—
Backward Euler	$(I-\Delta A)^{-1}$	—	$(I-\Delta A)^{-1}\,\Delta$	—
Trapezoidal	$(I-\frac{\Delta}{2}A)^{-1}\,(I+\frac{\Delta}{2}A)$	—	$(I-\frac{\Delta}{2}A)^{-1}\,\Delta$	S4
Zero-Order Hold	$\exp({\Delta A})$	—	$A^{-1}\,\left(\exp({\Delta A})-I\right)$	S4D, S5
Zero-Order Hold	$\exp({\Delta_{t}A_{t}})$	—	$A_{t}^{-1}\,\left(\exp({\Delta_{t}A_{t}})-I\right)$
Exponential-Euler	$\exp({\Delta_{t}A_{t}})$	—	$\Delta_{t}$	Mamba-1, -2 3 3 3 While the Mamba-1 paper reports ZOH discretization, the implementation follows https://github.com/state-spaces/mamba/issues/129 .

Table 1: Table of canonical linear-time invariant discretizations (top) and custom linear-time varying discretizations derived from our exponential-adjusted framework (bottom), along with their appearance in structured SSMs used in deep learning. Our theory formalizes the prior Mamba discretization as exponential-Euler and extends it with the more expressive exponential-trapezoidal method. The generalized discretization framework converts a continuous SSM

\dot{{\bm{h}}(t)}={\bm{A}}(t){\bm{h}}(t)+{\bm{B}}(t)x(t)

into the discrete recurrence

{\bm{h}}_{t}=\alpha_{t}{\bm{h}}_{t-1}+\beta_{t}{\bm{B}}_{t-1}x_{t-1}+\gamma_{t}{\bm{B}}_{t}x_{t}

, where various discretization methods yield different formulas for

\alpha_{t},\beta_{t},\gamma_{t}

Table 2: Arithmetic Intensity for (a) SISO, (b) MIMO. The batch and head dimensions cancel out. The arithmetic intensity of MIMO increases linearly with rank

R

, enabling better hardware utilization during memory-bound phases like decode. Here

N

is the state size (expansion factor) and

P

is the head dimension. For Mamba-3, typically

R\ll N,P

Figure 2: Contrasting Mamba-2 and Mamba-3 Architectures: Key updates include exponential-trapezoidal discretization, data-dependent RoPE embeddings, MIMO projections, QK normalization, and learnable biases.

Model	FW-Edu	LAMB.	LAMB.	HellaS.	PIQA	Arc-E	Arc-C	WinoGr.	OBQA	Average
	ppl $\downarrow$	ppl $\downarrow$	acc $\uparrow$	acc_n $\uparrow$	acc $\uparrow$	acc $\uparrow$	acc_n $\uparrow$	acc $\uparrow$	acc $\uparrow$	acc $\uparrow$
Transformer-180M	$16.89$	$45.0$	32.5	$39.0$	67.1	$59.8$	$27.9$ 27.9	$51.2$	$21.8$	$42.8$
GDN-180M	$16.52$ 16.52	40.8	$31.3$	$40.2$ 40.2	$66.3$	62.3	28.2	$51.7$	$22.0$	$43.2$ 43.2
Mamba-2-180M	$16.76$	$41.8$	$30.9$	$40.1$	$66.8$ 66.8	$60.1$	$27.3$	52.0	23.2	$42.9$
Mamba-3-SISO-180M	16.59	37.7	32.5	40.8	$66.1$	$61.5$ 61.5	$27.9$ 27.9	52.0	$22.8$ 22.8	43.4
Mamba-3-MIMO-180M	16.46	32.1	34.0	41.0	$66.7$	$60.6$ 60.6	$27.7$	52.9	$22.0$	43.5

Table 3: Downstream language modeling evaluations on models trained with 100B FineWeb-Edu tokens. Best results for each size are bolded , and second best are underlined , excluding Mamba-3 MIMO variants. All models are trained with the same procedure. Mamba-3 SISO outperforms Mamba-2 and others at every model scale, and the MIMO variant with rank

R=4

further improves modeling capabilities.

Model (1.5B)		SWDE	SQD.	FDA	TQA	NQ	Drop	NIAH-Single-1			NIAH-Single-2			NIAH-Single-3
Context Length		2048						1024	2048	4096	1024	2048	4096	1024	2048	4096
Pure	Transformer	$48.9$	$46.6$	$58.4$	$67.5$	$31.7$	$26.4$	$100.0$	$100.0$	$0.0$	$92.2$	$100.0$	$0.0$	$98.6$	$99.4$	$0.0$
	GDN	32.7	$40.0$ 40.0	28.3	$63.5$	$25.7$ 25.7	$24.5$	100.0	100.0	99.8	100.0	$93.8$ 93.8	$49.8$ 49.8	$83.8$	$68.4$	34.2
	Mamba-2	$30.7$ 30.7	$39.1$	$23.7$ 23.7	$64.3$ 64.3	$25.1$	28.5	100.0	$99.6$	$62.0$	100.0	$53.8$	$11.8$	95.8	87.4	$13.4$
	Mamba-3 SISO	$28.5$	40.1	$23.4$	64.5	26.5	$27.4$ 27.4	100.0	100.0	$88.2$ 88.2	100.0	95.4	50.6	$92.4$ 92.4	$81.4$ 81.4	34.2
	Mamba-3 MIMO	36.3	41.7	29.3	64.5	$26.2$ 26.2	$26.3$	100.0	100.0	$93.0$ 93.0	100.0	$86.0$	$40.4$	95.8	$84.4$ 84.4	$25.6$

Table 4: Retrieval capabilities measured by a mixture of real-world and synthetic retrieval tasks. Real-world retrieval tasks utilize cloze variants of the original datasets and are truncated to 2K length. Mamba-3 demonstrates strong associative recall, question-answering, and length generalization on needle-in-a-haystack (NIAH), but suffers with information extraction of semi-structured and unstructured data. The Transformer baseline uses RoPE which may explain its length generalization issues, and hybrid models utilize NoPE (no positional embeddings). We find a pre-gate, grouped RMSNorm can be added to Mamba-3 SISO hybrid models to improve the length generalization of the NIAH tasks at a slight decrease in real-world retrieval performance.

Table 5: Left : Ablations on core modeling components of Mamba-3 SISO, results on test split of dataset. Right : Formal language evaluation (scaled accuracy, %). Higher is better. SISO models are trained on short sequences and evaluated on longer lengths to test length generalization. For GDN we report the variant with eigenvalue range

[-1,1]

Model	512 tokens		1024 tokens		2048 tokens		4096 tokens		16384 tokens
	Prefill	Prefill+Dec	Prefill	Prefill+Dec	Prefill	Prefill+Dec	Prefill	Prefill+Dec	Prefill	Prefill+Dec
vLLM (Llama-3.2-1B)	0.26	4.45 4.45	0.52	9.60	1.08	20.37	2.08	58.64	12.17	976.50
Gated DeltaNet	0.51 0.51	4.56	1.01 1.01	9.11 9.11	2.01 2.01	18.22 18.22	4.00 4.00	36.41 36.41	16.21 16.21	145.87 145.87
Mamba-2	0.51 0.51	4.66	1.02	9.32	2.02	18.62	4.02	37.22	16.22	149.02
Mamba-3 (SISO)	0.51 0.51	4.39	1.01 1.01	8.78	2.02	17.57	4.01	35.11	16.22	140.61
Mamba-3 (MIMO $R=4$ )	0.60	4.74	1.21	9.48	2.42	18.96	4.76	37.85	19.44	151.81

Table 7: Prefill and Prefill+Decode latency across sequence lengths. Mamba-3 adds minimal overhead to its forward-pass and retains competitive decode latencies. Details in Appendix G .

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Parameterization	Form of $\lambda_{t}$	ppl $\downarrow$
Default	$\sigma(u_{t})$	15.72
Fixed $1/2$	$\tfrac{1}{2}$	15.76
No trapezoidal (Euler)	$1$	15.81

Appendix

Table 8: Ablations on

\lambda_{t}

parameterization in the exponential-trapezoidal update.

Model	$180\text{M}$	$440\text{M}$	$880\text{M}$	$1.5\text{B}$
SISO MLP dim	1,500	2,048	3,072	4,096
MIMO MLP dim ( $R=4$ )	1,264	1,792	2,800	3,824

Appendix

Table 11

Figure 4: Pretrained 1.5B models’ performance on the held-out FineWeb-Edu test set at varying context lengths. Mamba-3 exhibits strong length extrapolation while Mamba-2 falters at longer contexts.

Mamba-3 Norm Type	LM Avg.	SWDE	SQD.	FDA	TQA	NQ	Drop	NIAH-Single-1			NIAH-Single-2			NIAH-Single-3
Context Length	—	2048						1024	2048	4096	1024	2048	4096	1024	2048	4096
No Norm	$56.4$ 56.4	$58.5$ 58.5	$47.0$ 47.0	$65.9$ 65.9	$64.8$	33.4	$27.0$	100.0	100.0	$36.2$	100.0	100.0	$9.4$	99.8	100.0	$8.8$
Post-Gate Default RMS	56.5	$54.5$	$46.6$	$61.9$	$65.4$ 65.4	$31.9$	29.2	100.0	100.0	100.0	100.0	$99.8$	$49.2$	$87.6$	$94.0$	$62.0$ 62.0
Pre-Gate Default RMS	$55.9$	$55.4$	$46.9$	67.3	$65.4$ 65.4	$33.0$	$28.1$	100.0	100.0	$86.2$	100.0	100.0	97.8	$99.2$	$97.8$ 97.8	90.2
Post-Gate Grouped RMS	$56.2$	$51.4$	$46.7$	$56.8$	$64.2$	$30.4$	$27.6$	100.0	100.0	$79.4$	100.0	100.0	$65.8$	$93.8$	$97.0$	$9.6$
Pre-Gate Grouped RMS	$56.1$	58.6	47.3	$52.4$	65.7	$33.3$ 33.3	$28.5$ 28.5	100.0	100.0	100.0	100.0	100.0	$96.0$ 96.0	99.8	$97.2$	$56.8$

Appendix

Table 9: Ablations of optional norm type (grouped vs default) and placement (pre- vs post-gate) on pretrained hybrid Mamba-3 SISO models at the 1.5B scale. All models have BCNorm. No additional norm demonstrates the strongest in-context retrieval performance on average, while pre-gate, grouped RMS results in the best performance on synthetic retrieval, especially on lengths longer than its training context.

Figure 5: Mamba-3 demonstrates better pretraining performance compared to strong baselines like Mamba-2 and Gated DeltaNet. These are the validation perplexity on FineWeb-Edu of our fully pretrained 1.5B models.

Figure 6: Exploration of state size (inference speed proxy) versus pretraining perplexity (performance proxy). Mamba-3 and Mamba-3 MIMO continue to set the Pareto frontier.

Table 10: Ablations on

B,C

bias initialization (left) and presence (right) for Mamba-3.

Model (Forward)	Kernel DSL	Fusion Level
Mamba-2	Triton	IP , Conv , SSM , Gate , OP
Gated DeltaNet	Triton	IP , Conv , Chunked Delta , Gate , OP
Mamba-3 (SISO)	Triton	IP , SSM+Rotary+Gate , OP
Mamba-3 (MIMO)	TileLang	IP , SSM+Rotary+Gate , OP

Appendix

Table 11: Kernel DSL and fusion structure for forward (prefill) kernels.

Model (Decode)	Kernel DSL	Fusion Level
Mamba-2	Triton	IP , Conv , SSM , Gate , OP
Gated DeltaNet	Triton	IP , Conv , Recurrent Delta , Gate , OP
Mamba-3 (SISO)	CuTe + Triton	IP , Rotary , SSM+Gate , OP
Mamba-3 (MIMO)	CuTe + Triton	IP , Rotary , SSM+Gate , OP

Appendix

Table 12: Kernel DSL and fusion structure for decode kernels.

Sequence Modeling

Sequence modeling studies how to represent, predict, classify, and generate ordered data such as text, time series, events, and biological sequences.

See all 246 papers

From Markov to Laplace: How Mamba In-Context Learns Markov Chains
Feb 14, 2025Marco Bondaschi, Nived Rajaraman, Xiuying Wei +5Mamba Selective SSMs
EPFL · UC Berkeley · Google DeepMind +1
TurboMamba: Faster Inference, Faster Training, and Better Prediction with Lexical Memory
Oct 14, 2026Hanxin ZhangHybrid Transformer-State Space Model Architectures Autoregressive Language Modeling
Department of Neurology, University of Pittsburgh
Ternary Mamba: Grouped Quantization-Aware Training of W1.58A16 State Space Models
Jun 16, 2026Ramprasath Ganesaraja, Sahil Dilip Panse, Swathika NLLM Quantization Ternary Quantization
EdgeVerve Systems Limited

Feb 14, 2025J/K move · Enter open · S save

From Markov to Laplace: How Mamba In-Context Learns Markov Chains

Marco Bondaschi, Nived Rajaraman, Xiuying Wei, Kannan Ramchandran+4

EPFL · UC Berkeley · Google DeepMind+1

Mamba Selective SSMs In-Context Learning

Open paper AskPDF arXiv

While transformer-based language models have driven the AI revolution thus far, their computational complexity has spurred growing interest in viable alternatives, such as structured state space sequence models (SSMs) and Selective SSMs. Among these, Mamba (S6) and its variant Mamba-2 have shown remarkable inference speed-ups over transformers while achieving comparable or superior performance on complex language modeling tasks. However, despite these architectural innovations and empirical successes, the fundamental learning capabilities of Mamba remain poorly understood. In this paper, we address this gap by studying in-context learning (ICL) on Markov chains and uncovering an interesting phenomenon: even a single-layer Mamba efficiently learns the in-context Laplacian smoothing estimator, which is both Bayes and minimax optimal. To explain this, we theoretically characterize the representation capacity of Mamba and reveal the fundamental role of convolution in enabling it to represent the optimal Laplacian smoothing. These theoretical insights align strongly with empirical results and, to the best of our knowledge, represent the first formal connection between Mamba and optimal statistical estimators. Finally, we outline promising research directions inspired by these findings.

Abstract

Sequence Modeling

Explore similar work

From Markov to Laplace: How Mamba In-Context Learns Markov Chains

TurboMamba: Faster Inference, Faster Training, and Better Prediction with Lexical Memory

Ternary Mamba: Grouped Quantization-Aware Training of W1.58A16 State Space Models