Prior Knowledge-enhanced Spatio-temporal Epidemic Forecasting

1 Introduction

Population-level epidemic forecasting is of great importance to public health management, since accurate forecasting of new infections at each location can help governments and healthcare providers allocate resources, design intervention policies, and evaluate the effectiveness of control measures [Qian et al., 2021].

In the early years, constrained by data scarcity, epidemic forecasting primarily relied on mechanistic models, e.g., SIR [Kermack and McKendrick, 1927]. With the rise of deep learning, data-driven approaches have been increasingly adopted in this field [Arora et al., 2020]. Simultaneously, realizing the spatial dependencies of disease transmission among regions, spatio-temporal forecasting across multiple regions jointly has garnered great attention [Deng et al., 2020].

* Xiaoxiao Wang and Hanning Yuan are corresponding authors.

However, due to the limited epidemiological data, these models remain prone to overfitting. In recent years, there has been a surge focused on hybrid models that integrate mechanistic principles with deep learning [Cao et al., 2022; Wang et al., 2022; Han et al., 2025], e.g., employing deep neural networks to estimate the parameters of mechanistic models as shown in Figure 1(a). Nevertheless, existing methods still encounter three primary limitations:

• Insensitivity to Weak Epidemic Signals. Unlike traffic forecasting, epidemic data often contain weak signals as shown in Figure 1(b). That is, the confirmed cases in most of the time are small, but infectious diseases can surge to high prevalence levels in a short time. Existing deeplearning modules [Wu et al., 2019] are often insensitive to such weak yet critical signals, leading to missed early warnings.
• Over-simplified Adjacency Relation. Existing methods capture the spatial dependency by using the human mobil-

ity intensities among regions [Wang et al., 2018; Cao et al., 2022]. Though the vital factor, i.e., human mobility for disease transmission, is considered, they omit the intrinsic similarities among regions as shown in Figure 1(c).

• Unstable Parameter Estimation. As mentioned, hybrid models rely on neural networks to produce vital epidemic parameters. However, unconstrained network outputs often fluctuate violently in data-scarce scenarios as shown in Figure 1(d), resulting in mechanically implausible and unstable parameter estimation.

An intuitive idea to tackle those issues is to incorporate the external epidemic knowledge to regulate the modeling process. However, related data may not always be available. To this end, in this paper, we propose Spatio-Temporal priOraware Epidemic Predictor (STOEP), which incorporates two types of implicit knowledge without the requirements of external data, i.e., spatio-temporal priors and expert priors, to facilitate the epidemic modeling. The spatio-temporal priors are derived from data to enhance the representation of each region even if the signals are weak, and adjust the adjacency relation among regions in addition to the mobility data, while the expert priors are obtained from domain experts to regularize the output of estimated epidemic parameters to stabilize the epidemic forecasting. More specifically, STOEP is composed of three modules: (1) Case-aware Adjacency Learning (CAL), which dynamically fuses mobility data with patterns of confirmed cases by a learnable pattern memory to capture temporal epidemic priors dynamically; (2) Space-informed Parameter Estimating (SPE), which introduces a specialized spatial enhancement block with learnable spatial priors to explicitly amplify signals; (3) Filter-based Mechanistic Forecasting (FMF), which employs an adaptive thresholding strategy guided by expert priors to suppress noise in data-scarce regions. The main contributions are summarized as follows:

• We inject spatio-temporal priors into the epidemic forecasting model, which makes the adjacency relation aware of temporal patterns via CAL, and makes model inputs enhanced by spatial knowledge via SPE.
• We propose an expert prior-guided adaptive thresholding strategy via FMF to stabilize the epidemic forecasting by regularizing the output of estimated epidemic parameters.
• Extensive experiments on two real-world epidemic datasets, i.e., COVID-19 and influenza, demonstrate that STOEP outperforms the best baseline by 11.1% on average in RMSE. We also released the code for public use 1 .
• A system based on STOEP is deployed and used at a provincial CDC in China to facilitate downstream applications.

2 Preliminaries

2.1 Definitions

Definition 1 (Daily Confirmed Cases). We denote the number of daily confirmed cases in region n at day τ as xτn , and the number of daily confirmed cases over all regions as Xτ ∈ RN , where N is the number of regions.

1 https://github.com/stdi-lab/STOEP

Definition 2 (Historical Observations). To make the epidemic forecasting model work effectively, we divide historical observations into two types, i.e., essential factors and optional factors. The essential factors include historical daily confirmed cases Xt-(Tin-1):t , susceptible population St-(Tin-1):t ∈ RN×Tin , infected population It-(Tin-1):t ∈ RN×Tin and recovered population Rt-(Tin-1):t ∈ RN×Tin , where Tin is the given lookback time window. Other related features can also be incorporated as optional factors, e.g., day of week. We denote historical observations as Ot-(Tin-1):t ∈ RN×Tin×C , where C ≥ 4 is the number of factors considered.

Definition 3 (Historical Mobility Data). We use historical mobility data M t-(Tin-1):t ∈ RN×N×Tin to characterize the mobility intensity among different regions, where entry mτij ≥ 0 denotes the flow intensity from region i to region j at time τ ∈ {t - (Tin - 1), . . . , t}.

2.2 Problem Statement

Given the historical observations Ot-(Tin-1):t ∈ RN×Tin×C and mobility data M t-(Tin-1):t ∈ RN×N×Tin , predict the daily confirmed cases in the future Tout timesteps, i.e., ˆXt+1:t+Tout ∈ RN×Tout , which can be formulated as:

3 Methodology

STOEP follows the paradigm of hybrid spatio-temporal epidemic forecasting, which essentially leverages a mechanistic model to forecast future confirmed cases, but the epidemic parameters of the model are estimated by neural networks. The architecture of STOEP is depicted in Figure 2, which consists of three modules:

• Case-aware Adjacency Learning (CAL), which takes historical mobility data as well as daily confirmed cases, and produces a case-aware adjacency matrix. It would be used to indicate the adjacency intensities.
• Space-informed Parameter Estimating (SPE), which takes historical observations, learned spatial priors and case-aware adjacency matrix, to estimate the epidemic parameters of the mechanistic model.
• Filter-based Mechanistic Forecasting (FMF), which takes estimated epidemic parameters as well as the human mobility data, predicts the future daily confirmed cases via a filter-based mechanistic model.

The complexity and scalability of STOEP is practical for real-world use. Next, we elaborate on each module in detail.

3.1 Case-aware Adjacency Learning

Case-aware adjacency learning (CAL) aims to learn the caseaware adjacency intensities among different regions A ∈ RN×N based on the historical mobility data M t-(Tin-1):t and daily confirmed cases Xt-(Tin-1):t .

Main Idea. As mentioned, to adequately model the adjacency relation among regions, we fuse the mobility-based intensities with data-derived intensities. The main idea is to adjust the mobility-based intensities using data-derived adjacency relationship, which are detailed as follows.

Mobility-based Adjacency Learning. We first perform a linear transformation from past to future onM to obtain the future mobility prediction Ht+1:t+Tout ∈ RN×N×Tout . And the mobility-based adjacency matrix Am ∈ RN×N can be derived by mean-pooling over the prediction horizons:

Case-aware Adjacency Adjustment. For each region, its recent observation of confirmed cases may contain some noise. To quantify its correlation with other regions based on the recently observed cases, we first perform a soft clustering to obtain its robust representation, and then use the representation to calculate the intensity with other regions.

More specifically, for a historical daily confirmed case of region n, i.e., xt-(Tin-1):t n , we take its most recent S days (S ≤ Tin) to form a recent observation xt-(S-1):t n . xt-(S-1):t n would be normalized to eliminate the effect of population disparity to form an epidemic pattern ˜xt-(S-1):t n . Then the robust representation rn ∈ RdA of the current epidemic pattern of region n can be derived by attention-based pattern retrieval from a learnable pattern set P ∈ RP×S :

where αn are the attention weights of region n over P patterns, K ∈ Rd and V ∈ Rd are transformed key and value matrices of P. After that, rn would be compared with vectors in a learnable region embedding table E ∈ RN×dA to obtain its recent correlation with other regions. Formally, the case-based adjacency intensity ∆Gnm of region n to region m is:

where em ∈ RdA is entry m of E, ⟨·, ·⟩ denotes the inner product. α is a learnable scalar that scales the correction strength, which is initialized as 0 to prevent training perturbations. Finally, the case-based adjacency ∆G ∈ RN×N serves as a residual term to adjust the mobility-based adjacency Am to obtain the final adjacency matrix A:

3.2 Space-informed Parameter Estimating

Space-informed parameter estimating (SPE) aims to learn epidemic parameters of the mechanistic model. In this work, we leverage MetaSIR [Wang et al., 2018] as the mechanistic model, since it can jointly model multiple regions. In the context of MetaSIR, STOEP needs to estimate the infection rate β ∈ RN×Tout and recovery rate γ ∈ RN×Tout of different regions in future time steps.

Main Idea. To amplify the weak signals in epidemic data, we aim to explicitly learn some stable spatial dependency, which would be fused with dynamic dependency derived from historical observations, to comprehensively enhance the input features. After that, the enhanced features are fed into GraphWaveNet [Wu et al., 2019] to estimate β and γ, which are detailed as follows.

Dynamic Dependency Generation. We first transform the historical observations Ot-(Tin-1):t to a hidden representation Xc ∈ RN×Tin×Cc by a 1 × 1 convolution. We compute the inter-region dynamic dependency attention map Anode ∈ RN×N by applying multi-head self-attention to Xc for every input step, and averaging the scores over time and heads:

where H is the number of heads.

Spatial Prior Incorporation. To capture stable and longrange spatial dependency, we introduce a learnable score matrix S ∈ RN×N that encodes a global spatial prior. S is initialized as an identity matrix and optimized during the training. We apply a row-wise softmax on S to obtain a static dependency attention map Astruct = softmaxrow(S).

We use a gating mechanism to fuse the dynamic dependency and the static dependency to obtain a fused attention map Afuse ∈ RN×N as follows:

where [·; ·] is the concatenation, ϕ(·) is a linear layer to produce a gate Γ ∈ RN×N , σ(·) is sigmoid activation, and ⊙ is Hadamard product. After that, we regularize Afuse by symmetrizing, enforcing nonnegativity, and row-normalizing to obtain a refined attention map Afuse norm ∈ RN×N .

Feature Enhancement. Based on Afuse norm , we enhance the features Xc in a residual connection manner:

where ˜Ot-(Tin-1):t ∈ RN×Tin×Cc is the feature enhanced input, and ϵ ∈ [0, 1] is a learnable scalar (initialized as 0). Parameter Estimation. After node features are enhanced, we feed ˜Ot-(Tin-1):t as well as the adjacency matrix A obtained in CAL into GraphWaveNet-style backbone [Wu et al., 2019] for spatio-temporal modeling to obtain the hidden representation Z ∈ RN×Tout×d :

Finally, two linear heads take the hidden representation and predict the epidemiological parameters β ∈ RN×Tout and γ ∈ RN×Tout for each node on future time steps:

3.3 Filter-basedMechanistic Forecasting

Filter-based mechanistic forecasting (FMF) aims to forecast future daily confirmed cases ˆXt+1:t+Tout based on the estimated epidemic parameters β and γ from SPE via a filterbased MetaSIR.

Main Idea. To mitigate the instability of epidemic parameter estimation, by discussing with the domain experts, we suppress the disease transmission process if estimated epidemic parameters or infected population are small. A straightforward idea is to set a constant threshold to suppress the epidemic parameters. However, due to the heterogeneity across space and time, we propose an adaptive thresholding strategy to flexibly decide whether the disease transmission process should be suppressed or not, which are detailed as follows.

Adaptive Thresholding Strategy. Adaptive thresholding strategy aims to learn an adaptive threshold given a set of numbers, e.g., {vr}Rr=1 . We first compute the quantile κ = Quantile ({vr}Rr=1 , κ), where κ is the desired quantile ratio. During the training, κ would be smoothed by Exponential Moving Average (EMA) to control the adaptivity-stability trade-off. The adaptive threshold is defined as

where δmin is a safeguard lower bound to avoid vanishing thresholds when the quantile becomes very small.

Small Parameter Detection. Small parameter detection aims to detect whether the estimated epidemic parameters βn ∈ RTout and γn ∈ RTout of a region n are sufficiently small. For each region n, we first obtain its maximum parameter estimates over the prediction horizon, i.e., βmax n and γmax n . We then apply the adaptive thresholding over the maximum parameter set of all regions, i.e., {βmax n }N and {γmax n }N with quantile κβ and κγ , to obtain adaptive thresholds βthr and γthr. A region with both small βmax n and γmax n would be detected:

Low Infection Detection. Low infection detection aims to detect whether the infected population of a region is sufficiently small. We first apply the adaptive thresholding on It-(Tin-1):t with quantile κI to obtain a global nearzero threshold Iabs. Then, for each region n, the proportion of near-zero infection days can be defined as ρn = 1 Tin ∑tτ=t-(Tin-1) I(I τn ≤ Iabs). We then apply the adaptive thresholding over {ρn}N with quantile κρ to obtain ̂κZR. To maintain suppression sensitivity, we cap the threshold by κmax ρ , i.e., ̂κZR ← min(̂κZR, κmax ρ ). A region with proportion of near-zero infection days larger than ̂κZR would be detected:

Suppressed Metapopulation SIR. The final suppression filter is defined as the union of two detection results, i.e., fn = f quiet n ∨ f small n . We apply the filter to adaptively suppress the disease transmission process by modifying the infection parameter βn in the MetaSIR model [Wang et al., 2018], which is achieved as follows:

where ψ ∈ (0, 1) is a small constant for downscaling. Based on the revised infection update equation, we can obtain the daily confirmed cases prediction according to MetaSIR [Wang et al., 2018]. During training, the model is optimized via Mean Absolute Error (MAE) loss.

4 Experiments

We evaluate the performance of the proposed method on two real-world epidemic datasets.

4.1 Datasets

To validate the applicability of the proposed model, we conducted evaluation on two types of epidemic datasets with unique characteristics, which are summarized as follows:

• COVID-19 Dataset. This dataset contains daily confirmed COVID-19 cases in 47 prefectures in Japan covering the period from April 1, 2020, to September 21, 2021 (539 days)

from the NHK COVID-19 2 . Other related observations are from the Ministry of Health, Labour and Welfare, etc. The mobility data are from Facebook Movement Range Maps 3 . The entire dataset is publicly available 4 .

• Flu Dataset. This dataset contains daily confirmed flu cases in 11 cities of Zhejiang Province, China, during 2023-2024 (731 days). We can derive the daily susceptible/infected/recovered population from it in a way similar to that used for the COVID-19 dataset. The mobility data are from Baidu Migration Map 5 .

For each dataset, we partition it into non-overlapping training, validation, and test sets chronologically at a 6:1:1 ratio.

4.2 Experimental Settings

Baselines. We compare STOEP with the following baselines:

• Mechanistic Models: SIR [Kermack and McKendrick, 1927] and MetaSIR [Wang et al., 2018].
• Deep Learning Models: (1) Time-series Models: AMD [Hu et al., 2025], TimeKAN [Huang et al., 2025], DUET [Qiu et al., 2025], LightGTS [Wang et al., 2025]; (2) Spatio-temporal Graph Models: DCRNN [Li et al., 2018], STGCN [Yu et al., 2018], GraphWaveNet [Wu et al., 2019], ColaGNN [Deng et al., 2020], MTGNN [Wu et al., 2020], AGCRN [Bai et al., 2020], CovidGNN [Kapoor et al., 2020], PopNet [Gao et al., 2022], STDDE [Long et al., 2024].

2 https://www3.nhk.or.jp/news/special/coronavirus/data/

3 https://data.humdata.org/dataset/movement-range-maps

4 https://github.com/deepkashiwa20/MepoGNN

5 https://qianxi.baidu.com/

• Hybrid Models: STAN [Gao et al., 2021], MPSTAN [Mao et al., 2024], CausalGNN [Wang et al., 2022], MepoGNN [Cao et al., 2022], EISTGNN [Han et al., 2025].

Evaluation Methods. Following previous work [Tang et al., 2023; Cao et al., 2022], we evaluate the 3d, 7d, 14d, and overall prediction performance of our model by four metrics, i.e., Root Mean Squared Error (RMSE), Mean Absolute Error (MAE), Symmetric Mean Absolute Percentage Error (SMAPE), and Relative Absolute Error (RAE).

Implementation & Training Settings. We implement STOEP and baselines using PyTorch on an NVIDIA RTX 5060 GPU. The model is optimized via Adam with a curriculum learning strategy.

4.3 Overall Evaluation

Table 1 summarizes the results of different methods on the COVID-19 dataset against all baselines. STOEP consistently achieves the best performance across almost all horizons and metrics. On the overall results, it obtains 10.3% lower RMSE, 8.9% lower MAE, and 8.7% lower RAE than MepoGNN, highlighting a substantial improvement in both accuracy and robustness. The advantage is particularly evident on the 7day horizon, where STOEP achieves 15.9% lower RMSE and 13.3% lower MAE, demonstrating its superiority in mid-term forecasting.

Table 2 summarizes results on the Flu dataset. STOEP again delivers the best performance across all horizons and metrics. Compared with MepoGNN on the overall results, it achieves 11.8% lower RMSE, 9.0% lower MAE and 7.5% lower RAE. On the 7-day horizon, STOEP reduces RMSE and MAE by 10.6% and 7.5%, respectively, underscoring stable mid-term benefits in a real operational setting.

The significance test confirms that STOEP's improvements are statistically significant, with a permutation p-value < 2 × 10-4 across all baselines.

4.4 Ablation Study

We conduct ablation experiments on both datasets to quantify the contribution of each component in STOEP by comparing the following variants: w/o SPE keeps the architecture of SPE module but routes its inputs through a 1×1 convolution whose output is directly fed into GraphWaveNet for parameter estimation. w/o CAL removes the case-based adjacency intensity. CAL-C uses k-shape clustering [Paparrizos and Gravano, 2015] within the CAL module to obtain temporal patterns instead of self-learned patterns. w/o FMF replaces the FMF module with a classic MetaSIR for forecasting.

As shown in Figure 3, SPE contributes most, and STOEP lowers RMSE by 5.6% and RAE by 5.5% on the COVID-19 dataset, while on the Flu dataset, it lowers RMSE by 7.3% and RAE by 8.4%, compared to w/o SPE. The CAL module significantly reduced the RMSE by 8.3% on the Flu dataset. While the CAL-C variant performed better than the w/o CAL variant, it was inferior to STOEP on both datasets. This suggests that the self-learned patterns are more effective than using clustering. After the disease transmission suppression in FMF is removed, i.e., w/o FMF, RMSE is increased by 1.4% on the Flu dataset, which also demonstrates its effectiveness in suppressing unstable parameter estimation.

4.5 Hyperparameter Sensitivity

We further analyzed the sensitivity of the proposed model with respect to the number of learned patterns in the CAL module. Figure 4 reports the results for RMSE and RAE as the number of patterns varies from 7 to 11 on both the COVID-19 dataset and the Flu dataset. Overall, the performance is relatively stable across different settings, but performance is generally optimal at 9 patterns. These results suggest that a moderate number of patterns is sufficient to capture the temporal dynamics without introducing redundancy.

4.6 Case Study

To further demonstrate the effectiveness of STOEP, we have conducted a case study on COVID-19 for STOEP and the most competitive forecasting model, i.e., MepoGNN. We visualize the prediction results (7d ahead) of STOEP and MepoGNN from July 18, 2021, to September 19, 2021, on COVID-19 dataset in Figure 5, where the blue line represents the ground-truth. It can be observed that our model is closer to the ground-truth than MepoGNN and more capable of capturing the occurrence of peaks. For instance, STOEP successfully predicted a peak around August 13-14, a period when Delta variant was prevalent in Japan and the epidemic was expanding rapidly 6 . This indicates that our model can reflect real-world scenarios more accurately than MepoGNN.

4.7 SystemDeployment

A system based on STOEP has already been deployed at Zhejiang CDC in China, providing daily influenza forecasts for 11 cities. As shown in Figure 6, the system utilizes our model to visualize real-time trends and risk levels. Public health practitioners can use the system to trigger the early warning of infectious diseases and precisely allocate medical resources among regions in advance.

6 https://japan.kantei.go.jp/99 suga/statement/202108/ 00009. html

5 RelatedWork

Epidemic forecasting can be mainly categorized into three types: mechanistic models, data-driven models, and hybrid models.

Mechanistic Models. Mechanistic models model the infecting dynamics across different groups of populations, i.e., compartments, using differential equations, e.g., SIR model [Kermack and McKendrick, 1927; Rizoiu et al., 2018], SEIR model [Wang et al., 2024]. Considering the spatial heterogeneity, MetaSIR [Wang et al., 2018] is further proposed to jointly forecast for several regions. However, they usually rely on experts' knowledge and strong assumptions on disease transmission, which limit their prediction capability.

Data-driven Models. Data-driven models can flexibly learn to forecast based on past observations, which are widely studied in epidemic forecasting and other domains, e.g., traffic forecasting. Auto-Regressive model [Perrotta et al., 2017], ARIMA [Kufel, 2020] are initially adopted. In the era of deep learning, Bi-LSTM [Arora et al., 2020] is commonly used. For spatio-temporal forecasting, the spatial dependency among different regions can be constructed based on region adjacency [Deng et al., 2020; Kapoor et al., 2020], mobility intensity [Gao et al., 2022], or self-attention [Shen et al., 2023; Zou et al., 2026]. Though data-driven models show superior performance, they learn parameters purely based on data, which limits their generalization capability.

Hybrid Models. Hybrid models combine the mechanistic models and powerful modeling ability of deep learning. STAN [Gao et al., 2021] and PISID [Fujita and Akutsu, 2025] regulate the network training loss with SIR constraints, while MPSTAN [Mao et al., 2024] adds the MetaSIR module into the recurrent cell. MepoGNN [Cao et al., 2022], CausalGNN [Wang et al., 2022] and EISTGNN [Han et al., 2025] estimate epidemic parameters of mechanistic models by neural networks. The hybrid idea is also used in other domains, for example, CTENet [Zhang et al., 2026] fuses the physical and chemical terms into the decoder for air quality forecasting, while GravityFormer [Wang et al., 2026] fuses the gravity model for human activity intensity prediction. Our proposed method also belongs to this category. However, different from previous work, we enhance the models with spatio-temporal priors via CAL and SPE, and expert priors via FMF.

6 Conclusion

In this paper, we propose STOEP, a spatio-temporal prioraware epidemic predictor, which enhances the hybrid model by prior knowledge. Spatio-temporal priors are injected into the adjacency learning and amplify the input signals, while expert priors guide the design of parameter adaptive thresholding. Experiments on COVID-19 and influenza datasets show that STOEP consistently outperforms the best baseline by 11.1% on average in RMSE. A system based on STOEP is also deployed at Zhejiang CDC in China for forecasting to facilitate downstream applications. Further investigation into the cross-regional and cross-disease transferability of the proposed model is warranted.

Acknowledgements

This work was supported by the National Key R&D Program of China (No. 2023YFC2308703), the National Natural Science Foundation of China (No. 62306033, 42371480).