Identifiable Token Correspondence for World Models

2026-05-15

AI Summary

This work addresses the temporal inconsistency issues—such as object duplication, disappearance, or abrupt changes—that arise in Transformer-based world models during long-horizon rollouts due to the absence of explicit cross-timestep token correspondences. To mitigate this, the authors propose a structured probabilistic inference framework that formulates next-frame prediction as a binary decision process for each token: either copying it from the previous frame or generating it anew. This approach introduces, for the first time in visual world models, an identifiable cross-temporal token correspondence mechanism, enabling explicit modeling of token-level state evolution. The method achieves state-of-the-art performance across four challenging benchmarks, notably attaining a 72.5% return (+5.1%) and a 35.6% score (+7.7%) on Craftax-Classic, significantly outperforming existing approaches.

Abstract

Transformer-based 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 explicitly modeling correspondence between tokens across time. We formulate next-frame prediction as a structured probabilistic inference problem with latent token correspondence variables, deriving a model in which each next-frame token is explained either by copying a token from the previous frame or by generating a new token. 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.

Problem

• temporal inconsistency
• token correspondence
• world models
• next-frame prediction
• visual reinforcement learning

Innovation

• token correspondence
• world models
• temporal consistency
• structured probabilistic inference
• visual reinforcement learning