Pre-training Epidemic Time Series Forecasters with Compartmental Prototypes

TL;DR

CAPE is a novel pre-trained epidemic forecasting model that leverages compartmental prototypes from simulation data to improve zero-shot predictions across multiple diseases, addressing data scarcity and distribution shifts.

Contribution

It introduces the first open-source pre-trained epidemic model using compartmental prototypes, bridging epidemiological knowledge with time series forecasting.

Findings

Outperforms strong baselines on 17 disease benchmarks

Enables zero-shot epidemic forecasting with improved accuracy

Models epidemic dynamics as mixtures of latent compartmental states

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 and the design of various compartmental models 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 compartmental population states, termed \textit{compartmental prototypes}. It models a flexible dictionary of compartment prototypes directly from a large collection of simulation 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 adopts the next-token-prediction paradigm during pre-training with lightweight epidemic-aware regularization that aligns the learned prototypes with epidemiological semantics. On a comprehensive benchmark spanning 17 diseases, CAPE significantly outperforms strong baselines with zero-shot forecasting. This work represents a principled step toward pre-trained epidemic models that are both transferable and epidemiologically grounded. We provide our code in: https://github.com/nuuuh/CAPE.