M$^2$RNN: Non-Linear RNNs with Matrix-Valued States for Scalable Language Modeling

Dit paper introduceert M²RNN, een schaalbaar taalmodel met matrix-gebaseerde niet-lineaire recurrente lagen dat superieure prestaties en langere contextgeneralisatie bereikt ten opzichte van bestaande hybride architecturen, terwijl het tegelijkertijd efficiënter is in gebruik van tensorcores en kleinere staten vereist.

De kern

Stel je voor dat een kunstmatige intelligentie (zoals een chatbot) een enorm boek leest. Om het verhaal te begrijpen, moet het niet alleen de woorden onthouden, maar ook begrijpen wie wie is, wat er gebeurt, en hoe de zinnen met elkaar verbonden zijn.

Dit artikel introduceert een nieuwe manier om deze "geheugenmachine" te bouwen, genaamd M2RNN. Hier is de uitleg in simpele taal, met wat creatieve vergelijkingen:

1. Het Probleem: De twee uitersten

Tot nu toe hadden we twee soorten geheugens voor AI:

• De "Super-Snelle maar Flauwe" (Transformers): Dit is de huidige standaard. Het is als een super-snel fototoestel dat elke zin in één keer kan bekijken. Het is geweldig voor snelheid, maar het heeft een beperking: het kan niet goed "rekenen" of complexe logica volgen (zoals het bijhouden van wie wie is in een lang verhaal of het uitvoeren van code). Het is te simpel voor moeilijke taken.
• De "Slimme maar Traage" (Oude RNN's): Dit zijn de oude, slimme modellen die één woord per keer lezen. Ze zijn heel goed in logica en het onthouden van complexe verbanden, maar ze zijn ontzettend traag om te trainen. Het is alsof je een boek moet lezen door één letter per seconde te scannen, terwijl de rest van de wereld al vooruit is.

2. De Oplossing: M2RNN (De "Matrix-Geheugen")

De auteurs van dit paper hebben een nieuwe uitvinding gedaan: M2RNN.

Stel je voor dat je een gewone notitieblok hebt (een lijst met woorden). Dat is wat de oude, slimme modellen deden. Ze schreven alles op in één lange lijn.
M2RNN doet iets anders: het gebruikt een gigantisch, dubbelzijdig notitieblok (een matrix).

• De Analogie: In plaats van alleen te schrijven "De man loopt", schrijft M2RNN een heel schema op: "De man (persoon A) loopt (actie B) in de tuin (locatie C) met zijn hond (object D)".
• Door dit "dubbelzijdige" geheugen te gebruiken, kan het model veel meer informatie tegelijk vasthouden zonder dat het langzamer wordt. Het is alsof je van een klein notitieblok overstapt op een enorme whiteboard waar je alles tegelijk kunt zien.

3. Waarom is dit zo cool? (De 3 grote voordelen)

A. Het kan eindelijk "nadenken" (State Tracking)

Oude slimme modellen waren traag. Nieuwe snelle modellen (zoals Mamba) waren te simpel.
M2RNN is snel én slim. Het kan taken doen die de snelle modellen niet aankunnen, zoals het bijhouden van wie wie is in een lang verhaal of het uitvoeren van programmeercode.

• Vergelijking: Het is alsof je een snelle auto hebt die ineens ook een racefiets kan besturen. Het kan de complexe bochten nemen die de andere auto's niet aankunnen.

B. Het onthoudt alles (Langere context)

Vaak vergeten AI-modellen wat er aan het begin van een lang gesprek gebeurde als je aan het einde bent.
Omdat M2RNN zo'n groot "whiteboard" (matrix) heeft, kan het veel meer feiten tegelijk opslaan.

• Vergelijking: Stel je voor dat je een spreekbeurt moet geven. Een oude AI vergeet de eerste zin voordat je klaar bent. M2RNN heeft een geheugen dat groot genoeg is om je hele speech, plus de reacties van het publiek, en de notities van de zaalmanager tegelijk te onthouden.

C. Het werkt efficiënt op moderne computers

Een groot probleem met oude slimme modellen was dat ze de kracht van moderne videokaarten (GPU's) niet goed benutten. Ze maakten veel "leeg werk" (zoals het vullen van lege plekken in een rij).
M2RNN is ontworpen om perfect te passen in de "machines" die we vandaag hebben.

• Vergelijking: Het is alsof je een vrachtwagen bouwt die precies past in de laadruimte van een container. Geen ruimteverspilling, dus alles gaat sneller en goedkoper.

4. De "Hybride" aanpak: Het beste van twee werelden

De auteurs zeggen niet dat we alles moeten vervangen. Ze stellen voor om M2RNN te gebruiken als een krachtige toevoeging in bestaande modellen.

• De Analogie: Stel je een team voor. Je hebt een super-snelle renner (de huidige AI) en een super-slimme strateeg (M2RNN).
  • Als je alleen de renner hebt, ben je snel, maar maak je fouten in de strategie.
  • Als je alleen de strateeg hebt, ben je slim, maar te traag.
  • De oplossing: Je laat de renner de meeste weg lopen, maar op de lastige bochten (de moeilijke logica) schakel je de strateeg in.
  • Het resultaat: Je bent bijna net zo snel als de renner, maar je maakt veel minder fouten en lost veel complexere problemen op.

Conclusie

Dit paper toont aan dat we niet hoeven te kiezen tussen snelheid en slimheid. Met M2RNN kunnen we AI-modellen bouwen die:

1. Net zo snel zijn als de huidige topmodellen.
2. Veel beter zijn in het onthouden van details en het oplossen van moeilijke puzzels.
3. Efficiënter werken op de hardware die we al hebben.

Het is een belangrijke stap naar AI die niet alleen tekst kan voorspellen, maar echt kan "denken" en complexe taken kan uitvoeren, zonder dat het dagen duurt om te trainen.

Technische samenvatting

Probleemstelling

Hoewel Transformers de huidige standaard zijn voor taalmodellen vanwege hun parallelle trainbaarheid, hebben ze fundamentele beperkingen:

1. Expressiviteit: Transformers vallen binnen de complexiteitsklasse TC0, wat betekent dat ze bewezen onvoldoende expressief zijn voor complexe taak zoals entiteitsvolging (state tracking), code-uitvoering en het oplossen van permutatieproblemen (zoals de $S_5$-groep).
2. Efficiëntie: De kwadratische tijdscomplexiteit tijdens training en lineair groeiende geheugenvereisten tijdens inferentie maken ze minder schaalbaar voor zeer lange contexten.

Alternatieven zoals Lineaire RNN's (bijv. Mamba-2, Gated DeltaNet) en State Space Models (SSM's) bieden lineaire complexiteit en efficiënt inferentie, maar lijden onder twee grote tekortkomingen:

• Beperkte State Tracking: Ze zijn minder expressief dan niet-lineaire RNN's en kunnen complexe state-tracking taken niet oplossen.
• Slechte Context Retrieval: Door hun beperkte statengrootte (vaak vector-gebaseerd) kunnen ze slechte prestaties leveren bij "needle-in-a-haystack" taken en het terugvinden van specifieke informatie uit lange contexten.

Aan de andere kant hebben traditionele Niet-lineaire RNN's (zoals LSTM en GRU) wel de nodige expressiviteit, maar presteren ze slecht in taalmodellering en contextuele retrieval. Dit wordt toegeschreven aan hun kleine statengrootte (vector-gebaseerd) en inefficiënte hardwaregebruik (geen parallelisatie over de sequentielengte, veel padding-overhead op GPU's).

Methodologie: M2RNN

De auteurs introduceren M2RNN (Matrix-to-Matrix RNN), een niet-lineaire RNN-architectuur die de voordelen van niet-lineariteit combineert met de schaalbaarheid van matrix-gebaseerde staten.

Kerncomponenten:

1. Matrix-gebaseerde Hidden States: In plaats van een vector $h_t \in \mathbb{R}^d$, gebruikt M2RNN een matrix $H_t \in \mathbb{R}^{K \times V}$. Dit wordt bereikt via een outer product expansie ($k_t v_t^\top$), vergelijkbaar met lineaire attention, maar toegepast in een niet-lineaire recursie. Dit vergroot de statencapaciteit drastisch zonder een evenredige toename in het aantal parameters.
2. Niet-lineaire Transities: De state-update volgt de formule:
   $$Z_t = \tanh(H_{t-1}W + k_t v_t^\top)$$
   $$H_t = f_t H_{t-1} + (1 - f_t)Z_t$$
   Waarbij $W$ een transitie-matrix is en $f_t$ een vergeet-poort (forget gate).
3. Onafhankelijke Vergeet-Poort (Forget Gate): De poort $f_t$ is een functie van de input $x_t$ en is onafhankelijk van de vorige state $H_{t-1}$. Dit in tegenstelling tot LSTM/GRU, waardoor de poort parallel kan worden berekend over de sequentie, wat de trainingskosten verlaagt.
4. Hybride Architectuur: De auteurs testen M2RNN zowel als homogene laag als in hybride modellen waar recurrente lagen worden afgewisseld met Attention-lagen (bijv. 1 Attention-laag per 7 recurrente lagen). Ze onderzoeken ook het vervangen van slechts één of enkele lagen in bestaande hybride modellen (bijv. Gated DeltaNet) door M2RNN.

Systeem-Optimalisaties:

• Hardware-efficiëntie: De matrix-vormige recursie elimineert de noodzaak tot padding van de batch-dimensie om Tensor Cores te gebruiken (een probleem bij vector-RNN's zoals FlashRNN). Dit zorgt voor efficiënter gebruik van NVIDIA Tensor Cores zonder verspilde FLOPs.
• Distributed Training: Er worden twee strategieën voor Tensor Parallelism (TP) gepresenteerd:
  • Topology-aware: Groepeert value-heads per GPU, vereist geen extra communicatie maar koppelt parameters aan de TP-configuratie.
  • Topology-independent: Deelt parameters over GPUs met extra AllReduce-communicatie, maar behoudt het parameteraantal ongeacht het aantal GPUs.

Belangrijkste Bijdragen

1. Matrix-gebaseerde niet-lineaire RNN: Een nieuwe architectuur die de expressiviteit van niet-lineaire RNN's combineert met de schaalbaarheid van matrix-staten.
2. Bewijs van State Tracking: M2RNN lost perfect state-tracking taken op (zoals $S_3$ en $S_5$ permutaties) en generaliseert perfect naar onzichtbare sequentielengten, iets waar lineaire RNN's en zelfs sommige geavanceerde SSM's moeite mee hebben.
3. Hardware-geoptimaliseerde Kernels: Implementatie in Triton die Tensor Core-gebruik maximaliseert zonder padding-overhead, wat trainingsinefficiëntie oplost.
4. Hybride Strategie: Het bewijs dat het vervangen van slechts één recurrente laag in een hybride model door M2RNN leidt tot vergelijkbare prestaties als een volledig M2RNN-model, met minimale impact op de trainingsdoorvoer.

Resultaten

De auteurs trainden modellen van 410M parameters (dense) en 7B parameters (MoE) op 100B tokens.

• Taalmodellering:
  • In homogene settings presteert M2RNN vergelijkbaar met Mamba-2 en Gated DeltaNet, maar overtreft traditionele RNN's (LSTM/GRU) aanzienlijk.
  • In hybride settings (M2RNN + Attention) overtreft het modellen zoals Hybrid Mamba-2 en Hybrid Gated DeltaNet met 0,4–0,5 perplexity punten op de 7B MoE-modellen.
• State Tracking:
  • M2RNN bereikt >99,5% nauwkeur op de $S_3$ permutatietaken over contextlengtes tot 512, terwijl andere modellen (zoals Gated DeltaNet) falen bij generalisatie naar langere sequenties.
• In-Context Retrieval & Lange Context:
  • M2RNN lost het probleem van slechte retrieval op. In hybride modellen verbetert het de prestaties op LongBench met tot 8 punten ten opzichte van state-of-the-art hybride lineaire attention-architecturen.
  • Het vervangen van één laag in een Hybrid Gated DeltaNet model resulteert in significante verbeteringen op retrieval-benchmarks (zoals RULER en real-world datasets).
• Efficiëntie:
  • Hoewel M2RNN duurder is dan lineaire RNN's in een homogene setting, is de doorvoer van een hybride model met slechts één M2RNN-laag binnen 6% van het oorspronkelijke model, terwijl de nauwkeurigheid aanzienlijk stijgt.

Betekenis

Dit paper is significant omdat het de mythe doorbreekt dat niet-lineaire RNN's per definitie inefficiënt of slecht presterend zijn voor taalmodellering. De sleutel tot succes is niet de niet-lineariteit zelf, maar de statengrootte.

M2RNN biedt een nieuw paradigma voor schaalbare taalmodellen:

1. Het combineert de expressiviteit die nodig is voor complexe redeneertaken (code, logica) met de efficiëntie van lineaire RNN's.
2. Het biedt een praktische oplossing voor het "state-tracking" probleem dat Transformers en lineaire SSM's niet kunnen oplossen.
3. Het demonstreert dat het hybride gebruik van deze lagen (zeldzaam toegepast in een groter model) de beste balans biedt tussen prestatie en rekenefficiëntie, waardoor het een veelbelovende bouwsteen is voor de volgende generatie Large Language Models (LLM's).