Technische Zusammenfassung: Identifizierbare Token-Korrespondenz für Weltmodelle

1. Problemstellung

Transformer-basierte Weltmodelle haben starke Leistungen im visuellen Reinforcement Learning (RL) gezeigt, indem sie Sequenzen vergangener Zustände und Aktionen als Token-Ströme behandeln, um zukünftige Zustände vorherzusagen. Diese Modelle leiden jedoch häufig unter temporaler Inkonsistenz während Rollouts mit langem Horizont. Häufige Fehlermodi umfassen die Verdopplung von Objekten, das Verschwinden von Objekten und deren Verwandlung in andere Objekte.

Die Autoren identifizieren eine Hauptursache für diese Fehler: Bestehende Ansätze behandeln die Vorhersage des nächsten Bildes rein als Problem der Token-Generierung. Sie versäumen es, die Korrespondenz zwischen Token über die Zeit hinweg explizit zu modellieren. In realistischen visuellen Umgebungen entsprechen Token in aufeinanderfolgenden Bildern oft denselben zugrunde liegenden Entitäten, die persistieren und sich bewegen. Wenn ein Modell gezwungen wird, jeden Token ohne Verfolgung der Identität neu zu generieren, hat es Schwierigkeiten, persistente Strukturen zu bewahren, was zu einer fragilen Identitätserhaltung und sich über die Zeit aufkumulierenden Fehlern führt.

2. Methodik: Identifizierbare Token-Korrespondenz (ITC)

Die Arbeit schlägt Identifiable Token Correspondence (ITC) vor, ein Transformer-Weltmodell, das die Vorhersage des nächsten Bildes als strukturiertes probabilistisches Inferenzproblem mit latenten Token-Korrespondenzvariablen formuliert.

Kernmechanismus

Anstatt jeden Token des nächsten Bildes neu zu generieren, bestimmt ITC für jeden Token des nächsten Bildes, ob er:

1. Von einem Token im vorherigen Bild kopiert werden soll (zur Bewahrung der Identität).
2. Als neuer Token generiert werden soll (zur Handhabung von Änderungen oder neuem Inhalt).

Diese Zuweisung wird als Optimal Transport (OT)-Problem modelliert.

Technische Architektur

1. Tokenisierung & Einbettung:

   • Visuelle Beobachtungen werden unter Verwendung eines Tokenizers für die Nachbarschaftssuche auf Patches in diskrete Token umgewandelt.
   • Token-Einbettungen werden mit 3D-Rotary-Positionseingebungen (RoPE) angereichert. Im Gegensatz zur standardmäßigen 1D-RoPE kodiert diese Formulierung zeitliche, vertikale und horizontale Koordinaten unabhängig voneinander, was es dem Modell ermöglicht, lokalisierte relationale Strukturen sowohl im Raum als auch in der Zeit zu erfassen.

2. Optimal Transport-Löser:

   • Das Modell konstruiert einen Graphen, wobei Quellknoten Token des vorherigen Zustands ($u$) und Kandidaten-Token für den nächsten Zustand, die vom Transformer vorhergesagt werden ($\tilde{p}$), repräsentieren. Zielknoten repräsentieren die finalen Token des nächsten Zustands ($\hat{u}$).
   • Eine Affinitätsmatrix wird berechnet, um die Beziehung zwischen vorherigen Token und vorhergesagten Token zu bewerten. Diese Matrix umfasst:
     • $A^{(prev)}$: Bewertet die Affinität zwischen vorherigen Token und Vorhersagen, bestraft durch den räumlichen Abstand, um physikalische Plausibilität zu erzwingen (z. B. können Objekte nicht willkürlich weit teleportieren).
     • $A^{(gen)}$: Bewertet den „Bonus" der Zulassung eines neu generierten Tokens (Wildcard) anstelle der Wiederverwendung eines vorherigen Tokens, unterliegt jedoch einer konstanten Strafe.
   • Der Löser (unter Verwendung des Sinkhorn-Algorithmus) berechnet einen Transportplan, der die Kosten des Transports von Masse von Quellen zu Zielen minimiert.

3. Binärisierung & Decodierung:

   • Der kontinuierliche Transportplan wird mittels eines gierigen Binärisierungsverfahrens (Algorithmus 2) in eine binäre Zuweisungsmatrix umgewandelt. Dies stellt eine strikte Eins-zu-Eins-Zuordnung sicher.
   • Decodierungsregel: Für jede Position im nächsten Bild wird, wenn der Transportplan einen vorherigen Token zuweist, dieser Token kopiert. Wenn der Plan einen „generieren"-Token zuweist, sampelt das Modell aus der vom Transformer vorhergesagten Verteilung.

Trainingsloop

Das Weltmodell und die Policy werden gleichzeitig trainiert. Das Weltmodell generiert „Vorstellungen" (synthetische Trajektorien), die zusammen mit realen Umgebungsdaten verwendet werden, um die Policy über Proximal Policy Optimization (PPO) zu aktualisieren.

3. Hauptbeiträge

• Formulierung der Token-Korrespondenz: Die Arbeit stellt die Vorhersage des nächsten Bildes neu dar, indem sie die Token-Korrespondenz explizit als Optimal-Transport-Problem modelliert und die Aufgaben des „Wiederverwendens" persistenter Entitäten und des „Generierens" neuen Inhalts trennt.
• Integration von Optimal Transport: Es wird ein differenzierbarer (via Sinkhorn) jedoch diskreter (via Binärisierung) Mechanismus eingeführt, um zwischen dem Kopieren vorheriger Token und dem Generieren neuer Token zu wählen, was Halluzinationen effektiv reduziert.
• 3D-Positionscodierung: Die Implementierung von 3D-RoPE ermöglicht es dem Modell, zwischen zeitlichen und räumlichen Unterschieden zu unterscheiden, eine notwendige Komponente für das Verfolgen sich bewegender Objekte in Videosequenzen.

4. Experimentelle Ergebnisse

Die Autoren bewerteten ITC auf vier herausfordernden Benchmarks: Craftax-classic, Craftax, MinAtar und Atari 100K.

• Craftax-classic: ITC erzielte eine Rendite von 72,5 % und eine Punktzahl von 35,6 % und übertraf damit deutlich den vorherigen State-of-the-Art (Dedieu et al., 2025), der eine Rendite von 67,4 % und eine Punktzahl von 27,9 % erreichte. Es zeigte zudem eine schnellere Konvergenz.
• Craftax: In dieser komplexeren Umgebung erzielte ITC eine Rendite von 7,09 % und eine Punktzahl von 2,40 % und übertraf Baselines einschließlich Simulus und Dedieu et al.
• MinAtar: ITC übertraf den vorherigen State-of-the-Art für modellbasiertes RL in allen vier Spielen (Asterix, Breakout, Freeway, SpaceInvaders).
• Atari 100K: Auf einer Suite von 26 Spielen angewandt, erzielte ITC neue State-of-the-Art-Leistungen im Interquartile Mean (IQM) von 1,092 und einer Optimalitätslücke von 0,376 und übertraf dabei das token-basierte Weltmodell Simulus.
• Vorhersagegenauigkeit: Auf einem Datensatz von 10.000 Übergängen verbesserte ITC die Gesamtvorhersagegenauigkeit um 3,39 % im Vergleich zur Baseline, mit einer bemerkenswerten Verbesserung von 3,45 % bei Übergängen, die sich bewegende Kreaturen betrafen.
• Qualitative Verbesserungen: Die visuelle Analyse bestätigte, dass ITC häufige Artefakte wie die Verdopplung und das Verschwinden von Objekten, wie sie bei Baseline-Modellen wie DreamerV3 und Dedieu et al. zu sehen sind, eliminiert.
• Effizienz: Die Methode führte zu einem vernachlässigbaren Overhead und erhöhte die gesamte End-to-End-Trainingszeit nur um 2,8 %.

5. Bedeutung und Behauptungen

Die Arbeit behauptet, dass durch die explizite Modellierung der Korrespondenz zwischen Token über die Zeit hinweg ITC ein weit verbreitetes Problem in Weltmodellen löst, bei dem Agenten über lange Horizonte hinweg die Objektidentität nicht aufrechterhalten können. Die Autoren stellen fest, dass dieser Ansatz es dem Modell ermöglicht, „Token selektiv wiederzuverwenden, die eine starke Korrespondenz mit dem vorherigen Bild aufweisen", wodurch Ähnlichkeiten zwischen den Bildern genutzt werden, anstatt sich ausschließlich darauf zu verlassen, dass der Transformer jeden Token neu generiert.

Die Bedeutung liegt in der Erzielung von State-of-the-Art-Leistungen in diversen Umgebungen (von gitterbasierten offenen Welten bis hin zu Standard-Atari-Spielen) bei gleichzeitiger Wahrung der Recheneffizienz. Die Autoren schließen, dass ITC effektiv zeitliche Inkonsistenzen wie Verdopplung und Verschwinden verhindert, was zu Vorstellungen höherer Qualität und einer verbesserten Policy-Leistung führt.