Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Pre-training Epidemic Time Series Forecasters with Compartmental Prototypes

Zewen Liu, Juntong Ni, Max S. Y. Lau, Wei Jin

TL;DR

CAP E tackles brittle epidemic forecasters by learning a dictionary of $K$ latent compartment prototypes and expressing each outbreak as a time-varying mixture over these prototypes. It combines self-supervised pre-training with epidemic-aware regularizers to promote robust generalization under data scarcity and distribution shifts. Across a comprehensive benchmark spanning 17 diseases and 50+ regions, CAP E achieves superior zero-shot, few-shot, and full-shot forecasting, demonstrating strong transferability and epidemiological grounding. The work provides an open-source framework and insights into how latent compartment dynamics can yield transferable epidemic forecasts with practical public-health impact.

Abstract

Accurate epidemic forecasting is crucial for outbreak preparedness, but existing data-driven models are often brittle. Typically trained on a single pathogen, they struggle with data scarcity during new outbreaks and fail under distribution shifts caused by viral evolution or interventions. However, decades of surveillance data from diverse diseases offer an untapped source of transferable knowledge. To leverage the collective lessons from history, we propose CAPE, the first open-source pre-trained model for epidemic forecasting. Unlike existing time series foundation models that overlook epidemiological challenges, CAPE models epidemic dynamics as mixtures of latent population states, termed compartmental prototypes. It discovers a flexible dictionary of compartment prototypes directly from surveillance data, enabling each outbreak to be expressed as a time-varying mixture that links observed infections to latent population states. To promote robust generalization, CAPE combines self-supervised pre-training objectives with lightweight epidemic-aware regularizers that align the learned prototypes with epidemiological semantics. On a comprehensive benchmark spanning 17 diseases and 50+ regions, CAPE significantly outperforms strong baselines in zero-shot, few-shot, and full-shot forecasting. This work represents a principled step toward pre-trained epidemic models that are both transferable and epidemiologically grounded.

Pre-training Epidemic Time Series Forecasters with Compartmental Prototypes

TL;DR

CAP E tackles brittle epidemic forecasters by learning a dictionary of latent compartment prototypes and expressing each outbreak as a time-varying mixture over these prototypes. It combines self-supervised pre-training with epidemic-aware regularizers to promote robust generalization under data scarcity and distribution shifts. Across a comprehensive benchmark spanning 17 diseases and 50+ regions, CAP E achieves superior zero-shot, few-shot, and full-shot forecasting, demonstrating strong transferability and epidemiological grounding. The work provides an open-source framework and insights into how latent compartment dynamics can yield transferable epidemic forecasts with practical public-health impact.

Abstract

Accurate epidemic forecasting is crucial for outbreak preparedness, but existing data-driven models are often brittle. Typically trained on a single pathogen, they struggle with data scarcity during new outbreaks and fail under distribution shifts caused by viral evolution or interventions. However, decades of surveillance data from diverse diseases offer an untapped source of transferable knowledge. To leverage the collective lessons from history, we propose CAPE, the first open-source pre-trained model for epidemic forecasting. Unlike existing time series foundation models that overlook epidemiological challenges, CAPE models epidemic dynamics as mixtures of latent population states, termed compartmental prototypes. It discovers a flexible dictionary of compartment prototypes directly from surveillance data, enabling each outbreak to be expressed as a time-varying mixture that links observed infections to latent population states. To promote robust generalization, CAPE combines self-supervised pre-training objectives with lightweight epidemic-aware regularizers that align the learned prototypes with epidemiological semantics. On a comprehensive benchmark spanning 17 diseases and 50+ regions, CAPE significantly outperforms strong baselines in zero-shot, few-shot, and full-shot forecasting. This work represents a principled step toward pre-trained epidemic models that are both transferable and epidemiologically grounded.

Paper Structure

This paper contains 36 sections, 1 theorem, 28 equations, 15 figures, 10 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work and Problem Definition
  3. Proposed Method
  4. Modeling latent compartmental prototypes
  5. Self-supervised Pre-training
  6. Epidemic-aware Regularization
  7. Optimization Scheme
  8. Experiment
  9. Experiment Setup
  10. Simulation on mechanistic model
  11. Fine-Tuning (Full-Shot Setting)
  12. Data Scarce Scenario, and Ablation Study
  13. Deeper Analysis
  14. Conclusion
  15. Appendix
  16. ...and 21 more sections

Key Result

Theorem A.1

Let $F,V \in \mathbb{C}^{n\times n}$ with $V$ invertible, and define where $\rho(\cdot)$ denotes the spectral radius, $\|\cdot\|_2$ the operator $2$-norm, and $\sigma_{\max}(\cdot), \sigma_{\min}(\cdot)$ the maximal and minimal singular values, respectively. Then $R_0$ satisfies the bounds

Figures (15)

  • Figure 1: (a) CAPE encoder with latent compartment prototypes; (b) Hierarchical contrasting for temporal representations; (c) Random masking and reconstruction; (d) Optimizing the encoder and prototype representations alternatively. (e) Epidemic-aware regularization, including losses for monotonic and non-monotonic dynamics.
  • Figure 2: Simulation on SIRD model. Left: Ground-truth trajectory; Right: Inferred compartment contributions.
  • Figure 3: Downstream performance vs. compartments distribution(See \ref{['app: ED']}).
  • Figure 4: Downstream performance when the model is pre-trained with either respiratory or non-respiratory data.
  • Figure 5: DBI between the embeddings of each pair of downstream datasets from the pre-trained model. (See Figure \ref{['fig: representation_']} for visualization.).
  • ...and 10 more figures

Theorems & Definitions (3)

  • Theorem A.1: Bounds for $R_0$
  • proof
  • proof