Technische Samenvatting: MapFormers

Probleemstelling

Huidige op Transformer gebaseerde AI-systemen vertonen, ondanks indrukwekkende prestaties op taken binnen de verdelingsgrenzen (in-distribution), aanzienlijke broosheid wanneer ze geconfronteerd worden met scenario's buiten de verdelingsgrenzen (out-of-distribution, OOD). Ze falen vaak in het generaliseren naar nieuwe situaties die mensen en dieren met gemak hanteren. De auteurs schrijven dit gat toe aan een gebrek aan structurele bias die biologische hersenen toestaat om cognitieve kaarten te construeren: interne modellen die de abstracte, invariant structuur van een omgeving ontkoppelen van zijn specifieke, variabele inhoud.

Bestaande pogingen om cognitieve kaarten te modelleren, zoals de Tolman-Eichenbaum Machine (TEM) die actie-afhankelijke matrices in RNN's gebruikt, lijden onder twee kritieke beperkingen:

1. Gebrek aan Schaalbaarheid: Ze vereisen sequentiële matrixvermenigvuldiging om posities bij te werken, wat parallelle verwerking verhindert.
2. Vooraf Gespecificeerde Acties: Ze vereisen vaak dat de set acties (bijv. kardinaalrichtingen) hardgecodeerd is in plaats van uit data te worden geleerd.

Bovendien zijn standaard positionele coderingen (zoals RoPE) statisch of gebaseerd op heuristieken die de dynamische compositie van acties, vereist voor padintegratie in onbekende omgevingen, niet ondersteunen.

Methodologie

Het paper introduceert MapFormers, een klasse van Transformer-architecturen die zijn ontworpen om cognitieve kaarten te leren en padintegratie uit te voeren via voorspelling van de volgende token zonder toezicht. De kerninnovatie ligt in het gebruik van invoer-afhankelijke matrices om positionele coderingen bij te werken, afgeleid van Lie-groepentheorie.

Theoretisch Kader

• Lie-Groep Formalisme: De auteurs modelleren acties als elementen van een Lie-groep (specifiek compacte groepen zoals $SO(2)$ voor rotaties). In plaats van een aparte matrix te leren voor elke mogelijke actie, leert het model Lie-algebra-generatoren.
• Padintegratie via Exponentiële Afbeeldingen: Acties worden weergegeven als infinitesimale generatoren in de Lie-algebra. Padintegratie (het compositen van een reeks acties) wordt uitgevoerd door deze generatoren in de algebra op te tellen (lineaire accumulatie) en vervolgens een enkele exponentiële afbeelding toe te passen om de eindige transformatie in de Lie-groep te verkrijgen.
  • Voor commutatieve kaarten (bijv. navigatie op een 2D-rooster) maakt dit padintegratie parallel berekenbaar via cumulatieve sommatie, gevolgd door een enkele matrix-exponentiële.
  • Deze benadering verlaagt de inductieve bias van het aantal acties naar de structurele complexiteit van de algebra.

Architecturale Variaties

Het paper stelt twee variaties van MapFormers voor, die absolute en relatieve positionele codering verenigen met distincte geheugenmodellen:

1. MapEM (Episodisch Geheugen):

   • Gebruikt absolute positionele embedding om een cognitieve kaart te representeren.
   • Ontkoppelt inhoud (observaties) en structuur (positie) in aparte neurale populaties.
   • Berekent aandacht op de conjunctie van inhoud en positie ($Q_G = \text{vec}(Q^T \cdot P)$), wat effectief fungeert als een aandachtsmasker waarbij positionele coderingen onafhankelijk van inhoud worden bijgewerkt.
   • Geschikt voor taken die een robuuste terugroeping van specifieke toestanden vereisen.

2. MapWM (Werkgeheugen):

   • Gebruikt relatieve positionele codering (analoog aan RoPE, maar dynamisch).
   • Representeren structuur impliciet door neurale activiteit (sleutels en queries) te roteren naar distincte deelruimten via invoer-afhankelijke rotatiematrices ($R_{\theta PI}$).
   • De rotatiehoek $\theta$ is afgeleid van een laag-rang projectie van de invoer-token, waardoor het model kan leren welke tokens fungeren als acties (die de positie bijwerken) en welke observaties zijn (die de positie onveranderd laten).

Invoer-Afhankelijk Mechanisme

Beide architecturen maken gebruik van een twee-traps projectie om de "integratieduur" ($\Delta_t$) voor elke token te bepalen:

1. Binnenprojectie: Mappt de invoer-token naar een laag-dimensionale actie-representatie.
2. Buitenprojectie: Projecteert deze representatie naar een integratieduur per blok.
   De uiteindelijke rotatiehoek is $\theta_t = \omega \odot \Delta_t$, waarbij $\omega$ een leerbare hoeksnelheidsvector is. Padintegratie wordt berekend als de cumulatieve som van deze hoeken ($\theta_{0 \to t} = \text{cumsum}(\theta_t)$).

Belangrijkste Bijdragen

1. MapFormers Architectuur: Een nieuwe Transformer-variant die cognitieve kaarten leert door structuur te ontkoppelen van inhoud met behulp van invoer-afhankelijke positionele coderingen gebaseerd op Lie-algebra-generatoren.
2. Parallelle Padintegratie: Door gebruik te maken van commutatieve Lie-groepen (bijv. $SO(2)^k$) bereiken de modellen parallelle padintegratie, waarmee ze de sequentiële bottleneck van eerdere op RNN gebaseerde cognitieve kaartmodellen overwinnen.
3. Unificatie van Geheugenmodellen: Het werk biedt een theoretische link tussen absolute/relatieve positionele codering en episodisch/werkgeheugen, en toont aan hoe beide kunnen worden afgeleid uit dezelfde Lie-theoretische principes.
4. Schaalbaarheid naar Niet-Commutatieve Kaarten: Het kader wordt uitgebreid naar niet-commutatieve problemen (bijv. stambomen) met behulp van meerdere generatoren, hoewel dit sequentiële padintegratie vereist.

Experimentele Resultaten

De auteurs hebben MapFormers geëvalueerd op drie formele taken die verschillende cognitieve capaciteiten richten:

1. Selectief Kopiëren (Gating):

   • Taak: Kopieer een reeks terwijl je afleidende tokens negeert.
   • Resultaat: MapFormers (zowel EM als WM) bereikten perfecte OOD-generalisatie, terwijl standaard baselines (RoPE, TAPE) faalden. Dit demonstreert het vermogen om dynamische gating-mechanismen te leren.

2. Roosternavigatie (1D en 2D):

   • Taak: Voorspel observaties in een reeks van afwisselende acties en willekeurige observaties, waarbij het model de latente topologie moet infereren en de positie moet integreren.
   • Resultaat: MapFormers bereikten bijna perfecte nauwkeurigheid op OOD-reeksen (variërende lengte en dichtheid), terwijl baselines (inclusief CoPE en PathAtt) faalden in generalisatie.
   • Neurale Analyse: De modellen leerden observatietokens te projecteren naar nulrotatie (identiteit) en actietokens naar specifieke rotatierichtingen, waardoor inhoud effectief werd ontkoppeld van structuur.

3. Dyck-2 (Geneste Hiërarchieën):

   • Taak: Voorspel geldige voortzettingen van haakjesreeksen, waarbij het bijhouden van een stack vereist is.
   • Resultaat: MapFormers generaliseerden perfect naar langere en diepere reeksen dan gezien tijdens training, gebruikmakend van slechts één laag. Baselines vereisten meerdere lagen en faalden toch op OOD-gegevens.
   • Mechanisme: De modellen embedden de stack direct in de positionele codering door orthogonale dimensies te gebruiken voor verschillende haakjestypes en tegenovergestelde richtingen voor open/sluit-haakjes.

4. Natuurlijke Taal Schaling:

   • Taak: Pretraining op OpenWebText.
   • Resultaat: Een 12-laags MapWM-model toonde consistente verbeteringen in perplexiteit ten opzichte van een RoPE-baseline ($18,79$ versus $19,14$) en betere lengte-extrapolatie. Echter, op de BLIMP syntactische benchmark toonde MapWM geen significante winst ten opzichte van RoPE, wat suggereert dat commutatieve padintegratie alleen mogelijk niet voldoende is voor alle syntactische modellering.

Betekenis en Claims

Het paper beweert dat MapFormers een kritieke structurele bias bieden die systematische OOD-generalisatie mogelijk maakt door latente structuren te vangen via padintegratie in plaats van te vertrouwen op heuristieken.

• Robuustheid: De modellen tonen aan dat het ontkoppelen van abstracte relaties van inhoud robuuste generalisatie mogelijk maakt naar ongeziene reekslengtes, dieptes en omgevingsdichtheden.
• Schaalbaarheid: Het gebruik van parallelle berekening op commutatieve kaarten stelt deze principes in staat te schalen naar grote, real-world domeinen, zoals bewezen door de resultaten voor natuurlijke taal.
• Biologische Plausibiliteit: De aanpak overbrugt het gat tussen biologische cognitieve kaarten (die "wat" factoriseren van "waar") en modern deep learning, en biedt een weg naar AI-systemen die kunnen redeneren door nieuwe problemen door de onderliggende geometrie van hun omgeving te begrijpen.

De auteurs blijven bescheiden wat betreft de beperkingen, en merken op dat niet-commutatieve padintegratie niet paralleliseerbaar is en dat de huidige commutatieve beperking de modellering van bepaalde complexe syntactische structuren in taal kan beperken. Zij suggereren dat toekomstig werk hybride architecturen of niet-commutatieve extensies voor complexer relationeel redeneren kan verkennen.