Research map / Image generation / Autoregressive and latent generation
StateSpaceDiffuser: Bringing Long Context to Diffusion World Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Adds a Mamba state-space long-context branch to the DIAMOND diffusion world model so generated frames stay consistent with long interaction histories, evaluated on MiniGrid forward-backward recall and CSGO user studies.
Problem
Diffusion-based world models condition only on a short window of recent frames, so they lose long-term context and drift from previously observed content, hurting temporal consistency in action-conditioned generation.
Approach
Process the entire interaction history with a Mamba state-space model over per-frame Cosmos features, then fuse its last features into a DIAMOND diffusion generator via a two-stage training scheme (SSM first, then frozen-SSM conditioned diffusion), giving diffusion access to long context at almost no extra compute.
Key results
- On MiniGrid context length 50, final-frame PSNR of 39.32 for StateSpaceDiffuser vs 25.14 for DIAMOND; average PSNR improvement of 51.9% over the diffusion baseline.
- In a CSGO user study with 12 participants, preference ratings of 0.20 at frame 15 and 0.24 at frame 17 favored StateSpaceDiffuser over DIAMOND.
- Zeroing the state-space features drops performance below the DIAMOND baseline, and the Long-Context Branch adds only about 0.6% to 1.8% of inference FLOPs.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add state-space features to the DIAMOND diffusion model | StateSpaceDiffuser · MiniGrid, context length 50 | final-frame PSNR: 25.14 -> 39.32 | Helped |
| Add state-space features to the DIAMOND diffusion model | StateSpaceDiffuser · MiniGrid, context length 50 | average PSNR: 51.9% | Helped |
| Add state-space features to the DIAMOND diffusion model | StateSpaceDiffuser · CSGO, 15th frame | user-study preference rating: 0.20 | Helped |
| Add state-space features to the DIAMOND diffusion model | StateSpaceDiffuser · CSGO, 17th (last) frame | user-study preference rating: 0.24 | Helped |
| Add state-space features to the DIAMOND diffusion model | StateSpaceDiffuser · CSGO example | average PSNR: 20.77 -> 19.36 | Hurt |
| Add state-space features to the DIAMOND diffusion model | StateSpaceDiffuser · CSGO example | final-frame PSNR: 16.17 -> 16.11 | No clear effect |
| Zero out the Long-Context Branch features passed to the Generative Branch (ablation) | StateSpaceDiffuser · MiniGrid | PSNR | Hurt |
| Use Mamba instead of LSTM in the Long-Context Branch (ablation) | State-Space World Model · MiniGrid, context length 50 | average PSNR: 26.80 -> 32.64 | Helped |
| Use Mamba instead of GRU in the Long-Context Branch (ablation) | State-Space World Model · MiniGrid, context length 50 | average PSNR: 27.40 -> 32.64 | Helped |
| Use Mamba instead of S4 in the Long-Context Branch (ablation) | State-Space World Model · MiniGrid, context length 16 | average PSNR: 24.29 -> 27.09 | Helped |
| Add Gaussian noise (std 2.5) to 11 middle context frames (ablation) | StateSpaceDiffuser · MiniGrid, context length 50 | PSNR | Mixed |
Surprising findings
- PSNR did not match human judgments of content preservation on CSGO, so the authors relied on a user study instead of fidelity metrics.
- End-to-end training was unstable, so a two-stage scheme with a frozen pretrained SSM was needed.
Limitations
- Generative branch inherits DIAMOND artifacts under strong actions and small lightweight decoder yields visual artifacts in long rollouts.
- Single-layer SSM with 256-dim state can lose fine detail in extended rollouts in complex environments.
- Performance degrades on longer contexts and PSNR-based evaluation is unreliable for CSGO content preservation.
Experimental setup
Datasets: MiniGrid (custom randomly generated mazes) CSGO
Benchmarks: MiniGrid forward-backward recall CSGO forward-backward imagination
Metrics: PSNR SSIM user preference rating
Baselines: DIAMOND LSTM GRU S4 State-Space World Model
Builds on: DIAMOND Mamba EDM Cosmos tokenizer S4WM
Closest papers by meaning
- Video World Models with Long-term Spatial Memory (NeurIPS 2025 · similarity 0.88)
- Learning 3D Persistent Embodied World Models (NeurIPS 2025 · similarity 0.87)
- WorldMem: Long-term Consistent World Simulation with Memory (NeurIPS 2025 · similarity 0.87)
- Trimming the Long-Tail of Visual World Modeling Evaluation (NeurIPS 2026 · similarity 0.87)
- EDELINE: Enhancing Memory in Diffusion-based World Models via Linear-Time Sequence Modeling (NeurIPS 2025 · similarity 0.86)
- Learning Latent Dynamic Robust Representations for World Models (ICML 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.