Memory-enhanced Invariant Prompt Learning for Urban Flow Prediction under Distribution Shifts

TL;DR

The paper addresses the challenge of distribution shifts in urban flow prediction by introducing Memory-enhanced Invariant Prompt learning (MIP), which leverages a memory bank to extract invariant (causal) and variant (environment-dependent) prompts and performs latent-space interventions on the variant prompts. An invariant learning objective, aided by a memory-augmented semantic graph and a spatial-temporal backbone, disentangles causal from spurious signals to yield robust predictions under OOD conditions. Comprehensive experiments on METR-LA and NYCBike show that MIP consistently outperforms state-of-the-art baselines across multiple test sets with distribution shifts, while ablation studies confirm the contribution of memory, adaptive adjacency, and invariant learning. The framework offers scalable, real-time capable predictions and introduces a practical approach to robust urban forecasting without enumerating multiple explicit environments.

Abstract

Urban flow prediction is a classic spatial-temporal forecasting task that estimates the amount of future traffic flow for a given location. Though models represented by Spatial-Temporal Graph Neural Networks (STGNNs) have established themselves as capable predictors, they tend to suffer from distribution shifts that are common with the urban flow data due to the dynamics and unpredictability of spatial-temporal events. Unfortunately, in spatial-temporal applications, the dynamic environments can hardly be quantified via a fixed number of parameters, whereas learning time- and location-specific environments can quickly become computationally prohibitive. In this paper, we propose a novel framework named Memory-enhanced Invariant Prompt learning (MIP) for urban flow prediction under constant distribution shifts. Specifically, MIP is equipped with a learnable memory bank that is trained to memorize the causal features within the spatial-temporal graph. By querying a trainable memory bank that stores the causal features, we adaptively extract invariant and variant prompts (i.e., patterns) for a given location at every time step. Then, instead of intervening the raw data based on simulated environments, we directly perform intervention on variant prompts across space and time. With the intervened variant prompts in place, we use invariant learning to minimize the variance of predictions, so as to ensure that the predictions are only made with invariant features. With extensive comparative experiments on two public urban flow datasets, we thoroughly demonstrate the robustness of MIP against OOD data.

Figures (11)

• Figure 1: Causal graph of spatial-temporal prediction. Previous STGNNs learn the embeddings containing both invariant and variant features and make predictions based on the relation line 2, but they are always misled to learn the spurious relation line 3 because of the variant features. In this study, we try to separate the invariant and variant features and make predictions based on the causal relationship, line 1.
• Figure 2: The sampled traffic speed recorded by three sensors in the METR-LA dataset. The records correspond to two Wednesdays in Weeks 3 and 13 of the dataset.
• Figure 3: The workflow of MIP.
• Figure 4: Ablation study: RMSE of MIP and its variants.
• Figure 5: RMSE of MIP with different intervention rate.