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

Probleemstelling

Moderne sequentiemodellen (zoals RNN's en Transformers) bestaan uit een complex samenspel van architecturale keuzes: parameterisatie van updates, stabilisatielagen, geheugenvoorzieningen en leesmechanismen. Bestaande methoden om recurrente dynamica te stabiliseren (zoals orthogonale of unitaire RNN's) beperken zich vaak tot het construeren van de operator (de transitiematrix), terwijl de verborgen staten zelf in de Euclidische ruimte ($\mathbb{R}^n$) blijven. Dit leidt tot een disconnectie tussen de geometrie van de toestand en de dynamica van de updates. Er is behoefte aan een fundamenteel raamwerk waarin de toestand zelf een groepselement is, wat natuurlijke beperkingen en geometrische eigenschappen biedt zonder extra componenten toe te voegen.

Methodologie

Het artikel introduceert een raamwerk voor sequentiemodellen waarbij zowel de token-representaties als de verborgen staten direct worden gedefinieerd op een gesloten ondergroep $G$ van de unitaire groep $U(d)$.

1. Axiomatische Basis:
Het model rust op drie minimale aannames:

• Causaliteit: De nieuwe toestand $H_{t+1}$ hangt alleen af van de huidige toestand $H_t$ en het token $x_t$.
• Groeps-sluiting: Als $H_t \in G$, dan moet ook $H_{t+1} \in G$.
• Lokale updates: Elke stap wordt gegenereerd door een element in de Lie-algebra $\mathfrak{g}$ van de groep $G$.

2. Het Update-mechanisme:
In plaats van lineaire transformaties in $\mathbb{R}^n$, wordt de update uitgevoerd via de exponentiële afbeelding:
$$H_{t+1} = H_t \exp(\Phi_\theta(H_t, M_{x_t}))$$
Hierbij is $\Phi_\theta$ een meetbare map die een element in de Lie-algebra $\mathfrak{g}$ genereert. Door de eigenschappen van de exponentiële afbeelding en de groepssluiting, blijft de toestand automatisch binnen de groep $G$.

3. Architecturale Templates:
Het artikel definieert twee gedeelde templates die groep-onafhankelijk zijn, waarbij de keuze van de ondergroep $G$ fungeert als een "drop-in" component:

• RNN-template: Een recurrente update waarbij de toestand wordt bijgewerkt via een tangent-vector in $\mathfrak{g}$.
• Transformer-template: Een laag-voor-laag structuur die twee sub-stappen gebruikt:
  1. Attention-stap: Berekening van aandachtsgewichten op basis van groepssimilariteit ($\Re\text{tr}(H_i^* H_j)$) en een update naar het gewogen gemiddelde in de tangentruimte.
  2. Grounding-stap: Een aanpassing richting de huidige token-embedding.
     In deze templates worden geen aparte Query/Key/Value-maps gebruikt; de similariteit is inherent aan de groepstructuur.

4. Leesmechanisme (Readout):
Voorspellingen worden gedaan via prototype-matrices $P_v \in G$ en een logit-functie gebaseerd op het reële deel van de spoor (trace) van het product van de toestand en de prototype:
$$\ell_v(H_t) = \tau \Re\text{tr}(H_t^* P_v) + b_v$$
Dit is een natuurlijke, groep-intrinsieke methode voor similariteitsmeting.

5. Tangent Mixing:
Een belangrijke extensie is de optionele toepassing van een lineaire map $\Gamma: \mathfrak{g} \to \mathfrak{g}$ in de tangentruimte voordat de exponentiële stap wordt genomen. Dit stelt het model in staat om updates te herschikken tussen "voorspellende" en "geheugen"-richtingen in de tangentruimte.

Belangrijkste Bijdragen

1. Unificatie van RNN en Transformer: Het biedt een enkel raamwerk dat beide architecturen afleidt uit dezelfde groep-gebaseerde principes, waarbij de ondergroepkeuze de enige variabele is die de interne geometrie bepaalt.
2. Geometrische State-Representatie: In tegenstelling tot eerdere werken die operators construeren, plaatst dit werk de toestand zelf op de groep (bijv. $O(d)$), wat zorgt voor gebonden dynamica per constructie.
3. Drop-in Componenten: Het toont aan dat het wisselen van ondergroep (bijv. van $U(d)$ naar $O(d)$ of $SU(d)$) eenvoudigweg gaat door het aanpassen van de projectie naar de tangentruimte en de update-regels, zonder de externe trainingsinterface te veranderen.
4. Empirische Validatie: Het artikel presenteert experimenten met de orthogonale ondergroep $O(d)$, wat resulteert in "Orthogonal-State Models" (OSM).

Resultaten

De auteurs evalueren de $O(d)$-instantiatie (OSM-RNN en OSMFormer) op twee benchmarks: Tiny Shakespeare en Penn Treebank (PTB), vergeleken met standaard baselines (LSTM, Transformer) onder gelijke parameterbudgetten.

• Prestaties:
  • Op Tiny Shakespeare (100K parameters) behaalde de OSMFormer (met lineaire tangent mixing) een test BPC (bits per character) van 2.464, vergeleken met 2.583 voor de standaard Transformer.
  • Op Penn Treebank (100K parameters) behaalde de OSMFormer een test BPC van 1.614, vergeleken met 1.664 voor de baseline.
  • Bij RNN's (OSM-RNN vs. LSTM) op Tiny Shakespeare (50K parameters) behaalde OSM-RNN een test BPC van 2.539 versus 2.594 voor LSTM.
• Scalabiliteit: De prestaties van OSMFormer verbeteren consistent bij het verhogen van het parameterbudget (tot 500K), waarbij het de baselines overtreft op zowel validatie- als testset.
• Robuustheid: De modellen tonen zich robuust tegenover verschillende optimizer-instellingen (leer snelheid, clipping, weight decay), met een smaller bereik van beste validatieprestaties dan de baselines, wat wijst op stabielere training.
• Impact van Tangent Mixing: De "lineaire tangent mixing" variant bleek de beste prestaties te leveren, zelfs met minder parameters dan sommige niet-mixing varianten.

Betekenis en Toekomstperspectief

Dit werk biedt een fundamentele verschuiving in hoe we sequentiemodellen ontwerpen: van het construeren van stabiele operatoren in Euclidische ruimte naar het definiëren van staten op compacte groepen.

• Interpreteerbaarheid: Omdat updates plaatsvinden in de tangentruimte, biedt dit een directe manier om de geometrie van updates te analyseren (bijv. magnitude en richting).
• Stabiliteit: De compacte aard van de groepen garandeert per constructie gebonden dynamica, wat de training stabiliseert.
• Geheugen: Het raamwerk suggereert een nieuwe dimensie voor geheugenontwerp door het vergelijken van abelse versus niet-abelse ondergroepen.

De conclusie is dat dit raamwerk een praktische en effectieve route biedt voor het bouwen van stabiele, geometrisch onderbouwde sequentiemodellen, met $O(d)$ als een veelbelovende eerste instantiatie.