Riepilogo Tecnico: Corrispondenza di Token Identificabili per Modelli del Mondo

1. Enunciato del Problema

I modelli del mondo basati su Transformer hanno dimostrato prestazioni elevate nell'apprendimento per rinforzo visivo (RL) trattando sequenze di stati e azioni passati come flussi di token per prevedere stati futuri. Tuttavia, questi modelli soffrono frequentemente di incoerenza temporale durante roll-out a lungo orizzonte. I modi di fallimento comuni includono duplicazione di oggetti, scomparsa e trasmutazione in oggetti diversi.

Gli autori identificano una causa primaria di questi errori: gli approcci esistenti trattano la previsione del fotogramma successivo puramente come un problema di generazione di token. Non riescono a modellare esplicitamente la corrispondenza tra token attraverso il tempo. In ambienti visivi realistici, i token in fotogrammi successivi spesso corrispondono alle stesse entità sottostanti che persistono e si muovono. Quando un modello è costretto a rigenerare ogni token da zero senza tracciare l'identità, fatica a preservare strutture persistenti, portando a una preservazione dell'identità fragile e ad errori cumulativi nel tempo.

2. Metodologia: Corrispondenza di Token Identificabili (ITC)

Il documento propone la Corrispondenza di Token Identificabili (ITC), un modello del mondo Transformer che formula la previsione del fotogramma successivo come un problema di inferenza probabilistica strutturata con variabili latenti di corrispondenza dei token.

Meccanismo Principale

Invece di generare ogni token del fotogramma successivo ex novo, l'ITC determina per ogni token del fotogramma successivo se dovrebbe essere:

1. Copiato da un token del fotogramma precedente (preservando l'identità).
2. Generato come un nuovo token (gestendo cambiamenti o nuovi contenuti).

Questa assegnazione è modellata come un problema di Trasporto Ottimale (OT).

Architettura Tecnica

1. Tokenizzazione & Incorporamento:

   • Le osservazioni visive vengono convertite in token discreti utilizzando un tokenizer di ricerca per patch basato sul vicino più prossimo.
   • Gli incorporamenti dei token sono arricchiti con Incorporamenti di Posizione Rotatoria 3D (RoPE). A differenza del RoPE 1D standard, questa formulazione codifica le coordinate temporali, verticali e orizzontali indipendentemente, permettendo al modello di catturare strutture relazionali localizzate sia nello spazio che nel tempo.

2. Soluzione del Trasporto Ottimale:

   • Il modello costruisce un grafo dove i vertici sorgente rappresentano i token dello stato precedente ($u$) e i token candidati dello stato successivo predetti dal transformer ($\tilde{p}$). I vertici di destinazione rappresentano i token finali dello stato successivo ($\hat{u}$).
   • Viene calcolata una matrice di affinità per valutare la relazione tra i token precedenti e i token predetti. Questa matrice include:
     • $A^{(prev)}$: Valuta l'affinità tra token precedenti e predizioni, penalizzata dalla distanza spaziale per imporre plausibilità fisica (ad esempio, gli oggetti non possono teletrasportarsi arbitrariamente lontano).
     • $A^{(gen)}$: Valuta il "bonus" di ammettere un token appena generato (segnalibro) invece di riutilizzare uno precedente, soggetto a una penalità costante.
   • Il risolutore (utilizzando l'algoritmo di Sinkhorn) calcola un piano di trasporto che minimizza il costo di spostamento della massa dalle sorgenti alle destinazioni.

3. Binarizzazione & Decodifica:

   • Il piano di trasporto continuo viene convertito in una matrice di assegnazione binaria utilizzando una procedura di binarizzazione greedy (Algoritmo 2). Questo garantisce una mappatura rigorosamente uno-a-uno.
   • Regola di Decodifica: Per ogni posizione nel fotogramma successivo, se il piano di trasporto assegna un token precedente, quel token viene copiato. Se il piano assegna un token "genera", il modello campiona dalla distribuzione predetta dal transformer.

Ciclo di Addestramento

Il modello del mondo e la politica vengono addestrati concorrentemente. Il modello del mondo genera "immaginazioni" (traiettorie sintetiche) che vengono utilizzate insieme ai dati reali dell'ambiente per aggiornare la politica tramite Ottimizzazione della Politica Prossima (PPO).

3. Contributi Chiave

• Formulazione della Corrispondenza dei Token: Il documento riformula la previsione del fotogramma successivo modellando esplicitamente la corrispondenza dei token come un problema di trasporto ottimale, separando i compiti di "riutilizzo" di entità persistenti e "generazione" di nuovi contenuti.
• Integrazione del Trasporto Ottimale: Introduce un meccanismo differenziabile (tramite Sinkhorn) ma discreto (tramite binarizzazione) per selezionare tra la copia di token precedenti e la generazione di nuovi, riducendo efficacemente le allucinazioni.
• Codifica Posizionale 3D: L'implementazione del RoPE 3D permette al modello di distinguere tra differenze temporali e spaziali, un componente necessario per tracciare oggetti in movimento nelle sequenze video.

4. Risultati Sperimentali

Gli autori hanno valutato l'ITC su quattro benchmark impegnativi: Craftax-classic, Craftax, MinAtar e Atari 100K.

• Craftax-classic: L'ITC ha ottenuto un ritorno del 72,5% e un punteggio del 35,6%, superando significativamente lo stato dell'arte precedente (Dedieu et al., 2025) che ha ottenuto un ritorno del 67,4% e un punteggio del 27,9%. Ha anche dimostrato una convergenza più rapida.
• Craftax: Su questo ambiente più complesso, l'ITC ha ottenuto un ritorno del 7,09% e un punteggio del 2,40%, superando le linee di base inclusi Simulus e Dedieu et al.
• MinAtar: L'ITC ha superato lo stato dell'arte precedente per il RL basato su modelli in tutti e quattro i giochi (Asterix, Breakout, Freeway, SpaceInvaders).
• Atari 100K: Applicato a una suite di 26 giochi, l'ITC ha ottenuto nuove prestazioni allo stato dell'arte nella Media Interquartile (IQM) di 1,092 e nel Gap di Ottimalità di 0,376, superando il modello del mondo basato su token Simulus.
• Accuratezza di Previsione: Su un dataset di 10.000 transizioni, l'ITC ha migliorato l'accuratezza complessiva di previsione del 3,39% rispetto alla linea di base, con un miglioramento notevole del 3,45% nelle transizioni che coinvolgono creature in movimento.
• Miglioramenti Qualitativi: L'analisi visiva ha confermato che l'ITC elimina artefatti comuni come la duplicazione e la scomparsa di oggetti osservati in modelli di base come DreamerV3 e Dedieu et al.
• Efficienza: Il metodo ha introdotto un sovraccarico trascurabile, aumentando il tempo totale di addestramento end-to-end solo del 2,8%.

5. Significato e Affermazioni

Il documento afferma che modellando esplicitamente la corrispondenza tra token attraverso il tempo, l'ITC risolve un problema diffuso nei modelli del mondo in cui gli agenti non riescono a mantenere l'identità degli oggetti su orizzonti lunghi. Gli autori dichiarano che questo approccio permette al modello di "riutilizzare selettivamente i token che hanno una forte corrispondenza con il fotogramma precedente", sfruttando così le somiglianze fotogramma-per-fotogramma invece di affidarsi esclusivamente al transformer per rigenerare ogni token.

Il significato risiede nel raggiungimento di prestazioni allo stato dell'arte in ambienti diversi (dai mondi aperti basati su griglia ai giochi Atari standard) mantenendo l'efficienza computazionale. Gli autori concludono che l'ITC previene efficacemente le incoerenze temporali come duplicazione e scomparsa, portando a immaginazioni di qualità superiore e a prestazioni della politica migliorate.