Diffusion Transformers as Open-World Spatiotemporal Foundation Models

Dear Reviewer NF7R,

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

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.

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.

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

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

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.

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:

• 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.
• We will also visualize the specific spatio-temporal patterns memorized by each unit.

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

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?

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.

Thank you again for your positive ratings and constructive feedback. We are happy to address any remaining questions.

Sincerely,

The Authors

Dear Reviewer XMSo,

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

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?

Response. We appreciate the reviewer’s comment. Regarding the justification for unified prompt learning, we emphasize the following key points.

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

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

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.

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.

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

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

Thank you again for your constructive suggestions. We are happy to address any remaining questions.

Sincerely,

The Authors

Dear Reviewer EsKL,

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

W1. Stringent Data Requirements: The model mandates both grid and graph-structured inputs for unified training , restricting deployment in tasks lacking either data modality.

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.

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

We will clarify this point in $\underline{\text{Section 3.2 ("Unification of Data and Tasks")}}$ in the revised version to prevent potential misunderstandings.

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.

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.

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

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.

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.

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.

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.

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.

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

Q1. During model training, does a single task require simultaneous use of both graph-based and grid-based data?

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.

Q2. How do embeddings of graph-based and grid-based data achieve semantic alignment in the spatio-temporal domain?

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.

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

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

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

Sincerely,

The Authors

Thank you again for your continued engagement and thoughtful feedback.

We appreciate your effort to differentiate what the model can do from how it achieves it through technical innovation. To clarify our contributions from a methodological perspective, we note that recent breakthroughs in fields like time series, human mobility, cellular network, and Earth system modeling [1-4] consistently highlight three axes of innovation: (i) unified architectures across modalities, (ii) scalable training across heterogeneous data, and (iii) flexible task adaptation. These axes precisely frame UrbanDiT’s core design.

We argue that in domain-specific foundation modeling, technical innovation often lies not in inventing new architectural primitives, but in designing a scalable, extensible system that successfully solves the grand unification challenge of the target domain. This view is supported by seminal works across scientific fields [1–4], where innovation arises from integrating powerful components into cohesive frameworks that generalize broadly across real-world scenarios. For instance, AlphaFold3 [5] represents a major leap beyond AlphaFold [6] not because of increased architectural complexity, but because of a shift in modeling paradigm: it simplifies the network backbone while expanding to richer multi-modal inputs and adopts diffusion models. This trend underscores that in domain-specific foundation models, system-level unification and data-driven scalability are central to technical innovation.

While UrbanDiT builds on diffusion and Transformer backbones, just as Sundial, Aurora and UniTraj do, it introduces a holistic framework for cross-format, cross-task, and cross-domain urban modeling. By this standard, it makes substantive contributions as the first foundation model tailored to the complexity of open-world urban spatio-temporal learning.

We hope this clarifies our position and would greatly appreciate your reconsideration.

[1]Bodnar, C., Bruinsma, W. P., Lucic, A., Stanley, M., Allen, A., Brandstetter, J., ... & Perdikaris, P. (2025). A foundation model for the Earth system. Nature, 1-8.

[2] Liu, Y., Qin, G., Shi, Z., Chen, Z., Yang, C., Huang, X., ... & Long, M. Sundial: A Family of Highly Capable Time Series Foundation Models. In Forty-second International Conference on Machine Learning.

[3] Zhu, Y., Yu, J. J., Zhao, X., Wei, X., & Liang, Y. (2024). Unitraj: Learning a universal trajectory foundation model from billion-scale worldwide traces. arXiv preprint arXiv:2411.03859.

[4] Chai, H., Zhang, S., Qi, X., Qiu, B., & Li, Y. (2025). UoMo: A Universal Model of Mobile Traffic Forecasting for Wireless Network Optimization. In Proceedings of the 31st ACM SIGKDD Conference on Knowledge Discovery and Data Mining

[5] Abramson, J., Adler, J., Dunger, J., Evans, R., Green, T., Pritzel, A., ... & Jumper, J. M. (2024). Accurate structure prediction of biomolecular interactions with AlphaFold 3. Nature, 630(8016), 493-500.

[6] Jumper, J., Evans, R., Pritzel, A., Green, T., Figurnov, M., Ronneberger, O., ... & Hassabis, D. (2021). Highly accurate protein structure prediction with AlphaFold. nature, 596(7873), 583-589.

Dear Reviewer DoyW,

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

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.

Response. We respectfully argue that our contributions go beyond engineering design and offer meaningful technical innovations for the spatio-temporal modeling community.

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.

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.

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.

[1] Timedit: General-purpose diffusion transformers for time series foundation model. arXiv preprint. [2] Text2PDE: Latent Diffusion Models for Accessible Physics Simulation, ICLR 2025 [3] Diffusion on Language Model Encodings for Protein Sequence Generation, ICML 2025

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.

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:

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.

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

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.

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.

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.

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

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

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

Thank you again for your constructive feedback. We are happy to address any remaining questions.

Sincerely,

The Authors