PaperlayerView exact revision on arXiv

Source-linked AI summary

Long-Context State-Space Video World Models

Ryan Po, Yotam Nitzan, Richard Zhang, Berlin Chen, Tri Dao, Eli Shechtman, Gordon Wetzstein, Xun Huang

arXiv:2505.20171v1cs.CV

TL;DR

Existing video world models have limited temporal memory because attention over extended contexts is computationally expensive, hindering persistent long-horizon simulation. The paper introduces a causal SSM architecture with block-wise scanning and dense local attention, and reports superior long-term memory across Memory Maze and Minecraft evaluations. The method maintains constant per-frame inference time, although it remains limited beyond its training context and to low-resolution synthetic videos.

  • Problem

    Existing video world models have limited temporal memory from restricted attention contexts, while extending context makes training and inference increasingly expensive.

  • Method

    The method uses causal SSMs with block-wise scanning for temporal memory and dense local causal attention for neighboring-frame coherence.

  • Results

    The model outperforms sub-quadratic baselines across spatial retrieval and reasoning metrics on Memory Maze and Minecraft, while maintaining constant per-frame inference time.

  • Takeaways & Limitations

    Causal SSMs provide a scalable route to long-term memory for autoregressive video world models while preserving efficient per-frame inference.

  • Takeaways & Limitations

    The method cannot effectively handle memory longer than its training context and is currently limited to low-resolution synthetic videos.

Abstract

from arXiv · show

Video diffusion models have recently shown promise for world modeling through autoregressive frame prediction conditioned on actions. However, they struggle to maintain long-term memory due to the high computational cost associated with processing extended sequences in attention layers. To overcome this limitation, we propose a novel architecture leveraging state-space models (SSMs) to extend temporal memory without compromising computational efficiency. Unlike previous approaches that retrofit SSMs for non-causal vision tasks, our method fully exploits the inherent advantages of SSMs in causal sequence modeling. Central to our design is a block-wise SSM scanning scheme, which strategically trades off spatial consistency for extended temporal memory, combined with dense local attention to ensure coherence between consecutive frames. We evaluate the long-term memory capabilities of our model through spatial retrieval and reasoning tasks over extended horizons. Experiments on Memory Maze and Minecraft datasets demonstrate that our approach surpasses baselines in preserving long-range memory, while maintaining practical inference speeds suitable for interactive applications.

1. Introduction

Existing video world models struggle to preserve long-term temporal memory, causing inconsistent revisits and costly scaling with longer contexts. The proposed SSM-based architecture combines long-term memory with efficient training and constant per-frame inference.

  • Motivation: Restricted attention context limits temporal memory, so revisiting previously observed regions can produce altered or inconsistent environments.The issue undermines persistent-world simulation for interactive applications.
  • Motivation: Extending attention context incurs quadratic training cost and linearly increasing per-frame inference time.These costs make longer-context generation impractical for realtime, infinite-length applications.
  • Approach: The proposed architecture uses SSMs for causal temporal dynamics and world-state tracking, supplemented by local causal attention.This design targets long-term memory while preserving spatial and temporal coherence.
  • Approach: The hybrid architecture maintains long-term memory with constant per-frame inference speed, unlike sliding-window causal attention that sacrifices memory.This efficiency is presented as suitable for persistent-world interactive generation.
  • Evaluation: Experiments evaluate long-term memory through spatial retrieval and spatial reasoning, with reported gains on Memory Maze and Minecraft.The metrics test revisiting observed locations and inferring unseen locations from prior viewpoints.

2. Related Work

Video generation research has moved from fixed-length diffusion and recurrent approaches toward autoregressive and transformer-based models, but long sequences remain computationally difficult. This work applies causal, unidirectional SSMs to video world modeling rather than using bidirectional scans that forgo efficient autoregressive inference.

  • Video Generation: Video diffusion models generally produce higher-quality videos than discretized autoregressive models but face prohibitive computational demands for long sequences.Standard diffusion generation also lacks online incremental video generation.
  • Video World Models: World-model research has shifted from recurrent neural networks and variational autoencoders toward more scalable transformer-based diffusion and autoregressive models.These models learn visual state transitions for planning and interactive simulation.
  • Linear Attention: SSMs reduce runtime complexity through efficient state updates and can use more expressive hidden states than earlier linear recurrent formulations.They form a prominent family of linear RNNs developed as alternatives to self-attention.
  • SSM-Based Video Models: The paper’s method reports sharp, consistent long-horizon predictions with linear training scaling and constant inference costs, unlike attention-based alternatives.DFoT can look consistent over long horizons but remains bounded by its training context and has quadratic training scaling.
  • SSM-Based Video Models: Prior video-generation approaches typically scan SSMs bidirectionally across entire token sequences, foregoing their efficient autoregressive inference capabilities.The proposed work instead uses unidirectional SSMs for temporal dynamics and world-state transitions.

3. Preliminaries

This section introduces diffusion-based video generation and state-space models as foundations for autoregressive world modeling. Diffusion forcing enables sequential generation by independently varying frame noise, while SSMs provide efficient causal sequence processing through latent states.

  • Diffusion models: Diffusion models corrupt observed data with Gaussian noise and are trained to predict that noise or equivalent targets.The noisy data follows x_t = α_t x_0 + σ_t ε, with α_t and σ_t set by a predefined noise schedule.
  • Diffusion models: Video diffusion models commonly encode videos into latent space with a 3D VAE before learning the diffusion model there.
  • Diffusion forcing: Diffusion forcing samples noise levels independently across frames, enabling sequential denoising conditioned on previously generated clean frames.This supports autoregressive video generation conditioned on streaming controls for applications such as gaming and robotic learning.
  • Attention: Self-attention lets sequence tokens attend to one another through learned query, key, and value projections, whereas causal attention restricts access to previous tokens.Video models can apply block-wise causal masking so tokens attend only within the same or earlier frames, enabling autoregressive generation with KV-caching.
  • State-space models: SSMs model sequence dynamics with latent states and update those states recursively, while Mamba makes the transition parameters content-aware through input projections.At inference, SSM generation requires only the latest latent state, improving time and space complexity over attention at the cost of compressed memory.

4. Methods

The method combines block-wise SSM scanning, local attention, and long-context training to preserve temporal memory while retaining spatial coherence and efficient autoregressive inference. Clean prefixes during training encourage distant-context use, while fixed caches and block states keep generation costs constant per frame.

  • Architecture goals: The architecture addresses shortcomings of naive SSM replacement for autoregressive video generation while targeting long-term spatial memory and constant per-frame inference cost.
  • Block-wise SSM scan: Block-wise SSM scanning reorders spatio-temporal tokens into independently scanned spatial blocks, trading temporal correlation against spatial coherence through block size.Different layers use different block sizes to combine the benefits of smaller and larger blocks.
  • Block-wise SSM scan: Separate states for token blocks effectively increase the SSM state dimensionality at each layer.
  • Long-context training: The training scheme keeps a random-length clean prefix, noises later frames independently, and computes loss only on noised frames to encourage long-term dependencies.Standard diffusion forcing can favor nearby frames because local context is often more useful for denoising.
  • Frame local attention: Frame-wise local attention follows every Mamba scan to recover precise local information and short-term temporal consistency.The method applies block-wise causal attention during training, with the frame window defined by k.
  • Inference and action conditioning: During inference, the model tracks a fixed-length KV-cache for the previous k frames and one SSM state per block, maintaining constant memory usage.Actions corresponding to each frame condition autoregressive video generation; continuous actions enter through noise-level embeddings, while discrete actions use learned embeddings.

5. Experiments

Experiments evaluate long-term spatial memory on Memory Maze and TECO Minecraft through retrieval and reasoning tasks. The method outperforms sub-quadratic baselines across metrics while maintaining favorable scaling and long-horizon predictions.

  • Datasets: Memory Maze contains 2,000 action-frame pairs per trajectory, while TECO Minecraft contains 1,501 action-observation pairs from randomized gameplay.Ground-truth maze layouts are excluded during Memory Maze training.
  • Evaluation tasks: The retrieval task asks models to reverse a trajectory back to its starting position using observations from the original path.Evaluation uses the Maze dataset because some Minecraft actions, such as jumping, are not invertible.
  • Evaluation tasks: The reasoning task continues a trajectory with random actions and tests whether the model reconstructs observations from previously memorized environment regions.The task assumes sufficient context to commit the entire environment to memory.
  • Quantitative results: Our model outperforms all sub-quadratic baselines across SSIM, LPIPS, and PSNR on retrieval and reasoning, while approaching full-context causal transformers.Tables 2 and 3 report the quantitative comparisons for the 400-frame retrieval and 224-frame reasoning tasks.
  • Long-horizon behavior: Our method maintains high retrieval accuracy across frame distances, unlike limited-context transformers and linear-complexity models that degrade or lack sufficient expressivity.Its retrieval performance remains comparable to a causal transformer trained on the full context.
  • Efficiency and ablations: Training time scales linearly with context length, while inference maintains constant memory and computational costs; each architecture and training component is important for reasoning accuracy.Ablations show that removing the block-wise scan or long-context training adjustment harms long-horizon memory and spatial reasoning.

6. Limitations and Future Work

The method has practical and scope limitations despite constant inference time. Future work targets faster generation, memory beyond the training context, and scaling to high-resolution realistic videos.

  • Limitations and future work: The method does not yet support interactive frame rates despite achieving constant inference time.Timestep distillation is proposed as a possible way to accelerate generation.
  • Limitations and future work: Memory longer than the training context is not handled effectively, although length extrapolation methods for Mamba architectures may extend capacity.This is a stated boundary of the current memory capability.
  • Limitations and future work: Experiments are limited to low-resolution synthetic videos, leaving high-resolution realistic video scaling for future work.The scope boundary is attributed to computational constraints.

S1. Additional Results.

Supplementary results cover maze retrieval and spatial reasoning, plus long-horizon Minecraft generation. Under 240 context and 560 generated frames, the method achieves the lowest FVD, including against causal full-context attention.

  • Additional evaluations: The supplementary results include spatial reasoning and retrieval evaluations on the maze dataset, alongside long-horizon Minecraft videos conditioned on short context.The maze results appear in Supplementary Figures S1 and S2, while Minecraft generations appear in Supplementary Figure S3.
  • Long-term generation: The long-horizon generation evaluation uses 240 context frames and 560 generated frames to compare FVD across the method and relevant baselines.FVD is reported for the resulting 560 generated frames.
  • Long-term generation: The method achieves the lowest FVD, outperforming even a causal transformer trained on the entire context.The comparison is reported in Table S1 for long-term generation.

S2. Implementation Details.

The implementation uses latent video representations, matched 200M-parameter model sizes, staged long-context training, and frame-local attention. Supplementary figures provide additional maze retrieval and reasoning results and Minecraft long-video examples.

  • Latent diffusion: Maze videos use an internal VAE with spatial and temporal compression, while Minecraft uses the same image VAE as DFot for fair comparison.The choice reflects the Maze dataset’s large frame count and Minecraft’s low number of raw frames.
  • Models: All baseline architectures replace attention blocks with their relevant mechanisms while matching parameter counts at 200M per model.Layer counts are adjusted to maintain comparable parameter budgets.
  • Training: Training proceeds on shorter then longer maze videos, with 400-frame and 800-frame stages for reasoning and retrieval, while Minecraft training uses 300 frames.The cited schedule includes 150K plus 250K iterations for maze reasoning and 100K plus 50K steps for maze retrieval.
  • Training: Long-context training samples a random un-noised prefix with probability p = 0.5, requiring the prefix to exceed half the training sequence.When no prefix is sampled, all tokens remain noised, making the procedure equivalent to diffusion forcing.
  • Frame local attention: Frame-local attention uses a window k = 10 and groups frames into chunks of 5, with each chunk attending bidirectionally within itself and to the previous chunk.This design provides an effective frame window of 10 and yields a speedup over a fully causal mask.
  • Additional results: Supplementary Figures S1 and S2 show additional maze reasoning and retrieval results, while Figure S3 shows long Minecraft videos sampled from 150 output frames.Figure S3 uses 25 context frames and 125 random actions, with context frames marked by a red border.
Loading 2505.20171v1…