Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues

TL;DR

This work addresses the state-tracking limitations of linear recurrent networks by extending the allowed eigenvalue range of state-transition matrices from [0,1] to [-1,1], enabling parity and broader regular-language recognition. It provides theoretical results showing that positive-eigenvalue LRNNs cannot solve parity and that non-diagonal structures and GH products expand expressivity, including sufficiency for any regular language under finite precision. Empirically, extending to negative eigenvalues improves parity and modular counting performance for Mamba and DeltaNet, and allows scalable pretraining up to 1.3B parameters with competitive language modeling and promising code/math task improvements, especially for DeltaNet. The findings highlight the critical role of non-diagonal state transitions in state-tracking and suggest a viable path for scaling LRNNs to long sequences and complex tasks without Transformer-style quadratic costs.

Abstract

Linear Recurrent Neural Networks (LRNNs) such as Mamba, RWKV, GLA, mLSTM, and DeltaNet have emerged as efficient alternatives to Transformers for long sequences. However, both Transformers and LRNNs struggle to perform state-tracking, which may impair performance in tasks such as code evaluation. In one forward pass, current architectures are unable to solve even parity, the simplest state-tracking task, which non-linear RNNs can handle effectively. Recently, Sarrof et al. (2024) demonstrated that the failure of LRNNs like Mamba to solve parity stems from restricting the value range of their diagonal state-transition matrices to $[0, 1]$ and that incorporating negative values can resolve this issue. We extend this result to non-diagonal LRNNs such as DeltaNet. We prove that finite precision LRNNs with state-transition matrices having only positive eigenvalues cannot solve parity, while non-triangular matrices are needed to count modulo $3$. Notably, we also prove that LRNNs can learn any regular language when their state-transition matrices are products of identity minus vector outer product matrices, each with eigenvalues in the range $[-1, 1]$. Our experiments confirm that extending the eigenvalue range of Mamba and DeltaNet to include negative values not only enables them to solve parity but consistently improves their performance on state-tracking tasks. We also show that state-tracking enabled LRNNs can be pretrained stably and efficiently at scale (1.3B parameters), achieving competitive performance on language modeling and showing promise on code and math tasks.

Figures (14)

• Figure 1: Extending the eigenvalue range of the state transition matrices of diagonal LRNNs improves performance from random guessing (range $[0,1]$) to perfect score (range $[-1,1]$) on learning parity. Trained on sequences up to length 40; Tested on lengths 40–256 (3 seeds).
• Figure 2: Parity requires negative eigenvalues. States of one-layer LRNNs with the sequence $1111\ldots$ as input. If the eigenvalues of $\mathbf{A}(1)$ are nonnegative, the states either diverge or converge monotonically, and so, for large enough $t$ and in finite precision, cannot be distinguished. In contrast, the LRNN with $a(1) = -1$ alternates between two states like the parity automaton.
• Figure 3: A permutation of $k$ elements is also a composition of at most $k{-}1$ swaps. This maps to a product of $k{-}1$ Hoseholders, each representing a swap. Illustrated for $k=3$. $\mathbf{v}_1^\top {=} \left(\frac{1}{\sqrt{2}}, -\frac{1}{\sqrt{2}}, 0\right)$, $\mathbf{v}_2^\top {=} \left(0, \frac{1}{\sqrt{2}}, -\frac{1}{\sqrt{2}}\right)$.
• Figure 4: Sequence accuracy for varying sequence lengths on $S_5$ after 100 epochs of training. We report the best of 3 seeds for each method (in \ref{['fig:s5_large']} we report all seeds). The dashed vertical line indicates the sequence length used during training (32 except for the third plot from the left where it is 64). Each method is labeled with name, eigenvalue range, and number of layers. The dashed vertical line indicates the sequence length used during training. "Full matrix simple" is a one-layer baseline where the state update matrices are full and we have no control over the eigenvalue range.
• Figure 5: Performance vs sequence length of DeltaNet variants (340M (top) and 1.3B (bottom) parameters) on four datasets. DeltaNet with eigenvalue range $[-1, 1]$ improves perplexity in coding and math compared to the $[0, 1]$ baseline. Dashed vertical line at training context length (2048).

Theorems & Definitions (11)

• Theorem 1: Parity
• Theorem 2: Modular Counting
• Proposition 1: Expressivity of products of GH matrices
• Theorem 3
• Theorem 4
• Lemma 1
• proof
• Definition 1: Cascade product
• Theorem 5: Krohn-Rhodes, Theorem 4 in maler1994cascaded
• Theorem 6: Modular addition with reflections