HURST: Learning Heterogeneity-Adaptive Urban Foundation Models for Spatiotemporal Prediction via Self-Partitional Mixture-of-Spatial-Experts

ICLR 2026 Conference Submission15422 Authors

Keywords: Spatial-Temporal Prediction, Mixture of Experts, Pre-trained Model, Spatial Heterogeneity

Abstract: Urban foundation models (UFMs) are pre-trained spatiotemporal (ST) prediction models with the ability to generalize to different tasks. Such models have the potential to transform urban intelligence by reducing domain-specific models and generalizing to tasks with limited data. However, building effective UFMs is a challenging task due the existence of spatial heterogeneity in ST data, i.e., data distribution and relationship between attributes vary over space. Existing UFMs lack sufficient consideration of this important issue and thus have unsatisfactory performance over spatially heterogeneous urban settings. To address this limitation, this paper proposes $\textbf{HURST}$, a $\underline{\textbf{H}}$eterogeneity-Adaptive $\underline{\textbf{UR}}$ban Foundation Model for $\underline{\textbf{S}}$patio-$\underline{\textbf{T}}$emporal Prediction, that is capable of capturing the spatial pattern of heterogeneity underlying the urban setting to enhance the UFM's performance. HURST presents two key technical innovations: (1) a self-partitional Mixture-of-Spatial-Experts (MoSE) network that automatically learns to stratify urban areas into partitions, where region-specific expert networks are trained in a hierarchical manner, and (2) an error-guided adaptive spatio-temporal masking strategy that dynamically adjusts masking patterns based on region-specific training feedback. A prompt-tuning strategy is also designed to facilitate the above innovations. Comprehensive experiments on over ten datasets from three urban areas of varying sizes show that HURST achieves up to 46.9\% performance gain over SOTA baselines.

Primary Area: foundation or frontier models, including LLMs

Submission Number: 15422