StateSpaceDiffuser: Bringing Long Context to Diffusion World Models
Authors: Nedko Savov
Naser Kazemi
Deheng Zhang
Danda Pani Paudel
Xi Wang
Luc Van Gool
Venue: NeurIPS 2025 poster ()
Keywords: world model
state-space model
long context
generative model
computer vision
Primary area: deep_learning
TL;DR: We enhance diffusion models to be able to recall content from long back in a sequence in order to produce consistent content

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/

StateSpace Diffuseer:: Bringing Long Context to Diffusion World Models
Nedko Savov!!
Naser Kazemi' Xi Wang 12,3 INSAIT, Sofia University "St. Kliment Ohridski"
Deheng Zhang'
Danda Pani Paudel'
Lue Van Gool!
2 ETH Zurich
3 TTU Munich
History
Diffusion
StateSpaceDiffuser (Ours)
Ground Truth
I1
It
Iz
I1
Figure 1: Recalling content long the past. Given a history images I,,... ..., IT and accompanying actions, we navigate all the way back to the beginning- /1. The task is to generate frames along the way, consistent to what is seen in the history, given the actions. predictions of the first frame f1. Can a generative model recall the content 17 long back in the sequence? Diffusion models fall short (X), our model correctly recalls the content of 1 (V).
an example, we show the
Abstract
World models have recently gained prominence for action-conditioned visual prediction in complex environments. However, relying on only few recent observations causes them lose long-term context. Consequently, within few steps, the generated scenes drift from what was previously observed, undermining temporal coherence. This limitation, common state-of-the-arr world models, which are iaatussionbbsed,, stems from the lack of a lasting environment state.
To address this problem, introduce StafeSpaceDiffuuser where diffusion model is enabled to perform long-contex: tasks by integrating features from a state-space model, representing the entire interaction history. This design restores long-term memory while preserving the high- high-fidelity synthesis of diffusion models.
To rigorously measure temporal consistency, we develop an evaluation protocol that probes model's ability reinstantiate seen content in extended rollouts. Comprehensive experiments show that StateSpaceDiffuser significantly outper- forms 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 complex 3D environment. These results establish that bringing state-space representations into diffusion models is highly effec- tive in demonstrating both visual details and long-term memory. Project page: https::/insait institutee gi thub. Teabb.oo/ttateee io/StateSpaceDifffse!
"Corresponding author: nedko. savovéinsait.ct
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Diffusion
StateSpaceDiffuer
Diffusion
Diffusion
State-Space
Model
Figure 2: Our Approach. While diffusion models are limited to a short sequence input, our approach enables long-context processing for diffusion models with a state-space representation.
Introduction
World models have gained popularity for the production of visual consequences of given past observations and actions. These models can learn to generate environment observations entirely by training many interactions with the environment. Simply by observation, they capable of handling complex environments, such as car driving [27 29
41
66), 3D virtual environments navigation 2 84 They
77 191, platformer games g (88. ego-centric action videos 82
enable interactivity without the burden of hand-coding complex environments, but also offer feature representations for roboties and reinforcement learning agents for planning.
For long interaction with world models, essential that the generated video remains consistent with previously observed or generated content. Revisited areas should preserve their appearance, and objects observed again should keep their properties. However, as shown in Fig. current high-fidelity world models, which are mostly based on diffusion, cannot preserve context outside of a short time window, most often directly limited by their input window size [2 T9 8. This leads to increasing drift in content over time, where earlier information is forgotten overwritten. The inability to retain persistent memory of the environment poses a major challenge, especially for real-world applications such agent planning and virtual interaction, where coherent, temporally consistent environments are essential. Therefore, in this work our task is to stay consistent with a long history of past inputs, even if generating single future frame. This is in contrast long generation which focuses on producing extended realistic sequences, even without prior context.
To improve content consistency in diffusion-based world models, make of persistent long- context representation. Specifically, we leverage features from a discrete state-space model (Mamba), which has been shown be very effective at capturing long-term context in prior work 20 31) We summarize our system in Fig. 2 Although these models were previously applied to language and visually simple environments [12], our goal is preserve the long-term context in modern diffusion-based world models, targeting environments with higher visual complexity such as CSGO
600 In contrast, other state-based models, such as those using L.STMs GRUs (38 39 40] have
limited generative capacity and are mainly used for agent planning. Our approach combines the strong generative power of diffusion models with the long-term context tracking of state-space representations.
Importantly, the state-space model (SSM) is computationally efficient, which allows it process arbitrarily long sequences. This is achieved by maintaining a compact state that is updated at every sequence step. During training. SSMs have linear complexity in sequence length [31], further improved by parallelization. Unlike them, CNNs have a fixed receptive depth, and transformers characterize with heavy quadratic complexity. At inference, in streaming fashion, the SSM be executed with constant per-step latency and constant memory footprint, whereas transformers and CNNs, best, still grow linearly per step. As we show, test time, our proposed model scales far beyond its trained horizon, while the SSM contributes less than 2% of the total inference compute of the full model.
Our proposed model, StateSpaceDiffuuser, summarized in Fig. consists of state-space model that operates the long sequence, and diffusion model conditioned on both a short window of observations and state-space model features. The latter enables the diffusion branch to generate the content of the next frame conditioned long context rather than the last few frames.
To evaluate the consistency the generated long context of StateSpaceDifftuser, we design and develop an evaluation framework that involves navigating environments to then return back to the

---

initial position. We evaluate on two environments. (1) A simpler 2D maze environment (MiniGrid), in which we establish the presence or absence of memory ability by remembering the maze layout given partial observations. And (2) 3D first-person shooter game (CSGO), which serves to show the performance of our method on a visually challenging interactive environment with many factors at play. Our quantitative and qualitative evaluation results show that StateSpaceDiffusee produces content significantly more consistent with long history than a diffusion-only method. Evaluation in the maze environment yields 51.9% PSNR improvement over the baseline average (56.3% improvement on the most memory challenging cases). A user study confirms that our method produces images closer to previously observed content in the CSGO dataset compared to baselines. More details are shown in Sec. S
Our contributions are as follows.
We propose StateSpaceDiffusee, which integrates a state-space model with a diffusion model for visual world modeling. It is capable of generating consistent content in long-horizon generation, with almost no extra computational cost.
We develop an evaluation protocol to test the content preservation abilities of a world model and perform extensive evaluations of world models on long-horizon generation tasks.
Our evaluation shows significant quantitative improvement and a strong user preference over the baseline in the case of long contexts. Furthermore, our studies attribute the improvements to our model design and confirm generalization to longer contexts.
2
Related Work
2.1 World Models
Generative environment models. Initially developed as imagination-based models for training
mosel-based reinforcement learning (MBRL) agents 13
38 40 70), world models have evolved into
powerful generative systems that condition on actions to produce future frames [T] 57 [33 Early work by (36] demonstrates that training recurrent latent dynamics on VAE image representations can enable agents to plan in imaginative rollouts. Extensions such as SimPle 47 and Dreamer 337 refine this approach by improving reconstruction quality and stability, culminating in DreamerV2 and DreamerV3 [39 400 systems that achieve human-level performance on Atari and demonstrate the ability to generalize across diverse domains. More recent efforts, such as IRIS (88) TWM (64), STORM [86), and DayDreamer [79]. employ Transformer-based hybrid backbones and focus on sample efficiency, long-horizon coherence, or robotic control. However, many these methods rely on discrete latent tokens and relatively short contexts, which limits visual fidelity in complex scene motion or when extended rollouts are required.
World models are also central to realistic video generation conditioned on actions. Genie S leverages video tokenizer and a Latent Action Model for dynamic next-frame generation, whereas GAIA-1 46). GAIA-2 [66] tackle autonomous driving by autoregressively predicting image tokens from multi-modal inputs. Recent works highlight broader applicability and complex generative capabilities. DINO-WM [87 uses pretrained visual features for zero-shot planning, GameFactory (84) adapts game environment actions to realistic environments, while allow video generation control by periodical text instructions. Both illustrate how world models can transcend traditional RL frameworks and support open-ended content creation.
Diffusion- approaches. Parallel to these developments, diffusion models [73 44 75 have emerged as powerful class of generative methods for high-fidelity image and video synthesis. They have been applied to text-to-vides viveo [7I]. space-time video generation [. and broad world simulation tasks [6]. Within MBRL, DIAMOND I uses diffusion model to generate high-quality frames for Atari, making for playable environments and enhancing agent performance. Methods like Pandora 80) and LCT 35 generate video based on periodic text instructions. Nonetheless, current diffusion- based world models, typically transformer-bases condition on only short window of past frames to handle the quadratic computational complexity, making long-horizon dependencies difficult to maintain. This makes it challenging to maintain long-horizen horizon dependencies.
3

---

2.2 Sequence Modeling
RNNs and Transformers. Sequential modeling has historically been dominated by recurrent neural networks (RNN) such as LSTM and GRU 45 17] which process input tokens step by step and are able to handle moderate contexts. However, RNNs often struggle with extremely long sequences due to vanishing gradients and limited memory capacity [59]. Transformers [78 addressed these issues by employing self-attention, making them effective in capturing Long-range dependencies. Beyond world modeling, Transformers have become the backbone for a broad range of tasks, including language modeling [21
7 63 and computer vision [2
28), due to their ability to handle global context.
Various Transformer variants have attempted to reduce the quadratic cost of self-attention for long
sequences [50] 16 85 [4]. Vision-specific models like Swin [56] or MVIT [24] adopt hierarchical
or local attention, yet scaling them to long video horizons remains computationally prohibitive.
Previously, DFoT
I0 addressed the ability for long future prediction. However, the long-context
consistency problem has only been recently addressed by few concurrent works. [51 74 81
improve context abilities by proposing strategies to sample number of historical observations to use as conditioning. Instead, our approach involves summarizing information from the entire history automatically through state-space models.
State-Space Models (SSMs). As an alternative, SSMs 5 53 76 can process sequences in linear time by learning continuous dynamics in a latent state. Representative structured state- space models include S4 [33 34] and H3 [18] that generalizes the recurrence in Linear Attention 49]. S4, S5 722, and S6
232 leverage carefully designed operators (e.g. HiPPO matrices (32) to efficiently capture long-range dependencies. Mamba [31 introduces selective gating to improve expressiveness without sacrificing linear scalability. S4WM [2] has shown that applying SSMs as world models shows promise for maintaining coherence over hundreds of imagined steps while preserving computational tractability.
Hybrid Architectures. As Transformers excel at local interaction with low computational cost and SSMs can capture long-horizon dependencies efficiently, hybrid designs have been proposed for vision tasks. Mamba Vision 42] incorporates state-space models into transformer, and Dimba [26]. DiS (25 into a diffusion network backbone. for computationally cheaper image discriminative and generative tasks, operating on image patches. (54) modify the softmax in attention to emulate a forget gate and improve transformer context abilities. Mamba VLT [55] and Samba (6) exploit state-space models for better object tracking with long-range consistency.
3 Data
We design an experimental protocol to evaluate long-term content consistency in diffusion world models, comprising of three experimental setups with rising level of complexity, based on controlled maze environment (MiniGrid) and a complex 3D first-person environment (CSGO).
We create a dataset based on the partially observed MiniGrid maze environment [2]. In this setup, each maze consists of a grid where each cell can be a wall, an empty space, or colored marker. Markers act like empty spaces, but are visually distinct. An agent navigates the maze. but at each time step, only sees a 7x7 window centered around itself rather than the full 85x85 maze (see Fig. 6 (b) for example). We use modified version of MiniGrid with randomly generated mazes, allowing us to adjust the size, wall complexity, and number of color-coded markers. each episode, the agent is tasked with visiting sequence of 40 random markers via the shortest path. Once halfway through the episode, the agent stops following the path and retraces its steps back to the starting point. Each episode is 100 steps long (50 forward, 50 backwards). We evaluate different context lengths by selecting subsequences around the long sequence center. Notably, the second half of each sequence depends heavily the model's ability to recall earlier frames, making it ideal for testing long-context reasoning.
We also design simplified dataset called MiniGrid Simple, consisting of just 34 samples without walls and a single marker placed behind the starting position. The agent moves three steps forward and three steps back, returning to its initial position. Since the context window of our baseline is just 4 steps, this setup provides a minimal but effective test of long-term recall. We use this to compare the performance of our baseline and state-space-chhaneed models in reconstructing the marker color.

---

Generative Branch
Output
Short Context
In.t
Diffusion
Tz
In-g
PT-2
Action Enbedding
Io.g
21-3
Fusion
Long-Context Branch
Fusion
History
10-5
4--5
Concat
1
22-6
In.7
Ag-7
Action Embedding
Noise
LayerNorn
1--8
State-Space Model
Encoder
Input
Figure 3: Architecture of our StateSpaceDiffuser model. consists of: a state-space model for processing long context information; diffusion model generating high-fidelity context-aware next observation, conditioned on state-space features.
To evaluate our method in a more visually complex setting, we use CSGO [60] a dataset human gameplay 3D first-person shooter. It includes 51 action types, such 23 rotational commands, movement directions, jumping, and various special actions (e.g., firing, changing weapons), To adapt the dataset for long-context testing, create mirrored sequences: for each original sequence, append its reversed version, ensuring that actions are also reversed. We use a corresponding one-hot encoding (e.g. turning left becomes turning right), or create new one correspondence does not exist (e.g. jump, shoot). This setup forces the model to rely on information from earlier in the sequence when generating later frames.
Methodology
Given sequence of environment interactions 01, a2 ..., 0T-1, the resulting observations 41, 12, ....... (with initial frame I1) and the current action ar, the objective of a world model F(-) to produce the next image T F([T,, ..., Ir]. (1, o...))rr)).. Recently, the best- perttrrefoomngggg generative ar- chitectures for modeling F(-) are diffusion models based on transformers or UNet. In training, transformers computationally intensive O CNNs have fixed receptive fields and are ill- fitted to long-term dependencies. Therefore, these models take a short history window of observations: ..., ur)). (e.g. K 41] K 16(8). With long and growing sequence, the short history causes the loss of long-term temporal coherence. Instead, we propose to efficiently process the long sequence (O(T) with model designed for this purpose a state-space model. Such models maintain and update a state with each sequence step, and the serves summary the sequence so far. Extracting long-context features in this way and integrating them into the diffusion pipeline yields the proposed model StateSpaceDiffuser.
IT+1 F[Tr-K+1, IT], [ar- K+ [at-k+1,-.
4.1 StateSpaceDifftsser Architecture
Our architecture is shown in Fig. It is conceptually divided into two branches: Long-Context Branch and Generative Branch. The Long-Contex Branch preserves information over long sequences, and the Generative Branch uses this context to render high-quality images.
5

---

Long-Context Branch In contrast to transformer and CNN architectures, state-space models (SSMs) are designed specifically to efficiently process long sequences. Although SSMs are generally designed for continuous input signals, discrete SSMs maintain an internal state representation h. that updated with each time step in an input sequence of mieddimesionnl feature vectors f1,-.., fT., through the parameter matrices A, B and C, which are learned in training time:
he Ahh-1 + Bfc, m, Che
We denote a state-space model with Mi,... MrT M(fi... fr), with my denoting the model's output. To bring it into the world model setting, we define fa to be compact feature representation of 1 and train a model that predicts future observations: f2.....r+ fr+1 is common in existing work to apply SSMs at the patch or image token level as common in previous work 42 26]. Instead, avoid conflating spatial and temporal dependencies by temporally processing full frames. Each frame is encoded into a single compact feature ft used as a single step of our sequence. The encoding is obtained by the continuous Cosmos tokenizer [23) with scale 16 (CI16). The resulting patch tokens (dimension 16) are flattened to form the single feature vector ft per image. Alongside it, we incorporate a discrete action dt, which indexes learnable embedding of dimension 16. We concatenate fi and a to create the input at each step. We adopt the Mamba architecture due to its dynamic selection mechanism and efficient parallelism, which led to superior performance compared to other SSM variants.
M [fr, arl).
One key benefit of state-space models is their computational efficiency. As only a single state is maintained in inference, memory remains constant regardless of context length. As the same update is applied linearly a growing sequence, Mamba only updates the state from the previous step, making for a constant per-step latency. In contrast, CNNs and transformers do not maintain a state and have to reprocess the growing sequence at every step. In training, Mamba further parallelizes the sequence processing.
sequence, computational complexity remains O(T). When presented with
Provided a long input sequence of high- resolutio images and corresponding actions, the model predicts the Cosmos features corresponding to the next observation with an MSE loss. In those features, we expect relevant to the time-step information, recalled from the sequence. While context cues are to be preserved, the state is oow-dimensional, and this model is not generative. Therefore, the generative branch, designed for the generation of high-quality images, is intended to render the final output.
Generative Branch. To generate high-quality images in complex environments, we employ diffusion model. Our choice is the DIAMOND world model M. a UNet-based EDM diffusion model 488 designed for visual prediction in sequential environments. Therefore, our Generative Branch conditions on only four low-resolution frames and their corresponding actions, represented as 51 512-dimensional action embeddings. Despite this minimal context, it can produce high-quality predictions with just three denoising iterations per output frame. The architecture consists of two diffusion models: a primary model that predicts the next observation at low resolution, and secondary upsampler that refines these predictions to a resolution of 280 x 150. As the model predicts one frame a time, generating longer sequences is achieved through a sliding-window strategy, and each newly generated frame is appended to the input history for the next prediction. In isolation, this strategy causes short-context limitation to this model.
Fusion of features. To address the context limitation of the Generative Branch, we build a fusion module to integrate state-space features into it, in order to provide long-context information. To that end, we process the entire sequence with the Long-Context branch and obtain the last 4 output features fc. We fuse those features with the corresponding action embeddings from the Generative Branch. These features are first normalized and then passed through a two-layer MLP with SiLU activation. where the input size matches the feature dimensions. Similarly, the action embeddings, perturbed with noise, are processed by an MLP with the same architecture. To form the final conditioning vector, we concatenate the outputs of the two MLPss Empirically, we discovered that processing the memory and action conditions independently before concatenation yielded better performance than fusing them earlier in the pipeline.

---

4.2 Training Protocol
At each step, batch consists of sequences of actions, reversed actions and observations. Training is performed in two stages. Firstly, the Long- Context Branch is trained on long context length 50 or 16. The produced features decode to images with artifacts, but with important context cues. Then, freezing the Loog-CConter Branch, we train the Generative Branch, conditioned on the compressed long-context features, with sequence size of 4. This branch produces the final high-quality images with the correct context. Training details are given in App. A
We found that this two-stage training is crucial for stability. Direct end-to-end training is unstable, as diffusion gives noisy gradients to the SSM, and the SSM gives constantly changing features diffusion. In turn, diffusion learns ignore the SSM features. Therefore, stable features of pretrained SSM worked best in this architecture. Moreover, the training separation enabled to swap out in test time the Long-Contes Branch with another independently trained model, without having to further fine-tune the heavier Generative Branch.
5 Experiments
5.1 Experimental Setup
Baselines. We establish two baselines. The first is a pure diffusion model without state-space features: the DIAMOND model. Our second baseline is the State-Space World Model. It is the Long-Context branch of StateSpaceDiffuser and its training is equivalent to the first stage of training, as described in Sect. 4.2 At inference time, the predicted feature ft is decoded into an image I using the decoder from the Cosmos tokenizer. In App. [B.2 we present comparisons of sequence models to solidify our choice of Mamba as our backbone.
This model enables us to assess the memory capacity of state-space models (SSMs) in sequential visual prediction. Although its outputs tend to be blurry and contain artifacts in complex scenes, due to the absence of variational component and limited generative expressiveness compared to modern diffusion models, the SSM exhibits a strong ability to model long sequences and retain information from earlier in the trajectory. The strengths and shortcomings observed in this baseline directly inform and motivate the design of StateSpaceDiffuser.
Testing Protocol. Our evaluation protocol matches our mirrored action setup- we take actions and n reverse actions, and expect to generate the same observations for the second half of the sequence as seen in the first. On MiniGrid, we have a fixed sizeable visual difference per step, while for CSGO continuous motion often results in small per-step changes. Therefore, MiniGrid, we generate one frame in the future at time, while in CSGO we sequentially generate the whole second half of the sequence. On MiniGrid we evaluate with PSNR and SSIM on varying future horizons- the further in the sequence, the longer the memory required. In CSGO we perform a user study, more aligned to the visual complexity of the environment. We motivate this diffcrence with the known mismatch between perceived quality and fidelity metrics in continuous video [65] 67 30 (App. D.3). Although the baseline performs well when context is not essential, our protocol exposes its inability to model long-term context, resulting in degraded quality in this scenario.
5.2 Results and Analysis
Simple MiniGrid Evaluation. In this experiment, test the recall ability of the baseline, the State-Space World Model and StateSpace on a simple toy setup, as described in Sect. 3 We train and test on the same set of 34 samples. The goal is recall a color the final frame from the first frame in the sequence with a length of 7 frames. Two random samples (colors) from the results are shown in Fig. 4 with the corresponding model predictions. With input size 4, the baseline processes the sequence in a sliding window fashion and, as within the 3 steps the color information is lost, cannot reconstruct the correct color. Despite the small training set size, the baseline fails because of lack of long congecttext abilities. In contrast, our State-Space World Model, based on a computationally efficient state-space model, is able to predict the correct color. Finally, is demonstrated that our StateSpaceDiiffsser is also able to recall the correct content by effectively combining both paradigms. Notably, our methods perform equivalently on context length of 50 frames when predicting the 51st, StateSpaceDifftssr recalls the color from 50 steps ago.
7

---

StateSpace
Ground Truth
World Model
Diffuser
Figure 4: Long-Context Simple Demonstration. An agent starts moving to the right covering the color for the next T/2 frames, then the agent moves back the same amounts of steps. Diffusion baseline fails to recover the color, StateSpaceDiffuser and State-Space World Model successfully recover it through the state-space representation. Both T=7 and T=50 generates such results.
1.00
Model
Avg. PSNRI
Fin. PSNR SSIM
Conrext Length
0.50
DIAMOND
27.
25.44
0.95
Baseline.
State-Space Model
33.40
33.17
DtifeSpaceifffueer (Ours, w/o state)
23.68
20.95
0.92
0.00
StateSpaceDifftuerr (Ours)
41.01
40.55
0.98
Context Lengrh
0.50
DIAMOND
26.13
25.15
0.95
State-Space World Model
32.64
32.44
0.96
StateSpaceDiffuser (Ours)
39.68
39.32
-1.00
Horizon
Table 1: MiniGrid Quantitative Evaluation of Long-Context Awareness, Our StateSpaceDiffuser outperforms the baselines.
Figure 5: CSGO User study
results.
Forward-Backwand Evaluation MiniGrid. We compare the long context abilities of our diffusion (DIAMOND) and state-space (State-Space World Model) baselines our MiniGrid test set. We evaluate our models trained on context length 50 on context lengths 16 and 50 (demonstrating generalizability). We follow the protocol outlined in Sect. 5.I To evaluate, we compute the Peak Signal-to-Noise Ratio (PSNR) for each predicted frame in the reverse trajectory, reporting both the mean score and the PSNR at the final time step, which requires the longest-term memory. As shown in the Tab. 1 model significantly outperforms both baselines, particularly the end of the sequence, where successful recall the first frame critical. This highlights the model's ability to retain and reinstantiate long-ter visual context. In App. [B.4 B.5] show the stability and robustness of these results, in App. B.I performance gain analysis over computational cost. Compared to the State-Space World Model, our method achieves higher fidelity output, benefiting from the superior
generative capacity of the Generative Branch (examples in App. C.I C.2). Fig. 6 (b) presents example rollouts generated by model and the diffusion-only baseline. In MiniGrid, predictions are made one step at a time using the ground truth sequence. As a result, most content is carried over from the previous frame, with only the newly revealed area requiring inference. Our method excels at filling in these newly revealed regions, even when the relevant context originates far back in the sequence, In contrast, the diffusion baseline struggles recover such long long-range dependencies.
Recall Across a Context Length. In this experiment study the accuracy of our models over the varying context length of the cowward-backwe evaluation on MiniGrid. When predicting future observations, the last frame's content depends on the first frame's 's content, and the further back we go in the sequence the smaller the context length required for a good reconstruction. This is direct consequence of the mirror style the observations in our setup. In Fig. we show the PSNR each predicted time step. The first few predicted frames are easily predicted by all models as the solution the diffusion baseline quickly falls
falls within the short input window. However, performance as no form of information is preserved from the long context, while state-space model is able to harvest this information. Our Staffupace model gets the best of both worlds long-context awareness and high-fidelity predicted images, and performs the best.

---

S+22
5+24
(b) MiniGrid
Figure 6: Qualitative comparison of StateSpaceDiffuser and the baseline (DIAMOND). Top row: input frames and actions. Botton rows: generated frames under reversed actions. In CSGO last frames autoregressive prediction, in MiniGrid next frame prediction.
60
Midel
Bostaa
55-
Morld Moddd
50
SuoeSpanntttteer
45-
40
35
30-
25
26
20
15
10
12
14
16
17
15
30
35
40
65
50
Context Length
Context Length
Context Length 17
(b) Context Length 51
Figure 7: Recall performance on MiniGrid for two context lengths.
Input: Frames
Prediction
Ground Truth
StateSpaceDif fuser
State-Space World Model
Statespaccaitfuser (na state)
Figure 8: CSGO Ablations. We show that without the state-space features from the Long Context Branch, StateSpace loses its context preservation ability. We also show that while State-Spact World Model demonstrates long-context memory, the produced images are of subpar visual quality.
Borward-Backward Imagination Evaluation CSGO. Similarly MiniGrid, evaluated the recall abilities our model on the CSGO dataset, a visually complex environment a 3D world. In CSGO most actions are gradually executed over a sequence, and there is compounding effect on content change (e.g. jump unfolds over many frames). For a high impact evaluation we decide to give only the first half the sequence and continuously produce the second half (reverse) by feeding generated frames. As actions are motions at varying levels, the final frame may contain the correct content memorized but with low PSNR, as the camera position and scene geometry might be slightly shifted. Therefore, instead of fidelity metrics we perform user study where the 12 participants judge whether images produced by StateSpace
are closer content to the ground truth compared to the diffusion baseline (details in App. D.4]. Our rating is in the range [-1, 1]. with 0 being borderline, preference toward the baseline, preference for StateSpace The results shown in Fig. demonstrate clear preference of the users for StateSpaceDiffusee over the bascline

---

Model
Avg. PSNR
Fin. PSNR
SSIM
Model
Avg. PSNR
Fin. PSNR
SSIM
Context Lengch 100
Contert Lengti 1.50
DIAMOND
26.39
26.24
0.95
DIAMOND
24.35
24.20
0.94
State-ppcce World Mode)
31.65
30.89
State-Space World Model
0.94
StuleSpaceDiflueer (Ours)
37.99
35.87
0.98
StateSpaceDifluser (Ours)
30.75
28.93
0.96
Table 2: Generalization to Longer Context. Our model, trained on context length 50, generalizes to longer sequences (context 100 and 150).
for both prediction in the 15th frame (rating 0.20) and 17th (last) frame (rating 0.24). Fig. 6 (a) shows sample of CSGO imagination in different time steps, demonstrating that while the baseline fails to recall the correct content, the StateSpaceDiffuser correctly produces the details. (More App. D.I
State Features Ablation. We study the utilityy of the state-space features provided to the Generative Branch our StateSpaceDiffuser model. We take a trained model and perform a MiniGrid evaluation by replacing the output features of the Long-Context Branch with zeros before passing them to the Generative Branch. In Tab. T we show that this causes the performance to quickly drop even below baseline performance, clearly demonstrating that the features we demonstrate the same effect on CSGO. Without state features, the model hallucinates; without diffusion, the state-space model remembers but produces poor visual quality. (More in App. B.7) B.8
highly utilized. In Fig.
Generalization to Longer Context. In this experiment we show that StateSpaceDiffuser operates on much longer contexts without finetuning. We evaluate our model trained on context length 50 on lengths 100 and 150 using new MiniGrid test with longer sequences. Tab. shows that Dtafffse successfully generalizes to longer context, keeping a significant gain over the baselines. Analogously, in App. B.3 we show generalization from context length 16 length 50.
5.3 Strengths, Limitations and Scalability
Apart from the already established generalization across context length, via extra experiments, we find that StateSpaceDiffuser is able generalize across visual complexity (App. B.6) and can recover from strong motion artifacts (App. D.2). Our model can recover from input noise in future steps, but is clearly affected by it the current steps (App. B.5). Our lightweight StateSpaceDiffuuee was trained under a fixed compute budget. The lightweight diffusion decoder (no large pretrained backbone) can yield visual artifacts in long rollouts. Replacing the decoder with a better, larger one, can improve visual sharpness without changing the method. Our lightweight single-layer Long- Context Branch compresses the context into diw-dimeniional state (256), which can cause loss of detail in extended rollouts, especially in complex environments (App. [.2)). Scaling the SSM is expected to reduce nigh-feeuency decay over time. The separation in training enables separately scaling each branch before combining them.
(state watpss///hheadaaaaaeeeee/aaaaaae
6 Conclusion
We introduced StateSpaceDiffuser, a hybrid model that combines state-spact representations with diffusion to enable long-hore horizon visual world modeling. By decoupling global context modeling (via a state-space backbone) from high-fidelity synthesis (via diffusion), our model retains global context over many steps essentially additional computational cost. The resulting representation alleviates the drift and inconsistency that plague conventional diffusion-only systems in long sequences.
Experiments MiniGrid and CSGO validate method's consistency and fidelity across long sequences. the PSNR by 51.9% over the diffusion baseline and achieves final-frame PSNR of 39.32 versus 25.14 for DIAMOND on a long context length of 50 frames. Human raters also favor our generations for long-context consistency (Fig. 55.
protocol with horizon 50, StateSpaceDiffuser improves average
Our results establish state-space diffusion as a scalable and consistent solution for long-context visual generation. We believe that bridging state-space reasoning with diffusion generation is promising direction for robust, long-horizon world modeling, and we hope this work lays a solid foundation for future research in temporally coherent visual prediction.
10

---

7 Acknowledgments
This research was partially funded by the Ministry of Education and Science of Bulgaria (support for INSAIT, part of the Bulgarian National Roadmap for Research Infrastructure). This project was supported with computational resources provided by Google Cloud Platform (GCP).
References
[1] Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos Storkey, Tim Pearce, and François Fleuret. Diffusion for world modeling: Visual details matter in atari. Advances Neural Information Processing Systems, 37:58757-58791, 2025.
[2] Amir Bar, Gaoyue Zhou, Danny Tran, Trevor Darrell, and Yann LeCun. Navigation world models. arXiv preprint arXi::2412.03577 2024.
[3] Omer Bar-Tal, Hila Chefer, Omer Tov, Charles Herrmann, Roni Paiss, Shiran Zada, Ariel Ephrat. Junhwa Hur, Guanghui Liu, Amit Raj, et al. Lumiere: A space-time diffusion model for video generation. In SIGGRAPH Asia 2024 Conference Papers, pages 2024.
2024.
[4] Iz Beltagy, Matthew Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.051 2020.
[5] James Bradbury. Stephen Merity, Caiming Xiong, and Richard Socher. Quasi-recurrent neural networks. In International Conference on Learning Representations, 2017.
[6] Tim Brooks, Bill Peebles, Connor Holmes, Will DePue, Yufei Guo, Li Jing, David Schnurr, Joe Taylor, Troy Luhman, Eric Luhman, Clarence Ng. Ricky Wang, and Aditya Ramesh. Video generation models as world simulators. 2024.
Tom Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Ncelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot leamers.. Advances in neural information processing systems, 33:1 19771999,,,, 2020.
[8] Jake Bruce, Michael D Dennis, Ashley Edwards. Jack Parker-Holder, Yuge Shi, Edward Hughes, Matthew Lai, Aditi Mavalankar, Richie Steigerwald, Chris Apps, al. Genie: Generative interactive environments. In Forty-first International Conference on Machine Learning, 2024.
[9] Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF international conference computer vision, pages 9650-9660, 2021.
[10] Boyuan Chen, Diego Martí Monsó, Yilun Du, Max Simchowitz, Russ Tedrake, and Vincent Sitzmann. Diffusion forcing: Next-token prediction meets tull-sequence diffusion. Advances in Neural Information Processing Systems,
1225 2024.
[11] Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. Transdreamer: Reinforcement learning with transformer world models. arXiv preprint arXiv:2202.09M0 2022.
[12] Maxime Ahevalier-Boisver Bolun Dai, Mark Towers, Rodrigo de Lazcano, Lucas Willems, Salem Lahlou, Suman Pal, Pablo Samuel Castro, and Jordan Terry. Minigrid & miniworld: Modular & customizable reinforcement lcarning environments for goal- goal-oriented tasks. CoRR, abs/2306. 13831, 2023.
[13] Silvia Chiappa. Sébastien Racaniere, Daan Wierstra, and Shakir Mohamed. Recurrent environment simulators. arXiv preprint arXiv:1004.02254, 2017.
[14] Rewon Child, Scott Gray, Alee Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019.
[15] Kyunghyun Cho. B van Merrienboer, Caglar Gulcehre, F Bougares, I Schwenk, and Yoshua Bengio. Learn- ing phrase representations using encoder-decoder for statistical machine translation. In Conference on Empirical Methods in Natural Language Processing (EMNLP 2014), 2014.
[16] Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song. Andreea Gane, Tamas Sarlos, Peter Hawkins. Jared Quiney Davis, Afroz Mohiuddin, Lukasz Kaiser. et al. Rethinking attention with performers. In International Conference Learning Representations, 2020.
11

---

[17] Junyoung Chung, Caglar Gulcehre, Kyunghyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. In NIPS 2014 Workshop on Deep Learning, December 2014, 2014.
[18] Tri Dao, Daniel Y Fu, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. In Proceedings of the 11th International Conference on Learning Representations (ICLR), 2023.
[19] Etched Decart. Oasis world model. https:/ htt.//ossssmmmll..e. github. Accessed: October 31, 2024.
[20] Fei Deng, Junyeong Park, and Sungjin Ahn. Facing off world model backbones: Rnns, transformers, and
s4. Advances in Neural Information Processing Systems, 36, 2024.
[21] Jacob Devlin, Ming- Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidi- rectional transformers for language understanding. In Proceedings of the 2019 conference of the North American chapter of the association for computaaional linguistics: human language technologies, volume flong and short papers), pages 41 4171-4186 2019.
[22] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, G Heigold, S Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2020.
[23] NVIDIA et. al. Cosmos world foundation model platform for physical ai. arXiv preprint 2rX:::500.03355. 2025.
[24] Haoqi Fan, Bo Xiong. Karttikeya Mangalam, Yanghao Li, Zhicheng Yan, Jitendra Malik, and Christoph Feichtenhofer. Multiscale vision transformers. In 2021 IEEE/CVF International Conference on Computer Vision (ICCV), pages 6814-6615, JEEE, 2021.
[25] Zhengcong Fei. Mingyuan Fan, Changgian Yu, and Junshi Huang. Scalable diffusion models with state space backbone. arXiv preprint 246600000.0 2024.
[26] Zhengcong Fei, Mingyuan Fan, Changqian Yu, Debang Li, Youqiang Zhang, and Junshi Huang. Dimba: Transformer-matte diffusion models, 2024.
[27] Ruili Feng. Han Zhang, Zhantao Yang. Jie Xiao, Zhilei Shu, Zhiheng Liu, Andy Zheng, Yukun Huang. Yu Liu, and Hongyang Zhang. The matrix: Infinite-horizon world generation with real-time moving control. arXiv preprint 24.0035688, 2024.
[28] Yuqian Fu, Runze Wang, Yanwei Fu, Danda Pani Paudel, Xuanjing Huang, and Luc Van Gool. Objectrelator: Enabling cross-view object relation understanding in ego-centric and exo-centric videos. ICCV, 2025.
[29] Shenyuan Gao, Jiazhi Yang, Li Chen, Kashyap Chitta, Yihang Qiu, Andreas Geiger, Jun Zhang, and Hongyang Li. Vista: A generalizable driving world model with high fidelity and versatile controllability. arXiv preprint arXiv:2405. 17398, 2024.
[30] Rohit Girdhar, Mannat Singh, Andrew Brown, Quentin Duval, Samanch Azadi, Sai Saketh Rambhatla, Akbar Shah, Xi Yin, Devi Parikh, and Ishan Misra. Emu video: Factorizing text-to-video generation by explicit image conditioning, arXiv preprint arXiv:2311.10709 2023.
[31] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In First Conference on Language Modeling, 2024.
[32] Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. Advances in neural information processing systems, 33:1474-1493 2020.
[33] Albert Gu Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.0033 2021.
[34] Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré. On the parameterization and initialization of diagonal state space models, 2022.
[35] Yuwei Guo, Ceyuan Yang, Ziyan Yang, Zhibei Ma. Zhijie Lin, Zhenheng Yang. Dahua Lin, and Lu Jiang. Long context tuning for video generation. arXiv preprint urXiv:2503. 10589, 2025.
[36] David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Garnett, editors. Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018
12

---

[37] Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Represeniations, 2020.
[38] Danijar Hafher, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International conference on machine learning, pages 2555-2565. PMI.R, 2019.
[39] Danijar Hafner, Timothy Lillicrap. Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference on Learning Representations, 2021.
[40] Danijar Hafner. Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301. 04104, 2023.
[41] Mariam Hassan, Sebastian Stapf, Ahmad Rahimi, Pedro Rezende, Yasaman Haghighi, David Brüggemann, Isinsu Katircioglu, Lin Zhang, Xiaoran Chen, Suman Saha, et al. Gem: A generalizable ego-vision multimodal world model for fine-grained ego-motion, object dynamics, and scene composition control. arXiv preprint arXiv:2412.1 1119, 198, 2024.
[42] Ali Hatamizadeh and Jan Kautz. Mambavision: A hybrid caaba-transformer vision backbone, 2024.
[43] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pages 16000-16009, 2022.
[44] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models, 2020.
[45] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory, Neural computation, 988):1735 153911575-780 1997.
[46] Anthony Hu, Lloyd Russell, Hudson Yeo, Zak Murez, George Fedoseev, Alex Kendall, Jamie Shotton, and Gianluca Corrado. Gaia-1: A generative world model for autonomous driving. arXiv preprint arXiv:2309.15 17080. 2023.
[47] Lukasz Kaiser, Mohammad Babacizadeh, Piotr Milos, Blažej Osiński, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model based reinforcement learning for atari. In 8th International Conference on Learning Representations, ICLR 2020, 2020.
[48] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. Advances in neural information processing systems, 35:26565-26577, 2022.
[49] Angelos Katharopoulos, Apoorv Vyas. Nikolaos Pappas, and François Fleuret. Transformers are rans: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5156-5165. PMLR, 2020.
[50] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. International Conference on Learning Representations, 2020.
[51] Soonwoo Kwon, Jin- Young Kim, Hyojun Go. and Kyungjune Back. Toward stable world models: Measuring and addressing world instability in generative environments. arXiv preprint arXiv:2503.08122. 2025.
[52] Sangyoun Lec, Juho Jung, Changdae Oh, and Sunghee Yun. Enhancing temporal action localization: Advanced s6 modeling with recurrent mechanism. arXiv preprint arX 2024.
[53] Tao Lei, Yu Zhang, Sida I Wang, Hui Dai, and Yoav Arzzi Simple recurrent units for highly parallelizable recurrence. In Proceedings the 2018 Conference on Empirical Methods Natural Language Processing. Association for Computational Linguistics, 2018
[54] Zhixuan Lin, Evgenii Nikishin, Xu Owen He, and Aaron Courville, Forgetting transformer: Softmax attention with a forget gate. arXiv preprint arXiv:2503.00 0300 2025.
[55] Xinqi Liu, Li Zhou, Zikun Zhou, Jianqiu Chen, and Zhenyu He. Mambavlt: Time-evolving multimodal state space model for vision-language tracking. arXiv preprint arXiv:2411. 2024.
[56] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference computer vision, pages 10012-10022, 2021.
13

---

[57] Taiming Lu, Tiammin Shu, Alan Yuille, Daniel Khashabi, and Jieneng Chen. Generative world explorer. arXiv preprint vrXiv:2411.11444 2024.
[58]
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world models. In The Eleventh International Conference on Learning Representations, 2023.
[59]
Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrett neural networks. In International conference on machine learning, pages
Pmlr, 2013.
[60] Tim Pearce and Jun Zhu. Counter-strike deathmatch with large-scale behavioural cloning. In 2022 IEEE Conference on Games (CoG), pages 104-111. IEEE, 2022.
[61] Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Leon Derczynski, et al. Rwkv: Reinventing rans for the transformer era. Findings of the Association for Computational Linguistics: EMNLP 2023, 2023.
[62] Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. In International Conference on Machine Learning, pages 28043 28078. PMLR, 2023.
[63] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lec, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21(140):1-66 167,, 2020.
[64] Jan Robine, Mare Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer-based world models happy with 100k interactions. The Eleventh International Conference on Learning Representations. 2023.
[65] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10684-10695, 2022.
[66] Lloyd Russell, Anthony Hu, Lorenzo Bertoni, George Fedoscev, Jamie Shotton, Elahe Arani, and Gianluca Corrado. Gaia-2: A controllable multi-view generative world model for autonomous driving. arXiv preprint wwXi::2503.20522 2025.
[67] Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David J Flect, and Mohammad Norouzi. Image super-resolution via iterative refinement. IEEE transactions puttern analysis and machine intelligence, 45(4)::7713-422 2022.
[68] Nedko Savov, Naser Kazemi, Mohammad Mahdi, Danda Pani Paudel, Xi Wang, and Lue Van Gool. Exploration-drive generative interactive environments. In Proceedings of the Computer Vision and Puttern Recognition Conference, pages 27597-27607 2025.
[69] Mattia Segu, Luigi Piccinelli, Siyuan Li, Yung- Hsu Yang, Bernt Schiele, and Lue Van Gool. Samba: Synchronized set-of--equences modeling for multiple object tracking. arXiv preprint arXiv:2410.01866 2024.
[70] Ramanan Sekar, Oleh Rybkin, Kostas Daniildiis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via self-supervised world models. In International conference on machine learning. pages 8583-8592. PMLR, 2020.
[71] Uriel Singer, Adam Polyak, Thomas Hayes, Xi Yin. Jie An, Songyang Zhang, Qiyuan Hu, Harry Yang, Oron Ashual, Oran Gafni, et al. Make-a-video: Text-to-video generation without text-video data. In The Eleventh International Conference on Learning Representations.
[72] Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling. In ICLR, 2023.
[73] Jascha -Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilirium thermodynamics. In International conference on machine learning, pages 2256-2265. pmlr, 2015.
[74] Kiwhan Song, Boyuan Chen, Max Simchowitz, Yilun Du, Russ Tedrake, and Vincent Sitzmann. History- guided video diffusion. arXiv preprint arXiv:2502.0676, 06764, 2025.
[75] Yang Song, Jascha Sohl-Dickstein Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
14

---

[76] Yutao Sun. Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang. and Furu Wei. Retentive network: A successor transformer for large language models. arXiv preprini arXiv:2307.08621, 2023.
[77] Dani Valevski, Yaniv Leviathan, Moab Arar, and Shlomi Fruchter. Diffusion models are real-time game engines, arXiv preprint arXiv:2408. 2024.
[78] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention you need. Advances neural information processing systems, 30, 2017.
[79] Philipp Wu. Alejandro Escontrela. Danijar Hafner, Pieter Abbeel, and Ken Goldberg, Daydreamer: World models for physical robot learning, In Conference on robot learning, pages 2226-2240. PMLR, 2023.
[80] Jiannan Xiang. Guangyi Liu, Yi Gu, Qiyue Gao, Yuting Ning, Yuheng Zha, Zeyu Feng. Tianhua Tao, Shibo Hao, Yemin Shi, et al. Pandora: Towards general world model with natural language actions and video states. arXiv preprint 0945555406. 2024.
[81] Zeqi Xiao, Yushi Lan, Yifan Zhou, Wenqi Ouyang, Shuai Yang, Yanhong Zeng and Xingang Pan. Worldmem: Long-term consistent world simulation with memory. arXiv preprint 1533:25544422699 2025.
[82] Mengjiao Yang. Yilun Du, Kamyar Ghasemipour, Jonathan Tompson. Dale Schuurmans, and Pieter Abbeel. Learning interactive real-world simulators. arXiv preprint arXiv:2310.061 2023.
[83] Sherry Yang. Jacob Walker, Jack Parker-Holder, Yilun Du, Jake Bruce, Andre Barreto, Pieter Abbeel, and Dale Schuurmans. Video as the new language for real-world decision making, arXiv preprint arXiv:2402.17139, 2024.
[84] Jiwen Yu, Yiran Qin, Xintao Wang, Pengfei Wan, Di Zhang, and Xihui Liu. Gamefactory: Creating new games with generative interactive vidcos. arXiv preprint 2501:2501100822. 2025.
[85] Manzil Zaheer. Guru Guruganesh. Kumar Avinava Dubey. Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang. Li Yang, al. Big bird: Transformers for longer sequences. 2020.
Advances neural information processing systems, 33:17283-1723
[86] Weipu Zhang, Gang Wang, Jian Sun, Yetian Yuan, and Gao Huang. Storm: Efficient stochastic transformer based world models for reinforcement learning. Advances in Neural Information Processing Systems, 36:27147-22 166,, 2023.
[87] Gaoyue Zhou, Hengkai Pan, Yann LeCun, Lerrel Pinto. Dino-wm: World models on pre-trained visual features enable zero-shot planning, 2025.
15

---

Appendix
Training Details
16
StateSpaceDiffuseer Propertri
17
B.
Performance/Cost Tradeoff
17
B.2
Long-Conterd Architecture Comparison
17
Generalization Across Context Length
18
B .4
Stability Across Seeds
18
B.5
Robustness Noise
18
B.6
Generalization Across Visual Complexity
19
7
State Features Ablation
20
B.8
Comparison with State-Space World Model
20
MiniGrid Evaluations
21
C.1
MiniGrid Qualitative Evaluation
Imagination Qualitative Results
CSGO Evaluations
24
D.I
CSGO Qualitative Evaluation
24
D.2
Strengths and Limitations of StateSpace
25
D.3
Quantitative vs. Perceptual Consistency
26
D.4 User Study Details
27
A Training Details
For MiniGrid, DIAMOND baseline takes 30x30 images and upscales them 144x14; 144, for CSGO 30x56, upscaled to 150x280. The Long-Contex. Branch contains the Cosmos tokenizer, which decodes in powers of two. Therefore, The Long Context Branch takes 32x32 for MiniGrid and 144x272 for CSGO. In training, the downscaling is done from the high-resolution image with bicubic interpolation DIAMOND).
In our Long-Conter Branch, the Cosmos Tokenizer tokens are flattened to produce features of size 1296 (MiniGrid) 2448 (CSGO). Action dimensions in the Long Context Branch (and the State Space World Model) is 16 they are concatenated to the visual features. We use a single Mamba layer with state size 256, the input is expanded 4 times with MLP inside the Mamba layer, the internal convolution dimensions are 4.
We use 8 A100 GPUs for all models except for the MiniGrid State Space World Model, which was trained on 4 A100 GPUs for MiniGrid models. All models use the Adam optimizer. The State Space World Model is trained with learning rate of 5e and batch size 136 (MiniGrid) or batch size 272 (CSGO). Both the MiniGrid and CSGO models are trained for 70k iterations on sequence size 16. StateSpaceDiffuser includes 600M parameters, and is trained with a learning rate le weight decay le -2 grad norm clip 10 and batch size 64. The MiniGrid model is trained for 77k iterations, CSGO 220k iterations. For upscaling in MiniGrid training the denoiser for predicting rew-resouttion next frame. In CSGO we achieved our best results with the upsampler, part of the weights originally provided by DIAMOND.
train the sampler for 27k iterations after
For all models, we loaded the weights of a pre-trained State-Space World model into StateSpaceDiff- fuser. For MiniGrid, while models trained sequence length 16 performed well, in inference,
16

---

Backbone
Avg. PSNR?
Fin. PSNR1
SSIM
Context Length
DIAMOND
27.13
25.44
0.95
Mamba [31
34.41
37.81
0.97
Length
DIAMOND
26.13
25.15
0.95
Mamba 31
27.62
29.75
0.94
Table 3: Performance Gains, Normalized By Computational Cost. Strong gains still observed,
suggesting that the gains are much higher than the cost.
Backbone
Avg. PSNRA
Fin. PSNR1
SSIMf
Content Length
LSTM
26.80
26.91
0.93
27.40
26.68
0.93
Mamba
32.64
32.44
0.96
Table 4: Evaluating Sequence Models. The superior performance of Mamba leads to our choice of an SSM in our Long-Context Branch.
Backbone
Arg. PSNR
Fin. PSNR
SSIMT
Context Length 16
S4 [3]
24.29
24.31 27.87
0.91
Mamba
27.09
0.94
Table 5: Evaluating State-Space Backbones. The superior performance of Mamba leads to our choice to use this method as our Long-Context backbone.
achieved our best results for both context length 16 and 50 by using the model trained on context length 50, For CSGO, we used the State-Space World Model weights for context length and also evaluated on context length 50.
When we evaluate DIAMOND and StateSpaceDiffuser, we use denoising steps denoise the next observation and 10 to upscale it.
B StateSpaceDifftsee Properties
B.1 Performance/Cost Tradeoff
Our Long-Context Branch adds very little computation to the diffusion model, and in exchange it offers significant improvements in consistent generation. In inference with batch size of 1, measure DIAMOND to require 909.515 GFLOPS (4 input frames). The Long Context Branch with context length 16, requires only 5.5 GFLOPS for context length 16 and 16.741 GFLOPS for context length 50. That is only 0.6% of all inference computations of the full model, for sequence size 16, and 1.8% for sequence size 50.
We show that the gains from the Long-Context Branch surpass its computational cost by a large margin. To account for the cost in StateSpaceDiffuser scores, normalize them by multiplying by 0.006 for sequence size 16 and - 0.018 for sequence size 50. The normalized scores are reported on Tab. They confirm that there is still significant gain in performance despite the normalization.
B.2 Long- Long-Context Architecture Comparison
We compare the choice of the SSM in State-Space World Model with other popular sequence processing models LSTM and GRU. We train on context size 50 with our MiniGrid setup. We add linear layer on the input and output (dim 256, with ReLU activation) of the models. Tab, 4hoows that Mamba outperforms the other sequence models in long-range temporal dependencies, while remaining computationally efficient. This confirms the conclusion the original paper [31].
In addition, consider the choice of an SSM model itself. We consider the S4WM model the closest in spirit model in literature. However, as their model code is not available, comparing an S4 backbone to choice of Mamba serves the closest
get to comparison. We train State-Space
17

---

Model
Avg. PSNRT
Fin. PSNRT
SSIMT
Context Length
DIAMOND
27.13
25.44
0.95
State-Space World Model
29.71
31.34
0.96
StateSpaceDiffuser (Ours)
38.04
Length
DIAMOND
26.13
25.15
0.95
State-Space World Model
27.25
24.49
StateSpaceDiffusee (Ours)
28.12
30.30
0.96
Table 6: MiniGrid Quantitative Evaluation of Long Context Awareness, Trained on Context
Length 16. Our models generalize their performance from smaller longer sequence.
Noise
28
31
34
37
SSM
15.22
16.62
15.36
15.18
14.85
16.55
23.44
26,60
28.42
29.35
29.75
29.75
30.00
Full
15.50
16.90
15.49
15,30
14.94
16.65
23.28
26.51
28.35
29.33
29.75
29.95
29.85
Noise
40
41
42
43
44
45
46
47
48
49
50
51
SSM
29.93
30.06
29.99
30.12
30.14
30.45
30.11
30.18
30.19
30.08
30.19
30.19
Full
30.09
29.88
30.11
30.08
30.41
30.12
30.13
30.17
30.12
30.20
30.13
30.15
Table 7: Noise robustness across steps. Our method is affected on the noisy frames but quickly recovers in further steps. SSM denotes the noise added only to the Long-Context Branch; Full denotes the noise added on the Generative Branch as well.
World Model with S4 and Mamba MiniGrid using context length of 16 and comparing their performance. We keep the context small to account for the larger amount of training iterations usually needed by SSMs. As average. As Mamba introduces dependency on the input of the state update, this clearly benefits its ability perform in long context.
Tab. [ Mamba outperforms S4, achieving 9.1 PSNR improvement
B.3 Generalization Across Context Length
In Tab. 6. we show the evaluation results of models trained in context size 16, on both context length 16 and 50. A reasonably good performance demonstrates that the models do not overfit particular sequence size and still perform well a context length longer than it has been trained with. While the State World Model exhibits uncertainty in its predictions for longer context without training (blurriness, color deviations), its features prove useful for StateSpaceDiffuss with the Generative Branch producing fighhe-fidelity result. StateSpaceDiffuser noticeably outperforms the baseline context length 50.
B.4 Stability Across Seeds
To show the stability of our results, we perform evaluation of our MiniGrid model, on context length 16, under different seeds. We obtain 41.00 + 0.008 Avg. PSNR, 40.52 + 0.019 Fin. PSNR, 0.98 + 0.0004 SSIM. We also perform evaluation over seeds of larger-scale, more expensive evaluation our new 100 context length experiment from Sect. B.3 This results in: 37.99 0.002 Avg. PSNR, 35.88 = 0.029 Fin. PSNR. 0.98 = 0.000004 SSIM. In both cases, the obtained metrics are extremely stable and consistent between seeds, with a low standard deviation.
B.5 Robustness to Noise
We test the robustness of our method by adding noise in the middle section of the rollouts that serve as context. specific, we consider context length 50 in MiniGrid and add Gaussian noise (std 2.5) to the 11 frames in the middle. We consider two cases: 1) adding noise to the SSM input only; 2) adding noise both to the SSM input and the diffusion model input. Results are shown on Tab. 7 at different steps of prediction after the middle frame. We observe that for the specific frames with added noise, the performance decreases. On those frames content can disappear and context not correctly recalled. However, in both cases, within steps after the noisy frames (after frame 34
18

---

Model
Avg. PSNR
Fin. PSNR
SSIM
Low Complexity
Baseline (low complexity)
26.09
25.60
0.95
Ours (low complexity)
36.72
35.78
0.97
Middle Complexity
Baseline (middle complexity)
27.27 39.68
26.70
0.94
StateSpaceDiffuser (Ours)
39.32
0.98
High Complexity
Baseline (high complexity)
23.09
22.87
0.93
StateSpaceDiftuser (Ours)
31.67
30.87
0.97
Table Generalization Across Visual Complexity. Our approach is shown to consistently outper- form the baseline on multiple levels of visual complexity without finetuning.
Input:
Frames
Prediction
Ground Truth
StateSpaceDi fuser
State Space World Model
StateSpaceeiiffser
Ground Truth
StateSpaceDif fuser
State-Span World Model
StateSpaceOiffuser (na state)
Ground Truth ...
StatesppaceDifuser
State- Space World Model
SteteSpaceOiffuser (no state)
Figure 9: CSGO Ablations. We show that without the state-space features from the Long Context Branch, StateSpaceDiffuser loses context preservation ability. We also show that while State-Spar Space World Model demonstrates long-context memory, the produced images are of subpar visual quality.
noisy frames window size 4), the memory and content recovers and is correctly predicted, with stable performance until the last frame. The lower scores suggest some loss in performance. However, the fact that memory recovers after the noise suggests a certain level robustness to noise.
B.6 Generalization Across Visual Complexity
In this work, we have evaluated different environment setups with increasing level of complexity the very constrained Simple MiniGrid, free navigation in maze (MiniGrid), and 3D first-persod environment (CSGO). To further compare the performance of StateSpaceDiffuser across environments with different visual complexities, we generate 2 more variants of our MiniGrid dataset based on visual complexity. We define complexity as number of markers and complexity of the maze walls (values in the range [1, 5)). We generate a dataset with low complexity (200 markers, difficulty 3), and with high complexity (450 markers, difficulty 5). Our original MiniGrid dataset lies between in terms of complexity (360 markers, difficulty 4). We take our StateSpaceDiffuuse pretrained model on context length 50 (middle complexity) and evaluate it the new datasets with varying difficulties (without any finetuning).
19

---

Input: Franes
Ground Truth
StateSparolk TTuser
State Space World Macel
Ground Truth
StoteSpareO1 Fuuser
State Worts Morel
DEAMORD
StateSpace01 Tuueer
Space
Figure 10: Qualitative Results with Context Length 50 MiniGrid. StateSpaceDiffuser demon- strates long-context preservation compared to DIAMOND and better visual fidelity compared to State-Space World Model.
Ground Truth
StateSpaceDi ffuser
Full Sequence Sample with Context Length 50 on MiniGrid. Demonstrates a full 50-frame context and the per-frame predictions. In red are shown the frames from the second half returning to the initial position.
Figure 11:
The results in Tab. suggest that our model better able generalize to lower complexity rather than higher complexity. However, in all cases, the performance remains higher than the baseline (DIAMOND).
B.7 State Features Ablation
To expand on the state ablation experiment results from Sect. in Fig. are shown more examples comparing StateSpace with and without a state. The results reconfirm that StateSpace loses its special ability memorize long context after we zero out the states that were produced before passing them to the Generative Branch. Without the state, the model can no longer recall long-context information and hallucinates new locations.
B.8 Comparison with State-Space World Model
In Fig. are shown extra examples of the State-Space World Model (used as a Long Context Branch StateSpaceDifffuse) versus the result StateSpaceDiffuser is clearly seen that the State Space World Model produces more artifacts and more visually unappealing images. However, its important property to reconstruct previously observed content leads to features being crucial for StateSpaceDiffuser to exploit as part of its Long Context Branch. In the shown examples, the images are predicted in an autoregressive manner by re-encoding StateSpaae output from the previous steps. We have observed that this approach is much more stable than feeding the prediction
20

---

of State-Space World Model to itself. The latter produces very blurry images with no discernible detail. Note that the State-Space World Model has not specifically been trained with an autoregressive approach. This confirms that in StaffS,,ae the diffusion component stabilizes the state-space component, which in turn improves the long-context abilities of the diffusion component.
Compared to MiniGrid, here we see much more pronounced artifacts on the results of State-Space World Models, which are then cleaned up by the Generative Branch in StateSpaceDiffuser. This confirms the potential of StateSpaceDifftser particularly for complex environments.
C MiniGrid Evaluations
In this section, we expand our discussion on our model's performance on the simpler MiniGrid dataset and offer additional qualitative and quantitative results.
C.1 MiniGrid Qualitative Evaluation
In Fig. 12 and Fig 10 we show visual samples with the forward-backward evaluation used in our quantitative results. We assess the StateSpaceDiffuser model, trained on sequence length 51 and presented in the main paper and its corresponding State-Space World Model.
Note that in our standard evaluation on MiniGrid, at each time step we give the ground truth sequence up to that step and predict the next observation. In a single time step the agent takes fixed motion in one of four directions and reveals exactly 1 row or column of the environment depending on the direction. As most of the content of the next frame is present in the previous frame, the unknown content is only contained within the new revealed area. In the first half of the sequence, the agent explores by navigating the maze. In Fig. 1 we we show one such traversal and the predicted next frame for each time step from StateSpaceDifffsee
is observed that the new revealed area is predicted far from the ground truth in the first half of the sequence. This is expected as this area has not yet been observed. However, in the second half, the return along the path, the ability to recall the content from the long given sequence helps to predict the correct content of the repeatedly revealed areas. Therefore, in all our evaluations we have made the decision to only consider the prediction quality of the second half of the second half of the sequence.
In context length 16- Fig. 12 we clearly observe poor performance of the diffusion-only baseline, DIAMOND, as this model has no method to take into account long context. Looking closely at the State-Space World Model's output, we observe shifts in color and inconfidence in the content of the new revealed areas. The effect is subtle and varies in the sequence, but it is noticeable. This effect causes a drop in fidelity metrics and a direct consequence of the non-variational approach of predicting the next frame from State-Space World Model. In contrast, StateSpace is free of such artifacts and predicts closer to the ground-truth images. Still, as is conditioned on the state-space features, it can be affected by significant uncertainty in the content of particular grid cells.
The observations are even more pronounced for context length 51 Fig. 10 Later in the sequence, the State-Space World model tends to increase its shifts from the ground truth (dar squares appear darker, grey squares tend to fade). While on simpler dataset like MiniGrid such artifacts are less noticeable, for more complex setup like CSGO this becomes more apparent.
C.2 Imagination Qualitative Results
Additionally, instead of giving the ground truth sequence at every step, we also attempt to give only the first half and feed already predicted frames for time steps in the second half of the sequence (imagination). In this way the ground-truth frames in the second half are never seen by the model (same the setup we have in the CSGO dataset). This is more challenging setup, in which the content cannot be copied from the previous ground truth frame, and any errors in the current frame prediction propagate into the next frame
We show qualitative results in imagination Fig. 13 It is observed that in this more complex setup the diffusion baseline quickly drifts away from the context, and the entire image no longer corresponds to past context. In contrast, because of the incorporation of state-space features, the StateSpaceDiffuseer is noticeably better preserving the content of previous steps.
21

---

Input: 9 Frames
Prediction
Ground Truth
DIAMOND
State-Space World Model
faatespaceDiffuser
Ground Truth
DIAMOND
State-Space World Model
StateSpaceDi ffuser
Ground Truth
DIAMOND
State-Space World Model
StateSpaceDi ff fuser
Ground Truth
DIAMOND
State-Sppe Space World Model
StateSpaceDi ffuser
Ground Truth
DIAMOND
State- Space World Model l.
StateSpaceD1 fuser
Ground Truth
DIAMOND
State- Space World Model
StateSpaceDi ffuser
Qualitative Results on MiniGrid. Compared to DIAMOND, State-Space World Model is able recall past content better but lacks in certainty and visual fidelity. However, StateSpaceDiffuser is able to both to consider long context and to produce a high quality image.
Figure 12:
22

---

Prediction
Input: 9 Frames
Ground Truth
DIAMOND
StateSpaceDD1 ffuser
Ground Truth
DIAMOND
StateSpaceDDi of fuser
Ground Truth
DIAMOND
StateSpaceDiffuser
Ground Truth
DIAMOND
CrateSpaceDiffusee
Ground Truth
DIAMOND
StateSpaceDi wwweeiiffuuuue fuser
Ground T
STAMOND
Figure 13: Imagination Qualitative Results on MiniGrid. Frames are consecutively generated given previously generated frames. Sacceieeeeeeeetteeee shows clear superiority context preservation. previooi
23

---

D CSGO Evaluations
D.1 CSGO Qualitative Evaluation
We show additional qualitative examples in Fig. 14 We show the last 8 frames (predicted without the use of the last 8 ground-truth frames), given frames of the mirrored version of the sequence. The last predicted frame uses context length 16. In the results, it is observed that DIAMOND our unable to recall a previous context that was left beyond the input frames
diffusion baseline accepts as input per step. In contrast, StateSpaceDiffuser is able to predict the correct content through its Long Context Branch.
In addition, we evaluated our model (trained on context length 16) in a longer context of 50 frames. In Fig. 13 we show examples, where the last frames are predicted in an autoregressive manner, while the first 43 frames are given as ground truth (their corresponding predicted frames are also depicted). Therefore, the context for the last frame is of 43 ground-trute -tuth images and 7 generated images. We show the last 26 frames in the sequence, as the first 25 are a mirrored version of them. In this challenging setting, we observe more artifacts and significantly less memory capabilities. However, the model is capable of recalling some visual cues on this context length that were visible
Input: 9 Frames
Prediction
Ground Truth
DIAMOND
StateSpaceDiffuser
Ground Truth
DIAMOND
StateSpaceDiffuut
Ground Truth
DIAMOND
StateSpaceDi
Ground Truth
StateSpaceDi fuser
Ground Truth
DIAMOND
StateSpaceOiffuser
Ground Truth
DIAMOND
StateSpaceDiffuse
Figure 14: Our StateSpaceDiffuser model is able to recover from insufficient information in the short context, while the diffusion baseline DIAMOND, has no mechanism to do so.

---

Input: 25 Franes
Prediction
Ground Trush
Prediction
Prediction
Ground Truch
Prediction
Figure 15: StateSpaceDiffuser on Context Length 50 CSGO. In red are marked the imagined frames and the corresponding ground truths. The model is able to reconstruct frames seen at the start of the sequence.
Input: 9 Frames
Prediction
Ground Truth
DIAMOND
StateSpaceD1 fffusr
Ground Truth
DIAMOND
StateSpaceDiffuser
Ground Truth
DIAMOND
StateSpaceDi ffuser
Figure 16: Recovery from strong action. Strong actions cause significant artifacts in DIAMOND. While this behavior is inherited by StafeSpace in contrast, it quickly recovers from artifacts using the state-space features.
only at the start of the sequence, visible on the examples the red roof, looking through the scope, revealing a door.
D.2 Strengths and Limitations of StateSpaccDifftser
DIAMOND has limitation, which causes significant artifacts when the action is strong and results in larger visual change (e.g., large turn.. Our Generative Branch is based on DIAMOND and hence has similar limitation. However, while DIAMOND's artifacts tend affect the entire predicted sequence, the StaffSp has the property recover in subsequent steps by making use of the state-space representation recover content. This is visible throughout Fig. [4 but also particularly in the examples shown in Fig. 16
25

---

Input: 9 Frames
Prediction
Ground Truth
DIAMOND
StateSpace0iffuser
Ground Truth
DIAMOND
StateSpaceOiffuser
Figure 17: Limitations. Our method sometimes only reconstructs coarse features and leaves details out. Even in these cases, the content appears closer to the ground truth than the diffusion baseline.
Input: 9 Frames
Prediction
Ground Truth
DIAMOND
StateSpaceDi iffser
Figure 18: Measuring Content Preservation. On the shown sample, visually our method preserves the content better than the baseline. However, the PSNR fidelity metrics do not match this conclusion. We conclude that fidelity metrics do not match well the goal of content preservation.
The effect of the strong action is one of our main motivations to use autoregressive-styy. prediction for the CSGO models rather than a single frame prediction given a full context (as done with MiniGrid), While MiniGrid has actions with constant measurable visual effect, CSGO's s actions vary in strength. observe that many sequences consist of almost no motion, divided by few strong actions. single-action prediction task, weak motion would not require much new content to be generated. This previous frame, and not much improvement expected compared baseline. Conversely, with strong actions, the baseline and StateSpaceDiffuser exhibit inconfidence in the next frame. This affects the short context window, filling it with artifacts. comparison to the baseline, the StateSpaceDiffuser is able recover in the following frames.
often results in a copy
While our method's 's content consistently closer the ground truth than the diffusion baseline, sometimes is only able to preserve coarse features (colors, general shape of the scene) and is less effective with finer details. We show this in Fig. 17 visible that sometimes our method misses showing object (crate in the first example) or just preserves the general shape of scene (second example). Even with those limitations, the model often outperforms the baseline. As higher level of detail requires larger memory capacity, we believe that with more computational resources, scaling the Long-Conter Cong-Context Branch can aid these issues.
D.3 Quantitative Perceptual Consistency
As previously discussed, contrast to MiniGrid, the complexity of the CSGO environment makes fidelity metrics such as PSNR unsuitable for evaluating content consistency. CSGO is characterized by actions with variable motion unfolding over multiple frames. Long rollouts accumulate small motion mismatches into camera-view drift, so later frames need not match ground-truth pixels. Content often remains perceptually similar while viewpoint and details differ, and PSNR under-reports this similarity. We demonstrate this by computing the fidelity metries of an example. For the baseline, we obtain 20.77 Avg. PSNR and 16.17 Fin. PSNR. For StateSpaceDiffuss and 16.1 Fin. PSNR. Given the metrics, in this example, our model does not differ significantly from the baseline in terms of quality of the last frame and is somewhat worse average. However, when looking the visual results in Fig. T8] we clearly see more similarity to the ground truth the content of our method than in the baseline. The viewpoint, details, object proportions, and parts of
19.36 Avg. PSNR
26

---

Sample
Which image metches bener the grey tordered image:
Figure 19: User Study Sample Question.
the scene do not match exactly, causing the metric to be worse. Due to this inadequacy, and following common practice, we opt for a user study to evaluate the quality of the CSGO results.
D.4 User Study Details
The user study is performed by first selecting 40 examples where long-context memory would be important for example, content at the beginning of sequence covered up by wall. The examples are picked from the ground truth images in the test set; the predicted images are not observed when selecting, order not to bias the process. After generating the predictions from the diffusion baseline and StateSpaaeDiiffseee
we build two triples of images per prediction for prediction horizons 15 and 17. This results in 80 total triplets. For each of them, we ask participants to determine if the baseline or our model is better. The participants are coleagues and students from our institute, external to this project. In order to avoid bias, we shuffle the order of the baseline and StateSpaceDiffuser results for each sample, marking the first blue and the second red. The user is asked to compare the match blue and images the ground truth image. An example question is shown in Fig. 19
The text that the users saw is visible below:
Thank you for taking part in
study! For each question, you will see row of 3 frames:
First frame (grey border) our ground truth
Second frame (blue border) a frame that attempts to resemble the first frame
Third frame (red border) a frame that attempts to resemble the first frame
Your task is to judge if the red or blue frame resembles more closely the grey frame. Rate each pair by selecting: Blue Blue frame resembles the grey frame better; Red Red frame resembles the grey frame better. Please base your decision on both the content and not the visual quality of the images.
27

---

1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
It does
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Discussion on limitations can be found supplementary material.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer:
[NA]
Justification: Not applicable
Guidelines:
28

---

The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We share details in methodology and supplementary material, we will share the code
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of model checkpoint, other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model limited in some way (e.g., registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
29

---

Justification:
Data generation and model code will be provided
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.c cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ siaips.cc/ubbllccguuddeess /CoeeS/mmidebbcee ssi onPolicy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer:
[Yes]
Justification:
We include those details mostly in supplementary material
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question:
Does the paper report error bars suitably and correctly defined or other appropriate
information about the statistical significance of the experiments?
Answer: [No]
Justification: Experiments are too heavy be executed multiple times, we have computa- tional limitations
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
30

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative eтror rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of exccution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We mention the gpu models and number of gpus. The details about workers etc will also be visible in code.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics htttp::/neurips http:///aaarp..coooo/occccccaaaa
Answer: [Yes]
Justification: Yes
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations their jurisdiction..
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
No direct societal impact
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
31

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Not high risk
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We cite every dataset and model paper used
Guidelines:
The answer NA means that the paper does
existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
32

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: No extra assets
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [Yes]
Justification: Described in detail in the Appendix
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
33

---

Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy for what should or should not be described.
34


Reviews:
[{"review_id": "Wp08G9ZHsc", "reviewer": "Reviewer_VVHr", "summary": "This 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.", "questions": "Have you tried end-to-end training of the state and diffusion parts together, instead of training the state model separately first?", "limitations": "They mentioned their limitations in supplementary material that the reconstruction quality is still not so good right now.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Long-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", "was_revised": true, "final_justification": "I maintain my recommendation to accept this paper. Long-term consistency is a crucial and underexplored challenge in world models. While the area is still in its early stages, this paper makes a meaningful contribution by proposing a novel architecture that demonstrates promising capabilities in long-term memorization, as evidenced by experiments on a Counter-Strike dataset. This is a valuable step forward and has the potential to inspire future work in this important direction.\n\nAfter reading the other reviews, I feel that some of the critiques may be overly harsh given the scope and maturity of the problem. It is unrealistic to expect a single paper to solve such a fundamental and long-term research challenge in its entirety. While I acknowledge the current limitations—such as evaluation on a single domain-specific dataset—I believe the contributions are solid for a first step. These limitations can reasonably be addressed in future work.", "rebuttal": "We thank the reviewer for recognizing the importance of our topic, experimental results and the potential for use in future models. We address the reviewer's comments below.\n\n**Long Sequence Models Ablation.**\nAs the reviewer suggested, we have trained versions of the State Space World Model, but by replacing Mamba with LSTM and GRU for a context size 50 on our Minigrid setup. We add a linear layer on the input and output (dim 256, with ReLU activation) of the models. We are unable to use a standard transformer for a long image sequence because of the rising computational cost with each step - $O(N^2)$ complexity of the model in the general case. Furthermore, after trained, the transformer cannot generalize across a longer sequence than trained with. Results are shown below:\n\n| Model    | Avg. PSNR | Fin. PSNR |  SSIM |\n| :------- | --------: | --------: | ----: |\n| LSTM     |    26.80 |    26.906 | 0.93 |\n| GRU      |    27.40 |    26.680 | 0.93 |\n| Mamba    |    32.64  |   32.44  | 0.96 |\n\n\nMamba outperforms the other sequence models, not having to perform BackPropagation Through Time (BPTT), among other advantages. In addition, as shown in Tab. 2 in the main paper, Mamba also outperforms S4 - another popular state space model. Compared to other methods, Mamba excels at modeling long-range temporal dependencies, while remaining not computationally demanding. This is confirmed in the original paper that proposed Mamba - \"Mamba: Linear-Time Sequence Modeling with Selective State Spaces\", where Mamba is compared with other state space models (Table 1) and transformer language models (Table 3).\n\n\n**End-to-end Training.**\n We trained the SSM and the diffusion model end-to-end in two setups: (1) from scratch, and (2) with an SSM pretrained on sequence dynamics followed by joint training. However, the best results were achieved by first train the SSM to capture long-range structure and then freezing it while training the diffusion decoder.\n\nJoint training is challenging, as the components optimize on different timescales. The SSM needs many iterations of sequence-level supervision to stabilize a long-horizon representation, whereas the diffusion decoder optimizes a local denoising objective and rapidly finds a short-horizon solution that ignores a still-noisy SSM state (conditioning collapse). When we unfreeze a pretrained SSM for joint finetuning, decoder gradients quickly drift the SSM away from the representation that carries long-range cues; during this adaptation, the decoder again defaults to visual inputs and weakens reliance on the state.\nFreezing the SSM preserved its long-horizon information and led the diffusion decoder to make consistent use of it, yielding the strongest qualitative and quantitative results in our compute budget. \n\nWe find that this approach of separately training the two components has the important benefits of drop-replacing one of them without finetuning the entire system. For example, it allows us to use the same diffusion decoder that we trained for context size 16 with a new state space model, that we trained for sequence size 50, and the new resulting StateSpaceDiffuser model works better for longer context, without requiring any diffusion training. (Note that even without training a new state space model, StateSpaceDiffuser can generalize to much longer context out of the box - check answer of Reviewer dCL4 for more details.)\n\n**Reconstruction Quality.**\nAs part of StateSpaceDiffuser we train a lightweight conditional diffusion decoder under a fixed compute budget (no large pretrained backbone), which can yield frames artifacts late in long rollouts. Our contribution - the long-context conditioning mechanism via the SSM state, is agnostic to the choice of diffusion model. Replacing the decoder with a newer/larger one or training longer would improve visual sharpness without changing our method. Additionally, scaling the SSM (state dimension/heads/parameters) is expected to reduce high-frequency decay over many steps. Importantly, our primary objective is content consistency at long horizons. The qualitative results shown come from a forward–backward experimental setup with sufficient motion to occlude and later reveal content. On these challenging cases the baseline (DIAMOND) often becomes inconsistent or predicts incorrect content, while our model recalls the previously seen content - at some quality cost. For completeness, in the final Sup.Mat. we will also include qualitative examples from randomly sampled test sequences without the forward–backward constraint to show that visual quality in the general case does not degrade.\n\n**Writing Suggestions.**\nWe thank the reviewer for the suggested work - we will add it to our related work."}, {"review_id": "LFy5NAsxrU", "reviewer": "Reviewer_dCL4", "summary": "To 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.", "questions": "To 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?", "limitations": "Yes, discussed in detail in supplementary.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths: 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", "was_revised": true, "final_justification": "I appreciate the author's efforts in introducing additional experiments and evaluations during the rebuttal phase. I have increased my score to reflect this (2-3).", "rebuttal": "**Computational Complexity.** Mamba is known to have a linear complexity proportional to the sequence length, while trainsformers - a quadratic complexity over the (token) sequence length. While Mamba maintains a state, transformers and CNNs don't. To process the full sequence it has to be stored in its entirety in memory and processed at once. Even with KV-caching, the resulting linear complexity is only available in inference, and is linear per call - it does not support a streaming over a constantly increasing sequence. This is the case for current world models, including DIAMOND - a diffusion-based world model.\n\nMeasurements of the computational cost of the SSM and the diffusion model in our particular case can be seen at Lines 25-29 in The Supplementary Materials. Our SSM takes only 0.6\\% of all inference computations of the full model, for sequence size 16, and 1.8\\% for sequence size 50. To account for this computational cost in our StateSpaceDiffuser scores, we normalize them by multiplying by $1-0.006$ for sequence size 16 and $1-0.018$ for sequence size 50. We report the normalized scores and compare with the baseline:\n\n| Model                     | Avg. PSNR (Norm.) | Fin. PSNR (Norm.) | SSIM (Norm.) |\n| :------------------------ | ----------------: | ----------------: | -----------: |\n| **Context Size 16**       |                   |                   |              |\n| DIAMOND                   |             27.13 |             25.44 |        0.950 |\n| StateSpaceDiffuser (Ours) |             34.41 |             37.81 |        0.974 |\n| **Context Size 50**       |                   |                   |              |\n| DIAMOND                   |             26.13 |             25.15 |        0.950 |\n| StateSpaceDiffuser (Ours) |             27.62 |             29.75 |        0.943 |\n\n\nEven after normalizing, significant gains are observed.\n\n**Generalization Over Longer Context.** We train on fixed-length sequences, which reasonably raises the reviewer's generalization concerns. StateSpaceDiffuser operates on much longer contexts without finetuning: in Sup.Mat. Sect. 2.1 (L50–56), a model trained only at length 16 already outperforms the baseline at 50. At the reviewer’s request, we further evaluate a model trained at 50 on lengths 100 and 150 using a new MiniGrid test set (up to 150). Results are below:\n\n\n| Model                           | Avg. PSNR | Fin. PSNR | SSIM  |\n| ------------------------------- | --------- | --------- | ----- |\n| Context Length 100              |           |           |       |\n| Baseline (DIAMOND)              | 26.386    | 26.240    | 0.947 |\n| State Space World Model (Mamba) | 31.648    | 30.885    | 0.957 |\n| StateSpaceDiffuser (SSW)        | 37.993    | 35.866    | 0.977 |\n| Context Length 150  |   |     |   |\n| Baseline (DIAMOND)      | 24.346    | 24.197    | 0.938 |\n| State Space World Model (Mamba) | 27.927    | 26.980    | 0.935 |\n| StateSpaceDiffuser (SSW)        | 30.750    | 28.930    | 0.962 |\n\nIt can be seen that StateSpaceDiffuser successfully generalizes to 100 and 150 frames in the future, keeping a significant gain over the baselines. These results could be further improved by finetuning on longer sequences. \n\n\n**Evaluation Across Seeds.**\nAs the reviewer suggested, we perform evaluation multiple times under different seeds. We run 4 times with different seeds for MiniGrid with context size 16 (computationally cheaper) and obtain $41.00\\pm0.008$  Avg. PSNR, $40.52\\pm0.019$ Fin. PSNR, $0.98\\pm0.0004$ SSIM. We also perform evaluation over 4 seeds of a larger scale more expensive evaluation - our new 100 context length experiment from above. This results in: $37.99\\pm0.002$ Avg. PSNR, $35.88\\pm0.029$ Fin. PSNR, $0.98\\pm0.000004$ SSIM. In both cases, the obtained metrics are extremely stable and consistent in between seeds (low std).\n\n**Evaluation Over Varying Visual Complexity.**\nIn our main paper, we evaluate on 3 different environment setups with increasing level of complexity - the very constrained Simple MiniGrid (Line 245), free navigation in a maze (Line 257), 3D first-person environment (Line 274). We keep the model scale (hyperparameters) the same across these environments and evaluate the performance on each. To further compare the performance of StateSpaceDiffuser across environments with different visual complexities, we generate 2 more variants of our MiniGrid dataset based on visual complexity. We define complexity as number of markers and complexity of the maze walls (values in the range $[1,5]$). We generate a dataset low complexity (200 markers, difficulty 3), and one with high complexity (450 markers, difficulty 5). Our original MiniGrid dataset lies in between in terms of complexity (360 markers, difficulty 4). We take our StateSpaceDiffuser pretrained model on context length 50 (middle complexity) and evaluate it on the new datasets with varying difficulties (without any finetuning):\n\n| Model                            | Avg. PSNR | Fin. PSNR | SSIM |\n| :------------------------------- | --------: | --------: | ---: |\n| Baseline (low complexity)    |     26.09 |     25.60 | 0.95 |\n| Baseline (middle complexity) |     27.27 |     26.70 | 0.94 |\n| Baseline (high complexity)   |     23.09 |     22.87 | 0.93 |\n| Ours (low complexity)            |     36.72 |     35.78 | 0.97 |\n| Ours (middle complexity)         |     39.68 |     39.32 | 0.98 |\n| Ours (high complexity)           |     31.67 |     30.87 | 0.97 |\nIt is observed that our model is better able to generalize to lower complexity rather than higher complexity. Still, in all cases, the performance remains higher than the baseline (DIAMOND).\n\n**User Study.**\nCSGO is characterized by actions with variable motion unfolding over multiple frames. Long rollouts accumulate small motion mismatches into camera-view drift, so later frames need not match ground-truth pixels. As shown in Sup.Mat. (L127–137, Fig. 8), content often remains similar while viewpoint/details differ, and PSNR/LPIPS under-report this similarity. This is a known issue - see High-Resolution Image Synthesis with Latent Diffusion Models (p. 8); Image Super-Resolution via Iterative Refinement (p. 2). From Latent Diffusion: \"simple image regression model achieves the highest PSNR and SSIM scores; however these metrics do not align well with human perception and favor blurriness over imperfectly aligned high frequency details.\". Following common practice (e.g. EmuVideo, Stable Diffusion, Latent Diffusion), we therefore include human preference evaluations.\nWe thank the reviewer for pointing out the mistake in the question's response. The exact instructions, shown to our volunteers is available on Line 150-156 of Sup.Mat.\n\n**MiniGrid Exploration Policy.**\nThe policy for generating the rollouts in MiniGrid is deterministic. We choose a sequence of 40 random markers (colors) and find the shortest valid path between them in the maze. This results in a very long path to follow - the actions are deterministic. Once 50 actions are complete into this path, the policy stops following the rest of the path and retraces the steps back to the beginning. This constitutes one rollout in our dataset. We will clarify this in our methodology section.\n\n**Separate Action Embeddings.**\nOur training consists of two separate stages: 1) training the Long-Context Branch; 2) training the Generative Branch, conditioned on the Long-Context Branch. Therefore, we keep separate action embeddings in each branch to be trained in each stage for the needs of the particular branch. In this way we are able to preserve the modularity - e.g. replacing the state space model without sabotaging any shared action embeddings. We keep the dimensions in the Generative Branch as chosen in the baseline (DIAMOND), for our Long-Context Branch we use more compact dimensionality (16) to account for the lower amount of parameters and computational budget for this branch.\n\n**Robustness to Noise.**\nWe perform the experiment of adding noise in the middle section of the rollouts, suggested by the reviewer. In specific, we consider context length 50 in MiniGrid and add Gaussian noise (std 2.5) to the 11 frames in the middle (so noise remains symmetric in the sequence). We consider two cases: 1) adding noise to the SSM input only; 2) adding noise both to the SSM input and the diffusion model input. Results are shown below at different steps of prediction after the middle frame:\n\n\n| Model                 |    27 |    28 |    29 |    30 |    31 |    32 |    33 |    34 |    35 |    36 |    37 |    38 |    39 |    40 |    41 |    42 |    43 |    44 |    45 |    46 |    47 |    48 |    49 | 50    | 51    |\n| :-------------------- | ----: | ----: | ----: | ----: | ----: | ----: | ----: | ----: | ----: | ----: | ----: | ----: | ----: | ----: | ----: | ----: | ----: | ----: | ----: | ----: | ----: | ----: | ----: | ----- | ----- |\n| SSM noise   | 15.22 | 16.62 | 15.36 | 15.18 | 14.85 | 16.55 | 23.44 | 26.60 | 28.42 | 29.35 | 29.75 | 30.00 | 29.93 | 30.06 | 29.99 | 30.12 | 30.14 | 30.45 | 30.11 | 30.18 | 30.19 | 30.08 | 30.19 | 30.10 | 30.19 |\n| SSM + diffusion noise | 15.50 | 16.90 | 15.49 | 15.30 | 14.94 | 16.65 | 23.28 | 26.51 | 28.35 | 29.33 | 29.75 | 29.95 | 29.85 | 30.09 | 29.88 | 30.11 | 30.08 | 30.41 | 30.12 | 30.13 | 30.17 | 30.12 | 30.20 | 30.13 | 30.15 |\n\nWe observe that for the specific frames with added noise, the performance decreases. We observe on those frames that content can disappear and long context is not correctly recalled. However, in both cases, within 4 steps after the noisy frames (after frame 34 - 5 noisy frames + window size 4) the memory and content recovers and is correctly predicted, with stable performance until the last frame. Still, the lower scores suggest some loss in performance. The fact that memory recovers after the noise suggests a certain level of robustness to noise. \n\n**Writing Suggestions.** Citation issues will be fixed in the final manuscript."}, {"review_id": "PKImrcudAl", "reviewer": "Reviewer_eFWm", "summary": "Existing 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.", "questions": "Q1) 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?", "limitations": "Please see the Weaknesses part.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "# 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.", "was_revised": true, "final_justification": "I thank authors for their efforts to address my concerns. The detailed rebuttal has partially resolved my questions. However, I still believe the evaluation are too simple and not generic enough, and the paper lacks a comparison with some published baselines to let the audience have a sense of the true advantages of this method. I am leaning rejection.", "rebuttal": "We thank the reviewer for their invaluable feedback and for acknowledging the importance of our topic, the comprehensiveness of our evaluation, and the clarity of our writing. We address the reviewer's comments below.\n\n**Computational Cost.** In inference, the state space model maintains a state in a streaming fashion, which allows for live processing of a growing sequence. At every time step of the sequence, the state is updated, used to generate the output, and kept for the next step. As we only need to keep the updated state, the latency remains constant - memory complexity and computational complexity per step remain the same and is not dependent on the sequence size. In contrast, transformers and CNNs (diffusion backbones) have increasing latency per step - even with KV-caching, the computational and memory complexity for a single step is $O(T)$, where T is the runtime. The complexity per step quickly becomes computationally infeasible with the growing context, making it necessary to use a cache of a few recent frames, preventing long-term dependencies. Conditioning the world model on the SSM state therefore preserves long-horizon information while keeping both memory and per-frame compute constant.\nIn addition, we remind that while training, transformers still have a quadratic complexity, which is not the case with SSMs.\n\n**Baselines.** We thank the reviewer for suggesting Diffusion Forcing. It is a strong transformer diffusion model that excels at long video generation. However, we find that it too has the same context limitation, innate to transformers. We used their official online demo and performed a long context test on all 16 of the images they made available to try (Demo 3 - Image -> Extremely Long Video). In our setup we do a 60 degree turn in place to the right, followed by a 60 degree turn to the left. This resulted in 10 generated frames. In all of our 16 tries we observed the content from the start that was hidden by panning the camera to be forgotten - in its place completely new content is hallucinated at the end of the generated video. We are not allowed to upload images or add links, but we encourage the reader to attempt this quick test.\n\nDiffusion Forcing is targeted on long output generation, while we target consistency with long input (context) - even if generating a short future sequence. Our experiments with DIAMOND give us conclusions about the long-context abilities of our model. However, we recognize that, in future work, combining both paradigms can create a more capable model - both handling long generation and long context. \n\n**CSGO Qualitative Results.**  (1) As part of StateSpaceDiffuser, we train a light diffusion world model, under a fixed compute budget. As a result, frames in the long rollouts can occasionally appear with lower quality. Our contribution - the long-context mechanism via the SSM state, is decoder-agnostic: if a newer/larger conditional diffusion model is substituted, or if training time and sampling budget are increased, visual sharpness can improve without changing our method. Orthogonally, scaling the SSM (state dimension/heads and number of parameters) reduces the decay of high-frequency details carried over many steps, which also aids perceived sharpness. (2) Our primary objective is content consistency at long horizons. Therefore, the qualitative panels on Fig.6 and Sup. Mat. are in the challenging setting of forward-backward evaluation, with ensured enough motion to hide and reveal content. In these cases, the baseline produces inconsistent results while our method recalls the previously seen content, with a visual quality trade-off on the output. We will include qualitative examples in the final version of Sup. Mat. to demonstrate that the visual quality is better in the general case - on random sequences, without mirroring the action sequence (forward-backward).\n\n**Context Length.** In training, we always use training sequences of a fixed size. At a first glance, it may seem that this could lead to overfitting and the reviewer rightfully raises the question about generalizeability on different context lengths. StateSpaceDiffuser is capable of operating on much longer context lengths than it has been trained on, without any fine-tuning required. We have demonstrated this in Sup. Mat. Sect. 2.1 (Lines 50-56), where we show that a model, trained exclusively on context length 16, significantly outperforms the baseline on context length 50 - without any finetuning. To expand this discussion, as the reviewer asked, we take our model, trained on context length 50 and test it on much longer context sizes. To this end, we generate a new Minigrid test set with sequence length up to 150. Results are shown below:\n\n\n| Model                           | Avg. PSNR | Fin. PSNR | SSIM  |\n| ------------------------------- | --------- | --------- | ----- |\n| Context Length 100              |           |           |       |\n| Baseline (DIAMOND)              | 26.386    | 26.240    | 0.947 |\n| State Space World Model (Mamba) | 31.648    | 30.885    | 0.957 |\n| StateSpaceDiffuser (SSW)        | 37.993    | 35.866    | 0.977 |\n| Context Length 150  |   |     |   |\n| Baseline (DIAMOND)      | 24.346    | 24.197    | 0.938 |\n| State Space World Model (Mamba) | 27.927    | 26.980    | 0.935 |\n| StateSpaceDiffuser (SSW)        | 30.750    | 28.930    | 0.962 |\n\n\nIt can be seen that StateSpaceDiffuser successfully generalizes to hundreds of frames in the context.  We observe visually reasonable visual and context recall quality - examples will be included in Sup. Mat.\n\n**Writing feedback.** We thank the reviewer for pointing out the double reference issue, which we will make sure to fix in the final manuscript."}, {"review_id": "PHbUQ4Zvzk", "reviewer": "Reviewer_1SP4", "summary": "This 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.", "questions": "1. 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.", "limitations": "A 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.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**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", "was_revised": true, "final_justification": "Thanks to the authors for being active in the discussion phase. They addressed all my concerns. \n\nI agree with Reviewer VVHr that some of the other reviews may be overly harsh in insisting on evaluations on more complex datasets as a requirement for acceptance. As long as the paper is well positioned with respect to concurrent and prior work, I think it is reasonable for the evaluation to be limited given the inevitable resource constraints, since we cannot expect a single paper to solve every problem. The authors have said they will include these discussions, which I believe is more than sufficient for acceptance.\n\nI also think the way SSMs and diffusion models have been combined in this paper is non-trivial. After discussing with the authors, the approach seems well-motivated; it simply needs to be presented more clearly in the methods section, which the authors have agreed to do. So I would recommend the paper for acceptance.", "rebuttal": "We thank reviewer for the feedback and the comments about importance of our topic, the suitability of our evaluation framework, the effectiveness of our approach and the diversity of our test settings. We address the reviewer's comments below.\n\n**Computational Limitations of Existing Methods.**  \nEven if 50 frames fit in memory, the real barrier is computational complexity. During training, our SSM updates a fixed-size state once per frame, so compute and activations grow linearly with sequence length $T$. A Transformer-based diffusion must still form full self-attention, which is quadratic $O(T^{2})$ in compute and memory. At inference the gap remains: the SSM’s per-step cost is constant, whereas even KV-cached attention or large-receptive-field CNNs scale linearly with the window. In other words, the SSM can process an increasing sequence in a streaming manner with constant latency. Thus a 50-frame cache alleviates memory but not the quadratic compute bottleneck, and long contexts are typically truncated, forfeiting long-range information.\n\nLong Context Tuning for Video Generation (Guo et al.)  trains with full attention over an expanded window (O(N²)), and at inference the per-step cost still scales with the window despite KV cache. We instead compress history into a fixed-size SSM state, enabling linear complexity and constant latency (even with an LCT decoder). Moreover, LCT’s asynchronous shot generation is ill-suited for interactive environments with per-frame actions.\n\nFurthermore, in contrast to methods which finetune/train new models to extend the context, we show that our method is able to extend to much longer sequences than trained without finetuning. We refer to Sup.Mat. Sect. 2.1 where we generalize from context size 16 to context size 50, and an experiment that can be seen in the response of Reviewer dCL402, where we generalize from context size 50 to up to 150.\n\n**SSM/Diffusion Training Separation.** We confirm that our training consists of two stages: 1) Training the SSM (State Space World Model) 2) Load the weights of the SSM, freeze them, and train the Generative Branch of StateSpaceDiffuser. \n\nJoint (end-to-end) training - from scratch or finetuned - underperformed. Training separately worked best. The two branches learn on different time scales: the SSM needs many sequence-level iterations to stabilize long-horizon state, while the diffusion model quickly fits a local denoising loss and ignores a still-adapting SSM. Finetuning didn’t help: diffusion gradients disrupt the SSM and the decoder reverts to immediate image cues. Freezing the SSM prevents drift and lets the decoder use long-term context. Importantly, the diffusion model does not just \"render the SSM’s output\": in Sup.Mat. Sec. 3.6 Fig. 11, an ablation shows the SSM alone produces low-quality frames with the right content, while StateSpaceDiffuser—given the same SSM features - produces visibly higher-quality, context-aware results, not a copy of the SSM output. Thus, the diffusion model isn’t bottlenecked by SSM visual quality.\n\nThe reviewer is concerned if performing the training separately could be a scalability issue.\nWe view the modular training as a scalability advantage, not a limitation. It decouples learning an efficient, fixed-size memory for long-range semantics (the SSM) from high-fidelity rendering (the diffusion decoder). This lets us upgrade either module without retraining the other (e.g., drop-in a stronger decoder, or train an SSM on longer sequences or more data), making it possible to focus on larger scale training on a single component at a time. As we discuss in the next section, we believe the scalability to more complex environments can improve with larger-scale models - decoupling training does not affect the scalability of the models.\n\n**Scalability.** We agree that MiniGrid is a controlled setting. This is why we also evaluate on CSGO, which exhibits high-entropy first-person motion, long occlusions/re-appearances, and large viewpoint changes - failure modes typical of real video. Our goal is long-horizon content consistency, which goes beyond maximizing single-frame sharpness. The SSM state is a compact summary by design, while a diffusion decoder recovers visual detail each frame. We alotted most of our limited model compute budget for a lightweight diffusion model, complementing it with a cost-effective SSM. Under a larger compute budget, scaling up the SSM (number of parameters, state size, heads) and the diffusion model (training iterations, more parameters, newer model) is the most straightforward way to improve the visual quality and the capacity of the context preservation, including for more complex environments. For completeness, the final supplement will also show random rollouts (non forward-backward), clarifying that the low quality of the baseline results is most pronounced in the hardest, long-context cases.\n\n**Limitations.**\n In Sup.Mat. Sect. 3.2 we have discussed our limitations and attributed them: (1) artifacts after strong action, attributed to the diffusion baseline; (2) decay of context details, attributed to the lossy SSM state (potentially scaling up the SSM can help). As CSGO is a pre-collected dataset of human play, it is difficult to separate clean subsets of failure modes in order to collect failure rate statistics. We can however comment that the decay seems to increase with longer-range dependencies. We can also demonstrate this decay on the predictions of MiniGrid context generalization up to context size 150, which is in contrast to the stable performance for context size up to 50 shown on Fig. 7(b). We show the average PSNR in ranges of 5 for the extended prediction:\n\n| PSNR (mean) |  1–5  | 6–10  | 11–15 | 16–20 | 21–25 | 26–30 | 31–35 | 36–40 | 41–45 | 46–50 | 51–55 | 56–60 | 61–65 | 66–70 | 71–75 |\n| :---------: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |\n|             | 35.55 | 32.30 | 32.02 | 31.75 | 31.31 | 31.00 | 30.69 | 30.42 | 30.16 | 29.87 | 29.61 | 29.45 | 29.15 | 29.04 | 28.94 |\n\n**Contribution Importance.**\n We agree that, at a high level, leveraging longer history should help. Our contribution is to make this practically effective for diffusion world models: a compact SSM state that carries long-horizon information with constant per-frame latency/memory, an architecture that feeds that information to a diffusion model, and a training scheme that allows for the diffusion model to make use of this information. Naively training end-to-end led to conditioning under-use. While others have used fixed image caches and explicit representations like pointclouds, this is the first work that explores exploiting a state space model for the benefit of a diffusion world model and makes it possible in practice. We consider this a significant contribution. We will improve our methodology and introduction sections to include these points more clearly.\n\nWe also contribute a forward–backward evaluation with sufficient motion to test long-horizon consistency. Under this protocol, diffusion-only baselines often predict incorrect content, whereas our model reinstates previously seen content - showing the benefit is not just using more history, but how it is summarized and used under realistic compute. We further observe (see Sup.Mat. Sect. 3.2, Sect 3.6, context generalization in this rebuttal) robustness to large motion, generalization to much longer sequences, and that diffusion renders higher-quality frames later in the rollout than the SSM while still relying on SSM features.\n\n**Evaluation Clarifications.** We justify the split in Sup.Mat. L111-119. MiniGrid is an environment where an action results in a fixed amount of motion over 1 step. In MiniGrid, each action induces a fixed 1-step motion, so single-frame prediction cleanly isolates context recall from rendering quality. In CSGO, motion varies and unfolds over many frames; one-step prediction mostly probes short-context and both methods look strong. However, multi-step prediction highlights long-horizon consistency. Our model preserves context while the baseline hallucinates, so we adopt multi-step evaluation for CSGO and include MiniGrid multi-step qualitatives in Sup.Mat. Fig. 4.\n\nIn CSGO multi-step rollouts accumulate camera/viewpoint drift: even with correct actions, slight motion mismatches de-align later frames (see Sup.Mat.). Content often matches while viewpoint/details differ. Thus PSNR/LPIPS fail to reflect content similarity, which we show experimentally in Sup.Mat. L127–137, Fig. 8. This is a known issue - see High-Resolution Image Synthesis with Latent Diffusion Models (p. 8); Image Super-Resolution via Iterative Refinement (p. 2). From Latent Diffusion: \"simple image regression model achieves the highest PSNR and SSIM scores; however these metrics do not align well with human perception and favor blurriness over imperfectly aligned high frequency details.\" . Following common practice (EmuVideo/Stable Diffusion/Latent Diffusion), we therefore include human preference evaluations.\n\nThe Simple MiniGrid setup is an easily reproducible minimal example designed entirely to showcase the the problem and our approach.  Even on a minimal color-recall task (two actions, almost no visual complexity, tiny dataset prone to overfitting), the diffusion baseline, limited by its context window, fails. Processing the long history with an SSM lets our model reliably recover the correct color. We then scale difficulty in MiniGrid and CSGO, showing that long-context recall persists as environment complexity rises.\n\n**Writing Suggestions.** We thank the reviewer for the detailed notes. We will add the requested clarifications on diffusion world model limitations (with the suggested citation), include a Data section figure illustrating MiniGrid and CSGO with example actions, and fix the noted typographical and caption issues."}]

Author rebuttal:

Meta-review: This paper addresses the lack of long-term memory in diffusion-based world models. It introduces StateSpaceDiffuser, a hybrid model that uses a State-Space Model (SSM) to maintain a compressed history of an environment, which then conditions a diffusion model for high-fidelity frame generation. Experiments in 2D and 3D environments demonstrate that this approach significantly improves long-term visual consistency and recall compared to a diffusion-only baseline.

The paper's strengths are clear: it tackles a critical problem in generative modeling, proposes a novel and effective hybrid architecture, provides strong empirical evidence of its success, and uses a well-designed evaluation protocol to directly test its core claim of long-term recall. Initial weaknesses concerning evaluation depth and methodology were comprehensively addressed during the rebuttal with new experiments. The primary remaining limitation, noted by one reviewer, is the scope of the evaluation environments, suggesting that scalability to more complex, real-world domains is an important direction for future work. 

For these reasons, I recommend acceptance. The paper makes a solid, novel contribution to a significant, unsolved problem. The authors rebuttal was critical to solidify the paper's quality, with new experiments that directly addressed initial reviewer concerns regarding evaluation depth and methodology. These should be incorporated in the camera ready. While the evaluation scope can be expanded, the current work provides a valuable proof of concept and opens a compelling new research direction. The AC also agrees with the positive reviewers, that given the maturity of this area, the progress demonstrated in this paper is sufficient for publication, despite the limited evaluation score.
Decision: Accept (poster)