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

arXiv: https://arxiv.org/abs/2412.05534 · HTML: https://mlsanity.com/paper/2412.05534/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 53 references that are in the map.
- First posted on arXiv 7 Dec 2024.

## Builds on (most cited first, 29 of 29)

- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Convolutional Neural Networks on Graphs with Fast Localized Spectral Filtering](https://mlsanity.com/paper/1606.09375/) (2016)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Spatio-Temporal Graph Convolutional Networks: A Deep Learning Framework for Traffic Forecasting](https://mlsanity.com/paper/1709.04875/) (2017)
- [Graph WaveNet for Deep Spatial-Temporal Graph Modeling](https://mlsanity.com/paper/1906.00121/) (2019)
- [Adaptive Graph Convolutional Recurrent Network for Traffic Forecasting](https://mlsanity.com/paper/2007.02842/) (2020)
- [GMAN: A Graph Multi-Attention Network for Traffic Prediction](https://mlsanity.com/paper/1911.08415/) (2019)
- [Deep Spatio-Temporal Residual Networks for Citywide Crowd Flows Prediction](https://mlsanity.com/paper/1610.00081/) (2016)
- [Invariant Risk Minimization Games](https://mlsanity.com/paper/2002.04692/) (2020)
- [Discrete Graph Structure Learning for Forecasting Multiple Time Series](https://mlsanity.com/paper/2101.06861/) (2021)
- [Handling Distribution Shifts on Graphs: An Invariance Perspective](https://mlsanity.com/paper/2202.02466/) (2022)
- [Spatial-Temporal Identity: A Simple yet Effective Baseline for Multivariate Time Series Forecasting](https://mlsanity.com/paper/2208.05233/) (2022)
- [CoST: Contrastive Learning of Disentangled Seasonal-Trend Representations for Time Series Forecasting](https://mlsanity.com/paper/2202.01575/) (2022)
- [PDFormer: Propagation Delay-Aware Dynamic Long-Range Transformer for Traffic Flow Prediction](https://mlsanity.com/paper/2301.07945/) (2023)
- [AdaRNN: Adaptive Learning and Forecasting of Time Series](https://mlsanity.com/paper/2108.04443/) (2021)
- [Dish-TS: A General Paradigm for Alleviating Distribution Shift in Time Series Forecasting](https://mlsanity.com/paper/2302.14829/) (2023)
- [Spatio-Temporal Meta-Graph Learning for Traffic Forecasting](https://mlsanity.com/paper/2211.14701/) (2022)
- [Spatio-Temporal Self-Supervised Learning for Traffic Flow Prediction](https://mlsanity.com/paper/2212.04475/) (2022)
- [Deciphering Spatio-Temporal Graph Forecasting: A Causal Lens and Treatment](https://mlsanity.com/paper/2309.13378/) (2023)
- [STDEN: Towards Physics-Guided Neural Networks for Traffic Flow Prediction](https://mlsanity.com/paper/2209.00225/) (2022)
- [Spatial-Temporal Hypergraph Self-Supervised Learning for Crime Prediction](https://mlsanity.com/paper/2204.08587/) (2022)
- [Taming Local Effects in Graph-based Spatiotemporal Forecasting](https://mlsanity.com/paper/2302.04071/) (2023)
- [Graph Out-of-Distribution Generalization via Causal Intervention](https://mlsanity.com/paper/2402.11494/) (2024)
- [Long-Term Occupancy Grid Prediction Using Recurrent Neural Networks](https://mlsanity.com/paper/1809.03782/) (2018)
- [TESTAM: A Time-Enhanced Spatio-Temporal Attention Model with Mixture of Experts](https://mlsanity.com/paper/2403.02600/) (2024)
- [Causality and Independence Enhancement for Biased Node Classification](https://mlsanity.com/paper/2310.09586/) (2023)
- [Towards Out-of-Distribution Sequential Event Prediction: A Causal Treatment](https://mlsanity.com/paper/2210.13005/) (2022)
- [Attention-based Spatial-Temporal Graph Convolutional Recurrent Networks for Traffic Forecasting](https://mlsanity.com/paper/2302.12973/) (2023)
- [Physics-guided Active Sample Reweighting for Urban Flow Prediction](https://mlsanity.com/paper/2407.13605/) (2024)

## Built on by (most cited first, 1 of 1)

- [Learning A Universal Crime Predictor with Knowledge-guided Hypernetworks](https://mlsanity.com/paper/2511.02336/) (2025)

## Abstract (arXiv)

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.