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

Summary

Nunley asks what sequence architectures follow if hidden states and token embeddings are elements of a compact matrix group, specifically a closed subgroup of U(d), rather than Euclidean vectors. Three axioms define the framework: causality, staying inside the group, and updates generated by Lie-algebra (tangent) elements. From these, one skeleton yields both an RNN and a transformer. Each step builds a tangent update, projects it into the algebra, optionally passes it through a learned linear "mixing" map, and applies it multiplicatively through the matrix exponential. Logits come from trace similarity to learned prototype group elements (USIM readout). Swapping the subgroup changes only the projection and update internals. The experiments cover only O(d), using character-level Tiny Shakespeare and Penn Treebank at matched parameter budgets. There, the orthogonal transformer (OSMFormer) and OSM-RNN match or beat a standard ALiBi transformer and an LSTM in bits per character. The learned tangent mixing gives the largest gains. The work matters as a compact, geometry-first design that is stable by construction, though all results are single-seed.

Key points

• Tiny Shakespeare at about 100K parameters: OSMFormer (USIM, ALiBi, linear mixing, 91,429 params) reaches 2.464 test BPC, compared with 2.583 for a 2-layer, 1-head ALiBi transformer (104,111 params).
• Penn Treebank (character level) at about 100K parameters: OSMFormer gets 1.614 test BPC against 1.664 for the baseline.
• RNN comparison at about 52K parameters on Tiny Shakespeare: OSM-RNN scores 2.539 test BPC, while the LSTM scores 2.594.
• On Tiny Shakespeare, OSMFormer stays ahead at every scaling budget tested from 100K to 500K. At 500K it reaches 2.365 test BPC, versus 2.493 for the baseline.
• Mixing ablation at 100K:
  • no mixing gives 2.525;
  • a BCH-style Lie-bracket correction gives 2.532;
  • linear tangent mixing gives 2.464.
• In a 9-setting optimizer sweep at 500K, OSMFormer trained in every setting and showed a narrower validation spread than the baseline. All runs are single-seed, and only O(d) was tested.

Interesting ideas

• I1. The paper drops learned query, key and value maps entirely and scores attention as the trace inner product between two states, which stays unchanged under a global change of basis; this is a parameter-cheap similarity that works with almost no extra machinery.
• I2. Each transformer layer has two sub-steps, first moving toward the attention-weighted context and then pulling back toward the token's own embedding, so lexical identity is re-anchored at every depth.
• I3. The derivative of the readout splits each update into directions that change the prediction and directions the output cannot see, which the paper reads as predictive versus memory channels.
• I4. Classes are scored by similarity to learned prototype elements that live in the same space as the states, instead of through a generic linear output layer.
• I5. The evaluation protocol fixes parameter budgets, puts each component in its own ablation row, and sweeps optimizer settings to compare stability ranges, not just best scores.

For Mova

How these ideas could serve Mova, our small open language model — its files, stages and tests:

• I1 (tree): gparse could test a head-dependent scoring variant built on a basis-invariant similarity of token states, as a small-budget ablation that also leaves fewer opaque projection matrices to explain.
• I2 (tree): a parser layer with an explicit "re-ground to the word" step could keep word-specific evidence visible for arc decisions and make it possible to measure how much context versus the word itself drove each attachment.
• I3 (world): this suggests explicitly separating state changes that affect the current answer from latent bookkeeping such as holdings or event order, and logging both so every update to the text's world can be explained.
• I4 (tokens): POS or morphology tags could be predicted as nearest learned prototypes, so a wrong tag can be explained by showing which prototype the token's state sat closest to.
• I5 (learn): parser and tagger variants could be compared at matched size with a fixed optimizer sweep, reporting the spread across settings so that a gain from a tuning fluke is not mistaken for a real one.