UniFlow: A Foundation Model for Unified Urban Spatio-Temporal Flow Prediction

TL;DR

This work introduces UniFlow, a foundation model that unifies grid-based and graph-based urban spatio-temporal flow prediction under a single, transformer-based framework. It employs multi-view spatio-temporal patching to standardize inputs, a spatio-temporal transformer for sequential modeling, and Spatio-Temporal Memory Retrieval Augmentation (ST-MRA) to learn and reuse shared patterns across data types. ST-MRA builds explicit memories and uses learned time-domain, frequency-domain, and adaptive spatial prompts to augment predictions, enabling cross-learning and improved robustness, including few-shot and zero-shot scenarios. Extensive experiments across nine real-world datasets show UniFlow outperforming specialized baselines by substantial margins, with strong robustness and interpretability through memory prompts, demonstrating the practical potential of a universal urban flow model.

Abstract

Urban spatio-temporal flow prediction, encompassing traffic flows and crowd flows, is crucial for optimizing city infrastructure and managing traffic and emergency responses. Traditional approaches have relied on separate models tailored to either grid-based data, representing cities as uniform cells, or graph-based data, modeling cities as networks of nodes and edges. In this paper, we build UniFlow, a foundational model for general urban flow prediction that unifies both grid-based and graphbased data. We first design a multi-view spatio-temporal patching mechanism to standardize different data into a consistent sequential format and then introduce a spatio-temporal transformer architecture to capture complex correlations and dynamics. To leverage shared spatio-temporal patterns across different data types and facilitate effective cross-learning, we propose SpatioTemporal Memory Retrieval Augmentation (ST-MRA). By creating structured memory modules to store shared spatio-temporal patterns, ST-MRA enhances predictions through adaptive memory retrieval. Extensive experiments demonstrate that UniFlow outperforms existing models in both grid-based and graph-based flow prediction, excelling particularly in scenarios with limited data availability, showcasing its superior performance and broad applicability. The datasets and code implementation have been released on https://github.com/YuanYuan98/UniFlow.

Figures (7)

• Figure 1: High-level comparison between our proposed model and separate modeling for spatio-temporal data. Our solution unifies grid and graph data formats.
• Figure 2: The overall architecture of UniFlow. For the input data, we first perform spatio-temporal patching to construct sequential structures. Then the transformer model performs sequential modeling, which is prompted by spatio-temporal memory retrieval augmentation.
• Figure 3: ST-MRA workflow in UniFlow.
• Figure 4: Performance comparison of few-shot predictions between UniFlow and baseline models on the CrowdSH dataset. The red dashed line represents zero-shot performance.
• Figure 5: Results of ablation studies of UniFlow in terms of the number of units in each memory.