Subgroups of $U(d)$ Induce Natural RNN and Transformer Architectures

Joshua Nunley

cs.LGFeb 20, 2026

Subgroups of $U(d)$ Induce Natural RNN and Transformer Architectures

Authors: Joshua Nunley

Organizations: Cognitive Science Program & Department of Informatics, Indiana University

Abstract

This paper presents a direct framework for sequence models with hidden states on closed subgroups of U(d). We use a minimal axiomatic setup and derive recurrent and transformer templates from a shared skeleton in which subgroup choice acts as a drop-in replacement for state space, tangent projection, and update map. We then specialize to O(d) and evaluate orthogonal-state RNN and transformer models on Tiny Shakespeare and Penn Treebank under parameter-matched settings. We also report a general linear-mixing extension in tangent space, which applies across subgroup choices and improves finite-budget performance in the current O(d) experiments.

Figures & tables

Subgroup	Tangent space	Tangent projection	Practical effect
$\mathrm{U}(d)$	$\mathfrak{u}(d)$	$A\leftarrow\tfrac{1}{2}(A-A^{\ast})$	full unitary updates
$\mathrm{SU}(d)$	$\mathfrak{su}(d)$	skew-Hermitian with trace removal	phase-free unitary updates
$\mathrm{O}(d)$ / $\mathrm{SO}(d)$	$\mathfrak{so}(d)$	$A\leftarrow\tfrac{1}{2}(A-A^{\top})$	real orthogonal updates
$T^{k}$	diagonal imaginary algebra	keep diagonal tangent terms	abelian channel dynamics

Table 1: Subgroup-specific drop-in components for the shared RNN/transformer templates.

Model	Params	Val BPC	Test BPC
Transformer (ALiBi, 2L1H)	104,111	2.259	2.583
OSMFormer (USIM, ALiBi, linear mixing)	91,429	2.254	2.464

Table 2: Tiny Shakespeare, approximately 100K parameters (single-seed; OSM row uses USIM readout).

Model	Params	Val BPC	Test BPC
Transformer (ALiBi, 2L1H)	101,816	1.716	1.664
OSMFormer (USIM, ALiBi, linear mixing)	103,482	1.665	1.614

Table 3: Penn Treebank, approximately 100K parameters (single-seed; OSM row uses USIM readout).

Model	Params	Val BPC	Test BPC
LSTM	52,305	2.407	2.594
OSM-RNN (USIM, linear mixing)	55,545	2.366	2.539

Table 4: Tiny Shakespeare, approximately 52K parameters (single-seed RNN comparison; OSM row uses USIM readout).

Budget	Baseline Params	Baseline Val/Test	OSM Params	OSM Val/Test
100K	104,111	2.259 / 2.583	91,429	2.254 / 2.464
300K	299,047	2.207 / 2.552	277,357	2.184 / 2.377
400K	392,987	2.210 / 2.559	380,757	2.184 / 2.373
500K	499,727	2.229 / 2.493	511,749	2.176 / 2.365

Table 5: Tiny Shakespeare scaling (single-seed): baseline transformer vs OSMFormer (USIM, linear mixing).

Budget	Baseline Params	Baseline Val/Test	OSM Params	OSM Val/Test
300K	295,072	1.596 / 1.551	310,002	1.570 / 1.524
400K	408,616	1.568 / 1.523	423,754	1.557 / 1.512
500K	494,552	1.558 / 1.512	491,454	1.549 / 1.503

Table 6: Penn Treebank scaling (single-seed): baseline transformer vs OSMFormer (USIM, linear mixing).

Figure 1: Tiny Shakespeare scaling curve corresponding to Table 5 .

Figure 2: Penn Treebank scaling curve corresponding to Table 6 .

Mixing mode	Params	Val BPC	Test BPC
No tangent mixing	103,501	2.355	2.525
BCH-style intrinsic mix	103,503	2.349	2.532
Linear tangent mixing	91,429	2.254	2.464

Table 7: Tiny Shakespeare, 100K OSMFormer mixing ablation (single-seed, USIM readout fixed across rows).

Figure 3: Optimizer robustness summary on Tiny Shakespeare at 500K parameters (2-layer and 4-layer settings).

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Subgroup	Raw parameters	Projection to $\mathfrak{g}$	On-group map
$\mathrm{U}(d)$	$B_{v},C_{v}\in\mathbb{C}^{d\times d}$	$\mathrm{skew}_{\mathbb{C}}$	$\Exp$
$\mathrm{SU}(d)$	$B_{v},C_{v}\in\mathbb{C}^{d\times d}$	$\mathrm{tr0}\circ\mathrm{skew}_{\mathbb{C}}$	$\Exp$
$\mathrm{SO}(d)$	$B_{v},C_{v}\in\mathbb{R}^{d\times d}$	$\mathrm{skew}_{\mathbb{R}}$	$\Exp$
$\mathrm{O}(d)$	$B_{v},C_{v}\in\mathbb{R}^{d\times d}$	$\mathrm{skew}_{\mathbb{R}}$	$F^{\eta_{v}}\Exp(\cdot)$ for $M_{v}$ , $F^{\xi_{v}}\Exp(\cdot)$ for $P_{v}$
$T^{k}$	$\beta_{v},\gamma_{v}\in\mathbb{R}^{k}$	$i\,\operatorname{diag}(\cdot)$	$\Exp$ (equivalently elementwise complex phase)

Appendix

Table 8: Embedding/prototype parameterization by subgroup (operator form). For

\mathrm{O}(d)

F=\operatorname{diag}(-1,1,\dots,1)

and

\eta_{v},\xi_{v}\in\{0,1\}

. Exact token-wise formulas are given in the recipe list below.

Recurrent Neural Networks

Recurrent neural networks process sequences by updating a persistent hidden state, enabling temporal modeling, memory, and streaming computation.

See all 303 papers

M $^2$ RNN: Non-Linear RNNs with Matrix-Valued States for Scalable Language Modeling
Mar 15, 2026Mayank Mishra, Shawn Tan, Ion Stoica +2Recurrent Neural Networks Language Modeling
UC Berkeley · Learning Machine · MIT-IBM Watson Lab +2
Tracking States or Tracking Cosets? An Algebraic Account of Learned State Tracking
Sep 24, 2026Zhiyu Zhang, Yupeng LiTransformer Recurrent Neural Networks
Conflux Labs Ltd · Michigan State University
RotRNN: Modelling Long Sequences with Rotations
Jul 9, 2024Kai Biegun, Rares Dolga, Jake Cunningham +1Recurrent Neural Networks Long-Context Modeling
UCL AI Centre London, UK · UCL AI Centre & UiPath London, UK

Mar 15, 2026J/K move · Enter open · S save

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

Mayank Mishra, Shawn Tan, Ion Stoica, Joseph E. Gonzalez+1

UC Berkeley · Learning Machine · MIT-IBM Watson Lab+2

Recurrent Neural Networks Language Modeling Autoregressive Language Modeling Long-Context Language Modeling

Open paper AskPDF arXiv

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.

Abstract

Recurrent Neural Networks

Explore similar work

M2^22RNN: Non-Linear RNNs with Matrix-Valued States for Scalable Language Modeling

Tracking States or Tracking Cosets? An Algebraic Account of Learned State Tracking

RotRNN: Modelling Long Sequences with Rotations

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