Combining mechanistic models with machine learning
Estimated time: 2-3 hours | Advanced session | Requires: Comfort with ODE-based models (e.g. the first session) and basic familiarity with neural networks (what layers and activation functions are)
Do this session if you want to discover unknown mechanisms from data using neural networks, or you have dynamics that simple parametric models can’t capture. Assumes basic familiarity with neural networks.
UDEs rely on differentiable programming — automatic differentiation through ODE solvers and neural networks — which is not practically available in R. The Julia ecosystem (SciML, Lux.jl) is currently the most mature platform for this kind of work. If you are an R user, this session is still worth reading for the concepts, but the code cannot be straightforwardly translated.
We discuss when UDEs are appropriate (and when simpler methods suffice) in When to use UDEs.
In the pMCMC session, we formally compared two mechanistic models for the Tristan da Cunha two-wave outbreak: SEITL and SEIT4L. Both encode the same biological hypothesis - temporary immunity followed by potential reinfection - but with different assumptions about timing. The original study compared them with an information criterion and found SEIT4L (Erlang-distributed immunity) the better fit.
But what if we didn’t know the mechanism? What if we were seeing a two-wave pattern for the first time and had no theory to guide us?
We might hypothesise alternative explanations:
All of these would manifest as time-varying transmission β(t) rather than structural changes to the compartmental model.
Universal Differential Equations (UDEs) combine mechanistic differential equations with neural networks to learn unknown functions from data. In our case, we’ll ask: “What transmission pattern β(t) would a simple SIR model need to explain these dynamics?”
Beyond this, UDEs can discover missing terms in equations which, with careful interpretation, can amount to learning unknown mechanisms (see Going further).
We keep the known epidemiological structure, but let a neural network learn the unknown parts:
dS/dt = -β(t) * S * I / N # ← Learn β(t) from data!
dI/dt = +β(t) * S * I / N - γ * I
dR/dt = +γ * IReplace the unknown β with a neural network that respects constraints:
β(t) = NeuralNetwork([I/N, t]) # Must be non-negativeWhat does “replace β with a neural network” actually mean? Instead of writing β = 0.5 (a fixed number) or even β(t) = β₀ * (1 + sin(t)) (a parametric function we specify), we let the neural network learn the function from data. The network takes inputs (current prevalence, time) and outputs a transmission rate. During training, we adjust the network’s internal weights so that the resulting epidemic trajectory matches observed data. The network is just a flexible function approximator - we’re not assuming any particular form for how β varies, we’re letting the data tell us.
By the end of this session, you will:
using DifferentialEquations
using Distributions
using DataFrames
using CSV
using Plots
using Random
using LinearAlgebra
using DrWatson
using Lux # Modern neural network framework (https://lux.csail.mit.edu/)
using ComponentArrays
using Optimization
using OptimizationOptimisers
using Optimisers # For Adam optimizer
Random.seed!(1234)TaskLocalRNG()flu_tdc = CSV.read(datadir("flu_tdc_1971.csv"), DataFrame)
scatter(
flu_tdc.time, flu_tdc.obs,
xlabel = "Time (days)", ylabel = "Daily incidence",
label = "Observed cases", markersize = 4, color = :red,
title = "Tristan da Cunha 1971: Two waves?",
legend = :topright
)
vspan!([25, 35], alpha=0.2, color=:blue, label="Second wave?")
The Tristan da Cunha dataset has only 59 data points. As we note in When to use UDEs, UDEs really benefit from hundreds to thousands of observations — for 59 points, simpler approaches such as Gaussian processes or splines are likely more appropriate. We use TdC here because it is familiar from earlier sessions and its two-wave pattern poses an interesting modelling challenge, not because it is a good candidate for UDEs in practice. Keep this limitation in mind when interpreting results.
In a real analysis you should always validate a UDE on synthetic data first — generate observations from a known β(t), train the UDE, and check that it recovers the true function. This confirms the method works before you apply it to data where the truth is unknown. We skip this step here for time, but see Going further → Validation with synthetic data for an outline.
Notice the two distinct periods of transmission:
Question: What transmission pattern β(t) would produce this? Can we learn it from the data?
UDEs use gradient-based optimisation rather than Bayesian inference, so the prior/posterior predictive workflow from Model Checking doesn’t directly apply. In practice, you should still validate UDE predictions against held-out data and check sensitivity to initialisation (as we do below).
First, let’s establish what a constant transmission rate gives us:
"""
Standard SIR model with constant transmission.
"""
function sir_constant_beta(du, u, p, t)
S, I, R = u
β, γ = p
N = S + I + R
du[1] = -β * S * I / N
du[2] = β * S * I / N - γ * I
du[3] = γ * I
end
# Literature-informed values for influenza (R_0 ≈ 2.5, D_inf ≈ 2 days)
R_0_guess = 2.5
D_inf_guess = 2.0
γ_guess = 1.0 / D_inf_guess
β_guess = R_0_guess * γ_guess
## The island as the ship lands, at the start of day 1. Of the five islanders it
## brought back, three had been ill during the eight day voyage and are past their
## infectious period, so they start recovered; the two who fell ill on landing
## start infectious. N = 284, the island population
u0 = [279.0, 2.0, 3.0]
tspan = (0.0, maximum(flu_tdc.time))
## `times` starts at 0.0, the left edge of the first observation day, because the
## incidence over that day is S(0) - S(1) and there is no way to form it without a
## point at 0. Differencing the 60 saved values then gives one incidence per
## observed day, already lined up with the days the observations count
times = [0.0; flu_tdc.time]
prob_constant = ODEProblem(sir_constant_beta, u0, tspan, [β_guess, γ_guess])
sol_constant = solve(prob_constant, Tsit5(), saveat=times)
# Extract incidence
S_values = sol_constant[1, :]
incidence_constant = -diff(S_values) ## 59 daily values from the 60 grid points
# Plot
plot(flu_tdc.time, incidence_constant, label="Constant β", linewidth=2)
scatter!(flu_tdc.time, flu_tdc.obs, label="Observed",
markersize=4, color=:red,
xlabel="Time (days)", ylabel="Cases",
title="Baseline: Can we do better?")
Question: The fit isn’t perfect. Could a time-varying β(t) do better?
We’ll use a small network - remember, we only have 59 data points!
# The time normalisation used by the network's second input. It is a constant of
# the problem, so it is defined here rather than carried in the parameter vector.
const TMAX = maximum(times)
# Network architecture: 2 inputs → 8 hidden → 1 output
# Total parameters: ~33 (keeps training fast)
rng = Random.default_rng()
Random.seed!(rng, 1234)
# Define network structure - small for fast ForwardDiff autodiff
nn_architecture = Chain(
Dense(2 => 8, tanh), # Input: [I/N, t/TMAX]
Dense(8 => 1, softplus) # Output: β > 0 (softplus ensures positivity)
)
# Initialise parameters
p_nn, st_nn = Lux.setup(rng, nn_architecture)
p_nn = ComponentArray(p_nn) # Convert to ComponentArray for optimization
println("Neural network parameters: ", length(p_nn))
println("Data points: ", length(flu_tdc.obs))
println("Ratio: ", round(length(flu_tdc.obs) / length(p_nn), digits=2), " data points per parameter")Neural network parameters: 33
Data points: 59
Ratio: 1.79 data points per parameterChain(...) - combines multiple neural network layers in sequence. Input flows through each layer in order.Dense(in => out, activation) - a fully-connected layer with in inputs, out outputs, and an activation function applied element-wise.tanh - activation function that squashes values to [-1, 1]. Helps learn smooth functions and keeps gradients bounded.softplus - defined as log(1 + exp(x)), ensures outputs are always positive. Smoother than ReLU at zero.Lux.setup(rng, model) - initialises the neural network parameters p and state st. Lux separates parameters from architecture for functional programming.ComponentArray - wraps nested parameter structures into a flat array while preserving named access. Needed for gradient-based optimisation.A few things to note about this architecture (expand the code explainer above for definitions of each function):
Why 8 hidden units? It’s a balance: more units = more flexibility but more parameters to learn from limited data. With 59 data points, we keep the network small.
Why softplus output? It ensures β is always positive — transmission rates can’t be negative.
Why 2 inputs? We feed the network both prevalence (I/N) and normalised time (t/TMAX). This lets β depend on both “how bad is the epidemic right now” and “what day is it” — capturing both behavioural responses to prevalence and time-dependent changes like weekends or policy changes.
Parameter count: The first layer has 2×8 weights + 8 biases = 24 parameters. The second has 8×1 weights + 1 bias = 9 parameters. Total: 33 parameters for 59 data points — roughly 2 data points per parameter. In a standard neural network, this would be a recipe for overfitting, but in a UDE the network output has to flow through an ODE solver and produce a physically plausible epidemic trajectory. The network can’t just memorise noise because arbitrary weight values would produce nonsensical dynamics (Rackauckas et al. 2020).
"""
SIR model where β is learned by neural network.
β = NN([I/N, t/TMAX]) predicts transmission rate.
"""
function sir_ude!(du, u, p, t)
S, I, R = u
N = S + I + R
γ = γ_guess # Fixed here for simplicity; in practice you'd estimate γ jointly
# Neural network predicts β(t, I)
# Inputs: prevalence (I/N) and normalised time (t/TMAX). Both γ and TMAX are
# constants of the problem rather than things to be learned, so they are
# closed over here rather than carried in the parameter vector, which would
# hand them to the optimiser.
nn_input = [I/N, t/TMAX]
# Evaluate neural network
β_learned, _ = nn_architecture(nn_input, p.nn_params, st_nn)
β = β_learned[1] # Extract scalar
# Standard SIR dynamics with learned β
du[1] = -β * S * I / N
du[2] = β * S * I / N - γ * I
du[3] = γ * I
endMain.Notebook.sir_ude!# Built once and reused; training calls this a few thousand times
ude_prob = ODEProblem(sir_ude!, u0, tspan, ComponentArray{Float64}(nn_params = p_nn))
"""
Predict incidence using current UDE parameters.
"""
function predict_ude(p_all)
# Promote u0 to the parameter element type so ForwardDiff Duals
# propagate through the ODE solve into the loss; without this the
# solver returns a plain Float64 trajectory and gradient extraction
# trips a Lux/ComponentArrays edge case.
T = eltype(p_all)
u0_promoted = T.(u0)
prob_ude = remake(ude_prob, u0 = u0_promoted, p = p_all)
sol_ude = solve(prob_ude, Tsit5(), saveat=times, dense=false)
# `retcode` is a SciMLBase enum, not a Symbol: comparing it against
# `:Success` is true even for a successful solve, which would send every
# prediction to NaN and flatten the loss.
if !SciMLBase.successful_retcode(sol_ude)
return fill(T(NaN), length(times) - 1)
end
S_pred = sol_ude[1, :]
return -diff(S_pred) ## 59 daily values from the 60 grid points
end
"""
Loss function using Poisson negative log-likelihood.
"""
function loss_ude(p_all, _)
pred = predict_ude(p_all)
# Return large loss if ODE solver failed (pred contains NaN)
any(isnan, pred) && return eltype(pred)(1e10)
# Poisson negative log-likelihood (clamp predictions to avoid log(negative))
ε = 1e-6
obs = flu_tdc.obs
pred_safe = max.(pred, ε)
nll = sum(pred_safe .- obs .* log.(pred_safe))
return nll
endMain.Notebook.loss_udeWhy Poisson and not sum of squared errors? As we learned in the observation models session, choosing a likelihood is really choosing a distance metric. SSE treats all residuals equally — a mismatch of 10 cases at the peak (expected ~40) counts the same as at the tail (expected ~2). Poisson NLL instead penalises relative errors, so predicting 0 when 5 cases occurred is heavily penalised. This forces the model to explain the second wave, where counts are small but non-zero.
As discussed above, the ODE structure constrains the network — the learned β(t) has to produce a valid epidemic trajectory, which limits what the weights can do.
Training UDEs is hard, especially with limited data like ours (59 points, 33 parameters). You should expect:
If training produces NaN values, reduce the learning rate. If the loss doesn’t decrease, try different initialisation. See Practical challenges for detailed troubleshooting.
Training means finding network weights that minimise the loss function. This requires computing gradients of the loss with respect to all 33 parameters - a task handled by automatic differentiation (autodiff).
# The network weights are the only things being optimised. Lux initialises them
# in Float32; the ODE state is Float64, so promote them to match and avoid a
# mixed-precision solve.
p_all = ComponentArray{Float64}(nn_params = p_nn)
println("Starting UDE training...")
println("="^60)
# Initial loss (before any training)
println("Initial loss: ", round(loss_ude(p_all, nothing), digits=2))
# Set up optimisation with ForwardDiff for gradient computation
opt_func = OptimizationFunction(loss_ude, Optimization.AutoForwardDiff())
opt_prob = OptimizationProblem(opt_func, p_all, nothing)
# Train using Adam optimiser — store loss history for convergence check
loss_history = Float64[]
println("\nTraining with Adam optimiser (10,000 iterations)...")
@time result = solve(opt_prob, Optimisers.Adam(0.01), maxiters=10_000,
callback = (state, l) -> begin
push!(loss_history, l)
if state.iter % 1000 == 0
println(" Iter $(state.iter): loss = $(round(l, digits=2))")
end
return false
end)Starting UDE training...
============================================================
Initial loss: 475.89
Training with Adam optimiser (10,000 iterations)...
Iter 1000: loss = -423.65
Iter 2000: loss = -423.23
Iter 3000: loss = -431.72
Iter 4000: loss = -433.82
Iter 5000: loss = -436.6
Iter 6000: loss = -440.65
Iter 7000: loss = -448.64
Iter 8000: loss = -456.75
Iter 9000: loss = -458.11
Iter 10000: loss = -458.66
Iter 10000: loss = -458.66
17.897786 seconds (207.55 M allocations: 23.108 GiB, 14.65% gc time, 47.62% compilation time)retcode: Default
u: ComponentVector{Float64}(nn_params = (layer_1 = (weight = [-1.8306004801690894 -3.183675704382478; 14.352842540463072 -12.93781372882589; … ; -5.295699326252252 -7.040691979242648; -0.33851293031150453 9.436611655661913], bias = [-0.6253509613825384, 3.6928387505473923, 3.9874005978652884, 0.11321867824862619, -0.15365428924985025, -0.13147378910313037, -0.17315140443217308, -1.6055414242017825]), layer_2 = (weight = [-0.9615651584362392 -1.7953644344400714 … -1.0721860122109812 -3.435608254050385], bias = [0.3063998943247169])))
OptimizationFunction(loss, ad_backend) - wraps a loss function with automatic differentiation for gradient computation.Optimization.AutoForwardDiff() - uses ForwardDiff.jl for forward-mode automatic differentiation. Efficient when the number of parameters is small (~33 here). For larger networks, reverse-mode (Zygote.jl) would be faster.Adam(learning_rate) - adaptive gradient descent optimiser that adjusts per-parameter learning rates and uses momentum.callback - function called after each iteration. Returns true to stop early, false to continue. Useful for monitoring progress.
println("\n" * "="^60)
println("Final loss: ", round(result.objective, digits=2))
p_trained = result.u
============================================================
Final loss: -458.66ComponentVector{Float64}(nn_params = (layer_1 = (weight = [-1.8306004801690894 -3.183675704382478; 14.352842540463072 -12.93781372882589; … ; -5.295699326252252 -7.040691979242648; -0.33851293031150453 9.436611655661913], bias = [-0.6253509613825384, 3.6928387505473923, 3.9874005978652884, 0.11321867824862619, -0.15365428924985025, -0.13147378910313037, -0.17315140443217308, -1.6055414242017825]), layer_2 = (weight = [-0.9615651584362392 -1.7953644344400714 … -1.0721860122109812 -3.435608254050385], bias = [0.3063998943247169])))
Always check the loss curve, and be honest with yourself about what it shows.
Ours drops steeply over the first ten iterations, rebounds, then descends slowly and unevenly before flattening out. Running to 10,000 iterations gets the loss to within about three and a half units of where 40,000 leaves it, so this is close enough to converged to interpret. Stopping early here would have been actively misleading rather than merely imprecise: at 500 iterations the loss is around -340 and the fitted trajectory peaks at about 11 cases, which is a different qualitative answer from the one below.
The middle of the curve is worth a second glance too. Progress slows enough around iterations 50 to 200 to look like a plateau, and stopping there would have given a much worse fit while appearing converged.
plot(loss_history, xlabel="Iteration", ylabel="Loss",
label="Training loss", linewidth=2,
title="Convergence check")
A few things to understand about this training process:
Automatic differentiation computes exact gradients of the loss with respect to parameters. We use ForwardDiff.jl, which is efficient for small networks like ours (33 parameters).
The Adam optimiser is a gradient descent variant that adapts learning rates for each parameter and uses momentum to smooth updates. The learning rate (0.01) controls step size: too large and training becomes unstable; too small and convergence is slow. Adam is widely used for neural network training because it works well across many problems without much tuning.
What to expect: the loss should fall steadily and substantially. Do not expect it to fall by orders of magnitude here, and do not use that as your check: a Poisson negative log-likelihood is not bounded below by zero, and ours goes from about +476 to about -460, so ratios of it mean nothing. Watch the shape instead. A loss that sits exactly flat means no gradient is reaching the parameters, which is a bug rather than a converged fit; NaN usually means the learning rate is too high; and a curve still falling steeply at the end means run for longer.
# Get UDE predictions
pred_ude = predict_ude(p_trained)
# Plot comparison
plot(flu_tdc.time, incidence_constant, label="Constant β", linewidth=2,
xlabel="Time (days)", ylabel="Cases",
title="UDE vs Constant β")
plot!(flu_tdc.time, pred_ude, label="UDE (learned β)", linewidth=2)
scatter!(flu_tdc.time, flu_tdc.obs, label="Observed",
markersize=4, color=:red)
Key question: what did the UDE buy, and what did it pay?
The constant-β baseline nearly matches the height of the first peak, about 40 against an observed 47, and then collapses to nothing by day 15. It has one wave to give and it spends it on the first.
The UDE reaches a similar peak, around 41, and then keeps generating cases through days 20 to 50, where the data really do have them and the baseline has none. Over that window it produces about 90 cases against the 104 observed. A time-varying β lets one model do both, which is the whole claim of the approach and the reason to prefer it here.
The loss is doing the work that makes the tail matter. Predicting nothing on a day when cases occurred is close to unbounded under a log-likelihood, so the forty days the baseline abandons cost it more than anything it gains at the peak. Swap in sum of squared errors and the incentive reverses: an empty tail is cheap and a missed peak is not, so the fit stops bothering with the second wave.
What transmission pattern did the network learn? Since the network takes both prevalence (I/N) and time as inputs, we need to evaluate it along the actual epidemic trajectory — the β(t) the ODE solver used during integration, not at some arbitrary fixed prevalence.
# Solve the trained UDE at high resolution to get the trajectory
prob_trained = ODEProblem(sir_ude!, u0, tspan, p_trained)
sol_trained = solve(prob_trained, Tsit5(), dense=true)
# Evaluate β along the actual trajectory
t_plot = range(0.0, maximum(times), length=100)
N_total = sum(u0)
β_along_trajectory = [begin
S_t, I_t, R_t = sol_trained(t)
nn_input = [I_t / N_total, t / maximum(times)]
β_val, _ = nn_architecture(nn_input, p_trained.nn_params, st_nn)
β_val[1]
end for t in t_plot]
# β on its own is hard to read, so convert to reproduction numbers: R_0(t) is
# what the learned β implies about transmissibility, and R_eff(t) accounts for
# the susceptibles actually left
R0_along = β_along_trajectory ./ γ_guess
Reff_along = [R0_along[i] * sol_trained(t_plot[i])[1] / N_total
for i in eachindex(t_plot)]
plot(t_plot, R0_along,
label="R₀(t) implied by learned β", linewidth=2,
xlabel="Time (days)", ylabel="Reproduction number",
title="What the network learned, in reproduction numbers",
legend=:topleft)
plot!(t_plot, Reff_along, label="R_eff(t)", linewidth=2)
hline!([1.0], label="R_eff = 1", linestyle=:dash, color=:black)
vspan!([25, 35], alpha=0.2, color=:blue, label="Second wave period")
Read R_eff first, because it is the quantity with a threshold. It rises to about 3 as the first wave takes off, falls to about 0.3 by day 15, and is then held a little above 1 through days 20 to 30, which sustains the second wave. Both of those moves are the network’s doing rather than the structure’s: the fall is mostly the learned β collapsing, with depletion contributing about two fifths of it, and the recovery is β climbing again while the susceptible fraction barely changes. After that it settles a little below 1 and the tail decays.
Now read R₀ against it, over the window where R_eff is actually near 1. Between days 20 and 30 the susceptible fraction falls from about 35% to 21%, and the implied R₀ climbs from about 3 to about 5 to hold R_eff up. The network is not describing a change in behaviour there; it is buying back, through transmissibility, what susceptible depletion takes away.
Past day 31, where R_eff crosses back below 1, it stops keeping up. R₀ goes on climbing, but not fast enough, and R_eff drifts down to about 0.6 by day 59 as the tail decays.
The underlying β is not monotone either: it rises through the first wave, falls back as that wave dies, and only then climbs. Plotting it alone hides both the threshold and the implausibility, which is why the reproduction numbers are the better lens.
Be careful about what that shape can be attributed to. Susceptible depletion is already in the model: the S factor in -β * S * I / N accounts for it, and R_eff is R_0 multiplied by that same susceptible fraction. So a β that has to climb is not the network describing depletion. It is compensating for what the structure cannot do once the susceptibles are gone, which is the argument the next section takes up.
Resist reading behaviour into the early shape as well. β climbs from 0.28 to about 1.7 over the first six days while prevalence is still rising, turns down before prevalence peaks on day 8, and is back to roughly its starting value by day 14 once the wave has largely passed. Transmission rising as the epidemic grows and falling as it fades is not the signature of a population reacting to visible cases.
In any case the training data cannot separate the two candidate explanations. They trace a single path through the network’s two-dimensional input space, one prevalence for each day, so dependence on prevalence and dependence on time are indistinguishable from this fit. Ask the network what β would be at a prevalence the epidemic never reached on a given day and it will answer, but nothing in the data informed that answer.
Treat the values at the far right with care. By then susceptibles are nearly exhausted, so β has very little effect on incidence and the data barely constrain it: training longer moves the final value substantially while leaving the fitted trajectory and the loss almost unchanged. A parameter the data cannot see is a parameter you should not interpret.
The exact shape varies between runs (see Uncertainty quantification below), so treat these numbers as this run’s, not the model’s.
Before interpreting this curve, look back at Step 5. The fit is good, so the question is not whether the network found something but whether what it found is believable. The next section takes that up.
Interpreting a learned β(t), in general:
A dip followed by a recovery suggests the network found that transmission has to vary to explain the shape of the data, which is consistent with behavioural relaxation: initial caution, then a return to normal contact patterns. Note that our own curve has this shape, and that the section above gives reasons not to read it that way here. A shape alone cannot establish the reading: it takes data that separate prevalence from time, which a single outbreak does not.
An implied R₀ that climbs without limit, as ours does after day 15, is a different signal. It says the network is compensating for something the structure cannot do, rather than describing a behaviour: here, a susceptible pool that only ever shrinks, which a single closed SIR compartment forces on it. Real reintroduction or waning immunity would show up this way too, and the shape alone cannot tell you which you have.
A roughly constant β(t) would say the network did not find time-varying transmission useful, which might mean the network is too small, that more iterations are needed, or genuinely that time-varying β is not the right mechanism for these data.
Everything above is a single optimisation run — the learned β(t) is a point estimate with no uncertainty attached. This is a serious limitation for drawing scientific conclusions. A quick way to get a rough sense of uncertainty is to re-run the training from several different random seeds (e.g. 5-10) and overlay the resulting β(t) curves. Where the curves agree, the signal is robust; where they diverge, the data cannot constrain the answer. See Going further → Ensemble uncertainty for more detail.
Training UDEs is difficult, especially with limited data. Understanding why helps set realistic expectations.
The fundamental tension is between flexibility and data. Our network has 33 parameters and we have 59 data points - roughly 2 points per parameter. In classical statistics, you’d want at least 10-20 observations per parameter for reliable inference. With neural networks embedded in ODEs, the situation is even more constrained because the parameters interact nonlinearly through the differential equation.
Identifiability is another challenge. Many different β(t) curves could produce similar epidemic trajectories. The data doesn’t uniquely determine β(t) — there’s a family of solutions. The ODE structure constrains this family (not every β(t) produces a valid epidemic), but it doesn’t eliminate the ambiguity entirely.
Initialisation matters because optimisation can get stuck in local minima. Different random seeds for the network weights lead to different solutions. Running multiple times and comparing results is good practice.
Numerical stability can be problematic. The combination of ODE solving and gradient computation can produce NaN values if the optimiser takes steps that are too large, or if the ODE becomes stiff (rapidly changing dynamics requiring small time steps).
The fit is decent, but look at what it cost: β has to climb to several times the constant-β baseline to keep the tail going. That price is itself a signal: maybe time-varying β is not the full story.
Consider what happens during the trough (days 13-21):
The neural network can only learn β(t), but the SIR structure itself constrains what’s possible:
In the pMCMC session, we saw that the SEIT4L model — which adds an exposed class and Erlang-distributed temporary immunity — explains the two waves through structural changes, not time-varying transmission. That is the more parsimonious explanation for this data.
What other structural changes could help?
UDEs can only discover mechanisms within the constraints we impose. If we say “it must be SIR with learned β(t)”, we’re already limiting what can be learned.
The fit itself is not the problem here, since time-varying transmission does reproduce both waves. The problem is the price it pays for that. By day 30, where the data still constrain the answer, the implied R_0 is already about 5, twice the literature value the constant-β baseline uses, and it is still rising. That is a strong hint that a different structure would explain the data more plausibly.
Overfitting occurs when the training loss becomes very low but the predictions look unrealistic — perhaps fitting noise or producing wildly oscillating β(t). With UDEs this is less common than in standard neural networks because the ODE structure constrains the output, but it can still happen with very flexible networks. The main solution is reducing network size (fewer hidden units).
Underfitting is the opposite: the UDE behaves almost identically to constant β, meaning the network hasn’t learned anything useful. The loss may still be high. Solutions include increasing network capacity (more hidden units) or running more training iterations.
NaN explosion happens when gradients become infinite, usually because the optimiser took too large a step that pushed the ODE into an unstable region. You’ll see NaN values in the loss. Solutions include reducing the learning rate (0.01 to 0.001), using gradient clipping, or switching to a more robust ODE solver that handles stiff equations.
Different random initialisations of the network weights can lead to different local minima and different learned β(t) curves. Let’s check how much this matters.
# Train from 3 different random seeds and compare learned β(t)
plot(title="Learned β(t) across random seeds",
xlabel="Time (days)", ylabel="β", legend=:topright)
for seed in [1234, 5678, 9012]
rng_test = Random.default_rng()
Random.seed!(rng_test, seed)
p_nn_test, _ = Lux.setup(rng_test, nn_architecture)
p_nn_test = ComponentArray(p_nn_test)
# Only the network weights are optimised, as above
p_test = ComponentArray{Float64}(nn_params = p_nn_test)
opt_func_test = OptimizationFunction(loss_ude, Optimization.AutoForwardDiff())
opt_prob_test = OptimizationProblem(opt_func_test, p_test, nothing)
result_test = solve(opt_prob_test, Optimisers.Adam(0.01), maxiters=10_000,
callback = (state, l) -> false)
# Extract the learned β(t) along the trajectory for comparison
p_test_trained = result_test.u
prob_test = ODEProblem(sir_ude!, u0, tspan, p_test_trained)
sol_test = solve(prob_test, Tsit5(), dense=true)
β_test = [begin
S_t, I_t, R_t = sol_test(t)
nn_input = [I_t / sum(u0), t / maximum(times)]
β_val, _ = nn_architecture(nn_input, p_test_trained.nn_params, st_nn)
β_val[1]
end for t in t_plot]
plot!(t_plot, β_test, label="Seed $seed (loss=$(round(result_test.objective, digits=1)))",
linewidth=2)
endCompare not just the final loss values but the shape of the β(t) curves. Similar losses do not guarantee similar β(t) — two different transmission patterns can produce comparably good fits. Where the curves agree, you can be more confident in the result; where they diverge, the data cannot distinguish between explanations.
Try changing one thing at a time and observe the effect on the learned β(t) curve and the final loss:
Dense(2 => 8, tanh) to Dense(2 => 4, tanh) (smaller) or Dense(2 => 16, tanh) (larger). How does the number of parameters change? Does the fit improve or does overfitting become visible?Adam(0.01) with Adam(0.001) or Adam(0.05). What happens to training stability and convergence speed?Random.seed!(1234) to a different value and re-run. How much does the learned β(t) vary between seeds? What does this tell you about identifiability?For each change, re-run from Step 1 through Step 6 and compare the β(t) plot to your original result.
In the section on structural constraints, we argue that an SIR model with learned β(t) can only sustain the second wave by driving β ever higher. Can you explain, in your own words, why sustaining the second wave requires R_0 to climb the way it does — considering both the depleted susceptible pool and near-zero prevalence during the trough? What structural change to the model would let it produce a second wave without an implausible R_0, and how does this relate to the SEIT4L model from the pMCMC session?
UDEs aren’t always appropriate. Here’s when they add value.
| Method | Flexibility | Interpretability | Data needs | Computational cost |
|---|---|---|---|---|
| Constant β | Low | High | Low (10-50 points) | Very low |
| Time-varying β with splines | Medium | Medium | Medium (50-200) | Low |
| GP for β(t) | Medium | Medium (but with UQ) | Medium (50-200) | Medium |
| UDE | High | Low | High (100-1000+) | High |
For TdC data (59 points): GPs or splines are likely more appropriate than UDEs.
Real-data applications of UDEs are still relatively few. In epidemiology, Kuwahara & Bauch (2024) used a UDE to predict COVID-19 second waves from case and mobility data, Schmid et al. (2026) combined a UDE-based SEIR model with wastewater surveillance to project case counts, and Rojas-Campos et al. (2023) demonstrated the UDE → SINDy pipeline on simulated regional transmission data. Outside epidemiology, UDEs have been applied to glucose-insulin dynamics (de Rooij et al. 2025) and battery degradation modelling (Kuzhiyil et al. 2024). Philipps et al. (2025) review the broader state of UDEs in systems biology and note that most published work still uses synthetic data.
The neural network architecture affects what patterns can be learned:
Try different architectures to understand what your data can support.
See Uncertainty quantification in the main analysis for the basic idea of training from multiple random seeds. For a more rigorous version, train 10-20 runs and compute pointwise confidence bands on β(t) from the ensemble. This is particularly important for scientific claims about mechanisms.
Before trusting UDE discoveries on real data, test on synthetic data where you know the truth:
This section introduces a related but distinct topic — symbolic regression — that extends the UDE workflow. It is included for interest and can be safely skipped.
A major limitation of neural network UDEs is interpretability - you get a function that works, but you can’t write down what it is. Symbolic regression addresses this by searching for equations rather than neural network weights.
The idea: instead of learning that β(t) is “some neural network”, discover that β(t) = β₀(1 + α·sin(ωt)) or β(t) = β₀·exp(-κI/N). The output is a human-readable formula that can be interpreted, published, and debated.
Two main approaches exist:
SINDy (Sparse Identification of Nonlinear Dynamics) builds a library of candidate terms (polynomials, trigonometric functions, etc.) and uses sparse regression to find which terms are needed. It’s fast and produces interpretable equations, but limited to the function library you define. In Julia, DataDrivenDiffEq.jl provides SINDy implementations.
Genetic programming / symbolic regression evolves equations by combining mathematical operations, selecting those that fit well. Tools like PySR (accessible from Julia via PythonCall) or SymbolicRegression.jl can discover surprising functional forms you wouldn’t have thought to include in a library.
The workflow often combines both approaches: first use a neural network UDE to confirm that a time-varying or state-dependent term improves fit, then apply symbolic regression to find an interpretable equation for that term.
For policy applications, the difference matters in principle. A hypothetical finding like “transmission increases when prevalence drops below 5%” would be actionable; “transmission follows neural network with weights [0.23, -0.41, …]” is not. In practice, recovering clean symbolic expressions from noisy epidemic data is difficult, but symbolic UDEs offer a path from discovery towards interpretation.
For more on this approach, see the SciML tutorial on missing physics and Philipps et al. (2025) for a recent review of UDEs in systems biology.
This session and the variational inference session round out the course’s coverage of inference methods. Together with the MCMC, particle filter and pMCMC sessions, you now have a range of approaches for fitting mechanistic models to data, and some sense of when each is appropriate.