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

Mayank Mishra; Shawn Tan; Ion Stoica; Joseph Gonzalez; Tri Dao

cs.LGMar 15, 2026

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

Authors: Mayank Mishra, Shawn Tan, Ion Stoica, Joseph Gonzalez, Tri Dao

Organizations: UC Berkeley · Learning Machine · MIT-IBM Watson Lab · Princeton University · Together AI

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, and show 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 language models.

Figures & tables

Figure 1: Forget gate (

f_{t}

) behavior as a function of the input (

x_{t}

) for different values of

\alpha_{n}

and

\beta_{n}

n

denotes the

n^{\textrm{th}}

head.

Panel of Figure 2 — Figure 2: Visualization of the Matrix-to-Matrix RNN layer. This block replaces attention and is combined with MLP and RMSNorm modules as in the Transformer.

Figure 3: Accuracy on the permutation group

S_{3}

state-tracking task for M 2 RNN, Gated DeltaNet

[-1,1]

, Gated DeltaProduct

[-1,1]

and GRU. The vertical line at

128

denotes the training context length for all models. For the Gated DeltaProduct, we use product of 2 HouseHolder matrices since it can solve the

S_{3}

task in theory Grazzi et al. (2024) . However, as the evaluation context length is increased beyond the training context length, we observe accuracy degradation.

Panel of Figure 4 — Figure 4: TP topology-aware M 2 RNN layer running with TP on 2 GPUs. Note that RMSNorm 1 and RMSNorm 2 have different weights for the RMSNorm module on both GPUs.

Panel of Figure 5 — Figure 5: TP topology-independent M 2 RNN layer. Note that RMSNormTP have weights sharded along the model width dimension for the RMSNorm module on both GPUs requiring a synchronization in both the forward and backward computation.

Model	Wiki PPL	LMB PPL	LAMBADA	HellaSwag	PIQA	ARC-E	ARC-C	WinoGrande	BoolQ	OBQA	COPA	SciQ	Avg Acc
Mamba-2	22.93	38.56	32.78	47.70	70.29	62.04	30.55	53.28	58.69	33.20	69.00	85.50	54.30
Gated DeltaNet	22.70	34.17	34.21	47.16	71.00	61.24	28.07	51.07	58.84	32.80	70.00	85.50	53.99
RNN	33.74	317.38	13.00	44.12	69.70	62.54	30.63	49.72	45.75	33.40	63.00	77.10	48.90
GRU	25.80	63.83	26.10	48.70	71.49	64.90	30.97	49.64	57.68	33.80	71.00	82.90	53.72
M 2 RNN	22.92	33.63	33.98	47.76	71.60	62.29	30.29	53.43	50.21	34.80	72.00	86.00	54.24
Transformer++	23.32	42.87	33.36	45.02	68.88	59.76	29.01	50.43	54.56	32.00	72.00	84.00	52.90

Table 1: Language Modeling on the 410M parameter dense model across commonsense reasoning benchmarks. All models are evaluated in the 0-shot setting.

Model	Wiki PPL	LMB PPL	LAMBADA	HellaSwag	PIQA	ARC-E	ARC-C	WinoGrande	BoolQ	OBQA	COPA	SciQ	Avg Acc
Mamba-2	13.73	11.03	50.09	67.43	76.82	74.33	43.52	61.64	64.01	41.00	80.00	93.40	65.22
Gated DeltaNet	13.89	11.10	48.92	66.94	77.26	73.82	40.53	60.14	65.41	39.60	78.00	92.20	64.28
RNN	17.65	43.68	30.25	64.49	76.61	72.64	41.04	55.64	58.35	40.20	79.00	87.60	60.58
GRU	14.80	15.33	43.94	67.36	76.82	75.67	43.52	58.96	63.79	41.00	82.00	91.30	64.44
M 2 RNN	13.80	11.48	49.18	67.57	76.71	75.51	44.97	58.17	64.16	41.40	82.00	92.00	65.17
Transformer++	14.94	17.26	43.90	62.99	75.63	71.97	38.91	57.54	63.12	37.20	75.00	91.40	61.76

Table 2: Language Modeling on the 7B (1B active) MoE model across commonsense reasoning benchmarks. All models are evaluated in the 0-shot setting.

Figure 6: Zero-shot in-context retrieval performance on the RULER benchmark for M 2 RNN layers and Gated DeltaNet models on 410M parameter scale. The vertical lines at 4096 indicate the training context length.

Figure 8: Zero-shot in-context retrieval performance on the RULER benchmark for M 2 RNN layers and Gated DeltaNet models for the 7B MoE model. The vertical lines at 4096 indicate the training context length.

Model	SQuAD	NQ	DROP	TriviaQA	FDA	SWDE	Avg
Mamba-2	31.2	15.0	19.5	50.5	15.6	32.4	27.4
Gated DeltaNet	34.8	15.4	24.9	51.7	16.3	41.8	30.8
RNN	9.6	5.9	14.6	29.7	0.0	5.0	10.8
GRU	15.8	8.8	17.2	41.5	0.7	7.7	15.3
M 2 RNN	34.0	13.0	22.9	52.3	10.6	28.5	26.9
Transformer++	14.0	15.5	28.2	52.4	46.0	51.7	34.6

Table 3: In-context retrieval performance on real-world data at the 410M parameter scale.

Model	SQuAD	NQ	DROP	TriviaQA	FDA	SWDE	Avg
Mamba-2	41.7	26.5	30.0	65.0	33.0	54.4	41.8
Gated DeltaNet	40.3	24.2	29.9	64.3	32.4	61.8	42.2
RNN	20.0	12.2	20.4	53.1	0.8	9.8	19.4
GRU	27.0	16.7	24.6	59.3	5.2	18.4	25.2
M 2 RNN	37.9	21.1	28.2	62.5	27.6	49.9	37.9
Transformer++	22.1	24.0	27.6	62.2	67.0	70.1	45.5

Table 4: In-context retrieval performance on real-world data at the 7B (1B active) MoE scale.

Model	GovReport	QMSum	MultiNews	TREC	TriviaQA	SAMSum	LCC	RepoBench-P	Avg
Mamba-2	6.6	14.3	9.9	26.5	23.2	11.5	13.7	8.9	14.3
Gated DeltaNet	7.6	18.2	12.0	19.0	34.1	21.3	14.9	11.8	17.3
RNN	9.3	15.2	9.2	0.5	12.2	5.0	17.9	14.9	10.5
GRU	7.0	16.5	10.8	1.0	21.4	8.5	12.5	9.0	10.8
M 2 RNN	6.3	16.4	10.8	22.0	24.3	13.2	10.7	10.1	14.2
Transformer++	5.6	5.7	6.2	11.0	9.7	4.6	12.0	9.9	8.1

Table 5: Performance comparison on the 410M model across long-context summarization, coding and few-shot learning benchmarks.

Model	GovReport	QMSum	MultiNews	TREC	TriviaQA	SAMSum	LCC	RepoBench-P	Avg
Mamba-2	5.1	16.0	13.4	42.0	59.7	22.0	13.4	8.7	22.5
Gated DeltaNet	4.8	20.3	12.2	47.0	61.6	36.2	9.7	7.9	25.0
RNN	9.3	16.7	10.5	4.0	27.5	13.0	17.1	15.8	14.2
GRU	5.7	18.0	10.8	9.0	52.4	24.6	20.9	13.9	19.4
M 2 RNN	8.1	19.1	11.8	26.5	51.3	30.7	20.4	19.3	23.4
Transformer++	6.1	6.1	17.0	16.0	11.4	10.0	13.1	10.2	11.2

Table 6: Performance comparison on the 7B (1B active) MoE model across long-context summarization, coding and few-shot learning benchmarks.

Model	Wiki PPL	Lambada PPL	state size
RNN-406M	33.74	317.38	1,360
GRU-474M	25.80	63.83	1,360
M 2 RNN-410M	22.92	33.63	86,016

Table 7: Perplexity comparison for RNN, GRU and M 2 RNN models with different state sizes.

Figure 10: Training throughput measured in billion tokens per day on a machine with 8x NVIDIA HGX H100 GPUs for our 7B MoE configuration.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Attention Pattern	$\textbf{W}_{q}$	$\textbf{W}_{k}$	$\textbf{W}_{v}$	$\textbf{W}_{g}$	$\textbf{W}_{f}$	$\textbf{W}_{o}$	Total ( $P$ )	State Size ( $NKV$ )
Multi-Head	$NKd_{\textrm{m}}$	$NKd_{\textrm{m}}$	$NVd_{\textrm{m}}$	$NVd_{\textrm{m}}$	$Nd_{\textrm{m}}$	$NVd_{\textrm{m}}$	$\approx N(2K+3V)d_{\textrm{m}}$	$\approx\frac{PKV}{(2K+3V)d_{\textrm{m}}}$
Multi-Query	$NKd_{\textrm{m}}$	$Kd_{\textrm{m}}$	$Vd_{\textrm{m}}$	$NVd_{\textrm{m}}$	$Nd_{\textrm{m}}$	$NVd_{\textrm{m}}$	$\approx N(K+2V)d_{\textrm{m}}$	$\approx\frac{PKV}{(K+2V)d_{\textrm{m}}}$
Multi-Key	$Kd_{\textrm{m}}$	$NKd_{\textrm{m}}$	$Vd_{\textrm{m}}$	$NVd_{\textrm{m}}$	$Nd_{\textrm{m}}$	$NVd_{\textrm{m}}$	$\approx N(K+2V)d_{\textrm{m}}$	$\approx\frac{PKV}{(K+2V)d_{\textrm{m}}}$
Multi-Value	$Kd_{\textrm{m}}$	$Kd_{\textrm{m}}$	$NVd_{\textrm{m}}$	$NVd_{\textrm{m}}$	$Nd_{\textrm{m}}$	$NVd_{\textrm{m}}$	$\approx 3NVd_{\textrm{m}}$	$\approx\frac{PK}{3d_{\textrm{m}}}$

Appendix

Table 8: Number of parameters in all linear projections.

N

denotes the number of heads and

d_{\textrm{m}}

denotes the model’s embedding dimension. For all our models, we use

K=64

and

V=16

which results in the largest state size for the multi-value attention formulation for constant number of total parameters (

P

Explore similar work

Olmo Hybrid: From Theory to Practice and Back
Apr 3, 2026William Merrill, Yanhong Li, Tyler Romero +19Language Modeling Hybrid Attention
Allen Institute for AI · ETH Zürich · Lambda ML Team +2
Triadic Linear Attention: Three-Dimensional Recurrent States for Long-Context Sequence Modeling
Sep 29, 2026Oliver Sieberling, Bharat Runwal, David Jin +3Linear Attention Long-Context Language Modeling
Massachusetts Institute of Technology · MIT-IBM Computing Research Lab
2Mamba2Furious: Linear in Complexity, Competitive in Accuracy
Feb 19, 2026Gabriel Mongaras, Eric C. LarsonLinear Attention State Space Models
Lyle School of Engineering Southern Methodist University Dallas, TX 75205

Apr 3, 2026J/K move · Enter open · S save

Olmo Hybrid: From Theory to Practice and Back

William Merrill, Yanhong Li, Tyler Romero, Anej Svete+18

Allen Institute for AI · ETH Zürich · Lambda ML Team+2

Language Modeling Hybrid Attention Language Model Scaling Language Model Pretraining Transformer Expressivity

Open paper AskPDF arXiv

Recent work has demonstrated the potential of non-transformer language models, especially linear recurrent neural networks (RNNs) and hybrid models that mix recurrence and attention. Yet there is no consensus on whether the potential benefits of these new architectures justify the risk and effort of scaling them up. To address this, we provide evidence for the advantages of hybrid models over pure transformers on several fronts. First, theoretically, we show that hybrid models do not merely inherit the expressivity of transformers and linear RNNs, but can express tasks beyond both, such as code execution. Putting this theory to practice, we train Olmo Hybrid, a 7B-parameter model largely comparable to Olmo 3 7B but with the sliding window layers replaced by Gated DeltaNet layers. We show that Olmo Hybrid outperforms Olmo 3 across standard pretraining and mid-training evaluations, demonstrating the benefit of hybrid models in a controlled, large-scale setting. We find that the hybrid model scales significantly more efficiently than the transformer, explaining its higher performance. However, its unclear why greater expressivity on specific formal problems should result in better scaling or superior performance on downstream tasks unrelated to those problems. To explain this apparent gap, we return to theory and argue why increased expressivity should translate to better scaling efficiency, completing the loop. Overall, our results suggest that hybrid models mixing attention and recurrent layers are a powerful extension to the language modeling paradigm: not merely to reduce memory during inference, but as a fundamental way to obtain more expressive models that scale better during pretraining.

Abstract

Explore similar work

Olmo Hybrid: From Theory to Practice and Back

Triadic Linear Attention: Three-Dimensional Recurrent States for Long-Context Sequence Modeling

2Mamba2Furious: Linear in Complexity, Competitive in Accuracy