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

1. Problem Statement

Modern sequence models (RNNs and Transformers) rely on a complex interplay of architectural choices, including update parameterization, stabilization layers, and memory mechanisms. While orthogonal and unitary RNNs have been explored to stabilize dynamics by constraining operators, the hidden states themselves typically remain Euclidean vectors ($\mathbb{R}^n$).

The paper addresses the following questions:

• What happens if we treat hidden states and token representations as elements of a closed subgroup of the unitary group $U(d)$ rather than Euclidean space?
• Can a unified, minimal framework be derived that naturally generates both RNN and Transformer architectures from this single assumption?
• Do these "group-valued" models offer competitive performance and improved stability compared to standard baselines?

2. Methodology

A. Theoretical Framework

The author proposes a minimal axiomatic framework where the state space $H_t$ and token representations $M_{x_t}$ belong to a closed Lie subgroup $G \leq U(d)$ with Lie algebra $\mathfrak{g}$.

Core Axioms:

1. Causality: $H_{t+1}$ depends only on $(H_t, x_t)$.
2. Group Closure: If $H_t \in G$, then $H_{t+1} \in G$.
3. Local Update: Updates are generated by tangent elements in $\mathfrak{g}$.

The Update Mechanism:
Instead of additive updates ($H_{t+1} = H_t + \Delta$), the model uses a multiplicative update via the exponential map:
$$H_{t+1} = H_t \exp(\Phi_\theta(H_t, M_{x_t}))$$
where $\Phi_\theta$ maps the current state and token to the Lie algebra $\mathfrak{g}$. This ensures the state remains on the manifold $G$ by construction.

Readout:
Predictions are generated using a "subgroup-native" readout based on the real part of the trace inner product (similarity) between the state and prototype matrices $P_v \in G$:
$$\ell_v(H_t) = \tau \Re \text{tr}(H_t^* P_v) + b_v$$

B. Unified Architectural Templates

The paper defines a shared skeleton for both RNNs and Transformers where the choice of subgroup $G$ acts as a "drop-in" component.

• RNN Template: Updates the state sequentially using the tangent projection of the token and current state.
• Transformer Template:
  • Attention: Computes attention weights using subgroup similarity scores $s_{ij} = \tau \Re \text{tr}((H_i)^* H_j)$. This is coordinate-invariant under global basis changes.
  • Update Steps: The state update is split into two sub-steps:
    1. Attention Step: Updates the state toward the attention-weighted aggregate of previous states.
    2. Grounding Step: Adjusts the state toward the current token embedding.
  • Both steps utilize tangent space projections and optional linear mixing.

C. Tangent Mixing

A key innovation is the Linear Tangent Mixing component ($\Gamma: \mathfrak{g} \to \mathfrak{g}$).

• Base Case: $\Gamma = \text{Id}$ (Identity).
• Learned Variant: A learned linear map in tangent coordinates.
• Function: It reweights and routes updates between "predictive" (readout-visible) and "memory" (readout-null) directions within the tangent space before exponentiation. This allows the model to reallocate capacity without altering the group structure.

D. Subgroup Instantiations

The framework is generic but the paper focuses on specific subgroups:

• $U(d)$ / $SU(d)$: Full unitary or phase-free updates.
• $O(d)$ / $SO(d)$: Real orthogonal updates (the primary focus of experiments).
• $T^k$: Abelian channel dynamics.

3. Key Contributions

1. Unified Group-Valued Framework: A minimal axiomatic setup that derives both RNN and Transformer architectures directly from the assumption that states live on a Lie subgroup of $U(d)$.
2. Drop-in Subgroup Design: Demonstrates that changing the subgroup (e.g., from $U(d)$ to $O(d)$) acts as a modular replacement for the state space, tangent projection, and update map, without changing the outer model interface.
3. Orthogonal-State Models (OSM): Introduces specific instantiations for $O(d)$, termed OSM-RNN and OSMFormer, which maintain orthogonal states throughout the network.
4. Linear Tangent Mixing: Proposes a general extension applicable to any subgroup that significantly improves finite-budget performance by learning how to mix tangent directions.

4. Experimental Results

The authors evaluated $O(d)$ models (OSM) on character-level language modeling tasks: Tiny Shakespeare and Penn Treebank (PTB). Models were trained with parameter-matched settings against standard baselines (Vaswani et al. Transformers, LSTMs).

Key Findings:

• Performance:
  • Tiny Shakespeare (100K params): OSMFormer achieved 2.464 Test BPC, outperforming the baseline Transformer (2.583 BPC).
  • Penn Treebank (100K params): OSMFormer achieved 1.614 Test BPC, outperforming the baseline Transformer (1.664 BPC).
  • RNN Comparison: OSM-RNN (55K params) outperformed a standard LSTM (52K params) on Tiny Shakespeare (2.539 vs 2.594 Test BPC).
• Scaling: OSMFormer consistently outperformed the baseline Transformer across parameter budgets ranging from 100K to 500K on both datasets.
• Ablation on Mixing: The Linear Tangent Mixing variant was the best-performing configuration. Interestingly, it achieved superior results with fewer parameters (91K) compared to the non-mixing version (103K) due to efficient hidden dimension allocation.
• Robustness: In a 500K parameter sweep with varying learning rates and optimizer settings, OSMFormer demonstrated a narrower range of best-validation performance and remained trainable across settings where the baseline was more sensitive.

5. Significance and Future Directions

Significance:

• Theoretical Elegance: The work provides a mathematically rigorous, minimal path to sequence modeling where stability (bounded dynamics) and geometry (Lie algebra updates) are inherent to the architecture, not added via regularization.
• Practical Competitiveness: It proves that constraining states to a manifold (specifically $O(d)$) is not just a theoretical curiosity but a practical strategy that yields state-of-the-art or competitive results with fewer parameters.
• Modularity: The "drop-in" nature of the subgroup choice suggests a new design axis for future models, allowing researchers to swap algebraic structures (e.g., from orthogonal to symplectic) to suit specific memory or stability needs.

Future Work (as outlined by the author):

• Interpretability: Leveraging the tangent space structure to track update magnitudes and attention contributions for better model diagnostics.
• Stability Analysis: Formalizing finite-step sensitivity bounds and conducting multi-seed robustness sweeps.
• Memory Dynamics: Investigating how different subgroup families (Abelian vs. Non-Abelian) affect long-range dependency retention and transfer.

In conclusion, this paper establishes a direct pipeline for building sequence models on Lie subgroups, showing that Orthogonal-State Models (OSM) with linear tangent mixing are a highly effective, stable, and parameter-efficient alternative to standard Euclidean sequence architectures.