# Pre-training Epidemic Time Series Forecasters with Compartmental Prototypes

arXiv: https://arxiv.org/abs/2502.03393 · HTML: https://mlsanity.com/paper/2502.03393/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 90 references that are in the map.
- First posted on arXiv 5 Feb 2025.

## Builds on (most cited first, 30 of 33)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [A Survey of Large Language Models](https://mlsanity.com/paper/2303.18223/) (2023)
- [Informer: Beyond Efficient Transformer for Long Sequence Time-Series Forecasting](https://mlsanity.com/paper/2012.07436/) (2020)
- [A Time Series is Worth 64 Words: Long-term Forecasting with Transformers](https://mlsanity.com/paper/2211.14730/) (2022)
- [Autoformer: Decomposition Transformers with Auto-Correlation for Long-Term Series Forecasting](https://mlsanity.com/paper/2106.13008/) (2021)
- [Are Transformers Effective for Time Series Forecasting?](https://mlsanity.com/paper/2205.13504/) (2022)
- [FEDformer: Frequency Enhanced Decomposed Transformer for Long-term Series Forecasting](https://mlsanity.com/paper/2201.12740/) (2022)
- [TimesNet: Temporal 2D-Variation Modeling for General Time Series Analysis](https://mlsanity.com/paper/2210.02186/) (2022)
- [Chronos: Learning the Language of Time Series](https://mlsanity.com/paper/2403.07815/) (2024)
- [Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution](https://mlsanity.com/paper/2202.10054/) (2022)
- [An Empirical Study of Catastrophic Forgetting in Large Language Models During Continual Fine-tuning](https://mlsanity.com/paper/2308.08747/) (2023)
- [Unified Training of Universal Time Series Forecasting Transformers](https://mlsanity.com/paper/2402.02592/) (2024)
- [MOMENT: A Family of Open Time-series Foundation Models](https://mlsanity.com/paper/2402.03885/) (2024)
- [TS2Vec: Towards Universal Representation of Time Series](https://mlsanity.com/paper/2106.10466/) (2021)
- [A Transformer-based Framework for Multivariate Time Series Representation Learning](https://mlsanity.com/paper/2010.02803/) (2020)
- [The Vendi Score: A Diversity Evaluation Metric for Machine Learning](https://mlsanity.com/paper/2210.02410/) (2022)
- [Self-Supervised Contrastive Pre-Training For Time Series via Time-Frequency Consistency](https://mlsanity.com/paper/2206.08496/) (2022)
- [Foundation Models for Time Series Analysis: A Tutorial and Survey](https://mlsanity.com/paper/2403.14735/) (2024)
- [Towards Robust Blind Face Restoration with Codebook Lookup Transformer](https://mlsanity.com/paper/2206.11253/) (2022)
- [Deep Time Series Models: A Comprehensive Survey and Benchmark](https://mlsanity.com/paper/2407.13278/) (2024)
- [Central Moment Discrepancy (CMD) for Domain-Invariant Representation Learning](https://mlsanity.com/paper/1702.08811/) (2017)
- [GMNN: Graph Markov Neural Networks](https://mlsanity.com/paper/1905.06214/) (2019)
- [Learning Fast and Slow for Online Time Series Forecasting](https://mlsanity.com/paper/2202.11672/) (2022)
- [A Review of Graph Neural Networks in Epidemic Modeling](https://mlsanity.com/paper/2403.19852/) (2024)
- [EINNs: Epidemiologically-informed Neural Networks](https://mlsanity.com/paper/2202.10446/) (2022)
- [T-Rep: Representation Learning for Time Series using Time-Embeddings](https://mlsanity.com/paper/2310.04486/) (2023)
- [Epidemiology-Aware Neural ODE with Continuous Disease Transmission Graph](https://mlsanity.com/paper/2410.00049/) (2024)
- [Gaussian kernel smoothing](https://mlsanity.com/paper/2007.09539/) (2020)
- [EpiLearn: A Python Library for Machine Learning in Epidemic Modeling](https://mlsanity.com/paper/2406.06016/) (2024)

## Built on by (most cited first, 3 of 3)

- [STReasoner: Empowering LLMs for Spatio-Temporal Reasoning in Time Series via Spatial-Aware Reinforcement Learning](https://mlsanity.com/paper/2601.03248/) (2026)
- [PPG-Distill: Efficient Photoplethysmography Signals Analysis via Foundation Model Distillation](https://mlsanity.com/paper/2509.19215/) (2025)
- [EpiEvolve: Self-Evolving Agents for Streaming Pandemic Forecasting under Regime Shifts](https://mlsanity.com/paper/2606.05513/) (2026)

## Abstract (arXiv)

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 textitcompartmental 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.