Scientific machine learning: a neural network inside a differential equation

E09 fitted the Lotka-Volterra equations to the lynx and hare records. That model assumes the form of the interaction: hares are eaten at a rate proportional to hares times lynx ("mass action"). What if you do not trust that term? The idea behind universal differential equations (Rackauckas et al. 2020) is to keep the terms you trust and let a small neural network stand in for the one you do not. Bayesian neural ODEs (Dandekar et al. 2020) add the obvious next step: be Bayesian about the network weights as well as the physical parameters, so that the discovered mechanism comes with uncertainty instead of being a point fit. The plain neural ODE, $\dot z = \text{NN}(z)$ with no physics at all, is Chen et al. (2018).

This notebook builds all three points of that spectrum on 21 years of data, and it is frank about what happened. The short version: the textbook hybrid (a network learning the interaction from scratch) is easy to optimise and close to impossible to sample in the time budget of this notebook; a reformulation that puts the physics into the prior samples cleanly in under a minute; and when the models are asked to forecast six years they have not seen, every bit of flexibility costs accuracy. It assumes you have read E09: the wrap_jax route, the RK4 lax.scan solver, the log-scale state and the advice on starting values are used here without being re-taught.

#1 · The data and a spectrum of three models

The same 21 years as E09: two cycles of roughly ten years, lynx lagging hares.

With $H$ hares and $L$ lynx, the three models differ only in what they claim to know about the right-hand side.

$g$, $c$ and $f$ are small neural networks. Why is it reasonable to call the linear terms "known"? They are statements about one species alone - hares without predators multiply, lynx without prey starve - that can be measured without ever watching the two interact. The interaction is the part an ecologist argues about: mass action, or a predator that gets full (Holling's type II response), or prey that hide when predators are many.

As in E09 everything is integrated on the log scale, $x = \log H$, $y = \log L$, where the three right-hand sides read $\dot x = \alpha - \beta e^{y}$ (mechanistic), $\dot x = \alpha - g_1$ (from scratch), $\dot x = \alpha - \beta e^{y + c_1}$ (physics in the prior) and $\dot x = f_1$ (black box), and the solver output is the location of a LogNormal likelihood.

One deliberate difference from E09: a single noise scale $\sigma$ for both species. E09 gave each species its own and both came out near 0.25. Give a flexible model two noise scales and it spends its freedom fitting the smoother lynx series almost exactly ($\sigma_{\text{lynx}} \approx 0.08$ at the optimum when I tried it) while explaining the hares as noise - a funnel that also wrecked the sampling. The measurement process is the same for both species (trappers, one company's ledgers), so one $\sigma$ is the better model, and all three models share it to keep the comparison fair.

#2 · The network and the solver, in JAX

The network is a pure function of one flat weight vector. Equinox would work (wrap_jax accepts pytrees, as E09 showed), but it would buy nothing at this size and cost clarity. A flat vector means PyMC owns the weights as one random variable with a weight dimension, there is exactly one pt.specify_shape at the boundary, the prior scale is one number, and the symmetries of section 6 can be written as index permutations. The layers are unpacked by slicing, which costs nothing under jit.

Both networks have two inputs (the standardised log populations), one hidden layer of N_HIDDEN = 3 tanh units and two outputs. Three units is small even by the standards of this literature; section 4 explains why more would not fit in this notebook.

• mlp: the ordinary network, $W_2 \tanh(W_1 s + b_1) + b_2$, 17 weights. Used by the black box and the from-scratch hybrid.
• correction: $W_2\,[\tanh(W_1 s + b_1) - \tanh(b_1)]$, 15 weights. No output bias, and zero at $s = 0$ by construction, so that $\beta$ and $\delta$ keep a meaning: they are the interaction rates at the centre of the data. It is the network analogue of centring a regressor.

The solver is E09's fixed-step RK4 in a lax.scan, generalised to take the right-hand side as an argument. One change: 2 steps per year instead of 20. Every gradient evaluation costs a forward and a backward pass through the whole scan, NUTS will want millions of them, and a ten-year cycle does not need a step of 18 days. The check:

Two steps per year are wrong by a few thousandths on the log scale - about seventy times smaller than the observation noise ($\sigma \approx 0.25$). One step per year would not be good enough. Section 9 repeats the check on the fitted networks, where it matters more.

#The PyMC model

One builder for all three models. n_train is there for the forecast in section 9: the model then integrates and sees only the first n_train years. Priors on the rates, the initial state and the noise are E09's. The weights get Normal(0, w_scale); the choice of w_scale is the subject of section 5.

#Check the gradient, and price it

wrap_jax gets its gradient from jax.vjp, so there is little to get wrong - but "little" is not "nothing" (a stop_gradient left in, a non-differentiable branch), and the check is cheap. verify_grad compares the pullback with finite differences along random projections, here through 40 RK4 steps and the network. The weak part of such a check is the finite differences, not the gradient: with the rates on their natural scale ($\beta \approx 0.03$, and a trajectory that is very sensitive to it) my first version failed with a relative error of $2 \times 10^{-4}$ against a tolerance of $10^{-4}$, at a point where forward- and reverse-mode differentiation agree to sixteen digits. So the check is done where the sampler works, on the log-rates, and backed by that second, finite-difference-free comparison.

What NUTS pays per step is one evaluation of the log-density and its gradient. The JAX backend (the whole model as one XLA program, E09's Route 3) is what the fits below use.

A fraction of a millisecond each (timings vary from run to run; read ratios). The network is not free: forty RK4 steps with four network evaluations each, forwards and backwards, make a hybrid gradient several times dearer than a mechanistic one. But at these prices the per-gradient cost is not what decides whether a fit is affordable. What decides it is how many gradients the sampler needs per draw. For the mechanistic model that is about five. Hold that thought.

#The baseline: E09's model, again

Same recipe as E09 - start from the physically motivated point, no jitter, low-rank mass matrix - wrapped in a helper that every fit in this notebook goes through. The helper records the diagnostics that make sense for every model here: those of the predicted trajectory log_mu, which is a function-space quantity. maxdepth=8 caps a NUTS trajectory at 255 leapfrog steps (the default cap is 1023); the reason will become clear.

A few seconds, no divergences, a handful of leapfrog steps per draw, and the same rates as E09 ($\alpha \approx 0.55$, $\beta \approx 0.028$, $\gamma \approx 0.79$, $\delta \approx 0.024$) with one shared $\sigma \approx 0.25$. Neither the coarser solver nor the shared noise scale changed the answer. This is the yardstick for everything below.

#3 · First attempt: learn the interaction from scratch

The textbook universal differential equation: $\dot x = \alpha - g_1(x, y)$, $\dot y = -\gamma + g_2(x, y)$ with $g = \text{softplus}(\text{mlp})$ and Normal(0, 1) weights. Before sampling anything, optimise it - this is what most of the literature does, and we need a starting point anyway, because a randomly initialised neural ODE is nowhere near an oscillator.

The log-posterior is written directly in JAX (it is five lines) so that the whole optimisation can be vectorised: jax.vmap over 32 random initialisations of an Adam loop inside one lax.scan. Thirty-two complete trainings of a neural ODE take a few seconds.

Two lessons already. The objective is multimodal in the way E09 warned about, only more so: half of the runs end in a poor local optimum - trajectories that damp out to a level line through the data, or that catch the first cycle and lose the second - 25 nats or more worse than the best. A single training run would have had an even chance of reporting failure. And the good runs agree: about a dozen end within a few nats of each other, their trajectories are nearly indistinguishable, and they reproduce both cycles.

Did they learn the same mechanism, and is it mass action? The identifiable quantity is the net per-capita growth rate, $\alpha - g_1$ for hares and $-\gamma + g_2$ for lynx ($\alpha$ and the level of $g_1$ can be traded against each other). Mass action makes two sharp predictions: hare growth falls linearly in $L$ and does not depend on $H$; lynx growth rises linearly in $H$ and does not depend on $L$. Below, one abundance varies over its observed range while the other species is held at its median; each good run is one line, against the 90% posterior band of the mechanistic model.

Read the top-right and bottom-left panels first: hare growth falls as lynx increase and lynx growth rises with hares, with the right sign, nearly the right curve (a straight line in the abundance is an exponential on this log axis), and the zero crossing - the equilibrium - where the mechanistic model puts it. Nothing told the network any of that beyond the signs. The other two panels are the test mass action could fail: the growth of a species should be flat in its own abundance. The optimised networks stay close to the flat band. For hares they fan out at the low-abundance end, some up and some down; for lynx most of them tilt gently downwards at high lynx numbers (a hint of crowding among predators, or nothing). So a network that knows nothing about predation recovers the shape of the mass-action law from 21 years of pelts, and the optima disagree with each other mostly where that law makes its most specific claim.

"The optima disagree" is a poor man's uncertainty statement. A spread of point estimates from different initialisations (a "deep ensemble") is not a posterior: it says nothing about the width of each mode, and the runs were selected by me. What we want is the posterior.

#4 · Inference, honestly: why I could not sample that model here

The natural next cell would wrap rhs_scratch in the model builder, start NUTS at the best optimum and report a posterior. I did that, many times. None of those fits is in this notebook, because none was both trustworthy and affordable. What I measured while preparing it (nutpie, low-rank mass matrix, 4 chains started at the best optimum, on the busy 8-core laptop that built this page):

("Two $\sigma$" is E09's noise model, one scale per species; section 1 explains why I dropped it.) Shrinking the network to two hidden units did not rescue it, nor did fixing $\alpha$ and $\gamma$, an exponential instead of a softplus output, centring the network, NumPyro with vectorised chains, or a tighter weight prior (which made the optimisation fail more often). The step size always adapted to about 0.03 - 0.06 and most draws ran into the depth cap. The reason is visible in the curvature at the optimum:

(Adam had not quite converged, which is why the optimum is polished with L-BFGS first; a Hessian away from a stationary point means little.) The stiffest direction is pinned down to a posterior standard deviation of about 0.0002 and the softest are as wide as the prior, of order one: a ratio of several thousand in a local Gaussian approximation. The stiff direction has a simple origin. On the log scale a constant error $\epsilon$ in a growth rate grows into an error $\epsilon t$ in the trajectory, so twenty years of data constrain the average growth rate along the orbit to a few parts in a thousand, and the period almost as tightly. In the mechanistic model that constraint is a fixed linear combination of four log-rates, and the low-rank mass matrix absorbs it - five leapfrog steps per draw. In the network the same constraint is a curved sheet in a 17-dimensional weight space, because which combination of weights sets the average growth rate depends on where in weight space you are. No linear preconditioner straightens a curved sheet. NUTS has to walk along prior-wide directions with steps sized for the thin one: hundreds of gradient evaluations per draw, and even then the chains diffuse too slowly to agree.

Longer trajectories do fix it. The one network fit with a prior this loose that I let run to the default depth cap (the model of the next section at w_scale = 1) mixed well - trajectory $\hat R$ 1.01 - at about 800 leapfrog steps per draw: ten minutes for one fit. That is the honest price of a loosely constrained Bayesian neural ODE with NUTS, on a problem with 42 data points. It is why the published examples use long runs, stochastic-gradient samplers or variational approximations - and why the black-box fit below, which has the same geometry and no physics to lean on, will come with a warning label.

#5 · The fix: put the physics in the prior

The from-scratch network spends most of its capacity re-deriving what we already believed. The alternative is to keep mass action and let the network multiply it:

$$ \dot x = \alpha - \beta\,e^{y}\,e^{c_1(x, y)}, \qquad \dot y = -\gamma + \delta\,e^{x}\,e^{c_2(x, y)} $$

With all weights zero, $c \equiv 0$ and this is Lotka-Volterra. A Holling type II response would show up as $c_1, c_2$ falling with $H$; prey refuge as $c$ falling with $L$. The network is now a residual on trusted physics (the usual role of the network in a universal differential equation), and the prior scale of the weights acquires a meaning a scientist can argue about: how far from mass action am I prepared to go before seeing data?

#Priors over weights are priors over functions

Nobody has prior beliefs about W1[2,lynx]. The way to choose w_scale is to look at what it implies in function space: the correction factor $e^{c}$ over the region the data visit, and the trajectories. Priors this simple can be drawn with NumPy and pushed through the same JAX solver with vmap; w_scale = 0 is the mechanistic prior of E09. The figure shows, for three scales, 60 prior hare trajectories (top) and the prior distribution of the correction factor $e^{c_1}$ at one corner of the region the data visit - many hares and many lynx (bottom).

Blow-ups first. A neural ODE with random weights is the classic source of NaNs, and the handling here has three layers. Design: the state is a log-population, so no solver error can produce a negative population, and the hidden units are tanh, so the correction factor is bounded for any input however far a trajectory strays; the exact solution of such a system cannot reach infinity in finite time. (E09's section 10 shows what the natural scale does instead: negative populations, then NaN in the likelihood.) Count, never silently drop: the table reports non-finite draws, and statistics are computed on the finite ones. At w_scale up to 0.3 there are none. At 1.0, ten draws in a thousand fail - and every one of them is finite when solved again with 100 steps per year. Those failures are the solver, not the dynamics: a wild network makes the dynamics fast somewhere, and RK4 with two steps a year is only stable while nothing happens on a time scale much shorter than a year. In the sampler: nutpie treats a non-finite log-density as a rejected, divergent proposal and carries on, so an unstable corner of parameter space costs a divergence, not a crash - and section 9 checks the step size on the posterior draws that matter.

Then the functions. The milder kind of blowing up - trajectories that leave any plausible range, above a million pelts or below a hundred - happens for 29% of draws from the mechanistic prior (E09 showed those and accepted them), for about the same share at w_scale = 0.1, 45% at 0.3 and 62% at 1.0. The mechanism tells the same story. At 0.1 the network rarely moves the interaction by more than 6% anywhere in the region the data visit: a Lotka-Volterra model with extra steps. At 1.0 the typical draw changes predation by a factor of 1.8 and one draw in twenty by a factor of 13 or more, within the range of populations actually observed - the prior has all but forgotten mass action, and we are back in the geometry of section 4. w_scale = 0.3 says: mass action is right to within about 10% in a typical place, and to within a factor of 1.5 nineteen times in twenty. That is a statement I can defend - loose enough for a type II response to show itself, tight enough to be a correction - and it is what models["hybrid"] was built with. It is also a knob you can turn; "Try it yourself" asks you to.

#6 · Fitting the hybrid, and how to check a network posterior

The residual network has a natural starting point that needs no optimisation at all: E09's physical start for the rates and weights near zero. (Exactly zero would be a poor choice - with $W_2 = 0$ the gradient with respect to the first layer vanishes - so the weights start at a small random value, the same for every chain.)

Under a minute rather than ten. A hundred leapfrog steps per draw is still twenty times the mechanistic model's appetite - the curved sheet of section 4 has not gone away, the prior has only made it small enough to cross - and a dozen divergences (under 1% of the draws) are a reminder of that. No draw reached the depth cap. On the trajectory, $\hat R$ and ESS are fine.

#$\hat R$ on weights is the wrong question

A one-hidden-layer tanh network computes exactly the same function if you permute its hidden units, or flip the sign of all weights into and out of one unit (tanh is odd). With three units that is $3! \times 2^3 = 48$ weight vectors per function, and the Gaussian prior is blind to the difference, so the posterior has 48 identical copies of everything. Four chains that each explore a different copy have sampled the same posterior over functions, and their weight-space $\hat R$ says they disagree completely. We can manufacture that situation without running anything: map three of our four chains to other copies, which are equally valid posterior draws, and recompute.

The relabelled chains define the same functions to machine precision - and here the weight $\hat R$ does not move either. That is not what I expected when I wrote the cell, and it is instructive. At w_scale = 0.3 the 48 copies are not separate modes: they overlap around zero and merge into one symmetric blob, and every chain already wanders through all of them, so relabelling a chain changes nothing statistically. The black box of section 8, whose weights must be large to produce an oscillation at all, shows the other case. In both cases the weight marginals are useless, as a look at them makes plain:

Every weight has posterior mean zero to two decimals and a standard deviation of 87% to 107% of its prior's 0.3. Read weight by weight, the data taught this network next to nothing - which is false, as section 7 shows. What the data constrain are products and combinations of weights, in a sign- and permutation-symmetric way that leaves almost no trace in any marginal. Weights are not parameters in the sense of E01-E10: they have no meaning, no identifiability, and no business in a summary table. A good $\hat R$ on them proves as little as a bad one.

#Convergence in function space

What must converge is what you will report: the predicted trajectory at each year, the learned right-hand side $(\dot x, \dot y)$ on a grid of states covering the data, and the log-likelihood (a scalar summary of fit that is invariant under every symmetry). The black box goes through the same checks in section 8.

Every function-space quantity has $\hat R$ of about 1.01 or better and at least 700 effective draws out of 2000: good enough for the means and 90% bands reported below. The usual caveat from E09 applies with more force: all chains started at the same point, so this shows that they agree with each other, not that there is no second, functionally different mode they all missed. Section 9 meets one.

#Cheap approximations, judged in function space

If NUTS is this expensive, do the usual shortcuts work? Mean-field ADVI and Pathfinder, both from the same starting point. Two practical notes, both learned the hard way. ADVI's default initial standard deviation of 1 on the unconstrained scale throws its first samples into regions where the ODE gradient is astronomically large, and the optimisation returns NaN within a few iterations; start_sigma of 0.05 and gradient clipping fix that (full-rank ADVI has no start_sigma and produced NaN regardless, so it is not shown). And pmx.fit_pathfinder forks worker processes by default, which deadlocks against the initialised JAX runtime exactly as PyMC's own NUTS did in E09: parallel=False.

Both shortcuts run in seconds and put the trajectory roughly where NUTS does (the mean is off by up to one to one and a half posterior standard deviations in some years). Both get the uncertainty wrong, and in two opposite directions at once.

• The trajectory band is too wide: 1.7 times for mean-field ADVI, 1.2 times for Pathfinder. The folklore says mean-field ADVI underestimates uncertainty; that is a statement about parameter marginals. Here the data pin down a few combinations of parameters extremely tightly (section 4), a factorised Gaussian cannot represent a combination, and so each parameter wobbles independently and the period and the average growth rate wobble with them - visibly, in the second cycle. Pathfinder's low-rank covariance captures part of it.
• The mechanism band is far too narrow: the standard deviation of the learned correction is a quarter of the NUTS value for both. The correction is a product of first- and second-layer weights whose posterior is the symmetric blob of the last section; a single Gaussian centred near zero with shrunken widths (ADVI's weight standard deviations come out at about half of NUTS's, Pathfinder's at about a quarter) turns that into a network that is confidently almost switched off.

So an approximation can look conservative on the quantity you plotted and be badly overconfident on the one you did not. Either shortcut is a reasonable way to find a starting point or to screen model variants. Neither is a substitute for the posterior when the width of a band is the result, which in this notebook it is.

#7 · The scientific payoff: what did the network learn about predation?

The honest framing first. Because the prior is centred on mass action, this model cannot "rediscover" mass action - section 3 is the closest we get to that, with optimisation only. The question a Bayesian hybrid answers is the one a scientist would actually pose: do the data call for a departure from mass action, where in state space, and how sure are we? Three views.

Where did the data teach the network anything? The ratio of posterior to prior standard deviation of the correction $c(H, L)$, over the region of state space the data visit. A ratio of 1 means "still the prior".

Does the interaction look like mass action where it is constrained? The per-capita predation rate $\beta L\,e^{c_1}$ as lynx vary, and the per-capita lynx gain $\delta H\,e^{c_2}$ as hares vary (the other species at its median), against the straight lines $\beta L$ and $\delta H$ of the mechanistic fit. The ticks along the bottom mark the observed abundances.

A symbolic summary with uncertainty. If the interaction were a power law, $U_1 \propto H^{a} L^{b}$, then $a = 1 + \partial c_1 / \partial x$ and $b = 1 + \partial c_1 / \partial y$; mass action is $a = b = 1$, a saturating type II response has $a < 1$. jax.jacobian gives those elasticities for free. We average them over the observed states, once per posterior draw, and do the same for draws from the prior.