Mamba

Overview

Mamba is a sequence model built on selective state-space layers. It reads a sequence one step at a time, keeps a compact hidden state, and lets the input decide what to keep and what to forget. Because the state has a fixed size, the cost of processing a sequence grows in a straight line with its length, not with the square of its length as in standard attention.

The model was introduced in late 2023 as an answer to a clear weakness of earlier state-space models: they were fast but could not choose what to remember. Mamba fixes this by making the key parameters depend on the current input. A hardware-aware algorithm keeps the method fast on a GPU. The result is a model that competes with Transformers on language and other long-sequence tasks while using far less memory at inference time. It builds on S4 and shares ideas with Neural ODE models.

What it is

A model with a memory that can choose

Imagine reading a long book while keeping a single page of notes. You cannot write down everything, so for each new sentence you decide what to add, what to erase and what to ignore. Mamba does this with a hidden state. The state is the page of notes. The input at each step controls how much is written, how much is erased and what is read out.

Earlier state-space models used the same fixed rule for every step. They were like a note-taker who copies every tenth word no matter what the text says. Mamba is a note-taker who pays attention to the content.

The block

A Mamba block has a small number of parts:

• A linear projection that widens the input.
• A short causal convolution that mixes nearby positions.
• A selective state-space layer that carries the long-range memory.
• A gating branch that multiplies the output.
• A linear projection back to the model width.

Blocks are stacked with normalisation and residual connections. There is no attention layer and no separate feed-forward block. The gating and widening in the block play the role of the feed-forward part.

What "selective" means

In a classic linear time-invariant state-space model, the matrices that control the state are the same at every step. In Mamba, three quantities are computed from the input at each step:

• The step size, which sets how far the state moves.
• The input matrix, which says how the new input is written into the state.
• The output matrix, which says how the state is read.

Because these change with the input, the model can ignore filler tokens and hold on to important ones. This is called selectivity.

What it is not

• It is not a recurrent network in the older sense. It can be computed with a parallel scan during training.
• It is not just attention made cheaper. It stores a fixed-size summary, while attention keeps every past item available.
• It is not tied to language. It has been applied to audio, genomics, images, time series and other data.

Where it came from

Early state-space ideas

State-space models come from control theory, where a hidden state evolves under linear rules and an output is read from it. In deep learning, the line of work began with a method for remembering long histories through polynomial projections (the HiPPO framework, 2020). That led to the structured state-space sequence model, S4, published in 2021 and presented in 2022.

From S4 to Mamba

S4 was strong on long-range benchmarks but weak on language, where choosing what to remember matters. Follow-up models simplified the structure, added gating and tried to close the gap with attention. Mamba, from late 2023, added selection and a new algorithm to compute it fast. The paper is titled "Mamba: Linear-Time Sequence Modeling with Selective State Spaces".

Mamba-2 and what followed

A 2024 follow-up, often called Mamba-2, showed a close mathematical link between state-space models and a form of attention. This link is known as structured state-space duality. It let the authors design a faster layer that uses matrix multiplication more heavily, which suits modern accelerators. Hybrid designs that mix state-space layers with a few attention layers also appeared and showed good results.

A growing family

Since then, many variants have been proposed for vision, speech, graphs, recommendation and tabular data. They share the same core idea: a selective recurrence that is computed with a parallel scan. Related work on test-time memory, such as Miras, explores other ways to write and read a learned memory.

How it works

Step 1: a continuous view

Start with the continuous form of a state-space model:

• The state changes by a linear rule, dh/dt = A h + B x.
• The output is y = C h.

Here x is the input, h is the hidden state, and A, B and C are matrices. This is the same kind of system a Neural ODE solves, but with a linear rule.

Step 2: discretise

Data come in steps, so the continuous system must be converted. A step size, usually called delta, turns A and B into discrete versions. In simple terms, a large delta means "move a lot in this step" and a small delta means "barely move".

After discretisation the update is a plain recurrence: the new state equals a decay factor times the old state plus a write term times the input.

Step 3: make the parameters depend on the input

In Mamba, delta, B and C come from small linear layers that read the current input. This is the selection mechanism.

• A large delta makes the state reset toward the current input. Past content fades.
• A small delta makes the state ignore the current input. Past content is kept.

Think of a door that opens wide for important guests and stays shut for noise.

Step 4: compute it fast

If the parameters depend on the input, the model can no longer be written as one big convolution, which is the trick that made S4 fast. Mamba instead uses a scan: a method that combines the steps in a tree so many can be done in parallel.

The authors also designed the algorithm around how GPUs work. The big expanded state is kept in the fast on-chip memory and is never written in full to slow memory. This is the same spirit as FlashAttention. The speed comes from careful memory movement, not from fewer arithmetic operations.

Step 5: training versus inference

During inference, Mamba does not need a growing cache of past keys and values. It only needs the fixed-size state. This is the main practical gain for long outputs.

Step 6: stack and predict

Blocks are stacked. The last layer feeds a head. For language, this is a softmax over tokens. For classification, the states are pooled and a linear layer predicts a label. For regression, a linear layer predicts a number.

A simple analogy for the whole pipeline

A river carries leaves downstream. The riverbed is the fixed structure of the model. The speed of the water at each point is set by what is happening at that point. Leaves that matter get caught in an eddy and stay. Leaves that do not matter float away. At the end of the river, you look at what is in the pool.

Key concepts and terms

Time-invariant versus selective

Variants and extensions

Bidirectional and multi-directional scans

A causal scan only sees the past. For data with no natural order, such as the cells of a row or the patches of an image, designers run one scan forward and one backward and merge the results. Some vision variants scan in four directions.

Mamba-2 and the dual view

The dual view writes the same layer as a masked matrix product, which looks like linear attention with a decaying mask. This lets one implementation switch between a recurrent form for inference and a matrix form for training.

Hybrid stacks

A hybrid stack uses mostly state-space layers and inserts an attention layer every few blocks. The attention layers give exact lookup. The state-space layers keep cost low. Reported results suggest that even a small share of attention closes much of the recall gap.

Mixture-of-experts combinations

State-space layers can be paired with sparse expert feed-forward layers, as in Mixture of experts. This raises capacity without raising the cost per token much.

Domain adaptations

• Audio and speech: long waveforms are a natural fit.
• Genomics: sequences of hundreds of thousands of symbols.
• Vision: patches scanned in several orders.
• Time series: long histories at low cost.
• Tables: cells or rows scanned as sequences.

Deployment notes

• Inference state size is fixed, so serving cost per user is predictable.
• Batching is simple, since every sequence carries a same-sized state.
• Long prompts can be processed in chunks, passing the state from one chunk to the next.
• Quantising the projections works well. Be careful with the recurrence itself, which is sensitive to rounding.
• Keep a regression test that compares the fused kernel against the slow reference, so a driver or library update does not change results silently.

A worked example

A scalar numeric walk-through

Take a one-dimensional state with a simple rule: new state = decay * old state + write * input. In Mamba the decay is exp(-delta * a) for some positive a, and the write term scales with delta.

Let a = 1. Suppose the input sequence is 1, 0, 0, 5, 0 and delta is chosen per step by the model.

When the input is zero and delta is small, the state barely decays, so the memory of step 1 and later of step 4 survives. When delta is large, the state absorbs the new input. A fixed delta could not do both. That is the value of selection.

A minimal selective scan in PyTorch

This is a slow reference version. It shows the idea and is not a replacement for the fast kernel.

import torch
import torch.nn as nn
import torch.nn.functional as F

class TinySelectiveSSM(nn.Module):
    def __init__(self, d_model=16, d_state=8):
        super().__init__()
        self.d_model, self.d_state = d_model, d_state
        self.A_log = nn.Parameter(torch.log(torch.arange(1, d_state + 1).float()).repeat(d_model, 1))
        self.to_delta = nn.Linear(d_model, d_model)
        self.to_B = nn.Linear(d_model, d_state)
        self.to_C = nn.Linear(d_model, d_state)

    def forward(self, x):                      # x: (batch, length, d_model)
        b, L, d = x.shape
        A = -torch.exp(self.A_log)             # (d, n), negative for stability
        delta = F.softplus(self.to_delta(x))   # (b, L, d)
        B = self.to_B(x)                       # (b, L, n)
        C = self.to_C(x)                       # (b, L, n)
        h = torch.zeros(b, d, self.d_state)
        ys = []
        for t in range(L):
            dA = torch.exp(delta[:, t].unsqueeze(-1) * A)          # (b, d, n)
            dB = delta[:, t].unsqueeze(-1) * B[:, t].unsqueeze(1)  # (b, d, n)
            h = dA * h + dB * x[:, t].unsqueeze(-1)
            ys.append((h * C[:, t].unsqueeze(1)).sum(-1))          # (b, d)
        return torch.stack(ys, dim=1)

layer = TinySelectiveSSM()
out = layer(torch.randn(2, 32, 16))
print(out.shape)   # torch.Size([2, 32, 16])

Using a library implementation

Open-source packages provide the fast fused kernel. The pattern is short.

import torch
from mamba_ssm import Mamba            # package name used by the public reference code

block = Mamba(d_model=256, d_state=16, d_conv=4, expand=2).to("cuda")
x = torch.randn(4, 2048, 256, device="cuda")
y = block(x)                           # same shape as x

What to try

• Feed a long synthetic task, such as copying a token after a gap of thousands of steps.
• Compare memory use at length 2k, 8k and 32k against an attention block.
• Compare speed of the slow loop above with the fused kernel to see why the kernel matters.

A second example: memory of a flag

Suppose a table column holds a daily status, and a rare flag on day 1 changes the meaning of every later row. A model must carry that flag across thousands of rows.

• With selection, the step size is tiny on ordinary rows, so the flag stays in the state.
• On the flag row, the step size is large, so the flag is written strongly.
• On a later reset row, the step size is large again and the old flag is replaced.

Read the table as a stream and the state is a running answer to the question "which regime am I in".

Strengths and limitations

Strengths

• Linear time and memory in sequence length.
• Constant memory per generated token, since no growing cache is needed.
• Selection lets the model filter noise and keep what matters.
• Very good throughput on long sequences.
• Simple block design, with no attention and no separate feed-forward layer.
• Works across modalities.
• Hybrid designs can combine its efficiency with attention's exact recall.

Limitations

• A fixed-size state is a lossy summary. Tasks that need exact lookup of an earlier item can be harder than for attention.
• In-context copying and retrieval tasks often show a gap against attention at the same size.
• The fast kernel is specialised. Support on older or non-standard hardware is weaker.
• Tooling and tuning knowledge are less mature than for Transformers.
• Long-range gains depend on data. On short sequences, attention is already cheap and strong.
• Interpretability tools built for attention maps do not apply directly.
• Sensitive to initialisation of the decay parameters and to numerical precision.

The retrieval gap in plain words

Attention keeps every item on a shelf and can pull out any of them. Mamba keeps a summary on a single page. For tasks like "find the phone number mentioned 10,000 words ago" the shelf is better. For tasks like "follow the gist of a long story" the page is often enough, and it is much cheaper.

How it compares with alternatives

Mamba versus attention

Attention compares every pair of positions. This gives strong recall but costs time and memory that grow with the square of length. Mamba compresses history into a state. It scales better but can lose detail. Many recent systems use a mix: mostly state-space layers with a small number of attention layers for exact lookup.

Mamba versus recurrent networks

Older recurrent networks such as LSTM and GRU are also fixed-state models. They were slow to train because each step depends on the last. Mamba keeps the recurrence but removes the nonlinearity between steps so that a scan can parallelise training. It also uses a much larger and better-structured state.

Mamba versus S4

S4 has fixed dynamics, which allow a fast convolution, but it cannot select content. Mamba trades the convolution trick for selection and uses a hardware-aware scan to stay fast.

Role in tabular foundation models

Why a sequence model matters for tables

A tabular foundation model learns from raw tables: rows and columns, numbers and categories mixed together, missing cells and often no column names. The core design question is how to turn a two-dimensional grid into something a neural network can read. The usual answer is a Transformer with attention across rows and across columns. See TabPFN, TabICL and FT-Transformer. Mamba is an alternative engine for the same job, and the trade is about cost and about long contexts.

Row and column structure

A table can be read in two directions.

• Across the columns of one row, which is short, since most tables have tens or hundreds of columns.
• Down the rows of one table, which can be very long, with thousands to millions of rows.

Attention across columns is cheap because the count is small. Attention across rows is where cost explodes. A selective state-space layer is a good fit for the row direction. It can scan a long list of rows in linear time and carry a summary of what has been seen. In an in-context learning setup, the labelled example rows form the context, and the model must use them to predict a new row.

Order is a problem to manage

Rows in a table have no natural order, and neither do columns. A recurrence reads in a fixed order. This means designers must handle order with care:

• Shuffle rows during pre-training so the model does not learn order artifacts.
• Average predictions over several random orders at test time.
• Use column embeddings that do not depend on position, or feed columns in several permutations.
• Use bidirectional scans, which read forward and backward and merge the results.

Mixed types and missing values

Each cell is first turned into a vector. Numbers go through a small numeric encoder, categories go through an embedding or hashing step, and missing cells get a learned mask vector plus an indicator. After that, a Mamba layer sees only vectors, so mixed types are not its concern. Selection helps here: the model can learn to give little weight to cells flagged as missing and to carry forward the informative ones.

No column names

A foundation model for raw tables often must work when column names are absent or meaningless. In that case, the model must infer a column's role from its values and from its relations to other columns. A sequence model over column vectors can learn this from the data statistics alone. The state acts as a running summary of "what kind of table am I looking at".

Pre-training on synthetic and real tables

Pre-training uses large numbers of generated tables, drawn from priors such as random causal graphs, and often real tables as well. A state-space backbone can be trained on the same material. The main gain is the ability to include far more example rows in context for the same memory. See Synthetic data generation for how such tables are made.

In-context learning

In-context learning is the central trick of tabular foundation models: show labelled rows, then ask for a prediction on new rows, with no weight updates. Attention-based models do this very well. Fixed-state models have a known weakness at exact recall, but tabular prediction usually needs statistical summaries of the context, not exact lookup. This makes state-space layers a reasonable candidate for large contexts. Hybrid layouts, with a few attention layers on top of a state-space stack, are a sensible compromise.

Continual learning

Because inference needs only a fixed-size state, new rows can be absorbed by updating that state, without recomputing a growing cache. This suits streaming tables, where data arrive over time. It does not by itself stop forgetting in the weights, which still needs replay or regularisation. See Continued pre-training.

Imputation and anomaly detection

For imputation, the model reads the observed cells and predicts masked ones, in the style of Masked cell modeling. A selective scan can read a long row of cells or a long list of neighbouring rows to inform each guess. For anomaly detection, a model that predicts the next row or cell gives a natural score: large prediction error means a surprising row. See Anomaly detection and Imputation networks.

Sector and semantic classification

Deciding what domain a table belongs to, or what a column means, requires reading many cell values and summarising them. A state-space encoder pools a long column into a single vector in linear time, which a classification head then reads. See Semantic sector heads. For a view of how this fits with full model designs and measured results, see /model/architecture, /model/research and /model/benchmarks.

An honest summary of the role

Mamba is not the standard backbone of current tabular foundation models. Attention remains the default because tables are often small and exact set-like reasoning over rows is valuable. Mamba becomes attractive when the context is very long, when memory at inference is tight or when data stream in. It is best seen as a tool for scaling the row dimension.

Practical guidance and best practices

When to choose Mamba

• Sequences are long, from several thousand to millions of steps.
• Inference memory or latency per token is a main constraint.
• The task needs a gist more than exact recall.
• You can use a hybrid design if recall matters.

When to stay with attention

• Sequences are short or moderate.
• The task needs precise lookup or copying.
• You need mature tooling and well-known recipes.

Training tips

• Use mixed precision, but keep the state recurrence in higher precision where the kernel allows.
• Keep the decay parameter initialisation as in the reference design.
• Use a learning rate similar to or a bit lower than for a Transformer of the same size.
• Apply weight decay to projections but not to the decay and step-size parameters.
• Use gradient clipping.

Handling tabular inputs

• Embed each cell to a vector first and add column identity features.
• Shuffle row order every batch during pre-training.
• Use both scan directions if order should not matter.
• Provide a mask channel for missing cells.

Evaluating

• Test at lengths longer than seen in training to check length generalisation.
• Include a retrieval-style test to see the recall gap.
• Report memory and throughput along with accuracy.

A short decision table

Quick sanity checks before a long run

• Overfit a tiny batch to confirm the loss can reach near zero.

• Print the range of the step-size values. If they collapse to one value, selection is not working.

• Check that outputs match between the recurrent and the parallel mode on a short input.

• Confirm that shuffling rows changes predictions only slightly for a table task.

• Log memory at several lengths to confirm linear growth.

Common mistakes

• Expecting attention-like exact recall from a fixed-size state.
• Using the slow reference loop in production instead of the fused kernel.
• Feeding table rows in a fixed order and learning order artifacts.
• Ignoring numerical precision in the recurrence.
• Treating missing cells as zeros without a mask.
• Comparing against a weak attention baseline, or against one with no efficient kernel.
• Testing only at the training length and assuming it holds at longer lengths.
• Applying weight decay to the state-decay parameters.
• Assuming a state-space layer needs no normalisation or residual path.
• Choosing it for short inputs where its advantage does not appear.

Frequently asked questions

What is Mamba in simple terms?

Mamba is a sequence model that reads data one step at a time and keeps a small, fixed-size memory. What makes it special is that the input decides how much to remember and how much to forget at each step. This makes it fast and light on memory for very long sequences.

How is Mamba different from a Transformer?

A Transformer compares every position with every other position, which gives strong recall but costs time and memory that grow with the square of the length. Mamba compresses history into a fixed state and costs time that grows in a straight line with length. The trade is some loss of exact recall in return for speed and low memory.

What does selective mean in a selective state-space model?

Selective means that the parameters controlling the memory, such as the step size and the read and write maps, are computed from the current input. The model can therefore ignore unhelpful content and hold on to useful content. Older state-space models used the same fixed parameters at every step.

Why is Mamba fast if it is a recurrence?

Training uses a parallel scan that combines many steps at once, and the kernel is written to keep the large internal state in fast on-chip memory. Inference runs as a simple recurrence with constant work per step. Both modes avoid the quadratic cost of attention.

Can Mamba replace attention completely?

For many long-sequence tasks it can match or beat attention at lower cost. For tasks that need exact recall of earlier items, attention still has an edge, which is why hybrid models that combine both are popular. The best choice depends on the task and on the length of the input.

What is Mamba-2?

Mamba-2 is a later version that reveals a mathematical link between state-space layers and a form of attention. That link allowed a redesigned layer that makes heavier use of matrix multiplication and runs faster on modern accelerators. It keeps the same overall idea of a selective, fixed-size state.

Can Mamba be used on tabular data?

Yes, in two main ways. It can read the cells of a row as a short sequence, or it can scan a long list of rows as context for in-context learning. Attention is still the common backbone for tabular foundation models, and a state-space layer is mainly useful when the row context is very long.

Does Mamba handle missing values?

The model itself sees only vectors, so missing values must be handled before it, with a mask and a learned placeholder for missing cells. Selection then lets the model learn to give little weight to masked cells. It is the encoding step that does the main work.

How does Mamba relate to S4?

S4 is an earlier structured state-space model with fixed dynamics that can be computed as a convolution. Mamba keeps the state-space idea but makes the dynamics depend on the input, which gives selectivity and better language results, and it replaces the convolution trick with a hardware-aware scan.

Do I need a GPU to use Mamba?

Training at any useful scale needs a GPU, and the fast kernels target recent NVIDIA hardware. Small reference versions in plain PyTorch run on a CPU for learning and testing, but they are slow because they loop over steps. Inference on a CPU is possible for small models since it only needs the fixed state.

Glossary

• Causal convolution: a short filter that only looks at current and past positions.
• Decay: the factor by which old state is scaled at each step.
• Delta: the learned step size that controls how strongly the state moves.
• Gating: multiplying by a learned signal to pass or block information.
• Hardware-aware algorithm: a method designed around memory movement on a GPU.
• Hidden state: the fixed-size summary carried from step to step.
• HiPPO: a framework for remembering long histories with polynomial projections.
• Hybrid model: a stack mixing state-space and attention layers.
• Linear time: cost grows in proportion to length.
• Parallel scan: a way to compute a recurrence using a tree of combines.
• Selection: making model parameters depend on the input.
• State expansion: using a hidden state larger than the number of channels.
• State-space model: a model built on a linear hidden-state recurrence.
• Structured duality: the link between state-space layers and attention.
• Throughput: tokens processed per second.

Summary

Mamba is a selective state-space model. It keeps a fixed-size hidden state, lets the input decide what to write and what to forget, and uses a hardware-aware scan to stay fast. The outcome is linear-time training, constant-memory inference and strong results on long sequences, at the price of weaker exact recall than attention.

For tabular foundation models, Mamba is an option for scaling the row dimension, handling streaming rows and keeping inference memory low, while attention remains the common default for short and medium contexts. Hybrid designs are a practical middle path. Related reading: S4, Neural ODE, Transformers, /model/architecture and /model/benchmarks.

Related

It belongs to Architectures in my stack, under AI and machine learning. Others in the same group: