Technical Summary: Identifiable Token Correspondence for World Models

1. Problem Statement

Transformer-based world models have demonstrated strong performance in visual reinforcement learning (RL) by treating sequences of past states and actions as token streams to predict future states. However, these models frequently suffer from temporal inconsistency during long-horizon rollouts. Common failure modes include object duplication, disappearance, and transmutation into different objects.

The authors identify a primary cause for these errors: existing approaches treat next-frame prediction purely as a token generation problem. They fail to explicitly model the correspondence between tokens across time. In realistic visual environments, tokens in successive frames often correspond to the same underlying entities that persist and move. When a model is forced to regenerate every token from scratch without tracking identity, it struggles to preserve persistent structures, leading to fragile identity preservation and compounding errors over time.

2. Methodology: Identifiable Token Correspondence (ITC)

The paper proposes Identifiable Token Correspondence (ITC), a transformer world model that formulates next-frame prediction as a structured probabilistic inference problem with latent token correspondence variables.

Core Mechanism

Instead of generating every next-frame token de novo, ITC determines for each next-frame token whether it should be:

1. Copied from a token in the previous frame (preserving identity).
2. Generated as a new token (handling changes or new content).

This assignment is modeled as an Optimal Transport (OT) problem.

Technical Architecture

1. Tokenization & Embedding:

   • Visual observations are converted into discrete tokens using a nearest-neighbor patch lookup tokenizer.
   • Token embeddings are augmented with 3D Rotary Position Embeddings (RoPE). Unlike standard 1D RoPE, this formulation encodes temporal, vertical, and horizontal coordinates independently, allowing the model to capture localized relational structures across both space and time.

2. Optimal Transport Solver:

   • The model constructs a graph where source vertices represent previous state tokens ($u$) and candidate next-state tokens predicted by the transformer ($\tilde{p}$). Destination vertices represent the final next-state tokens ($\hat{u}$).
   • An affinity matrix is computed to score the relationship between previous tokens and predicted tokens. This matrix includes:
     • $A^{(prev)}$: Scores the affinity between previous tokens and predictions, penalized by spatial distance to enforce physical plausibility (e.g., objects cannot teleport arbitrarily far).
     • $A^{(gen)}$: Scores the "bonus" of admitting a newly generated token (wildcard) instead of reusing a previous one, subject to a constant penalty.
   • The solver (using the Sinkhorn algorithm) computes a transport plan that minimizes the cost of moving mass from sources to destinations.

3. Binarization & Decoding:

   • The continuous transport plan is converted into a binary assignment matrix using a greedy binarization procedure (Algorithm 2). This ensures a strict one-to-one mapping.
   • Decoding Rule: For each position in the next frame, if the transport plan assigns a previous token, that token is copied. If the plan assigns a "generate" token, the model samples from the transformer's predicted distribution.

Training Loop

The world model and policy are trained concurrently. The world model generates "imaginations" (synthetic trajectories) which are used alongside real environment data to update the policy via Proximal Policy Optimization (PPO).

3. Key Contributions

• Formulation of Token Correspondence: The paper reframes next-frame prediction by explicitly modeling token correspondence as an optimal transport problem, separating the tasks of "reusing" persistent entities and "generating" new content.
• Optimal Transport Integration: It introduces a differentiable (via Sinkhorn) yet discrete (via binarization) mechanism to select between copying previous tokens and generating new ones, effectively reducing hallucinations.
• 3D Positional Encoding: The implementation of 3D RoPE allows the model to distinguish between temporal and spatial differences, a necessary component for tracking moving objects in video sequences.

4. Experimental Results

The authors evaluated ITC on four challenging benchmarks: Craftax-classic, Craftax, MinAtar, and Atari 100K.

• Craftax-classic: ITC achieved a return of 72.5% and a score of 35.6%, significantly surpassing the previous state-of-the-art (Dedieu et al., 2025) which achieved 67.4% return and 27.9% score. It also demonstrated faster convergence.
• Craftax: On this more complex environment, ITC achieved a return of 7.09% and a score of 2.40%, outperforming baselines including Simulus and Dedieu et al.
• MinAtar: ITC surpassed the previous state-of-the-art for model-based RL in all four games (Asterix, Breakout, Freeway, SpaceInvaders).
• Atari 100K: Applied to a suite of 26 games, ITC achieved new state-of-the-art performance in Interquartile Mean (IQM) of 1.092 and Optimality Gap of 0.376, outperforming the token-based world model Simulus.
• Prediction Accuracy: On a dataset of 10,000 transitions, ITC improved overall prediction accuracy by 3.39% compared to the baseline, with a notable 3.45% improvement in transitions involving moving creatures.
• Qualitative Improvements: Visual analysis confirmed that ITC eliminates common artifacts such as object duplication and disappearance seen in baseline models like DreamerV3 and Dedieu et al.
• Efficiency: The method introduced negligible overhead, increasing total end-to-end training time by only 2.8%.

5. Significance and Claims

The paper claims that by explicitly modeling the correspondence between tokens across time, ITC resolves a widespread issue in world models where agents fail to maintain object identity over long horizons. The authors state that this approach allows the model to "selectively reuse tokens that have a strong correspondence with the previous frame," thereby leveraging frame-to-frame similarities rather than relying solely on the transformer to regenerate every token.

The significance lies in achieving state-of-the-art performance across diverse environments (from grid-based open worlds to standard Atari games) while maintaining computational efficiency. The authors conclude that ITC effectively prevents temporal inconsistencies like duplication and disappearance, leading to higher quality imaginations and improved policy performance.