State space models (SSMs) have been shown to possess the theoretical capacity to model both star-free sequential tasks and bounded hierarchical structures Sarrof et al. (2024). However, formal expressivity results do not guarantee that gradient-based optimisation will reliably discover the corresponding solutions. Existing benchmarks probe either monotonic state tracking, as in the standard Flip-Flop task, or structural nesting, as in the Dyck languages, but neither isolates reversible semantic state retrieval. We introduce the UNDO Flip-Flop task to fill this gap. By extending the standard Flip-Flop with an UNDO, the task requires a model to maintain an implicit bounded stack and recover historical states under non-monotonic update sequences. We evaluate one-layer and two-layer Mamba-2 under this framework. Both variants fail to acquire the provably expressible stack-based rollback mechanism, converging instead on a local toggle heuristic that inverts the current state rather than retrieving stored history. Under an adversarial retraction pressure test held within the training length distribution, the two-layer model collapses to 41.10% accuracy, which is below random chance. The results confirm systematic rather than incidental failure. Causal ablation shows that the bottleneck lies in retrieval, not storage. These results draw a clear line between what an architecture can in principle represent and what gradient descent reliably learns, a distinction that theoretical expressivity analyses alone cannot capture.
Figures & tables
Step
1
2
3
4
5
6
7
Operation
W(1)
W(0)
I
W(1)
UNDO
I
R
Stack
[1]
[1,0]
[1,0]
[1,0,1]
[1,0]
[1,0]
Table 1: Example of the [UNDO] Flip-Flop. Final output: 0 .
Model Variant
Epoch
Train Acc. (%)
ID Test Acc. (%)
OOD Test Acc. (%)
Baseline standard Flip-Flop, One-Layer
25
100
100
96.25
Baseline standard Flip-Flop, Two-Layer
5
100
100
99.05
[UNDO] Flip-Flop, One-Layer
100 (max)
95.95
95.95
79.10
[UNDO] Flip-Flop, Two-Layer
100 (max)
98.55
98.55
87.35
Table 2: Performance on standard and [UNDO] Flip-Flop tasks.
Figure 1: Asymmetric generalisation failure under UNDO
Figure 2: Mechanistic comparison between the robust stack mechanism (left) and the observed toggle heuristic (right). The shortcut fails to resolve consecutive identical bits because it treats [UNDO] as a bitwise inversion of the current state rather than a historical retrieval.
Probes are routinely paired with an intervention: ablate the direction the probe found, run the model, and read the change in task accuracy, taking a large drop as evidence that the computation depends on what the probe read and a near-zero drop as evidence that it does not. Either inference requires that the ablation have removed the target from the layer. We find that the ablation does not remove what it targets. A probe refitted on the ablated activations recovers its original accuracy in every cell we test, and keeps recovering when the probe's entire row space is deleted rather than a single axis, because the quantity survives in the orthogonal complement. Because a refitted probe recovers, neither a large task drop nor a near-zero one establishes whether the model needed the target, and one probe fit detects this. Replacing the ablation with iterative nullspace projection, scored against random subspaces of matched dimension, reverses the conclusion: representations that looked causally inert carry most of the task. The correction also separates where a variable is most readable from where deleting it does most damage, and those are not the same layer in any pretrained model we study. The erasure is defined by a linear probe family, so removing a nonlinearly encoded quantity remains open.
State-space models (SSMs) have emerged as an efficient strategy for building powerful language models, avoiding the quadratic complexity of computing attention in transformers. Despite their promise, the interpretability and steerability of modern SSMs remain relatively underexplored. We take a major step in this direction by identifying activation subspace bottlenecks in the Mamba family of SSM models using tools from mechanistic interpretability. We then introduce a test-time steering intervention that simply multiplies the activations of the identified bottlenecks by a scalar. Across 7 SSMs and 6 diverse benchmarks, this intervention improves performance by an average of 8.27%, without requiring any task-specific tuning. Finally, we validate that the identified bottlenecks are indeed hindering performance by modifying them to yield an architecture we call Stable-Mamba, which achieves long-context performance gains when retrained from scratch.
Vamshi Sunku Mohan, Kaustubh Gupta, Aneesha Das +1
AppViewX, Bangalore · Independent Researcher · Microsoft Research, Redmond
State-Space Models (SSMs) have recently been shown to achieve strong empirical performance on a variety of long-range sequence modeling tasks while remaining efficient and highly-parallelizable. However, the theoretical understanding of their expressive power remains limited. In this work, we study the expressivity of input-Dependent Complex-valued Diagonal (DCD) SSMs on sequential state-tracking tasks. We show that single-layer DCD SSMs cannot express state-tracking of any non-Abelian group at finite precision. More generally, we show that k-layer DCD SSMs can express state-tracking of a group if and only if that group has a subnormal series of length k, with Abelian factors. That is, we identify the precise expressivity range of k-layer DCD SSMs within the solvable groups. Empirically, we find that multi-layer models often fail to learn state-tracking for non-Abelian groups, highlighting a gap between expressivity and learnability.