Source-linked AI summary
Can Video World Models Track Unobserved World States?
Joonghyuk Shin, Yicong Hong, Jaesik Park, Xun Huang
TL;DR
The paper asks whether video world models maintain hidden world states rather than merely render plausible frames. Using an action-conditioned visual Shell Game, it finds that standard backbones fit five swaps but fail on longer chains, while recurrent transitions with negative eigenvalues and nonlinear TTT fast weights extrapolate.
Problem
Visual plausibility does not establish that a video world model remembers, updates, and recovers states that are no longer visible.
Method
The paper evaluates action-conditioned video models on a five-cup Shell Game that visually analogizes swap-only S5 state tracking while separating rendering from hidden-state composition.
Results
Standard Transformers, Mamba, and nonnegative-eigenvalue linear attention fit five swaps but fall toward chance on longer chains, whereas negative-eigenvalue recurrence and nonlinear TTT fast weights extrapolate.
Takeaways & Limitations
Extrapolation requires carrying state across chunks, revising it in place with an expressive transition, and providing a gradient path from writes to later use.
Takeaways & Limitations
The visual Shell Game evaluates only the ball’s position, a single coordinate of the composed permutation, rather than the entire hidden permutation.
Abstract
from arXiv · showhide
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 gap with an action-conditioned video Shell Game, a visual analog of $S_5$ state tracking that decouples visual rendering from compositing the hidden state underneath. Bidirectional and autoregressive Transformers, Mamba, and linear attention restricted to nonnegative transition eigenvalues all fit the training horizon of 5 swaps and then fall toward chance on longer swap chains (extrapolation) while still rendering plausible video with additional denoising steps providing no benefit. The pixel-based diffusion target never supervises the unseen hidden state, so the generated frames cannot carry it and the state has to live inside the architecture rather than in the tokens. For a Transformer, that architectural state is only an append-only KV cache, so the model has to re-derive the hidden arrangement from the whole history at every chunk. We find two mechanisms that do extrapolate, and both carry a state across chunks and revise it in place. Linear attention succeeds once its transition eigenvalues may be negative, and TTT with a nonlinear fast weight succeeds by updating the feature map through which it reads its own state. We further examine harder cases in dynamic world exploration tasks, and discuss the broader implications for building stateful video world models.
1 INTRODUCTION
The paper tests whether action-conditioned video models maintain hidden world states rather than merely render plausible frames. Standard backbones fit short swap chains but fail to extrapolate, while recurrent states with revisable, expressive transitions succeed.
- Motivation: Video world models must update and recover hidden states when partial observations make the next frame depend on earlier actions.Plausible rendering alone does not establish state tracking.
- Motivation: Fixed-depth parallel backbones can fit their effective horizon but are expected to fail on longer NC1-hard hidden-state tracking chains.TC0 contains fixed-depth Transformers and related models, whereas the S5 word problem is NC1-hard under the stated separation.
- Empirical study: On the video Shell Game, Transformers, Mamba-based models, Gated DeltaNet variants, and bidirectional Transformers fit training length but fail to extrapolate while rendering plausible frames.Autoregressive DiT rollout and teacher/diffusion forcing do not produce a reusable scratchpad.
- Mechanisms: Length extrapolation requires carrying state across chunks and revising it through an expressive transition capable of realizing swaps.The two successful mechanisms are negative-eigenvalue linear transitions and nonlinear TTT fast weights.
- Mechanisms: Negative transition eigenvalues enable linear attention to represent reflection-like swaps, while TTT updates a nonlinear feature map so its effective kernel becomes history-dependent.The TTT update reaches W0 and W2 rather than only the readout W1.
- Broader settings: The paper extends the discussion to static and dynamic exploration, where state must also be corrected from observations and unresolved questions remain.These settings lack the Shell Game’s clean correspondence to compositional S5 tracking.
2 RELATED WORK
Related work spans video world models for interactive tasks, planning, and driving, alongside efficient recurrent sequence models and theoretical analyses of state-tracking expressivity. These studies motivate examining hidden-state maintenance under partial observation and long-horizon composition.
- Video World Models: Video world models support interactive applications, robot planning, and autonomous driving by predicting future frames from observations and actions.
- Video World Models: Architectures include bidirectional and causal Transformers, Mamba, and DeltaNet-based backbones trained with diffusion or flow-matching objectives.
- Video World Models: Recent work studies consistency as stable 3D structure and long-context memory, while harder issues include partial observation, precise physics, and multi-agent coherence.
- Linear Attention, Fast Weights, and Test-Time Training: Linear attention replaces a growing softmax-attention cache with a fixed-size recurrent state written and read online.
- Linear Attention, Fast Weights, and Test-Time Training: Delta-rule updates overwrite stale associations with residuals, while gated variants add data-dependent decay for memory management.
- Linear Attention, Fast Weights, and Test-Time Training: Test-time training treats recurrent state as small-network weights updated online, with LaCT applying nonlinear fast-weight updates over large chunks.
- State Tracking and Expressivity: State tracking maintains a hidden configuration through sequential operations, with the non-solvable-group word problem lying beyond fixed-depth TC0 models.
- State Tracking and Expressivity: Visual research connects these questions to analyses of autoregressive models and iterative denoising, plus benchmarks showing video models below human visual state-tracking performance.
3 FROM S5 TO THE VISUAL SHELL GAME
The paper turns S5 compositional state tracking into an action-conditioned visual Shell Game where the ball’s position depends on swaps performed while hidden. It evaluates video architectures and recurrent-state designs on both hidden-state accuracy and visual generation.
- 3.1 COMPOSITIONAL STATE TRACKING: State tracking composes a stream of swaps as a permutation, using S3 and the harder non-solvable S5 benchmark.The visual Shell Game asks only for the ball’s final position, a single coordinate of the composed permutation.
- 3.1 COMPOSITIONAL STATE TRACKING: Figure 1 measures Transformer depth and length extrapolation for S3 and S5, including minimum depth at N=128.The experiments compare architectural requirements and generalization after training at a fixed length.
- 3.3 THE VISUAL SHELL GAME: BENCHMARK, MODELS, AND EVALUATION PROTOCOL: The Shell Game reveals the ball, hides it during N swaps, and reveals its position afterward, with one swap aligned to each five-frame autoregressive chunk.The setup uses five identical cups, fixed-camera rendering, and 256×192 episodes.
- 3.3 THE VISUAL SHELL GAME: BENCHMARK, MODELS, AND EVALUATION PROTOCOL: The evaluated models share a latent video diffusion backbone while varying temporal mixers, attention, recurrent, and linear-attention components.The study includes bidirectional and causal Transformers, Mamba2, DeltaNet variants, LaCT, and an explicit scratchpad.
- 3.3 THE VISUAL SHELL GAME: BENCHMARK, MODELS, AND EVALUATION PROTOCOL: Most standard backbones fit the training length but fail to extrapolate, motivating architectures that carry and revise recurrent state across chunks.Figure 4 also separates hidden-state tracking from visual fidelity and tests denoising-step effects.
- 3.3 THE VISUAL SHELL GAME: BENCHMARK, MODELS, AND EVALUATION PROTOCOL: Models train on 30,000 five-swap episodes and are evaluated through 30 swaps using final ball-position accuracy and visual-fidelity metrics.Chance for state tracking is 1/5, and evaluation uses 64 episodes at chain lengths up to 30 swaps.
4 DO CURRENT VIDEO BACKBONES TRACK HIDDEN STATE?
Standard video backbones fit the five-swap training horizon and render plausible frames, but fail to track hidden state on longer chains. Their caches and readouts preserve visible content more reliably than the unobserved arrangement.
- Standard backbones fail: 100% in-distribution accuracy is achieved by full-context attention and recurrent variants, but none extrapolates robustly beyond five swaps.These models learn bounded-depth shortcuts rather than a compositional solution to the hidden permutation chain.
- Standard backbones fail: Plausible visual frames persist throughout action sequences even when models fail to track the hidden state.Visual fidelity therefore does not establish state accuracy.
- Standard backbones fail: Additional denoising steps do not change state accuracy, while four-step sampling matches 50-step sampling for working models.Denoising primarily affects visual fidelity in this setting.
- Inspecting the autoregressive cache: SWA1 through SWA4 fail even at five swaps, while SWA5 and wider windows fit training but still fail to extrapolate.A wider cache exposes more history but does not create a revisable running state.
- Inspecting the autoregressive cache: Current readouts decode visible swap identity at 1.0, whereas hidden-ball accuracy remains 0.2–0.3 across shown models and about 0.2 at N = 10.The probe contrast separates rendering-related information from hidden-state information.
- Inspecting the autoregressive cache: Successful recurrent mechanisms maintain a separate state variable for the hidden swap sequence, unlike backbones that mostly carry action- and retrieval-bound content.This distinction accompanies the observed extrapolation gap.
5 MECHANISMS THAT ENABLE STATE TRACKING
State tracking requires both a carried state and an expressive update that revises it in place. Negative-eigenvalue recurrent transitions and nonlinear fast weights provide two successful realizations, while additive memory and linear inner updates remain limited.
- Mechanisms that enable state tracking: State tracking requires a state carried across chunks and an update rule expressive enough to compose hidden transitions in place.The paper tests negative-eigenvalue transitions, nonlinear fast weights, and an additive-token control against this hypothesis.
- Negative-eigenvalue linear attention: β_t widened from [0, 1] to [0, 2] enables transition eigenvalues in [−1, 1], producing consistent gains across linear-attention variants.GDN-neg essentially solves the Shell Game with 0.95 accuracy at N = 30.
- Nonlinear fast weights: LaCT uses the parameters of a small inner network as recurrent state and updates its fast weights online across chunks.The matrices persist across autoregressive chunks and are updated inside each layer.
- Nonlinear fast weights: Four or more fast-weight heads extrapolate, while one or two heads fail beyond the training length.The experiment indicates that independently controlled heads matter more than raw fast-weight state size.
- Nonlinear fast weights: A single-layer linear inner model preserves recurrent storage and fits training but fails to extrapolate because its update remains additive like plain linear attention.SwiGLU instead produces a history-dependent feature map outside the same closed-form associative reduction.
- Nonlinear fast weights: Updating W0 and W2 while fixing W1 retains most tracking, whereas updating W1 alone fails.The evolving feature map, rather than the final readout alone, supplies the history-dependent kernel.
- Token-space control: One additive memory token reaches perfect in-distribution accuracy and generalizes to N = 10, but all token counts fall to chance by N = 20.Carrying a state is insufficient when the update rule cannot revise it expressively.
6 BEYOND THE CANONICAL SHELL GAME
The Shell Game findings extend imperfectly to dynamic exploration, where state must be updated from observations as well as actions. In Memory Maze, local-action conditioning makes pose tracking the remaining bottleneck, and LaCT helps more consistently than the negative-eigenvalue switch.
- Beyond the canonical Shell Game: Shell Game backbones fail when correct frames depend on long swap chains among five or more hidden objects.This motivates testing whether the mechanisms transfer to broader world-model tasks.
- Beyond the canonical Shell Game: General world models must update state from observations as well as actions, including clean context frames, generated frames, and dynamics evolving out of sight.This makes dynamic exploration harder than a pure action-stream permutation problem.
- Memory Maze: Under local-action conditioning, Memory Maze requires tracking camera pose before retrieving visual memory tied to that location.The model receives discrete controls and an initial camera pose rather than global pose information.
- Memory Maze: Every backbone performs well with global pose conditioning but degrades sharply with local actions, showing that retrieval is not the bottleneck.LaCT is usually among the strongest backbones, while the negative-eigenvalue switch gives no consistent gain.
7 CONCLUSION
The paper argues that hidden-state tracking in video requires an internal state that persists across chunks, can be revised in place, and remains trainable through later use. Linear attention with negative transition eigenvalues and nonlinear fast weights satisfy these requirements and extrapolate beyond the training horizon.
- Standard video backbones fit the Shell Game training horizon but fall toward chance on longer swap chains while still rendering plausible videos.The study includes autoregressive and bidirectional Transformers, Mamba-based models, and nonnegative-transition linear attention.
- Reliable extrapolation requires a state carried across chunks, an expressive transition that can realize swaps, and a gradient path from writes to later use.The paper identifies these requirements as jointly necessary for maintaining and revising hidden arrangements.
- A Transformer KV cache carries information append-only across layers, forcing each swap to read the latest arrangement at one layer and write its update at a higher layer.Internal recurrences instead apply each transition to the previous state within the same layer.
- Pixel-based diffusion targets supervise rendered content rather than hidden ball positions, so generated chunks are not trained to carry state tokens.Teacher forcing and diffusion forcing provide teacher latents, preventing later reveal errors from training earlier chunks through downstream effects.
- Linear TTT updates connect the fast state to linear attention or DeltaNet depending on the inner loss, with MSE recovering the DeltaNet rule.For chunked updates, the learning-rate matrix determines the batched transition, whose stability depends on the largest eigenvalue of G_m.
- Freezing a nonlinear fast-weight feature map leaves fixed-kernel linear attention, whereas updating the feature map makes the effective kernel history-dependent and restores tracking.In the Shell Game ablation, dot-product binding preserves tracking while MSE fails even at the training horizon.
C TOY STATE-TRACKING EXPERIMENT DETAILS
The toy experiment evaluates hidden permutation tracking in a rendered five-cup Shell Game, training models on short swap sequences and testing length extrapolation. It compares recurrent, attention, and fast-weight architectures under standardized visual and parameter settings.
- Task: Each action is a transposition, and the target is the running permutation after every swap without intermediate hidden-state supervision.The experiment uses the swap-only symmetric-group problems S3 and S5.
- Models and Depth Sweep: Models include RNNs, GRU, LSTM, Transformer, Mamba2, Gated DeltaNet, Kimi Delta Attention, and negative-eigenvalue variants.The “-neg” variants widen the transition spectrum from [0, 1] to [−1, 1].
- Training and evaluation: Models are trained on N = 128 swaps and considered solved when final-step accuracy reaches 0.95, against chance levels of 1/6 for S3 and 1/120 for S5.Extrapolation evaluates minimum-depth models at N ∈ {128, 256, 512, 1000}.
- Visual Shell Game: The renderer uses 256×192 videos, five cups, 10 swap actions, and three global controls for a total of 13 actions.Episodes begin with a reveal, apply occluded pairwise swaps, and end with a final reveal.
- Visual Shell Game: Training uses autoregressive chunks of five latent frames, with N = 5 yielding 8 chunks and N = 30 yielding 33 chunks.The corresponding episodes contain 157 and 657 pixel frames, respectively, making length extrapolation the primary diagnostic.
- Video scaffold: Each chunk contains five latent frames on a 32×24 grid, tokenized into 960 tokens by 1×2×2 patches.The scaffold uses a frozen Wan2.1 VAE with 4× temporal and 8× spatial compression.
- Persistent state: LaCT and the explicit scratchpad carry persistent per-layer state across chunks, with LaCT using fast matrices and the scratchpad using additive vectors.The configurations are summarized in the accompanying architecture tables.
D.2 TRAINING PROTOCOL
The training protocol uses diffusion forcing over independently noised latent chunks and evaluates hidden-state accuracy from the final reveal. Optimization comparisons include NorMuon, AdamW, and a final-chunk loss upweighting intervention.
- Training protocol: Each five-latent-frame chunk receives an independently sampled noise level under a flow-matching schedule with shift s = 3.All models use the Wan VAE as tokenizer and diffusion forcing as a diffusion variant of teacher forcing.
- Optimization: Models are optimized for 200K steps with NorMuon at a learning rate of 2 × 10^-4, using mixed precision except for recurrent layers kept in full precision.NorMuon converges faster than AdamW in the reported comparisons.
- Optimization: Figure A1 compares NorMuon, AdamW, and AdamW with a 10× final-chunk loss on Shell Game accuracy and Memory Maze LPIPS.All panels use Mamba2 + SWA1.
- Evaluation: Final-reveal state accuracy uses a red-ball detector over five equal-width cup regions and averages results over 64 episodes.Episodes without a detected ball count as incorrect, so zero accuracy can reflect rendering failure.
E MEMORY MAZE EXPERIMENT DETAILS
Memory Maze experiments test whether models retain and use visual state over long explorations under global-state, local-action, deterministic, and collision-dependent dynamics. Recurrent models benefit from longer observed prefixes, but state tracking becomes harder when actions require visual confirmation.
- Setup: Memory Maze models train on 400-frame clips from 29K training and 1K validation trajectories, using frozen Wan2.1 VAE latents.Evaluation reports a 220-frame continuation from either 177 or 577 clean context frames.
- Deterministic variants: In deterministic variants, fixed translations and rotations make pose a deterministic function of the initial pose and action sequence, bringing global-state and local-action results closer.This isolates local-action conditioning from drift caused by variable MuJoCo dynamics.
- Original Memory Maze: At 800F on the original Memory Maze, causal DiT reaches PSNR 29.25 versus 28.8 for Po et al.’s causal DiT, while Mamba2 + SWA1 reaches 32.18 versus 28.2 for their SSM model.These results use models trained only at 400F.
- Original Memory Maze: Most recurrent variants, including LaCT, improve with the longer prefix despite training only at 400F.The two evaluation settings are called 400F and 800F for simplicity.
- Conditioning: Global-state conditioning provides the exact pose at every frame, whereas local-action conditioning requires models to infer pose from discrete actions and visual dynamics.Models perform substantially worse under local-action conditioning on the original benchmark.
- Collision-dependent dynamics: Collisions make transitions depend on visual evidence because blocked moves leave pose unchanged, and performance drops relative to collision-free variants.The drop is larger under local-action conditioning.
- Architecture comparison: LaCT is strongest across most 800F settings, while negative-eigenvalue linear transitions provide no consistent gain in Memory Maze.GDN-neg performs similarly to GDN, suggesting different structural demands from the Shell Game.
F 3D BLOCK WORLD EXPERIMENT DETAILS
The textured 3D Block World tests local-action-conditioned rollouts beyond the training horizon, comparing recurrent updates and attention-window sizes. LaCT performs best among SWA1 recurrent backbones, while longer recent context substantially improves Mamba2.
- Evaluation: The evaluation supplies an 80-frame clean prefix and compares 80-frame training-length generations with 240-frame continuations at 320 frames.Metrics are SSIM, PSNR, and LPIPS using 200K-step EMA weights.
- Results: At 320 frames, LaCT reaches a PSNR of 29.36 and an LPIPS of 0.0288, compared with 21.32 and 0.1060 for Mamba2 under SWA1.LaCT uses a nonlinear fast-weight update without an explicit target for tracking out-of-view objects.
- Results: GDN performs similarly to Mamba2, while widening its transition spectrum with GDN-neg yields little additional benefit.
- Attention-window sweep: Widening Mamba2 attention from SWA1 to SWA10 improves 320-frame PSNR from 21.32 to 32.29 and reduces LPIPS from 0.1060 to 0.0167.The improvement peaks at SWA10, with performance declining again at SWA15.
- Cross-model comparison: In the longer comparison, Causal SWA5, Mamba2 with SWA10, and LaCT with SWA1 reach PSNR scores of 37.60, 32.29, and 29.36.The setting uses 80 context frames and a 240-frame continuation, compared with the earlier 70/210-frame evaluation.
G ADDITIONAL DISCUSSION AND LIMITATIONS
The discussion contrasts global-state and local-action conditioning, examines how video architectures carry hidden state, and identifies unresolved limits in learning and correcting state over long horizons. The paper also emphasizes that state expressivity and learnability both matter.
- Global State Conditioning vs Local Action Conditioning: Recovering the physics engine’s state from observations and discrete local actions alone, while correcting drift over long horizons, remains an open problem.Under MuJoCo, the same action can move the agent by different amounts, causing integrated poses to drift.
- State Carriers and Learnability: Pure video training provides no hidden-state target, so generated frames cannot serve as a scratchpad and state must remain inside the video backbone.The flow-matching target reconstructs frozen Wan VAE latents that carry no label for the hidden state.
- State Carriers and Learnability: LaCT tests a state carrier that updates across chunks, but its chunk size controls how much swap composition is delegated to recurrence versus fixed-depth backbone computation.For sequence length L and chunk size c, LaCT receives roughly L/c serial fast-weight updates.
- Limitations: The broader exploration settings require correcting state from observations, and the paper does not yet know which state-tracking abstraction is appropriate for video world models.
- Limitations: The mechanisms are validated only on synthetic length extrapolation tasks at roughly 200M parameters with one flow-matching recipe, leaving transfer to real-world or larger-scale video untested.
H QUALITATIVE SAMPLES
The qualitative samples visualize long-horizon Shell Game and 3D Block World rollouts alongside Memory Maze variants. Successful models preserve hidden arrangements or object positions, whereas other models may remain visually coherent while their underlying contents drift.
- Figure conventions: The qualitative appendix shows each experiment as one shared-context episode whose generated continuation is produced independently across rows.
- Shell Game: At twice the Shell Game training length, GDN-neg and LaCT8 reveal the same cup as ground truth, while the causal DiT and GDN do not.Models trained at N=5 are rolled out over N=10 swaps; the causal DiT retains its training-context SWA75 window.
- Textured 3D Block World: In textured 3D Block World, LaCT8 keeps blocks and the ball near ground-truth positions, while other backbones render a coherent room with drifted contents.All rows use SWA1 and therefore differ only in the recurrent update.
- Memory Maze: The Memory Maze figures compare ground truth and backbone rollouts across global-state and local-action conditions, including original and deterministic datasets.
- Memory Maze: The deterministic Memory Maze local-action variant keeps every backbone except GRU close to ground truth through rollout because pose is a pure function of the action stream.
- Memory Maze: With collisions under local-action conditioning, a blocked move leaves the pose unchanged, making the update depend on layout unavailable to the model’s direct view.