Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues
Journal:
arXiv
Published Date:
Nov 19, 2024
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.
