Efficient Transition Matrices to Enable State Tracking in State-Space Models
Aleksandar Terzic, Nicolas Menet, Michael Hersche, Thomas Hofmann, Abbas Rahimi
Most modern state-space models (SSMs) rely on diagonal transition matrices, which ensure high computational efficiency but pose significant restrictions on their state-tracking abilities. In contrast, unstructured transition matrices allow SSMs to track the states of arbitrary finite-state automata (FSAs); however, at a prohibitively high compute and memory cost even for moderate state sizes. To reconcile the advantages of the two approaches, we propose an efficient parametrization of the transition matrices in selective SSMs that enables state-tracking at the level of unstructured SSMs while keeping the computational cost comparable to that of diagonal SSMs. Our PD-SSM parametrizes the transition matrix as the product of a column one-hot matrix ($P$) and a complex-valued diagonal matrix ($D$).  Consequently, the computational cost of parallel scans over PD matrices scales linearly with the state size. Theoretically, PD-SSM can learn to track any FSA and reaches the lower bound on state dimension necessary to do so. Empirically, our model learns to track the states of FSAs of various complexities, even if the state transitions are obscured by variable-length sentences in natural language, by integrating it into a hybrid Transformer-SSM architecture.