Technische Samenvatting: Identificeerbare Token-correspondentie voor Wereldmodellen

1. Probleemstelling

Wereldmodellen op basis van transformatoren hebben sterke prestaties aangetoond in visuele versterkende leer (RL) door sequenties van vorige toestanden en acties te behandelen als tokenstromen om toekomstige toestanden te voorspellen. Deze modellen lijden echter vaak aan temporele inconsistentie tijdens rollouts met een lange horizon. Veelvoorkomende faalmodi omvatten objectverdubbeling, verdwijning en transmutatie in verschillende objecten.

De auteurs identificeren een primaire oorzaak voor deze fouten: bestaande benaderingen behandelen voorspelling van de volgende frame puur als een token-generatieprobleem. Ze modelleren de correspondentie tussen tokens over de tijd niet expliciet. In realistische visuele omgevingen corresponderen tokens in opeenvolgende frames vaak met dezelfde onderliggende entiteiten die blijven bestaan en bewegen. Wanneer een model gedwongen wordt om elke token opnieuw te genereren zonder identiteit te volgen, worstelt het met het behoud van persistente structuren, wat leidt tot kwetsbare identiteitsbehoud en cumulerende fouten naarmate de tijd vordert.

2. Methodologie: Identificeerbare Token-correspondentie (ITC)

Het artikel stelt Identificeerbare Token-correspondentie (ITC) voor, een transformer-wereldmodel dat voorspelling van de volgende frame formuleert als een gestructureerd probabilistisch inferentieprobleem met latente token-correspondentievariabelen.

Kernmechanisme

In plaats van elke token van de volgende frame de novo te genereren, bepaalt ITC voor elke token van de volgende frame of deze moet worden:

1. Gekopieerd van een token in het vorige frame (identiteit behoudend).
2. Genereerd als een nieuwe token (omgaan met veranderingen of nieuwe inhoud).

Deze toewijzing wordt gemodelleerd als een Optimaal Transport (OT) probleem.

Technische Architectuur

1. Tokenisatie & Embedding:

   • Visuele waarnemingen worden omgezet in discrete tokens met behulp van een tokenizer voor zoekopdrachten naar patches op basis van de dichtstbijzijnde buur.
   • Token-embeddings worden verrijkt met 3D Rotary Position Embeddings (RoPE). In tegenstelling tot standaard 1D RoPE, codeert deze formulering temporele, verticale en horizontale coördinaten onafhankelijk, waardoor het model gelokaliseerde relationele structuren over zowel ruimte als tijd kan vastleggen.

2. Optimaal Transport Oplosser:

   • Het model construeert een graaf waarbij bronvertices vorige toestands-tokens ($u$) vertegenwoordigen en kandidaat-tokens van de volgende toestand die door de transformer zijn voorspeld ($\tilde{p}$). Bestemmingsvertices vertegenwoordigen de definitieve tokens van de volgende toestand ($\hat{u}$).
   • Er wordt een affiniteitsmatrix berekend om de relatie tussen vorige tokens en voorspelde tokens te scoren. Deze matrix omvat:
     • $A^{(prev)}$: Scoreert de affiniteit tussen vorige tokens en voorspellingen, bestraft door ruimtelijke afstand om fysieke plausibiliteit af te dwingen (bijv. objecten kunnen niet willekeurig ver teleporteren).
     • $A^{(gen)}$: Scoreert de "bonus" van het toelaten van een nieuw gegenereerd token (wildcard) in plaats van een vorige token opnieuw te gebruiken, onderworpen aan een constante straf.
   • De oplosser (met gebruik van het Sinkhorn-algoritme) berekent een transportplan dat de kosten minimaliseert van het verplaatsen van massa van bronnen naar bestemmingen.

3. Binarisatie & Decoding:

   • Het continue transportplan wordt omgezet in een binaire toewijzingsmatrix met behulp van een greedige binarisatieprocedure (Algoritme 2). Dit zorgt voor een strikte één-op-één mapping.
   • Decoderegels: Voor elke positie in de volgende frame, als het transportplan een vorige token toewijst, wordt die token gekopieerd. Als het plan een "genereer"-token toewijst, steekt het model een steekproef uit de voorspelde verdeling van de transformer.

Trainingslus

Het wereldmodel en het beleid worden gelijktijdig getraind. Het wereldmodel genereert "imaginaties" (synthetische trajecten) die samen met echte omgevingsdata worden gebruikt om het beleid bij te werken via Proximal Policy Optimization (PPO).

3. Belangrijkste Bijdragen

• Formulering van Token-correspondentie: Het artikel herformuleert voorspelling van de volgende frame door token-correspondentie expliciet te modelleren als een optimaal transportprobleem, waarbij de taken van "hergebruik" van persistente entiteiten en "genereren" van nieuwe inhoud worden gescheiden.
• Integratie van Optimaal Transport: Het introduceert een differentieerbare (via Sinkhorn) maar discrete (via binarisatie) mechanisme om te kiezen tussen het kopiëren van vorige tokens en het genereren van nieuwe, wat hallucinaties effectief reduceert.
• 3D Positionele Encodering: De implementatie van 3D RoPE stelt het model in staat om onderscheid te maken tussen temporele en ruimtelijke verschillen, een noodzakelijk onderdeel voor het volgen van bewegende objecten in videosequenties.

4. Experimentele Resultaten

De auteurs evalueerden ITC op vier uitdagende benchmarks: Craftax-classic, Craftax, MinAtar en Atari 100K.

• Craftax-classic: ITC behaalde een return van 72,5% en een score van 35,6%, aanzienlijk beter dan de vorige state-of-the-art (Dedieu et al., 2025) die 67,4% return en 27,9% score behaalde. Het toonde ook snellere convergentie.
• Craftax: Op deze complexere omgeving behaalde ITC een return van 7,09% en een score van 2,40%, en presteerde beter dan baselines, waaronder Simulus en Dedieu et al.
• MinAtar: ITC overtrof de vorige state-of-the-art voor modelgebaseerde RL in alle vier de spellen (Asterix, Breakout, Freeway, SpaceInvaders).
• Atari 100K: Toegepast op een reeks van 26 spellen, behaalde ITC nieuwe state-of-the-art prestaties in het Interkwartielgemiddelde (IQM) van 1,092 en een Optimaliteitskloof van 0,376, en presteerde beter dan het token-gebaseerde wereldmodel Simulus.
• Voorspellingnauwkeurigheid: Op een dataset van 10.000 transities verbeterde ITC de algehele voorspellingnauwkeurigheid met 3,39% ten opzichte van de baseline, met een opmerkelijke verbetering van 3,45% in transities die bewegende wezens betrekken.
• Kwalitatieve Verbeteringen: Visuele analyse bevestigde dat ITC veelvoorkomende artefacten elimineert, zoals objectverdubbeling en verdwijning die worden gezien in basismodellen zoals DreamerV3 en Dedieu et al.
• Efficiëntie: De methode introduceerde verwaarloosbare overhead, waardoor de totale end-to-end trainingstijd slechts met 2,8% toenam.

5. Betekenis en Claims

Het artikel beweert dat door de correspondentie tussen tokens over de tijd expliciet te modelleren, ITC een wijdverbreid probleem in wereldmodellen oplost waarbij agenten falen om objectidentiteit te behouden over lange horizons. De auteurs stellen dat deze benadering het model in staat stelt "selectief tokens opnieuw te gebruiken die een sterke correspondentie hebben met het vorige frame", waardoor frame-tot-frame-lijken worden benut in plaats van uitsluitend te vertrouwen op de transformer om elke token opnieuw te genereren.

De betekenis ligt in het behalen van state-of-the-art prestaties in diverse omgevingen (van roostergebaseerde open werelden tot standaard Atari-spellen) terwijl de rekenefficiëntie wordt behouden. De auteurs concluderen dat IFC temporele inconsistenties zoals verdubbeling en verdwijning effectief voorkomt, wat leidt tot imaginaties van hogere kwaliteit en verbeterde beleidsprestaties.