Tiny Recursive Reasoning with Mamba-2 Attention Hybrid

Dit onderzoek toont aan dat het vervangen van Transformer-blokken in een recursief redeneermodel door Mamba-2-hybride operatoren de redeneercapaciteit behoudt en de prestaties op abstracte taken zoals ARC-AGI-1 verbetert door een bredere dekking van correcte oplossingen.

De kern

De Kernvraag: Kan een slimme "snelle" motor een "denkende" auto worden?

Stel je voor dat je een heel slimme, maar kleine robot hebt (een kunstmatige intelligentie) die moet oplossen: "Hoe kom ik van A naar B in dit doolhof?" of "Wat is het volgende patroon in deze puzzel?"

Vroeger dachten mensen: "Om slim te zijn, moet de robot enorm groot zijn (miljarden parameters)." Maar recent onderzoek heeft laten zien dat je ook heel slim kunt zijn als je langzamer en dieper denkt, zelfs met een heel klein brein. Dit heet recursief redeneren.

In plaats van direct het antwoord te schreeuwen, laat je de robot in zijn hoofd een paar keer "nadenken" over zijn eigen gedachten, zonder dat hij hardop iets zegt. Hij verfijnt zijn ideeën in het geheim.

Het Experiment: De Motor wisselen

De onderzoekers (Wang en Reid van Intercom) hadden een heel succesvol klein model genaamd TRM. Dit model gebruikte een specifieke soort "motor" (de Transformer-blokken) om te denken. Deze motor is goed, maar misschien niet de enige optie.

Ze vroegen zich af: "Wat gebeurt er als we deze motor vervangen door een ander type, genaamd Mamba-2?"

De Analogie:

• De oude motor (Transformer): Denk aan een team van onderzoekers die allemaal tegelijk naar een whiteboard kijken en alles bespreken. Ze zijn goed in het zien van verbanden overal, maar het is een beetje rommelig en traag als het team groeit.
• De nieuwe motor (Mamba-2): Denk aan een zeer efficiënte, snelle solopionier die informatie één voor één verwerkt, alsof hij een lange rij poststukken doorloopt. Hij is razendsnel en slim in het onthouden van wat hij net zag, maar hij kijkt niet altijd terug naar alles wat hij al heeft gezien.

De onderzoekers wilden weten: Als we de "solopionier" (Mamba-2) in het "denkproces" van de robot stoppen, wordt hij dan dommer, of juist slimmer?

Wat deden ze?

Ze bouwden een hybride versie. Ze namen de snelle, efficiënte Mamba-2-motor, maar lieten hem samenwerken met een klein beetje van de oude "whiteboard"-techniek (Attention). Zo kregen ze het beste van twee werelden: snelheid én het vermogen om verbanden te zien.

Ze hielden de grootte van de robot exact hetzelfde (ongeveer 6,8 miljoen "hersencellen"), zodat het eerlijk was.

De Resultaten: Meer opties, net zo goed kiezen

Toen ze dit nieuwe model testten op de beroemde ARC-AGI-puzzels (een soort IQ-test voor AI), gebeurde er iets fascinerends:

1. Beter "veilig spelen" (Coverage): Het nieuwe model (Mamba-hybrid) bedacht meer verschillende goede oplossingen. Als je het model 100 keer laat denken, had hij in 100% van de gevallen de juiste oplossing ergens in zijn lijstje staan. De oude robot deed dit iets minder vaak.

   • Analogie: Stel je voor dat je een sleutelbos hebt. De oude robot had 10 sleutels, waarvan 1 goed was. De nieuwe robot had 15 sleutels, waarvan 1 goed was. De nieuwe robot heeft dus meer kans dat je de juiste sleutel vindt als je er eentje uitpakt.

2. Net zo goed kiezen (Selection): Het interessante deel is dat de nieuwe robot net zo goed was in het kies van de allerbeste oplossing als de oude robot.

   • Analogie: Hoewel de nieuwe robot meer sleutels had, was hij net zo slim in het direct de juiste sleutel uit het bos te halen als de oude robot. Hij werd niet "verward" door al die extra opties.

De winst: Op de officiële test (pass@2) was het nieuwe model 2% beter. Dat klinkt klein, maar in de wereld van AI-puzzels is dat een enorme sprong.

Waarom werkt dit?

Het onderzoek suggereert dat de "solopionier" (Mamba-2) op een andere manier door de puzzel "wandelt" dan de "teamwerker" (Transformer).

• De oude robot denkt op één manier en komt snel tot een conclusie.
• De nieuwe robot denkt op een iets andere manier, waardoor hij andere paden in het doolhof ontdekt die de oude robot over het hoofd zag.

Hierdoor heeft de nieuwe robot een diverser assortiment aan goede antwoorden. Hij is niet per se "slimmer" in het kiezen van het allerbeste antwoord, maar hij is veel beter in het verzamelen van goede ideeën.

De Belangrijkste Lessen

1. Klein kan slim zijn: Je hoeft geen gigantische computer te hebben om slim te zijn; als je slim denkt (recursief), werkt het ook met kleine modellen.
2. Mamba-2 is een goede kandidaat: De nieuwe, snellere technologie (Mamba-2) kan de plaats innemen van de oude technologie in denk-processen zonder de intelligentie te verliepen. Sterker nog, het maakt het model breder in zijn denkvermogen.
3. Stabiliteit is key: Ze ontdekten ook dat een specifieke technische instelling (het gebruik van "post-normalisatie") cruciaal was om te voorkomen dat de robot "dwaas" werd na het denken van te veel stappen.

Conclusie

Dit paper laat zien dat we de "motor" van onze denkende robots kunnen vervangen door snellere, efficiëntere versies. Het resultaat is een robot die net zo goed kiest, maar die veel meer goede ideeën in zijn hoofd heeft. Het is alsof je een slimme detective vervangt door een nog snellere detective die tegelijkertijd 10 verschillende theorieën kan bedenken, waardoor de kans dat hij de dader vindt, groter wordt.

Dit is een eerste stap naar het bouwen van AI's die niet alleen groter worden, maar slimmer denken met minder energie.

Technische samenvatting

Probleemstelling

Recent onderzoek naar recursieve redeneermodellen, zoals de Tiny Recursive Model (TRM), heeft aangetoond dat zeer kleine netwerken (rond de 7 miljoen parameters) sterke prestaties kunnen leveren op abstracte redeneertaken door middel van latente recursie. Dit proces omvat iteratieve verfijning in de verborgen representatieruimte zonder tussenliggende tokens te genereren.

De kernvraag die dit paper onderzoekt, is of de keuze van de operator binnen dit recursieve raamwerk kritiek is. TRM gebruikt standaard Transformer-blokken (met aandachtmechanismen). Echter, Mamba-2 (een State Space Model) heeft een inherente vorm van iteratieve verfijning via zijn state-space recurrentie ($h_t = a_t h_{t-1} + b_t x_t$). Het is onduidelijk of het vervangen van de Transformer-blokken door Mamba-2-hybride operatoren de redeneercapaciteit behoudt of zelfs verbetert, zonder de efficiëntie te verliezen die Mamba-2 biedt.

Methodologie

De auteurs presenteren een variant van TRM waarbij de Transformer-blokken worden vervangen door een Mamba-2 + Attention hybride operator, terwijl het aantal parameters exact gelijk blijft (parameter parity).

1. Architectuur:

   • De recursieve structuur van TRM wordt behouden: $H$ buitenste cycli en $L$ binnenste cycli met twee latente staten ($z_H$ en $z_L$).
   • De update-functie $f$ wordt gewijzigd. In plaats van alleen Attention-blokken, wordt een stack gebruikt bestaande uit: Mamba-2 $\rightarrow$ Mamba-2 $\rightarrow$ Attention $\rightarrow$ MLP.
   • Er wordt ook een variant getest met MLP-t (transposed MLP) in plaats van Attention voor kruis-punt communicatie, om de rol van sequentiële versus dichte mixing te isoleren.
   • Post-Normalisatie: Een cruciaal implementatiedetail is het gebruik van post-norm (RMSNorm na de residual add) in plaats van pre-norm. Dit voorkomt dat de grootte van de residual stream exponentieel groeit tijdens het "unrollen" van de recursie, wat stabiliteit garandeert en NaN-fouten voorkomt.

2. Parameter Matching:

   • TRM-attn (Baseline): 6.83M parameters.
   • TR-mamba2attn (Hybride): 6.86M parameters.
   • Dit zorgt ervoor dat verschillen in prestatie puur toe te schrijven zijn aan de operatorkeuze en niet aan modelgrootte.

3. Evalueringsprotocollen:

   • ARC-AGI-1: Abstracte redeneerpuzzels. De evaluatie gebruikt een uitgebreid augmentatieprotocol (dihedrale transformaties en kleurenperrmutaties) en meet Pass@K (of het juiste antwoord in de top-K voorspellingen zit). De officiële metric is Pass@2.
   • Sudoku & Maze: Extra taken om de robuustheid te testen op verschillende soorten ruimtelijke en constraint-gebaseerde problemen.

Belangrijkste Bijdragen (Contributions)

• C1: Het is het eerste Mamba-hybride model dat wordt toegepast op recursief latente redeneren.
• C2: Empirische validatie op ARC-AGI-1 die een verbetering van +2.0% toont op de officiële Pass@2-metric, met nog grotere voordelen bij hogere K-waarden.
• C3: Een analyse van het trade-off tussen dekking (coverage) en selectie. De hybride operator verbetert de diversiteit van kandidaatoplossingen zonder de kwaliteit van de top-1 selectie te verlagen.

Resultaten

1. ARC-AGI-1 Prestaties:

• Pass@2 (Officiële Metric): De hybride variant (TR-mamba2attn) scoort 45.88%, vergeleken met 43.88% voor de baselines (TRM-attn). Dit is een verbetering van +2.0%.
• Pass@K Trend: Het voordeel neemt toe bij hogere K-waarden. Bij Pass@100 is de hybride variant +4.75% beter.
• Pass@1: Beide modellen presteren bijna gelijk (40.50% vs 40.75%), wat aangeeft dat de hybride variant de beste oplossing even goed kan selecteren, maar wel meer alternatieve correcte oplossingen genereert.

2. Sudoku en Maze:

• Sudoku: Op kleine roosters (9x9) presteren modellen met dichte MLP-t-blokken het best (87.4%), wat suggereert dat constraint-satisfaction baat heeft bij all-to-all communicatie. De hybride Mamba-variant scoort hier iets lager (84.2%) dan de pure MLP-t, maar beter dan de pure Attention-variant.
• Maze (30x30): Op grotere ruimtelijke taken falen MLP-t-varianten volledig (0.0%), terwijl de hybride Mamba-Attention variant 80.6% bereikt (tegenover 60.8% voor de baseline). Dit benadrukt dat sequentiële verwerking (Mamba) essentieel is voor grotere ruimtelijke afhankelijkheden.

3. Analyse van Dekking vs. Selectie:
De auteurs analyseren de voorspellingsstatistieken en vinden:

• De hybride variant genereert meer unieke kandidaten per puzzel (339.5 vs 266.6) en heeft een hogere entropie in de stemverdeling.
• De baselines (TRM-attn) hebben een sterkere concentratie op de top-1 kandidaat (41.1% van de stemmen), wat leidt tot een hogere "margin" maar minder diversiteit.
• Conclusie: Mamba-2 zorgt voor een bredere verkenning van de oplossingsruimte (beter voor Pass@K), terwijl de Attention-variant iets selectiever is (behoudt Pass@1). Op moeilijke puzzels (waar het juiste antwoord zeldzaam is in de kandidatenpool) wint de hybride variant duidelijk.

Significantie en Conclusie

Dit paper valideert dat Mamba-2 hybride operatoren een levensvatbare en zelfs superieure keuze zijn binnen het ontwerpruimte van recursieve redeneermodellen.

• Efficiëntie en Capabiliteit: Het bewijst dat de inherente recurrentie van State Space Models (SSM) goed samenwerkt met de externe recursieve lus van TRM, zonder de redeneercapaciteit te degraderen.
• Nieuwe Richting: Het stelt een nieuwe richting voor voor "test-time compute" scaling. In plaats van alleen tokens te genereren, kan men de interne staat van een SSM gebruiken voor iteratieve verfijning.
• Toekomstperspectief: De resultaten suggereren dat de volgende stap het "internaliseren" van de recursieve lus in de SSM-state updates zelf zou kunnen zijn, waardoor de buitenste iteratie overbodig wordt en de efficiëntie verder toeneemt.

Kortom, het paper toont aan dat het combineren van de efficiëntie van Mamba-2 met de kracht van latente recursie leidt tot robuustere modellen die een bredere reeks mogelijke oplossingen kunnen vinden, wat essentieel is voor complexe abstracte redeneertaken.