World Modeling in Transformers

Pierre Beckmann; Matthieu Queloz; Andre Freitas

cs.AISep 18, 2026

World Modeling in Transformers

Authors: Pierre Beckmann, Matthieu Queloz, Andre Freitas

Organizations: EPFL · IDIAP Research Institute · MATS · University of Bern ·

Abstract

Behavioral failures can make a transformer appear to lack a world model even when it has learned faithful representations of its environment. We demonstrate this in TaxiGPT, a transformer trained on random walks through Manhattan whose failures have been interpreted as evidence of an incoherent internal map. Through mechanistic analysis and causal interventions, we show that the model represents intersections and streets, tracks its position, and uses a goal compass to navigate. We trace its failures to interference between superposed intersection features, which disrupts localization within the internal map. Affordance packing, which groups representations of intersections with the same legal moves, helps limit the consequences of these errors. Finally, we propose mechanistic indicators that we use to compare models and show that world-modeling capacities emerge at different stages of training. Our findings motivate a shift from asking whether a model has a world model to mechanistically studying its world modeling: the interacting capacities through which it represents its environment and uses those representations to guide behavior.

Figures & tables

Fig. 1: The map encoded inside TaxiGPT. (a) The Manhattan map recovered by the street steering test: we inject an intersection feature, feed a legal move, and measure which intersection feature grows most afterward (color shows the correct next intersection’s rank); (b) the intersection features are superposed in a small subspace of the residual stream; (c) and are packed by affordance: intersections with the same legal next moves cluster together (main affordance classes shown).

Fig. 2: Using the map. (a) Localization : the model carries a running position, read from a look-back window of past positions. Erasing past positions hurts current-intersection decode while erasing the past moves does not (measured on rides of 27–34 moves). (b) Navigation : a goal compass encodes the bearing to the goal. It is also causal: clamping it north or south makes the taxi follow that direction. Two mechanisms write into logits: a legal-move mechanism (activated by the intersection feature) favors legal moves while the goal-compass mechanism favors goalward moves.

Fig. 3: A weak write and noise in the position code are enough to cause off-graph moves. (a) When the true-position write is weak, noise can make a wrong intersection feature most active. (b) We causally reproduce this phenomenon by weakening the write and adding noise to correctly localized states, causing the model to emit illegal moves (

700

stress rides with initially legal top moves, depth

\geq 60

; edits at layer 18). (c) The farther the wrong feature lies from the true direction, the fewer legal moves the two intersections share and the more likely an illegal move becomes. The nearby region protected by affordance packing is shaded green in (a) and (c).

Fig. 4: Two stress rides where a wrong node becomes most active, leading to an illegal move. The shared fluctuations reflect how superposition works: when the true node is written more strongly, wrong nodes with superposed features also become more active.

Fig. 5: TaxiGPT’s failure modes. a The first row shows how true-node (dashed) and wrong-node (solid) activations change over successive moves, comparing clean rides, recovering slips and four failure modes. b In these failure modes, weaker writes (write

\downarrow

) and stronger noise (noise

\uparrow

) let the wrong node lie at a larger angle (cos

\downarrow

) and share fewer legal moves (aff

\downarrow

), making illegal moves more likely. c As a causal check, we remove the wrong-node activation (

-

wrong), remove position noise (

-

noise), or strengthen the true-position write (

+

write), and show the fraction of illegal-move probability removed (gray marks: controls with edits of the same size along random directions).

Fig. 6: Factors affecting write strength and noise. (a) In distribution, the write (measured at layer 18) falls with distance to the goal but rises with depth. Under stress the pattern is broadly the same, except in the out-of-distribution deep-and-far regime (red ring), where the write drops sharply. The write is also weaker when the current intersection’s representation overlaps with many others (crowding), and when the route is less predictable (route surprise). (b) Noise grows with one factor: depth. (Lines are medians; bands are the interquartile spread.)

	Map / localization					Navigation		Behavior
	decode	causal	streets	super-	by	legal	goal	stress	detour	compr.
			probe	position	afford.?	moves	compass	test	test
	% int $\geq$ .9	telep $>$ 50%	held out	angle / dims	own class	% int.	steer err	on graph	success	score
SP·NTP [1pt] 12L $\times$ 768d $\times$ 12h	12.7% [1pt] L6	14.9% [1pt] L10	23.5% [1pt] L7	35.9 ∘ [1pt] 162d, L6	41.0% [1pt] L6	18.1% [1pt] L10	25.8 ∘ [1pt] L12	71.7%	0.0%	.101
NSP·NTP [1pt] 48L $\times$ 1600d $\times$ 25h	17.1% [1pt] L43	16.9% [1pt] L46	12.2% [1pt] L36	38.0 ∘ [1pt] 251d, L43	41.1% [1pt] L43	22.3% [1pt] L44	23.2 ∘ [1pt] L47	74.7%	0.2%	.054
RW·NTP [1pt] 48L $\times$ 1600d $\times$ 25h	99.6% [1pt] L18	94.0% [1pt] L11	89.2% [1pt] L15	48.1 ∘ [1pt] 244d, L18	91.4% [1pt] L18	99.4% [1pt] L35	18.1 ∘ [1pt] L15	91.5%	63.1%	.524

Table 1: Comparing world modeling capacities across architectures, data and training regimes. For each indicator, we select the layer where it scores best in its corresponding sweep. Rows are dataset · objective: SP shortest paths, NSP noisy shortest paths, RW random walks; NTP next-token prediction, NextLat next-latent prediction ( Teoh et al., 2026 ) . (Protocol, controls, and training budgets in App. H.1 .)

Fig. 7: Emergence of world-modeling capacities through training (RW·NTP,

384

d). Navigation capacities (compass and legality) emerge first, localization (decode and causal) second (a). The position code expands (nearest-neighbor angles grow); affordance packing rises early, then relaxes as intersections differentiate (b). Mechanistic indicators can be compared with behavioral proxies (c). Values are normalized to the final checkpoint. (Protocol and controls in App. H.2 .)

Appendix figures & tables38 assets

Supplementary material from the paper’s appendix.

Appendix

Fig. 8: Reading and steering position across layers. (a) Held-out top-1 decode accuracy. (b) One-move success after minimal teleportation on

200

scenes, with

\alpha=\beta=1.5

; end competes with the eight moves. Dotted guides mark the selected reading layer (18) and steering layer (11). The dashed curve is the random-node edit control for diff-means steering.

criterion	intersections	share
teleport succeeds on any scene (L11)	4,172	99.3%
teleport succeeds on $>50\%$ of scenes	3,951	94.0%
teleport succeeds on every scene	2,764	65.8%
decode accuracy $\geq 0.9$ (L18)	4,478	99.6%
decode accuracy $=1$	4,026	89.5%

Appendix

Table 3: Most intersection features are causal. Teleportation succeeds for

99.3\%

of the

4{,}202

testable intersections on at least one scene, and for

94.0\%

on more than half. Decoding accuracy is reported below for the

4{,}497

intersections observed in held-out rides. Rows within each group overlap.

layer	6	8	10	12	14	16	18	20	random-init
latitude	.988	.984	.977	.971	.968	.964	.961	.960	.534
longitude	.988	.985	.979	.975	.972	.969	.966	.965	.474

Appendix

Table 4: Coordinates are linearly readable from intersection features across layers. Latitude/longitude

R^{2}

on the

20\%

of intersections held out from fitting. The random-init control uses layer 6.

Fig. 10: A linear readout of intersection features recovers Manhattan’s spatial layout. Actual locations (left) and coordinates predicted from layer 6 features in the trained (middle) and random (right) models, on the same scale. Colors mark three geographic bands along Manhattan’s long axis, kept fixed across panels. Held-out latitude/longitude $R^{2}$ is $0.988/0.988$ for the trained model and $0.534/0.474$ for the random model. Maps include intersections used to fit the readout.

layer	18	20	24	28	32	36	44
linear probe	.989	.997	.998	.999	.999	.999	.999
logit-lens	.937	.962	.993	.999	.999	.999	.999

Appendix

Table 6: Legal-move F1 by layer. The linear probe is fitted on training intersections and scored on held-out intersections; logit-lens uses the model’s output projection. Each entry uses the threshold maximizing F1 on its reported scores. These scores pool intersection–move pairs, rather than requiring the complete legal-move set to be correct.

Fig. 12: Move logits from centered intersection directions at layer 32, split by legality. We apply the logit lens to

c_{n}-\bar{c}

. Legal moves generally receive positive logits and illegal moves negative logits.

Fig. 13: Two tests of street connectivity. (a) The steering test asks whether an intersection feature and a move activate the feature of the correct next intersection. We inject an intersection feature into the averaged state 0 and feed move

m

, then measure which intersection feature grows most at state 1. The other prompt states are averaged (beige) to avoid tying the test to a particular ride; state 1 evolves normally (blue). (b) The probing test asks whether street connections can be read from the model’s state during an ordinary ride. Eight linear probes, one per move, predict the feature of the intersection that move would reach. We match each output to the closest intersection feature and test on held-out source intersections.

Fig. 14: Street steering test. Each tested move is drawn from its source intersection to the street midpoint, so opposite directions can be colored separately. Color shows the correct next intersection’s rank among all candidates; the right panel enlarges the boxed area.

Fig. 15: Without past positions, decoding worsens as the ride gets longer. We ablate past position features while preserving the origin, destination and final token, then decode the current intersection at layer 18.

Fig. 16: The goal compass encodes goal direction and guides travel. (a) Error in the goal bearing decoded from held-out ride states, by layer. (b) The model follows the imposed direction across the

16

tested bearings. Each wedge points in a direction imposed on the compass. Its radius shows how closely the model travels in that direction: mean cosine alignment, with

1

for perfect alignment. Alignment is highest near Manhattan’s long axis (dashed), estimated from the map coordinates.

	$\beta{=}0$	$\beta{=}0.5$	$\beta{=}1$	$\beta{=}1.5$	$\beta{=}2$
$\alpha{=}0.25$	.701	.796	.834	.847	.845
$\alpha{=}0.5$	.904	.913	.892	.899	.894
$\alpha{=}1$	.905	.919	.920	.922	.917
$\alpha{=}1.5$	.843	.846	.857	.857	.826
$\alpha{=}2$	.505	.513	.512	.484	.434

Appendix

Table 7: Choosing the compass edit strength. At layer 15,

\beta

scales removal of the existing compass component and

\alpha

scales the added target direction. Entries give mean

\cos(\text{heading}-\text{target})

1

means aligned,

0

perpendicular, and

-1

opposite. We use

60

start–goal pairs and target bearings

45^{\circ}

and

225^{\circ}

; the

16

-bearing sweep uses

100

pairs. Bold marks the chosen setting.

start distance (moves)	3–6	7–11	12–17	18–32
reaches goal:
with compass	1.00	1.00	0.99	0.94
compass ablated	0.68	0.54	0.38	0.23
random-plane ablated	1.00	1.00	0.98	0.94
moves that are legal:
compass ablated	0.995	0.997	0.996	0.995

Appendix

Table 8: Compass removal reduces goal-reaching while moves remain legal. Starting distance is the minimum number of legal moves to the goal. The four distance bins contain

31

50

115

, and

466

tasks. Removing a random plane of the same dimension serves as a control.

Fig. 17: Without the compass, the model follows streets but misses the goal. Four selected rides after compass removal; circles mark origins and stars mark goals. The corresponding baseline rides reach their goals.

Fig. 18: The at-goal feature detects arrival and can trigger stopping. (a) Held-out accuracy when the goal is set to the current intersection or a distant node. (b) Stop-token argmax rate and mean probability after injection at states away from the goal; the random control shows its argmax rate.

outcome	feature fires	$P(\text{end})$	$n$
stops at the goal	0.93	0.80	2411
wrong stop or generation limit	0.13	0.89	383

Appendix

Table 9: At-goal readout at final states of self-generated rides. Generation-limit outcomes reach the

128

-move cap without a stop token or illegal move.

Fig. 19: One direction controls stopping and how much the model meanders. (a) Stop probability after adding the continue direction at terminal training-ride states. (b) Path length relative to the shortest path (clay), conditional on ending at the goal, and the fraction ending there (slate). Negative strengths favor commitment; positive strengths favor exploration. Bands show one standard deviation across three generation seeds; dashed lines show random-direction controls.

Fig. 20: Steering changes the route on a fixed origin–goal pair. One selected origin–goal pair at five steering strengths, with a

200

-move cap. Circles mark the origin and stars the goal. Negative strength shortens this route; positive strength adds loops.

Fig. 21: Commit-to-goal projection by depth on stress rides. The line shows the median and the band the interquartile range; higher values point toward committing to the goal.

Fig. 22: The stress test shifts both pair distances and visited states. (a) Origin–goal distances in the released stress reference and held-out training-distribution rides, with reference failure rate on the right axis. (b) Visited-state occupancy for held-out rides and our generated stress set. Heatmaps show depths up to $100$ ; their top distance bin includes all distances of $57$ moves or more.

kind	condition	share
low-mass unlucky draw	$P_{\rm illegal}<0.001$	$6.0\%$
give-up	else $v\cdot u_{\rm stop}\geq 38.8$	$16.2\%$
silent slip	else $L=T$ and a supplier exists	$23.9\%$
fatal slip	else $L\neq T$ , $u_{L}\cdot u_{T}>0.2$	$24.7\%$
full corruption	else $L\neq T$ , $u_{L}\cdot u_{T}\leq 0.2$	$28.7\%$
unclassified	else $L=T$ without a supplier	$0.5\%$

Appendix

Table 10: Failure categories, applied in the order shown. Here

v=h-\bar{\mu}

u_{n}=(\mu_{n}-\bar{\mu})/\|\mu_{n}-\bar{\mu}\|

T

is the true intersection, and

L=\arg\max_{n}v\cdot u_{n}

. The same direction scores define the strongest wrong node in the figures. A supplier is a non-true node among the

13

highest direction scores where the highest-logit illegal move is legal. Shares use all

1{,}681

cases.

Fig. 23: Stress trips illustrating four failure modes. Examples are selected near category medians at the final state. Maps mark the start (black dot), goal (star), and illegal-move attempt (red cross). Matching colored squares mark every 20 moves on the route and time axis. Traces show the layer-18 true-position write (green) and strongest wrong-node activation (red). True- and wrong-node activations often rise and fall together. This is expected under superposition: intersection features have overlapping directions, so strengthening the true-position feature also raises the activation of wrong features that overlap with it.

Fig. 26: Give-up slips combine stopping activity with enlarged residuals. (a) Adding the give-up direction makes successful-ride states predict end . (b) Residual norms across failure categories, relative to clean states. (c) On give-up prefixes, reducing excess residual activity lowers illegal-move probability more than suppressing the stopping signal alone. Random edits match the reduction’s magnitude at each layer and position; error bars are

95\%

bootstrap intervals over prefixes.

Fig. 27: Detours first move away from the goal, then return. (a) Median goal distance over

10{,}000

rides, with the interquartile range; the clay line shows the percentage of ongoing rides going off-graph in each five-move bin (right axis). (b) Two selected successful trajectories illustrating the reachability-constrained return; dots mark origins and stars goals.

Fig. 28: Individual forced moves do not immediately disrupt the position code, but repeated forcing creates deep-and-far, unlikely rides on which the position write weakens. (a) Move-to-move changes in write and noise after imposed versus self-chosen moves. (b) Median write decreases with goal distance within depth bands. (c) Detour routes are less likely than stress routes. (d) Lower route probability accompanies weaker writes; shading shows the interquartile range. Route log-probability is averaged over moves taken so far. States above the give-up threshold are excluded throughout.

Fig. 29: Detours exhibit the same failure modes: weak position writes and noise, with recovery under the same position edits. Columns compare failure categories and reference states. Rows show position readouts around failure or recovery, scaled write/noise/cosine/affordance measurements, and the fraction of illegal probability removed. Gray marks show matched random edits.

Fig. 30: Detour trips illustrating four failure modes. Examples are near category medians at failure. Maps mark the start (dot), goal (star), and illegal-move attempt (cross); squares mark every

20

moves on the route and time axis. Traces show layer-18 true-position write (green) and strongest wrong-node activation (red).

Fig. 31: Two routes to one intersection, scored through their continuations. Prefixes A and B have equal length and the same goal. Continuations sampled after A are scored after B; a trial passes only if every token in all

30

continuations clears

\epsilon=0.01

. The drawing shows a selected subset.

at the end of a route	pair passes		pair fails
	weaker	stronger	weaker	stronger
distance to the goal (moves)	$27.5$ (21–40)		$54$ (41–67)
length of the route (moves)	$20.5$ (10–34)		$18.5$ (9–33)
write on the true intersection	$426$	$457$	$405$	$438$
noise in the position code	$51.3$	$51.9$	$51.0$	$50.9$
log-probability of the route	$-0.870$	$-0.832$	$-0.833$	$-0.853$

Appendix

Table 11: Prefix states in passing and failing compression pairs. The

146

sampled-prefix trials contain

52

passing and

94

failing pairs. Each pair is ordered by write strength. Entries are medians; parentheses give the first and third quartiles for goal distance and prefix length.

Fig. 32: Compression generates rides far from the training distribution, where position writes weaken. (a) Origin–goal distances and off-graph ride rates. (b) Visited-state occupancy by moves taken and remaining goal distance, compared with held-out training-distribution rides. (c) Median true-position write by goal distance, grouped by moves taken in the full prefix + suffix ride; states above the give-up threshold are excluded.

Fig. 33: Compression rides exhibit the same failure modes: weak position writes and noise, with recovery under the same position edits. Columns compare failure categories and reference states. Rows show position readouts around failure or recovery, scaled write/noise/cosine/affordance measurements, and the fraction of illegal probability removed. Gray marks show matched random edits.

		median length (moves)
goal distance (moves)	pairs	route	continuation	illegal /1k	diff-means decode	score
sampled-prefix diagnostic run
$1$ – $20$	18	25.5	–	0.16	100%	0.722
$21$ – $40$	45	29	–	0.89	100%	0.622
$41$ – $60$	49	23	–	4.69	100%	0.184
$61+$	34	9.5	–	8.76	100%	0.059

Appendix

Table 12: Compression scores fall as the goal becomes more distant. Upper block:

146

supplementary trials with sampled prefix lengths. Lower block: the same

60

pairs of

12

-move routes in each band, changing only the goal (seed 0). Illegal /1k is the off-graph rate per thousand checked moves. Decode requires both prefixes to identify the shared intersection. Dashes indicate unrecorded continuation lengths.

Map / localization

Navigation

Behavior

decode

causal

streets

tracking

super-

legal

goal

stress

detour

compr.

steer

probe

position

afford.?

moves

compass

test

% int

\geq

telep

>

50%

top5

held out

-

past pos

angle / dims

own class

% int.

steer err

on graph

success

score

SP·NTP [1pt] 12L

\times

768d

\times

12h

12.7% [1pt] L6

14.9% [1pt] L10

8.2% [1pt] L8

23.5% [1pt] L7

.19

\rightarrow

.09 [1pt] L6

35.9 ∘ [1pt] 162d, L6

41.0% [1pt] L6

18.1% [1pt] L10

25.8 ∘ [1pt] L12

71.7%

0.0%

.101

NSP·NTP [1pt] 48L

\times

1600d

\times

25h

17.1% [1pt] L43

16.9% [1pt] L46

8.0% [1pt] L45

12.2% [1pt] L36

.21

\rightarrow

.02 [1pt] L43

38.0 ∘ [1pt] 251d, L43

41.1% [1pt] L43

22.3% [1pt] L44

23.2 ∘ [1pt] L47

74.7%

0.2%

.054

RW·NTP [1pt] 48L

\times

1600d

\times

25h

99.6% [1pt] L18

94.0% [1pt] L11

93.5% [1pt] L17

89.2% [1pt] L15

1.00

\rightarrow

.15 [1pt] L18

48.1 ∘ [1pt] 244d, L18

91.4% [1pt] L18

99.4% [1pt] L35

18.1 ∘ [1pt] L15

91.5%

63.1%

.524

Appendix

Table 13: World-modeling measurements across models. SP , NSP and RW denote shortest paths, noisy shortest paths and random walks; NTP and NextLat denote next-token and next-latent prediction. Street steering reports top-five next-intersection accuracy; tracking reports current-position accuracy before and after removing past position codes. Selected layers appear below each score.

Model	Batch size	Corpus tokens (B)	Epochs	Training tokens (B)
SP	$48$	$\sim 0.12$	—	—
NSP	$48$	$\sim 1.68$	$5$	$\sim 8.39$
Large RW	$48$	$\sim 4.74$	$1$	$\sim 4.74$
Small RW·NTP	$256$	$\sim 4.74$	$6$	$\sim 28.41$
Small RW·NextLat	$256$	$\sim 4.74$	$6$	$\sim 28.41$

Appendix

Table 14: Reported training budgets of the compared models. Corpus sizes and epochs follow Vafa et al. (2024) and Teoh et al. (2026) . Total training tokens are corpus tokens multiplied by epochs, counting repeated passes. SP uses the best validation checkpoint, so its epoch count and total exposure are unspecified. Token counts are in billions. Global batch sizes count rides for Vafa models and packed

256

-token sequences for small RW models.

model		slips	recovers
		% of steps	within 5 steps
RW·NTP	48L $\times$ 1600d	1.52%	67.8%
RW·NTP	48L $\times$ 384d	0.56%	91.2%
RW·NextLat	48L $\times$ 384d	0.24%	77.4%

Appendix

Table 15: Smaller random-walk models slip less often under stress.

	Decode (L40)		Causal (L33)				Tracking (L40)			Streets: steering test (L24)		Streets: probe test (L39)
step	$\geq.9$	$\geq.5$	$>.5$	mean	floor	dose	full	$-$ pos	$-$ rand	top1	top5	next	current
50	0.3	0.3	8.6	0.31	13.6	0.28	–	–	–	–	–	–	–
200	0.4	1.4	12.4	0.37	9.2	0.64	0.03	0.02	0.04	0.6	2.0	1.6	2.2
700	3.8	20.4	26.5	0.48	10.8	1.03	0.27	0.03	0.27	7.8	18.9	12.3	22.8
1500	8.8	43.5	32.6	0.51	10.0	1.07	0.51	0.02	0.47	10.4	25.3	20.5	43.6
3000	21.8	80.4	50.3	0.59	6.8	1.08	0.71	0.01	0.64	10.8	27.6	34.4	65.4

Appendix

Table 16: World-modeling capacities develop at different stages of training. Raw measurements at twelve checkpoints of the small RW·NTP model, using fixed layers. The upper block covers localization and street encoding; the lower block covers packing, navigation, and behavior. Legal-move prediction and the goal compass improve before reliable position decoding. For causal measurements, floor is teleportation success without an edit and dose is the edit-to-residual norm ratio. Tracking compares the full history with removal of past position information (

-

pos) or a random-subspace edit (

-

rand); street probes predict the next or current intersection. Packing reports the nearest-neighbor feature angle, the dimensions explaining

90\%

of variance (

d_{90}

), and the percentage of features closest to their own legal-move group’s mean (affordance). Accuracy and success rates are percentages, except causal mean, tracking, and legal-set accuracy, which are fractions. Compression is also a fraction; angle and compass error are in degrees. Dashes mark unavailable measurements.

World Model Learning

World model learning builds predictive models of how environments change under actions to support planning, simulation, and decision-making.

See all 419 papers

Identifiable Token Correspondence for World Models
May 15, 2026Youngin Kim, Ray Sun, Inho Kim +2World Model Learning Transformer-Based RL
Interdisciplinary Program in Artificial Intelligence, Seoul National University · Department of Computer Science and Engineering, Seoul National University · KRAFTON
The Intervention Gap in Latent World Models
Aug 30, 2026Donna VakalisWorld Model-Based Planning Model-Based Planning
Mila – Quebec AI Institute · University of Montreal
Transformers Struggle to Use Their Emergent World Models: Revisiting the Tower of Hanoi, and the Illusion of Thinking
Aug 7, 2026Devin Pereira, Willem ZuidemaLLM World Models LLM Planning
1Artificial Intelligence Program, University of Amsterdam · ELLIS Unit Amsterdam · Institute for Logic Language and Computation, University of Amsterdam

May 15, 2026J/K move · Enter open · S save

Identifiable Token Correspondence for World Models

Youngin Kim, Ray Sun, Inho Kim, Bumsoo Park+1

Interdisciplinary Program in Artificial Intelligence, Seoul National University · Department of Computer Science and Engineering, Seoul National University · KRAFTON

World Model Learning Transformer-Based RL World Models Video Prediction Visual RL

Open paper AskPDF 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.

Abstract

World Model Learning

Explore similar work

Identifiable Token Correspondence for World Models

The Intervention Gap in Latent World Models

Transformers Struggle to Use Their Emergent World Models: Revisiting the Tower of Hanoi, and the Illusion of Thinking