Mamba: The State Space Model That Challenges the Transformer
Mamba introduces selective state space models with linear scaling in sequence length and constant-time inference, offering a genuine alternative to the Transformer for long-sequence tasks.
Standard attention scales as O(N²) in both memory and compute with sequence length N. A 1M-token context with standard attention requires roughly 1 trillion attention operations per layer. Even with FlashAttention's IO efficiency, this remains compute-expensive and limits practical context lengths. Mamba (arXiv:2312.00752) proposes a fundamentally different architecture.
State Space Models: The Foundation
SSMs model sequences as continuous dynamical systems. The core recurrence is:
h'(t) = Ah(t) + Bx(t)
y(t) = Ch(t)
where h is a hidden state, x is input, y is output, and A/B/C are learned matrices. Discretized to sequences, this becomes a linear recurrence that can be computed as a convolution during training (parallel, like attention) but executed as a recurrence during inference (constant compute per step, unlike attention which must attend to all previous tokens).
Team workspace
Ship faster with chat, meetings, and projects in one place — Zlyqor.
Prior SSMs like S4 used time-invariant parameters - A, B, C were fixed regardless of the input token. This made them less effective at content-aware selection (remembering the right things and forgetting others). Mamba makes B and C input-dependent: the model learns to selectively update its hidden state based on the current input. This selectivity is what enables Mamba to match Transformer quality on language tasks.
Hardware-Parallel Scan Algorithm
Making B and C input-dependent breaks the convolution computation path. Mamba uses a parallel scan algorithm (also known as prefix sum) to compute the recurrence in O(N log N) operations in a way that is highly parallelizable on modern GPUs. The scan is implemented with custom CUDA kernels similar in spirit to FlashAttention's tiling approach.
Mamba-2 (arXiv:2405.21060) showed that the selective SSM is a special case of a broader class called structured state space duality (SSD), which connects SSMs and linear attention. Mamba-2 uses a simplified block structure (scalar instead of matrix selective parameters) that runs 2-8x faster than Mamba-1 on modern hardware.
Where Mamba Wins vs Transformers
Mamba excels at: very long sequences (>100k tokens), streaming inference (constant memory per step), tasks requiring selective state tracking (DNA sequences, audio, time series). Transformers remain superior for tasks requiring precise in-context learning and retrieval from a fixed context window.
Mamba in Production
Jamba (AI21 Labs) is a hybrid model combining Transformer layers and Mamba layers, using attention for local context and SSM for long-range dependencies. This hybrid approach captures the strengths of both architectures.
Practical deep-dives on LLMs, developer tools, and AI engineering. No filler. Unsubscribe any time.
// written byFIG. AUTH-01
530
Mahmudul Haque Qudrati
CEO & ML Engineer
CEO and ML Engineer at Pristren. Builds AI-powered software for teams and writes about machine learning, LLMs, developer tools, and practical AI applications.
Toolformer: Teaching LLMs to Use Tools Without Human Annotations
Toolformer learns to call external APIs - calculators, search engines, calendars - by self-supervising on when API calls improve prediction, requiring no human-labeled examples of tool use.
