Integrating Kolmogorov-Arnold Networks with Ordinary Differential Equations for Efficient, Interpretable and Robust Deep Learning: A Case Study in the Epidemiology of Infectious Diseases

In this study, we extend the universal differential equation (UDE) framework by integrating Kolmogorov-Arnold Network (KAN) with ordinary differential equations (ODEs), herein referred to as KAN-UDE models, to achieve efficient and interpretable deep learning for complex systems. Our case study centers on the epidemiology of emerging infectious diseases. We develop an efficient algorithm to train our proposed KAN-UDE models using time series data generated by traditional SIR models. Compared to the UDE based on multi-layer perceptrons (MLPs), training KAN-UDE models shows significantly improves fitting performance in terms of the accuracy, as evidenced by a rapid and substantial reduction in the loss. Additionally, using KAN, we accurately reconstruct the nonlinear functions represented by neural networks in the KAN-UDE models across four distinct models with varying incidence rates, which is robustness in terms of using a subset of time series data to train the model. This approach enables an interpretable learning process, as KAN-UDE models were reconstructed to fully mechanistic models (RMMs). While KAN-UDE models perform well in short-term prediction when trained on a subset of the data, they exhibit lower robustness and accuracy when real-world data randomness is considered. In contrast, RMMs predict epidemic trends robustly and with high accuracy over much longer time windows (i.e., long-term prediction), as KAN precisely reconstructs the mechanistic functions despite data randomness. This highlights the importance of interpretable learning in reconstructing the mechanistic forms of complex functions. Although our validation focused on the transmission dynamics of emerging infectious diseases, the promising results suggest that KAN-UDEs have broad applicability across various fields.