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

Abstract: Transformers are highly parallel but are limited to computations in the TC$^0$ complexity class, excluding tasks such as entity tracking and code execution that provably require greater expressive power. Motivated by this limitation, we revisit non-linear Recurrent Neural Networks (RNNs) for language modeling and introduce Matrix-to-Matrix RNN (M$^2$RNN): an architecture with matrix-valued hidden states and expressive non-linear state transitions. We demonstrate that the language modeling performance of non-linear RNNs is limited by their state size. We also demonstrate how the state size expansion mechanism enables efficient use of tensor cores. Empirically, M$^2$RNN achieves perfect state tracking generalization at sequence lengths not seen during training. These benefits also translate to large-scale language modeling. In hybrid settings that interleave recurrent layers with attention, Hybrid M$^2$RNN outperforms equivalent Gated DeltaNet hybrids by $0.4$-$0.5$ perplexity points on a 7B MoE model, while using $3\times$ smaller state sizes for the recurrent layers. Notably, replacing even a single recurrent layer with M$^2$RNN in an existing hybrid architecture yields accuracy gains comparable to Hybrid M$^2$RNN with minimal impact on training throughput. Further, the Hybrid Gated DeltaNet models with a single M$^2$RNN layer also achieve superior long-context generalization, outperforming 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 LLMs.

• The paper introduces M²RNN, employing matrix-valued states to enhance non-linear recurrence and overcome limitations of traditional vector RNNs.
• It utilizes outer-product expansion and efficient gating to improve state tracking and achieve superior language modeling performance.
• M²RNN demonstrates hardware efficiency through optimized tiling and custom kernels, enabling competitive throughput in hybrid large language models.

Matrix-to-Matrix RNNs: Non-Linear RNNs with Matrix-Valued States for Scalable Language Modeling

Architectural Foundations and Motivations

Modern LLMs primarily leverage attention-based architectures, notably Transformers, due to sequence parallelism, hardware efficiency, and robust in-context retrieval capabilities. Quadratic compute and memory overheads drive ongoing research into more efficient alternatives—state space models (SSMs), linear recurrent architectures (RNNs), and their hybridizations. However, linear RNNs exhibit expressivity limitations in hard state-tracking tasks and poor in-context retrieval performance. Conversely, classical non-linear RNNs (e.g., LSTM, GRU) have superior state tracking but underperform in language modeling and retrieval benchmarks due to constrained state sizes and poor training efficiency. The paper proposes the Matrix-to-Matrix RNN (M$^2$RNN), a non-linear recurrent architecture employing matrix-valued states and outer-product expansions, combining the expressivity of non-linear RNNs with scalable state size and efficient hardware tiling.

Matrix-to-Matrix RNN Layer Design

M$^2$RNN builds upon the outer product state expansion found in linear attention and SSMs, enabling much larger state representations without quadratic parameter growth. The architecture incorporates a forget gate, analogously to LSTM/GRU, but with a scalar per-head formulation independent of the recurrent state, enabling efficient parallel computation across the sequence.

The core update equations:

• State update: $$\mathbf{Z}_t = \tanh(\mathbf{H}_{t-1}\mathbf{W} + k_t v_t^\top)$$,
• Hidden state: $$\mathbf{H}_t = f_t \mathbf{H}_{t-1} + (1 - f_t)\mathbf{Z}_t$$,
• Output: $y_t = \mathbf{H}_t^\top q_t + w_r \odot v_t$.

The M$^2$RNN forget gate is parameterized as $$f_{t,n} = \psi(x_t) = [1 + e^{x_t + \beta_n}]^{-\alpha_n}$$, with learnable $\alpha_n, \beta_n$ controlling decay and saturation (Figure 1).

Figure 1: Forget gate ($f_t$) behavior as a function of input ($x_t$) for different $\alpha_n, \beta_n$ parameters.

The sequence mixing block incorporates short depthwise convolutions for key, query, value, and output gating, RMS normalization, and essential residual connections for stable training. Orthogonal and identity initializations are found superior for transition matrices, aligning with prior non-linear RNN literature.

Figure 2: Matrix-to-Matrix RNN block replaces attention, integrating MLP and RMSNorm modules akin to Transformers.

Computational Expressivity and State Tracking

Prior theoretical results anchor the expressivity of saturated non-linear RNNs in their ability to simulate DFAs and recognize regular languages. M$^2$RNN inherits this power: the matrix-valued recurrence can subsume vector-valued non-linear RNNs, thus achieving the full regular language class, which is unattainable by linear RNNs with diagonal/input-independent transitions. Empirical results on the $S_3$ permutation group problem demonstrate robust length generalization—M$^2$RNN and GRU generalize accurately well beyond training context lengths, while linear competitors degrade rapidly.

Figure 3: Accuracy on $S_3$ permutation state-tracking task: M$^2$RNN and GRU achieve near-perfect accuracy across context lengths.

Language Modeling, Retrieval, and Long-Context Performance

State size is the principal driver of recurrent language modeling efficacy, not merely the presence of non-linearity. M$^2$RNN's matrix expansion vastly increases capacity over vector RNNs, leading to substantial performance improvements across benchmarks, even under fixed parameter budgets. Ablation studies confirm dramatic gains: M$^2$RNN outperforms vector RNN and GRU (with equivalent parameter count) on WikiText and LAMBADA perplexity by 10–280 points.

Hybrid models—interleaving M$^2$RNN with attention and/or linear RNNs—yield best-in-class results, outperforming Transformer++, Mamba-2, and Gated DeltaNet hybrids. Even a single M$^2$RNN layer in these hybrids produces notable accuracy improvements with negligible throughput impact.

Zero-Shot In-Context Retrieval and Long-Context Generalization

Zero-shot evaluations on RULER NIAH variants demonstrate enhanced retrieval accuracy and generalization across context lengths well beyond training boundaries, particularly in hybrid M$^2$RNN models, compared to both linear and attention-based baselines.

Figure 4: M$^2$RNN and Gated DeltaNet zero-shot retrieval performance at 410M scale (vertical: training context length).

Figure 5: Zero-shot retrieval at 7B MoE scale for M$^2$RNN and Gated DeltaNet (vertical: training context).

In real-world retrieval tasks and LongBench summarization/coding/few-shot benchmarks, M$^2$RNN hybrids consistently outperform linear and Transformer architectures, with hybrid combinations yielding up to 8 points higher accuracy in long-context evaluation.

Hardware Efficiency and Systems Optimizations

Matrix-valued state expansion allows optimal tiling over batch and head, utilizing tensor cores efficiently without wasteful padding along the batch dimension, unlike vector RNNs (FlashRNN). The multi-value attention formulation maximizes state size under fixed parameter budgets. Custom kernels developed in Triton, with ongoing CUTLASS optimization, provide scalable linear-time training. M$^2$RNN enables both topology-aware and topology-independent tensor parallelism strategies; the former incurs no extra communication, while the latter is parameter-count preserving and requires lightweight synchronization (Figures 4 and 5).

Figure 6: Topology-aware TP—independent RMSNorm weights per GPU.

Figure 7: Topology-independent TP—RMSNorm weights are sharded and synchronized.

Training throughput remains competitive: inserting M$^2$RNN layers into hybrids results in only a minor throughput degradation (within 6%) compared to linear RNN baselines, making the architecture practical for production scale.

Figure 8: Training throughput for 7B MoE configuration (linear: context length); M$^2$RNN throughput is near hybrid baselines.

Conclusion

M$^2$RNN establishes matrix-valued non-linear recurrence as a potent and scalable component for LLMs. It resolves key limitations of both linear and classical non-linear RNNs—expressivity, retrieval, state tracking, and hardware utilization. Hybrid architectures with sparing use of M$^2$RNN layers yield substantial accuracy and generalization improvements with minimal computational overhead. Distributed training is supported via robust TP strategies, enabling flexible deployment across heterogeneous hardware.

Future work should focus on further kernel optimization to lower constant cost, and scaling M$^2$RNN to larger parameter regimes and longer contexts, solidifying its potential as a foundational sequence modeling block for next-generation LLMs.