Research map / Architectures and efficiency / RNNs and normalization
M$^2$RNN: Non-Linear RNNs with Matrix-Valued States for Scalable Language Modeling
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces M²RNN, a non-linear RNN with matrix-valued states and an input-independent forget gate, plus Triton kernels and tensor-parallel strategies, showing competitive LM and retrieval results and gains when used as a few layers in hybrids.
Problem
Non-linear RNNs underperform linear RNNs in language modeling, long-context retrieval, and training efficiency due to small vector-valued states and sequential, poorly hardware-utilized computation, while linear RNNs lack state-tracking expressivity.
Approach
A non-linear RNN with matrix-valued state updated by tanh(H W + k v^T), gated by an input-only forget gate (so gates are parallelizable), with a multi-value outer-product state expansion that avoids padding waste on tensor cores; used sparingly within hybrid linear RNN/attention models.
Key results
- At 410M, Hybrid Gated DeltaNet + M2RNN-3 reaches 55.23 avg accuracy and 28.93 LAMBADA-style LMB PPL vs 54.04 for Hybrid Gated DeltaNet.
- At 7B MoE, Hybrid Gated DeltaNet + M2RNN-5 gets 12.85 Wikitext PPL and 31.2 LongBench avg vs 23.1 for Hybrid Gated DeltaNet.
- Standalone M2RNN (410M) matches Mamba-2 average accuracy (54.24 vs 54.30) and reaches perfect S_3 state tracking up to length 512.
Surprising findings
- Adding a single M2RNN layer to a hybrid matches the accuracy of a full hybrid M2RNN model, and gives throughput within 6% of Hybrid Gated DeltaNet at 16k context.
- Gated DeltaProduct with two Householder matrices, despite theoretical S_3 expressivity, fails to generalize beyond training length 256 while GRU and M2RNN generalize to 512.
Limitations
- Standalone M2RNN is slower than linear baselines due to sequential non-parallelizable recurrence and an unoptimized Triton backward kernel.
- Homogeneous M2RNN trails Gated DeltaNet on Wikitext perplexity, attributed to a 3x smaller state size.
- Adding M2RNN layers degrades S-NIAH-1 accuracy in hybrids.
Experimental setup
Datasets: Nemotron-CC-v2
Benchmarks: LAMBADA HellaSwag PIQA ARC WinoGrande BoolQ OpenBookQA COPA SciQ WikiText RULER SQuAD
Metrics: perplexity accuracy F1
Baselines: Mamba-2 Gated DeltaNet RNN GRU Transformer++ Gated DeltaProduct Hybrid Mamba-2 Hybrid Gated DeltaNet
Builds on: Mamba-2 Gated DeltaNet FlashRNN GRU Transformer++
Closest papers by meaning
- FlashRNN: Unlocking Parallel Training of Nonlinear RNNs for Large Language Models (ICLR 2026 · similarity 0.90)
- Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues (ICLR 2025 · similarity 0.89)
- Fast-Slow Recurrent Neural Networks (NeurIPS 2017 · similarity 0.89)
- Implicit Language Models are RNNs: Balancing Parallelization and Expressivity (ICML 2025 · similarity 0.88)
- A Tensor Decomposition Perspective on Second-order RNNs (ICML 2024 · similarity 0.88)
- DeltaProduct: Improving State-Tracking in Linear RNNs via Householder Products (NeurIPS 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.
