Physics-Informed Neural ODEs with Scale-Aware Residuals for Learning Stiff Biophysical Dynamics

TL;DR

This work tackles learning stiff, oscillatory biophysical dynamics with neural differential equations by introducing PI-NODE-SR, which normalizes physics residuals across variables evolving on disparate timescales and trains an explicit, low-order solver within a physics-informed objective. The method stabilizes training and enables long-horizon extrapolation on Hodgkin–Huxley dynamics, recovering both oscillation frequency and near-accurate amplitudes, while also capturing fast morphological features in gating variables. Through extensive ablations, the study demonstrates that scale-aware residuals, appropriate adjoint choices, and controlled physics loss are essential to outperform vanilla Neural ODEs and PINNs, even under noisy observations. The approach suggests a principled path toward robust, data-efficient learning of stiff biological dynamics and motivates extensions to partial observability, uncertainty quantification, and network-level models with solver–adjoint innovations.

Abstract

Neural differential equations offer a powerful framework for modeling continuous-time dynamics, but forecasting stiff biophysical systems remains unreliable. Standard Neural ODEs and physics informed variants often require orders of magnitude more iterations, and even then may converge to suboptimal solutions that fail to preserve oscillatory frequency or amplitude. We introduce PhysicsInformed Neural ODEs with with Scale-Aware Residuals (PI-NODE-SR), a framework that combines a low-order explicit solver (Heun method) residual normalisation to balance contributions between state variables evolving on disparate timescales. This combination stabilises training under realistic iteration budgets and avoids reliance on computationally expensive implicit solvers. On the Hodgkin-Huxley equations, PI-NODE-SR learns from a single oscillation simulated with a stiff solver (Rodas5P) and extrapolates beyond 100 ms, capturing both oscillation frequency and near-correct amplitudes. Remarkably, end-to-end learning of the vector field enables PI-NODE-SR to recover morphological features such as sharp subthreshold curvature in gating variables that are typically reserved for higher-order solvers, suggesting that neural correction can offset numerical diffusion. While performance remains sensitive to initialisation, PI-NODE-SR consistently reduces long-horizon errors relative to baseline Neural-ODEs and PINNs, offering a principled route to stable and efficient learning of stiff biological dynamics.

Figures (7)

• Figure 1: Extrapolation behaviour of vanilla NeuralODE models trained on Hodgkin–Huxley dynamics using different solvers. Top: spike raster, voltage $V(t)$, and phase portrait. Bottom: gating variable predictions ($n$, $m$, $h$).
• Figure 2: Vanilla PINN performance on Hodgkin–Huxley dynamics.Top Left: Spike raster shows zero predicted spikes beyond training. Top Middle: Voltage traces collapse to a flat resting potential. Top Right: Phase portrait shows no learned oscillatory structure. Bottom Row: Gating variables remain flat and outside biological bounds.
• Figure 3: PI-NODE-SR extrapolation results on Hodgkin-Huxley dynamics. Top left: spike raster comparison; Top center: predicted vs. ground truth membrane potential $V(t)$; Top right: phase portrait ($V$–$n$); Bottom: gating variable trajectories ($n$, $m$, $h$).
• Figure 4: Impact of physics loss coefficient $\lambda$ on long-term extrapolation. Increasing $\lambda$ improves structural fidelity, but overly high values smooth out action potentials and slightly degrade timing/precision. $\lambda=1.5$ achieves the best trade-off.
• Figure 5: PI-NODE-SR with backsolve adjoint extrapolation results on Hodgkin-Huxley dynamics. Top left: spike raster comparison; Top center: predicted vs. ground truth membrane potential $V(t)$; Top right: phase portrait ($V$–$n$); Bottom: gating variable trajectories ($n$, $m$, $h$).