The UNDO Flip-Flop: A Controlled Probe for Reversible Semantic State Management in State Space Model
AI 摘要
论文提出了UNDO Flip-Flop任务,揭示了SSM理论表达能力与实际学习能力的差距。
主要贡献
- 提出了UNDO Flip-Flop任务,用于评估模型的可逆语义状态管理能力
- 发现Mamba-2在UNDO Flip-Flop任务中无法学习到stack-based rollback机制
- 通过因果消融实验确定瓶颈在于检索而非存储
方法论
设计UNDO Flip-Flop任务,评估Mamba-2模型在非单调更新序列下的状态恢复能力,并进行因果消融分析。
原文摘要
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.