Learning Epidemiological Dynamics via the Finite Expression Method

TL;DR

The paper presents the Finite Expression Method (FEX), a symbolic-learning framework that uses reinforcement learning to extract explicit, interpretable governing equations for epidemiological dynamics from data. By approximating the vector field in $\frac{d\mathbf{x}}{dt}=f(\mathbf{x})$ with a surrogate $\phi_{\rm FEX}(\mathbf{x})$, trained via Euler integration to minimize a per-step loss, FEX yields compact, human-readable expressions for disease spread. Across synthetic SIR/SEIR/SEIRD data and real-world COVID-19 data, FEX demonstrates competitive predictive accuracy while delivering explicit equations, enabling interpretability and rapid model development. The approach shows particular strength against neural baselines and offers a practical path to interpretable, data-driven epidemiological modeling, with future work addressing computational costs and solution nonuniqueness.

Abstract

Modeling and forecasting the spread of infectious diseases is essential for effective public health decision-making. Traditional epidemiological models rely on expert-defined frameworks to describe complex dynamics, while neural networks, despite their predictive power, often lack interpretability due to their ``black-box" nature. This paper introduces the Finite Expression Method, a symbolic learning framework that leverages reinforcement learning to derive explicit mathematical expressions for epidemiological dynamics. Through numerical experiments on both synthetic and real-world datasets, FEX demonstrates high accuracy in modeling and predicting disease spread, while uncovering explicit relationships among epidemiological variables. These results highlight FEX as a powerful tool for infectious disease modeling, combining interpretability with strong predictive performance to support practical applications in public health.

Figures (5)

• Figure 1: The computational structure is represented using binary trees, where each node is assigned either a binary or unary operator. Expressions are recursively constructed, starting from depth-1 trees. Binary operators include $\mathbb{B}:=\{+, -, \times, \div, \ldots\}$, and unary operators include $\mathbb{U}:=\{\sin, \exp, \log, \text{Id}, (\cdot)^2, \int\cdot\text{d} x_i, \frac{\partial\cdot}{\partial x_i}, \ldots\}$.
• Figure 2: Flowchart of the FEX method. The process consists of an iterative search loop (a), weight optimization (b), and expression generation to identify solutions for the target ODEs or PDEs. Key components include the Expression Tree, Controller, and Candidate Pool, which collaboratively refine expressions through sampling, scoring, and optimization mechanisms.
• Figure 3: Illustration of two tree structures used in the FEX implementation.
• Figure 4: Comparison of MSE over time for three methods (FEX, RNN, NN) on (a) SIR model, (b) SEIR model, and (c) SEIRD model over 250 time steps.
• Figure 5: Comparison of actual and predicted COVID-19 cases using the FEX and SEIQRDP methods: (a) active cases ($Q$), (b) deceased cases ($D$), and (c) recovered cases ($R$). FEX demonstrates superior data fitting during the training phase and accurate predictions in the testing phase, compared to SEIQRDP.