# Identifiable Token Correspondence for World Models

arXiv: https://arxiv.org/abs/2605.16457 · HTML: https://mlsanity.com/paper/2605.16457/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Latent dynamics models and model-based RL, one of 11 directions in World models.
- Builds on 13 of its 40 references that are in the map.
- First posted on arXiv 15 May 2026.

## Placed in

- World models / [Latent dynamics models and model-based RL](https://mlsanity.com/wm/latent-dynamics-models-and-model-based-rl/)
- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 13 of 13)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Qwen2-VL: Enhancing Vision-Language Model's Perception of the World at Any Resolution](https://mlsanity.com/paper/2409.12191/) (2024)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Cosmos World Foundation Model Platform for Physical AI](https://mlsanity.com/paper/2501.03575/) (2025)
- [Transformers are Sample-Efficient World Models](https://mlsanity.com/paper/2209.00588/) (2022)
- [Benchmarking the Spectrum of Agent Capabilities](https://mlsanity.com/paper/2109.06780/) (2021)
- [Searching for Activation Functions](https://mlsanity.com/paper/1710.05941/) (2017)
- [Transformer-based World Models Are Happy With 100k Interactions](https://mlsanity.com/paper/2303.07109/) (2023)
- [Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein](https://mlsanity.com/paper/2201.12324/) (2022)
- [VideoRoPE: What Makes for Good Video Rotary Position Embedding?](https://mlsanity.com/paper/2502.05173/) (2025)
- [MinAtar: An Atari-Inspired Testbed for Thorough and Reproducible Reinforcement Learning Experiments](https://mlsanity.com/paper/1903.03176/) (2019)
- [Simulus: Combining Improvements in Sample-Efficient World Model Agents](https://mlsanity.com/paper/2502.11537/) (2025)

## Abstract (arXiv)

Token-based transformer world models have shown strong performance in visual reinforcement learning, but often suffer from temporal inconsistency in long-horizon rollouts, including object duplication, disappearance, and transmutation. A key reason is that most existing approaches treat next-frame prediction purely as a token generation problem, without considering the persistence of tokens across time. We introduce Identifiable Token Correspondence (ITC), a decoding step for token-based transformer world models that formulates next-frame prediction as a structured assignment problem with latent token correspondence variables: each next-frame token is explained either by copying a token from the previous frame or by generating a new one. ITC leaves the transformer architecture and training procedure unchanged and can be added on top of existing backbones. Our experiments show state-of-the-art performance on 4 challenging benchmarks. The proposed method achieves a return of 72.5% and a score of 35.6% on the Craftax-classic benchmark, significantly surpassing the previous best of 67.4% and 27.9%. We release our source code on https://github.com/snu-mllab/Identifiable-Token-Correspondence.