Research map / Video / Video segmentation and tracking
Diffusion Transformers as Open-World Spatiotemporal Foundation 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.
Introduces UrbanDiT, a diffusion-transformer foundation model trained from scratch on multiple urban spatio-temporal datasets that handles grid and graph data and several tasks (prediction, interpolation, extrapolation, imputation) via unified prompt learning.
Problem
Existing urban spatio-temporal models are largely restricted to specific data types (grid or graph) and mainly to prediction tasks, lacking a single open-world foundation model that unifies diverse data sources, data types, and tasks.
Approach
Scales diffusion transformers for urban spatio-temporal learning by converting grid and graph data into a unified sequence, using masks to specify tasks, and adding data-driven prompts (retrieved from time-, frequency- and space-domain memory pools) and task-specific mask prompts concatenated to the input, trained with rectified flow.
Key results
- On grid-based forward prediction, UrbanDiT achieves the best MAE/RMSE on most datasets with a single unified model (e.g., TaxiBJ MAE 12.61 vs. UniST 14.04), reported as 11.3% relative improvement.
- On spatial extrapolation, UrbanDiT reports TaxiBJ MAE 8.10 / RMSE 12.23 vs. the best baseline CSDI at 36.66 / 75.89 (as reported in the paper's table).
- Zero-shot UrbanDiT is reported to outperform nearly all baselines that had training data on PopSH, per the few-shot/zero-shot figure (values not given in text).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use UrbanDiT instead of the strongest listed baseline | Grid-based forward prediction · TaxiBJ, FlowSH, CrowdNJ, and PopBJ | MAE: 14.04 -> 12.61; 8.77 -> 5.61; 0.094 -> 0.092; 0.078 -> 0.077 | Helped |
| Use UrbanDiT instead of the strongest listed baseline | Grid-based forward prediction · TaxiNYC | MAE: 5.05 -> 5.58 | Hurt |
| Use UrbanDiT instead of the strongest listed baseline | Grid-based forward prediction · TaxiBJ, FlowSH, TaxiNYC, CrowdNJ, and PopBJ | RMSE: 23.67 -> 21.09; 19.95 -> 14.44; 16.37 -> 15.53; 0.168 -> 0.166; 0.136 -> 0.129 | Helped |
| Use UrbanDiT instead of the strongest listed baseline | Spatial extrapolation · TaxiBJ, FlowSH, TaxiNYC, CrowdNJ, and PopBJ | MAE: 36.66 -> 8.10; 15.53 -> 5.44; 19.56 -> 4.91; 0.34 -> 0.099; 0.18 -> 0.084 | Helped |
| Use UrbanDiT instead of the strongest listed baseline | Spatial extrapolation · TaxiBJ, FlowSH, TaxiNYC, CrowdNJ, and PopBJ | RMSE: 75.89 -> 12.23; 34.77 -> 10.17; 49.85 -> 12.52; 0.74 -> 0.155; 0.32 -> 0.146 | Helped |
| Add the frequency-domain prompt (ablation) | UrbanDiT · TaxiBJ | RMSE | Helped |
| Add the time-domain prompt (ablation) | UrbanDiT · TaxiBJ | RMSE | Helped |
| Add the spatial prompt (ablation) | UrbanDiT · TaxiBJ | RMSE | Helped |
| Add the task-specific prompt (ablation) | UrbanDiT · TaxiBJ | RMSE | Helped |
| Add prompt design rather than removing all prompts (ablation) | UrbanDiT · TaxiBJ | RMSE | Helped |
| Use around 20 diffusion inference steps (ablation) | UrbanDiT · TaxiBJ and TaxiNYC | RMSE | Helped |
| Use Rectified Flow with 20 inference steps rather than original DDPM | UrbanDiT · Spatio-temporal generation | efficiency: 25x | Helped |
Surprising findings
- Video prediction models (SimVP, MAU, MIM, TAU) and several deep traffic models perform poorly, suggesting urban dynamics differ from conventional video data.
- Time-series forecasting models (PatchTST, iTransformer, Time-LLM) struggle with urban spatial interactions.
Limitations
- Currently focuses on human activity data such as mobility and traffic, lacking environmental variables like air pollution or climate.
- Baselines are trained on the same multi-dataset mix, but training time is longer than per-dataset baselines.
- Some reported results show inconsistencies (e.g., UrbanDiT not best on all metrics; text claims about zero-shot rest on figures).
Experimental setup
Datasets: TaxiBJ FlowSH TaxiNYC CrowdNJ PopBJ PopSH SpeedSH SpeedBJ SpeedNJ
Benchmarks: TaxiBJ FlowSH TaxiNYC CrowdNJ PopBJ SpeedBJ SpeedSH SpeedNJ
Metrics: MAE RMSE
Baselines: UniST PromptST STID PatchTST iTransformer Time-LLM CSDI SimVP
Builds on: DiT InstaFlow UniST PromptST CSDI
Closest papers by meaning
- Get Rid of Isolation: A Continuous Multi-task Spatio-Temporal Learning Framework (NeurIPS 2024 · similarity 0.91)
- Spatio-Temporal Few-Shot Learning via Diffusive Neural Network Generation (ICLR 2024 · similarity 0.89)
- GeoPro-Net: Learning Interpretable Spatiotemporal Prediction Models Through Statistically-Guided Geo-Prototyping (AAAI 2025 · similarity 0.87)
- Drawing Informative Gradients from Sources: A One-stage Transfer Learning Framework for Cross-city Spatiotemporal Forecasting (AAAI 2025 · similarity 0.86)
- Causal Deciphering and Inpainting in Spatio-Temporal Dynamics via Diffusion Model (NeurIPS 2024 · similarity 0.86)
- Causal Spatio-Temporal Prediction: An Effective and Efficient Multi-Modal Approach (NeurIPS 2025 · similarity 0.85)
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.