StarWM: Self-Supervised Trained Attention Routing for Robust World Models

A robust world model must strike the balance between faithfully capturing environmental dynamics and abstracting away from irrelevant content. While reconstruction-based world models ensure faithful supervision, they misallocate representational capacity by pixel area rather than dynamics relevance, which can cause task-irrelevant content to dominate the learned representation. Alternatively, reconstruction-free methods avoid this bias but risk discarding possibly relevant information. We propose StarWM, which combines advantages of both paradigms through a cross-attention module gated by self-supervised dynamics signals that decides where reconstruction applies, while a dual-stream decoder lets reconstruction shape only what is extracted at attended locations, with stop-gradient barriers preventing the two objectives from interfering. This allows reconstruction to supervise the visual content of attended regions without contaminating the latent with non-predictive information. On DeepMind Control with dynamic video backgrounds, StarWM outperforms reconstruction-based baselines while keeping competitive with the concurrent reconstruction-free approach, and is the only method that performs consistently well across clean, sequential-video, and random-video distractor settings. Mechanistic probing confirms StarWM preserves state attributes with near-perfect fidelity through long-horizon imagination while systematically discarding distractors.