Source-linked AI summary
The Illusion of State in State-Space Models
William Merrill, Jackson Petty, Ashish Sabharwal
TL;DR
The paper asks whether SSMs offer RNN-like expressive power for state tracking. It proves that common bounded-depth SSMs remain limited like transformers, then shows experimentally and constructively that an input-dependent extension can overcome this limitation.
Problem
Whether SSMs truly provide RNN-like expressive power for sequential state tracking remains unresolved despite their recurrent-like architecture.
Method
The authors analyze generalized linear SSMs using circuit-complexity arguments and evaluate S4, S6, transformers, RNNs, and an input-dependent SSM on permutation composition.
Results
Common SSM variants, including S4 and S6, are limited to L-uniform TC0 and cannot solve NC1-hard permutation composition; experiments show fixed-depth SSMs struggle while one-layer RNNs succeed.
Takeaways & Limitations
Current SSMs do not have a general state-tracking advantage over transformers, while input-dependent transitions can provide greater expressive power.
Takeaways & Limitations
The chess state-tracking hardness result depends on source-target notation, while standard notation may be easier.
Abstract
from arXiv · showhide
State-space models (SSMs) have emerged as a potential alternative architecture for building large language models (LLMs) compared to the previously ubiquitous transformer architecture. One theoretical weakness of transformers is that they cannot express certain kinds of sequential computation and state tracking (Merrill & Sabharwal, 2023), which SSMs are explicitly designed to address via their close architectural similarity to recurrent neural networks (RNNs). But do SSMs truly have an advantage (over transformers) in expressive power for state tracking? Surprisingly, the answer is no. Our analysis reveals that the expressive power of SSMs is limited very similarly to transformers: SSMs cannot express computation outside the complexity class $\mathsf{TC}^0$. In particular, this means they cannot solve simple state-tracking problems like permutation composition. It follows that SSMs are provably unable to accurately track chess moves with certain notation, evaluate code, or track entities in a long narrative. To supplement our formal analysis, we report experiments showing that Mamba-style SSMs indeed struggle with state tracking. Thus, despite its recurrent formulation, the "state" in an SSM is an illusion: SSMs have similar expressiveness limitations to non-recurrent models like transformers, which may fundamentally limit their ability to solve real-world state-tracking problems.
1. Introduction
The paper argues that common SSMs do not provide the hoped-for RNN-like expressive power for hard state tracking. Theory and experiments show bounded-depth SSMs face limitations resembling transformers, while an input-dependent extension can recover greater expressivity.
- Motivation: SSMs were proposed as alternatives to transformers partly because their recurrent-like design might improve sequential state tracking.The paper frames permutation composition as a simple state-tracking problem that RNNs can express naturally.
- Main result: The authors prove that linear and Mamba-style SSMs, like transformers, cannot express inherently sequential problems such as permutation composition.The result applies to realistic models with a bounded number of layers.
- Implications: These limitations cover state-tracking tasks including chess under source-target notation, Python evaluation, and entity tracking.The chess hardness qualification is notation-dependent: source-target notation is hard, while standard notation may be easier.
- Empirical evidence: Experiments find that S4, S6, and transformers fail to learn permutation composition at fixed depth, whereas a one-layer RNN succeeds.The experiments support the theoretical separation between true recurrent models and the evaluated SSMs.
- Extension: An input-dependent transition-matrix extension of a linear SSM can express and learn permutation composition with one layer while retaining similar parallelizability.Its practical viability for large-scale language modeling remains open.
2. Background
SSMs are sequence-processing architectures resembling simplified RNNs, proposed partly to handle stateful computation while retaining practical advantages such as faster inference. Their analysis uses circuit complexity, algebraic formal language theory, and assumptions about numerical precision to characterize expressivity.
- Motivation: SSMs were proposed as transformer alternatives that might better express inherently sequential or stateful problems, alongside faster inference.The paper notes that SSMs have also shown strong empirical performance in some long-context settings.
- Architecture: SSMs process sequences through recurrent-style layers and include S4 and Mamba’s S6 layer as special cases.A generalized linear SSM layer has recurrent and convolutional forms, producing hidden states and outputs from input-dependent projections.
- Scope: The paper’s expressivity results apply to generalized linear SSMs, including S4 and S6, independently of initialization.H3 is outside the analyzed definition because its context is not represented by a single vector.
- Numeric Datatype: Circuit-complexity analysis treats neural computations over a datatype D whose available bit precision affects formal expressivity.The paper distinguishes fixed finite-precision representations from parameterized or log-precision floating-point models.
- Numeric Datatype: The analysis assumes iterated addition, iterated product, and fixed-size matrix powering are computable in L-uniform TC0.These operations are straightforward for finite-precision datatypes and are also known or established for the paper’s log-precision setting.
- Circuit Complexity: TC0 consists of constant-depth, polynomial-size threshold-circuit families and represents extremely parallel computation.Problems outside TC0 are inherently sequential under the paper’s framing, including several NC1-hard problems unless TC0 = NC1.
3. State Tracking
State tracking can be formalized as reducing a sequence of state updates to a product in a finite monoid. This algebraic view distinguishes easy problems from NC1-complete ones and establishes that certain chess-state sequences are NC1-complete.
- State Tracking: A state-tracking problem applies a sequence of updates to a world state and asks for the resulting state.Finite-world state tracking can be represented as a word problem over a finite monoid.
- Monoid Word Problems: A monoid word problem reduces an arbitrarily long sequence of elements under associative multiplication to one resulting element.The parity example models a light switch, showing that some word problems are simple state-tracking tasks.
- Monoid Word Problems: Word-problem difficulty depends on the underlying algebraic structure: some problems lie in TC0, while others are NC1-complete and require recurrent processing.The paper uses this distinction to connect algebraic structure with neural-network expressivity.
- Permutation Composition: The S5 word problem composes permutations over five objects, and its NC1-completeness connects it to state tracking such as chess moves, code evaluation, and entity tracking.The paper notes that solving arbitrary swap sequences requires determining the final location of an object.
- Chess State Tracking: Chess state tracking in UCI notation takes a board and move sequence as input and returns the resulting board, or a null state for an illegal move.The task ignores draws and differs from SAN because UCI records source and target squares as a tuple.
- Chess State Tracking: S5 reduces to chess state tracking through fixed move sequences that encode permutations, making chess state tracking NC1-complete under AC0 reductions.The construction maps permutation elements to chess moves and decodes the final board through a finite lookup.
4. SSMs Can be Simulated in TC0
The paper shows that the convolutional forms of common non-gated and diagonal SSMs can be simulated by L-uniform TC0 circuits. Therefore, assuming recurrent and convolutional forms compute the same function, these SSMs cannot express inherently sequential problems.
- General Simulation Framework: The analysis reduces SSM expressivity to whether transition-matrix products and parameter computations can be performed in L-uniform TC0.A lemma supplies sufficient conditions involving interval matrix products and input-dependent matrices.
- Non-gated SSMs: For non-gated SSMs with input-independent matrices, the convolutional form is computable by an L-uniform TC0 circuit family.The proof uses matrix powering for repeated transition matrices and treats the remaining parameters as constants.
- S4: S4 satisfies the non-gated conditions, so its convolutional form has an L-uniform TC0 circuit simulation.This follows from the non-gated SSM theorem.
- Diagonal SSMs: Diagonal SSMs also have L-uniform TC0 simulations because iterated multiplication of diagonal matrices reduces to iterated scalar multiplication.The theorem assumes diagonal transition matrices and TC0-computable parameter functions.
- S6 and Mamba: S6, the layer used by Mamba, satisfies the diagonal conditions and therefore has an L-uniform TC0 convolutional-form simulation.Its transition matrix is diagonal because it is defined by exponentiating a fixed diagonal matrix.
- Expressivity Consequences: Common SSM variants are limited to TC0 expressivity and therefore cannot solve NC1-hard problems such as S5 composition under the stated assumption.The paper states this formally for S4 and S6 assuming TC0 ≠ NC1.
- Comparison with RNNs: RNNs can express S5 through finite-state constructions, whereas bounded-depth SSMs cannot simulate the recurrent state that enables this capability.The paper contrasts its realistic bounded-layer analysis with prior SSM-to-RNN simulation arguments using infinite depth.
5. Extending the Expressive Power of SSMs
Minimal SSM extensions can recover expressive power for state tracking, but they introduce practical trade-offs in parallelization and learning dynamics.
- S4 and S6 remain limited to TC0, which excludes state-tracking problems such as S5.Their seemingly stateful design does not suffice for problems requiring expressive power beyond TC0.
- Adding recurrent nonlinearities turns an SSM into an RNN and enables recognition of every regular language with one layer when k = |Σ|.The RNN-SSM applies the nonlinearity at each recurrent update.
- Input-dependent transition matrices make SSMs compute iterated products of general matrices, which can exceed TC0.Unlike diagonal-matrix products, general-matrix products cannot be computed in TC0.
- A one-layer log-precision IDS4 SSM can recognize the word problem for S5 by simulating a deterministic finite automaton through transition-matrix composition.The construction represents transition-monoid elements as boolean matrices and uses matrix multiplication to compose them.
- RNN-SSM and WFA-SSM variants can express the word problem for S5, and these models cannot be simulated in TC0.The two extensions therefore close the expressive gap between ordinary SSMs and recurrent models for this task.
- These extensions raise practical concerns: IDS4 requires iterated matrix products for parallelization, while recurrent nonlinearities are harder to parallelize and may affect learning dynamics.Matrix products may cause vanishing or exploding gradients, and nonlinear recurrent computation is less compatible with SCAN-style parallelization.
6. Can SSMs Learn Permutations in Practice?
Experiments test whether SSMs learn group word problems as sequence length grows. RNN and IDS4 solve arbitrarily long sequences at one layer, whereas transformers, S4, and Mamba require increasing depth on non-commutative tasks.
- The experiments train models on progressively longer sequences and measure full-sequence test accuracy and the minimum depth reaching 90% accuracy.The task uses token tagging, providing supervision at every prefix product.
- Single-layer RNN and IDS4 models learn word problems for arbitrarily long sequences across A5, A4 × Z5, and Z60.These results match the expected state-tracking ability of recurrent and input-dependent-transition models.
- Transformers, S4, and Mamba require depth that monotonically increases with sequence length on non-commutative groups.Longer A5 sequences therefore require deeper models, consistent with their TC0 limitation.
- Transformers, S4, and Mamba also require greater depth on A4 × Z5, despite this problem being theoretically expressible in TC0.The authors suggest either model-specific expressivity limitations or difficulty learning a constant-depth solution.
- S4 and Mamba appear empirically better than transformers at approximate state tracking on the non-commutative tasks.For length-n sequences from A4 × Z5 or A5, transformers require at least as many and often more layers.
7. Conclusion
Common SSMs, including S4 and Mamba, share transformers’ TC0 expressiveness limits and cannot truly track hard sequential state. An input-dependent SSM similar to Liquid S4 offers evidence that these limitations can be overcome.
- S4, Mamba, and related SSMs are confined to L-uniform TC0 and cannot solve inherently sequential problems such as permutation composition.Permutation composition captures hard state tracking because it is NC1-complete.
- One-layer RNNs can easily learn an S5 word problem, whereas linear SSMs require depth growing with sequence length.
- An input-dependent SSM similar to Liquid S4 can both express and learn the S5 word problem.This provides evidence that current SSM expressiveness limitations can be overcome.
Impact Statement
The paper’s formal analysis relies on log-precision floating-point arithmetic and TC0-computable operations, including iterated sums, products, and matrix powers. It defines canonical flattened evaluation to preserve associativity for compositional expressions, with that assumption needed for diagonalizable SSM analysis.
- Log-precision floats use c log n bits, with mantissa and exponent together representing a floating-point value.The model follows the log-precision floating-point framework used to analyze transformers.
- Iterated matrix summation treats entries as reals, performs exact arithmetic, and casts the result back to datatype D.This construction adapts approximate iterated addition over floats to matrices.
- Iterated D-products are formalized for square matrices, extending the arithmetic framework needed to analyze SSM computations.
- Iterated scalar multiplication and matrix powering over log-precision floats are shown to be computable in L-uniform TC0.These are the two special matrix-product cases used in the analysis.
- Flattened expression evaluation combines adjacent sums or products into single operations, giving compositional arithmetic expressions a canonical value.This makes the resulting operation associative; the assumption is needed for the diagonalizable SSM analysis.
A.2. Complexity of Matrix Powering
The matrix-powering analysis reduces float matrix powers to integer matrix powers through exact rescaling, applies an L-uniform TC0 construction, then reverses the scaling and casts back.
- Integer powers of fixed-dimensional n-bit matrices with exponent z ≤ n can be computed in L-uniform TC0.
- The D-matrix-power lemma extends integer matrix powering to datatypes whose entries are equivalently representable as n-bit integers.Casting into and out of the datatype remains within L-uniform TC0.
- Float matrix powers with fixed dimensions and c log n precision are computable in L-uniform TC0.
- The proof rescales a float matrix so every entry represents an integer, applies the D-matrix-power result, reverses the scaling, and casts back.
- The construction computes the rescaled power with a threshold circuit of size poly(m), where m ≤ 2nc, and then divides by 2^qz before recasting.
- Polynomial division is available in L-uniform TC0 by combining earlier P-uniform results with later results for integer division and iterated multiplication.
B. S6 Parameterization
S6 uses fixed and input-derived parameters, then an input-dependent discretization factor to produce discretized state-space matrices. These matrices satisfy the diagonalizability condition and are computable in L-uniform TC0.
- S6 parameterizes a fixed invertible diagonal matrix A and sets D_i = I.
- S6 discretizes A and B using an input-dependent factor δ_i, then defines the resulting discretized matrices.
- The discretized A matrix is diagonal, satisfying the diagonalizability condition used by the analysis.
- The matrices Ā_i, B̄_i, C_i, and D_i can all be computed from x_i in L-uniform TC0.The argument uses linear transformations, diagonal inversion, softplus, and exp.
C. Diagonalizable SSMs
Theorem 4.6 extends the TC0 analysis to SSMs whose transition matrices are simultaneously diagonalizable, showing that their convolutional form remains computable in L-uniform TC0.
- C. Diagonalizable SSMs: Theorem 4.6 covers SSMs with simultaneously diagonalizable transition matrices rather than only diagonal transition matrices.The extension relies on numerical substitutability between a diagonal matrix and its diagonalized decomposition over log-precision floating-point representations.
- C. Diagonalizable SSMs: The transition matrix ¯Ai is computable to log precision as W diag(¯ai)W−1, with ¯ai ∈ R^d.
- C. Diagonalizable SSMs: Each of ¯ai, ¯Bi, Ci, and Di is computable from xi in L-uniform TC0.
- C. Diagonalizable SSMs: The resulting SSM convolutional form is computable by an L-uniform TC0 circuit family.The proof uses associativity of D-matrix products and reduces iterated diagonal-matrix multiplication to iterated scalar multiplication.
C.1. Diagonalizable S6
The diagonalizable S6 construction preserves L-uniform TC0 computability: its discretized parameters and convolutional form can be computed within that class.
- C.1. Diagonalizable S6: The paper extends S6 with continuous-time parameters satisfying the conditions needed for diagonalizable transition matrices.The construction specifies D = I and uses an input-dependent discretization factor δi for A and B.
- C.1. Diagonalizable S6: A diagonalizable matrix can be replaced by W diag(a)W−1 in log-precision computation without meaningful error.The decomposition is A = W diag(a)W−1 for an invertible W and diagonal diag(a).
- C.1. Diagonalizable S6: Scalar multiplication and matrix exponentiation preserve simultaneous diagonalizability through the same decomposition.The matrix exponential satisfies exp(A) = W exp(diag(a))W−1 over the stated representation.
- C.1. Diagonalizable S6: Diagonalizable matrices over log-precision floats, including the needed inverses, are computable in L-uniform TC0.The inverse is formed from W−1, diag(a)−1, and W using a finite number of additions and multiplications.
- C.1. Diagonalizable S6: The discretized parameters ¯Ai and ¯Bi are L-uniform TC0-computable functions of xi, while Ci and Di inherit computability from standard S6.Because A is fixed, all transition matrices are simultaneously diagonalizable via W; therefore the convolutional form is also computable in L-uniform TC0.
- C.1. Diagonalizable S6: Diagonalizable S6 has an L-uniform TC0 circuit family computing its convolutional form.The fixed simultaneous diagonalization of the transition matrices lets the model satisfy Theorem 4.6.
D. Nonlinearities in L-Uniform TC0 .
The paper argues that nonlinearities used in SSMs and transformers remain within L-uniform TC0 under bounded inputs, including broad classes of well-behaved functions.
- D. Nonlinearities in L-Uniform TC0: SSM and transformer parameterizations use nonlinearities such as exp and softplus, whose circuit complexity can be analyzed directly.
- D. Nonlinearities in L-Uniform TC0: Pre- and post-layer normalization keep SSM activations in a bounded domain, enabling the nonlinearities to be computed by the preceding lemma.
- D. Nonlinearities in L-Uniform TC0: A function on a bounded interval with a suitable convergent Taylor series can be approximated to log precision by L-uniform TC0 circuits.The approximation error is at most 2−nc for any c ≥ 1 when the Taylor coefficients meet the stated computability and magnitude conditions.
- D. Nonlinearities in L-Uniform TC0: The bounded-interval result is obtained by rescaling inputs to (−1, 1), applying the Reif–Tate construction, and mapping the result back.When B = 2^k, computing z = x/B is exact.
- D. Nonlinearities in L-Uniform TC0: The pointwise nonlinearities exp, log, and softplus are computable over (−B, B) in L-uniform TC0.Softplus is handled as the fixed composition log(1 + exp(x)).