Current video models render the visible world remarkably well, yet fail to track its unobserved state.

Ground Truth
Causal DiT (SWA1) · PSNR 40.4 / Acc 0.12
Causal DiT (SWA1) + LaCT · PSNR 41.1 / Acc 1.00

All models are trained with diffusion forcing on 5-swap episodes. At test time each model receives the clean frames of the opening reveal, then generates all 10 swaps and the final lift from the action sequence alone. Neither wider attention nor more denoising steps help. Proper state tracking requires a carried state and an expressive update rule to revise it in place (GDN-neg, LaCT).

Abstract

Video world models are increasingly used as simulators, yet visual fidelity alone does not show that a model maintains the hidden state of the world. We examine this difference with an action-conditioned video Shell Game, a visual analog of \(S_5\) state tracking that separates visual rendering from compositing the unobserved state. Trained on 5-swap chains, standard backbones (e.g., bidirectional and autoregressive Transformers, Mamba, and linear attention) render plausible videos and predict the location up to 5 swaps, but fail to learn the rule and generalize to longer sequences, with more denoising steps providing no improvement. As the pixel-based diffusion loss does not force the generated frames to hold the unseen ball position, output tokens cannot carry it, and the state has to live within the architecture. In a causal Transformer, this implicit state is an append-only KV cache, which is written once and never revised, so the model must re-compose the swaps at every chunk. Tracking \(S_5\) this way requires depth to grow with sequence length, and no fixed-depth Transformer can provide this. We study what enables learning the rule, and find that length generalization requires a state carried across chunks and revised in place through an expressive update. Linear attention achieves this by allowing the eigenvalues of its transition matrix to become negative, and a nonlinear TTT fast weight (SwiGLU) by updating the feature map through which it reads its own state. We further examine harder settings where the state is not fixed by actions alone, and discuss the implications for building stateful world models.

The \(S_5\) Swap Problem

Pixel-free state tracking on S3 and S5

\(S_5\) is the symmetric group of all \(5! = 120\) permutations of five objects. We use the swap-only variant, where each input is one transposition rather than an arbitrary permutation. Since any permutation of five objects decomposes into at most four swaps, this restriction remains \(\mathrm{NC}^1\)-complete. Under standard log-precision assumptions, fixed-depth Transformers and state-space backbones such as Mamba lie in \(\mathrm{TC}^0\), a constant-depth parallel class, while \(\mathrm{NC}^1\) permits depth that grows with sequence length (Merrill et al., 2024). Prior work supplies that serial depth through CoT scratchpad tokens (Merrill and Sabharwal, 2024) or repeated backbone iterations.

The figure first measures the depth each backbone needs, then tests length extrapolation after training at \(N=128\). Nonlinear RNNs can apply one swap per recurrent step. Transformers and Mamba2 instead learn bounded-length shortcuts. GDN and KDA carry matrix states, but their standard nonnegative transition eigenvalues cannot realize the reflection required by a swap. Allowing eigenvalues in \([-1,1]\), marked GDN-neg and KDA-neg, admits that reflection and enables much stronger extrapolation (Grazzi et al., 2024).

The Visual Shell Game

Shell Game episode construction

State tracking has mostly been studied in symbolic sequence models. We bring the same question to generative video world models. We adapt the Three.js and WebGL renderer of Liu and Lee, 2026, adding explicit lift and reveal phases plus action annotations. Five identical cups hide a ball across \(N\) swaps, and the final reveal tests the composed state without direct supervision. We train on 30,000 five-swap episodes and evaluate up to 30 swaps (chance \(1/5\)).

Standard Video Backbones Fail

Model architecture

Comparing Backbones in a Shared Video Diffusion Scaffold

Within the same roughly 200M-parameter Wan-style latent video diffusion scaffold, we test bidirectional and causal Transformers, nonlinear RNNs, Mamba2, GDN, KDA, LaCT fast weights, and explicit scratchpad tokens. The temporal mixer and carried-state mechanism vary, while the rendering pipeline, training objective, and parameter scale remain comparable.

Length generalization and denoising-step sweeps

Visual Fidelity Does Not Imply State Accuracy

Bidirectional and autoregressive DiTs, Mamba2, and plain linear-attention backbones such as GDN/KDA with nonnegative transition eigenvalues fit the five-swap training horizon but fail as the chain grows, even while their videos remain plausible. PSNR barely separates the models that track the ball from the ones that guess. Additional denoising steps do not recover accuracy, because they do not add serial depth to a fixed-depth backbone (Liu et al., 2025).

The KV Cache Is a Visual History, Not a Running State

SWA size ablation and hidden-state probes

A wider SWA window recovers the five-swap training horizon, but every window fails at ten, so more cached history does not yield extrapolation. Probing the caches tells the same story: the ball stays decodable from GDN-neg and LaCT8 states, but not from the append-only KV cache. See the paper for probing details.

Mechanisms That Enable State Tracking

Mechanisms that enable state tracking

State tracking requires two components, a state carried across chunks and an update rule expressive enough to compose the hidden transition in place. Autoregressive Transformers lack the carried state, while nonnegative-eigenvalue linear attention lacks the expressive update. The two mechanisms below address these limitations separately.

Negative-eigenvalue linear attention. A DeltaNet-style transition \(I - \beta_t k_t k_t^\top\) with \(\beta_t \in [0,1]\) can erase and overwrite but never reflect. Widening to \(\beta_t \in [0,2]\) reaches eigenvalues in \([-1,1]\), and at \(\beta_t = 2\) the transition becomes the exact reflection that realizes one swap. This consistently improves every linear-attention variant we test, with GDN-neg essentially solving the task (0.95 accuracy at 30 swaps).

Nonlinear fast weights (LaCT). A test-time-training layer whose recurrent state is a small SwiGLU MLP \(f_W(x)=W_1[\operatorname{silu}(W_0 x)\odot(W_2 x)]\), rewritten in place at every chunk. Ablations show the mechanism lives in the feature map: freezing the readout \(W_1\) retains extrapolation, while freezing \((W_0, W_2)\) or replacing the SwiGLU with a single linear layer destroys it. Because each swap updates the feature map through which the state is read and written, its effective kernel becomes history-dependent. With a frozen feature map, TTT-KVB reduces to additive linear attention (Liu et al., 2026). That additive form fails to extrapolate on \(S_5\), while nonlinear LaCT does. Head count matters as well: one or two fast-weight heads fail beyond the training length, while four or more extrapolate despite carrying less raw state. Visual rendering need not occupy every channel, so extra heads can decouple tracking from the pixels and specialize in the hidden state (we use eight, LaCT8).

Explicit scratchpad tokens, an additive token-space register carried across chunks, serve as a control: a single memory token already lifts the at-chance SWA1 baseline to perfect in-distribution accuracy, but the plain additive update lets the carried norm explode on longer rollouts and accuracy drops to chance by 20 swaps.

Visual Shell Game Rollouts

All backbones reach 1.00 accuracy at the five-swap training length.
Ground Truth (Ball Position: 4)
Causal DiT (SWA6) · Acc 1.00
Mamba2 · Acc 1.00
GDN · Acc 1.00
GDN-neg · Acc 1.00
LaCT8 · Acc 1.00
GDN-neg and LaCT8 remain at 1.00. Standard backbones begin to fail even though the videos still look plausible.
Ground Truth (Ball Position: 5)
Causal DiT (SWA6) · Acc 0.22
Mamba2 · Acc 0.44
GDN · Acc 0.25
GDN-neg · Acc 1.00
LaCT8 · Acc 1.00
GDN-neg and LaCT8 hold at 0.98. Standard backbones are near chance.
Ground Truth (Ball Position: 4)
Causal DiT (SWA6) · Acc 0.27
Mamba2 · Acc 0.11
GDN · Acc 0.22
GDN-neg · Acc 0.98
LaCT8 · Acc 0.98
GDN-neg reaches 0.95 and LaCT8 0.69. Standard backbones remain near chance.
Ground Truth (Ball Position: 5)
Causal DiT (SWA6) · Acc 0.30
Mamba2 · Acc 0.20
GDN · Acc 0.19
GDN-neg · Acc 0.95
LaCT8 · Acc 0.69

Beyond the Canonical Shell Game

Ground Truth
LaCT

Memory Maze with collisions (LaCT, 800 frames, 2× training length). The model tracks its pose, retrieves location-specific memory, and corrects blocked moves from observed collisions.

Ground Truth
LaCT

Dynamic Block World (LaCT, 320 frames, 2× training length). The model tracks blocks that keep moving off-screen, so the hidden state evolves independently of its own actions.

In the real world, state rarely follows deterministically from actions alone: it may need correction from what is observed, keep evolving out of sight, or change independently of the agent's actions. Memory Maze and Dynamic Block World probe two of these harder update structures. In both, LaCT, which we train from scratch as a video world model, is again the strongest backbone, while the negative-eigenvalue switch that is decisive on the Shell Game gives no consistent gain. The algebraic structure that composes an action-only permutation chain does not carry over, while the nonlinear fast-weight update appears to carry over in part. It remains open which architectures can cover the far broader range of updates real world models need.

BibTeX

@article{shin2026statefulvwm,
  title   = {{Can Video World Models Track Unobserved World States?}},
  author  = {Shin, Joonghyuk and Hong, Yicong and Park, Jaesik and Huang, Xun},
  journal = {arXiv preprint arXiv:2608.30692},
  year    = {2026}
}