Fixed-Point RNNs: Interpolating from Diagonal to Dense

2025/03/13 by Movahedi, Sajad, Sarnthein, Felix, Cirone, Nicola Muca +1 · 2 citations
#FOS: Computer and information sciences #Machine Learning (cs.LG)

paper · doi:10.48550/arxiv.2503.10799

Abstract

Linear recurrent neural networks (RNNs) and state-space models (SSMs) such as Mamba have become promising alternatives to softmax-attention as sequence mixing layers in Transformer architectures. Current models, however, do not exhibit the full state-tracking expressivity of RNNs because they rely on channel-wise (i.e. diagonal) sequence mixing. In this paper, we investigate parameterizations of a large class of dense linear RNNs as fixed-points of parallelizable diagonal linear RNNs. The resulting models can naturally trade expressivity for efficiency at a fixed number of parameters and achieve state-of-the-art results on the state-tracking benchmarks A₅ and S₅, while matching performance on copying and other tasks.

Cited by

• Structured Sparse Transition Matrices to Enable State Tracking in State-Space Models
• Learning State-Tracking from Code Using Linear RNNs

Related

• M²RNN: Non-Linear RNNs with Matrix-Valued States for Scalable Language Modeling
• DeltaProduct: Improving State-Tracking in Linear RNNs via Householder Products
• RWKV: Reinventing RNNs for the Transformer Era
• ParaRNN: Unlocking Parallel Training of Nonlinear RNNs for Large Language Models
• R-Transformer: Recurrent Neural Network Enhanced Transformer