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

Paper: 2602.18417 Authors: Joshua Nunley Categories: cs.LG, cs.CL

Abstract

This paper introduces a mathematical framework that unifies recurrent neural networks (RNNs) and transformers through the lens of Lie group theory. By constraining hidden states to closed subgroups of the unitary group $U(d)$, the authors derive both recurrent and attention-based architectures from a shared mathematical skeleton. The framework treats subgroup selection as a modular design choice that simultaneously determines the state space geometry, tangent space projections, and update mechanisms. Experimental validation focuses on orthogonal group $O(d)$ instantiations, demonstrating competitive performance on language modeling benchmarks while introducing a tangent-space linear mixing extension that enhances learning efficiency.

Key Contributions

• A minimal axiomatic framework deriving sequence models from closed subgroups of $U(d)$
• Unified mathematical skeleton generating both RNN and transformer architectures through subgroup choice
• Concrete $O(d)$-based implementations with orthogonal hidden states
• Tangent-space linear mixing mechanism applicable across different subgroup selections
• Empirical validation on Tiny Shakespeare and Penn Treebank under parameter-matched conditions

Mathematical Framework and Architecture Derivation

The core insight lies in recognizing that many successful neural architectures implicitly constrain their hidden states to structured manifolds. This paper makes these constraints explicit by working directly with closed subgroups $G \subset U(d)$. The framework begins with three axioms: (1) hidden states live in $G$, (2) updates occur via exponential maps from the tangent space $\mathfrak{g}$, and (3) learned parameters generate tangent vectors through projection operators.

From these axioms, the recurrent template emerges naturally. Given input $x_t$ and previous state $h_{t-1} \in G$, the model computes a tangent vector $v_t \in \mathfrak{g}$ via learned projection $\pi_\theta(x_t, h_{t-1})$, then updates the state through the exponential map: $h_t = h_{t-1} \exp(v_t)$. This formulation generalizes standard RNNs while maintaining geometric structure.

The transformer variant replaces sequential updates with parallel attention mechanisms. Query, key, and value projections are constrained to respect the subgroup structure, with attention weights computed over tangent-space representations. The self-attention operation becomes a structured mixing of tangent vectors before exponential mapping back to the group manifold. This geometric perspective reveals why certain architectural choices (like residual connections) emerge naturally from the group structure.

Orthogonal Instantiation and Experimental Results

Specializing to $O(d)$ provides concrete implementations with orthogonal hidden states. The tangent space becomes the space of skew-symmetric matrices, and the exponential map reduces to matrix exponentiation. This choice offers computational advantages: orthogonal matrices preserve norms, preventing gradient explosion/vanishing, and the skew-symmetric constraint reduces parameter count.

Experiments on Tiny Shakespeare (character-level) and Penn Treebank (word-level) compare orthogonal-state RNNs and transformers against standard baselines under matched parameter budgets. The orthogonal RNN achieves comparable perplexity to LSTM while using fewer parameters, suggesting the geometric constraint acts as effective regularization. The orthogonal transformer shows competitive performance with standard transformers, with the gap narrowing as model depth increases.

The linear mixing extension introduces learnable combinations of tangent vectors before exponential mapping: $h_t = h_{t-1} \exp(\sum_i \alpha_i v_i)$ where $\alpha_i$ are learned coefficients. This mechanism, applicable to any subgroup choice, consistently improves performance across experiments, particularly in low-parameter regimes. The improvement suggests that richer tangent-space dynamics enhance the model’s expressive capacity without sacrificing geometric structure.

Theoretical Implications and Future Directions

This work bridges differential geometry and deep learning, suggesting that successful architectures may be understood as discrete approximations to flows on Lie groups. The framework raises several theoretical questions: What properties of a subgroup $G$ determine the expressiveness of the resulting architecture? Can we characterize the function classes learnable by different subgroup choices?

The modular nature of the framework invites systematic exploration of the subgroup zoo. Beyond $O(d)$, natural candidates include $SU(d)$ (special unitary), $Sp(d)$ (symplectic), and various product groups. Each choice induces different geometric constraints and computational properties. The paper’s methodology provides a principled way to evaluate these alternatives.

From a practical standpoint, the tangent-space mixing mechanism deserves further investigation. Its consistent improvements suggest that enriching the tangent algebra—while maintaining group structure—may be a general principle for architecture design. Future work might explore adaptive mixing schemes, hierarchical tangent spaces, or connections to other geometric deep learning frameworks.

Takeaways

1. Closed subgroups of $U(d)$ provide a unified mathematical foundation for deriving both RNN and transformer architectures through geometric constraints on hidden states.

2. The framework treats subgroup selection as a modular design choice that simultaneously determines state space, update mechanisms, and projection operators.

3. Orthogonal-state models ($O(d)$ instantiation) achieve competitive performance on language modeling benchmarks while offering computational advantages through norm preservation.

4. Linear mixing in tangent space improves performance across subgroup choices, particularly in parameter-constrained settings, suggesting a general principle for architecture enhancement.

5. The geometric perspective reveals deep connections between Lie group theory and neural architecture design, opening pathways for principled exploration of structured state spaces.