Skip to main content

Why Does Mamba-2 Resemble the Transformer? Reading SSMs Through SSD

· 7 min read
p4r4d0xb0x
p4r4d0xb0x
Rustacean, AI, OSS Enthusiast

Conceptual diagram connecting state space models and attention through a structured matrix

The bottleneck in models that handle long sequences is not simply the number of parameters. Training must process every position in a long sequence in parallel, while generation must read past information for every new token. The Transformer gains powerful global attention, but accepts computation and memory costs that grow with sequence length. State space models (SSMs) in the Mamba family maintain a recurrent state to keep the state size fixed during generation, yet their training and GPU utilization are harder to describe as straightforwardly as the Transformer's.

Tri Dao and Albert Gu's paper does not set these two families up as competing boxes. Its central question is whether SSMs and attention can be expressed in the same structural language. The paper uses structured semiseparable matrices to organize the relationship between the two families, then proposes the core Mamba-2 layer and the SSD algorithm from that perspective. The discussion below is based on the abstract, methods, and experiments in the arXiv paper and HTML; it does not add results from running reproduction code.

The Problem: Linear Complexity Alone Is Not Enough​

A typical SSM can be written as a state update:

h_t = A_t h_{t-1} + B_t x_t
 y_t = C_t^T h_t

The state h is working memory that compresses past inputs. During generation, the model can pass only this state to the next step instead of revisiting every token. When A, B, and C depend on the input, as in Mamba's selective SSM, the model can control token by token which information is retained or erased. However, updating one token at a time can fail to fully utilize a GPU's matrix-multiplication hardware. Therefore, “linear in sequence length” and “fast on real hardware” must be treated as separate claims.

The paper's perspective is that the same transformation can be viewed in three ways. The recurrent form shows sequential state updates; the quadratic form exposes interactions between all positions as a matrix; and the matrix form structurally connects the two. This connection places SSMs in the same problem space as attention computations, which already have extensive optimization, rather than treating them merely as a new kind of RNN.

Term explanation: State Space Model (SSM)

Simple definition: A sequence model that updates an internal state whenever it receives an input and produces an output from that state.

Example: Imagine keeping a one-line summary of a conversation in a notebook and consulting only that summary when reading the next sentence. Unlike an ordinary note, the model can learn the rules for summarizing and how much to retain.

The Core of SSD: A Matrix Structure Connects Two Computation Paths​

The paper shows that an SSM can be represented by a particular semiseparable matrix. This matrix does not store every relationship between positions without structure. It uses structures such as diagonal blocks and low-rank blocks to reduce the parameters and multiplications required. As a result, the same operation can be computed either as a sequential recurrence or as block matrix multiplications.

SSD (State Space Duality) names this duality framework. The important point is not the simplistic claim that “attention is exactly the same as an SSM.” The paper explicitly addresses the relationship between particular SSMs represented by structured matrices and kernel or masking variants of attention—not every form of attention. The intuition in the title is powerful, but its scope is bounded by mathematical conditions.

The execution algorithm divides a sequence into chunks, uses parallel matrix operations within each block, and passes a summarized recurrent state between blocks. This preserves the meaning of the sequential state computation while bringing in the large matrix multiplications that GPUs handle well. The hardware efficiency discussed by the paper comes precisely from being able to choose between these computation paths.

Term explanation: Semiseparable matrix

Simple definition: A structured matrix that may look dense overall, but compresses many blocks of relationships between positions into small low-rank representations.

Example: Instead of storing a road map for every pair of intersections in a city, store summaries for each district and rules for connections between districts. Necessary route calculations remain possible while storage and computation are reduced.

Term explanation: Recurrent state

Simple definition: A fixed-size internal value passed to the next input instead of the entire past.

Example: Bring a single page of meeting minutes to the next meeting instead of the full transcript. A shorter record is convenient, but the risk of losing important details must also be managed.

Mamba-2 Changes More Than the Algorithm​

Alongside the SSD layer, the paper adjusts the design of the Mamba block. It computes data-dependent projections in parallel at the beginning of the block and brings the head structure familiar from multi-head attention into the SSM design space. It also describes a design that reduces synchronization points to make tensor parallelism easier. For sequence parallelism, which divides long sequences across devices, it presents a way to pass the recurrent state between devices.

The practical meaning is that reducing the FLOPs of one new layer is not enough; the entire training system must fit together. Kernels, communication, variable-length batches, and memory layout must align before real latency falls. The paper also discusses avoiding padding for variable-length sequences, but does not guarantee the same effect for every hardware and framework combination.

Paper Structure and Scope of the Evidence​

The paper introduces the problems of Transformers and SSMs in the Introduction, then reviews SSMs, attention, and structured matrices in the Background. It represents SSMs as matrices, develops the theoretical relationship between structured masked attention and SSD, and then moves to efficient algorithms, the Mamba-2 design, system optimizations, and experiments. The progression is problem → common representation → algorithm → architecture → validation.

The abstract and body report that the core SSD layer is 2–8 times faster than Mamba's selective scan and achieves language-modeling performance competitive with the Transformer. The body includes a crossover with FlashAttention-2 under specific conditions, long-sequence speed comparisons, Pile training and downstream evaluation, scaling laws, and ablations. These numbers are results under the implementation, hardware, batch, and length conditions set by the paper. This article does not run an independent benchmark, so it does not generalize them to “2–8 times faster in every environment.”

What Changes in Practice​

  • Inference servers: Maintaining a fixed-size state instead of a full KV cache may simplify memory planning for services with long inputs or many concurrent users. It does not mean that the state preserves every detail of a long context.
  • Training stack: SSD's chunkwise and block structure creates opportunities to use GEMM and tensor parallelism. Before adoption, profile kernel implementation and communication costs together.
  • Model selection: It is more accurate to view Mamba-2 not as a one-line replacement for the Transformer, but as a layer offering a trade-off between recurrent inference and parallel training.
  • Evaluation: In addition to perplexity and downstream scores, measure long-context retrieval, state size, latency and throughput by batch size, and peak memory under identical conditions.

Limitations and Open Questions​

The paper presents a deep connection between SSMs and attention, but that connection does not mean an equivalent replacement for all softmax attention. Reported speedups also result from particular optimized implementations, so they must be checked again on new GPUs, with different precisions, small batches, and real serving workloads. Because a recurrent state compresses context, each application must also test what information-loss patterns arise in complex arbitrary retrieval and long-range interactions.

Within the scope supported by the provided materials, Mamba-2's most important contribution is not the slogan “linear complexity.” It is the ability to view the same sequence transformation both as a recurrence and as matrix multiplication, then connect that choice to the language of hardware and distributed training. Quantitative performance requires further verification, but the research direction is clear: a new architecture must explain both the elegance of its equations and the realities of kernels, memory, and communication.

Sources​

// COMMENTS

Comments