Source-linked AI summary

UniST: A Prompt-Empowered Universal Model for Urban Spatio-Temporal Prediction

Yuan Yuan, Jingtao Ding, Jie Feng, Depeng Jin, Yong Li

arXiv:2402.11838v5cs.LG

TL;DR

Urban spatio-temporal prediction still lacks a universal model that can handle diverse domains and cities despite the importance of forecasting urban dynamics. UniST combines diverse data, Transformer-based spatio-temporal modeling, pre-training, and knowledge-guided prompts, and demonstrates strong cross-scenario performance, especially in few-shot and zero-shot settings. Its scalability is limited by the relative lack of diversity in available spatio-temporal data.

  • Problem

    Urban spatio-temporal prediction lacks a universal model that generalizes across diverse domains and cities, where shared patterns must be identified despite substantial distribution differences.

  • Method

    UniST uses diverse spatio-temporal data, patching into a unified sequence, Transformer-based encoding and decoding, masked-token pre-training, and knowledge-guided prompt learning.

  • Results

    UniST demonstrates powerful prediction across scenarios, particularly in few-shot and zero-shot settings.

  • Takeaways & Limitations

    UniST supports the potential of a one-for-all model for urban spatio-temporal prediction across multiple scenarios.

  • Takeaways & Limitations

    UniST’s scalability may encounter marginal returns because spatio-temporal data has relatively less diversity than language or visual datasets.

Abstract

from arXiv · show

Urban spatio-temporal prediction is crucial for informed decision-making, such as traffic management, resource optimization, and emergence response. Despite remarkable breakthroughs in pretrained natural language models that enable one model to handle diverse tasks, a universal solution for spatio-temporal prediction remains challenging Existing prediction approaches are typically tailored for specific spatio-temporal scenarios, requiring task-specific model designs and extensive domain-specific training data. In this study, we introduce UniST, a universal model designed for general urban spatio-temporal prediction across a wide range of scenarios. Inspired by large language models, UniST achieves success through: (i) utilizing diverse spatio-temporal data from different scenarios, (ii) effective pre-training to capture complex spatio-temporal dynamics, (iii) knowledge-guided prompts to enhance generalization capabilities. These designs together unlock the potential of building a universal model for various scenarios Extensive experiments on more than 20 spatio-temporal scenarios demonstrate UniST's efficacy in advancing state-of-the-art performance, especially in few-shot and zero-shot prediction. The datasets and code implementation are released on https://github.com/tsinghua-fib-lab/UniST.

1 INTRODUCTION

Urban spatio-temporal prediction lacks a universal model that generalizes across domains and cities, partly because datasets vary in format and distribution. UniST addresses this gap by combining diverse data, spatio-temporal pre-training, and knowledge-guided prompts.

  • Motivation: Existing approaches are usually confined to specific domains within a single city and depend on extensive training data.These constraints limit their generalization potential across urban scenarios.
  • Challenges: Spatio-temporal datasets differ in dimensions, temporal durations, spatial coverage, and data distributions across domains and cities.These differences make standardization, adaptation, and generalization difficult.
  • Shared Patterns: Despite varied observed patterns, urban datasets may share underlying regularities arising from human activity, repetition, and relationships among functional zones.Traffic speed and communication networks, for example, are both influenced by human mobility.
  • UniST: UniST is designed for scalability across scenarios, effective capture of complex spatio-temporal relationships, and alignment of shared patterns through spatio-temporal prompts.These capabilities define the paper’s one-for-all modeling objective.
  • UniST: UniST targets universal prediction through extensive multi-domain and multi-city data, spatio-temporal patching, masked-token pre-training, and prompt learning.Its prompt network identifies shared patterns and adapts prompts to align distinct dataset distributions.

2 RELATED WORK

Urban spatio-temporal prediction has been approached with diverse deep-learning architectures, while newer work explores foundation models and prompt learning. However, spatio-temporal data’s intertwined spatial and temporal dependencies and limitations involving dataset formats and city transfer remain important considerations.

  • Urban Spatio-Temporal Prediction: CNNs, RNNs, ResNets, MLPs, GNNs, Transformers, and diffusion models have been introduced to capture urban spatio-temporal patterns.
  • Foundation Models for Spatio-temporal Data and Time Series: Foundation-model research applies large language models to urban systems and language-based tasks, including CityGPT, CityBench, and UrbanGPT.
  • Limitations: The paper identifies whether diverse datasets with diverse formats can be leveraged and notes restriction within the same city.
  • Foundation Models for Spatio-temporal Data and Time Series: Unlike straightforward time series, spatio-temporal data contains intertwined dependencies across spatial and temporal dimensions and is not inherently generated by language.
  • Prompt Learning: Prompt learning aims to improve pretrained models’ generalization to specific tasks or domains, with this work adaptively generating spatio-temporal prompts.

3 METHODOLOGY

UniST unifies diverse urban spatio-temporal data through a Transformer-based architecture, large-scale pre-training, and knowledge-guided prompt learning. Its design targets cross-scenario generalization by modeling spatial and temporal dependencies and adapting prompts to related patterns.

  • Problem formulation: UniST defines spatio-temporal data as variable-sized 4D tensors and formulates prediction from historical observations to future steps.The dimensions T, C, H, and W can vary across scenarios; few-shot adaptation uses limited target samples, while zero-shot prediction uses no fine-tuning.
  • Overall framework: A two-stage framework combines large-scale pre-training on data from multiple domains and cities with knowledge-guided prompt learning.The prompt network uses spatial hierarchy and temporal periodicity to guide in-context learning.
  • Base model: A Transformer encoder-decoder uses spatio-temporal patching to convert diverse data into a unified sequential format.Patching transforms channel-independent 3D tensors into smaller tensors, while convolutional embedding and positional encoding support sequential processing.
  • Self-supervised pre-training: Masked pre-training uses multiple strategies to capture spatio-temporal, spatial, and temporal dependencies.Tube masking targets spatial extrapolation, block masking targets spatial transferability, and temporal masking reconstructs future data from historical information.
  • Knowledge-guided prompt learning: Prompt learning represents spatial closeness, spatial hierarchy, temporal closeness, and temporal periodicity through four networks and learned memory pools.Representations query spatial and temporal memories to extract prompts that are integrated into the Transformer input; the approach assumes related fine-grained patterns can be identified across datasets.

4 PERFORMANCE EVALUATIONS

UniST is evaluated across diverse urban spatio-temporal datasets against broad baseline families under short-term, long-term, few-shot, and zero-shot settings. It consistently outperforms baselines, including an average 10.1% improvement in long-term prediction and strong transfer to unseen scenarios.

  • Experimental Setup: Experiments cover more than 20 spatio-temporal datasets spanning crowd flow, dynamic population, traffic speed, cellular usage, taxi trips, and bike demand.The datasets vary in cities, domains, variables, sampling frequencies, spatial scales, durations, and sizes.
  • Experimental Setup: UniST is compared with heuristic, deep urban prediction, video prediction, and multivariate time-series forecasting models using MAE and RMSE.Baselines include dedicated and one-for-all variants, while UniST is evaluated across all datasets.
  • Short-Term Prediction: UniST consistently outperforms all baselines across the reported short-term prediction datasets.The results also indicate that spatial dependency is important and that baseline performance is unstable across scenarios.
  • Long-Term Prediction: 10.1% average improvement over each dataset’s best baseline is achieved in long-term prediction.The long-term setup uses an input step and prediction horizon of 64 to assess extended temporal dependencies.
  • Few-Shot Prediction: UniST outperforms all baselines in few-shot settings and achieves larger relative improvements than in short-term and long-term prediction.Figure 4 reports few-shot evaluations using 1% and 5% of training data on Crowd and BikeNYC.
  • Zero-Shot Prediction: UniST achieves remarkable zero-shot performance, surpassing many baselines trained with data on entirely novel datasets.The authors attribute this result to spatio-temporal transfer and extraction of fine-grained similarities through spatial and temporal properties.

5 STUDY AND ANALYSIS ON UNIST

Ablation and interpretability studies examine UniST’s prompts, memory pools, embedding behavior, and scaling. Prompt properties improve performance, memory size helps up to 512 embeddings, and larger models improve performance with diminishing returns.

  • Ablation Study: Removing any spatial or temporal prompt property decreases performance across four traffic speed datasets.The evaluated properties are spatial closeness and hierarchy, temporal periodicity, and temporal closeness.
  • Ablation Study: Increasing memory-pool embeddings from 128 to 512 improves performance across all four datasets, while 1024 performs similarly to 512.The results suggest 512 embeddings is the optimal choice in this study.
  • Prompt Learner Analysis: Spatial and temporal memory embeddings diverge in different directions from initialization to their optimized states.The study uses t-SNE visualization to compare initial and final embedding states.
  • Prompt Learner Analysis: Memory embeddings reveal consistent memorized patterns across the Crowd and TrafficSH scenarios.Inputs assigning the highest attention weight to each embedding are aggregated and averaged for comparison.
  • Prompt Learner Analysis: Attention-weight distributions differ markedly between two datasets, indicating dataset-responsive prompt selection.The comparison examines the mean attention weight assigned to each memory embedding in each domain.
  • Scalability: Larger UniST models generally improve performance, but scaling exhibits diminishing marginal returns relative to language and vision models.The analysis compares training loss and testing RMSE across parameter sizes and attributes the weaker scaling trend to limited spatio-temporal data diversity.

6 CONCLUSION

UniST addresses universal urban spatio-temporal prediction by learning shared patterns from diverse data across scenarios. The authors report strong cross-scenario prediction, particularly in few-shot and zero-shot settings, while identifying broader foundation models for spatio-temporal data as future work.

  • Conclusion: UniST leverages diverse data from multiple sources to align shared spatio-temporal patterns across scenarios.The conclusion frames this alignment as the basis for prediction across urban scenarios.
  • Conclusion: UniST demonstrates prediction capability across all scenarios, particularly in few-shot and zero-shot settings.The conclusion presents this as the paper’s central supported outcome.
  • Future Work: Future work includes integrating broader foundation-model ideas while developing models trained specifically on pure spatio-temporal data.The paper notes that spatio-temporal data has intertwined spatial and temporal dependencies and is not inherently generated by language.

A.2 Data Preprocessing

UniST is trained through multi-dataset spatio-temporal pre-training and prompt tuning, with customized prompts generated from spatio-temporal knowledge. The experiments compare its scalability and training setup with diverse baselines.

  • Data preprocessing: Datasets are split chronologically into 70% training, 15% validation, and 15% testing periods, with intermediate sequences removed.All datasets are normalized to [−1, 1], while reported predictions are denormalized.
  • Evaluation and baselines: The study evaluates UniST with MAE and RMSE and compares it against statistical, spatio-temporal, Transformer, meta-learning, and transfer-learning baselines.Compared baselines include scalable methods such as PatchTST, MAML, and MetaST, alongside non-scalable approaches.
  • Optimization: Training uses Adam, with learning rates of 3e-4 for pre-training and 5e-5 for prompt tuning, selected through validation-based model selection.Performance is evaluated every ten epochs, and the best validation model is used for testing.
  • Pre-training: UniST pre-training samples datasets and mini-batches, applies one of four masking strategies, and reconstructs masked inputs with the base model.The pre-training algorithm is defined over multiple datasets, a base prediction model, and a loss function.
  • Prompt tuning: Prompt tuning trains a network that generates customized prompts from spatial closeness, spatial hierarchy, temporal closeness, and temporal period features.The objective minimizes prediction error with l2 regularization to reduce over-fitting.

D.5 Experimental Design

UniST is evaluated on short-term, long-term, few-shot, and zero-shot prediction tasks. The latter two assess generalization from source datasets to target datasets with limited or no target-data fine-tuning.

  • Task design: The experimental design includes short-term, long-term, few-shot, and zero-shot prediction tasks.This setup follows established evaluation practices for foundation models in time-series forecasting.
  • Task design: Short-term and long-term prediction evaluate models without transfer learning by training on datasets and testing on their corresponding test sets.These tasks measure prediction performance in non-transfer settings.
  • Generalization tasks: Few-shot and zero-shot prediction use source datasets to build a pretrained model and memory pool for prediction on target datasets.Few-shot prediction includes limited target fine-tuning, whereas zero-shot prediction differs through the absence of that fine-tuning process.
  • Evaluation scope: Together, the four tasks evaluate prediction performance and generalization across diverse spatio-temporal datasets.

E.1 Analysis of Distribution Shifts

The distribution-shift analysis examines generated prompt attention across familiar, related, and cross-city datasets. Similar spatio-temporal patterns receive similar prompts, supporting adaptive prompt generation across shifts.

  • Analysis settings: The analysis compares prompt attention weights on memory-pool embeddings across training/testing splits and datasets from related or different domains and cities.Figures 9 and 10 visualize these attention-weight distributions.
  • Analysis settings: The analysis is designed to assess within-dataset generalization, domain adaptation within a city, and reuse of previously learned knowledge across cities and domains.
  • Prompt similarity: Similar prompts are consistently generated for datasets with similar spatio-temporal patterns, including train/test splits and related domains within one city.The reported comparisons cover the CrowdNJ, TaxiBJ, and SpeedSZ scenarios.
  • Prompt similarity: Prompt generation adapts to the unique characteristics of each dataset, enabling the model to handle distribution shifts across diverse spatio-temporal contexts.

E.2 Performance under Noise Perturbations

UniST is tested under Gaussian noise perturbations and compared across positional encodings, computational costs, and dataset-pattern similarities. It retains comparative performance under noise, while sine-cosine encoding is reported as more robust.

  • Noise robustness: UniST is evaluated under Gaussian noise at multiple corruption levels to assess robustness to noisy or irregular input data.The perturbations simulate real-world data-quality problems.
  • Noise robustness: UniST consistently outperforms baseline models under noise perturbations, even when the best baseline is evaluated without noise.The comparison is reported in Table 7.
  • Positional encoding: Sine-cosine positional encoding shows more robust performance against noise perturbations than learnable positional embeddings.Table 7 reports sine-cosine results, while Table 8 reports learnable-embedding results.
  • Efficiency: UniST has model size and memory costs comparable to other approaches, but its multi-dataset pre-training requires longer training time.The study considers the additional training time acceptable alongside its reported performance advantage using one model.
  • Dataset similarities: t-SNE followed by k-means is used to visualize and cluster dataset patterns, revealing similarities within some city-domain pairs and shifts across most cities or domains.Crowd and Cellular data in Nanjing, and Bike and Taxi data in New York City, are identified as similar pairs.

E.5 Additional Ablation Studies

Additional ablations show that UniST benefits from using all four masking strategies and all designed spatial-temporal prompt properties. Their effects vary by task and dataset, while even less task-aligned masks contribute to broader dependency learning.

  • Masking Strategies: All four masking strategies achieved the best performance across prediction, imputation, and spatial extrapolation on TrafficCD.Removing temporal masking most affected prediction, random masking most affected imputation, and block masking most affected spatial extrapolation.
  • Prompt Properties: Figure 13 reports ablation studies for the four spatial and temporal knowledge-extraction properties.The figure concerns t_c, t_p, s_c, and s_h.
  • Masking Strategies: Less task-aligned masking strategies still improved other tasks, indicating broader effects on learning spatio-temporal dependencies and dynamics.Random masking can support robust features across time points, while temporal masking can improve temporal-dynamics understanding during spatial extrapolation.
  • Prompt Properties: The four prompt properties are spatial closeness, spatial hierarchy, temporal periodicity, and temporal closeness.These properties are denoted as s_c, s_h, t_p, and t_c, respectively.
  • Prompt Properties: Removing any prompt property decreased performance across four traffic speed datasets, although each property’s contribution varied by dataset.The result supports retaining all four properties in the spatio-temporal prompt design.

E.6 Additional Prediction Results

Additional prediction results cover short-term, long-term, few-shot, and zero-shot settings across multiple urban datasets. The reported metrics are average MAE and RMSE over all prediction steps.

  • Prediction Results: Tables 11 and 12 report short-term prediction comparisons on seven datasets using MAE and RMSE.Errors are averaged over all prediction steps.
  • Prediction Results: Table 13 reports long-term prediction comparisons on four datasets using average MAE and RMSE over all prediction steps.
  • Few-Shot and Zero-Shot Results: Table 14 evaluates few-shot and zero-shot prediction on Crowd using MAE and RMSE.The few-shot settings use 1%, 5%, and 10% of the training data, while zero-shot results are reported only for UniST.
  • Few-Shot and Zero-Shot Results: Table 15 evaluates few-shot and zero-shot prediction on BikeNYC using MAE and RMSE.The few-shot settings use 1%, 5%, and 10% of the training data, while zero-shot results are reported only for UniST.
  • Few-Shot and Zero-Shot Results: Table 16 evaluates few-shot and zero-shot prediction on TaxiBJ using MAE and RMSE.The few-shot settings use 1%, 5%, and 10% of the training data, while zero-shot results are reported only for UniST.
Loading 2402.11838v5…