Structured Sparse Transition Matrices to Enable State Tracking in State-Space Models

We thank the reviewer for their careful reading of our paper. We are glad that the reviewer finds the paper well written and that they appreciate the theoretical analysis and motivation. We are pleased that they find our benchmarks suitable and effective, even as indicators for real-world language modeling benefits.

Weaknesses 1 and 2: We agree that the paper would benefit from a more thorough comparison with the variety of modern SSMs, in particular those adopting BD and DPLR transition matrices. To address this, we provide significantly extended results to strengthen our contribution.

Firstly, the table below extends the experimental evaluation from Walker et al. (2025). We now evaluate +10 recurrent and parallelizable sequence models on four FSA emulation tasks. We fully conform to their experimental setup which adopts a two-layer architecture and uses a single, fixed set of hyperparameters on all tasks. While the authors used the best results obtained with state size 128 or 512, we fixed the state size to 128. Notably, our model sets the state-of-the-art performance with no hyperparameter tuning.

On the Pathfinder dataset, our model achieves 62.61% average accuracy. The average LRA accuracy of our method is now 74.43%, compared to the 66.44% achieved by Mamba and 73.89% achieved by S7.

To extend the evaluation of our method on more real-world datasets, we further add results on time-series classification with long-range dependencies, following Walker et al. (2025). To obtain the results, we ran a grid search over learning rates in {1e-5, 1e-4, 1e-3} and number of layers in {1, 2}. The state size is 16, and the embedding size is 64. Notably, this hyperparameter grid is significantly less extensive than the one used to evaluate the baseline methods. We happily report that PD-SSMs natively also achieve a new state-of-the art for SSMs on this dataset, notably without adopting a CDE structure.

In addition to the experimental results, we want to highlight that we provided theoretical arguments in the paper on why PD is favourable over the BD and DPLR structures, although we agree with the reviewer that we did not explicitly spell them out:

• We proved the existence of FSAs that cannot be emulated by an SSM with state size less than N-1.
• We proved the fact that a single-layer PD-SSM can emulate any N-state FSA with state size N.

Taken together, these results clearly speak against adopting block-diagonal (BD) transition matrices, since they are theoretically less expressive than PD-SSMs (they cannot mix between blocks), but more expensive (O(b^2 d L) vs our O(d L)). In addition, theoretical constructions for DPLR models to emulate FSAs rely either on an exponential number of layers and DPLR factors (Grazzi et al., 2024), or rely on MLPs with exponentially large hidden layers (Peng, Zhang, Goldstein et al, 2025), in stark contrast with the guarantees provided by our parametrization. We will expand on this analysis in a revision of our text.

Weakness 3: Due to the restricted number of pages in the initial submission, we were indeed unable to further open up this point in the main body. Still, as the reviewer points out, we provide a detailed break-down on the empirically observed computational cost in the Appendix, and further mention this in the limitations section. While our implementation is not optimized on the kernel level, Table S8 suggests that the generation of P matrices is responsible for most of the compute. Indeed, the table demonstrates that a parallel scan with PD transition matrices only has a 2x overhead compared to diagonal transition matrices. To address the reviewer's concern, we will switch to a log-log plot for Figure 4 in the main body of the text, move Table S8 from the Appendix to the main sections, and more thoroughly discuss the P matrix generation as a bottleneck. We suggest the following additional text:

As shown in Figure 4, PD-SSM scales significantly better than SD-SSM. PD-SSM’s higher runtime compared to diagonal SSM stems primarily from the additional operations in the generation of P: the simultaneous soft selection of transition matrices $M_i$ across the sequence during parallel execution incurs additional compute and memory. Table S8 confirms these findings and suggests potential improvements in future work.

Weakness 4: We appreciate the reviewer catching this inconsistency. The two tasks, C2xC30 and D30 were indeed trained on sequences of length up to 90.

Question: Hyperparameters

• Long-Range Arena: We performed a grid search defined by all combinations of the following values: Learning rate in {1e-5, 5e-5, 1e-4, 5e-4, 1e-3}, weight decay in {0.001, 0.01, 0.1, 0.05}, and dropout in {0.0, 0.1, 0.3}. On each task, we report the average accuracy over three random seeds. The embedding dimension and state size were both fixed to 128, with the exception of the Retrieval dataset in which we reduced it to 64. We used Adam with the default parameters (0.9, 0.999). The hyperparameters of the models we compare with were obtained via an extensive Bayesian hyperparameter search, as outlined in (Soydan, Zubic et al., 2025).

• FSA Results (Section 4.2): We ran a grid search over learning rates in {5e-5, 1e-4, 5e-4, 1e-3}, state sizes in {64, 128}, and the number of dictionary matrices (K) in {6, 8}. The reported baseline results are taken from Terzic et al. (2025), who ran similar hyperparameter grid searches for all baseline methods.

• New FSA Results: We re-used the hyperparameters from the referenced paper (Walker et al., 2025) with no tuning and with the state size fixed to 128.

• Time Series: We ran a grid search over learning rates in {1e-5, 1e-4, 1e-3} and number of layers in {1, 2}. The state dimensionality was fixed to 16, and the embedding dimensionality was fixed to 64. All other hyperparameters were re-used from the referenced paper (Walker et al. 2025). Notably, this hyperparameter grid is significantly less extensive than the one used to evaluate the baseline methods, see Walker et al. (2024).

Suggestions:

1. Comparison with BD and DPLR: Addressed in the two tables above, as well as in our theoretical arguments.

2. Log-scale y axis for runtime: Agreed, a logarithmic plot will provide more information than the current format.

3. Discussion of the Structured-LNCDE Paper: We became aware of the paper shortly after submission and thank the reviewer for bringing it up. We agree that it is very relevant related work and we now thoroughly engage with it through both experimental and theoretical claims, and accordingly made a benchmarking comparison with it (see above tables). Interestingly, the referenced work theoretically analyses sparse SSMs. The authors note that unstructured sparse SSMs do not provide any efficiency benefits over unstructured dense SSMs, and note that they anticipate that ongoing work on sparse matrices will enable future gains in efficiency. Our work exactly presents a structured sparse SSM with significant efficiency gains compared to unstructured variants.

Once again, we thank the reviewer for their thoughtful feedback, which enabled us to substantially improve the paper, and hope that by addressing all their concerns, we can convince them to improve their score.

Thank you very much for being considerate in this matter. The intended effect of these lines of code is ensuring that all models are trained and evaluated on the same data. The global setting of seeds to 1 happens only after model initialisation, meaning that it affects the randomly sampled training and evaluation data; it does, however, not affect the random model initialisation. To ensure that all models are trained on exactly the same data, the seed should have been set to 1 before the first training batch is sampled. We instead set it the manual seed at model validation, meaning that most but not all of the training data is the same across models. The evaluation data is shared across the models. Compared with standard practices which use the exact same data to train different models, our code can only increase the variance of the results. The referenced lines of code induce the effect of training models on a fixed dataset for a certain number of epochs.

We should also clarify that the code which manually sets the seed to 1 in range_evaluation.py was used to train models on FSA tasks in our section 4.2, in which we only reported the best-performing model out of three random seeds. To obtain the results we used in the rebuttal, which indeed exhibit a very low standard deviation, we imported our model into the training pipeline provided by the authors of the paper you referenced, "Structured Linear CDEs: Maximally Expressive and Parallel-in-Time Sequence Models" by Walker et al. 2025. We only made a single modification to their training pipeline, that of removing the padding of sampled sequences by a fixed length. Concretely, the authors initially padded all of the sampled sequences with 256 zeros, which very significantly and unnecessarily increased the training time (the maximum training sequence length is 40, to which we would then always append 256 inconsequential zeros). We instead pad each sampled batch such that all of the sequences are as along as the longest sampled sequence. We have made no other modification to their code, and the omission of the padding should have no consequence on the final results, given that the gradients are masked for the padded parts of the sequence, and the validation accuracy computation also naturally takes padding into account.

To be fully specific, in the PyTorch GitHub repo provided by the above referenced paper, in the train.py function, lines 364 and 367 have None as the padding length in our implementation, rather than 40 or 256 the authors used. This influences the collate_fn function in file data_dir/dataloaders.py, so that the three conditional lines of code 240-242 are no longer executed.

To be fully certain that this has no effect on the results, we have re-trained our model with padding length set to 40, the same as what was used in the baseline LNCDE methods, and we observed no effect on the results. All of the runs of PD-SSM converge to almost perfect accuracy with very low standard deviation.

Regular Language Results

You are correct to question the surprisingly low standard deviation. The models are not learning the same weights; our reported deviation numbers are in the wrong units and should be multiplied by a factor of 100 (only for the regular language table). We apologize for the confusion.

UEA Results

Thank you for pointing this out. The correct statement should have read "Notably, this hyperparameter grid is significantly less extensive than the one used to evaluate the baseline SSM methods, see Walker et al. (2024)". Indeed, SLiCEs did not undergo a hyperparameter optimization; the baseline methods that did undergo such a search include all of the SSMs and LNCDE in Walker et al. 2024, as well as the work on oscillatory SSMs by Rusch and Rus at ICLR 2025.

In an effort to equalize the scale of the hyperparameter search among the SSM methods, we now extended the hyperparameters search of our PD-SSM using the following grid: Learning rate in {1e-5, 1e-4, 1e-3}, number of layers in {2, 4, 6}, state dimension in {16, 64, 128} and embedding dimension in {16, 64, 128}. We report the results below:

Note that certain results we had previously reported are missing; these results were obtained using 1 layer, which are excluded as the baselines were not evaluated using 1 layer. The baselines were also not evaluated using state size 128, with the grid rather being {16, 64, 256}. We used {16, 64, 128} to more quickly obtain results relevant for this discussion.

Theoretical Expressiveness

Let us expand a bit on our statements. Our formal analysis is entirely based on Proposition 3, which states the following: For any $N>1$ there exists an FSA with $N$ states that cannot be emulated by any single-layer SSM with state size less than $N-1$ under unique state encodings. The proof is in the appendix, section C. The assumption of unique state encodings means that each state of the FSA gets assigned to a unique vector in $\mathbb{R}^d$.

The proposition makes no assumption on the structure of the transition matrix. The richer intra-block dynamics of dense (or block-diagonal) SSMs do not influence the conclusion.

This implies that a block-diagonal SSM would need blocks of size at least $N-1$ in order to emulate the $N$-state finite-state automaton. The cost of matrix-matrix multiplication would then be O($N^3$), as compared to O($N$) in our sparse parametrisation.

We know that our sparse parametrization can emulate any $N$-state FSA with state size $N$. To see this, notice that we can identify each FSA state with its one-hot representation. The mapping that we describe in section 2.2 then results in a matrix with a column one-hot structure.

You are correct to point out that this does not prove that the PD method is in general more expressive than a block-diagonal system. When we made that statement, we implicitly assumed that both SSMs operate using the same state size and that the block-diagonal system has more than one block. In this case, there exists the difference in expressivity described above. Concretely, under the assumption of unique state encodings, SSM hidden dimension N, there exists an N-state FSA that can be emulated by PD-SSM but cannot be emulated by a block-diagonal SSM with maximum block size less than N-1.

One can question the assumption of unique state encodings. However, if we do not assume unique state encodings and instead allow arbitrary state representations, a single-unit SSM can represent a very large number of states, bounded only by the numerical representational capacity, as discussed in footnote 2 on page 6. We thus believe that our assumption is realistic, but we are open to alternative analyses. One might, for example, assume that a cluster of vectors at Euclidean distance $\epsilon$ from a centroid encodes a state.

Of course, our analysis is relevant in the context of FSAs. One can easily imagine that, under a different setting, block-diagonal SSMs might have advantages exactly due to the intra-block dynamics which you mention. There might be patterns that the sparse transitions of PD-SSM cannot capture. Some insight may be gained from the Walker et al. (2025) analysis of sparse transition matrices; the analysed matrices are however not column one-hot.

We thank the reviewer for their exceptionally careful reading of our paper. We are glad that they recognize the importance of the studied problem and the theoretical and methodological contributions of our work. We are pleased that the reviewer recognizes the empirical analysis and practical considerations addressed. It is very encouraging to hear that our work would merit acceptance given an appropriate revision, and we are especially thankful for the constructive and comprehensive list of suggestions for improvements.

Insufficient Experimental Validation

Experimental Setup Details: The text we submitted indeed omits important information on the experimental setup, including the optimizer, loss functions, hyperparameter search procedures. The accidental omission of these details is unfortunate, and have made sure to include them in full detail in the final text. For example, the experimental details on LRA only listed the hyperparameters for each task. In contrast, the improved text now reads as follows:

The optimal hyperparameters for our model on the long-range arena dataset were obtained va grid search defined by all combinations of the following values: Learning rate in {1e-5, 5e-5, 1e-4, 5e-4, 1e-3}, weight decay in {0.001, 0.01, 0.1, 0.05}, and dropout in {0.0, 0.1, 0.3}. On each task, we report the average accuracy over three random seeds. The embedding dimension and state size were both fixed to 128, with the exception of the Retrieval dataset in which we reduced it to 64. We used Adam with the default parameters (0.9, 0.999). The hyperparameters of the models we compare with were obtained via an extensive Bayesian hyperparameter search, as outlined in (Soydan, Zubic et al., 2025).

We agree that clearly stating such information is central to gauge the fairness of our comparison, and have added analogous descriptions to all of the tasks considered.

Comparison with Recent Models and Limited Experimental Scope: We agree that a thorough comparison with the variety of modern SSMs is of central importance. While we did address the two extreme ends of the spectrum, diagonal and dense transition matrices, a more thorough comparison with a range of established models was missing. To address this, we now provide a significantly extended comparison with related methods.

Firstly, we extended the Table from Walker et al. (2025), by evaluating our proposed model on the full set of state tracking tasks from Deletang et al., (2023). We fully conform to their experimental setup which adopts a two-layer architecture and uses a single, fixed set of hyperparameters on all tasks. While the authors used the best results obtained with state size 128 or 512, we fixed the state size to 128. A summary of central results is shown below. Notably, our model sets the state-of-the-art performance with no hyperparameter tuning.

While the table above significantly extends the set of experiments and comparisons we already report on, it is still limited to FSA emulation. We extend the evaluation of our method to more real-world tasks by including additional experiments on time-series classification with long-range dependencies.

PD-SSMs achieve a new state-of-the-art among the evaluated state space models on this dataset (taken from Walker et al. 2025), notably even without building on controlled differential equation (CDE) architectures. S6 is the previous best-performing evaluated non-CDE SSM, with a more extensive overview of SSM results provided in Walker et al. (2024).

Presentation Issues

We thank the reviewer for pointing out the notational and conceptual inconsistencies. We aimed to provide a self-contained tutorial on the relevant concepts and frameworks. We are thus particularly grateful for the detailed proofread provided by the reviewer. In addition to their suggestions, we have found several unfortunate inconsistencies that we have since corrected. For example, the initial submission's Theorem 3 states the following:

Under the assumption that ${TC}^0 \neq {NC}^1$, constant depth and width diagonal log-precision SSMs can be emulated by a logspace-uniform $TC^0$ circuit.

The assumption is, however, not required for the stated conclusion, and the implications of the conclusion are not clearly spelled out. The improved version of Theorem 3 is now summarized as:

Under the assumption that ${TC}^0 \neq {NC}^1$, log-precision diagonal SSMs cannot emulate automata with non-solvable transformation groups.

Addressing the Suggestions:

Lines 79-82: The text now explicitly states: $\forall \varepsilon > 0, \forall A \in \mathbb{R}^{N\times N}, \exists A_\varepsilon \in \mathbb{R}^{N \times N}$, invertible $W \in \mathbb{C}^{N \times N}$ and diagonal $\Lambda \in \mathbb{C}^{N \times N}$ s.t. $A_\varepsilon=W\Lambda W^{-1}$ and $\||A - A_\varepsilon\||_F < \varepsilon.$

Lines 79-82: Assuming $A(u_t)$ are diagonalizable under the same change of basis already implies individual diagonalizability.

Lines 144-148: The specific choice of the P parametrization does in fact provide certain practical benefits compared to alternatives such as a linear generator by reducing the number of parameters and operations. We now include a paragraph elaborating on those.

Section 4.2 We agree that the proposed experiments are interesting ones, and we did in fact report results with these two ablations. Due to the page constraint, the ablations are part of the appendix (E.2 and E.3). We deeply appreciate the suggestion for improvement; a more extensive version of the ablations can be of benefit in the main text too.

Lines 227-231 The new version of the paper provides a full sketch of the architecture, highlighting architectural details such as positioning of the norms and the design of the residual connections.

Addressing the Appendix Issues:

1. We now refer to the cyclic group using Zn.
2. We now provide a definition of $\delta(\cdot, \sigma_{1,\dots,T}):Q \rightarrow Q$ as the composition of the single-step $Q \rightarrow Q$ functions $\delta(\cdot, \sigma_1), ..., \delta(\cdot, \sigma_T)$, with $\sigma_i \in \Sigma$.
3. We now denote the group set as G and denote the binary operation by \cdot. We have appropriately quantified all objects in the sense that it is always clear that the elements listed in the different properties belong to the group set G.
4. We now make no references to automata in this section.
5. The factor group should indeed be written as (Z₂ × Z₄)/Z₄. We have corrected this error and ensured consistent notation throughout.
6. See above.
7. We have removed the mention of uniformity in this definition.
8. The definition of $NC^1$ has been revised to correctly state that the gate fan-in is restricted to 2.
9. We have revised the text to better conform to standard notation, writing total derivatives using $d$ and partial derivatives using $\partial$.
10. We have corrected the typos in the proof and subsequent results to ensure consistency with the presented indexing.
11. The text now makes it clear that P is an element of the monoid H.
12. We have incorporated the helpful suggestion.

Due to space constraints, we could not address each individual suggestion in the rebuttal. Once again, we thank the reviewer for their thoughtful feedback, including their extensive list of comments and suggestions. Addressing their concerns and following their suggestions allowed us to substantially improve the paper, and by doing so, we hope to convince them to recommend acceptance.

We are glad that the reviewer has a positive impression of our paper, highlighting its clear motivation, rigor, and strong empirical evaluation, as well as the premise of our PD parameterization in enhancing expressivity with minimal memory overhead.

Extended evaluation, including standard datasets:

We share the reviewer's interest in how the increased expressiveness translates to other impactful tasks, such as e.g. language modeling, code generation, or mathematical problem solving, tasks typically executed by LLMs. Such an evaluation was unfortunately out of the scope of our paper, but it is a future work we are deeply interested in. Based on recent literature, increased SSM expressivity can translate to practical benefits in other real-world tasks. Grazzi et al. (ICLR 2025) demonstrate that increasing the expressivity of DeltaNet, which can be interpreted as a matrix-valued SSM, translates to significantly reduced perplexity on code generation and mathematical problem solving. Hu et al. (ACL 2025) consistently observe that models pre-trained on formal languages, specifically context-free and context-sensitive languages, exhibit increased performance on standard NLP tasks. A large-scale evaluation of the benefits of our concrete implementation is certainly something we plan to address in future work.

While we do not provide language modeling performance, we have extended the evaluation of our model to an extensive collection of baselines (+10 new models) on: 1) four tasks in FSA emulation (two new tasks compared to our submission); 2) a real-world dataset in long-range multi-class time-series classification.

The table below extends the experimental evaluation from Walker et al. (2025). We now evaluate +10 sequence models on four FSA emulation tasks. Notably, our model sets the state-of-the-art performance in this setup with no hyperparameter tuning.

The table below reports the performance of a variety of parallelizable sequence models on time-series classification with sequences of lengths up to over 17k. S6 is the previous best-performing evaluated SSM, with a more extensive overview of SSM results provided in Walker et al. (2024).

As we can see, our method also achieves a new state-of-the-art among the evaluated SSMs on this dataset (taken from Walker et al. 2025), outperforming approaches that rely on controlled differential equation (CDE) architectures.

Clarifying our statements on the scaling of related methods:

The reviewer is correct to point out that RWKV-7 and related methods have indeed been scaled to larger experimental setups. Our statement on unfavorable scaling preventing large-scale training does not refer to the family of matrix-valued SSMs that adopt a diagonal plus low-rank (DPLR) transition matrix structure, such as RWKV-7. The comment rather refers to unstructured transition matrices from Terzic et al. (2025), which do enable the emulation of any N-state FSA using a single layer with state size N, yet are prohibitively expensive in a larger setting. We have made sure to explicitly state these facts in a revised version of the paper.

We additionally acknowledge that our submission did not provide sufficient theoretical detail on RWKV-7 and related DPLR methods. We have referred to previous work which provides bounds on the model size necessary to emulate FSAs using such approaches, yet we have not clearly stated what those bounds are and how they compare to ours.

To clarify what we mean by the statement that DPLR-based methods are not expressive enough to encode any N-state FSA with a reasonably-sized model, we turn to Theorem 4 from Grazzi, Siems et al. (2024), which exactly considers SSMs with diagonal plus low-rank transition matrices. We paraphrase the theorem to focus on the relevant aspects of it as follows:

SSMs with state transition matrices that are repeated products of Generalized Householder matrices, each with eigenvalues in the range [−1,1], can recognize any regular language. In particular, every FSA A=(Σ,Q,q0,δ) can be implemented by a finite precision SSM with s≤2^|Q|layers, where the state is of size n≤|Q| and the number of Generalized Householder factors in the transition matrix is p≤s.

In this construction, the number of layers is exponential in the number of states, as is the number of factors which build up the transition matrix. In contrast, our model allows the emulation of any N-state FSA using one layer and a state size of N, where the size of the readout is linear in the number of automaton states.

In a separate result, Peng, Zhang, Goldstein et al. (2025) show that the RWKV-7 model is able to emulate an arbitrary FSA using four layers. However, the provided construction relies on MLP layers which implement lookup tables whose sizes are exponential in the size of the state set, which becomes infeasible for even moderately large FSAs.

We have now added a paragraph to the paper that contextualizes our theoretical results through comparison with these previous theoretical results on model expressivity, and thank the reviewer for bringing this up.

Clarifying the universal approximation property of diagonal SSMs

The referenced result on universal function approximation assumes a mapping from a sufficiently regular fixed-length sequence to another fixed-length sequence. The setting of tokens living in a finite set, such as what we have, is drastically different and not discussed in their work. In addition, the cited work makes no assumption on the numerical precision used by the SSM, and the latent dimension of the SSM and the subsequent MLP scale with the length of the sequence. As such, the framework and the assumptions used by the authors are vastly different from the empirically validated circuit-complexity bounds presented in Merill et al. (2024), which form the background of our work.

Questions:

1. We indeed refer to accurate linear classification of the SSM hidden states into the classes defined by the automaton states. In our tasks, the input to the SSM is a sequence of one-hot encoded entries from the FSA vocabulary, and the model, with a linear layer at the output, is trained to predict the automaton state given such a sequence of inputs.

2. The number of possible values that the entries of the hidden state can assume is very large, essentially limited only by numerical representation capacity. An interesting approach towards answering this question may make use of methods for automatic extraction of FSAs from the model, such as e.g. the one proposed in Weiss et al. (2019).

3. In this task, the states are fully defined by the underlying automaton. The difference between this task and the one we presented in Section 4.2 lies in the encoding of the inputs. Whereas the inputs in the "synthetic state tracking" task were one-hot encodings of a discrete set of inputs, the inputs we consider in "natural-language state tracking" are LLM-generated embeddings of meaningful English sentences.

4. This result is indeed expected and is found elsewhere in literature, e.g., Soydan, Zubic et al. (2025).

We thank the reviewer for their thoughtful review and hope that we were able to address all their questions and concerns.

We again appreciate the positive score.

With regards to the exponentially wide MLP hidden state, we are concretely referring to the final sentence in appendix D.2 of the arXiv version of the RWKV-7 paper, which states the following: "Our construction uses MLPs to implement lookup tables with sizes on the order of $|Σ|^{2n}$, which may require MLP layers that are exponentially wide in the number of states of the original DFA."

Thanks for your response! I appreciate the new empirical results and maintain my high score.

One clarification regarding the discussion of RWKV-7 in the rebuttal:

However, the provided construction relies on MLP layers which implement lookup tables whose sizes are exponential in the size of the state set

I don't think this is correct - A lookup table should be polynomial (not exponential) size in the number of states, since it maps each (state, token) pair to a state, and thus takes size roughly $\lvert Q \rvert^2 \cdot \lvert \Sigma \rvert$. However, it is possible for RNNs to do better than this, which is perhaps the intuition you're alluding to. See: https://arxiv.org/html/2402.15814v2

We thank the reviewer for their thoughtful feedback and are pleased they found the paper clear and appreciated our focus on the expressivity-efficiency trade-off in state-space models (SSMs), as well as the simplicity and theoretical depth of our proposed parameterization.

On the FSA focus and practical relevance

We agree that real-world applicability is critical and appreciate the concern about whether FSA-based analysis generalizes beyond synthetic settings. Our core contribution is a state-space model that can exactly emulate arbitrary finite-state automata (FSA) while maintaining the parallel scan efficiency of diagonal SSMs. This property significantly enhances the expressive power of parallel sequence models—an ability particularly relevant in domains requiring structured reasoning, such as code generation or formal language understanding.

Although large-scale NLP and vision tasks were beyond the scope of this submission, we note growing evidence that architectural expressivity grounded in formal structures like FSAs does translate into practical benefits. For example, Grazzi et al. (ICLR 2025) demonstrate that making DeltaNet—a matrix-valued SSM—more expressive in terms of its ability to emulate finite-state automata directly translates to improved performance on tasks like code generation and mathematical problem solving. Hu et al. (ACL 2025) similarly show that pretraining on formal languages enhances downstream NLP task performance. These works reinforce our belief that FSA emulation capabilities can be a building block for success in broader domains.

New real-world benchmarks

In our work, we introduce the first state-space model capable of perfectly tracking the states of finite-state automata (FSA) while maintaining the parallel scan cost at the same asymptotic level as diagonal state-space models. Through experiments on the Long Range Arena (LRA), we already observed practical benefits on text and image classification.

To address the reviewer's helpful suggestion for more practical evaluations, we added experiments on six multiclass time-series classification datasets (table below), covering long sequences (up to 17,000 time steps). This benchmark is widely used in real-world settings like healthcare, and it can be evaluated within the rebuttal window. To obtain the results, we ran a grid search over learning rates in {1e-5, 1e-4, 1e-3} and number of layers in {1, 2}. The state dimensionality was fixed to 16, and the embedding dimensionality was fixed to 64. All other hyperparameters were re-used from the referenced paper (Walker et al. 2025). Notably, this hyperparameter grid is significantly less extensive than the one used to evaluate the baseline methods, see Walker et al. (2024), yet we still achieve state-of-the-art results among the evaluated methods. S6 is the previous best-performing evaluated non-CDE SSM, with a more extensive overview of SSM results including LRU, S5, and Mamba provided in Walker et al. (2024).

PD-SSMs achieve a new state-of-the-art among the evaluated state space models on this dataset (taken from Walker et al. 2025), notably even without building on controlled differential equation (CDE) architectures.

These findings highlight a growing consensus around the practical benefits of expressive architectures, and they motivate our ongoing work. While the LRA benchmark offers valuable insight, we acknowledge its somewhat synthetic nature. Our newly added time-series experiments help address this by demonstrating real-world relevance. A comprehensive evaluation of our model on large-scale NLP tasks remains an exciting direction for future work.

Extended automaton emulation results

We also significantly expanded our comparison on FSA emulation tasks. We now include over 10 recent baselines across four automaton challenges—two more than in the original submission. Concretely, we extended the Table from Walker et al. (2025), by evaluating our proposed model on the full set of state tracking tasks from Deletang et al. (2023). We fully conform to their experimental setup which adopts a two-layer architecture and uses a single, fixed set of hyperparameters on all tasks. While the authors used the best results obtained with state size 128 or 512, we fixed the state size to 128. A summary of central results is reported in the table below. Notably, our model sets the state-of-the-art performance with no hyperparameter tuning.

Conclusion

While our focus is on improving the theoretical expressiveness of SSMs via FSA emulation, we have taken concrete steps to show this expressiveness translates into real-world benefits on long-range, structured sequence tasks. A comprehensive evaluation on large-scale NLP or vision datasets remains an exciting direction for future work, and we appreciate the reviewer’s encouragement in that direction.