🤖 machine learning

The Expressive Limits of Diagonal SSMs for State-Tracking

This paper establishes that the expressivity of $k$ -layer diagonal complex-valued State-Space Models (SSMs) for state-tracking is precisely limited to solvable groups with a subnormal series of length $k$ , revealing a fundamental theoretical barrier for non-Abelian groups and an empirical gap between expressivity and learnability.

Original authors: Mehran Shakerinava, Behnoush Khavari, Siamak Ravanbakhsh, Sarath Chandar

Published 2026-08-17

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Mehran Shakerinava, Behnoush Khavari, Siamak Ravanbakhsh, Sarath Chandar

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to follow a set of instructions that change its internal mood or position. In the world of artificial intelligence, this is called "sequence modeling." The robot reads a story one word at a time, and with every new word, it updates its understanding of where the story is going. For a long time, the most famous robots (called Transformers) were great at this, but they were also incredibly hungry for computer power. So, scientists invented a new, leaner type of robot called a State-Space Model, or SSM. Think of an SSM as a robot with a tiny, efficient memory bank that updates itself instantly as it reads, rather than re-reading the whole story every time.

But here's the tricky part: just because a robot is efficient doesn't mean it's smart enough to do everything. Some tasks are like simple math (adding numbers), while others are like complex puzzles involving swapping items around in a specific order. In math, these puzzles are often described using "groups," which are just fancy rules for how things can be combined or swapped. If the rules are simple and predictable (like adding numbers), the robot might handle them easily. But if the rules are messy and depend on the order of operations (like swapping three different colored balls where swapping red then blue is different than blue then red), the robot might get confused. Scientists want to know: What is the absolute limit of what these efficient robots can actually understand? Can they learn to solve the messy puzzles, or are they stuck with the simple ones?

This paper dives deep into that question by testing a specific, very efficient kind of robot called a "Diagonal SSM." The researchers wanted to see if these robots could learn to track complex state changes, specifically those involving non-Abelian groups (the messy, order-dependent puzzles). They discovered a hard mathematical wall: a single-layer diagonal robot simply cannot solve these messy puzzles, no matter how hard you try. However, if you stack two or more of these robots on top of each other, they theoretically gain the power to solve them, but only if the puzzle has a specific, layered structure (called a "solvable group").

Here is the twist that makes the story interesting: even though the stacked robots should be able to solve these puzzles in theory, the researchers found that in practice, they often fail to learn how to do it. It's like having a car with an engine powerful enough to climb a mountain, but the driver (the training algorithm) keeps getting stuck in a ditch and can't find the path up. The paper proves that the power exists within the architecture, but standard training methods struggle to unlock it, especially for the more complex, non-Abelian tasks.

The Robot's Memory Bank

To understand what the authors found, we first need to look at how these "Diagonal SSMs" work. Imagine a robot with a row of light switches. When it reads a new word, it flips these switches based on a simple rule. In a "Diagonal" SSM, the rule for flipping one switch doesn't depend on the other switches; each switch acts independently. This makes the robot incredibly fast and easy to train, like a team of workers each doing their own job without needing to talk to each other.

The researchers tested these robots on a game called "state-tracking." In this game, the robot is given a sequence of commands (like "swap the red ball," "rotate the blue ball") and must remember the final arrangement of the balls. If the commands are simple and commute (meaning the order doesn't matter, like adding 2 + 3 is the same as 3 + 2), the robot is a whiz. But if the commands don't commute (like putting on socks before shoes vs. shoes before socks), the robot faces a much harder challenge.

The Theoretical Wall: One Layer vs. Many

The paper starts with a strong mathematical proof about what a single-layer robot can do. The authors showed that a single-layer diagonal SSM is fundamentally limited to solving only "Abelian" puzzles—those where the order of operations doesn't matter. No matter how you tune the robot, if it only has one layer of these independent switches, it cannot track the state of a complex, non-Abelian group (like the permutations of three objects, known as the group $S_3$ ).

However, the story gets more exciting when you stack layers. The authors proved that if you stack $k$ layers of these diagonal robots, the team can solve a puzzle if and only if the puzzle can be broken down into a chain of $k$ simpler, Abelian steps. Think of it like a relay race. If the puzzle is too complex for one runner, you can pass the baton to a second runner, who passes it to a third. As long as the puzzle can be sliced into a specific number of simple, ordered steps, a team of that many robots can theoretically solve it. This means that depth (adding more layers) is the key to unlocking more complex intelligence, but it comes with a strict rule: you need exactly as many layers as the puzzle has "slices" of complexity.

The Learning Gap: Theory vs. Reality

This is where the plot thickens. The authors didn't just stop at the math; they built these robots and tried to train them. They tested single-layer and two-layer models on various tasks, including simple addition (Abelian) and the tricky $S_3$ permutation puzzle (non-Abelian).

The results were a mix of success and frustration. For the simple tasks, the robots learned quickly and could handle very long sequences. But for the complex $S_3$ task, even though the two-layer robot should have been powerful enough to solve it (according to their math), it mostly failed. The robot just couldn't learn the right pattern, even after thousands of tries.

The researchers found that the problem wasn't that the robot couldn't do it; the solution actually existed within the robot's design. It was more like the robot was looking for a needle in a haystack, and the standard training method (gradient descent) kept missing the needle. In one experiment, they tried to "help" the robot by starting it very close to the correct answer. When they did this, the robot suddenly learned the task and could handle sequences much longer than it had seen before. This suggests that the solution is there, hidden in the robot's "weight space," but the path to find it is incredibly narrow and difficult to navigate with standard training.

The Takeaway

The paper concludes that while diagonal SSMs are efficient and theoretically capable of solving complex state-tracking problems if you stack enough layers, there is a massive gap between what they can do and what they actually learn. The math says a two-layer robot can solve the $S_3$ puzzle, but in practice, it usually fails. This highlights a crucial lesson for AI research: just because a model has the potential to be smart doesn't mean it will easily learn to be smart. The authors suggest that to bridge this gap, we might need to change how we train these models or perhaps tweak their architecture slightly (like allowing a bit more interaction between the switches) to make the path to the solution less slippery.

In short, these efficient robots are like brilliant students who have the textbook knowledge to solve a hard problem but keep failing the exam because they can't figure out how to apply the knowledge. The paper maps out exactly what they are capable of, but also warns us that getting them to actually perform is a much harder challenge than the math suggests.

Technical Summary: The Expressive Limits of Diagonal SSMs for State-Tracking

Problem Statement

State-Space Models (SSMs) have emerged as efficient, parallelizable alternatives to Transformers for long-range sequence modeling. However, their theoretical expressive power remains poorly understood, particularly regarding their ability to perform state-tracking tasks. These tasks, which involve maintaining and updating a hidden state based on a sequence of inputs (e.g., parity checks, modular addition, or permutation tracking), are closely related to regular languages and are critical for applications like code execution and program analysis.

While Transformers are known to struggle with certain state-tracking tasks (specifically those requiring non-solvable group operations), SSMs were initially hypothesized to excel due to their recurrent nature. Recent work has shown that standard SSMs (including time-invariant and diagonal variants) fail on even simple solvable tasks like parity. This paper investigates the precise expressive limits of Input-Dependent Complex-valued Diagonal (DCD) SSMs on sequential state-tracking tasks, specifically focusing on the ability to track finite groups. The core question is: What classes of groups can a $k$ -layer DCD SSM theoretically represent, and can standard learning algorithms reliably find these representations?

Methodology

The authors employ a combination of group-theoretic analysis and empirical evaluation.

Theoretical Framework

Model Definition: The study focuses on DCD SSMs, where the transition matrix $A(x_t)$ is diagonal and complex-valued, and depends on the input $x_t$ . The state update follows an affine recurrence: $h_t = A(x_t)h_{t-1} + b(x_t)$ .
Group Theory Connection: The authors map state-tracking tasks to the simulation of finite groups. A model tracks a group $G$ if it can compute the cumulative product of a sequence of group elements.
Single-Layer Analysis: The authors prove that a single-layer DCD SSM cannot track any non-Abelian group at finite precision. They demonstrate that even with an input-dependent bias term $b(x_t)$ and a decoder, the commutative nature of diagonal matrices restricts the model to Abelian groups.
Multi-Layer Analysis: The paper extends this to $k$ $k$ -layer DCD SSMs. They establish a necessary and sufficient condition for expressivity: A $k$ $k$ -layer DCD SSM can track a group $G$ $G$ if and only if $G$ $G$ admits a subnormal series of length $k$ with Abelian factors.
- Mathematically, this means $G = G_k \trianglerighteq G_{k-1} \trianglerighteq \dots \trianglerighteq G_1 \trianglerighteq \{e\}$ , where each quotient $G_{i+1}/G_i$ is Abelian.
- This result identifies the precise expressivity range of $k$ -layer DCD SSMs as the class of solvable groups with a specific structural depth.
Constructive Proof: The authors provide a concrete construction for a 2-layer DCD SSM tracking the symmetric group $S_3$ (the smallest non-Abelian solvable group) by decomposing it into a cascade of two Abelian automata ( $C_2$ and $C_3$ ).

Empirical Evaluation

The authors evaluate single- and two-layer variants of several SSM architectures (Mamba, Negative Mamba, AUSSM, and simplified AUSSM) alongside RNNs.

Tasks: State-tracking on various groups, including Abelian groups ( $C_2, C_6, C_{24}, C_{60}$ ), products of cyclic groups ( $C_2 \times C_4, C_3 \times C_6$ ), and solvable non-Abelian groups ( $S_3, A_4$ ). The non-solvable group $A_5$ is included as a control.
Protocol: Models are trained on sequences up to length 60 using curriculum learning and tested for generalization on sequences up to length 1000.

Key Contributions

Theoretical Characterization: The paper provides the first precise characterization of the expressive power of diagonal SSMs. It proves that:
- Single-layer DCD SSMs are limited to Abelian groups.
- $k$ -layer DCD SSMs are limited to solvable groups with a subnormal series of length $k$ .
- This establishes a strict theoretical benefit of depth for modeling non-Abelian dynamics within the diagonal constraint.
Identification of the Learnability Gap: While the theory proves that 2-layer DCD SSMs can represent non-Abelian groups like $S_3$ , the empirical results show that standard gradient-based optimization fails to learn these solutions. Models consistently fail to generalize to long sequences for $S_3$ and $A_4$ , despite the theoretical capacity to do so.
Initialization Insight: The authors demonstrate that initializing a model near the analytical solution (derived from their theoretical construction) allows it to learn and generalize. This suggests that the failure is not due to a lack of capacity in the weight space, but rather the difficulty of the optimization landscape in finding these specific encodings without guidance.

Results

Abelian Groups: Single-layer models (specifically Negative Mamba and AUSSM) successfully learn and generalize on Abelian tasks (e.g., parity, modular counting).
Non-Abelian Solvable Groups: Multi-layer models (2-layer) theoretically capable of tracking $S_3$ and $A_4$ fail to generalize in practice. They achieve high accuracy on training lengths but collapse on longer test sequences.
Non-Solvable Groups: As expected, no model (including multi-layer DCD SSMs) learns the non-solvable group $A_5$ .
Comparison with RNNs: RNNs, which have no theoretical expressivity barrier for these tasks, generally succeed in learning and generalizing, highlighting that the failure of SSMs is specific to their architectural constraints and optimization dynamics, not just the difficulty of the task.

Significance and Claims

The paper claims to establish a sharp distinction between expressivity and learnability in diagonal SSMs.

Expressivity: Depth provides a disciplined, theoretically provable pathway to increase the expressive power of diagonal SSMs from Abelian to solvable groups. The authors define the exact boundary of this capability using group theory.
Learnability: The authors argue that while these models possess the theoretical capacity to solve complex state-tracking tasks, standard training algorithms struggle to discover the required weight configurations. This suggests that the "inductive bias" of diagonality, while efficient, may actively hinder the discovery of non-Abelian structures.
Implications: The work suggests that simply stacking diagonal layers is insufficient for reliable performance on non-Abelian tasks without better optimization strategies or architectural modifications (e.g., block-diagonal structures). It calls for a separation of concerns when evaluating new sequence architectures: a model's theoretical capacity does not guarantee its practical utility.

The authors conclude that future progress requires not only expanding the expressive frontier (e.g., via block-diagonal transitions) but also developing training methods capable of reliably navigating the loss landscape to find solutions that are guaranteed to exist in principle.