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.
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 4Figure 5
Figure 4: Hyperparameter sensitivity on ENERGY.
Figure 5: (a) Box plots of learned gate values across retrieval similarity and memory dispersion levels. (b) Pearson and Spearman correlations among learned gate values and these two signals.
Figure 6: Branch performance and memory reliance across MMD2 groups on ENERGY.
Memory branch
Deviation branch
Mean 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 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 ± 0.21
18.65 ± 0.28
17.29 ± 0.23
16.60 ± 0.20
16.63 ± 0.19
RMSE
27.82 ± 0.35
30.53 ± 0.42
30.31 ± 0.40
26.32 ± 0.34
26.88 ± 0.32
MAPE
18.00 ± 0.48
21.21 ± 0.65
17.69 ± 0.46
20.24 ± 0.55
17.76 ± 0.44
Appendix
Table 8: Comparison with Distribution Shift Methods on PEMS
Metric
ShiftTS
STONE
TFPS
STRAP
REMIND
MAE
1.76 ± 0.03
2.35 ± 0.04
1.96 ± 0.03
2.77 ± 0.05
1.72 ± 0.03
RMSE
2.43 ± 0.04
2.93 ± 0.05
2.49 ± 0.04
3.75 ± 0.06
2.39 ± 0.04
MAPE
23.82 ± 0.55
23.51 ± 0.58
22.16 ± 0.50
26.28 ± 0.70
21.49 ± 0.48
Appendix
Table 9: Comparison with Distribution Shift Methods on ENERGY
Table 15
Figure 8: Hyperparameter sensitivity on CAMELS.
STRAP
REMIND
Noise
MAE
RMSE
MAPE
MAE
RMSE
MAPE
Mean g
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.
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.
Zhaobo Hu, Vincent Gauthier, Mehdi Naima
SAMOVAR, Télécom SudParis · Institut Polytechnique de Paris · CNRS – LIP6 +1
Retrieval augmentation provides time series forecasters with historical continuations, yet even candidates that outperform the base forecast may fail to improve the final prediction. We propose ROAR, a Retrieval Opportunity-Aware Refinement framework for zero-shot time series forecasting. To better exploit these improvement opportunities, its training objective allocates additional emphasis across queries based on base-forecast difficulty and the relative improvement offered by retrieved candidates. Using this objective, ROAR first learns to aggregate aligned historical candidates and uses a learned gate to control their correction strength against a fixed base forecaster. It then jointly calibrates the forecasting module and gate to coordinate their contributions, while anchoring the combined prediction to the first-stage refined output. We derive exact decompositions of refinement gains and the opportunity-weighted training loss, and characterize additional gains from joint calibration under a local linearization. Experiments on seven benchmarks show that ROAR achieves the lowest average MSE among the compared methods. Further evaluations demonstrate average forecasting improvements across multiple backbone families and retrieval-augmented forecasters.
Jinjin Chi, Lei Feng, Yiming Wang +4
College of Computer Science and Technology, Jilin University, Changchun, China · College of Computing and Data Science, Nanyang Technological University, Singapore
We introduce DeRegiME -- Deep Regime Mixture of Experts -- a direct multi-horizon probabilistic forecaster that separates latent uncertainty regimes from the underlying signal and softly assigns each forecast location to learned recurring regimes using a sparse variational Gaussian process (GP) whose nonstationary regime-mixing kernel and Student-t likelihood combine per-regime sub-kernels and noise processes via a shared gate. This yields a single sparse-GP posterior, not a mixture of GP experts. DeRegiME addresses a key limitation of neural forecasters: point forecasts discard residual uncertainty, and probabilistic heads -- whether single marginals, uninterpreted mixtures, quantile sets, or diffusion samples -- rarely expose the regime structure of the residual. Yet distribution shift in noisy heteroskedastic time series may be abrupt, gradual, or horizon-dependent and often appears in residual uncertainty rather than the conditional mean. DeRegiME yields an interpretable mean-residual-noise decomposition with a direct-sum feature-space representation that anchors regimes as clusters of residual similarity whose transitions surface as implicit changepoints. The effective number of regimes is pruned by the stick-breaking gate. We prove kernel validity and predictive-density propriety, and across ten benchmarks and three encoder grids DeRegiME improves negative log predictive density (NLPD) by 20.3% over the strongest encoder-matched baseline, a DeepAR/GluonTS-style dynamic Student-t head, with parallel gains on CRPS (3.0%) and MSE (4.7%). Improvements are consistent across all datasets, which span abrupt, gradual, and seasonal shifts.
Kieran Wood, Stefan Zohren, Stephen J. Roberts
Machine Learning Research Group, University of Oxford · Oxford-Man Institute, University of Oxford