Memory-Enhanced Invariant Prompt Learning for Urban Flow Prediction Under Distribution Shifts

Abstract

While Spatial-Temporal Graph Neural Networks (STGNNs) excel at urban flow prediction, they struggle with distribution shifts caused by dynamic spatial-temporal environments. To improve generalizability to out-of-distribution (OOD) data, a typical solution is to disentangle invariant patterns that carry stable causal effects from variant ones that are environment-dependent. Existing OOD-robust methods attempt to model these environments but face challenges in quantifying dynamic changes and suffer from high computational costs. As a solution, we propose Memory-enhanced Invariant Prompt Learning (MIP), which enables environmental interventions directly within the latent space by learning a memory bank from the spatial-temporal urban flow graphs. Then, by performing spatial-temporal interventions on the variant prompts, diverse environments are constructed in the latent space to facilitate invariant learning. The invariant prompts, together with a memory-enhanced causal graph, are fed into an STGNN backbone to produce accurate predictions. Extensive experiments on two public urban flow datasets confirm MIP’s effectiveness in improving robustness against OOD data.