The Illusion of State in State-Space Models
William Merrill, Jackson Petty, Ashish Sabharwal
ICML 2024
cs.LG, cs.CC, cs.CL, cs.FL
2024-04-13
S4 and Mamba, like transformers, only express TC0. One-layer RNNs and IDS4 learn A5 at any length; S4/Mamba need depth that grows with sequence length.
Transformers can only express problems in TC0: constant-depth, polynomial-size threshold circuits. State tracking, applying a sequence of world updates in order, often lands in NC1-complete problems. The cleanest example is the word problem for A5, the alternating group on five elements: given a list of even permutations, compute the running product. A one-layer RNN can express it. A transformer cannot.
S4 and Mamba-style state-space models were sold as a parallel stand-in for RNNs. Gu et al. even argued that linear SSMs can simulate general recurrent models. Do they close the state-tracking gap transformers leave open? No.
The object of study is a generalized linear SSM layer. Hidden state updates as hi = Āi h{i-1} + B̄i xi, and the convolutional form unrolls that into a product of transition matrices plus a sum. For S4, Ā is input-independent. For Mamba's S6 layer, Ā is diagonal and may depend on the input (the "selective" part); diagonality reduces the product to per-coordinate scalar multiplies.
A lemma does the heavy lifting: if products of Ā over any interval sit in L-uniform TC0, the whole convolutional form does too. Two routes put production SSMs in that box.
S4, S6/Mamba, and transformers therefore express only TC0. Under the standard conjecture TC0 ≠ NC1 they cannot solve the S5/A5 word problem, nor tasks that encode it: chess move tracking in UCI (source, target) notation, certain formulations of Python evaluation, entity tracking in a long narrative.
The Gu et al. simulation of RNNs assumed unbounded depth. With a fixed number of layers, it does not go through.
Two minimal patches restore power. Putting a nonlinearity inside the recurrent step yields an RNN-SSM that recognizes any regular language in one layer, at the cost of SCAN-style prefix-sum parallelism. Letting Āi be a full input-dependent matrix, IDS4, in the spirit of Liquid S4, leaves TC0 because iterated products of general matrices are not in TC0; one layer can simulate a DFA and SCAN still applies.
The experiment casts word problems as token tagging, with the label at step i equal to the prefix product. Three groups with 60 elements each: abelian Z60, solvable non-abelian A4 × Z5, and non-solvable A5. Models: transformer, RNN, S4, Mamba, IDS4. The plot is minimum depth for 90% validation accuracy versus sequence length.
Figure 3 is blunt. One-layer RNNs and one-layer IDS4 handle arbitrarily long sequences on all three groups. Transformers, S4, and Mamba need depth that grows monotonically with length on A5.
They also need growing depth on A4 × Z5, which is theoretically in TC0. Either these architectures express a proper subset of TC0, or a constant-depth solution exists and they cannot learn it. S4 and Mamba use fewer layers than transformers on the non-commutative tasks; the saving is a constant, not an asymptotic gap.
Chess hardness holds for UCI (source, target) notation. Standard SAN may be easier. Entity-tracking hardness likewise depends on how the prompt is written.
"SSMs are more recurrent, so they track state better" is false as a claim about expressive power. Mamba's selectivity makes Ā input-dependent, but the diagonal constraint keeps the product inside TC0. Real state tracking needs either a full input-dependent Ā (IDS4) or a nonlinearity in the recurrent step. The former stays parallelizable. Whether it trains at language-model scale, and whether gradients explode, is left open.
Anyone betting on Mamba to execute code, track entities over long narratives, or keep a chess board, should not expect it to clear a hurdle transformers cannot.
The proofs identify the convolutional form with the recurrent form; over floats they are not exactly the same. The datatype is c log n-bit floats; finite precision would be weaker. The constructions that escape TC0 need input dimension larger than the alphabet. Experiments are synthetic group multiplication, with no language-model check that entities actually drift. IDS4 is shown to learn the toy task, not to language-model. H3 is out of scope because its context is not a single vector.