Title: Diffusion Transformers as Open-World Spatiotemporal Foundation Models
Venue: NeurIPS 2025
Accepted: yes
Keywords: Foundation model

Spatiotemporal learning

Open-world
Abstract: The urban environment is characterized by complex spatio-temporal dynamics arising from diverse human activities and interactions. Effectively modeling these dynamics is essential for understanding and optimizing urban systems. In this work, we introduce UrbanDiT, a foundation model for open-world urban spatio-temporal learning that successfully scales up diffusion transformers in this field. UrbanDiT pioneers a unified model that integrates diverse data sources and types while learning universal spatio-temporal patterns across different cities and scenarios. This allows the model to unify both multi-data and multi-task learning, and effectively support a wide range of spatio-temporal applications. Its key innovation lies in the elaborated prompt learning framework, which adaptively generates both data-driven and task-specific prompts, guiding the model to deliver superior performance across various urban applications. UrbanDiT  offers three advantages: 1)  It unifies diverse data types, such as grid-based and graph-based data, into a sequential format; 2) With task-specific prompts, it supports a wide range of tasks, including bi-directional spatio-temporal prediction, temporal interpolation, spatial extrapolation, and spatio-temporal imputation; and 3) It generalizes effectively to open-world scenarios, with its powerful zero-shot capabilities outperforming nearly all baselines with training data. UrbanDiT sets up a new benchmark for foundation models in the urban spatio-temporal domain. Code and datasets are publicly available at \url{https://github.com/tsinghua-fib-lab/UrbanDiT}.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents UrbanDiT, a foundation model designed for open-world urban spatio-temporal learning. Built on a diffusion transformer backbone, UrbanDiT unifies diverse urban datasets, including both grid-based and graph-based data, into a sequential format and supports a wide range of tasks such as bi-directional prediction, temporal interpolation, spatial extrapolation, and spatio-temporal imputation. Its core innovation lies in a unified prompt learning mechanism, incorporating both data-driven prompts and task-specific prompts. UrbanDiT demonstrates state-of-the-art performance across numerous datasets and tasks, including zero-shot and few-shot generalization.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n1.\tThis paper addresses an important and practical challenge in urban spatio-temporal modeling: the fragmentation of models across tasks and scenarios. By proposing a unified perspective, the work advances the goal of general-purpose urban AI systems.\n\n2.\tThe proposed method presents solid technical innovations. The unified data representation enables seamless integration of grid-based and graph-based data, while the unified prompt learning framework is novel and effectively adapts to heterogeneous urban data and task types.\n\n3.\tThe empirical evaluation is thorough and convincing, covering a wide range of cities, data modalities (grid and graph), and key urban spatio-temporal tasks including prediction, interpolation, extrapolation, and imputation.\n\n4.\tThe model demonstrates strong open-world capabilities, including zero-shot and few-shot generalization, which are crucial for real-world deployment in unseen cities or under limited data availability.\n\n5.\tThe paper is clearly written and well-structured, making the methodology and results easy to understand and follow.\n\n**Weaknesses**\n\n1. The scope of the current model is limited to human-related spatio-temporal data (e.g., mobility, traffic), without incorporating environmental or multi-modal factors such as weather, air quality, or climate variables.\n\n2. While the unified prompt mechanism is effective, its interpretability remains somewhat implicit; it would be beneficial to visualize or analyze how different prompt types contribute across tasks or domains.\n\nQUESTIONS:\n1.\tThe prompt ablation study is informative. Could the authors additionally evaluate prompt effectiveness under few-shot conditions? Specifically, do different types of prompts remain beneficial when training data is scarce?\n\n2.\tSee weakness\n\nLIMITATIONS:\nThe paper briefly mentions that UrbanDiT focuses on human activity data and does not yet model environmental factors.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-06-25"}

comments: SUMMARY:
The paper proposes UrbanDIT, which is a foundation model for urban prediction tasks, especially those dynamic ones such as traffic flows and taxi demand. The model incorporates three major designs, including a unification process of data and tasks, in which all types of data are converted into sequential data, easy to be trained using transformers; unified prompt learning, in which data-driven prompt and task-specific prompt are designed, to enhance the performance of different datasets and tasks; and difussion transformer as backbone network model. 

The model is tested in a comprehensive set of experiments including multiple study areas and datasets, making the effectiveness of the model apparent. In fact, some of the performance lifting is eye-catching, such as table 3.

In an overall manner I am supportive of this paper, a very solid and promising contribution to urban foundation models.

STRENGTHS AND WEAKNESSES:
Strengths: 1) an important application domain of machine learning and foundation models; 2) a smart and novel design of the model using the previously mentioned three major design choices; 3) experiements are credible and comprehensive, showing the very good performance of the model.

Weaknesses: 1) the justifications of many design choices are lacking, such as why unified prompt learning is needed, and some very concrete real-world examples would be helpful, e.g., maybe a key in a memory pool is linked to morning rush hours? 2) the discussion around other inference tasks in urban analysis is lacking, e.g. how this model can benefit the analysis of land use, house price etc.

QUESTIONS:
See the questions in weaknesses.

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 4
ORIGINALITY: 3
DATE: 2025-06-29

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nUrbanDiT pioneers a diffusion transformer (DiT)-based foundation model for unified urban spatio-temporal learning. It addresses the fragmentation in existing approaches by integrating diverse data types into a sequential format via spatio-temporal patching. The core innovation lies in its prompt learning framework: data-driven prompts generated from memory pools guide adaptation to heterogeneous urban data, while task-specific prompts enable a single model to handle five distinct tasks—bidirectional prediction, temporal interpolation, spatial extrapolation, and spatio-temporal imputation.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths：\n1. Unified Architecture: Seamlessly processes grid/graph data and multi-task learning, eliminating task-specific models.\n2. Efficient Inference: 20-step diffusion reduces inference time by 25 times vs. standard DDPM.\n\nWeaknesses:\n1. Stringent Data Requirements: The model mandates both grid and graph-structured inputs for unified training , restricting deployment in tasks lacking either data modality.\n\n2. Limited Spatio-Temporal Data Diversity: Experiments focus on common urban activities, omitting emerging data types like mobile check-ins, bike-sharing patterns, or social media mobility traces, which may reveal finer urban dynamics.\n\n3. Limited technical contributions. The proposed model does not provide much new insights or design ideas. Instead, it is a simple ensemble of existing techniques. For example, the memory pool design largely follows recent works such as UniST.  The model simply uses diffusion transformers as the major building blocks with no modification.\n\nQUESTIONS:\n1. During model training, does a single task require simultaneous use of both graph-based and grid-based data?\n2.  How do embeddings of graph-based and grid-based data achieve semantic alignment in the spatio-temporal domain?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 5, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes UrbanDiT, a foundation model for open-world urban spatio-temporal learning. It unifies diverse data types (grid and graph) into a sequential format and supports multiple tasks through a prompt learning framework combining data-driven and task-specific prompts. Built on diffusion transformers, UrbanDiT handles various flow prediction tasks in a unified way. Extensive experiments across multiple cities and domains show that UrbanDiT achieves state-of-the-art performance and strong zero-shot generalization. The model further demonstrates robust few-shot performance, effective masking strategies for task control, and scalability across model sizes and dataset scales.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1. Unified framework for heterogeneous data: The model handles both grid-based and graph-based spatiotemporal data in a unified sequential format, which allows it to be applied across a variety of urban datasets with different spatial structures.\n\n2. Consistent formulation of diverse tasks: By representing multiple spatiotemporal tasks through masking strategies within a diffusion transformer framework, the model offers a coherent way to address forward/backward prediction, interpolation, extrapolation, and imputation.\n\n3. Comprehensive empirical evaluation: The paper provides extensive experiments across multiple cities and task types. Results show that the model performs competitively compared to a range of strong baselines, including both task-specific and general-purpose methods.\n\nWeaknesses:\n\n1. The model largely builds upon existing ideas such as diffusion transformers and prompt-based conditioning. While the integration is systematic, the contribution is more in engineering design than in fundamental innovation.\n\n2. The use of multiple memory pools (temporal, spatial, frequency) increases model complexity, but the paper provides limited discussion on how to balance effectiveness with interpretability or potential overfitting.\n\n3. Lack of computational analysis: The paper does not report training time, inference latency, or memory usage, which are important for assessing the practicality of scaling the model in real-world urban systems.\n\nQUESTIONS:\nPlease see in Weakness\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}