UrbanDiT: A Foundation Model for Open-World Urban Spatio-Temporal Learning

1. The motivation of this paper is not well-justified. The choice of diffusion transformers seems arbitrary and does not convincingly address why they are the most suitable model, compared with STGNNs or standard transformers, for urban spatio-temporal learning tasks.
2. The technical novelty is limited compared with previous literature. The idea of sequential unification strategy seems to come from MOIRAI [1], and the design of masking and prompt are very similar to UniST [2], which also focuses on urban spatio-temporal learning. Besides, the authors directly employ diffusion transformer without specific modification for spatio-temporal data.
3. The experimental setup is unclear. The authors provide an overview of the datasets used, but it is not explicitly stated which datasets were used for pretraining and which were used solely for evaluation. This lack of clarity raises concerns about the fairness of the few-shot and zero-shot comparisons. If all datasets were included in the pretraining phase, the zero-shot and few-shot results may not provide an accurate reflection of the model's generalization capabilities.
4. UrbanDiT presents promising results on the zero-shot and few-shot capabilities, but these evaluations are limited to two datasets, TaxiBJ and PopSH, which are both grid-based flow datasets. To strengthen the generalization claims, it would be beneficial to include evaluations on more heterogeneous datasets, e.g., PEMS-BAY, PEMS-04, PEMS-07, and PEMS-08.
5. In zero-shot and few-shot experiments, the authors did not include comparisons against some of the latest large pre-trained models, such as TTM [3] and OpenCity [4]. Both of them offer open-source codes and pre-trained model weights, demonstrating promising results in zero-shot and few-shot scenarios.
6. The authors flatten the input data into a one-dimensional sequential format. For datasets with a large number of nodes, such as SpeedSH with over 20,000 nodes (i.e., 20,000 tokens), this approach could lead to significant computational challenges. Given the quadratic complexity of the transformer architecture with respect to the input sequence length, processing such large inputs could result in severe efficiency bottlenecks.
7. The authors mention that the datasets used in the experiments are publicly available. However, upon reviewing the provided link, I was unable to locate the datasets.
8. As an ICLR submission, the paper would benefit from stronger theoretical foundations. For instance, the authors claim that the prompts learned from the pretraining data can generalize well to unseen datasets. It would be interesting to see a theoretical explanation of why and how these learned prompts are able to generalize effectively to new, unseen urban environments.

[1] Woo, Gerald, Chenghao Liu, Akshat Kumar, Caiming Xiong, Silvio Savarese, and Doyen Sahoo. "Unified Training of Universal Time Series Forecasting Transformers." In Forty-first International Conference on Machine Learning.

[2] 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, pp. 4095-4106. 2024.

[3] Vijay, E., Arindam Jati, Pankaj Dayama, Sumanta Mukherjee, Nam Nguyen, Wesley Gifford, Chandra Reddy, and Jayant Kalagnanam. "Tiny Time Mixers (TTMs): Fast Pre-trained Models for Enhanced Zero/Few-Shot Forecasting of Multivariate Time Series." arXiv (2024).

[4] Li, Zhonghang, Long Xia, Lei Shi, Yong Xu, Dawei Yin, and Chao Huang. "OpenCity: Open Spatio-Temporal Foundation Models for Traffic Prediction." arXiv preprint arXiv:2408.10269 (2024).

1. The key idea that using prompt learning to achieve versatility on diverse datasets is almost the same as UniST [1]. Both works generate prompts by leveraging the input data to retrieve memory pools containing diverse spatio-temporal patterns. It makes this work’s technical innovation very limited.
2. The contribution is overclaimed. As this work is similar to UniST, which presents a unified model to address various urban spatio-temporal prediction problems, it is not the first work to explore a foundation model for urban spatio-temporal learning.
3. The zero-shot experiments are insufficient that only results of PopSH and TaxiBJ are provided. How about the zero-shot performance on other datasets and tasks? Additionally, as UrbanDiT’s zero-shot performance is not always better than other models with training data (refer to Figure 9), how about the zero-shot comparison with other baseline models?
4. The paper only provides ablation results on TaxiBJ. It’s expected to see the results on other datasets, which would be beneficial to conclude more compelling insights.

[1] Unist: a prompt-empowered universal model for urban spatio-temporal prediction. KDD, 2024.

(1) The technical approach and problems addressed in this paper are very similar to the published paper UniST[1]. Both use masked recovery as the pre-training task, employ spatio-temporal Transformer as the basic architecture, and enhance model generalization across different data and tasks through spatio-temporal prompt learning. The differences mainly lie in technical details, such as using Diffusion Transformer instead of Masked Autoencoder as the basic architecture. This results in insufficient contributions of the paper.

(2) While the paper mentions similar large models like GPD, UniST, UrbanGPT, and CityGPT, the comparative experiments lack most of these spatio-temporal models. Including these comparisons would make the results more convincing.

(3) Technical and experimental details are incomplete and require more detailed explanations. For example: How are the prototype prompts in the Memory Pool obtained? What are the data sources during pre-training? Is the 6:2:2 split only for downstream tasks? Does UrbanDiT require retraining on target data/tasks during the downstream phase, such as in the experiments in Section 4.2?

(4) Scalability needs further exploration. The paper only discusses data scalability but lacks parameter scalability experiments, which is crucial for foundation models.

1. The paper lacks information on the GPU used and the efficiency of the model. There are concerns about whether training across all datasets and tasks based on DiT could result in significantly high computational complexity.
2. Given the authors' data processing approach, relying solely on the reshape function cannot restore the data to its original shape. Therefore, an output layer capable of handling various data formats is essential. However, it is puzzling that neither the paper nor the model diagram addresses the design of the output layer.

W1: The lack discussion on limitation and future works to briefly elaborate the gap, or at least the next step, between the proposed novel FM as a benchmark to a product that provide utility to stakeholders in the urban environment domain.

W2: Table 1 listed 4 other existing models that is most similar to UrbanDiT. However, none of these 4 models appear in the tables in the result sections, and UniST is the only one listed in subsection C1. There might be valid reasons to not include them, but it is not provided.

W3: Since unified prompt learning is a key design, more analysis than ablation is required. It would be good to have more analysis to have some insight into the interpretability of the unified prompt learning module. For example, are there any semantic meaning behind the entries in the memory pool that we can discover?

W4: There are many possible zero-shot setup, but this paper only explored one, that is the novel dataset setup. Another straight forward one is the novel task. Even the current setup, testing on PopSH, is not a truly zero-shot, since model has learned something about the population dataset (from PopBJ), and the city Shanghai (from SpeedSH).

W5: There is a lack of analysis to show the bottleneck of the current setup. To truly "inspire future research in the rapidly evolving field of foundation models" in the urban domain, it would be good if the authors can share the lessons learned from running this extensive experiments regarding the bottle neck of the domain, i.e. is it data, or compute, or lack of standardized benchmark (data + setup) or something else?