Title: StateSpaceDiffuser: Bringing Long Context to Diffusion World Models
Venue: NeurIPS 2025
Accepted: yes
Keywords: world model

state-space model

long context

generative model

computer vision
Abstract: World models have recently gained prominence for action-conditioned visual prediction in complex environments. However, relying on only a few recent observations causes them to lose long-term context. Consequently, within a few steps, the generated scenes drift from what was previously observed, undermining temporal coherence. This limitation, common in state-of-the-art world models, which are diffusion-based, stems from the lack of a lasting environment state.

To address this problem, we introduce StateSpaceDiffuser, where a diffusion model is enabled to perform long-context tasks by integrating features from a state-space model, representing the entire interaction history. This design restores long-term memory while preserving the high-fidelity synthesis of diffusion models.

To rigorously measure temporal consistency, we develop an evaluation protocol that probes a model’s ability to reinstantiate seen content in extended rollouts. Comprehensive experiments show that StateSpaceDiffuser significantly outperforms a strong diffusion-only baseline, maintaining a 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 demonstrating both visual details and long-term memory. Project page: https://insait-institute.github.io/StateSpaceDiffuser/


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a method for incorporating long-context understanding into diffusion-based world models by combining a state-space model (SSM) with a diffusion model. The SSM processes long input sequences and maintains a latent environment state that accumulates information over time. This environment state, along with a short history of image frames, is then fed into the diffusion model to predict future world states in image space.\n\nThe key insight is that while diffusion models excel at generating high-quality images, they are limited to short sequences due to memory constraints. In contrast, SSMs can handle long sequences efficiently by recursively updating a compact internal state. The proposed approach leverages both strengths: the SSM captures long-term context, and the diffusion model handles high-fidelity image prediction.\n\nTo assess the model's long-term reasoning capabilities, the authors introduce an evaluation framework where the task is to observe a long sequence and then predict it in reverse. Experiments are conducted on two environments: a 2D maze setting (MiniGrid) and a 3D first-person shooter game.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n1. The paper addresses an important problem: how to build long-term context understanding and how to build in a memory mechanism into generative world models. As the authors rightly point out, this capability is essential for a range of downstream tasks, and current models are limited to only operating on short time scales.\n2. The evaluation task is intuitive and well-suited for assessing long-context reasoning. Predicting an episode in reverse after observing a few frames resembles a natural form of episodic recall.\n3. The proposed model outperforms a diffusion-only baseline, demonstrating the effectiveness of combining a state-space model with a diffusion model.\n4. The experiments are conducted in both a simple 2D environment (MiniGrid) and a challenging 3D first-person environment, providing a reasonable spectrum of test settings.\n\n**Weaknesses**\n\n1. While the paper tackles an important problem, the experimental settings feel somewhat limited. Both environments—MiniGrid and CSGO—are relatively constrained compared to real-world scenarios. This raises questions about the method's scalability and general applicability.\n2. The state-space model (SSM) operates on a compressed state representation and this likely contributes to the poor reconstruction quality observed in Figure 6. I think this calls into question the practical utility of the method, or atleast makes the method not super impressive. The authors acknowledge the PSNR metrics limitations in the supplementary material, but I strongly feel such discussion belongs in the main text to properly contextualize the model’s performance.\n3. Although the model outperforms a diffusion-only baseline, this result is not very surprising, given that the SSM has access to long-term history. The improvement itself doesn’t offer particularly surprising insights. In some sense it’s just a natural consequence of the architecture so it’s hard to say that it’s a major conceptual breakthrough.\n4. The paper claims strong results but lacks a detailed discussion of failure modes. It’s unclear how often or why the SSM fails, especially given its lossy state representation. Without this, it's difficult to understand the model's limitations or how future work might build on it.\n5. The need to train the SSM and diffusion model separately is briefly mentioned, but not adequately justified or explored. This separation could be a significant limitation, as it prevents end-to-end optimization and ties overall model performance to effectiveness of the SSM. It’s unclear whether the diffusion model is truly doing anything more than rendering from the SSM’s output.\n6. The inclusion of a user study feels unnecessary and potentially misleading. Visual artifacts in Figure 6 already make it clear that the baseline underperforms, and qualitative metrics like PSNR or LPIPS would have provided more objective insight. A user study comparison seems uninformative in this context and could be omitted.\n7. Writing quality could be improved, especially in making architectural choices and limitations more transparent. There are also quite a few typos and presentation issues which I discuss below. \n\n**Writing and Clarity Suggestions**\n\n*Introduction:*\n\n1. The authors make the statement that Diffusion models are most often directly limited by their input window size. This might need more clarification --- Is this limitation due to memory constraints, or are current architectures fundamentally incapable of working with longer contexts?\n2. A positioning relative to relevant recent work, such as Long Context Tuning for Video Generation (Guo et. al), which challenges the idea that diffusion models are strictly short-context could be useful.\n3. The paragraph at the end that discusses metrics and points to Section 5 disrupts the flow. \n4. Line 66: The statement; \"effectively brings a state-space model and diffusion\" is awkward.\n \n*Data:*\n\n1. A visual illustration of the task environment or gameplay (especially for the CSGO setup) would make it easier to understand what the model is actually trying to predict.\n\n*Typographical and Caption Issues:*\n\n1. Line166: Check for a possible typo \n2. Figure 6: Do downward arrows indicate that both frames and actions are provided, or just actions? Clarify this in either the figure or its caption.\n3. Figure 5 caption: Missing space after the title—should follow standard formatting\n\nQUESTIONS:\n1. If the full sequence of 50 frames can be held in memory, is the problem essentially solved with a standard diffusion model? Is the real limitation just memory? It would be helpful if the authors discussed this more explicitly, as it seems central to the motivation for introducing the state-space model.\n2. Clarification on MiniGrid Setup (L238): The paper states, \"On MiniGrid, we generate 1 frame in the future...\" Why is only one frame predicted here, whereas in CSGO the model generates the second half of the sequence? Section 3 does not clearly explain this discrepancy. \n3. In (L247) the paper says, \"We train and test on the same set of 34 samples, allowing for overfitting.\" Could the authors elaborate on the reasoning behind this choice? Is the goal to ensure that the state-space model is well-trained before evaluating the generative model and that is only possible on a small dataset? If so, how does this affect scaling to larger datasets?\n4. The authors note in the supplementary that the state-space model shows more artifacts in CSGO than in MiniGrid, suggesting room for improvement. However, since the SSM must be trained separately, how scalable is this method to more complex real-world environments? 5. Would the quality of the generative model always be bottlenecked by the expressiveness of the SSM? The paper states, \"We load the pre-trained weights and train the Generative Branch to obtain the final model.\" Is the SSM frozen during this stage, or is it fine-tuned jointly with the generative model? Additional details on the training pipeline would be very helpful.\n\nLIMITATIONS:\nA discussion on how the model would scale to larger datasets and failure models on CSGO were not sufficiently discussed. More details are mentioned in the Weaknesses section.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nExisting world models often face insufficient memory issues due to the intractable increase in computation. This paper addresses this challenge by introducing Mamba, which aggregates all historical actions and frames with linear complexity. Conditioned on the global state extracted by Mamba, a two-stage diffusion model is employed to predict the next frame. The model shows superior long-term consistency in both the MiniGrid maze environments and the CSGO benchmark created by the authors.\n\nSTRENGTHS AND WEAKNESSES:\n# Strengths\n\nS1) The paper aims to address a very challenging and intriguing topic in this field.\n\nS2) Both quantitative evaluations and a careful user study are conducted for a comprehensive assessment of the proposed model.\n\nS3) The writing is clear and easy to follow.\n\n# Weaknesses\n\nW1) **Computation concern still exists.** If my understanding is correct, the computation cost still grows linearly with the context length. How will the model manage context in real time when dealing with hundreds of frames?\n\nW2) **Limited baselines.** The proposed model is compared only with a few baselines, primarily variants of itself. It would be more convincing to include some state-of-the-art methods, such as Diffusion Forcing (NeurIPS 2024), in the comparison.\n\nW3) **Inferior qualitative results.** From Figure 6, the visualization after 17 frames looks quite blurry. Could you discuss potential reasons and solutions for that?\n\nW4) **Duplicated references.** References 20, 21 are cited twice in the submitted manuscript. Please fix it in the revision.\n\nQUESTIONS:\nQ1) It is unclear how the model handles various context lengths at test time. Does the model always use a fixed context length during training? If the maximum training length is 50, can the model generalize to longer contexts in a zero-shot manner?\n\nLIMITATIONS:\nPlease see the Weaknesses part.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nTo tackle the visual consistency problem for long sequences in diffusion world models, the paper proposes StateSpaceDiffuser with the aim to compress the interaction history into some latent space via a state-space model. Briefly, this consists of state-space model that operates over long sequences and learns to predict $\\hat{f_t},… \\hat{f_{T+1}}$ and the last four steps from this predicted are concatenated with action embeddings and passed as input to the diffusion model.  To evaluate temporal consistency by measuring how well can re-instantiate seen content on extended rollouts in 3 settings.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: The paper addresses a known limitation of diffusion-based world models i.e., their inability to maintain visual consistency over long sequences due to short context windows. To tackle this, it presents a rather straightforward approach to condition the diffusion model with state-space model features. The key strength of the paper is identifying an appropriate pipeline for training this; separate out the SSM training via Cosmos tokeniser features, and then load the pre-trained weights to train the diffusion. The evaluation protocol is well-designed, particularly the forward-backwards navigation setup that directly tests model capacity on MiniGrid and CSGO environments. The paper reports improvements on the baseline with 39.68 dB PSNR on MiniGrid for context length 50, and positive user study results on CSGO. Additionally, I appreciate the ablation studies to investigate how much the SSM is helping.\n\nWeaknesses: The paper feels like an application contribution of combining two known techniques (MAMBA and diffusion) for learning WMs, with limited evaluation. For starters, the evaluation setting, although a useful starting point, is quite contrived and needs to be evaluated with 1. Long-time horizon, 2. one or two more environments with different levels of visual details, e.g, NetHack,  3., across more seeds to show reliability in the results presented. Furthermore, the CSGO evaluation relies on subjective user studies rather than objective metrics\n\nQUESTIONS:\nTo better understand this paper, it would be good to have a clear exposition of the following:\n\n1.\tWhat is the computational cost of the space-space model augmentation? It would be good to normalise this with the performance boosts in SSIM and PSNR. Particularly because confidence intervals are not reported due to computational constraints. However, it is celebrated repeatedly in the paper that the improvements are at essentially no additional computational costs.\n\n2.\tGiven that the longest context window being used is 50 – I would like an evaluation with context window of 100 and 150 at the minimum. This would help understand at what point is the prediction appropriate? \n\n3.\tFor evaluation, it would be good to understand if the model is trained for context of 50 – how much further can the model be unrolled (example 10 steps more) and then go back in sequence to the starting point. \n4.\tThere seems to be research with human subjects but identified as N/A in the question. > ‘In CSGO we perform a user study, more aligned to the visual complexity of the environment’ \n5.\tSomewhere in the paper, we need a clear definition of the agent for the minigrid settings. Currently, the policy is specified as visit 40 markers in a sequence. What does that entail exactly? \n6.\tThe action embedding is separate in the generative (512) and long-context (16) branch? What is the reason for this? \n7.\tHave you tried adding noise in the middle of the sequence and evaluated how robust is the method for this?\n\nLIMITATIONS:\nYes, discussed in detail in supplementary.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper tackles a key challenge in diffusion-based world models: the lack of long-term consistency in generated sequences due to their reliance on short observation windows. The authors propose StateSpaceDiffuser, a novel architecture that integrates a state-space model (specifically, Mamba) with a diffusion model to maintain a persistent memory of the environment throughout long rollouts. Experiments on both a synthetic 2D maze environment (MiniGrid) and a visually complex 3D environment (CSGO) demonstrate substantial gains in long-term fidelity over baseline diffusion models.\n\nSTRENGTHS AND WEAKNESSES:\nLong-term consistency is a very important problem for building world models, and the paper makes a solid first step toward addressing it. It studies the use of an external state model to maintain long-term memory, and the experimental results clearly show that this strategy helps. At first glance, it might seem like a straightforward idea, but I think the valuable part of this paper is that it successfully demonstrates the feasibility of this approach and opens up many interesting research directions—for example, how to improve the state features to contain more detailed information, or how to adapt the approach to more complex tasks beyond simply reversing actions.\n\nOne weakness is that the paper doesn't include a very extensive ablation on the choice of state model. I’m curious to see how simpler baselines like RNNs, LSTMs, or even vanilla transformers (e.g., LLMs) would perform. These comparisons could provide more insight into what properties of the state model really matter in this setup.\n\nMissing citations:\nXiang et al. Pandora: Towards general world model with natural language actions and video states\n\nQUESTIONS:\nHave you tried end-to-end training of the state and diffusion parts together, instead of training the state model separately first?\n\nLIMITATIONS:\nThey mentioned their limitations in supplementary material that the reconstruction quality is still not so good right now.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-07-04"}