[ICLR 2025] M2 RNN: Breaking the Expressivity Ceiling of Linear RNNs with Matrix-Valued States

TL;DR

The reign of purely linear RNNs (like Mamba and DeltaNet) might be facing a challenge. While linear models are fast, they are mathematically limited in solving complex state-tracking problems. M2 RNN (Matrix-to-Matrix RNN) reintroduces non-linear transitions (tanh) but scales them up using matrix-valued hidden states. The result? A model that captures complex logic better than Transformers in some cases, while remaining efficient enough for 7B+ parameter production models.

The "TC0" Bottleneck: Why Linear Isn't Enough

Recent research has proven a sobering fact: Transformers and linear SSMs reside in the TC0 complexity class. This means they struggle with tasks like tracking entities in a long story or executing code (the $S_5$ permutation problem). Non-linear RNNs (like GRUs) are theoretically more expressive, but they have historically failed in the LLM era because:

1. State Size: Their "memory" (state) is too small (a simple vector).
2. Hardware Inefficiency: They don't utilize GPU Tensor Cores well, making them slow to train.

Methodology: The Matrix-to-Matrix Leap

M2 RNN solves the capacity problem by moving from a vector state $h_t$ to a matrix state $H_t$.

1. The Core Recurrence

Instead of a linear update, M2 RNN uses: $Z_t=anh(H_{t-1}W+k_t{v}_t^{o}p)$ $H_t=f_t{H}_{t-1}+(1-f_t)Z_t$

Here, $k_t{v}_t^{o}p$ is an outer product that writes new information into the matrix. The $anh$ provides the non-linear "reasoning" capability. Crucially, the forget gate ($f_t$) is independent of the previous state, allowing for better parallelism than traditional LSTMs.

2. Hardware-Aware Tiling

One of the paper's most impressive contributions is how it handles GPU kernels. Traditional "Fast RNN" implementations often pad the batch size to use Tensor Cores, leading to wasted FLOPs. M2 RNN's matrix-valued operation is inherently a GEMM (General Matrix Multiply) that fits perfectly onto NVIDIA Tensor Cores (using WGMMA instructions) regardless of the batch size.

Experiments: Quality and Scaling

The authors tested M2 RNN at two scales: 410M and 7B (Mixture-of-Experts).

SOTA Comparison

In "Hybrid" architectures (where most layers are recurrent but a few are Attention), M2 RNN shines. It beat both Mamba-2 and Gated DeltaNet across several benchmarks.

• Perplexity: Outperformed Gated DeltaNet by 0.4-0.5 points.
• Retrieval: On the RULER benchmark (Needle-in-a-Haystack), M2 RNN showed significantly better generalization to context lengths not seen during training.

Figure: M2 RNN and GRU (non-linear) generalize perfectly on the S3 state-tracking task, while linear models (DeltaNet) fail as the sequence length increases.

The "Power of One" Layer

Perhaps the most practical finding: you don't need a full M2 RNN model to see benefits. Replacing even a single layer of a Mamba model with an M2 RNN layer provides massive boosts in "LongBench" accuracy (up to 8 points) with only a 6% drop in training speed.

Critical Insight & Conclusion

The industry has spent two years trying to make RNNs "linear" to save on computation. This paper argues that we might have thrown the baby out with the bathwater. Non-linearity is the secret sauce of "reasoning" and state tracking. By using matrix-valued states, M2 RNN gives us the best of both worlds: the expressivity of LSTMs and the scalability of Mamba.

Future Impact

M2 RNN is a compelling candidate for the next generation of "Reasoning Models." As we push LLMs to do more complex logic and long-context code analysis, the "TC0" limit of current Transformers will become a hard wall. M2 RNN provides a way over that wall.

Key Takeaways:

• State size matters more than almost any other recurrent hyperparameter.
• Hybridization is the current SOTA: interleave different "mixing" strategies (Attention + Nonlinear RNN) to capture different data patterns.
• Distributed Training: The paper provides two strategies for Tensor Parallelism, making it ready for cluster-scale training.