Saved in:
Bibliographic Details
Main Authors: Grazzi, Riccardo, Siems, Julien, Zela, Arber, Franke, Jörg K. H., Hutter, Frank, Pontil, Massimiliano
Format: Preprint
Published: 2024
Subjects:
Online Access:https://arxiv.org/abs/2411.12537
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866915203345547264
author Grazzi, Riccardo
Siems, Julien
Zela, Arber
Franke, Jörg K. H.
Hutter, Frank
Pontil, Massimiliano
author_facet Grazzi, Riccardo
Siems, Julien
Zela, Arber
Franke, Jörg K. H.
Hutter, Frank
Pontil, Massimiliano
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2411_12537
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues
Grazzi, Riccardo
Siems, Julien
Zela, Arber
Franke, Jörg K. H.
Hutter, Frank
Pontil, Massimiliano
Machine Learning
Computation and Language
Formal Languages and Automata Theory
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.
title Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues
topic Machine Learning
Computation and Language
Formal Languages and Automata Theory
url https://arxiv.org/abs/2411.12537