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.
Abstract (translated)
状态空间模型（SSMs）已被证明在理论上具备对无正则序列任务和有界分层结构进行建模的能力（Sarrof 等人，2024）。然而，理论表达能力并不能保证基于梯度的优化能可靠地发现相应解决方案。现有基准测试要么探测单调状态追踪（如标准 Flip-Flop 任务），要么探测结构嵌套（如 Dyck 语言），但两者均未孤立考察可逆语义状态检索。我们引入 UNDO Flip-Flop 任务以填补这一空白。通过在标准 Flip-Flop 基础上增加 UNDO 操作，该任务要求模型维持隐式有界栈，并在非单调更新序列下恢复历史状态。我们在此框架下评估了一层与两层 Mamba-2 模型。两种变体均未能习得理论上可表达的基于栈的回滚机制，反而收敛于一种仅反转当前状态而非检索存储历史的局部切换启发式策略。在训练长度分布内进行的对抗性回撤压力测试中，两层模型准确率降至 41.10%，低于随机概率。这些结果证实了系统性而非偶然性的失败。因果消融分析表明瓶颈在于检索而非存储。这些结果清晰划清了架构在原则上能表示什么与梯度下降能可靠学习什么之间的界限——这是单纯理论表达能力分析无法捕捉的区分。
URL
https://arxiv.org/abs/2604.05923
