Source-linked AI summary
Universal Differential Equations for Scientific Machine Learning
Christopher Rackauckas, Yingbo Ma, Julius Martensen, Collin Warner, Kirill Zubov, Rohit Supekar, Dominic Skinner, Ali Ramadhan, Alan Edelman
TL;DR
Scientific machine learning needs methods that combine mechanistic knowledge with flexible learning when experiments are expensive and data are limited. This paper introduces UDEs and the SciML ecosystem as a unified computational framework, demonstrating efficient training and broad applicability across equation discovery, stochastic and stiff systems, and high-dimensional PDEs. The implementation supports scientific workflows at practical scale, with examples using hundreds of lines of code and training in under half an hour on a standard laptop.
Problem
Expensive scientific experiments limit deep-learning effectiveness, while mechanistic models impose simplifying assumptions; the paper addresses the need to merge both approaches.
Method
The paper formulates universal differential equations as differential equations defined fully or partly by universal approximators and trains them through SciML solvers, adjoints, and symbolic tools.
Results
The framework handles diverse UDE applications, including biological equation discovery, non-ODE sparse regression, stochastic high-dimensional PDEs, and models with stiffness, delays, and implicit constraints.
Takeaways & Limitations
SciML provides a modular and composable ecosystem for applying shared, optimized training mechanisms across varied scientific machine-learning problems.
Takeaways & Limitations
BacksolveAdjoint can diverge for universal PDEs and is not universally applicable, especially for stiff equations.
Abstract
from arXiv · showhide
In the context of science, the well-known adage "a picture is worth a thousand words" might well be "a model is worth a thousand datasets." In this manuscript we introduce the SciML software ecosystem as a tool for mixing the information of physical laws and scientific models with data-driven machine learning approaches. We describe a mathematical object, which we denote universal differential equations (UDEs), as the unifying framework connecting the ecosystem. We show how a wide variety of applications, from automatically discovering biological mechanisms to solving high-dimensional Hamilton-Jacobi-Bellman equations, can be phrased and efficiently handled through the UDE formalism and its tooling. We demonstrate the generality of the software tooling to handle stochasticity, delays, and implicit constraints. This funnels the wide variety of SciML applications into a core set of training mechanisms which are highly optimized, stabilized for stiff equations, and compatible with distributed parallelism and GPU accelerators.
1 Introduction
The paper introduces UDEs as a framework for combining mechanistic scientific knowledge with machine learning, supported by SciML software for efficient training across diverse differential-equation problems. The ecosystem spans symbolic modeling, differentiated solvers, adjoints, sparse equation discovery, stochastic PDEs, and applications involving delays, stiffness, and non-local operators.
- Motivation: SciML combines mechanistic models and data-driven learning to address small-data settings while retaining structural scientific information.The framework is intended to merge the predictive structure of mechanistic models with the flexibility of machine learning.
- Universal differential equations: UDEs define differential equations fully or partly through universal approximators, including neural networks and other parameterized function representations.This formalism encompasses neural ODEs, physics-informed methods, neural-network optimal control, and model augmentation.
- Training infrastructure: Differential-equation solvers, adjoint methods, and helper libraries provide an efficient and stable training framework for UDEs, including stiff and non-reversible problems.The ecosystem includes a broad set of forward and adjoint sensitivity methods with different applicability trade-offs.
- Equation discovery: UDE-enhanced symbolic regression learns unknown interactions while preserving known model structure, enabling equation recovery from limited data without derivative-smoothing requirements.The paper reports accurate recovery of Lotka–Volterra equations where the corresponding SINDy approach with derivative smoothing fails on the same limited data.
- Equation discovery: UDEs extend sparse regression beyond ODEs by representing PDE operators with learnable components such as convolutional stencils and local nonlinear operators.For Fisher-KPP, the learned CNN filter recovers the [1, −2, 1] one-dimensional Laplacian stencil and the missing nonlinear growth term has a parabolic form.
- Broader applications: The formalism also supports adaptive stochastic formulations for high-dimensional PDEs and learns non-local parameterizations from physical data.The paper describes adaptive Euler–Maruyama training for high-dimensional HJB equations and a universal diffusion-advection PDE for ocean convection parameterization.
3 Discussion
The discussion presents UDEs and SciML tools as a general framework for combining mechanistic structure with machine learning across scientific applications. The reported examples include broad adjoint-method coverage, strong performance improvements, and concise implementations.
- Discussion: UDEs combine mechanistic models with machine learning and support equation discovery, high-dimensional partial differential equations, and discretized physics-informed neural networks.The paper also describes UDEs as compatible with the SciML tooling.
- Discussion: The implementation includes the full spectrum of adjoint sensitivity methods needed for efficient and accurate training across universal differential equations.The reported generalizations include stiff equations, DAEs, and SDEs.
- Discussion: Each example can be implemented in hundreds of lines of code, with none taking more than half an hour to train on a standard laptop.The authors present this as evidence of software efficiency and potential scalability.
4 Code and Data Availability
The paper provides reproduction code through the core SciML libraries used for UDE training, sensitivity rules, differential equation solving, and symbolic regression. Experimental data are simulated in the example code.
- Code and Data Availability: The reproducibility code is organized around DiffEqFlux.jl, DiffEqSensitivity.jl, DifferentialEquations.jl, and DataDrivenDiffEq.jl.These libraries respectively cover UDE training, forward and adjoint rules, differential equation solvers, and symbolic regression.
- Code and Data Availability: All experimental data are simulated in the example codes rather than supplied as a separate dataset.The passage states this directly for the experiments.
6 DiffEqFlux.jl Pullback Construction
DiffEqFlux.jl constructs differential-equation solves as differentiable primitives and supports multiple forward and reverse sensitivity strategies. These methods trade memory, stability, and computational cost across ODE, DAE, and neural-network settings.
- Pullback Construction: The pullback treats a differential-equation solve as one differentiable primitive inside Flux.jl's reverse pass.The construction overloads Zygote.jl and related differentiation mechanisms.
- Pullback Construction: For scalar-output functions, the pullback computes gradients through a vector-Jacobian product, enabling recursive reverse differentiation of composed programs.The discussion connects the pullback to Jacobians and primitive backward rules.
- Sensitivity Methods: The implementation offers eight adjoint modes and four internal vector-Jacobian-product methods with distinct performance characteristics and limitations.The full combination yields 48 adjoint-method approaches.
- Sensitivity Methods: BacksolveAdjoint uses O(1) memory but can diverge or produce gradient errors, especially for stiff equations, so checkpoints are added for stability.The checkpointed variant trades additional checkpoints for reduced divergence risk.
- Sensitivity Methods: Adjoint sensitivities avoid multiplicative scaling with state variables and parameters, making them more efficient for large-parameter models.The paper contrasts their O(N + P) cost with forward sensitivity scaling of O(NP).
- Sensitivity Methods: Forward sensitivity is efficient only for small parameter counts, whereas adjoint approaches are suited to neural networks with many parameters.The forward method scales as O(NP), while adjoints scale as O(N + P).
7 Sensitivity Algorithm Decision Tree
The decision tree selects sensitivity methods according to model size, stiffness, memory, differentiability, and hardware characteristics. It favors forward sensitivities for small systems and stable adjoints for larger or stiff problems.
- Method Selection: For equations with 50 parameters and states or fewer, the decision tree recommends considering forward-mode sensitivities.ForwardDiffSensitivity is preferred when the function is compatible with ForwardDiff.
- Method Selection: For larger equations, BacksolveAdjoint and InterpolatingAdjoint are candidates, but BacksolveAdjoint should be skipped for stiff equations or DAEs because it is likely unstable.The recommendation is conditional on checking gradient correctness.
- Method Selection: QuadratureAdjoint is recommended for stiff solvers when memory use is low because it trades memory cost for asymptotic computational efficiency.This recommendation applies when other methods are unsuitable for the equation's stability requirements.
- Method Selection: ReverseDiffAdjoint on CPUs and TrackerAdjoint on GPUs are fallback choices when other methods are unstable.The guidance explicitly recommends checking gradients against each other.
- VJP Selection: After selecting a general sensitivity algorithm, the tree chooses the vector-Jacobian-product implementation based on branching, scalarization, vectorization, and device.The listed options include ReverseDiffVJP, ZygoteVJP, TrackerVJP, and forward-mode or numerical Jacobians.
8 Continuous vs Discrete Adjoints
Continuous and discrete adjoint methods involve a problem-dependent trade-off between stability, computational efficiency, and spurious oscillations. Adjoint stability can also depend on the chosen discretization method.
- Discrete adjoints can be more stable than continuous adjoints in some cases, while continuous adjoints can be more stable in others.
- Continuous adjoints can reduce spurious oscillations in some settings.
- The choice between adjoint approaches represents a trade-off between stability and computational efficiency.
- Adjoint stability may depend on the discretization method used.
9 Integration with Existing Code
DifferentialEquations.jl supports automatic differentiation through differential-equation solves by giving program steps well-defined pullbacks. This allows existing Julia simulation code to be differentiated without user-code modifications.
- DifferentialEquations.jl defines a well-defined pullback for every program step used in its differential-equation solvers.
- The software enables automatic differentiation through differential-equation solves without modifying user code.
- Existing DifferentialEquations.jl simulation software, including large software infrastructures, can therefore be used with automatic differentiation.
10 Benchmarks
The benchmarks evaluate SciML tooling on ODE, neural ODE, and SDE workloads spanning standard, chaotic, celestial-mechanics, and reaction-diffusion systems. Reported results show faster neural ODE and SDE execution, while one adjoint benchmark comparison is qualified by instability and forward-sensitivity effects.
- ODE Benchmarks: The ODE benchmark suite includes the Lorenz weather-prediction model, Pleiades celestial-mechanics simulation, and standard ODE IVP test problems.
- ODE Benchmarks: The remaining benchmarks derive from a discretized two-dimensional reaction-diffusion equation covering systems biology, combustion mechanics, spatial ecology, and spatial epidemiology.
- ODE Benchmarks: The reaction-diffusion ODE was non-stiff in the tested parameter regime, with diffusion constant D = 100 and other parameters set to 1.0.
- Gradient Benchmarks: 12,000x and 1,200x slower: torchdiffeq was slower than SciML on Lorenz and Pleiades derivative calculations when it did not diverge.The comparison is qualified because forward sensitivity analysis was more efficient for these equation sizes and unavailable in torchdiffeq.
- Neural ODE Training: 2.761669e-02 in 2.7 seconds: optimized SciML neural ODE training achieved this final loss and runtime, compared with torchdiffeq’s 0.0596 loss in 289 seconds.SciML defaults achieved a final loss of 4.895287e-02 in 7.4 seconds.
- SDE Benchmarks: Approximately 1,600x: SciML solvers outperformed optimized torchsde on 100 geometric-Brownian-motion SDE solves.SciML took 0.00115 seconds, while the optimized torchsde version took roughly 5 seconds.
11 Sparse Identification of Missing Model Terms via Universal Differential Equations
The UDE workflow trains a universal approximator, samples its induced dynamics, evaluates candidate-function libraries, and applies sparse regression to identify missing model terms. Lotka–Volterra, Hudson Bay, Fisher–KPP, and function-approximator experiments show successful recoveries alongside sensitivity to noise, initialization, and model choices.
- Identification Workflow: SINDy-based UDE identification uses time-series data to construct candidate libraries and solve a sparse regression problem for governing terms.The workflow modifies SINDy to operate on subsets of an equation and on trained neural networks when derivative data are not directly assumed.
- Identification Workflow: The workflow selects state variables, computes the universal approximator’s action, evaluates a basis Θ(X), and identifies sparse coefficients.
- Lotka–Volterra Recovery: In Lotka–Volterra Scenario 1, the UDE was trained on 31 noisy trajectory points while assuming known linear terms and parameters.The neural network used two hidden layers with five neurons and Gaussian radial-basis activations.
- Lotka–Volterra Recovery: In Scenario 2, the UDE combined 61 x measurements with only 6 y measurements and used a shooting-like loss with parameter regularization.
- Lotka–Volterra Recovery: β ≈ 0.9239 and γ ≈ 0.8145: sparse regression recovered the quadratic Lotka–Volterra terms after UDE training, whereas direct identification from interpolated data failed.
- Lotka–Volterra Recovery: (50.4 ± 25.7)%: successful recovery was achieved across 498 error-free Scenario 1 runs with varied noise levels.Two runs failed because of numerical instabilities involving trained parameters and the numerical integrator.
- Hudson Bay Recovery: The Hudson Bay experiment recovered mixed prey–predator terms after subsampling the UDE at 0.5-year intervals and post-fitting the symbolic model.The reported fitted parameters were α = 0.557, δ = 0.826, β = −1.70, and γ = 2.04.
- Fisher–KPP Recovery: Uθ(ρ) = 1.0ρ−1.0ρ2: the Fisher–KPP experiment recovered the nonlinear term using sequentially thresholded least squares.
12 Adaptive Solving for the 100 Dimensional Hamilton-Jacobi-Bellman Equation
The section solves a high-dimensional semilinear parabolic PDE by converting it into a stochastic UDE and learning its unknown functional and initial-value components. Once trained, the learned initial condition provides the PDE solution at the starting point.
- The target is the solution at initial time t = 0 and starting point x = ζ for a semilinear parabolic PDE with known µ, σ, and f.
- The PDE is reformulated through a forward-backward SDE, with the unknown functional σT(t, Xt)∇u(t, Xt) and initial value U(0) approximated by universal approximators.
- The training loss measures the expected difference between the approximating solution and the required terminal condition g(XT).
- After training, θ2(X0) is the solution to the PDE, and the approximation is evaluated against the true solution.
- For the 100-dimensional example, the stochastic UDE used neural networks of widths 110 and was trained for 500 iterations with relative and absolute tolerances of 1e−4.
13 Reduction of the Boussinesq Equations
The section reduces the Boussinesq equations by learning a parameterized model from sparse PDE data and compares its predictions with training trajectories and a high-resolution benchmark. The examples also include adaptive solution of a 100-dimensional Hamilton-Jacobi-Bellman equation.
- A neural network was trained on 30 data points sampled from a diffusion-advection PDE to identify a missing function in its parameterization.
- Figure 15 shows an adaptive solution of the 100-dimensional Hamilton-Jacobi-Bellman equation as the universal approximator converges to the solution.
- Figure 16 compares training data with the trained UPDE over space at the 10th fitting time point and over time at the spatial midpoint.
- The trained neural network had a forward pass that took around 0.9 seconds.
- The full Boussinesq benchmark used Oceananigans.jl with adaptive CFL-limited time stepping and 128×128×128 spatial points.
14 Automated Derivation of Closure Relations for Viscoelastic Fluids
The section learns an additional memory-field closure for viscoelastic-fluid dynamics from simulated FENE-P data. The resulting UDE extends the model with neural networks while fitting stress trajectories under a specified initial condition.
- The FENE-P model captures normal stresses that store memory of deformation, while the Oldroyd-B model provides an approximation.
- The UDE adds a single memory field and is trained against simulated FENE-P data with λ = 2, L = 2, and η = 4.
- The added model uses two neural networks, each with one hidden layer containing 4 neurons and tanh activation.
- Training minimizes stress discrepancies over 100 evenly spaced time points in ti ∈[0, 2π] using ADAM with learning rate 0.015.
- The fluid is assumed to be at rest before t = 0, so the initial stress is zero.