StateSpaceDiffuser: Bringing Long-Context Content to Diffusion World Models
Nedko Savov, Naser Kazemi, Deheng Zhang, Danda Pani Paudel, Xi Wang, Luc V Gool
World models have recently become promising tools for learning to simulate environments and  produce realistic visuals based on actions in complex settings. However, because they rely on short observation sequences, they quickly lose track of context. As a result, visual consistency breaks down after just a few steps, and generated scenes no longer reflect information seen earlier. This problem comes from a common design choice in state-of-the-art diffusion-based models: they do not keep track of a lasting environment state. To address this problem, we introduce StateSpaceDiffuser, where a diffusion model is enabled to perform on long-context tasks by integrating a rich representation from a state-space model (Mamba). The state-space branch maintains a compact latent summarizing the entire interaction history, while the diffusion branch conditions on this latent to render context-aware future frames. This design restores long-term memory without sacrificing the high-fidelity synthesis of diffusion models.To rigorously measure temporal consistency, we develop an evaluation protocol that probes a model’s ability to remember and re-instantiate previously seen content during extended rollouts. Through comprehensive experiments, we show that the StateSpaceDiffuser significantly outperforms a strong diffusion-only baseline, maintaining coherent visual context for an order of magnitude more steps. It delivers consistent views in both a 2D maze navigation and a complex 3D environment.  These results establish that bringing state-space representations into diffusion models is highly effective in both producing visually detailed and capable of long-term memory.