Riepilogo Tecnico: MapFormers

Enunciato del Problema

I sistemi di intelligenza artificiale basati su Transformer, nonostante le prestazioni impressionanti su compiti in-distribution, mostrano una fragilità significativa quando affrontano scenari fuori distribuzione (OOD). Spesso falliscono nel generalizzare a situazioni nuove che esseri umani e animali gestiscono con facilità. Gli autori attribuiscono questo divario alla mancanza di bias strutturali che permettono ai cervelli biologici di costruire mappe cognitive: modelli interni che disaccoppiano la struttura astratta e invariante di un ambiente dal suo contenuto specifico e variabile.

I tentativi esistenti di modellare le mappe cognitive, come la Macchina Tolman-Eichenbaum (TEM) che utilizza matrici dipendenti dall'azione nelle RNN, soffrono di due limitazioni critiche:

1. Mancanza di Scalabilità: Richiedono una moltiplicazione di matrici sequenziale per aggiornare le posizioni, impedendo l'elaborazione parallela.
2. Azioni Predefinite: Spesso richiedono che l'insieme delle azioni (ad esempio, le direzioni cardinali) sia codificato a livello di codice anziché appreso dai dati.

Inoltre, le codifiche posizionali standard (come RoPE) sono statiche o si basano su euristiche che non supportano la composizione dinamica delle azioni necessaria per l'integrazione del percorso in ambienti sconosciuti.

Metodologia

Il documento introduce MapFormers, una classe di architetture Transformer progettate per apprendere mappe cognitive ed eseguire l'integrazione del percorso tramite la previsione del token successivo senza supervisione. L'innovazione fondamentale risiede nell'uso di matrici dipendenti dall'input per aggiornare le codifiche posizionali, derivate dalla teoria dei gruppi di Lie.

Quadro Teorico

• Formalismo dei Gruppi di Lie: Gli autori modellano le azioni come elementi di un gruppo di Lie (in particolare gruppi compatti come $SO(2)$ per le rotazioni). Invece di apprendere una matrice separata per ogni possibile azione, il modello apprende i generatori dell'algebra di Lie.
• Integrazione del Percorso tramite Mappature Esponenziali: Le azioni sono rappresentate come generatori infinitesimali nell'algebra di Lie. L'integrazione del percorso (comporre una sequenza di azioni) viene eseguita sommando questi generatori nell'algebra (accumulo lineare) e applicando quindi una singola mappatura esponenziale per ottenere la trasformazione finita nel gruppo di Lie.
  • Per mappe commutative (ad esempio, navigazione su griglia 2D), ciò permette di calcolare l'integrazione del percorso in parallelo tramite somma cumulativa, seguita da una singola esponenziale di matrice.
  • Questo approccio riduce il bias induttivo dal numero di azioni alla complessità strutturale dell'algebra.

Varianti Architettoniche

Il documento propone due varianti di MapFormers, unificando la codifica posizionale assoluta e relativa con modelli di memoria distinti:

1. MapEM (Memoria Episodica):

   • Utilizza embedding posizionali assoluti per rappresentare una mappa cognitiva.
   • Disaccoppia il contenuto (osservazioni) e la struttura (posizione) in popolazioni neurali separate.
   • Calcola l'attenzione sulla congiunzione di contenuto e posizione ($Q_G = \text{vec}(Q^T \cdot P)$), agendo efficacemente come una maschera di attenzione dove le codifiche posizionali si aggiornano indipendentemente dal contenuto.
   • Adatto a compiti che richiedono un richiamo robusto di stati specifici.

2. MapWM (Memoria di Lavoro):

   • Utilizza codifica posizionale relativa (analoga a RoPE ma dinamica).
   • Rappresenta la struttura implicitamente ruotando l'attività neurale (chiavi e query) in sottospazi distinti tramite matrici di rotazione dipendenti dall'input ($R_{\theta PI}$).
   • L'angolo di rotazione $\theta$ è derivato da una proiezione a rango basso del token di input, permettendo al modello di apprendere quali token agiscono come azioni (aggiornando la posizione) e quali sono osservazioni (lasciando la posizione invariata).

Meccanismo Dipendente dall'Input

Entrambe le architetture impiegano una proiezione a due stadi per determinare la "durata di integrazione" ($\Delta_t$) per ogni token:

1. Proiezione Interna: Mappa il token di input a una rappresentazione di azione a bassa dimensionalità.
2. Proiezione Esterna: Proietta questa rappresentazione su una durata di integrazione per blocco.
   L'angolo di rotazione finale è $\theta_t = \omega \odot \Delta_t$, dove $\omega$ è un vettore di velocità angolare apprendibile. L'integrazione del percorso è calcolata come la somma cumulativa di questi angoli ($\theta_{0 \to t} = \text{cumsum}(\theta_t)$).

Contributi Chiave

1. Architettura MapFormers: Una variante Transformer innovativa che apprende mappe cognitive disaccoppiando la struttura dal contenuto utilizzando codifiche posizionali dipendenti dall'input basate su generatori dell'algebra di Lie.
2. Integrazione del Percorso in Parallelo: Sfruttando gruppi di Lie commutativi (ad esempio, $SO(2)^k$), i modelli raggiungono un'integrazione del percorso parallela, superando il collo di bottiglia sequenziale dei precedenti modelli di mappa cognitiva basati su RNN.
3. Unificazione dei Modelli di Memoria: Il lavoro fornisce un collegamento teorico tra codifica posizionale assoluta/relativa e memoria episodica/di lavoro, mostrando come entrambi possano essere derivati dagli stessi principi teoretico-lieiani.
4. Scalabilità a Mappe Non Commutative: Il framework è esteso a problemi non commutativi (ad esempio, alberi genealogici) utilizzando più generatori, sebbene ciò richieda un'integrazione del percorso sequenziale.

Risultati Sperimentali

Gli autori hanno valutato MapFormers su tre compiti formali mirati a capacità cognitive distinte:

1. Copia Selettiva (Gating):

   • Compito: Copiare una sequenza ignorando i token distrattori.
   • Risultato: MapFormers (sia EM che WM) hanno raggiunto una generalizzazione OOD perfetta, mentre le baseline standard (RoPE, TAPE) hanno fallito. Ciò dimostra la capacità di apprendere meccanismi di gating dinamici.

2. Navigazione su Griglia (1D e 2D):

   • Compito: Prevedere le osservazioni in una sequenza di azioni alternate e osservazioni casuali, richiedendo al modello di inferire la topologia latente e integrare la posizione del percorso.
   • Risultato: MapFormers hanno raggiunto un'accuratezza quasi perfetta su sequenze OOD (lunghezza e densità variabili), mentre le baseline (inclusi CoPE e PathAtt) non sono riuscite a generalizzare.
   • Analisi Neurale: I modelli hanno appreso a proiettare i token di osservazione su una rotazione zero (identità) e i token di azione su direzioni di rotazione specifiche, disaccoppiando efficacemente il contenuto dalla struttura.

3. Dyck-2 (Gerarchie Annidate):

   • Compito: Prevedere continuazioni valide di sequenze di parentesi, richiedendo il tracciamento di uno stack.
   • Risultato: MapFormers hanno generalizzato perfettamente a sequenze più lunghe e profonde rispetto a quelle viste durante l'addestramento utilizzando un solo strato. Le baseline richiedevano più strati e fallivano comunque su dati OOD.
   • Meccanismo: I modelli hanno incorporato lo stack direttamente nella codifica posizionale utilizzando dimensioni ortogonali per diversi tipi di parentesi e direzioni opposte per le parentesi di apertura/chiusura.

4. Scalabilità al Linguaggio Naturale:

   • Compito: Pre-addestramento su OpenWebText.
   • Risultato: Un modello MapWM a 12 strati ha mostrato miglioramenti coerenti della perplessità rispetto a una baseline RoPE ($18.79$ contro $19.14$) e una migliore estrapolazione della lunghezza. Tuttavia, sul benchmark sintattico BLIMP, MapWM non ha mostrato guadagni significativi rispetto a RoPE, suggerendo che l'integrazione del percorso commutativa da sola potrebbe non essere sufficiente per tutta la modellazione sintattica.

Significato e Affermazioni

Il documento afferma che MapFormers forniscono un bias strutturale critico che abilita una generalizzazione sistematica OOD catturando strutture latenti tramite integrazione del percorso invece di affidarsi a euristiche.

• Robustezza: I modelli dimostrano che disaccoppiare le relazioni astratte dal contenuto permette una generalizzazione robusta a lunghezze di sequenza, profondità e densità ambientali non viste.
• Scalabilità: L'uso del calcolo parallelo su mappe commutative permette a questi principi di scalare a domini su larga scala e del mondo reale, come dimostrato dai risultati sul linguaggio naturale.
• Plausibilità Biologica: L'approccio colma il divario tra le mappe cognitive biologiche (che fattorizzano "cosa" da "dove") e il deep learning moderno, offrendo una via verso sistemi di IA in grado di ragionare su problemi nuovi comprendendo la geometria sottostante del loro ambiente.

Gli autori rimangono modesti riguardo alle limitazioni, notando che l'integrazione del percorso non commutativa non è parallelizzabile e che il vincolo commutativo attuale potrebbe limitare la modellazione di certe strutture sintattiche complesse nel linguaggio. Suggeriscono che il lavoro futuro potrebbe esplorare architetture ibride o estensioni non commutative per un ragionamento relazionale più complesso.