Julia SciML in Practice: Using Universal Differential Equations (UDE) with Neural Networks to Learn Unknown Epidemiological Parameters

This article introduces an open-source project based on the Julia SciML ecosystem, which combines the SIRHD epidemiological model with neural networks using Universal Differential Equations (UDE) to estimate six unknown transmission rate parameters and compare the performance of three activation functions: ReLU, Tanh, and Sigmoid. The original author of the project is digvijay1992, released on GitHub on May 29, 2026, under the MIT license.

UDE is a hybrid modeling method in the intersection of scientific computing and machine learning, combining known physical laws (such as the structure of SIRHD differential equations) with data-driven neural networks. Its advantages include: maintaining physical interpretability, reducing data requirements, better generalization ability, and direct estimation of unknown parameters. Traditional SIR models assume constant parameters, while UDE uses neural networks to dynamically estimate parameters, adapting to parameter changes in reality.

The SIRHD model tracks five groups: Susceptible (S), Infected (I), Recovered (R), Hospitalized (H), and Deceased (D). A system of differential equations describes the population dynamics. In UDE, six rate parameters (τSI, τIR, τID, τIH, τHR, τHD) are estimated by independent neural networks, each with a two-layer structure (input → hidden layer of 10 neurons → output). Comparison of three activation functions:

• ReLU: Simple computation, fast convergence, but zero gradient in the negative interval
• Tanh: Output range (-1,1), suitable for capturing positive and negative changes
• Sigmoid: Output (0,1), suitable for proportional parameters, but prone to gradient vanishing.

The project leverages the Julia SciML ecosystem, with core libraries including:

• DifferentialEquations.jl: High-performance differential equation solver
• DiffEqFlux.jl: Bridges differential equations and Flux neural networks
• Lux.jl: Neural network framework supporting automatic differentiation and GPU acceleration
• Optimization.jl: Unified optimization interface
• ComponentArrays.jl: Efficient handling of structured parameters Julia's automatic differentiation mechanism (Zygote.jl) supports end-to-end training.

Training steps:

1. Data generation: Run the SIRHD model with known true parameters for 25 days to generate time series of infected and recovered individuals
2. Model training: Train UDE with three activation functions to minimize the difference between predictions and real data
3. Performance comparison: Record convergence speed, final error, and computation time
4. Visualization: Compare predicted curves with real data The training uses the adjoint sensitivity method to efficiently compute gradients, supporting adjoint methods such as InterpolatingAdjoint and ReverseDiffVJP.

Application scenarios of the UDE method include:

• Epidemiology: Capture dynamic changes in transmission rates of epidemics like COVID-19 to assist policy-making
• Industrial systems: Modeling chemical processes and energy systems, combining physical laws with data correction
• Financial modeling: Combining known structures of asset price dynamics with neural network patterns
• Interpretable AI: Retaining physical interpretability, suitable for high-risk fields such as healthcare and finance.

The project demonstrates the core concept of scientific machine learning: using neural networks to enhance physical models rather than replace them. UDE improves expressive power while maintaining interpretability. It serves as an introductory example for Julia users and provides ML practitioners with ideas for integrating domain knowledge into models. The choice of activation function has an impact, but the hybrid modeling mindset is more critical—combining human knowledge with ML capabilities is an effective path to solving complex scientific problems.