RotRNN: Modelling Long Sequences with Rotations

Technische Samenvatting: RotRNN – Modellering van Lange Sequenties met Rotaties

Probleemstelling
Het modelleren van lange sequenties blijft een uitdagend domein binnen machine learning vanwege de moeilijkheid om langetermijnafhankelijkheden te vangen. Terwijl Transformers lijden onder kwadratische computationele schaling en gebrekkige inductieve biases, en niet-lineaire RNN's worstelen met verdwijnende en exploderende gradiënten, hebben recente State Space Models (SSM's) en Linear Recurrent Units (LRU's) state-of-the-art prestaties bereikt. Deze modellen vertonen echter aanzienlijke nadelen:

1. Complexe Initialisatie: SSM's zoals S4 vertrouwen op complexe theoretische afleidingen (bijv. HiPPO-initialisatie gebaseerd op optimale polynomiale projecties) die tijdens het trainingsproces niet goed begrepen worden.
2. Mismatch tussen Implementatie en Theorie: De LRU is conceptueel eenvoudiger, maar weerspiegelt in de praktijk niet volledig de theoretische motivatie. Specifiek slaagt het er niet in om geconjugeerde paren voor eigenwaarden van reële matrices af te dwingen, mist beperkingen om reële outputs te garanderen, en vertrouwt op een geleerde normalisator die geen consistente magnitudes van de verborgen toestand gedurende de training garandeert.
3. Stabiliteit: Bestaande methoden vereisen vaak complexe normalisatieschema's of initialisatieprocedures die mogelijk niet robuust zijn over verschillende dieptes en random seeds.

Methodologie
De auteurs stellen RotRNN (Rotational Recurrent Neural Network) voor, een lineair recurrent model dat de recurrente toestandmatrix $A$ parametriseert als een rotatiematrix. De methodologie is gebouwd op drie kernzuilen:

1. Parametrisatie via Rotatiematrices:
   In plaats van direct een beperkte matrix te leren, leert RotRNN een algemene gewichtsmatrix $M$ en brengt deze naar de speciale orthogonale groep $SO(N)$ met behulp van de matrixexponentiaal van een scheef-symmetrische matrix ($A = \exp(M - M^\top)$). Dit zorgt ervoor dat $A$ gedurende de training een geldige rotatiematrix blijft.
   Om efficiënte berekening van machtsverheffingen van matrices mogelijk te maken (vereist voor de convolutionele vorm van lineaire recurrentie), decomponeren de auteurs de rotatiematrix $A$ in een blokdiagonaal vorm: $A = P \Theta P^\top$. Hierbij is $P$ een orthogonale matrix en bestaat $\Theta$ uit $2\times2$ rotatieblokken gedefinieerd door hoeken $\theta_i$. Deze structuur maakt snelle berekening van $A^k$ mogelijk zonder dure dense matrixoperaties.

2. Geprincipeerde Normalisatie:
   Door gebruik te maken van de orthogonaliteit van rotatiematrices (waarbij de eigenwaarden op de eenheidscirkel liggen), leiden de auteurs een normalisatieprocedure af die een constante verwachte verborgen toestand-norm ($E[\|x_t\|^2] = c$) handhaaft gedurende de training.
   De recurrentie wordt gedefinieerd als $x_t = \gamma A x_{t-1} + \xi B u_t$, waarbij $\gamma$ een leerbare scalaire vervalfactor is en $\xi$ een herschalingscoëfficiënt die is afgeleid om aan de constante normconditie te voldoen. In tegenstelling tot de LRU, die normaliseert op basis van oneindige sequentielimieten, houdt de normalisatie van RotRNN stand voor alle tijdstappen, waardoor de exponentiële afname van vroege inputs — die vaak voorkomt wanneer normalisatieconstanten tot de macht $t$ worden verheven — wordt voorkomen.

3. Multi-Head Architectuur:
   Om de beperking van een enkele scalaire vervalfactor $\gamma$ (die mogelijk niet goed generaliseert naar taken die verschillende geheugenhorizonten vereisen) aan te pakken, gebruikt RotRNN $H$ onafhankelijke laagdimensionale heads. Elke head werkt met een eigen rotatiematrix $A^{(h)}$, inputmatrix $B^{(h)}$ en vervalfactor $\gamma^{(h)}$. De outputs worden geconcateneerd en lineair gemengd, waardoor het model afhankelijkheden op meerdere schalen kan vangen.

Belangrijkste Bijdragen

• Nieuwe Lineaire RNN-laag: Het voorstel van een recurrente laag waarbij de toestandmatrix als een rotatie wordt geparametriseerd, wat ervoor zorgt dat het algoritme getrouw blijft aan zijn theoretische motivatie tijdens de training.
• Efficiënte Berekening: Een methode voor het efficiënt berekenen van machtsverheffingen van parametrische rotaties door ze te decomponeren in blokdiagonaal vormen, wat snelle lineaire recurrentie mogelijk maakt.
• Robuuste Normalisatie: Een geprincipeerde afleiding van een normalisatieprocedure die een constante verwachte grootte van de verborgen toestand behoudt, wat sterkere stabiliteitsgaranties biedt dan voorgaande werken die vertrouwen op asymptotische convergentie.
• Theoretisch Inzicht: Een afleiding die de wiskundige equivalentie aantoont tussen een speciaal geval van de LRU (waar eigenwaarden geconjugeerde paren vormen) en de multi-head RotRNN, wat potentieel licht werpt op de interne mechanica van de LRU.

Resultaten
De auteurs evalueerden RotRNN op verschillende benchmarks voor lange sequentiemodellering:

• Long Range Arena (LRA): RotRNN toonde competitieve prestaties tegenover state-of-the-art lineaire recurrente modellen (S4, S5, LRU). Het behaalde de hoogste scores onder de baselines op taken met discrete input zoals ListOps, Text (IMDB) en Retrieval. Echter, het presteerde iets onder sommige baselines op pixel-niveau beeldtaken (Path-X, CIFAR).
• Speech Commands: Op de 35-way Speech Commands classificatietaak presteerde RotRNN identiek aan de gereproduceerde LRU (95,2% nauwkeurigheid) met een vergelijkbaar aantal parameters, en bleef het competitief met complexere diepe state space modellen.
• Stabiliteitsanalyse: Empirische analyse van verborgen toestand-normen op de ListOps dataset toonde aan dat RotRNN een bijna constante magnitude handhaaft over training, laagdiepte en random seeds. In contrast hiermee vertoonde de LRU wilde variaties in verborgen toestand-normen en had deze aanzienlijk langer nodig om te convergeren naar een stabiele magnitude.

Significantie en Claims
Het artikel beweert dat RotRNN een conceptueel eenvoudig, efficiënt en wiskundig geprincipeerd alternatief biedt voor bestaande lineaire recurrente modellen. De primaire significantie ligt in het overbruggen van de kloof tussen theoretische motivatie en praktische implementatie. Door rotatiematrices te gebruiken, vermijdt RotRNN de complexe initialisatieschema's van SSM's en de implementatie-inconsistenties van de LRU. De auteurs stellen dat hun model een robuuste normalisatieprocedure biedt die stabiliteit garandeert zonder te vertrouwen op "black box" initialisatie, terwijl het competitieve prestaties levert op standaard benchmarks. Verder suggereert het werk dat het bekijken van de LRU door de lens van rotatie-gebaseerde blokstructuren kan helpen om de mechanica van lineaire recurrentie in eerdere modellen te verhelderen.