Video Prediction of Dynamic Physical Simulations With Pixel-Space Spatiotemporal Transformers

TL;DR

This work investigates end-to-end pixel-space video prediction using a pure transformer (PSViT) tailored for physics-driven dynamics. By disentangling spatial and temporal attention within a U-Net–style multiscale backbone and enforcing causal prediction, PSViT extends physically accurate horizons by up to 50% compared with latent-space baselines, while maintaining competitive SSIM/PSNR on standard benchmarks. It also demonstrates interpretability through internal representations and linear probes that recover sequence-specific PDE parameters, and shows that register tokens can capture global dynamics essential for PDE inference. The approach offers a simple, parameter-efficient alternative to latent encoders, with strong generalization to out-of-distribution physical parameters and potential for scalable, physics-aware video modelling.

Abstract

Inspired by the performance and scalability of autoregressive large language models (LLMs), transformer-based models have seen recent success in the visual domain. This study investigates a transformer adaptation for video prediction with a simple end-to-end approach, comparing various spatiotemporal self-attention layouts. Focusing on causal modeling of physical simulations over time; a common shortcoming of existing video-generative approaches, we attempt to isolate spatiotemporal reasoning via physical object tracking metrics and unsupervised training on physical simulation datasets. We introduce a simple yet effective pure transformer model for autoregressive video prediction, utilizing continuous pixel-space representations for video prediction. Without the need for complex training strategies or latent feature-learning components, our approach significantly extends the time horizon for physically accurate predictions by up to 50% when compared with existing latent-space approaches, while maintaining comparable performance on common video quality metrics. In addition, we conduct interpretability experiments to identify network regions that encode information useful to perform accurate estimations of PDE simulation parameters via probing models, and find that this generalizes to the estimation of out-of-distribution simulation parameters. This work serves as a platform for further attention-based spatiotemporal modeling of videos via a simple, parameter efficient, and interpretable approach.

Figures (14)

• Figure 1: Example outputs from our PSViT model trained on video simulation data exhibiting physical dynamics. Successful collision prediction time-step annotated by the red boxes.
• Figure 2: (a) Overview of our PSViT model. (b) Space-time transformer layer. $T$ video frames are patched ($N$ patches), linearly embedded, then processed by local and global space-time transformer blocks. Local blocks use restricted attention windows with patch merge/unmerge (window size 2). Global blocks operate on a smaller effective image resolution. Skip connections between merge/unmerge operations preserve features. Input dimensions are annotated.
• Figure 3: Spatiotemporal self-attention masking strategies for 16-patch video frames, including register tokens (dashed) for sequence-level information. For a query patch $q_{t,s}$ (yellow) at time-step $t$, patch $s$: red patches are for spatial attention, green for temporal attention, and white are unused/masked. Previous frames (processed similarly) and future causally-masked frames are omitted.
• Figure 4: PSViT sample outputs vs. ground truth for Fluid and CLEVRER datasets, conditioned on 12 input frames (intermediate time-steps omitted).
• Figure 5: PSViT sample outputs (representative of median divergence performance) vs. ground truth for Roller, Moon, and 3D Balls datasets, conditioned on 12 input frames (intermediate time-steps omitted).