Title: Structured Sparse Transition Matrices to Enable State Tracking in State-Space Models
Venue: NeurIPS 2025
Accepted: yes
Keywords: State-Space Models

Expressiveness

Efficiency

Matrix Parametrisation

State-Tracking

Finite-State Automata
Abstract: Modern state-space models (SSMs) often utilize structured transition matrices
which enable efficient computation but pose restrictions on the model’s expressivity,
as measured in terms of the ability to emulate finite-state automata (FSA). While
unstructured transition matrices are optimal in terms of expressivity, they come
at a prohibitively high compute and memory cost, even for moderate state sizes.
We propose a structured sparse parametrization of transition matrices in SSMs
that enables FSA state tracking with provably optimal state size and depth, while
keeping the computational cost of the recurrence comparable to that of diagonal
SSMs. Our method, \emph{PD-SSM}, parametrizes the transition matrix as the product
of a column one-hot matrix ($P$) and a complex-valued diagonal matrix ($D$). As
a result, the computational cost of parallel scans scales linearly with the state
size. Theoretically, the model is BIBO-stable and can emulate any $N$-state FSA
with one layer of dimension $N$ and a linear readout of size $N ×N$, significantly
improving on all current structured SSM guarantees. Experimentally, the model
significantly outperforms a wide collection of modern SSM variants on various FSA
state tracking tasks. On multivariate time-series classification, it outperforms neural
controlled differential equations, a paradigm explicitly built for time-series analysis.
Finally, we integrate PD-SSM into a hybrid Transformer-SSM architecture and
demonstrate that the model can effectively track the states of a complex FSA in
which transitions are encoded into sets of variable-length English sentences. The
code is available at https://github.com/IBM/expressive-sparse-state-space-model.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes PD-SSM, a structured state-space model with a novel input-dependent structured state-transition matrix, $A(u_t)=P(u_t)D(u_t)$, where $P(u_t)$ is a binary matrix with one non-zero element per column and $D(u_t)$ is a complex diagonal matrix. This structure fits within a large family of existing structures, including the diagonal structure of Mamba, the dense structure of IDS4, Linear NCDEs, and SD-SSM, the block-diagonal structure of block-diagonal input dependent linear RNNs, and the diagonal-plus-low-rank structure of DeltaNet. \n\nThe paper shows that PD-SSM has the same asymptotic computational complexity as using diagonal transition matrices and can solve any finite state automata with a hidden dimension that scales linearly with the number of states. Additionally, PD-SSM empirically solves a number of state-tracking benchmarks and outperforms Mamba and S7 on a subset of the long range arena benchmark.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n\n- The paper is well written.\n\n- The PD structure is well motivated, both from a theoretical cost stand point and the analysis showing that any FSA can be represented in an efficient manner.\n\n- The benchmark datasets are suitable, and the PD structure clearly demonstrates an increase in expressivity over diagonal matrices.\n\n- The natural language state-tracking benchmark provides a good indication that the results on the regular language benchmarks translate into real-world language modelling performance.\n\n## Weaknesses\n\n- The paper limits the theoretical and empirical results to only diagonal, dense, and PD structures. However, as noted by the authors on lines 43 and 44, both block-diagonal (BD) and diagonal-plus-low-rank (DPLR) structures exist in the literature. Table 1 and the empirical results neglect these additional structures, and no arguments are made as to why the PD structure should be favoured over BD and DPLR.\n\n- Both Table 2 and Table 3 are currently incomplete. Table 2 is missing cycle navigation and even pairs, both of which SD-SSM solves perfectly, and Table 3 is missing the pathfinder datasets (which the authors note is due to an issue with their compute cluste). The paper would be significantly improved by completing the missing datasets and expanding the empirical results to include BD and DPLR structured matrices.\n\n- Given that computational cost is a main motivating factor for the PD structure, the empirical runtime of the model is not sufficiently discussed in the main body of the text. As shown in the Appendix, despite having the same asymptotic computational cost, going from diagonal to a PD structure can lead to an order of magnitude longer training time per step. \n\n- The main body claims $C_2\\times C_{30}$ and $D_{30}$ are trained on sequences up to length $40$, but the Appendix claims they were trained on sequences up to length $90$.\n\nQUESTIONS:\n## Questions\n\n- How were the hyperparameters chosen for each dataset?\n\n## Suggestions\n\n- As mentioned, a theoretical and empirical comparison to the already existing structures in the literature is critical for a paper proposing a new structure. This addition would significantly improve the paper. Given an appropriate demonstration of the benefits of the PD structure over BD and DPLR, I would raise my score.\n\n- A log scale on the y-axis of figure 4 would make the comparison with the diagonal structure easier. \n\n- After the submission deadline, [this paper](https://arxiv.org/abs/2505.17761) introduced two more possible structures for the transition matrix, and introduced a unifying framework for input-dependent structured state transition matrices. I mention it here as I believe including this paper in the discussion of alternative structures would be benficial. However, there is naturally no requirement to engage with it given that it appeared after submission.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses a fundamental limitation in State Space Models (SSMs): the trade-off between expressiveness and computational efficiency. While diagonal SSM blocks are computationally tractable, they suffer from limited expressive power. Conversely, unstructured transition matrices offer full expressiveness but at prohibitive computational costs. The authors propose a novel middle-ground approach using structured matrices with a PD parameterization, where P is a one-hot matrix computed via a complex procedure and D is diagonal. This parameterization maintains efficiency while forming a monoid under multiplication, enabling parallel scan algorithms.\n\nThe paper establishes theoretical foundations by measuring expressive power through the class of Finite State Automata that SSMs can emulate (without post-SSM nonlinearity). \nTo address the non-trivial computations required for P, which present differentiation challenges, the authors provide implementation strategies and circumvention methods.\nThe theoretical claims are empirically validated through testing on state tracking problems.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n\n1. **Clear Problem Formulation**: The paper tackles one of the most important open problems in the SSM literature with a well-motivated theoretical framework.\n2. **Theoretical Rigor**: The proposed architecture is theoretically grounded and addresses fundamental expressiveness limitations of diagonal SSMs.\n3. **Practical Considerations**: The authors provide detailed analysis of training difficulties and propose concrete solutions for implementation challenges.\n4. **Novel Contribution**: The PD parameterization represents a principled approach to balancing expressiveness and computational efficiency.\n\n## Weaknesses\n\n1. **Insufficient Experimental Validation**: The experimental section has significant gaps that undermine the paper's contributions:\n    - **Missing experimental setup details**: Critical information about hyperparameter selection, training procedures, and experimental configurations is absent. Specifically: How were hyperparameters chosen? Were hyperparameter sweeps conducted? How was fair comparison between different models ensured? The learning rate appears constant across S2 and S3, and hyperparameter choices for Mamba are not visible.\n    - **Lack of comparison with recent models**: No benchmarking against important recent SSM variants including Mamba2, DeltaNet, GatedDnet, and RWKV-7, which are essential baselines for evaluating the proposed method\n    - **Limited scope of experiments**: While the focus on toy models and limited language modeling experiments may be acceptable for an initial architecture introduction (following similar literature precedents), the current experimental validation is insufficient to demonstrate the method's practical advantages\n2. **Presentation Issues**: Appendix A (Algebra and FSA) contains poorly organized algebraic content with numerous inconsistencies and unclear definitions that require substantial polishing for clarity. For example, Appendix A2 inconsistently switches between Z/nZ, Zn, and Cn to denote the same group. The mathematical exposition suffers from such notation inconsistencies, missing quantifications, and terminological confusion that hinder comprehension.\n\nQUESTIONS:\n## Minor Issues and Suggestions\n\n- **Lines 34-36**: The universal approximation claim is correct and actually supports your arguments even more strongly. Consider clarifying that in the cited result, the approximation is actually achieved by the neural network after the SSM, while the SSM itself primarily reconstructs the input. Since your approach focuses on the SSM without post-processing nonlinearity, making this distinction explicit would better highlight the significance of your theoretical contribution.\n- **Lines 79-82**: The invertibility of A, while implied by perturbation arguments, should be stated more explicitly.\n- **Lines 86-87**: (Pedantic comment) in addition to the stated condition the matrices need to be all diagonalizable by themselves as well.\n- **Lines 144-148**: The choice of P through post-processed mixture seems convoluted. Please elaborate on:\n    - Whether this was the original design choice\n    - Alternative approaches considered (e.g., input-dependent low-rank followed by hardmax)\n    - Comparative analysis of different parameterization strategies\n- **Line 165**: (Pedantic comment) Property 2 alone is insufficient to justify associative scan usage. A remark after line 157 demonstrating that matrices in H can be written as PD would improve rigor.\n- **Section 4.2**: Consider expanding the analysis with additional comparisons that would strengthen the paper's impact. I assume only SSM blocks are compared here - it would be valuable to explore whether including MLPs afterwards changes the relative performance. Additionally, reporting how many layers Mamba requires to reach 100% accuracy would provide an extremely compelling addition to demonstrate the efficiency gains of your approach.\n- **Lines 227-231**: The architecture description needs expansion, potentially in an appendix.\n\n\n### Appendix Issues \n\n- **Notation Consistency**: Inconsistent use of Z/nZ, Zn, and C_n for the same group throughout\n- **Lines 40-41**: T(A) notation unclear; specify the second element of delta in Q and define σ₁,...,T\n- **Lines 53-57**: Quantify all mathematical objects\n- **Lines 66-67**: Inconsistent switching between \"group\" and \"automaton\" terminology\n- **Line 141**: Verify if C₄/C₂ × C₄ should be (C₂ × C₄)/C₄; inconsistent Z₄ and C₄ usage\n- **Lines 144-146**: Notation {e}/C₄, {e}/Z₄, {e}/G should be reversed\n- **Line 173**: Reference Definition 12 when first using \"Uniform\" in Definition 10\n- **Line 181**: There appears to be an inconsistency regarding fan-in restrictions. The text states that \"NC1 allows for logarithmic depth, but restricts the fan-in of the gates to 2,\" while the definition only mentions that fan-in is bounded. Please clarify this discrepancy.\n- **Line 231**: Use standard notation instead of ∂⁺ for total derivative\n- **Line 267**: In the equation \"xₖ = Aₖ xₖ₋₁ + bₖ\", the subscript of A and b should be \"k-1\" for consistency with both the proof and the results presented below.\n- **Line 273**: Quantify P appropriately\n- **Equation below 269**: (Pedantic comment) Specify that α and β are not all zero in summation equalities\n\n\n\n## Recommendation\n\nThis paper addresses an important problem in the SSM literature with a theoretically sound and well-motivated approach. The core contribution—finding a middle ground between diagonal SSMs and unstructured matrices through the PD parameterization—is valuable and tackles a fundamental limitation in the field. However, the current submission has notable gaps that prevent clear acceptance at this stage.\n\nI would genuinely like to accept this paper as it shows strong theoretical promise and addresses a central problem in the SSM literature. However, it requires multiple small but targeted improvements: the appendix and proofs need polishing, the experiments need improvement, and experimental details must be included to ensure fair comparisons with competing models. With proper revision addressing these specific issues, this work would clearly merit acceptance.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 2, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nBuilding on prior work, this paper proposes a new parameterization of transition matrices for state-space models that achieve greater expressivity for state tracking/automata simulation with lower memory. Concretely:\n1. The paper proposes parameterizing transition matrices as column one-hot times diagonal (PD) form\n2. Theoretically, the paper shows this parameterization unlocks expressivity for state tracking with linear memory in the convolutional form, which is reduced compared to a dense parameterization\n3. Empirically, the authors show this parameterization attains better runtime, better performance on fully synthetic and natural-language state tracking, and strong Long Range Arena performance relative to other time-variant SSMs.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n \n1. The paper is well motivated, clear, and rigorous\n2. The proposed PD parameterization is a compelling method to achieve greater expressivity while minimizing memory overhead.\n3. The empirical evaluation is fairly extensive, spanning synthetic state tracking evaluations and more natural tasks\n\n## Weaknesses\n\nIn addition to the existing empirical evaluations, the paper could benefit from reporting language modeling performance (e.g., cross entropy or bits/byte) on a standard dataset.\n\nThe following claim is made regarding other approaches for other methods that balance efficiency and expressivity with SSM transition matrix parameterizations:\n\n> However, these approaches either scale unfavorably and thus prohibit large-scale training, or are not expressive enough to encode arbitrary N-state FSA with a reasonably-sized model such as a single layer with state dimensionality N.\n\nThis claim is made without rigorous support: you have not shown in what sense other methods like RWKV-7 scale poorly. E.g., RWKV-7 also uses a constrained parameterization where each matrix has rank 1. The products of those matrices may have higher rank, which suggests that SCAN might take more memory, but you have not shown this rigorously, so this sentence comes off as too strong. You should either elaborate a bit about the other methods you are comparing against or hedge this sentence appropriately.\n\nRegarding scalability, some of these other methods (e.g., RWKV-7) have scaled up beyond the experiments in this paper. So while they might \"scale unfavorably\" compared to your approach in principle, they have indeed been scaled up.\n\n> Theoretically, time-invariant (fixed) diagonal transition matrices suffice to guarantee universal approximation to any regular causal sequence-to sequence map (Orvieto et al. 2024).\n\nI don't understand this sentence in the intro. Clarify this and explain how it is consistent with the reduced expressivity of diagonal transition matrices for state tracking.\n\nQUESTIONS:\n“a finite-state automaton with N states that cannot be emulated by any single-layer SSM” What does emulation precisely mean? An SSM can classify the output state of the SSM using a linear layer?\n\nWhat about the other direction: Given a single layer of your PD SSMs, if we want to capture its output with a WFA, how many states would we need?\n\nRegarding your natural language state tracking task: what does a state correspond to? a physical location? Show a full example input.\n\nIn 4.4, why are time-invariant models better than time-invariant models? Is this expected/found elsewhere in the literature?\n\nLIMITATIONS:\nThe major limitation I see the lack of larger scale language modeling experiments: if these were present and validated that the architecture could be scaled up (e.g., with an appropriate hardware kernel) and achieve good performance, I would consider raising my score. I would like to see the cross entropy attained by an SSM with this parameterization or even scaling laws or measures of compute efficiency relative to other architectures. However, I understand these experiments could require significant resources, and could reasonably be relegated to future work motivated by this paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-07-01"}

comments: SUMMARY:
This paper proposes a novel parameterization of the transition matrices in selective state-space models (SSMs). Specifically, the authors replace diagonal transition matrices with a product of a column one-hot matrix (P) and a complex-valued diagonal matrix (D). Through both theoretical and empirical analysis, they argue that this new parameterization effectively models arbitrary finite-state automata (FSA). They also demonstrate improved efficiency by comparing the empirical runtime of their method against SD-SSM and Diagonal SSM baselines.

STRENGTHS AND WEAKNESSES:
**Strengths**

1. The paper is clearly written and easy to follow.
2. It presents a strong theoretical foundation, supported by empirical improvements in both performance and efficiency.
3. The work addresses an important problem in the expressivity and computational efficiency of SSMs.
4. The proposed changes are simple to implement yet yield meaningful benefits.

**Weaknesses**

The analysis is centered entirely around finite-state automata. While FSAs are a useful proxy, state-space models are increasingly applied in real-world tasks such as natural language processing and computer vision. It remains unclear whether the proposed parameterization would translate into practical benefits in these broader domains.

QUESTIONS:
How does the FSA-based analysis inform real-world applications? Do FSAs share structural or behavioral similarities with typical use cases of SSMs in NLP or vision tasks?

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-03