The UNDO Flip-Flop: A Controlled Probe for Reversible Semantic State Management in State Space Model

Hongxu Zhou

cs.LGApr 7, 2026

The UNDO Flip-Flop: A Controlled Probe for Reversible Semantic State Management in State Space Model

Authors: Hongxu Zhou

Organizations: Erasmus Mundus Language and Communication Technologies Saarland University Saarbrücken, Germany

Abstract

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

Panel of Figure 2 — 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.

Figure 3: Behavioural Decomposition Chart

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

State Space Models

State space models represent sequences through evolving latent states, enabling efficient temporal modeling, system identification, and forecasting.

See all 295 papers

Refit the Probe: Single-Direction Ablation Is Not a Necessity Test
May 30, 2026Yuhang JiangTransformer Interpretability Linear Probing
Independent Researcher
Interpreting and Steering State-Space Models via Activation Subspace Bottlenecks
Feb 26, 2026Vamshi Sunku Mohan, Kaustubh Gupta, Aneesha Das +1State Space Language Models State Space Models
AppViewX, Bangalore · Independent Researcher · Microsoft Research, Redmond
The Expressive Limits of Diagonal SSMs for State-Tracking
Mar 2, 2026Mehran Shakerinava, Behnoush Khavari, Siamak Ravanbakhsh +1State Space Models State Tracking
School of Computer Science, McGill University · Mila – Quebec AI Institute · Chandar Research Lab +1

May 30, 2026J/K move · Enter open · S save

Refit the Probe: Single-Direction Ablation Is Not a Necessity Test

Yuhang Jiang

Independent Researcher

Transformer Interpretability Linear Probing Causal Interventions in Language Models Concept Erasure

Open paper AskPDF arXiv

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.

Abstract

State Space Models

Explore similar work

Refit the Probe: Single-Direction Ablation Is Not a Necessity Test

Interpreting and Steering State-Space Models via Activation Subspace Bottlenecks

The Expressive Limits of Diagonal SSMs for State-Tracking