Source-linked AI summary
Stiff Neural Ordinary Differential Equations
Suyong Kim, Weiqi Ji, Sili Deng, Yingbo Ma, Christopher Rackauckas
TL;DR
Neural ODEs can struggle with stiff systems because widely separated time scales make solver integration and gradient-based training unstable. This paper analyzes those failures, proposes scaling and stabilized derivative calculations, and demonstrates successful learning on ROBER and POLLU. The results support using these techniques for stiff chemical dynamics across applications with widely varying time scales.
Problem
Learning neural ODEs for stiff systems is limited by computationally costly stiff solves, ill-conditioned gradients, and scale separations in time and state magnitudes.
Method
The paper combines equation and loss scaling, deeper rectified-activation networks, and stabilized adjoint or derivative calculations for stiff neural ODE training.
Results
The proposed approaches accurately reproduced all species profiles in the ROBER and POLLU benchmarks, including slow and fast time scales, whereas baseline models failed to learn the overall dynamics.
Takeaways & Limitations
The demonstrated techniques support applying neural ODEs to chemical, environmental, biological, and other systems with widely varying time scales.
Abstract
from arXiv · showhide
Neural Ordinary Differential Equations (ODE) are a promising approach to learn dynamic models from time-series data in science and engineering applications. This work aims at learning Neural ODE for stiff systems, which are usually raised from chemical kinetic modeling in chemical and biological systems. We first show the challenges of learning neural ODE in the classical stiff ODE systems of Robertson's problem and propose techniques to mitigate the challenges associated with scale separations in stiff systems. We then present successful demonstrations in stiff systems of Robertson's problem and an air pollution problem. The demonstrations show that the usage of deep networks with rectified activations, proper scaling of the network outputs as well as loss functions, and stabilized gradient calculations are the key techniques enabling the learning of stiff neural ODE. The success of learning stiff neural ODE opens up possibilities of using neural ODEs in applications with widely varying time-scales, like chemical dynamics in energy conversion, environmental engineering, and the life sciences.
I. INTRODUCTION
Neural ODEs learn dynamical systems from time-series data without prespecified functional forms, but stiffness creates computational and optimization challenges. This work investigates those challenges and develops strategies for learning stiff systems.
- Neural Ordinary Differential Equations: Neural ODEs use neural networks to model dynamical systems from time-series data without requiring expert-specified functional forms.The ODE is integrated to produce predictions, and network parameters are optimized from differences between observed and predicted states.
- Challenges of Stiff Dynamics: Learning neural ODEs for stiff dynamics is challenging because stiffness combines widely separated time scales with high solver costs and ill-conditioned gradients.Stiffness can also produce gradient pathologies and ill-conditioned optimization in other data-driven modeling approaches.
- Study Scope: The study examines neural ODE learning on ROBER and POLLU, classical stiff chemical-reaction systems used to benchmark stiff ODE integrators.ROBER and POLLU describe species concentrations in stiff chemical reaction systems.
- Proposed Direction: The proposed strategies address scale separation in time and state magnitudes while introducing stabilized adjoint techniques for stiff neural ODEs.The paper reports successful learning on both benchmark problems and discusses implications for other data-driven stiff-system methods.
- Training Formulation: Neural ODE training integrates the learned differential equations and compares predicted state variables with observations through a loss function.Differentiable ODE solvers provide gradients with respect to neural-network parameters for stochastic optimization.
B. Stiff ODE Systems
Stiff systems are numerically difficult for explicit methods and are characterized here through established stiff benchmarks rather than a single universal stiffness definition. The selected examples represent chemical systems with strongly separated scales.
- Stiffness Definition: Stiff equations are classically described as problems for which explicit numerical methods do not work adequately.The paper notes that multiple definitions exist and that no single stiffness index captures every numerically difficult problem.
- Benchmark Systems: ROBER and POLLU are chosen as representative highly stiff systems studied extensively in the stiff ODE solver literature.They are used because stiffness definitions are ambiguous and these equations are established benchmark problems.
- Benchmark Systems: ROBER models three species and five reactions, whereas POLLU models 20 species and 25 reactions in atmospheric air-pollution chemistry.Both systems describe species-concentration dynamics in stiff chemical reaction systems.
III. METHOD
The method analyzes why reversing stiff ODEs destabilizes adjoint gradients and shows that small forward or reverse-solve errors can be strongly amplified. It motivates stabilized derivative calculations instead of direct ODE reversal.
- Adjoint Stability: Reversing an ODE can be numerically unstable because errors in the forward solution propagate into the gradient during the reverse solve.The instability is associated with large Lipschitz constants and stiff Jacobians.
- Adjoint Stability: λ = 100 causes about 1% error in reverse integration, while λ = 10000 cannot be reversed numerically in double precision.These results come from the linear ODE dy/dt = −λy on t ∈ (0,1).
- ROBER Reversibility: For ROBER, tolerances of 10^-6, 10^-7, and 10^-8 produce reverse-solve errors of 72%, 38%, and 5%, respectively.An error of 0.1% is reached only for tol < 10^-10, below single-precision arithmetic.
- Error Mechanism: The linear ODE u′(t) = λu(t) demonstrates that forward-solve error is multiplied by e^-λ during the backward solution.The backward solution also contains its own numerical error, ε_b(0).
- Stabilized Derivatives: Interpolated checkpointing and discrete adjoint sensitivities avoid reversing the ODE and agree on derivatives to at least three decimal places at tol = 10^-6.The paper notes that discrete adjoint consistency can ensure stability.
B. Linear Scaling Adjoints
Stiff ODEs make adjoint gradient calculations numerically unstable and computationally expensive. The proposed QuadratureAdjoint, IMEX, and parameter-splitting strategies reduce these costs by separating state solves from parameter derivative calculations.
- Reverse-solving a stiff ODE can amplify forward numerical errors exponentially when the system’s Lipschitz constant is large.The linear test problem demonstrates error blow-up in the backward solve, motivating alternatives to reversing the trajectory directly.
- Stiff equations produce ill-conditioned linear systems because their Jacobians have widely separated eigenvalues.Krylov methods can converge slowly, while direct Jacobian solves generally require LU factorization with O(k^3) cost.
- Forward sensitivities scale as O(k^3m^3), whereas conventional adjoints scale as O((k + m)^3) for implicit stiff solves.The cubic dependence on the combined state and parameter dimensions becomes costly for neural ODEs with many parameters.
- QuadratureAdjoint sequentially solves the state and adjoint equations using dense output, then computes parameter derivatives through explicit quadrature.This trades computation for memory and reduces the cost to O(k^3 + m).
- IMEX splitting treats part of the adjoint equation implicitly and part explicitly, reducing peak memory while retaining O(k^3 + m) complexity.Specializing factorization to the implicit block avoids including all parameter derivatives in the stiff solve.
- Computing derivatives for only l ≪ m parameters at a time lowers complexity, and separate differentiation passes can run in parallel.The resulting costs are O(mk^3l^2) for forward sensitivities and O((m/l)(k + l)^3) for adjoints.
C. Equation Scaling
Equation scaling addresses the remaining training instability caused by separated time and state magnitudes. The approach normalizes neural-network outputs and loss components using characteristic scales.
- The proposed normalized neural ODE uses characteristic state and time scales to address stiffness-induced scale separation.yscale contains characteristic scales for each species, while tscale denotes the characteristic time scale.
- Loss components for different species are re-balanced by normalizing their loss functions.This complements output scaling to reduce imbalance among species with different magnitudes.
A. ROBER Problem
Robertson’s equations model a stiff reaction network involving three chemical species. The species concentrations are governed by rate constants spanning very different magnitudes.
- ROBER models a reaction network with three chemicals whose concentrations are represented by [y1,y2,y3].Its reaction rate constants are k1 = 0.04, k2 = 3·10^7, and k3 = 10^4.
1. Baseline Model
The baseline neural ODE uses a small tanh network with stiff-solver integration and standard optimization. On ROBER, it fails to reproduce the dynamics because training remains unstable under severe scale separation.
- The baseline model uses one hidden layer with 50 nodes and hyperbolic tangent activation.Training uses Rosenbrock23, hybrid Rosenbrock23–Tsit5 auto-switching, ForwardDiff.jl, and ADAM with learning rate 0.005.
- 10,000 epochs were insufficient to reduce the flatlined training loss.The training loss stopped decreasing despite continued optimization.
- The baseline fails to learn the actual dynamics or reproduce species profiles, while gradient norms fluctuate strongly.The imbalance is especially severe because y2 is approximately 10^-5 while y1 and y3 change at order unity, producing unequal loss components.
2. Mitigation of Gradient Pathologies via Scaling
Equation scaling and network design stabilize training for stiff ROBER neural ODEs, whereas the baseline fails to reproduce the dynamics. Scaling is essential, while deeper rectified networks improve successful learning and training stability.
- Scaled neural ODE: Equation scaling enabled successful learning of the Robertson problem with a six-layer, five-node-per-layer GELU network.The scaled model’s results are compared with the ground truth in Figure 6.
- Training stability: The scaled model’s loss functions steadily decreased while its gradient norm remained stable during training.Training took 2 hours on a workstation using a single CPU processor.
- Sensitivity analysis: Scaling was essential: the same network structure without equation scaling failed to learn the ROBER problem.Shallow tanh networks could also learn ROBER but required longer training, potentially because of gradient issues.
- Acceleration techniques: Temporal regularization and learning-rate annealing could accelerate training but were not required for successful learning.Temporal regularization introduces stochasticity by randomly truncating the integration time during optimization.
B. POLLU Problem
The POLLU air-pollution benchmark tests whether the proposed stiff neural ODE techniques generalize beyond ROBER. A learned model captures both the initial induction dynamics and near-equilibrium states, while the baseline again fails.
- Model and data: The POLLU neural ODE used three hidden layers with 10 nodes per layer and data sampled uniformly over [0,60].Other hyper-parameter settings matched those used for ROBER.
- Results: The trained neural ODE captured both the induction-period dynamics and the near-equilibrium states.Figure 8 compares predictions with label data generated from the POLLU equations.
- Results: The baseline model again failed for POLLU, similarly to its failure on ROBER.The baseline result is presented in the supplemental material.
V. CONCLUSION AND DISCUSSION
The paper combines stabilized derivative calculations, equation scaling, and deeper rectified networks to learn stiff neural ODEs, while identifying broader implications and remaining theoretical gaps.
- Conclusion: The proposed derivative calculations, equation scaling, and deeper rectified networks successfully learn stiff neural ODEs on the ROBER and POLLU systems.The approaches accurately reproduce species profiles across slow and fast time scales, whereas baseline models failed to learn the overall trend.
- Limitations and future work: Further theoretical analysis is needed to explain how stiffness affects training gradients, failures, equation scaling, and optimization itself.The paper also identifies stiff-aware optimization and analysis of constrained neural-network DAEs as open directions.
- Discussion: Errors in continuous adjoint handling can affect neural ODE training and may require analogous stabilization in other implicit-layer methods.The discussion connects this issue to nonlinear solvers, deep equilibrium methods, and differentiable optimization layers under ill-conditioning.
- Discussion: The techniques may support machine-learning architectures that satisfy arbitrary constraint equations during their evolution through differential-algebraic equation formulations.The paper describes mass-matrix DAEs as one route for imposing such constraints.