Diffusion Transformers as Open-World Spatiotemporal Foundation Models
Authors: Yuan Yuan
Chonghua Han
Jingtao Ding
Guozhen Zhang
Depeng Jin
Yong Li
Venue: NeurIPS 2025 poster ()
Keywords: Foundation model
Spatiotemporal learning
Open-world
Primary area: deep_learning
TL;DR: We propose a foundation model for open-world spatiotemporal learning.

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}.

Diffusion Transformers as Open- World
Spatiotemporal Foundation Models
Yuan Yuan Chonghua Han', Jingtao Ding', Guozhen Zhang', Depeng Jin', Yong Li
Department of Electronic Engineering, BNRist, Tsinghua University
2 TsingRoc.ai
Beijing, China
CCorresponding author: liyong07 @singhua.eduu.c
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 unified model that integrates diverse data sources and types while learning universal spatio-temporal patterns across different cities and scenar- ios. This allows the model to unify both multi-date and multi-task learning, and effectively support a wide range spatio-tempora applications. Its key innova- tion lies in the elaborated prompt learning framework. which adaptively generates both data-driven and task-specific prompts, guiding the model deliver superior performance across various urban applications. UrbanDiT offers three advantages: 1) unifies diverse data types, such as grid-based and graph-based based data, into a sequential format; 2) With task-specific prompts, it supports 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 benchmark for
nearly all baselines with training data. UrbanDiT sets up foundation models in the urban spatio-temporal domain. Code and datasets are publicly available at https: http::////tthuu com/ htmm/tsinghuaaaa fib-laabUrrbnDD1
Introduction
The urban environment is characterized by complex spatio-temporal dynamics arising from diverse human activities and interactions within the city. These dynamics reflected in different types of data. For example, grid-based data divides urban space into regular cells, often used to track crowd flows. In contrast, graph-based data represents spatial structures like road networks nodes and edges, such as traflic speeds on roads. The data from different cities are usually with unique layouts, infrastructures, and planning strategies. Effectively modeling their spatio-temporal dynamics is crucial for optimizing urban services and understanding how cities function. Therefore, it raises an essential research question: can develop foundation model, similar those in natural language processing [1 and computer vision 3 26 I0), that learns universal apatio-tempont patterns and serves general- purpose model for various urban applications?
ceneral-purpose
In the context of urban spatio-temporal modeling, recent advancements such as GPD 53 Ur-
banGPT 23), and UniST [51] have opened exciting avenues for understanding complex urban dynamics. As compared in Table [ these models either utilize LLMSs 23 or develop unified models from scratch [5] (33) tailored for urban spatio-temporal predictions. By training on multiple datasets.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Diverse Urban Spatio-Temporal Data
Multiple tasks
Data Prompt
Tesk Prompt
Crowd
Bi-directional P Prediction
Taxi Demand
Temporal Interpolation
Grid data
Bike Usage
Graph data
Transportation
UrbanDiT
Spatial Extrapolation
Cellular
ST Imputation
Figure 1: A diagram of UrbanDiT:: foundation model that integrates diverse data sources while addressing multiple tasks.
they have showcased impressive generalization capabilities. However, their focus remains largely on prediction tasks, and they are often restricted specific data types-suue such as grid-based data (1] or graph-based traffic data [53] Thus, realizing the full potential of foundation models capable of seamlessly handling diverse data types, sources, and tasks in open-world scenarios remains an open and largely unexplored area of research.
Urban spatio-temporal data is usually defined by varying spatio-temporal resolutions and complex interactions among entities. Building a foundation model requires a scalable architecture capable of accommodating these complexities. Moreover, the intricate nature of urban spatio-tempol dynamics necessitates model that can learn from complex data distributions, Diffusion Transformers, exemplified by models like Sora 3. offer a compelling solution for this purpose. By combining the generative power of diffusion processes with the scalability and flexibility of transformer architectures, diffusion transformers present a promising backbone.
In this work, we introduce UrbanDiT, which unifies training across diverse urban scenarios and tasks, effectively scaling up diffusion transformers for comprehensive urban spatio-temporal learning. offers three appealing benefits: 1) It unifies diverse data types into a sequential format, allowing it to capture ppatio-temporc patterns across various cities and domains. 2) It supports wide range of tasks with a single model via task-specific prompts, without the need for re- re-training across different tasks. 3) It generalizes well open- scenarios, exhibiting powerful zero-shot performance. To build UrbanDiT, we first unify different input data by converting it into the sequential format. We transformer blocks as the denoising network, which are equipped with both temporal and spatial attention modules. To integrate diverse data types and tasks, we propose a unified prompt learning framework. It maintains memory pools to capture learned spatio-temporal patterns and generate data-driven prompts, while also create task-specific prompts for various spatio-temporal tasks. These prompts are concatenated into the unified sequential input before being fed into the transformer modules. The design of prompt learning serves as a flexible intermediary, adaptable to wide range of scenarios.
UrbanDiT, built on the DiT backbone with a prompt learning framework, is a pioneering open- world foundation model. It excels at handling diverse urban spatio-temporal data and a wide range of tasks, including bi-directional spatio-temporal prediction, temporal interpolation, spatial extrapolation, and urban spatio-temporal imputation. This makes UrbanDiT powerful and universal solution for various applications. We summarize our contributions as follows:
To the best of our knowledge, we are the first to explore a foundation model for general-purpose urban spatio-temporal learning, integrating diverse data types and multiple urabn tasks within single unified model.
We present UrbanDiT, an open-world foundation model built on diffusion transformers. Through our proposed prompt learning, UrbanDiT effectively brings together heterogeneous spatio-temporal data and tasks, using data-driven and task-specific prompts to enhance performance.
Extensive experiments demonstrate that UrbanDiT effectively captures complex urban spatio- temporal dynamics, achieving tfate-of-the-art performance across multiple datasets and tasks. also exhibits powerful zero-shot capabilities, proving its applicability in open-world settings. UrbanDiT marks significant step forward in the advancement of urban foundation models.

---

Table 1: Comparison between existing models and UrbanDiT across five aspects.
Method
Model Init.
Data Type
Diverse Sources Task Fiexibilitt
Zemossot
GPD 153
Seratch
Graph
UniST [1
Seratch
Grid
UrbunGPT 23
LLMs
Grid
CityGPT 12
LLMs
Languages
UrbanDi1
Seratch
Graph/Grid
2 Related Work
Urban
Learning. Urban spatio-temporal learning encompasses a variety of tasks
such as prediction [39 22 22 [54) interpolation [[]6], extrapolation [30] 29], and imputation [40
18, addressing how urban systems evolve across space and time. Deep learning has achieved
significant progress in these areas, with techniques ranging from CNNs 22
54 RNNs [43 42
24.. MLPs [36], GNNs 22 [5], and Transformers 7 19], to the more recent use of diffusion
Each of these approaches has been employed model complicated spatio- temporal relationships inherent to urban environments. However, most existing models are tailored to specific datasets and tasks. In contrast, our approach is designed to handle multiple tasks and generalize across diverse urban scenarios.
models 5
40 45
Urban Foundation Models. Foundation models have made significant progress in language mod- els
Recently, researchers have extended the concept of foundation models urban environments, aiming address unique challenges of urban spatio- temporal data. Some representative works in this area include UrbanGPT (23]. UniST (51). and CityGPT 133 UrbanGPT introduces LLMs designed for spatio-temporal predictions within urban contexts. UniST develops a foundation model from scratch specifically for urban prediction tasks, demonstrating zero-shot capabilities that allow the model to generalize new scenarios without ad- ditional training. CityGPT, on the other hand, focuses on enhancing the LLM's ability to comprehend and solve urban tasks by improving of key abilities across existing urban foundation models and UrbanDiT. As shown, UrbanDiT is trained from scratch, allowing it to fully leverage data diversity while offering flexibility across a wide range of tasks. Compared to previous efforts, UrbanDiT represents a significant advancement in developing urban foundation models.
and image generation 3 26 [0].
understanding of urban spaces. Table [ provides a comparison
Diffusion Models for Spatio-Tempores Data. Diffusion models, originally popularized in image generation, have recently gained attention handling spatio-temporal data and time series. They iteratively add and remove noise from data, allowing them to capture complex patterns across both temporal and spatial dimensions 44 52 45 34 In the context of time series, diffusion models
have been applied to tasks such as forecasting 21 34 and imputation [48 40), outperforming traditional methods by generating more accurate and coherent sequences. For spatio-temporal data, diffusion models have proven useful in a variety of tasks, including traffic prediction [45), environmental monitoring [52), and human mobility generation [60 55 By effectively modeling spatio-temporal dependencies, these models can capture both the spatial correlations and temporal dynamics inherent in urban systems. UrbanDiT leverages the generative power of diffusion models to capture complex urban spatio-temporal patterns, while its flexible conditioning mechanisms allow it to address a wide range of spatio-temporal tasks.
3 Method
3.1 Preliminary
Urban Spatio-Temporas Data. Urban spatio-temporal data typically falls into two categories: grid- based and graph-based data. Grid-based data is structured in uniform grid layout. Graph-based data, on the other hand, highlights connectivity, capturing the relationships between various urban entities like streets and intersections. For both different spatial organizations, the temporal dimension is characterized as time series data. The data can be denoted as X NNT, where N denotes the number of spatial partitions. For graph-based data, N corresponds to the number of nodes, while for grid-based data, is defined as the product of the height and width of the grid (N H W). This enables a unified representation of urban spatio-temporal data with different spatial organizations.
3

---

Output
Frequency palber
Reshape
Temporal patiem
ST
Data
Spabal palter
Urban Diffusion Transformer
membrry
Data
Data-driven prompt
Different data types
Sequential format
Alleetica
Concat
Mask
Task
Opatio-tempore patching
Embedding layer
memory spec-fi prompt
(d) Unified Prompt Learning
Noisy ST data
Timestep
Multiple Domains
Multiple Cities
t
(a) Datasets
(b) Urban Diffusion Transformer
(c) Masking Strategies
Figure 2:
Illustration of the whole framework of UrbanDiT, including four key components: a) Unifying different urban spatio-temporal data types; b) The diffusion pipeline of our UrbanDiT; c) Different masking strategies specify different tasks; d) Unified prompt learning with data-driven and task-specific prompts to enhance the denoising process.
Urban Opatio-Tempont Tasks. In addition the commonly recognized (1) forward prediction task, urban spatio-temport analysis encompasses several other critical tasks. (2) Backward Prediction involves estimating past states based current future data. It is essential for understanding historical trends and validating predictive models. (3) Temporal Interpolation aims to estimate values unobserved time points within known temporal range. (4) Spatial Extrapolation involves predicting values beyond the observed spatial domain. (5) Spatio-Tempore Imputation refers to the process of filling in missing values in spatio-temporal datasets.
3.2 UrbanDiT
Figure 2 illustrates the overall framework of UrbanDiT, which based on diffusion transformers.
This framework seamlessly integrates various data types and tasks into cohesive model.
Unification Data and Tasks. We convert data, characterized by three-dimensional structure (2D spatial and ID temporal dimensions), into unified sequential format. For the temporal dimension, we employ patching techniques commonly used in foundational models for time series [32 For grid- based data, we apply 2D patching methods, which are widely utilized in image processing, to organize the data. This allows us to rearrange the dinee-dimesiional grid data into one-dimensional sequential format. For graph-based data, node and integrate it with the temporal dimension to reshape the data into a one-dimensional format well. More details of data unification can be found in Appendix B.1
Graph Convolutional Networks (GCN) (55) to process each
To adapt to various tasks, we employ a unified masking strategy. These tasks can be framed as reconstructing missing parts of the data, with distinct masking strategies tailored to each task. For Forward Prediction, we mask future time steps while utilizing past and present data points to predict the missing values. Conversely, for Backward Prediction, we mask past time steps to estimate historical values based on current and future observations. In the case of temporal interpolation tasks, we apply masks to specific time points within continuous series, allowing the model to fill in these gaps. For spatio-temporal imputation, we randomly mask missing values across both spatial and temporal dimensions, enabling the model to leverage surrounding context for accurate estimations. Finally, in spatial extrapolation tasks, we mask areas outside the observed spatial domain to predict values for unobserved regions based on existing spatial patterns. Consequently, the input of the denoising network X represented as the concatenation of noise features and unmasked spatio-temporal data (conditional observations):
X²= x (1- M) x M

---

where X denotes the noise features, M is the mask that controls the availability of values for downstream tasks, and X represents the clean values of the spatio-temporal data. In this way, we can modulate different masks M to facilitate various urban spatio-temporal applications.
Sequential Input of Spatio Data. We first apply temporal patching process time series data each spatial location, represented X CONV N MxTXD), where T" T and pt the temporal patch size. Next, for grid-based data, we implement 2D spatial patching, resulting HxW xT" In this way, we effectively Ps PPXPxXpi
pt
in Xp = CoNV2D (X (x# xxWxXX XXxxxx'XX xD), where Xp E R L reorganize the data into a format well-suited for transformers.
Spatio-Temporal Transformer Block. The overall model is composed of multiple spatio-temporal transformer blocks. Each block employs both temporal attention and spatial attention, with spatial and temporal attention mechanisms operating independently. This design choice is made to enhance computational efficiency, as the complexity of attention scales with the square of the sequence length.
Diffusion Transformer. We adopt the diffusion transformer model, which integrates a denoising network designed to process complex inputs effectively. The inputs the denoising network consist of three key components: the noisy spatio-temporal data, the timestep, and the prompt. For the timestep
utilize them for layer normalization following previous practices [33 28 which helps stabilize and standardize the input features at each timestep. The prompt, which provides contextual information guidance for the model, is concatenated with the input data to enhance the model's ss understanding of the data and task at hand. This concatenation is straightforward due to the transformer's capability to manage variable sequence lengths, providing flexibility in processing diverse inputs. By incorporating these elements, the diffusion transformer model effectively learns to denoise and generate robust desired results in spatio-temporal contexts.
3.3 Unified Prompt Learning
Data-Driven Prompt.
The data-driven prompt is crucial for training a unified model with multiple and diverse datasets, as such datasets often exhibit significant variations in patterns and distributions. In this context, the prompt acts as a guiding mechanism, helping the model to effectively navigate these differences and generate accurate results. Similar to retrieval augmenttd generation, prompts retrieve the most relevant information, enhancing the model's s ability to contextualize and interpret spatio-temporal data. By aligning the model's 's learning process with the specific characteristics of patterns, data prompts ensure that UrbanDiT can adaptively respond to
various spatio-tempon wide range of urban spatio-temporal scenarios.
M Memory Pool
To achieve this goal. we employ memory net- works, specifically utilizing three memory pools designed to capture the time-domain, frequency- domain and spatial patterns of spatio-temporal dy- namics. For different input data, the prompt net- work retrieves prompts from these memory pools based on the respective time-domain, frequency- domain, and spatial patterns. As shown in Figure 3 each memory pool is structured
Key
Value
Key
Value
Prompe
Key
Value
Cosine similarity
W Weighted sum
key-value store (K,, V.) = {(k), u²), ..., (Re uz =x²),, Vi) {(k)). v/).... (k'), 2 o/)); )} (Kg, V.) = [(k] v.). ... (&& e) where both keys and values are learnable embeddings and randomly initialized. The data-driven prompts are generated follows:
Figure Structure of memory pools.
0, = SOFTMAX K,),
P. Σ a Vv,
0f = SOFTMAX(X K2),
Pf = Σ aa aa..Vz Vi,
a SOFTMAX(X,, Kz),
P. Σ a, Va.
X = CONCAT X),
Task-Specific Prompt. We design task sse-sppeiiiic prompts to enhance the model's performance across different tasks. These prompts are generated from the mask, and we employ attention mechanisms
5

---

to obtain the mask prompt Pm from the mask map as Pm - ATTENTON((FATTENNN))) The learned pattern Pm is then concatenated with the input sequence, resulting in X CONCAT(Pm X). This enables the model to effectively incorporate task-specific information. We provide details of data-driven and task-specific prompts in Appendix B.2
3.4 Training and Inference
The training process alternates between multiple datasets and tasks. In each iteration, we randomly select a dataset and corresponding task to perform gradient descent training. This approach enhances the model's robustness by exposing it to diverse scenarios and helps prevent overfitting by ensuring the model learns from a wide range of inputs and objectives. Let D (D..D2... Dm represent the set of datasets, and T' {Ti- TT....... I denote the set of tasks. Let L(d,, t.) be the loss function for the chosen dataset d, and task &, with the model parameters denoted 0. Overall, the training process is summarized as follows:
For = 1 to N d; Uniform(D),
t Uniform(T)
0 - 0 - VV.d
where is the total number of training iterations and 17 is the learning rate.
For the training of the UrbanDiT model, we adopt a novel diffusion training approach proposed by the InstaFlow (26), which significantly improves the efficiency of employing rectified flow, it is an ordinary differential equation (ODE) -baedd framework that aligns the noise and data distributions through straightened trajectory, as opposed the curved paths often seen in traditional models.
data generation. By
4 Performance Evaluations
Datasets. We utilize diverse set of datasets from multiple domains and cities to evaluate urban spatio-temporal applications, which include taxi demand, cellular network traffic, crowd flows, transportation traffic, and dynamic population, reflecting broad spectrum of urban activities. The datasets are sourced from different cities such New York City, Beijing, Shanghai, and Nanjing, each representing unique urban characteristics. These datasets vary significantly in their spatial structures (e.g., grid or graph formats), the number of locations, and their spatial and temporal resolutions. These variations are influenced by differences city structures, urban planning strategies, and data collection methodologies across regions. For a detailed summary of the datasets, please refer to Table 4 and Table in Appendix A We split the datasets into training, validation, and testing sets along the temporal dimension, using 6:2:2 ratio. To ensure overlap between them, carefully remove any overlapping points, ensuring clear separation across the temporal splits for evaluation.
Baselines. To evaluate the performance of UrbanDiT, we establish a comprehensive benchmark, comparing against state-of-the-an models across different urban tasks. For prediction tasks, we include both traditional time series models such as Historical Average (HA) and ARIMA, as well as advanced deep barning-based spatio-temporal models like STResNet [54]. ACFM (55). STNorm (9) Additionally,
STGSP (58), MC-STL (56), PromptST [57]. STID (36). and UniST [51]
compare
against leading video prediction models, including SimVP 14) TAU
38
MAU (6, and MIM [44),
as well recent time series forecasting approaches such PStchTST 33
iTransformer (27] Time-
LLM [20), and the diffusion- based model CSDI [40.. For graph- based datasets, we evaluate Dib against cutting-edee
apatio-temporal graph models, including STGCN (50). DCRNN 222. GWN (7), MTGNN [46]. AGCRN (2]. GTS [35], and STEP [37]. Furthermore, for spatio-tempore imputation model with state-of-ther offthe-rtt baselines such CSDI, ImputeFormer [3I], Grin [ and BriTS 5. adapting these methods for temporal interpolation and spatial extrapolation
tasks, we compare
tasks. We provide more details of baselines in Appendix C.
4.1 Comparison to the State-of-the-art
Bi-directional Ppatio-Temporal Prediction. For this task, we set both the historical input window and prediction horizon to 12 time steps. Depending on the dataset, the temporal granularity varies-P1 -12 steps may correspond to hour for datasets with 5-minute intervals, 6 hours for datasets with 30- minute intervals, and 12 hours for those with 1-hour intervals. For basclines that cannot handle

---

TaxiBJ
FlowSH
TaxiNYC
CrowdNJ
PopBJ
Model
MAE
RMSE
MAE
RMSE
MAE
RMSE
MAE
RMSE
MAE
RMSE
HA
53.03
91.55
13.43
38.92
26.49
77.10
0,48
0.93
0.232
0.343
ARIMA
57.5
291
9.15
26.70
23.91
99.22
0.443
0.989
0.236
0.404
STResNet
26.55
37.96
45.63
59.82
14.81
26.88
0.511
0.718
0.546
0.751
ACFM
19.87
30.95
24.95
46.92
9.85
20.82
0.284
0.468
0.141
0.200
STNorm
19.00
31.21
11.88
28.46
10.43
26.94
0.231
0.384
0.132
0.198
STGSP
17.54
27.31
17.54
38.77
10.52
25.94
0.263
0.410
0.157
0.229
MC-STL
28.51
38.50
33.83
46.06
26.01
36.75
0.727
0,504
0.235
0,311
MAU
46.37
71.07
21.38
45.04
21.79
49.15
0.402
0.648
0.166
0.256
MIM
42.40
68.18
22.49
47.29
9.151
24.53
0.399
0.715
0.214
0.298
SimVP
21.67
35.58
15.87
28.59
9.08
19.69
0.191
0.282
0.148
0.213
TAU
15.86
26.43
15.22
26.04
9.08
19.46
0.219
0.326
0.135
0.196
PrompiST
16.12
27.42
9.37
23.01
8.24
22.82
0.161
0.306
0.099
0.171
UniST
14.04
23.67
9.10
19.95
5.85
17.55
0.119
0.191
0.106
0.172
STID 16.36
25.55
12.92
21.19
8.32
18.49
0.160
0.234
0.203
0.262
PatchTST
30.55
53.36
10.69
28.17
17.03
50.45
0.223
0.465
0.189
0,291
PatchTST-all
33.62
60.55
12.16
31.79
21.27
58.61
0.403
0.811
0.176
0,279
iTransformer
24.05
42.17
10.19
25.91
45.19
45.19
0.216
0.466
0.154
0,249
Time-LLM
29.55
51.20
10.57
28.19
17.65
52.94
0.210
0,405
0.115
0.195
CSDI 14.76
25.87
8.77
23.37
5.05
16.37
0.094
0.168
0.078
0.136
UrbanDiT 12.61
21.09
5.61
14.44
5.58
15.53
0.092
0.166
0.077
0.129
Table 2: Performance comparison for grid-based forward prediction evaluated using MAE and RMSE. The results are the average prediction errors across all prediction steps. The best result is highlighted in bold, and the second-best is indicated with underlining.
TaxiBJ
FlowSH
TaxiNYC
CrowdNJ
PopBJ
Model
MAE
RMSE
MAE
RMSE
MAE
RMSE
MAE
RMSE
MAE
RMSE
CSDI
36.66
75.89
15.53
34.77
19.56
69.10
0.34
0.18
0.32
Imputeformer
37.13
77.53
17.67
38.96
20.28
49.85
0.39
0.21
0.34
Grin
41.73
92.61
22.56
47.76
22.44
58.15
0.51
0.23
BriTS
59.94
112.34
33.74
59.10
23.39
58.47
0.50
0.70
0.54
0.75
UrbunDiT (ours)
8.10
12.23
5.44
10.17
4.91
12.52
0.099
0.155
0.084
0.146
Table 3: Performance comparison for spatial extrapolation evaluated using MAE and RMSE. The results represent the average errors across different extrapolation steps.
datasets with different shapes, we train individual models for each dataset.For more flexible models like UniST and PatchTST, we train a single unified model across multiple datasets.
Table 2 provides comprehensive benchmark for forward prediction on grid- data. Appendix Table illustrates the results for graph-based data. As observed, traditional deep learning models such STResNet, ACFM, and MC-STL, do not deliver competitive performance. Similarly, video prediction models, such as MAU, MIM, and Si, reveal limitations, suggesting the difference between urban spatio-temporal dynamics and those in conventional video data. UniST demonstrates relatively strong performance, suggesting that training universal model across different datasets holds potential for improving prediction accuracy. However, time-series forecasting models struggled to capture the complex spatial interactions inherent in urban environments, indicating that precisely modeling these interactions is critical for achieving better results urban spatio-tempore prediction. Notably, CSDI ranks second in most cases, showing the effectiveness of diffusion- based models in capturing complex patterns within urban ppaiooltemporal data. Our proposed model, UrbanDiT, delivers the best performance across different datasets using single unified model, achieving relative improvement of 11.3%.
We also compare the backward prediction performance of UrbanDiT with the second-best baseline, CSDI, as shown in Appendix Table Notably, CSDI specifically trained for backward prediction tasks. However, UrbanDiT not only excels in forward prediction but also surpasses specialized

---

UnIST PeompeST
TAU SmmPP
SemVP
0.25
RMSE
RMSE
0.15
0.10
0.10
(a) PopSH 5% few-shot
(b) PopSH- 10% few-shot
Figure 4: Evaluation of UrbanDiT and baseline models in 5% and 1 % few-shot scenarios the PopSH dataset. The red dashed line indicates UrbanDiT's zero-shot performance
models like CSDI in backward prediction by 30.4%. This result demonstrates UrbanDiT's ability to capture complex spatio-temporal patterns more effectively.
Temporal Interpolation. We set the missing ratio to 0.5, meaning that we only know the even- numbered time steps (e.g.. 4, 2n). and the model is required predict the odd numbered time steps (e.g.. 1, 3, 5, ... 2n-1), Appendix Table Table@deend demonstrates that UrbanDiT, employing unified model, outperforms baselines trained separately for different datasets most cases.
Spatial Extrapolation. We evaluate the models' ability to predict missing values in specific spatial regions by masking 50% of of spatial locations across the temporal sequence. The objective is to determine how effectively models extrapolate unobserved spatial information from the remaining visible data. As shown Table 3 UibanDiT achieves the best performance most cases.
Spatio-Temporal Imputation. This task assesses the models' capacity to impute missing values across both spatial and temporal dimensions. We randomly mask 50% of positions in the 3D spatio- temporal data, simulating real-world scenarios where urban data may be incomplete due to sensor failures irregularities in data collection. As shown in Appendix Table TO UrbanDiT achieves the best performance in most cases.
These results substantiate that UrbanDiT consistently delivers superior performance across diverse tasks and datasets using a single, unified model. This capability positions UrbanDiT as general- purpose foundation model, enabling practitioners to leverage optimized parameters directly, thereby simplifying deployment and enhancing applicability urban spatio-temporal applications.
4.2 Few-shot and Zero-shot Performance
key strength foundation models their ability to generalize easily. Therefore, we perform experiments in both few-shot and zero-shot scenarios, testing its adaptability to new datasets with little or no additional training. In the few-shot scenario, we train UrbanDiT on small portion of the target dataset- speecffiiyy using only 5% and 10% of the available data-and then evaluate its performance on the corresponding test set. This setup challenges the model to generalize well from sparse data. In the zero-shot scenario, data from the target dataset is provided for training. Instead, we directly evaluate UrbanDiT's performance on the target dataset, relying solely on its pretrained knowledge handle unseen data without any fine-tuning.
Figure 4 demonstrates the few-shot and zero-shot performance of UrbanDiT in comparison to baseline models. In the few-shot setting (with 5% and 1% of the training data), UrbanDiT consistently outperforms baselines, showing strong ability to learn from minimal data. Even more striking, in the zero-shot scenario, UrbanDiT exhibits exceptional inference capabilities, surpassing nearly all baseline models that had access to training data. This highlights its generalization ability without fine-tuning, reinforcing its effectivvness
open-world foundation model.
4.3 Ablation Studies
Prompt. Unified prompt learning is key design in UrbanDiT. To investigate the contribution of each prompt the final performance, we conduct ablation studies by systematically removing each type of prompt. Specifically, we identify four types of prompts: F for frequency-domain prompt,

---

T for time-domain prompt, S for spatial prompt, and M for task-specific prompt. We denote the removal of a prompt as w/o MM and indicate the absence of any prompt as w/o P.
Figure 2 presents the results of ablation studies. The findings reveal that removing any single prompt sig- nificantly degrades the model's performance. In the absence of prompt design altogether, the model ex- hibits the poorest performance. Among the four types of prompts, the removal of the Trequency-doman prompt has the most pronounced negative impact on the overall performance.
Full
w/o
w/o
w/o
7ZZ
w/o 5
F
w/o
25
24
RMSE
23
Inference Steps of Diffusion Models. We further investigate the effect of inference steps the perfor- mance of diffusion models. The number of inference steps a critical factor in balancing the model's accu- racy and efficiency. Appendix Figure 8 iillustrates the performance of the diffusion model across different numbers of inference steps for two datasets, TaxiBJ and TaxiNYC, measured using RMSE. Notably, we observe that around 20 inference steps provide the optimal balance between computational efficiency and model performance for both datasets. By setting the diffusion steps to 500 and the inference steps to 20, we achieve a 25x improvement in efficiency compared the original DDPM 177, without sacrificing accuracy.
22
21
Figure 5: Ablation study on the prompt design using RMSE on the TaxiBJ dataset.
4.4 Scalability
As a foundation model, it is crucial to under- stand how model performance evolves as the datasize scale varies across different model sizes. This information is valuable for practitioners to train and fine-tune the foundation model ef- fectively. In Figure 6 we explore the rela- tionship between model performance and data- size scale for three model sizes: UrbanDiT-S (small), UrbanDiT-M (medium), and UrbanDiT- (large). As observed, all three models demonstrate improved performance as the data size increases. However, when the dataset size increases from 0.8 to the large model, UrbanDiT-L, shows a notably steeper improve- ment (with slope of 0.011), compared to the medium (slope of 0.0015) and small models (slope of 0.0019). This pronounced scaling ef- fect for the large model indicates its potential to further enhance performance as more data becomes available. These results highlight the promising scalability of UrbanDiT-L, suggesting that it can handle larger datasets and achieve even better outcomes with increased data size.
Model Performance Datasize Scale for Different Model Sizes
0.080
UrbanDiT-S
UrbanDiT-M
Model Performance
0.075
UrbanDiT-L
0.070
0.065
0.060
0.055
0.2
0.4
0.6
0.8
1.0
Datasize Scale
Figure 6: The scalability of UrbanDiT.
Conclusion
In this paper, we present UrbanDiT, an open-world foundation model built on diffusion transformers and a unified prompt learning framework. UrbanDiT enables seamless adaptation wide range of urban ppatio-temporal tasks across diverse datasets from urban environments. Our extensive experiments highlight the model's exceptional potential in advancing the field of urban spatio- temporal modeling. We believe this work not only pushes the boundaries of urban apatio-temport modeling but also serves as an inspiration for future research the rapidly evolving field of foundation models. Nonetheless, UrbanDiT currently focuses on human activity data such as mobility and traffic. To support holistic urban modeling, future work should incorporate environmental variables like air pollution, climate indicators, and microclimate dynamics.

---

Acknowledgements
This work was supported by the National Natural Science Foundation of China (No. U24B20180, 62476152). We would also like to express our sincere gratitude to Yu Zheng for the insightful discussions, which were invaluable to this research.
References
[1] Pierre Aumond, Arnaud Can, Vivien Mallet, Bert De Coensel, Carlos Ribeiro, Dick Bottel- dooren, and Catherine Lavandier. Kriging-based spatial interpolation from measurements for sound level mapping urban areas. The journal of the acoustical society of America, 143(5):2847-2857, 2018.
[2] Lei Bai, Lina Yao, Can Li, Xianzhi Wang, and Can Wang. Adaptive graph convolutional recurrent network for traffic forecasting. Advances in neural information processing systems, 33:17804-1788 17815, 2020.
[3] Tim Brooks, Bill Peebles, Connor Holmes, Will DePue, Yufei Guo, Li Jing, David Schnurr, Joe Taylor, Troy Luhman, Eric Luhman, et al. Video generation models world simulators. 2024. URL https://openai. con/researcchvidee sewprearrrrviiide-ggeentttoooooo inneeaationmmodell aw-waasdwwwiii worlddsimulators, 3, 2024.
[4] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[5] Wei Cao, Dong Wang, Jian Li, Hao Zhou, Lei Li, and Yitan Li. Brits: Bidirectional recurrent imputation for time series. Advances in neural information processing systems, 31, 2018
[6] Zheng Chang, Xinfeng Zhang, Shanshe Wang, Siwei Ma, Yan Ye, Xiang Xinguang, and Wen Gao. Mau: A motion-aware unit for video prediction and beyond. Advances in Neural Information Processing Systems, 34:26950-26962, 2021.
[7] Changlu Chen, Yanbin Liu, Ling Chen, and Chengqi Zhang. Bidirectional spatial-temporal adaptive transformer for urban traffic flow forecasting. IEEE Transactions on Neural Networks and Learning Systems, 2022.
[8] Andrea Cini, Ivan Marisca, and Cesare Alippi. Filling the g_ap_s: Multivariate time series imputation by graph neural networks. In International Conference Learning Representations, 2022.
[9] Jinliang Deng, Xiusi Chen, Renhe Jiang, Xuan Song, and Ivor W Tsang. St-norm: Spatial and temporal normalization for malti-variate time series forecasting. In Proceedings of the 27th ACM SIGKDD conference knowledge discovery & data mining, pages 269-278, 2021.
[10] Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transform- ers for high-resolution image synthesis. In Forty-first International Conference on Machine Learning, 2024.
[11] Yuchen Fang, Hao Miao, Yuxuan Liang, Liwei Deng, Yue Cui, Ximu Zeng. Yuyang Xia, Yan Zhao, Torben Bach Pedersen, Christian S Jensen, et al. Unraveling spatio-temporal foundation models via the pipeline lens: A comprehensive review. arXiv preprint 1XXx:25506001644 2025.
[12] Jie Feng, Yuwei Du, Tianhui Liu, Siqi Guo, Yuming Lin, and Yong Li. Citygpt: Empowering urban spatial cognition of large language models. arXiv preprint arXiv:2406. 13948, 2024.
[13] Jie Feng, Yuwei Du, Tianhui Liu, Siqi Guo, Yuming Lin, and Yong Li. Citygpt: Empowering urban spatial cognition of large language models. arXiv preprint arXiv:2406. 13948, 2024.
[14] Zhangyang Gao, Cheng Tan, Lirong Wu, and Stan Z. Li. Simvp: Simpler yet better video prediction. In Proceedings of the JEFF/CVE Conference on Computer Vision and Pattern Recognition, pages 3170-3180, 2022.
10

---

[15] Xu Geng, Yaguang Li. Leye Wang, Lingyu Zhang, Qiang Yang, Jieping Ye, and Yan Liu. Spatiotemporal multi-graph convolution network for ride-hailing demand forecasting, In Proceedings of the AAAI conference artificial intelligence, volume 33, pages 3656-3663, 2019.
[16] Benedikt Gräler, Edzer J Pebesma, and Gerard BM Heuvelink. Spatio-temporal interpolation using gstat. R J., 8(1):204, 2016.
[17] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
[18] Junfeng Hu, Xu Liu, Zhencheng Fan, Yuxuan Liang, and Roger Zimmermann. Towards unifying diffusion models for probabilistic spatio-temporal graph learning. arXiv preprint arXiv:2310. 17360, 2023.
[19] Jiawei Jiang, Chengkai Han, Wayne Xin Zhao, and Jingyuan Wang. Pdformer: Propaga- tion delay-aware dynamic long-range transformer for traffic flow prediction. arXiv preprint arXiv:2301.0799 2023.
[20] Ming Jin, Shiyu Wang, Lintao Ma, Zhixuan Chu, James Y Zhang, Xiaoming Shi, Pin- Yu Chen, Yuxuan Liang, Yuan-Fang Li, Shirui Pan, et al. Time-Ilm: Time series forecasting by reprogramming large language models. In The Twelfth International Conference on Learning Representations.
[21] Marcel Kollovich, Abdul Fatir Ansari, Michael Wang, and Yuyang Bernie Wang. Predict, refine, synthesize: Self-guiding diffusion models for probabilistic time series forecasting. Advances in Neural Information Processing Systems, 36, 2024.
Jasper Zschiegner, Hao
[22] Yaguang Li, Rose Yu, Cyrus Shahabi, and Yan Liu. Diffusion convolutional recurrent neural net- work: Data-driven traffic forecasting. In International Conference on Learning Representations, 2018.
[23] Zhonghang Li, Lianghao Xia, Jiabin Tang, Yong Xu, Lei Shi, Long Xia, Dawei Yin, and Chao Huang. Urbangpt: Spatio-temporal large language models, 2024.
[24] Zhihui Lin, Maomao Li, Zhuobin Zheng, Yangyang Cheng, and Chun Yuan. Self-attention convistm for spatiotemporal prediction. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 11531-11538 2020.
[25] Lingbo Liu, Ruimao Zhang, Jiefeng Peng, Guanbin Li, Bowen Du and Liang Lin. Attentive crowd flow machines. In Proceedings of the 26th ACM international conference on Multimedia,
pages 1553-1566 1561, 2018.
[26] Xingchao Liu, Xiwen Zhang, Jianzhu Ma, Jian Peng, et al. Instaflow: One step is enough for high-quality diffusion-based text-to-image generation. The Twelfth International Conference on Learning Representations, 2023.
[27] Yong Liu, Tengge Hu, Haoran Zhang, Haixu Wu, Shiyu Wang, Lintao Ma, and Mingsheng Long. itransformer: Inverted transformers are effective for time series forecasting. arXiv preprint 20..0:2310.06625, 2023.
[28] Haoyu Lu, Guoxing Yang, Nanyi Fei, Yuqi Huo, Zhiwu Lu, Ping Luo, and Mingyu Ding. Vdt: Gereral-purpose video diffusion transformers via mask modeling. In The Twelfth International Conference on Learning Representations.
[29) Jun Ma, Yuexiong Ding, Jack CP Cheng, Feifeng Jiang, and Zhiwei Wan. A temporal-spatial interpolation and extrapolation method based geographic long short-term memory neural network for pm2.5 5. Journal of Cleaner Production, 237:1 17729, 2019.
[30] James R Miller, Monica G Tumner, Erica AH Smithwick, C Lisa Dent, and Emily H Stanley. Spatial extrapolation: the science of predicting ecological patterns and processes. BioScience, 54(4):310-300, 2004.
11

---

[31] Tong Nie, Guoyang Qin, Wei Ma, Yuewen Mei, and Jian Sun. Imputeformer: Low rankness- induced transformers for generalizable spatiotemporal imputation. In Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 2260-2271, 2024.
[32] Yuqi Nie, Nam H Nguyen, Phanwadee Sinthong, and Jayant Kalagnanam. A time series is worth 64 words: Long-term forecasting with transformers. arXiv preprint arXiv:2211.14730. 2022.
[33] William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 4195-4205, 2023.
[34] Kashif Rasul, Calvin Seward, Ingmar Schuster, and Roland Vollgraf. Autoregressive denois- ing diffusion models for multivariate probabilistic time series forecasting. In International Conference on Machine Learning, pages 8857-8868. PMLR, 2021.
[35] Chao Shang and Jie Chen. Discrete graph structure learning for forecasting multiple time series. In Proceedings of International Conference on Learning Representations, 2021.
[36] Zezhi Shao, Zhao Zhang, Fei Wang, Wei Wei, and Yongjun Xu. Spatial-tempora identity: simple yet effective baseline for multivariate time series forecasting. In Proceedings of the 31 ACM International Conference on Information Knowledge Management, pages 4454-4458, 2022.
[37] Zezhi Shao, Zhao Zhang, Fei Wang, and Yongjun Xu. Pre-training enhanced spatial-tempolt graph neural network for multivariate time series forecasting. In Proceedings the 28th ACM SIGKDD conference on knowledge discovery and data mining, pages 1567-1577, 2022.
[38] Cheng Tan. Zhangyang Gao, Lirong Wu, Yongjie Xu, Jun Xia. Siyuan Li, and Stan Z Li. Temporal attention unit: Towards efficient spatiotemporal predictive learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 18770-18782, 2023.
[39] Cheng Tan, Siyuan Li, Zhangyang Gao, Wenfei Guan, Zedong Wang. Zicheng Liu, Lirong Wu, and Stan Z Li. Openstl: A comprehensive benchmark of Advances in Neural Information Processing Systems, 36:69819 1118899668 9-69931, 2023.
predictive learning.
[40] Yusuke Tashiro, Jiaming Song, Yang Song, and Stefano Ermon. Csdi: Conditional score-based diffusion models for probabilistic time series imputation. Advances in Neural Information Processing Systems, 01224844248111 2021.
[41] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.0922 2023.
[42] Yunbo Wang, Zhifeng Gao, Mingsheng Long, Jianmin Wang, and S Yu Philip. Predmn++: Towards a resolution of the deep-in-time dilemma in spatiotemporal predictive learning. In International Conference on Machine Learning, pages 5123-5132. PMLR, 2018.
[43] Yunbo Wang, Mingsheng Long, Jianmin Wang, Zhifeng Gao, and Philip Yu. Predrnn: Recurrent neural networks for predictive learning using spatiotemporal Istms. Advances in neural information processing systems, 30, 2017.
[44] Yunbo Wang. Jianjin Zhang, Hongyu Zhu, Mingsheng Long, Jianmin Wang, and Philip S Yu. Memory in memory: A predictive neural network for learning higher-order non-stationarity from spatiotemporal dynamics. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 91 9154-9162. 2019.
[45] Haomin Wen, Youfang Lin, Yutong Xia, Huaiyu Wan, Qingsong Wen, Roger Zimmermann, and Yuxuan Liang. Diffstg: Probabilistic spatio-temporal graph forecasting with denoising diffusion models. In Proceedings of the 31 ACM International Conference on Advances Geographic Information Systems, pages 1-12, 2023.
12

---

[46] Zonghan Wu, Shirui Pan, Guodong Long, Jing Jiang, Xiaojun Chang, and Chengqi Zhang. Connecting the dots: Multivariate time series forecasting with graph neural networks. In Proceedings of the 26th ACM SIGKDD international conference on knowledge discovery & data mining, pages 753-763, 2020.
[47] Zonghan Wu, Shirui Pan, Guodong Long. Jing Jiang, and Chengqi Zhang. Graph wavenet for deep Conference on Artificial Intelligence, pages 1907-1913, 2019.
graph modeling. In Proceedings of the 28th International Joint
[48] Chunjing Xiao, Zehua Gou, Wenxin Tai, Kunpeng Zhang, and Fan Zhou. Imputation-based time- series anomaly detection with conditional weight-incremental diffusion models. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 2742-2751, 2023.
[49] Yiyuan Yang, Ming Jin, Haomin Wen, Chaoli Zhang, Yuxuan Liang, Lintao Ma, Yi Wang. Chenghao Liu, Bin Yang, Zenglin Xu, et al. A survey on diffusion models for time series and spatio-temporal data, arXiv preprint arXiv:2404. 18886, 2024.
[50] Bing Yu, Haoteng Yin, and Zhanxing Zhu, Spatio-temporal graph convolutional networks: a deep learning framework for traffic forecasting. In Proceedings of the 27th International Joint Conference on Artificial Intelligence, pages 3634-3640, 2018.
[51] Yuan Yuan, Jingtao Ding, Jie Feng, Depeng Jin, and Yong Li. Unist: a prompt-empowered universal model for urban spatio-temporal prediction. In Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 4095-4106, 2024.
[52] Yuan Yuan, Jingtao Ding, Chenyang Shao, Depeng Jin, and Yong Li. Spatio-temporal diffusion point processes. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 3173-33 1773-3184 3184, 2023.
[53] Yuan Yuan, Chenyang Shao, Jingtao Ding, Depeng Jin, and Yong Li. Spatio-temporal few-shot learning via diffusive neural network generation. In The Twelfth International Conference on Learning Representations, 2024.
[54] Junbo Zhang, Yu Zheng, and Dekang Qi. Deep spatio-tempord. residual networks for citywide crowd flows prediction. In Proceedings of the AAAI conference on artificial intelligence, volume 31, 2017.
[55] Si Zhang, Hanghang Tong, Jiejun Xu, and Ross Maciejewski. Graph convolutional networks: comprehensive review. Computational Social Networks, 6(1):1-23, 2019.
[56] Xu Zhang, Yongshun Gong, Xinxin Zhang, Xiaoming Wu, Chengqi Zhang, and Xiangjun Dong. Mask-and -nhaanedd spatio-temporal learning for urban flow prediction. In Proceedings of the 32nd ACM International Conference on Information and Knowledge Management, pages 2298-3307,2 2023.
[57] Zijian Zhang, Xiangyu Zhao, Qidong Liu, Chunxu Zhang, Qian Ma, Wanyu Wang, Hongwei Zhao, Yiqi Wang, and Zitao Liu. Promptst: Prompt-enhanced tpatioltemporal multi-attriben prediction. In Proceedings of the 32nd ACM International Conference on Information and Knowledge Management, pages
2023.
[58] Liang Zhao, Min Gao, and Zongwei Wang. St-gsp: Spatial-temporal global semantic represen- tation learning for urban flow prediction. In Proceedings of the Fifteenth ACM International Conference on Web Search and Data Mining, pages 1443-1451, 2022.
[59] Yuanshao Zhu, Yongchao Ye, Shiyao Zhang, Xiangyu Zhao, and James Yu. Difftraj: Generating gps trajectory with diffusion probabilistic model. Advances in Neural Information Processing Systems, 36:65168 65188, 2023.
[60] Yuanshao Zhu, James Jianqiao Yu, Xiangyu Zhao, Qidong Liu, Yongchao Ye, Wei Chen, Zijian Zhang, Xuetao Wei, and Yuxuan Liang. Controltraj: Controlabble trajectory generation with diffusion model. In Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 4676-4681 2024.
13

---

Table 4: Basic statistics of grid-based data.
Dataset
City
Type
Temporal Period
Spatial partition
Interval
Mean
Std
FlowS11
Shanghai
Mobility flow
2016/04/25 2016/05/01
20 x 20
15min
31.935
137.926
PupBJ
Beijing
Crowd flow
2021/10/25
2021/11/21
28 24
One hour
0.367
0.411
TaxiBI
Beijing
Taxi flow
201 310300001- 2013/10/30
32 32
Half an hour
97.543
122.174
CrowdNJ
Nanjing
Crowd flow
2021/2202-
2021/03/01
20 29
One hour
0.872
1.345
TaxiNYC
New York City
Taxi flow
2015/01/01 - -2015/03/01
10 20
Half an hour 38.801
103.924
PopSH
Shanghai
Dynamic population
2014X08/01
20144/0828
32 28
One hour
0.175
0.212
Table 5: Basic statistics of Graph-based data.
Dataset
City
Type
Temporal Period
Interval
#Nodes
WEdges
Mean
Std
SpeedSH
Shanghai
Traffic speed
2022/01/27 2022/02/27
15min
21099
39065
7,815
4.044
SpeedBJ
Beijing
Traffic speed
2022/03/05 2022/04/05
15min
13675
24444
6.837
3.412
SpeedNJ
Nanjing
Traffic speed
2022/03/05 2022/04/05
15min
13419
25100
6.699
4.253
A
Datasets
We provide detailed overview of the datasets utilized in our study to support future research in the field urban spatio-temporal modeling. The datasets are categorized into two distinct types: grid- based and graph-based spatio-temporal data. Each type of data reflects different spatial organizations and dynamics, enabling a comprehensive evaluation of model performance across varied urban scenarios.
Grid-based data represent spatial information a structured, uniform grid layout, where each grid cell corresponds to specific geographical area. Table Totineo the essential details and statistics for grid-based datasets, including spatial resolution, temporal resolution, temporal period, and size of each dataset.
Graph-based data, on the other hand, capture urban spatial relationships through network of nodes and edges, where nodes typically represent points of interest (e.g.. intersections or key locations), and edges represent the connections between them (e.g., roads or transit lines). This type of data well- suited for modeling scenarios that involve irregular spatial structures, such as transportation networks. Table providess comprehensive summary of the graph- based datasets, including information on the number of nodes, edges, temporal resolution, temporal period, and dataset size.
B Methodology Details
B.1 Sequential Format of Input Data
We provide detailed description of the data unification process for both grid-based and graph-based ppatio-temporal data. The key goal is to transform the data into unified sequential format suitable for the transformer's s input.
Grid-based data is structured in uniform grid layout, typically represented in three-dimensions form Xgrid TR. xx7xxxW with two spatial dimensions (height H and width W) and one temporal dimension T. To process this data, we utilize 3D Convolutional Neural Networks (3D CNN), which are widely used for capturing both spatial and temporal dependencies in spatio-temporal tasks. The process is formulated as follows:
X' CONv3D kernel size (Pe. Psy Po. Ps))
X, RESHAPE [N]
where N = T x x represents the total number of spatio-temporal partitions, effectively converting the data into a one-dimensional sequence for further processing by the transformer model.
Graph-based data is inherently non-uuclidean, capturing relationships between urban entities (e.g., streets and intersections). The spatial dimension represented by a graph structure with nodes and edges, and the temporal dimension is still captured as time series each node. The graph- based
14

---

data can be represented as tensor X ggaaph E RN*7 where N is the number of nodes in the graph, and the number of time steps. To handle the temporal dimension, first apply ID convolutional network (ID CNN) along the time axis to capture local temporal dependencies. Next, to capture spatial relationships, we apply Graph Convolutional Network (GCN) [55] on the graph structure. For each temporal patch, the GCN aggregates information from neighboring nodes using the graph s adjacency matrix A E igNxa X Finally, we reshape the graph-based data into a sequential format. The operations are formulated as follows:
X' CONV gr kernel size = Pt)
X' GCN(X", W)
X, RESHAPE(X', [MM]
where M represents the number of spatio-temporal patches, ensuring that the graph-based data is transformed into one-dimensional sequence, similar to the grid-based data. This unified sequential representation allows both data types to be processed consistently by the transformer model.
B.2 Unified Prompt Learning
We provide details how to obtain the data-driven and task task-specific prompts.
Time-domain patterns. Suppose the patched ppatioltemporal data is denoted as X T'''', where T" T. and N' H x W. we extract time-domain patterns by applying an attention mechanism p. along the temporal dimension. This is done independently for each spatial location, allowing us to capture temporal dependencies across different spatial patches follows:
X, = ArTENTION(X*), N'XT X RNXIID
where D is the embedding size.
Frequency-domen -omain patterns. In our work, employ four distinct approaches to compute features in the frequency domain, depending on the configuration of the Fast Fourier Transform (FFT) and thresholding mechanisms:
Without FFT Threshold: we directly compute the FFT of the input tensor. The tensor permuted along the appropriate dimensions, and the real and imaginary components of the FFT are concate- nated along the last dimension. This results in frequency domain representation of the data. is formulated follows:
XEFT FFT(X),
Xfreq = [9R(XFFT), ((FFr))
where R Xprr) represents the real part of the FFT, and 3(Xrr)) represents the imaginary part.
Basic FFT Threshold: we apply a basic threshold technique by computing the amplitude of the FFT and creating binary mask. The mask retains frequency components whose amplitude is greater than the mean amplitude, filtering out oowffrquency noise and preserving significant frequency components. The process formulated as follow:
X = FFT(X),
A | PA HXW T ΣA.
1
M T(A MA), XFTTAAherdd = X OM,
X freq [R FFFFF.T 3 ((FFT),
Quantile-based FFT Threshold: We further refine the frequency selection by applying threshold
based on 80t% of the amplitude distribution. This approach retains the most prominent frequency components, allowing for more flexible filtering compared to the mean-based threshold.
15

---

The selection process can be formulated as follows:
XFFT FF)(X),
A | qso (A,,
M - (IA (so), XFFT,teeeed - X @ M,
freg ((FF..ttt)ee)) 3(XFFT,itera))]
Top-k Frequency Filtering: We retain only the top k frequency components (e.g., the first three). We generate mask to preserve only these dominant components, filtering out the rest. It is formulated follows:
X FFT(X), A= |
indices = argsort( A, descending) kj, M mask XFFT.ileredd XFFr 3 M,
Fry (R((FFTLiteaa)) 3 3(XFFTiieeaa))
Spatial patterns. For the same patched spatio-temporal data e TT''''' we extract spatial patterns by applying an attention mechanism along the spatial dimension, independently on each temporal patch. This process allows us model spatial dependencies within each time patch as follows:
X, ATTENTION(X)) XX), T' X, E XXxxD
C Experiment Details
C.1 Baselines
HA: History Average is forecasting method that predicts future values by calculating the mean of historical data from the same time periods.
This model utilizes the difference in data between consecutive recurring states to address non-stationary characteristics. By stacking multiple MIM blocks, it can capture higher- order non-stationarity in the data.
MIM 4K
The Motion-aware Unit extends the temporal scope prediction units seize correla- tions in motion between frames. It encompasses an attention mechanism and a fusion mechanism, which are integral to video prediction tasks.
MAU 6
Sim VP A simple yet effective video prediction model is entirely based on convolutional neural networks and employs MSE loss as its performance metric, providing reliable benchmark for comparative studies video prediction.
TAU [8]
The Temporal Attention Module breaks down temporal attention into two parts: within- frame and between-frames, and employs differential divergence regularization manage variations across frames.
STResNet [54]: STResNet employs residual neural networks to detect proximity, periodicity, and trends in the temporal data.
ACFM 2 The Attentive Crowd Flow Machine model forecasts crowd movements by using an attention mechanism to dynamically integrate sequential and cyclical patterns.
STGSP (58): This model highlights the significance of global and positional temporal data for spatio-temporal forecasting. It incorporates a semantic flow encoder to capture temporal position cues and an attention mechanism to handle multi-scale temporal interactions.
MC-STL [56]: MC-STL utilizes mask-enhanced contrastive learning to efficiently identify spatio- temporal relationships.
STNorm 99: It introduces two distinct normalization modules: spatial normalization for handling high-frequency elements and temporal normalization for managing local components.
STID 36): This MI.P-based spatio-temporal forecasting model discerns subtleties within the spatial and temporal axes, showcasing its design's ' efficiency and efficacy.
16

---

PromptST [57| An advanced pre- vre-training and prompt-tuning methodology tailored for spatio- temporal forecasting.
UniST 51: : A versatile urban spatio-temporal prediction model that grid-based data. It
employs various spatio-temporal masking techniques for pre-training and fine-tuning with spatio- temporal knowledge-basee prompts.
The Spatio- nentioal Graph Convolutional Network is deep learning architecture for predicting traffic patterns, harnessing both spatial and temporal correlations. integrates graph convolutional operations with convolutional sequence learning capture multi-scale dynamics within traffic networks.
STGCN [50]:
Graph WaveNet is technique crafted to overcome the shortcomings of current spatial- temporal graph modeling methods. It introduces self-adjusting adjacency matrix and utilizes stacked dilated causal convolutions to efficiently capture temporal relationships.
GWN [47]:
MTGNN 46 MTGNN is framework tailored for multivariate time series analysis. au- tonomously identifies directional relationships between variables via a graph learning component and incorporates additional information such as variable attributes.
GTS an approach that concurrently learns the topology of graph alongside Graph Neural Network (GNN) for predicting multiple time series. It models the graph structure using neural network, allowing for the generation of distinct graph samples, and aims to optimize the average performance across the distribution of graphs.
GTS [35]
DCRNN 22 The Diffusion Convolutional Recurrent Neural Network is a deep learning frame- work for spatiotemporal prediction. It treats traffic flow a diffusion phenomenon directed graph, securing spatial interdependencies via two-way random walks and temporal interdependen- cies through encoder-decoder setup with scheduled sampling.
STEP [37] Spa:tiaaal--emppnne Graph Neural Network Enhanced by Pre-training is framework that uses a pre- trained model to enhance ppatial-tempore graph neural networks for better forecasting of multivariate time series data.
The AGCRN framework improves upon Graph Convolutional Networks by incorpo- rating two adaptive components: Node Adaptive Parameter Learning and Data Adaptive Graph Generation. This approach effectively captures nuanced spatial and temporal relationships within traffic data, functioning independently of pre-set graph structures.
AGCRN (2]:
It employs patching and supe-super learning techniques for forecasting multi- variate time series. By dividing the time series into segments, it captures long-term dependencies and analyzes each data channel separately using unified network architecture.
PStchTST [3]:
iTransformer [27] This state-of-the-ar model for multivariate time series utilizes attention mechanisms and feed-forward neural network layers on inverted dimensions to emphasize the relationships among multiple variables.
Time-LLM [00 : TIME-LLM represents advanced approach in applying large-scale language models to time series prediction. It employs a reprogramming strategy that adapts LLMs for forecasting tasks without altering the underlying language model architecture.
CSDI 40]: CSDI is explicitly trained for imputation and can exploit correlations between observed values, leading to significant improvements in performance over existing probabilistic imputation methods.
It introduces low-rank inductive bias into the Transformer framework to balance strong inductive priors with high model expressivity, making it suitable for wide range of imputation tasks.
Imputeformer (31]
Grin 8 : GRIN introduces a novel graph neural network architecture designed to reconstruct
missing data in different channels of a multivariate time series, outperforming state-of-the-art methods in imputation tasks.
BriTS 55: BRITS is method for imputing missing values time series data, utilizing bidirec- tional recurrent neural network (RNN) without imposing assumptions on the data's underlying dynamics.
is worth noting that the baselines, including UniST (51 and PatchTST can also be trained using multiple datasets. In our comparison experiments, we train these models in unified manner using
17

---

the same diverse datasets to ensure fair comparison. This approach ensures that the performance gains of UrbanDiT are not merely due dataset diversity, but reflect the model's ' true advantage.
C.2 Experiment Configuration
For UrbanDiT-S (small), the model consists of 4 transformer layers with a hidden size of 256. Both the spatial and temporal patch sizes are set and the number of attention heads UrbanDiT-M (medium) is composed of 6 transformer layers with hidden size of 384, maintaining the same spatial and temporal patch sizes of and 6 attention heads. UrbanDiT-L (large) includes 12 transformer layers, a hidden size of 384, spatial and temporal patch sizes of 2. and 12 attention heads. Each memory pool contains 512 embeddings, with the embedding dimension matching the model's hidden size. The learning rate is set to le-4, and the maximum number of training epochs is 500, with early stopping applied to prevent overfitting. The batch size is tailored for each dataset to maintain similar number of training iterations across them.
C.3 Metrics.
To assess the performance of Urb in urban spatio-temporal applications, we employ widely recognized evaluation metries: Root Mean Square Error (RMSE) and Mcan Absolute Error (MAE). Given that UrbanDiT operates probabilistic model, we conduct 20 inference runs and use the average result for comparison against the ground truth. We apply the same evaluation framework to the probabilistic baselines, ensuring a consistent and fair assessment of all models.
D Additional Results
D.1 Results of Multiple Tasks
Table Table illustraat additional results of multiple tasks.
D.2
Few-Shof and Zero-Shot Performance
Figure demonstrates UrbanDiT's few-shot and zero-shot capabilities on the TaxiBJ dataset.
D.3
Ablation Studies
D.4
Computational Analysis
table provides an overview of model efficiency in terms overall training time and inference time. While the training time of UrbanDiT is longer than that of the baseline models due to its inclusion of multiple datasets, is important to note that training separate models for each dataset and summing the total training time results in comparable times between UrbanDiT and the baseline methods. Furthermore, UrbanDiT achieves the best performance across all datasets with single, unified model, demonstrating its efficiency and effectiveness in delivering superior results without the need for multiple specialized models. This efficiency is crucial for real-world applications, where scalability key.
Regarding inference latency, UrbanDiT incurs slightly higher costs due to its diffusion- genera- tive framework, which involves iterative sampling in the denoising process and multiple sampling for probabilistic prediction. However, with Rectified Flow acceleration, inference significantly faster, notably outperforming CSDI. The resulting latency is reasonable and practically negligible given the substantial performance gains and unified deployment benefits of the model.
18

---

UtbanDIT
PacchTST
UrbanDiT
PatchTST
TAU
Zero-shot
UniST
TAU
Zero-shot
PromptST
SimvP
PromposT
SimVP
70
60
60
RMSE
RMSE
40
30
30
20
(a) TaxiB] 5% few-shot
(b) TaxiBJ 10% few-shot
Figure 7: Evaluation of UrbanDiT and baseline models in 5% and 1% few-shot scenarios on the TaxiBJ dataset. The red dashed line indicates UrbanDiT's zero-shot performance
25
21
24
20
TaxiNYC
TaxiB
19
18
10
20
50
100
# #inference Steps
Figure 8: Performance evaluation (RMSE) with varying numbers of inference steps on TaxiBJ and TaxiNYC datasets.
TaxiBJ
FlowSH
TaxiNYC
CrowdNJ
PopBJ
Model
MAE
RMSE
MAE
RMSE
MAE
RMSE
MAE
RMSE
MAE
RMSE
CSDI
17.40
33.98
10.65
31.88
4.83
15.43
0.094
0.16
0.082
0.14
UrbanDiT
11.57
20.08
5.996
14.37
4.71
15.07
0.16
0.099
0.071
0.117
Table 6: Performance comparison for grid-based backward prediction evaluated using MAE and RMSE.
19

---

SpeedBJ
SpeedSH
N]
Model
MAE
RMSE
MAE
RMSE
MAE
RMSE
HA
1.35
2.13
0.92
1.46
1.94
3.01
STGCN
1.81
2,44
0.99
1.35
1.63
2.31
CRNN
1.37
1.98
0.89
1.28
1.53
2.38
GWN
1.69
2.32
0.93
1.32
1.50
2.16
MTGNN
1.15
1.70
0.86
1.33
1.57
2.42
AGCRN
1.66
2.29
1.14
1.56
1.77
2.46
GTS
1.76
2,36
1.31
1.74
2.04
2.68
STEP
1.45
2.04
0.93
1.32
1.58
2.42
STID
1.08
1.69
0.83
1.26
1.56
2.38
PatchTST
1.27
1.99
0.87
1.37
1.83
2.74
PatchTST
1.55
2.44
1.08
1.70
2.19
3.34
iTransformer
1.26
1.97
0.90
1.40
1.70
2.62
Time-LLM
1.28
2.00
0.87
1.36
1.82
2.76
UrbanDiT 1.02
1.66
0.78
1.20
1.51
2.30
Table 7: Performance comparison of prediction across three traph-based traffic speed datasets.
TaxiBJ
FlowSH
TaxiNYC
CrowdNJ
PopBJ
Model
MAE
RMSE
MAE
RMSE
MAE
RMSE
MAE
RMSE
MAE
RMSE
CSDI
11.20
18.42
5.71
13.14
3.86
11.59
0.055
0.092
0.044
0.077
Imputeformer
11.99
19.83
6.72
15.69
5.61
16.72
0.079
0.16
0.066
Grin
13.69
23.45
9.61
8.10
21.32
0.10
0.18
0.083
0.16
BriTS
17.57
15.24
19.41
50.25
0.19
0.16
0.25
UrbanDiT (ours)
9.09
14.54
4.90
10.308
4.50
11.46
0.077
0.121
0.056
0.094
Table 8:
Performance comparison for temporal interpolation evaluated using MAE and RMSE. The results represent the average errors across different interpolation steps.
TaxiBJ
FlowSH
TaxiNYC
CrowdNJ
PopBJ
Model
MAE
RMSE
MAE
RMSE
MAE
RMSE
MAE
RMSE
MAE
RMSE
CSDI
12.29
22.07
7.94
21.86
4.33
13.09
0.071
0.12
0.055
0.094
Imputeformer
13.65
23.18
9.22
19.97
5.95
16.36
0.093
0.16
0.069
0.12
Grin
16.83
27.61
9.70
23.52
9.15
21.43
0.16
0.30
0.096
0.18
BriTS
22.57
38.39
17.14
38.82
19.93
50.47
0.26
0.41
0.18
0.29
UrbanDiT (ours)
9.38
15.19
5.03
11.52
4.62
12.16
0.083
0.13
0.061
0.101
Table 9: results represent the average errors across different imputation steps.
Performance comparison for temporal imputation evaluated using MAE and RMSE. The
TaxiBJ
FlowSH
TaxiNYC
CrowdNJ
PopBJ
Model
MAE
RMSE
MAE
RMSE
MAE
RMSE
MAE
RMSE
MAE
RMSE
CSDI
7.92
12.42
4.28
8.62
3.86
11.54
0.057
0.091
0.046
0.083
Imputeformer
9.70
13.80
5.50
10.30
4.79
15.35
0.076
0.12
0.061
0.11
Grin
11.96
19.62
9.21
19.68
9.62
20.77
0.11
0.19
0.080
0.14
BhTS
13.99
23.53
17.95
38.57
19.17
50.15
0.21
0.44
0.13
0.19
UrbanDiT (ours)
7.83
12.13
5.07
9.79
3.63
11.44
0.057
0.090
0.049
0.092
Performance comparison for grid- based spatio-temporal imputation evaluated using MAE and RMSE. The results represent the average prediction errors across different prediction steps.
Table 10:
20

---

Table 11: Model Training and Inference Times
Model
Train Time All
Inference Time
STGCN
17 min
2
DCRNN
77 min
8 s
GWN
16 min
1 s
MTGNN
14 min
0.8
AGCRN
min
2
GTS
126 min
17
STEP
177 min
27 s
STResNet
5.7 min
0.6
ACTM
56 min
0.9
STNorm
46 min
5
STGCP
min
4
MC-STL
31 min
7
MAU
82 min
13
MIM
84 min
TAU
22 min
6
PromptST
45 min
9
Imputeformer
28 min
6s
BriTS
82 min
10
Grin
17 min
2
UniST
5
s
STID
10 min
5
PatchTST
33 min
5
iTransformer
23 min
6
Time-LLM
(multiple datasets)
min
CSDI
5.5 h
38 min
UrbanDiT
h (multiple datasets)
57
21

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: See Abstract and Section
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how be expected to generalize other settings.
much the results
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: See Section S
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
22

---

Justification:
The paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We release all the code and data, as well as instructions for how to replicate the results. See abstract and Section 4
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
23

---

Answer: [Yes]
Justification:
We have submitted code and data anonymously as supplementary materials.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https guwdes/ddeeeeee Cooeeeoodeesssseeeeee //nips.com htte://alllllgiiiee///oeeeuaaass ssiniiiiccypyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification:
We provide sufficient information on experimental setting. See Section 4
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We report the the statistical significance of the experiments suitably and correctly. See Section[4 [
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide sufficient information on the computer resources. See Appendix D.4.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
We make sure that the presented research conforms with the NeurIPS Code of Ethics.
Justification:
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We provide thorough discussion about broader impacts of this work. See Section
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
25

---

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: All assets used in the paper are properly credited. The license and terms of use are explicitly mentioned and properly respected.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the Datasets.papersstt paperswi thcode. cww/ datasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
26

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: All new assets introduced in the paper are well documented and we provide the documentation alongside the assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Question:
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
27

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [Yes]
Justification: We provide sufficient information on the usage of LLMs.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.LM policy http:p/seeeaarrsssec////o/eeeeee for what should or should not be described.
28


Reviews:
[{"review_id": "fEIUKZ285h", "reviewer": "Reviewer_DoyW", "summary": "The 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.", "questions": "Please see in Weakness", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\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.", "was_revised": true, "final_justification": "I think my concerns have been addressed. What I appreciate is that the paper tackles an important and meaningful problem. Therefore, I will maintain my positive score.", "rebuttal": "Dear Reviewer DoyW,\n\nWe appreciate the positive feedback regarding the strengths of our approach, particularly the unified framework for heterogeneous data, the consistent formulation of diverse tasks, and the comprehensive empirical evaluation across multiple cities and domains.  We have carefully addressed the concerns raised in the review, and the following is a point-by-point response to each of them.\n\n**W1. 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\n**Response.** We respectfully argue that our contributions go beyond engineering design and offer meaningful technical innovations for the spatio-temporal modeling community.\n\n\n**First, UrbanDiT tackles a critical and underexplored challenge: building a true “one-for-all” urban spatio-temporal foundation model.** To the best of our knowledge, it is the first model to unify diverse data formats while spanning multiple application domains and task types, all within a single model trained from scratch. Unlike previous universal or foundation models, which mainly focus on prediction tasks and are often limited to grid-based spatio-temporal data, UrbanDiT achieves a level of generalization across formats, tasks, scenarios, and regions that has not been realized by existing approaches. We believe this represents a significant step forward in the development of open-world spatio-temporal modeling.\n\n**Second, we introduce key methodological innovations to the spatio-temporal domain.** The DiT+Prompt structure is a general-purpose framework that has been successfully applied in various fields [1,2,3]. Our contribution lies in building a foundational model based on this general architecture, but not simply applying it directly. Instead, we make novel designs to address the unique challenges of spatio-temporal modeling. Our data unification strategy transforms diverse data formats into a unified representation, enabling UrbanDiT to seamlessly process different types of urban data. Additionally, we introduce a unified prompt learning framework that generates both data-driven and task-specific prompts, allowing joint learning across heterogeneous data and tasks without conflicts. Unlike NLP or CV, urban spatiotemporal modeling faces greater challenges due to the highly distinct patterns across different scenarios and cities. The key innovation of our framework is ensuring that these diverse data sources benefit from each other during joint training, making UrbanDiT highly adaptable for a wide range of urban tasks.\n\n\n\n**Last, UrbanDiT exhibits emergent and scalable generalization abilities.** As shown in our experiments, UrbanDiT demonstrates emergent zero-shot generalization across cities and domains. These capabilities indicate that UrbanDiT captures underlying spatio-temporal regularities beyond specific tasks, which we view as a step toward a universal spatio-temporal foundation model.\n\n\n\n[1] Timedit: General-purpose diffusion transformers for time series foundation model. arXiv preprint.\n[2] Text2PDE: Latent Diffusion Models for Accessible Physics Simulation, ICLR 2025\n[3] Diffusion on Language Model Encodings for Protein Sequence Generation, ICML 2025\n\n**W2. 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\n\n**Response.** We appreciate the reviewer’s concern about the potential trade-off between model complexity, interpretability, and the risk of overfitting when using multiple memory pools in UrbanDiT. We address these points as follows:\n\n**First, each memory pool is designed to capture distinct, interpretable factors** that reflect the core elements of urban spatio-temporal dynamics. These include temporal patterns (e.g., daily/weekly cycles), spatial correlations (e.g., neighborhood effects), and frequency-based latent periodicities.  This modular approach improves model interpretability and enables targeted analysis of each component's contribution. As shown in our ablation studies (Fig. 6, Section 4.3), **each memory pool uniquely contributes to the model's performance**. Removing any pool leads to significant degradation, indicating that each captures valuable, non-redundant information.\n\nTo enhance interpretability, we will include visualizations in the final version ($\\underline{\\text{Section 4.5}}$) showing the spatio-temporal patterns stored in the memory pools. Since we cannot provide images in this response, we summarize our findings here. From the initial to the final state, the embeddings in the memory pools diverge, reflecting the storage of increasingly personalized information. The memorized patterns show consistency across urban scenarios, while attention weights vary across datasets, indicating the model dynamically adjusts its focus based on input data characteristics.\n\n**Second, regarding overfitting, UrbanDiT mitigates this risk through its robust generalization across domains, cities, and tasks**. Its training paradigm alternates between datasets and tasks, acting as a natural regularizer and preventing overfitting to any single dataset. Strong zero-shot and few-shot performance (Section 4.3, Fig. 5) demonstrates its ability to generalize to unseen cities, outperforming models trained on target-domain data. This empirical evidence supports the claim that UrbanDiT captures robust, transferable patterns rather than memorizing dataset-specific details.\n\n\n**W3. 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\n**Response.** Thanks for this valuable suggestion. We have now included a computational analysis. The following table provides an overview of model efficiency in terms of overall training time and inference time.\n\nWhile the training time of UrbanDiT is longer than that of the baseline models due to its inclusion of multiple datasets, it is important to note that training separate models for each dataset and summing the total training time results in comparable times between UrbanDiT and the baseline methods. Furthermore, UrbanDiT achieves the best performance across all datasets with a single, unified model, demonstrating its efficiency and effectiveness in delivering superior results without the need for multiple specialized models. This efficiency is crucial for real-world applications, where scalability is key.\n\nRegarding inference latency, UrbanDiT incurs slightly higher costs due to its diffusion-based generative framework, which involves iterative sampling in the denoising process and multiple sampling for probabilistic prediction.  However, with Rectified Flow acceleration, inference is significantly faster, notably outperforming CSDI. The resulting latency is reasonable and practically negligible given the substantial performance gains and unified deployment benefits of the model.\n\n\n| Model        | Train Time All | Inference Time | \n|--------------|----------------|----------------|\n| STGCN        | 17 min         | 2 s            |\n| DCRNN        | 77 min         | 8 s            |\n| GWN          | 16 min         | 1 s            |\n| MTGNN        | 14 min         | 0.8 s          |\n| AGCRN        | 21 min         | 2 s            |\n| GTS          | 126 min        | 17 s           |\n| STEP         | 177 min        | 27 s           |\n| STResNet     | 5.7 min        | 0.6 s          |\n| ACTM         | 56 min         | 0.9 s          |\n| STNorm       | 46 min         | 5 s            |\n| STGCP        | 8 min          | 4 s            |\n| MC-STL       | 31 min         | 7 s            |\n| MAU          | 82 min         | 13 s           |\n| MIM          | 84 min         | 14 s           |\n| TAU          | 22 min         | 6 s            |\n| PromptST     | 45 min         | 9 s            |\n| Imputeformer | 28 min         | 6 s            |\n| BriTS        | 82 min         | 10 s           |\n| Grin         | 17 min         | 2 s            |\n| UniST        | 5 h            | 19 s           |\n| STID         | 10 min         | 5 s            |\n| PatchTST     | 33 min         | 5 s            |\n| iTransformer | 23 min         | 6 s            |\n| Time-LLM     | 6 h (multiple datasets) | 5 min |\n| CSDI         | 5.5 h          | 38 min         |\n| UrbanDiT      | 4 h (multiple datasets) | 57 s |\n\n\nWe will include this analysis in the $\\underline{\\text{Appendix D}}$  in the revised version to provide a comprehensive view of UrbanDiT’s scalability and real-world practicality.\n\nThank you again for your constructive feedback. We are happy to address any remaining questions. \n\n\nSincerely,\n\nThe Authors"}, {"review_id": "RLJLub38Y5", "reviewer": "Reviewer_EsKL", "summary": "UrbanDiT 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.", "questions": "1. 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?", "limitations": "Yes", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths：\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.", "was_revised": true, "final_justification": "+ The authors addressed specific technical questions regarding the design and workflow of the proposed model in the rebuttal phase. \n+ Overall, the paper has some contributions, especially in terms of building a scalable system. \n- The paper's general idea and overall framework is quite similar to existing work namely UniST. The new contributions appear to be somewhat incremental for a new paper in NeurIPS. \n- The authors emphasized in the rebuttal the enhanced capabilities of their proposed model such as scalability and generalizability. However, the technical innovations behind such achievements are inadequately explained, leading to the suspicion that the advancement is mostly the result of engineering tweaks. \n\nI would like to keep the score, but would not be upset if the paper is accepted.", "rebuttal": "Dear Reviewer EsKL,\n\nWe sincerely thank you for your thoughtful review and constructive feedback. We appreciate your recognition of the strengths of our model. In response to the concerns you raised, we have addressed the issues regarding data requirements,  data diversity, and the technical contributions of UrbanDiT. Below is our detailed response to each of your comments.\n\n**W1. Stringent Data Requirements: The model mandates both grid and graph-structured inputs for unified training , restricting deployment in tasks lacking either data modality.**\n\n**Response.** We appreciate the reviewer’s comment. However, we would like to clarify that **UrbanDiT does not require the simultaneous use of both grid and graph-structured inputs**. Instead, it is designed to handle either grid or graph data independently, providing flexibility for practitioners to select suitable data type based on the specific task.\n\nThis flexibility is, in fact, one of our key contributions. The data unification strategy enables the model to seamlessly process diverse data types, whether grid-based  or graph-based. This addresses the concern of restricting deployment in tasks lacking either data modality, as UrbanDiT can function effectively with either type of data.\n\nWe will clarify this point in $\\underline{\\text{Section 3.2 (\"Unification of Data and Tasks\")}}$ in the revised version to prevent potential misunderstandings.\n\n\n**W2. 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\n**Response.** We appreciate the reviewer’s valuable suggestion. In response, we have incorporated additional urban scenarios, such as bike-sharing data (BikeNYC) and mobile check-in data (MobileSH), into our experiments. These additions contribute to enhancing data diversity and provide a more comprehensive view of urban spatiotemporal dynamics. The following tables show the updated results, which will be incorporated into $\\underline{\\text{Table 2}}$ in the final version of our paper.\n\n\n| Model        | MobileSH MAE | MobileSH RMSE | BikeNYC MAE | BikeNYC RMSE |\n|--------------|--------------|---------------|-------------|--------------|\n| HA           | 0.100        | 0.165         | 7.17        | 15.68        |\n| ARIMA        | 0.112        | 0.197         | 13.63       | 25.01        |\n| PatchTST     | 0.062        | 0.108         | 5.30        | 12.33        |\n| iTransformer | 0.045        | 0.072         | 4.50        | 9.86         |\n| STResNet     | 0.102        | 0.138         | 3.94        | 7.18         |\n| ATFM         | 0.055        | 0.078         | 3.09        | 5.99         |\n| STNorm       | 0.042        | 0.065         | 3.03        | 6.47         |\n| STGSP        | 0.040        | 0.057         | 7.38        | 14.20        |\n| TAU          | 0.044        | 0.063         | 2.89        | 5.98         |\n| PromptST     | 0.043        | 0.069         | 2.76        | 5.84         |\n| MAU          | 0.081        | 0.125         | 2.95        | 6.12         |\n| MIM          | 0.079        | 0.126         | 2.89        | 6.36         |\n| STID         | 0.040        | 0.062         | 2.71        | 5.70         |\n| UniST        | 0.039        | 0.055         | 2.56        | 5.50         |\n| CSDI         | 0.045        | 0.072         | 2.26        | 5.22         |\n| UrbanDiT     | 0.034        | 0.052         | 2.04        | 5.01         |\n\n\n\nRegarding individual mobility traces from social media, we currently do not support this data type due to the irregular time intervals and discrete location tokens. We recognize the potential value of incorporating such data in the future and are actively exploring ways to unify discrete and numerical data types within our framework. However, the absence of this feature does not hinder the model’s ability to effectively capture and analyze the numerical spatio-temporal dynamics of urban systems that are central to our focus.\n\n\n\n\n**W3. 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\n**Response.**  We respectfully disagree with the reviewer’s characterization of UrbanDiT as lacking novel insights or design ideas. We highlight the following key innovations that distinguish our work from UniST and represent significant advancements in spatio-temporal modeling.\n\n**First, UrbanDiT addresses a critical gap in spatio-temporal modeling  by creating a true “one-for-all” foundation model.** Unlike UniST, which is primarily focused on prediction tasks and limited to grid-based data, UrbanDiT is the first model to unify spatio-temporal modeling across four key dimensions: data format, city, domain, and task, within a single model trained from scratch. This unification enables UrbanDiT to handle diverse scenarios in open-world settings, offering greater flexibility and scalability compared to UniST.\n\n\n**Second, we introduce key methodological innovations to the spatio-temporal domain.** While UrbanDiT builds on general-purpose frameworks like the diffusion transformer architecture, our contribution lies in developing a foundational model based on this architecture, not simply applying it directly. We introduce novel designs to address the unique challenges of spatio-temporal modeling. Our data unification strategy transforms diverse data formats into a unified representation, enabling UrbanDiT to seamlessly process different types of urban data. Additionally, we propose a unified prompt learning framework that generates both data-driven and task-specific prompts, allowing joint learning across heterogeneous data and tasks without conflicts. Unlike NLP or CV, urban spatiotemporal modeling faces greater challenges due to the highly distinct patterns across different scenarios and cities. **The key innovation of our framework is ensuring that these diverse scenarios and tasks complement and benefit from each other during the joint training.**\n\n\n\n**Lastly, an obvious distinctive feature of UrbanDiT is its performance gain over UniST.**  UrbanDiT is capable of capturing broader spatio-temporal patterns across different data formats, cities, domains, and tasks, making it a more versatile foundation model. \n\n\n\nIn summary, UrbanDiT introduces substantial methodological advancements that set it apart from UniST. Our focus on cross-format, cross-task, and cross-region generalization represents a significant step forward in spatio-temporal modeling. We believe these contributions establish UrbanDiT as a more flexible, scalable, and generalizable model.\n\n\n**Q1. During model training, does a single task require simultaneous use of both graph-based and grid-based data?**\n\n\n**Response.**  To clarify, UrbanDiT does not require the simultaneous use of both graph-based and grid-based data as inputs for a single task.  Instead, UrbanDiT is designed to handle either graph-based or grid-based data independently, depending on the task at hand. This flexibility is a key feature of the model, allowing UrbanDiT to adapt to diverse urban spatio-temporal tasks without requiring both data types simultaneously.\n\n**Q2. How do embeddings of graph-based and grid-based data achieve semantic alignment in the spatio-temporal domain?**\n\n**Response.** Although UrbanDiT does not require the simultaneous input of both graph-based and grid-based data, it supports flexible integration of both, which inherently implies a form of semantic alignment. \n\nBefore explaining the alignment, we would like to clarify that graph and grid data have different spatial organizations. Graph-based data represents spatial relationships in a network structure, while grid-based data organizes spatial data in a regular, structured grid. Despite these differences, both share key spatial characteristics, such as neighborhood relationships. In the temporal dimension, alignment is more natural, as both data types share common temporal features, like temporal closeness and periodicity.\n\nTo achieve semantic alignment, we generate spatially-related and temporally-related prompts that align the underlying semantics of these data types. Specifically, for different input samples, the model adaptively retrieves the most relevant prompt, ensuring that the alignment between spatial and temporal semantics is maintained.  This adaptive retrieval process facilitates a unified learning process across heterogeneous data.\n\nThank you once again for your thoughtful and constructive feedback. We believe the additional experiments and clarifications regarding the technical contributions address your primary concerns. Should you find our response and the updated results satisfactory, we would be grateful if you would consider raising your rating. We are happy to address any further questions or provide additional clarifications as needed.\n\nSincerely,\n\nThe Authors"}, {"review_id": "xcP8NViRux", "reviewer": "Reviewer_XMSo", "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. \n\nThe 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.\n\nIn an overall manner I am supportive of this paper, a very solid and promising contribution to urban foundation models.", "questions": "See the questions in weaknesses.", "limitations": "yes", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 3, "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.\n\nWeaknesses: 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.", "was_revised": true, "final_justification": "I've read the discussion and I am aware of different opinions from other reviewers. I recognise that this is a solid contribution for an important sub-domain of foundation models. Therefore I would like to keep my rating.", "rebuttal": "Dear Reviewer XMSo,\n\nWe also greatly appreciate your constructive suggestions, which will help enhance the clarity and depth of our manuscript. In response, we have addressed the points raise, providing clearer justifications and a broader discussion on the model’s applicability. Below is our detailed response to your comments.\n\n**W1. 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?**\n\n**Response.** We appreciate the reviewer’s comment. Regarding the justification for unified prompt learning, we emphasize the following key points.\n\n- Unified prompt learning is essential because it faciliates UrbanDiT to seamlessly handle diverse spatio-temporal data and tasks within a single model. This flexibility is particularly important in urban settings, where cities, scenarios, and tasks can vary significantly. Direct joint training across such heterogeneous data could lead to contradictions and inefficiencies. The prompts serve as crucial signals, guiding the model’s learning process. They allow the model to adapt to various tasks and data types, making it possible for these diverse elements to complement and benefit from each other. This is achieved through semantic alignment, which is facilitated by decomposing the data into its spatio-temporal dimensions.\n- As for the example of the memory pools, we have conducted an analysis where different memory keys store varying spatio-temporal patterns. Our findings show that memory units effectively capture specific patterns, such as traffic flow during peak hours. In the revised version, we will visualize these results. While we cannot provide images here, we summarize our findings as follows: the memory embeddings evolve during optimization, progressively storing more personalized information. Importantly, the memorized patterns demonstrate consistency across different urban scenarios. However, attention weights vary across datasets, indicating that the model dynamically adjusts its focus based on the specific characteristics of the input data.\n\n\nWe thank you for your suggestion to improve the clarity and explainability of our approach. We will include a more detailed analysis of the technical design ($\\underline{\\text{Section 3.3}}$) and provide an in-depth exploration of the prompt learning framework ($\\underline{\\text{Section 4.5}}$) in the revised version.\n\n\n**W2. 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.**\n\n**Response.** We appreciate your suggestion and agree that a broader discussion on potential inference tasks in urban analysis would be valuable. UrbanDiT focuses on modeling urban spatio-temporal dynamics, which are closely tied to human daily activities. This capability offers significant benefits for urban analysis tasks such as land use prediction and house price forecasting.\n\nMobility patterns are crucial indicators of how land is being used and will evolve. By modeling the movement of people across different time scales and regions, UrbanDiT can identify areas experiencing increasing commercial or residential activity, as well as shifts in land use.  Moreover, urban spatio-temporal dynamics play a key role in influencing house prices. For example, areas with improved transportation connectivity often experience changes in housing demand, which directly impacts property values. UrbanDiT provides a predictive advantage in identifying emerging trends in housing markets.\n\nIn summary, UrbanDiT’s ability to model urban spatio-temporal dynamics enables it to generate actionable insights for a variety of applications, offering concrete benefits in urban planning and real estate analysis. We will incorporate this discussion into the revised manuscript, particularly in the $\\underline{\\text{discussion and conclusion}}$ section.\n\nThank you again for your constructive suggestions. We are happy to address any remaining questions. \n\nSincerely,\n\nThe Authors"}, {"review_id": "SS5CW3dJy9", "reviewer": "Reviewer_NF7R", "summary": "This 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.", "questions": "1.\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", "limitations": "The paper briefly mentions that UrbanDiT focuses on human activity data and does not yet model environmental factors.", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 4, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "My concerns have been resolved, and I recommend accepting the paper.", "rebuttal": "Dear Reviewer NF7R,\n\nWe sincerely thank you for the constructive and insightful feedback. We are encouraged that you acknowledged the technical contributions of our work and its thorough experimental validation. We would like to address the concerns you raised as follows.\n\n**W1. 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\n**Response.** We appreciate the reviewer’s comment. UrbanDiT focuses on urban spatio-temporal dynamics, which are inherently tied to human activities. Urban systems are largely shaped by human-driven behaviors, such as mobility and traffic, making them central to the scope of our model. UrbanDiT is designed to capture the spatio-temporal patterns most directly influenced by human actions within urban environments.\n\nThe environmental and climate factors you mentioned are part of broader, more external spatio-temporal systems, often referred to as natural systems. These systems are indeed important, and incorporating them is a natural direction for future research. However, for this study, we have focused primarily on human-related data, as it is more directly tied to the dynamics of urban systems.\n\nWe did attempt to integrate climate spatio-temporal data with urban data in our experiments, but the results did not yield significant improvements. This outcome makes sense, as the dynamics governing urban and natural systems (e.g., climate) are fundamentally distinct. Nevertheless, we recognize the value of exploring the coupling between these two domains in future work, and this remains an important area of ongoing research.\n\n\n**W2. 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\n**Response.** Thanks for your valuable suggestion! To enhance the interpretability, we will include visualizations in the final version to show the spatio-temporal patterns stored in the memory pool ($\\underline{\\text{Section 4.5}}$). Specifically, we will provide the following:\n\n- We will show the embeddings of the memory units both at the initial state and after optimization. These visualizations will reveal how the model’s memory evolves during training, allowing us to understand the process of how the memory pools adapt over time.\n- We will also visualize the specific spatio-temporal patterns memorized by each unit.\n\nSince we are unable to provide images in this response, we summarize our findings as follows. From the start state to the final optimized state, the embeddings gradually diverge in different directions. This suggests that, throughout the optimization process, the memory pools progressively store and encapsulate more personalized information. Additionally, we find that the memorized patterns revealed through the prompt tool exhibit notable consistency across different urban scenarios. However, the attention weight distributions associated with the memory units for different datasets manifest significant dissimilarities. This indicates that the model’s focus is dynamic and responsive, adjusting according to the specific characteristics of the input data.\n\n\n**Q1. The 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\n**Response.** Thank you for this important question. Investigating why the model performs well in few-shot and even zero-shot conditions is crucial for understanding its robustness. To address this, we evaluated the effectiveness of prompts under few-shot settings and visualized the patterns memorized by each unit in these conditions. We found that, even with limited or no training data, the model is capable of extracting useful prompts that are highly relevant to the input samples, which effectively guide the learning process. We will include these findings and visualizations in the revised version of the paper ($\\underline{\\text{Section 4.5}}$) to further highlight this key aspect of the model's performance.\n\nThank you again for your positive ratings and constructive feedback. We are happy to address any remaining questions.\n\nSincerely,\n\nThe Authors"}]

Author rebuttal:

Meta-review: This paper proposes UrbanDiT, a diffusion transformer–based foundation model for open-world urban spatio-temporal learning. The model unifies heterogeneous urban data into a sequential representation and introduces a prompt learning framework combining data-driven and task-specific prompts, enabling a single model to handle diverse tasks, such as prediction, interpolation, extrapolation, and imputation. Extensive experiments across multiple cities demonstrate strong performance, scalability, and promising zero-shot/few-shot generalization.

**Strengths:**
1. It addresses an important and practical problem in building general-purpose models for urban spatio-temporal tasks.
2. The unified data representation and prompt mechanism are well-motivated, making the approach flexible across tasks and modalities.
3. The paper provides comprehensive empirical validation, with consistent improvements over baselines, strong generalization ability, and clear presentation.

**Weaknesses:**
1. The model builds upon existing techniques, such as diffusion transformers, prompt learning and memory pooling, raising concerns that contributions are more engineering integration than methodological novelty.
2. The requirement of both grid and graph modalities may limit deployment, and experiments focus mainly on mobility/traffic datasets, without broader multimodal or environmental data.
3. The paper does not thoroughly report training/inference cost, memory efficiency, or scalability trade-offs, which are important for real-world applications.
4. The roles of different prompts and memory pools could be better justified and illustrated with concrete urban examples.

Overall, despite these weaknesses, I agree with most reviewers that the paper makes a solid and meaningful contribution to urban AI, with a well-designed framework and strong empirical results. Therefore, I recommend acceptance. Meanwhile, the final version should include further discussion on technical novelty, efficiency, and broader applicability as discussed during the rebuttal.
Decision: Accept (poster)