M$^2$RNN: Non-Linear RNNs with Matrix-Valued States for Scalable Language Modeling

Transformers are limited to computations in the TC$^0$ complexity class, excluding tasks such as entity tracking and code execution that provably require greater expressivity. Motivated by this, we revisit non-linear Recurrent Neural Networks (RNNs) and introduce Matrix-to-Matrix RNN (M$^2$RNN): an architecture with matrix-valued hidden states and non-linear state transitions. We show that (i) the language modeling gap of non-linear RNNs is primarily a state-size gap, not a non-linearity penalty, and (ii) outer-product state expansion enables efficient tensor-core utilization. Empirically, M$^2$RNN achieves perfect state-tracking generalization at sequence lengths beyond training. These gains transfer to large-scale language modeling: at 7B MoE, Hybrid M$^2$RNN outperforms Hybrid Gated DeltaNetby $0.5$ perplexity points using $3\times$ smaller recurrent states, and replacing even a single recurrent layer with M$^2$RNN matches Hybrid M$^2$RNN accuracy with minimal throughput cost. Hybrid Gated DeltaNetwith a single M$^2$RNN layer also outperforms state-of-the-art hybrid linear-attention architectures by up to $8$ points on LongBench. Together, these results establish non-linear RNN layers as a compelling building block for efficient and scalable language models.