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

• Non-stationarity creates a continuum of distribution shift, while coarse adaptation and unassessed retrieval reliability can leave forecasts miscalibrated or unreliable.Source: §1 · p. 2
• REMIND treats retrieval as an active signal of historical support, assessing retrieved-observation relevance and future-outcome consistency to guide probabilistic forecasting.Source: §1 · p. 2
• A gate using retrieval similarity and memory dispersion balances memory-calibrated prediction for mild shifts against deviation-aware prediction for severe shifts.Source: §1 · p. 2
• A hierarchical Gaussian mixture captures heterogeneous uncertainty across the branches, while memory perturbation training is introduced to improve generalization under shift.Source: §1 · p. 2
• REMIND consistently outperforms state-of-the-art methods in forecasting accuracy and probabilistic prediction quality on real-world hydrology, traffic, and energy datasets with spatiotemporal shifts.Source: §1 · p. 3

• Parametric models learn parameters of predefined distributions, whereas generative diffusion and flow models learn more expressive distributions but typically require repeated sampling.Source: §2 · p. 3
• Shift-handling approaches include invariant representations in ShifTS [20], expert selection in TFPS [15], graph interventions in STONE, and historical-memory retrieval in RAFT [21], RATD, and STRAP [5].Source: §2 · p. 3
• These approaches often do not characterize shift degree to adapt probabilistic predictions, and many focus on point estimates rather than full predictive distributions.Source: §2 · p. 3

• The spatial structure is represented by graph G = (V, E, A), where nodes are locations, edges capture their relationships, and A is the adjacency matrix.Source: §3 · p. 3
• XH denotes observations over the past T time steps, and XP denotes observations over the next T time steps across locations and variables.Source: §3 · p. 3
• Given historical observations XH and graph G, the forecasting function F produces the predictive mean and standard deviation for XP.Source: §3 · p. 3

4.1 Overview

• REMIND treats retrieval relevance and outcome consistency as signals of historical support, adapting memory use rather than treating retrieval as passive lookup.Strong support favors reusing retrieved targets; weak support shifts prediction toward the current observation with deviation-based corrections.Source: §4.1 · pp. 3–4

4.2 Coarse-to-Fine Spatiotemporal Retrieval

• Coarse-to-fine retrieval first narrows candidates globally, then matches node-level temporal patterns to capture local misalignment and enrich retrieval diversity.The framework balances retrieval quality and efficiency by combining global and node-wise temporal retrieval.Source: §4.2 · pp. 4–5

4.3 Memory-Calibrated Probabilistic Branch

• Under mild shifts, the memory-calibrated branch reuses retrieved targets through query-conditioned calibration and estimates uncertainty using representation features and memory dispersion.Its Gaussian-mixture formulation supports heterogeneous predictive distributions and subsequent fusion with the other branch.Source: §4.3 · pp. 5–6

4.4 Deviation-Aware Probabilistic Branch

• Under severe shifts, the deviation-aware branch predicts from current observations and uses differences from retrieved patterns as residual corrections.It uses more mixture components than the memory-calibrated branch to model greater ambiguity and diverse future outcomes.Source: §4.4 · p. 6

4.5 Shift-Adaptive Probabilistic Fusion

• The shift-adaptive gate balances the branches using retrieval similarity and memory dispersion; lower similarity and larger dispersion are associated with larger forecast errors and error variance.The two signals remain relatively dispersed, suggesting they capture different aspects of distribution shift.Source: §4.5 · pp. 6–7
• The final hierarchical Gaussian mixture combines branch-wise mixtures, modeling both within-branch uncertainty and disagreement between branches.The fusion weights adapt to retrieval similarity and memory dispersion.Source: §4.5 · p. 7

4.6 Training Objective and Memory Perturbation

• Training uses predictive-distribution negative log-likelihood and perturbs memory gates, retrieved observations, and targets to simulate weak retrieval support and reduce over-reliance on confident memories.The perturbations are intended to improve generalization under challenging shift conditions.Source: §4.6 · p. 8

5.1 Experimental Settings

• The experiments use CAMELS, PEMS, and ENERGY, with forecasts of the next 12 steps from the previous 12 observations.The datasets cover hydrology, transportation, and energy; results are averaged over three random seeds.Source: §5.1 · p. 8

5.2 Performance Comparison

• REMIND achieves the best performance across all datasets and evaluation metrics.The comparison includes deterministic forecasting and probabilistic prediction measures.Source: §5.2 · p. 8, Table 1 · p. 9

5.3 Ablation Study

• Removing any core branch or the shift-adaptive gate degrades performance, with the largest core-architecture drop when the deviation-aware branch is removed.Among key mechanisms, removing memory dispersion causes the largest drop; memory perturbation and Gaussian mixture modeling also contribute.Source: §5.3 · p. 9, Table 2 · p. 9, Table 3 · p. 9

5.4 Hyperparameter Study

• The hyperparameter analysis links fewer mixture components in the memory branch to structured retrieved future patterns and indicates that the deviation branch benefits from more components.The gate also benefits from soft shift-guided regularization without being overly constrained.Source: Figure 4 · p. 9, §5.4 · p. 10

5.5 Gate Behavior Analysis

• The gate generally gives more weight to memory at higher retrieval similarity and lower memory dispersion, though the relationship is not strictly monotonic.Under high similarity, increasing dispersion shifts the better-performing branch from memory to deviation, and the gate reduces memory weight; fusion outperforms both branches.Source: §5.5 · p. 10, Figure 5 · p. 10
• Across ENERGY shift tertiles, the memory branch performs better under milder shifts and the deviation branch under stronger shifts; fusion has the lowest MAE in every group.Fusion improves or matches both branches in CRPS, while memory weights decrease as MMD2—and thus shift—increases.Source: §5.5 · p. 10, Figure 6 · p. 10

• REMIND is a retrieval-guided framework for spatiotemporal forecasting.Source: §6 · p. 10
• Retrieval serves as a shift-degree signal for balancing memory reuse and deviation-aware prediction.Source: §6 · p. 10
• The paper concludes that REMIND improves probabilistic forecasting generalization under distribution shift.Source: §6 · p. 10

A Derivation of the Hierarchical Gaussian Mixture

• The derivation treats the gate and branch parameters element-wise, with the multivariate spatiotemporal case obtained by stacking per-element results.Source: A Derivation of the Hierarchical… · p. 14

A.1 From a Hierarchical Mixture to a Flat Mixture

• The gate-weighted combination of branch-wise mixtures is equivalent to a flat Gaussian mixture with Imem + Idev components.The flat form supports analysis and log-likelihood evaluation, while the hierarchical view makes intra-branch multimodality and between-branch adaptation interpretable.Source: §A.1 · p. 14, A Derivation of the Hierarchical… · p. 14

A.2 Predictive Mean

• The overall predictive mean is the gate-weighted combination of the memory and deviation branch means.Source: §A.2 · p. 15

A.3 Predictive Variance

• Predictive variance combines gate-weighted within-branch variances with a term for disagreement between branch means.The disagreement term is largest when the gate is most ambivalent, g → 1/2, so uncertainty can rise in transitional regimes even when both branches are individually confident.Source: §A.3 · p. 15

A.4 Explicit Form of the Negative Log-Likelihood

• The negative log-likelihood is evaluated on the equivalent flat mixture, with log-sum-exp used for numerical stability.The full objective adds similarity- and dispersion-aware regularizers to the negative log-likelihood.Source: §A.4 · pp. 15–16

B Discussion of Spatiotemporal Retrieval

• The four retrieval levels are global spatiotemporal, node-wise temporal, time-wise spatial, and point-wise matching.Source: B Discussion of Spatiotemporal Retrieval · p. 16, Figure 7 · p. 16
• Node-wise retrieval supports flexible local matching, while time-wise matching disrupts temporal continuity and point-wise matching ignores higher-level spatiotemporal structure.Independent node-wise retrieval can capture asynchronous regional dynamics, but increases computational and memory costs as node count grows.Source: B Discussion of Spatiotemporal Retrieval · p. 17
• REMIND’s coarse-to-fine strategy combines global and node-wise retrieval, achieving the best balance between retrieval quality and efficiency among the compared strategies.Global retrieval is efficient but imposes strict temporal alignment; node-wise retrieval is more flexible but costlier.Source: B Discussion of Spatiotemporal Retrieval · pp. 17–18, Table 6 · p. 17

C.1 Datasets

• CAMELS, PEMS, and ENERGY span hydrology, transportation, and energy, with observations at daily, five-minute, and ten-minute resolutions, respectively.The datasets cover long observation horizons and varied spatial scales, and exhibit shifts across time and conditions.Source: §C.1 · p. 18

C.2 Metrics; C.3 Implementation Details.

• REMIND is evaluated with MAE, RMSE, MAPE, CRPS, and ES, and forecasts the next 12 time steps from the previous 12 observations.Results are averaged over three random-seed runs; datasets use a 6:2:2 train-validation-test split and Z-score normalization.Source: §C.2 · p. 18, §C.3 · p. 19

C.2 Metrics

• CRPS evaluates the full predictive distribution against observed values, and lower CRPS indicates better alignment with ground truth.In practice, it is approximated using quantiles from 0.05 to 0.95.Source: §C.2 · pp. 18–19
• ES is a multivariate proper scoring rule that measures forecast accuracy and predictive dispersion; lower ES indicates more accurate and well-calibrated probabilistic forecasts.Its first term measures expected prediction-to-observation distance, while its second captures diversity among predictive samples.Source: §C.2 · p. 19

C.4 Additional Comparison with Distribution Shift Methods

• On ENERGY, REMIND consistently outperforms the compared shift-aware baselines, while remaining highly competitive across PEMS and ENERGY overall.The comparisons report deterministic metrics while REMIND also preserves probabilistic forecasting capability.Source: Table 9 · p. 19, Table 8 · p. 19, §C.4 · p. 19
• On PEMS, REMIND performs on par with the strongest deterministic baseline and achieves the best result on one reported metric.The competing methods are optimized for point prediction, whereas REMIND also learns predictive uncertainty.Source: §C.4 · pp. 19–20

C.5 Additional Ablation Study

• Removing any core component weakens CAMELS performance, with the memory branch especially important for probabilistic quality.The deviation-aware branch and gate improve robustness when basin dynamics differ from historical patterns.Source: §C.5 · p. 20, Table 10 · p. 20
• Memory dispersion supports uncertainty estimation, memory perturbation reduces sensitivity to retrieval errors, and Gaussian mixtures improve accuracy and calibration by representing multiple plausible hydrological responses.The reported gains reflect combining retrieval, shift-aware fusion, and flexible probabilistic modeling.Source: §C.5 · p. 20, Table 11 · p. 20

C.6 Additional Hyperparameter Study

• A moderate top-K performs best on CAMELS by retaining sufficient historical support without excessive irrelevant memory.Too few retrieved samples narrow the evidence; too many dilute memory relevance and increase inconsistency among retrieved futures.Source: §C.6 · p. 20, Figure 8 · p. 20
• The memory-calibrated branch favors fewer Gaussian components, while the deviation-aware branch benefits from more; moderate gate regularization best balances likelihood learning and shift-signal guidance.The memory branch can inherit modes from retrieved targets, whereas the deviation branch must construct a broader distribution from deviation information.Source: §C.6 · p. 20

C.7 Robustness to Memory Corruption

• The test-time robustness evaluation adds Gaussian noise of increasing strength to retrieved samples on ENERGY.Noise strength is measured by its standard deviation in normalized data units.Source: §C.7 · p. 20
• As retrieved-sample noise increases, forecasting errors rise for both REMIND and STRAP, but REMIND degrades more gradually.Table 12 reports forecasting errors; Figure 9 shows relative MAE increase and REMIND’s mean memory weight.Source: §C.7 · p. 21, Table 12 · p. 21, Figure 9 · p. 21
• As noise increases, REMIND’s memory weight decreases, indicating that its gate reduces direct reuse of corrupted historical targets.The authors report this as consistent with adapting memory use as retrieved-target quality declines.Source: §C.7 · p. 21

D Limitations and Broader Impacts

• Retrieval-guided adaptation may be less effective when sparse, rapidly evolving, or entirely unseen scenarios make suitable historical patterns difficult to retrieve.The method therefore depends on the availability and quality of historical retrieval memory.Source: D Limitations and Broader Impacts · p. 21
• More reliable uncertainty estimation may support safer decisions in traffic management, energy systems, and hydrology, but inaccurate estimates under unseen scenarios may affect downstream decisions.The authors emphasize careful evaluation and responsible real-world deployment.Source: D Limitations and Broader Impacts · p. 21, NeurIPS Paper Checklist · p. 26

NeurIPS Paper Checklist

• The checklist reports that experimental details are documented and an anonymous code repository is provided for review, with full code release planned upon acceptance.The paper states that dataset splits, metrics, optimization settings, hyperparameter selection, and implementation details are reported.Source: NeurIPS Paper Checklist · pp. 23–24
• Experimental results are averaged over three random-seed runs, with standard deviations reported where space permits and box plots provided to show distribution quantiles.The checklist attributes variability to random initialization and training stochasticity.Source: NeurIPS Paper Checklist · p. 25
• The authors answer N/A because the core method development does not involve LLMs as important, original, or non-standard components.The checklist also directs readers to the NeurIPS handbook’s LLM policy.Source: NeurIPS Paper Checklist · p. 28