Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality

What problem it addresses

Models built on attention carry two efficiency costs that grow with sequence length. During training, attention scales quadratically in sequence length, because each position is compared with every other; during autoregressive generation, the model keeps a cache of past keys and values that grows linearly in sequence length ^[1]. Structured state space models (SSMs) take a different route: they read the sequence one step at a time and compress what they have seen into a state of fixed size, so training scales linearly in sequence length and generation needs only a constant-size state ^[1]. The two families have often been presented as alternatives with different strengths, and the paper's question is how far that opposition actually holds.

The state space duality framework

The paper develops a state space duality (SSD) framework that connects structured SSMs and variants of attention through decompositions of a well-studied class of structured semiseparable matrices ^[1]. A semiseparable matrix is one whose off-diagonal blocks have low rank: a matrix with a great deal of internal structure, so that it can be stored and multiplied by a vector far more cheaply than a dense one. The observation behind the framework is that the computations of both families can be expressed as multiplication by matrices of this kind. An SSM applies them through a step-by-step recurrence, while attention-like layers apply them through a matrix or chunked form. The two are then not rival mechanisms but two ways of organising the same object: one is cheap at generation time, the other is convenient for parallel training on modern hardware. The central claim of the paper is accordingly that the two model families are in fact closely related and can be described within one framework ^[1].

Mamba-2

The framework is not only descriptive: it yields a new architecture, Mamba-2, whose core layer is a refinement of Mamba's selective SSM ^[1]. That core layer is 2 to 8 times faster than Mamba's, while remaining competitive with Transformers on language modelling ^[1]. Because the layer refines an existing design rather than replacing it, Mamba-2 sits on a continuous line running from Mamba, not in opposition to attention.

The title

The title is an homage to Katharopoulos et al. (2020), "Transformers are RNNs", which first showed the equivalence between autoregressive attention and a particular linear recurrence ^[1]. Both titles state an identity rather than an analogy: the point is not that one kind of model resembles the other, but that under a suitable description they are the same computation. Replacing recurrent networks with state space models extends that claim from one family to another.

Authors and publication

The paper is by Tri Dao of Princeton University's Department of Computer Science and Albert Gu of Carnegie Mellon University's Machine Learning Department, and was submitted to arXiv on 31 May 2024 as 2405.21060 ^[1]. It was published at ICML 2024 ^[1].

Why it matters

Read one way, the paper is about speed: making a state space layer fast enough to train on the same hardware and data as transformers. Read another way, it is about how the field draws its map. Attention and state space layers had been described as separate families, and arguments about their relative merits were often arguments about which family one belonged to. In a framework where attention variants appear as particular cases of a structured matrix computation, and SSM layers as another arrangement of the same computation, a comparison means something different: one can ask which arrangement of the computation suits which stage of training or generation, instead of asking which family wins. The 2 to 8 times speedup of the Mamba-2 core layer is a concrete result of that shift in description ^[1].