When to Trust Memory: Retrieval-Guided Probabilistic Spatiotemporal Forecasting under Distribution Shift

Haochen Lv; Jianhao Zhang; Zhichen Lei; Yang Jiang; Xiaowei Mao; Shengnan Guo; Youfang Lin; Huaiyu Wan

cs.LGOct 23, 2026

When to Trust Memory: Retrieval-Guided Probabilistic Spatiotemporal Forecasting under Distribution Shift

Authors: Haochen Lv, Jianhao Zhang, Zhichen Lei, Yang Jiang, Xiaowei Mao, Shengnan Guo, Youfang Lin, Huaiyu Wan

Organizations: School of Computer Science and Technology, Beijing Jiaotong University, Beijing, China. · Key Laboratory of Big Data and Artificial Intelligence in Transportation, Ministry of Education, Beijing, China. · Beijing Key Laboratory of Traffic Data Mining and Embodied Intelligence, Beijing, China.

Abstract

Spatiotemporal forecasting plays an important role in real-world applications, where probabilistic approaches provide uncertainty estimates for decision-making. Real-world systems are inherently non-stationary, leading to distribution shifts between training and test data. However, existing probabilistic methods mainly focus on modeling data distributions under standard settings, while distribution shift methods often handle such shifts in a coarse manner, failing to characterize shift degree and adapt prediction strategy accordingly. As a result, they often produce miscalibrated and unreliable forecasts. To address this issue, we propose REMIND, a retrieval-guided probabilistic framework for spatiotemporal forecasting under distribution shift. REMIND leverages retrieval to estimate how well current observation are supported by historical patterns, enabling adaptive prediction under different shift regimes. It employs spatiotemporal retrieval to construct a memory bank and introduces a dual branch architecture, combining a memory-calibrated probabilistic branch for mild shifts with a deviation-aware probabilistic branch for severe shifts. A shift-adaptive gate, guided by retrieval similarity and memory dispersion, balances the two branches according to shift degree. We further design a hierarchical Gaussian mixture to capture heterogeneous uncertainty over the dual branch predictions, and introduce a memory perturbation training scheme to improve generalization. Experiments on real-world datasets with distribution shifts demonstrate that REMIND consistently outperforms state-of-the-art methods in both forecasting accuracy and probabilistic prediction quality.

Figures & tables

Figure 1: The overall architecture of REMIND. REMIND performs coarse-to-fine spatiotemporal retrieval to construct a memory bank and employs a dual branch architecture for probabilistic forecasting under distribution shift. A shift-adaptive gate dynamically balances the branches based on retrieval similarity and memory dispersion, and a hierarchical Gaussian mixture models heterogeneous predictive distribution.

Figure 2: Empirical analysis of retrieval similarity and memory dispersion under distribution shift on the CAMELS dataset. We analyze their relationships with error (MAE/MAPE) between retrieved targets and ground-truth future targets, error variance, and their pairwise correlation.

Dataset	Metric	DeepAR	LODE	DER	SDER	PriSTI	DiffSTG	STGNF	STMPN	REMIND
CAMELS	MAE	9.85	10.52	9.68	9.32	9.05	8.93	8.96	9.23	8.73
	RMSE	38.95	34.87	38.73	37.95	36.84	34.12	35.93	35.93	32.26
	MAPE	29.53	45.14	37.85	38.24	35.65	30.96	29.47	34.25	25.30
	CRPS	0.3611	0.3252	0.4210	0.3056	0.2955	0.3028	0.2818	0.2577	0.2232
	ES	28.73	31.27	36.99	27.67	28.18	27.46	26.46	29.86	26.04
PEMS	MAE	19.82	21.45	20.25	19.64	19.38	19.13	18.55	17.91	16.63

Table 1: Overall performance comparison. The best is bold and the second best is underlined .

Table 4

Figure 4: Hyperparameter sensitivity on ENERGY.

Figure 6: Branch performance and memory reliance across

\text{MMD}^{2}

groups on ENERGY.

	Memory branch	Deviation branch	Mean $\boldsymbol{g}$	REMIND
Dispersion	MAE / CRPS	MAE / CRPS		MAE / CRPS
Low	1.38 / 0.0795	1.52 / 0.0860	0.91	1.35 / 0.0780
Medium	1.60 / 0.0925	1.57 / 0.0910	0.66	1.52 / 0.0875
High	1.92 / 0.1110	1.83 / 0.1045	0.38	1.78 / 0.0994

Table 5: Memory dispersion analysis on ENERGY. We select the top third of test samples by retrieval similarity and divide them into three groups based on the tertiles of

\boldsymbol{C}

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: Four levels of spatiotemporal retrieval.

	CAMELS			ENERGY			PEMS
Method	MAE	MAPE	Time	MAE	MAPE	Time	MAE	MAPE	Time
Global	11.73	43.65	1.53	8.85	109.76	920.53	27.27	43.54	11407.35
Node-wise	10.87	36.40	3.90	7.96	104.70	3569.83	19.29	28.12	21179.26
Time-wise	14.00	58.26	7.15	8.76	108.58	467.50	27.41	44.42	4526.34
Point-wise	16.09	60.87	8.66	10.97	139.52	1239.79	41.72	53.20	15527.86
Ours	10.92	37.81	1.83	8.04	104.02	957.53	19.02	27.93	12377.27

Appendix

Table 6: Comparison of different spatiotemporal retrieval strategies.

Dataset	Domain	#Nodes	#Steps	Resolution	Time Span	Feature
CAMELS	Hydrology	7	12699	1 day	1980-01-01 – 2014-10-07	Flow
PEMS	Traffic	581	106848	5 min	2011–2017 (07-10 to 08-31)	Flow
ENERGY	Energy	103	34560	10 min	240 days	Wind speed

Appendix

Table 7: Detailed statistics of datasets.

Metric	ShiftTS	STONE	TFPS	STRAP	REMIND
MAE	16.84 $\pm$ 0.21	18.65 $\pm$ 0.28	17.29 $\pm$ 0.23	16.60 $\pm$ 0.20	16.63 $\pm$ 0.19
RMSE	27.82 $\pm$ 0.35	30.53 $\pm$ 0.42	30.31 $\pm$ 0.40	26.32 $\pm$ 0.34	26.88 $\pm$ 0.32
MAPE	18.00 $\pm$ 0.48	21.21 $\pm$ 0.65	17.69 $\pm$ 0.46	20.24 $\pm$ 0.55	17.76 $\pm$ 0.44

Appendix

Table 8: Comparison with Distribution Shift Methods on PEMS

Metric	ShiftTS	STONE	TFPS	STRAP	REMIND
MAE	1.76 $\pm$ 0.03	2.35 $\pm$ 0.04	1.96 $\pm$ 0.03	2.77 $\pm$ 0.05	1.72 $\pm$ 0.03
RMSE	2.43 $\pm$ 0.04	2.93 $\pm$ 0.05	2.49 $\pm$ 0.04	3.75 $\pm$ 0.06	2.39 $\pm$ 0.04
MAPE	23.82 $\pm$ 0.55	23.51 $\pm$ 0.58	22.16 $\pm$ 0.50	26.28 $\pm$ 0.70	21.49 $\pm$ 0.48

Appendix

Table 9: Comparison with Distribution Shift Methods on ENERGY

Table 15

Figure 8: Hyperparameter sensitivity on CAMELS.

Noise	MAE	RMSE	MAPE	MAE	RMSE	MAPE	Mean $g$
	STRAP			REMIND
0.0	2.77	3.75	26.28	1.72	2.39	21.49	0.74
0.1	2.95	3.96	28.10	1.76	2.44	22.10	0.66
0.2	3.24	4.31	31.47	1.84	2.54	23.38	0.53
0.4	3.78	4.93	37.85	2.02	2.76	25.91	0.36
0.6	4.31	5.58	44.92	2.23	3.02	29.14	0.24

Appendix

Table 12: Memory corruption on ENERGY.

Figure 9: Robustness to memory corruption on ENERGY.

Spatiotemporal Forecasting

Spatiotemporal forecasting predicts how quantities change across locations and time, from traffic and weather to disease and environmental systems.

See all 418 papers

Reversible Residual Normalization Alleviates Spatio-Temporal Distribution Shift
Apr 17, 2026Zhaobo Hu, Vincent Gauthier, Mehdi NaimaInvertible Neural Networks Distribution Shift Robustness
SAMOVAR, Télécom SudParis · Institut Polytechnique de Paris · CNRS – LIP6 +1
ROAR: Retrieval Opportunity-Aware Refinement for Zero-Shot Time Series Forecasting
May 24, 2026Jinjin Chi, Lei Feng, Yiming Wang +4Zero-Shot Time Series Forecasting Time Series Forecasting
College of Computer Science and Technology, Jilin University, Changchun, China · College of Computing and Data Science, Nanyang Technological University, Singapore
DeRegiME: Deep Regime Mixtures for Probabilistic Forecasting under Distribution Shift
May 19, 2026Kieran Wood, Stefan Zohren, Stephen J. RobertsDistribution Shift Robustness Time Series Forecasting
Machine Learning Research Group, University of Oxford · Oxford-Man Institute, University of Oxford

Apr 17, 2026J/K move · Enter open · S save

Reversible Residual Normalization Alleviates Spatio-Temporal Distribution Shift

Zhaobo Hu, Vincent Gauthier, Mehdi Naima

SAMOVAR, Télécom SudParis · Institut Polytechnique de Paris · CNRS – LIP6+1

Invertible Neural Networks Distribution Shift Robustness Spatiotemporal Forecasting Temporal Distribution Shift Spatiotemporal GNNs

Open paper AskPDF arXiv

Distribution shift severely degrades the performance of deep forecasting models. While this issue is well-studied for individual time series, it remains a significant challenge in the spatio-temporal domain. Effective solutions like instance normalization and its variants can mitigate temporal shifts by standardizing statistics. However, distribution shift on a graph is far more complex, involving not only the drift of individual node series but also heterogeneity across the spatial network where different nodes exhibit distinct statistical properties. To tackle this problem, we propose Reversible Residual Normalization (RRN), a novel framework that performs spatially-aware invertible transformations to address distribution shift in both spatial and temporal dimensions. Our approach integrates graph convolutional operations within invertible residual blocks, enabling adaptive normalization that respects the underlying graph structure while maintaining reversibility. By combining Center Normalization with spectral-constrained graph neural networks, our method captures and normalizes complex Spatio-Temporal relationships in a data-driven manner. The bidirectional nature of our framework allows models to learn in a normalized latent space and recover original distributional properties through inverse transformation, offering a robust and model-agnostic solution for forecasting on dynamic spatio-temporal systems.

Abstract

Spatiotemporal Forecasting

Explore similar work

Reversible Residual Normalization Alleviates Spatio-Temporal Distribution Shift

ROAR: Retrieval Opportunity-Aware Refinement for Zero-Shot Time Series Forecasting

DeRegiME: Deep Regime Mixtures for Probabilistic Forecasting under Distribution Shift