Symbol-Equivariant Recurrent Reasoning Models

Deze paper introduceert Symbol-Equivariant Recurrent Reasoning Models (SE-RRM), een architectuur die symmetrie expliciet in het ontwerp verankert om prestaties en generalisatie op puzzels zoals Sudoku en ARC-AGI aanzienlijk te verbeteren met minder data-augmentatie en parameters.

De kern

Stel je voor dat je een heel lastig raadsel probeert op te lossen, zoals een Sudoku of een ingewikkeld labyrint. Voor een computer is dit vaak een nachtmerrie. Traditionele kunstmatige intelligentie (zoals de grote taalmodellen die we nu kennen) probeert dit op te lossen door simpelweg "te gissen" op basis van wat ze eerder hebben gelezen. Het is alsof je een Sudoku probeert op te lossen door blindelings nummers in te vullen en hoopt dat het klopt. Dat werkt vaak niet goed, vooral niet als het raadsel groter wordt.

De auteurs van dit paper hebben een slimme nieuwe aanpak bedacht: SE-RRM. Laten we dit uitleggen met een paar leuke vergelijkingen.

1. Het oude probleem: De "Gedetailleerde" Leerling

Stel je een oude manier van leren voor (de bestaande modellen). Als je een kind leert een Sudoku te spelen, leer je ze: "Dit is een 1, dit is een 2, dit is een 3..."
Maar als je het kind dan een nieuwe Sudoku geeft waarbij de '1' eigenlijk een 'rode bal' is en de '2' een 'blauwe bal', raakt het kind in de war. Het moet opnieuw leren dat de rode bal nu de rol van de '1' speelt.

In de computerwereld betekent dit dat de AI duizenden voorbeelden moet zien, waarbij ze de kleuren of cijfers telkens willekeurig verwisselen (data-augmentatie), zodat de AI leert dat het type symbool niet uitmaakt, maar alleen de relatie tussen ze. Dit is veel werk en kost veel tijd.

2. De nieuwe oplossing: De "Slimme" Leerling (SE-RRM)

De auteurs zeggen: "Waarom leren we de AI niet dat symbool A en symbool B eigenlijk uitwisselbaar zijn?"

Ze bouwen een nieuwe architectuur, de SE-RRM, die als een meester-kok werkt die niet kijkt naar de specifieke naam van een groente, maar naar de rol die de groente speelt in het gerecht.

• De Analogie: Stel je voor dat je een recept hebt. Een oude AI denkt: "Als ik aardappels gebruik, moet ik ze zo snijden. Als ik wortels gebruik, moet ik ze anders snijden."
• De nieuwe AI (SE-RRM) denkt: "Het maakt niet uit of het aardappel of wortel is; als het een 'knapperig blokje' is in dit gerecht, snijd ik het op dezelfde manier."

Dit noemen ze symmetrie. De AI begrijpt dat als je alle cijfers in een Sudoku verwisselt (bijvoorbeeld alle 1'en worden 2'en en alle 2'en worden 1'en), de oplossing precies hetzelfde blijft, alleen met andere cijfers. De nieuwe AI is zo ontworpen dat dit van nature klopt, zonder dat ze duizenden voorbeelden nodig heeft om dit te leren.

3. Hoe werkt het technisch (in simpele taal)?

Stel je een raamwerk voor met drie dimensies:

1. Waar zit het? (De positie in het raster).
2. Wat is het? (Het symbool of de kleur).
3. Hoe denken we erover na? (De interne berekening).

De oude modellen keken alleen naar "Waar" en "Wat" als losse lijnen. De nieuwe SE-RRM kijkt naar een 3D-ruimte waar ze "Waar" en "Wat" tegelijk kunnen vergelijken. Ze gebruiken een slimme techniek (Self-Attention) om te zeggen: "Kijk eens, deze 'rode blok' hier en die 'blauwe blok' daar spelen dezelfde rol in het patroon."

Hierdoor hoeft de AI niet te leren dat rood en blauw uitwisselbaar zijn; ze weten het al omdat hun bouwplaat (de architectuur) dit zo regelt.

4. Wat levert dit op? (De resultaten)

De auteurs hebben dit getest op drie dingen: Sudoku, ARC-AGI (een soort visuele puzzel) en labyrinten.

• Sudoku: De oude modellen konden alleen 9x9 Sudoku's goed oplossen. Als je ze een 4x4 of een gigantische 16x16 Sudoku gaf, faalden ze. De nieuwe SE-RRM kon dit direct oplossen, zelfs zonder dat ze er ooit een 16x16 hadden gezien! Ze hebben de regels geleerd, niet alleen de cijfers.
• Minder data: Voor de visuele puzzels (ARC-AGI) hadden de oude modellen duizenden variaties nodig om te leren. De nieuwe AI deed het met slechts 8 variaties. Dat is als het verschil tussen een student die 1000 oefentoetsen moet maken en een genie dat het concept in één keer snapt.
• Kleiner en sneller: De nieuwe AI is veel kleiner (slechts 2 miljoen parameters) dan de concurrenten, maar werkt beter.

Conclusie

Kortom, dit paper introduceert een manier om computers te laten denken zoals een mens die de logica van een spel begrijpt, in plaats van alleen de tekens te onthouden.

Het is alsof je iemand leert een taal spreken door de grammatica te geven, in plaats van duizenden zinnen uit het hoofd te laten leren. Hierdoor wordt de AI slimmer, heeft ze minder data nodig en kan ze zich aanpassen aan nieuwe, grotere problemen die ze nog nooit heeft gezien. Het is een grote stap in de richting van computers die echt kunnen "redeneren" in plaats van alleen maar te "gissen".

Technische samenvatting

Probleemstelling

Redeneringsproblemen zoals Sudoku, ARC-AGI en mazes vormen een grote uitdaging voor neurale netwerken, waaronder Large Language Models (LLMs). Bestaande architecturen voor gestructureerde probleemoplossing, zoals Recurrent Reasoning Models (RRM's) (bijv. HRM en TRM), hebben moeite met symmetrieën in symbolen.

• In veel van deze problemen zijn symbolen (bijv. cijfers in Sudoku of kleuren in ARC-AGI) onderling uitwisselbaar. Een oplossing moet invariant blijven als de symbolen worden omgedraaid (bijv. alle '1's vervangen door '2's en vice versa).
• Bestaande RRM's behandelen deze symmetrieën niet expliciet in hun architectuur. Ze vertrouwen in plaats daarvan op kostbare data-augmentatie (het trainen met vele permutaties van symbolen) om dit te leren. Dit verhoogt de steekproefcomplexiteit, beperkt de generalisatie naar ongeziene symbolen en maakt het model minder robuust.
• LLM's presteren vaak slecht op deze strikt beperkte, combinatorische taken en kunnen moeilijk extrapoleren naar grotere probleemgroottes of nieuwe symbolen.

Methodologie: SE-RRM

De auteurs introduceren Symbol-Equivariant Recurrent Reasoning Models (SE-RRM). Dit is een nieuwe klasse van recurrente architecturen die permutatie-equivariantie op het niveau van de architectuur afdwingen.

Kernprincipes:

1. Drie-dimensionale Representatie:

   • In tegenstelling tot standaard RRM's die een 2D-matrix gebruiken (posities x features), introduceert SE-RRM een derde dimensie: symbolen.
   • De invoer wordt gemapt naar een tensor van formaat $D \times I \times K$, waarbij $D$ het aantal features is, $I$ het aantal posities, en $K$ het totale aantal unieke symbolen (kleuren/cijfers).
   • Alle symbolen delen dezelfde embedding-vector (behalve speciale tokens zoals 'mask' of 'unknown'), wat zorgt voor symmetrie.

2. Architectuur van het Blok:

   • Het model gebruikt een vaste punt-iteratie (fixed-point iteration) waarbij een recurrente toestand $Z_t$ wordt bijgewerkt.
   • Het SE-RRM-blok bevat twee soorten Self-Attention lagen:
     • $T_{D,I}$: Zelf-attention over de positie-dimensie (zoals in standaard Transformers).
     • $T_{D,K}$: Zelf-attention over de symbool-dimensie. Dit zorgt ervoor dat het model de relaties tussen symbolen leert, ongeacht hun specifieke waarde.
   • Na de attention lagen volgen per-punt MLP's (SwiGLU) en normalisatielagen (RMSNorm).

3. Equivariantie:

   • Door de architectuur zo te ontwerpen dat de operaties uniform worden toegepast over de symbool-dimensie, is het model equivariant onder permutaties van symbolen. Als de invoer symbolen verwisselt, verandert de uitvoer op een voorspelbare, equivalente manier.
   • Dit elimineert de noodzaak om het model te trainen op alle mogelijke permutaties van symbolen.

4. Extrapolatie:

   • Omdat symbolen niet als discrete, unieke IDs worden gecodeerd maar via een gedeelde embedding, kan het model theoretisch extrapoleren naar grotere sets symbolen (bijv. van 9x9 naar 16x16 Sudoku) zonder nieuwe parameters te hoeven leren voor de nieuwe symbolen.

Belangrijkste Bijdragen

1. Introductie van SE-RRM: Een nieuwe architectuur die symmetrie expliciet encodeert in plaats van impliciet te leren via data-augmentatie.
2. Efficiëntie: Het model presteert beter met aanzienlijk minder data-augmentatie en veel minder parameters (slechts 2 miljoen parameters vergeleken met 7M bij TRM en 27M bij HRM).
3. Generalisatie en Extrapolatie: Het bewijst dat architecturale symmetrie het mogelijk maakt om te generaliseren naar probleemgroottes en symbolen die niet tijdens het trainingstijd zijn gezien (bijv. training op 9x9, testen op 4x4, 16x16 en 25x25 Sudoku).
4. Robuustheid: Het model toont verbeterde robustheid onder distributieverschuivingen, wat cruciaal is voor praktische toepassingen in planning en diagnose.

Resultaten

De auteurs evalueren SE-RRM op drie benchmarks: Sudoku, ARC-AGI en Maze.

• Sudoku:

  • 9x9: SE-RRM presteert aanzienlijk beter dan HRM en TRM (>11% verbetering in Fully Solved Rate).
  • 4x4 (Mini): SE-RRM extrapoleert perfect (95.46% FSR), terwijl HRM en TRM volledig falen (0% FSR) omdat ze de onderliggende regels niet hebben geleerd zonder specifieke data-augmentatie.
  • 16x16 en 25x25: Bestaande modellen kunnen deze niet oplossen zonder hertraining (ze kunnen geen nieuwe symbolen verwerken). SE-RRM bereikt een aanzienlijk hogere nauwkeurigheid (GPA) dan willekeur, wat aantoont dat het de regels heeft overgedragen, hoewel het niet perfect is op deze zeer grote schaal.
  • Test-time Scaling: SE-RRM profiteert het meest van het verhogen van het aantal iteratiestappen tijdens inferentie.

• ARC-AGI (1 en 2):

  • SE-RRM presteert concurrerend met de state-of-the-art (HRM en TRM) maar gebruikt slechts 8 data-augmentaties per voorbeeld in plaats van de duizenden die nodig zijn voor andere modellen.
  • Dit bevestigt dat het expliciet coderen van symmetrie de leerbaarheid en generalisatie verbetert.

• Maze:

  • Voor mazes is symmetrie van symbolen (bijv. muren vs. startpunt) niet wenselijk. De auteurs toonden aan dat SE-RRM ook hier werkt door de symbool-equivariantie te "breken" (door unieke embeddings toe te staan), wat resulteerde in de beste prestaties (88.8% FSR).

Betekenis en Conclusie

Dit paper toont aan dat het expliciet inbouwen van inductieve bias (in dit geval symmetrie) in neurale netwerken essentieel is voor effectief redeneren op gestructureerde problemen.

• Efficiëntie: Het reduceert de afhankelijkheid van enorme datasets en data-augmentatie.
• Schaalbaarheid: Het biedt een pad naar het oplossen van problemen met variabele grootte en complexiteit, waar traditionele LLM's en RRM's vaak vastlopen.
• Toekomst: SE-RRM biedt een compacte, parameter-efficiënte alternatief voor zware LLM's in domeinen die vereisen dat het model strikte logische regels en constraints respecteert, zoals planning, diagnose en wiskundig redeneren.

De code is beschikbaar gesteld, wat de reproduceerbaarheid en verdere ontwikkeling van deze architectuur faciliteert.