Universal Differential Equations for Scientific Machine Learning

1 Introduction

• UDEs merge mechanistic structure with universal approximators while retaining scientific solvers and adjoint methods, providing a common framework for neural ODEs, optimal control, model augmentation, and broader SciML applications.The SciML ecosystem supports adaptivity, stiffness, stochasticity, delays, distributed parallelism, and GPU-compatible training across a modular software stack.Source: §1 · pp. 2–4

2.1 Efficient and Stable Training of Universal Differential Equations

• The training framework combines DifferentialEquations.jl, DiffEqSensitivity.jl, and DiffEqFlux.jl to provide adjoint choices across more than 300 solvers and equation types, including stabilized methods for stiff systems.Reverse adjoints can be unstable for stiff equations, whereas checkpointed and quadrature adjoints trade memory and computational complexity for stability; compiler analysis can select an efficient method.Source: §2.1 · pp. 4, 6

2.2 Features and Performance of the SciML Ecosystem

• More than 10x faster ODE solving than torchdiffeq was observed for systems up to 1 million equations, while torchdiffeq’s adjoint diverged on all but two examples.Supplementary benchmarks also report 100x faster spiral neural ODE training than torchdiffeq and 1,600x faster geometric Brownian motion training than torchsde, with performance gains applying broadly to scientific-model equations rather than all neural ODEs.Source: §2.2 · p. 7, Table 2 · p. 8

2.3.1 Improved Identification of Nonlinear Interactions with Universal Ordinary Differential Equations

• UODE-enhanced symbolic regression recovers correct Lotka–Volterra equations from limited, unevenly observed data that fails under the corresponding SINDy procedure with derivative smoothing.The UODE learns unknown interactions while preserving known structure, then symbolic regression converts the learned universal approximator into equations and can improve performance under noise.Source: §2.3.1 · pp. 8–9, Figure 2 · p. 9, Figure 3 · p. 10

2.3.2 Generalizing Sparse Regression of Non-ODEs via UDEs

• UDEs extend sparse regression to DAEs, stochastic equations, and PDEs while reducing PDE discovery to a low-dimensional regression problem by encoding structural operators.For Fisher–KPP data, the method recovers the [1, −2, 1] Laplacian stencil, diffusion constant, and quadratic growth term; a Fourier-layer approximator trains an order of magnitude faster than a neural network in this case.Source: §2.3.2 · pp. 9–11, Figure 4 · p. 12

2.4 Computationally-Efficient Solving of High-Dimensional Partial Differential Equations

• Recasting deep BSDE methods as universal stochastic differential equations makes higher-order, adaptive, and implicit methods available for solving high-dimensional semilinear parabolic PDEs.This formulation is intended to handle stiff driving equations while retaining neural-network approximations for equations such as nonlinear Black–Scholes and Hamilton–Jacobi–Bellman problems.Source: §2.4 · pp. 12–13, §2.4.1 · p. 13

2.4.1 Adaptive Solution of High-Dimensional Hamilton-Jacobi-Bellman Equations

• For a 100-dimensional linear-quadratic Gaussian control problem, the universal stochastic formulation supplies the training framework used to solve the associated Hamilton–Jacobi–Bellman equation.The method is obtained by viewing the fixed-step Euler–Maruyama deep BSDE discretization as a special case of a more general USDE.Source: §2.4.1 · p. 13
• An adaptive Euler-Maruyama-trained USDE accurately solves high-dimensional Hamilton-Jacobi-Bellman equations while extending stable adaptive tooling beyond classical methods.The method addresses the exponential-memory barrier of classical approaches and generalizes high-order, adaptive, highly stable software tooling to high-dimensional PDEs.Source: §2.4.1 · p. 14

2.5.1 Automated Discovery of Large-Eddy Model Parameterizations

• A trained UPDE accurately reproduces the averaged Boussinesq dynamics and accelerates 3D simulation by around 15,000x.The approach learns physical-parameter-dependent convection parameterizations from averaged temperatures and known ocean-regime parameters using stabilized ROCK methods.Source: §2.5.1 · pp. 14–15

2.5.2 Data-Driven Nonlinear Closure Relations for Model Reduction in Non-Newtonian Viscoelastic Fluids

• For FENE-P fluids, the learned nonlinear constitutive closure improves accuracy by more than 50x and accelerates the original 6-state DAE by 2x.The UODE avoids ad-hoc history truncation by learning the closure relation from numerical solutions, outperforming a linear Oldroyd-B-like model.Source: §2.5.2 · pp. 15–16

3 Discussion

• UDEs encode scientific inductive biases in hybrid mechanistic–machine-learning models, supporting equation discovery, high-dimensional partial differential equations, and discretized physics-informed neural networks.The framework extends beyond equation discovery to multiple scientific problem classes and efficiently implements related methods through SciML tools.Source: §3 · p. 16
• The SciML implementation is the first deep-learning-integrated differential-equation library to provide the full spectrum of adjoint sensitivity methods needed for UDE training.It generalizes to stiff equations, DAEs, SDEs, and other training problems while outperforming earlier machine-learning-enhanced adjoint-sensitivity ODE software by orders of magnitude.Source: §3 · p. 16
• Each manuscript example can be implemented in hundreds of lines of SciML code, with none requiring more than half an hour to train on a standard laptop.These results demonstrate software and methodological efficiency while indicating potential scalability to substantially larger applications.Source: §3 · p. 18

Code availability

• The SciML ecosystem’s core libraries implement UDE training, forward and adjoint rules, differential-equation solvers, and symbolic regression.The listed repositories are DiffEqFlux.jl, DiffEqSensitivity.jl, DifferentialEquations.jl, and DataDrivenDiffEq.jl.Source: §4 · p. 18

Data availability

• All experimental data are simulated in the example codes rather than provided as separate datasets.Source: §4 · p. 18

6 DiffEqFlux.jl Pullback Construction

• DiffEqFlux.jl defines the differential-equation solve as one differentiable primitive and dispatches the adjoint according to the sensealg type.The construction integrates with ChainRules.jl-compliant reverse-mode AD packages by overloading Zygote.jl and ChainRules.jl rules.Source: §6 · p. 30

6.1 Pullbacks and Adjoints

• Pullbacks recursively decompose vector-Jacobian products into a primitively known set, while the pullback of a scalar-output cost function computes its gradient.Source: §6.1 · p. 30

1 DAEs with Constraint Equations

• The DAE adjoint solves λ backward in time, applies callback-based jumps at discrete data points, and computes parameter and initial-condition derivatives from λ.For discrete costs, Dirac-delta contributions are represented by adjoint changes at observation times, with dG/du0 = −λ(0).Source: §1 · p. 32
• A single linear solve guarantees consistent initialization for semi-explicit index-1 DAEs, improving on mass-matrix approaches with common initialization failures.The formulation combines differential and algebraic constraints while incrementing the differential adjoint component at callbacks.Source: §1 · p. 32

6.3 Current Adjoint Calculation Methods

• Combining eight adjoint methods with four Jacobian-product strategies yields 48 approaches, but the paper omits a full performance comparison because of their combinatorial tradeoffs.A follow-up study is planned to measure these combinations across families of UDEs.Source: §6.3 · pp. 32, 34, 35
• InterpolatingAdjoint is the default because it is stable and memory-efficient, using checkpointed forward interpolation and at most two forward ODE solutions.It supports ODEs, SDEs, and DAEs while storing interpolation only between consecutive checkpoints.Source: §6.3 · p. 33
• BacksolveAdjoint uses the least memory but can diverge backward, producing gradient errors and limiting stiff-equation use to tolerances such as 10−6.Checkpoints reset the reverse integration and add memory proportional to the checkpoints required for stability.Source: §6.3 · p. 33
• Forward sensitivity methods scale as O(NP), so they suit models with few parameters, whereas adjoint methods scale as O(N + P) and better accommodate neural networks.ForwardDiffSensitivity has similar scaling and performance characteristics, while ForwardSensitivity can be efficient for small universal approximators.Source: §6.3 · pp. 33–34

• For equations with 50 parameters or states or fewer, use forward sensitivities, choosing ForwardSensitivity or the more efficient ForwardDiffSensitivity when compatible.Source: §7 · p. 36
• For larger equations, try BacksolveAdjoint and InterpolatingAdjoint, favoring BacksolveAdjoint when its gradient is correct unless stiffness or DAEs make it unstable.Source: §7 · p. 36
• With limited memory and a stiff solver, QuadratureAdjoint can improve computational efficiency by trading increased memory cost for asymptotic savings.Source: §7 · p. 36
• If other methods are unstable, compare gradients and use ReverseDiffAdjoint on CPUs or TrackerAdjoint on GPUs as fallbacks.Source: §7 · p. 36
• After selecting InterpolatingAdjoint, QuadratureAdjoint, or BacksolveAdjoint, choose ReverseDiffVJP(true) for branchless scalarized functions and ReverseDiffVJP() for branched CPU-scalarized functions.Source: §7 · p. 36
• Use ZygoteVJP() for vectorized CPU or GPU calculations, TrackerVJP() when Zygote lacks support, and autojacvec=true or false for forward-mode or numerical Jacobians.Source: §7 · p. 36

• Discrete adjoints can be more stable in some cases, while continuous adjoints can be more stable in others and reduce spurious oscillations [53] [47].Source: §8 · pp. 36–37
• The choice between discrete and continuous adjoints can trade stability against computational efficiency, and adjoint stability depends on the discretization method [102].Source: §8 · p. 37
• The software supports switching among optimization approaches and hundreds of differential equation solver methods with a single line of code change.Source: §8 · p. 37

9 Integration with Existing Code

• Automatic differentiation works through DifferentialEquations.jl differential-equation solves without requiring user-code modifications, extending UDE techniques to existing simulation software.The approach supports full-scale infrastructures including the MIT-CalTech CLiMA climate modeling system [116] and QuantumOptics.jl.Source: §9 · p. 37

10.1 ODE Solve Benchmarks

• SciML was 12,000x faster than torchdiffeq on Lorenz and 1,200x faster on Pleiades when torchdiffeq did not diverge.The comparison is partly attributed to forward sensitivity analysis being more efficient for these equation sizes, while torchdiffeq lacks that capability; backsolve adjoints diverged on all PDE examples.Source: §10.1 · p. 38

10.2 Neural ODE Training Benchmark

• SciML’s optimized neural ODE achieved final loss 2.761669e-02 in 2.7 seconds, compared with torchdiffeq’s 0.0596 loss in 289 seconds.SciML defaults achieved final loss 4.895287e-02 in 7.4 seconds; all packages trained for 500 ADAM iterations at learning rate 0.05.Source: §10.2 · p. 39

10.3 SDE Solve Benchmark

• SciML solved the geometric Brownian-motion SDE 1,600x faster than optimized torchsde, with times of 0.00115 seconds and roughly 5 seconds, respectively.Both solvers used the SRI method for 100 runs with 20 evenly spaced steps over t ∈[0, 1].Source: §10.3 · p. 39

11 Sparse Identification of Missing Model Terms via Universal Differential Equations

• The modified SINDy procedure replaces unavailable derivative data with the trained UDE output Uθ(t), then solves ˙X = ΘΞ with L1-regularized sparse regression over a candidate library.The approach selects state variables, evaluates the universal approximator and basis Θ(X), and is implemented with UDE/SINDy variants in DataDrivenDiffEq.jl on ModelingToolkit.jl and Symbolics.jl.Source: §11 · pp. 39–40

11.1.1 Application to the Partial Reconstruction of the Lotka-Volterra

• The UDE-based sparse identification recovered β ≈0.9239 and γ ≈0.8145 from partial Lotka–Volterra observations, whereas interpolation-based SINDy failed to recover the predator dynamics.The recovered neural-network output was zero for non-quadratic terms, and the quadratic coefficients were subsequently refined by sparse regression.Source: §11.1.1 · p. 42, Table 3 · p. 43
• 50.4 ± 25.7% of 498 error-free runs recovered the missing Lotka–Volterra equations across noise levels, while two runs failed from numerical instabilities.The study used 500 Scenario 1 repetitions with noise levels 0.1, 0.5, 1.0, 2.5, and 5% of the trajectory mean; results are shown in Fig. 8.Source: §11.1.1 · p. 42, Figure 8 · p. 44

11.1.2 Application to the Reconstruction of the Lotka-Volterra Using the Hudson Bay Dataset

• For the Hudson Bay dataset, the study presents recovered equations before and after parameter fitting, along with a resulting model and its extrapolation.Source: Figure 12 · p. 47, Figure 13 · p. 47

11.1.3 Application to the Reconstruction of the Fisher-KPP Equations

• For noisy Fisher–KPP data, the recovered nonlinear reaction term was Uθ(ρ) = 1.0ρ−1.0ρ2 using sequentially thresholded least squares over monomial candidates.The simulation used r = 1, 2.5% mean-scaled normal noise, and 25 spatial measurements.Source: §11.1.3 · pp. 48–49, Figure 14 · p. 49

11.1.4 Analysis of Alternative Function Approximators in Fisher-KPP

• Fourier bases were more robust and approximately an order of magnitude faster than equally sized neural networks: 7-parameter models passed all five optimizations and took approximately 250 seconds versus approximately 2,500 seconds.The 4-parameter neural network failed all five runs, while 7- and 15-parameter neural networks passed; Fourier models with 3, 5, 7, and 15 parameters all passed.Source: §11.1.4 · pp. 49–50

11.2 Discovery of Robertson’s Equations with Prior Conservation Laws

• For Robertson’s equations, a one-hidden-layer UDAE trained on 10 trajectory points achieved a loss of 9e −6 while using the prior reaction parameters and initial conditions.The neural network had hidden-layer size 64 and was trained with BFGS over t ∈[0.0, 1.0].Source: §11.2 · p. 50

12.1 Forward-Backwards SDE Formulation

• The PDE solution is represented through a forward-backward SDE whose unknown gradient term and initial value are approximated by universal approximators.The trained initial-condition network output provides the PDE solution at t = 0 and x = ζ.Source: §12.1 · pp. 51–52
• Training minimizes the expected terminal mismatch between the approximated backward solution and g(XT ), yielding a BSDE that solves the PDE.The loss is l(θ1, θ2|XT , UT ) = E [∥g(XT ) −UT ∥].Source: §12.1 · p. 52

12.2 The LQG Control Problem

• For the 100-dimensional LQG problem, the stochastic differential equation was supplemented with neural networks for the initial condition and σT (t, Xt)∇u(t, Xt).The networks used one hidden layer of size 110 for the initial condition and two layers of size 110 for the gradient-related term.Source: §12.2 · p. 52
• The LQG system was trained with LambaEM from DifferentialEquations.jl for 500 iterations using 100 trajectories per epoch and relative and absolute tolerances of 1e−4.The setup used λ = 1, T = 1, and X0 = (0, ..., 0) ∈R100.Source: §12.2 · p. 52
• For arbitrary g, Itô’s formula supplies the true solution used to calculate the approximation error.This provides the reference for evaluating the learned solution in the LQG example.Source: §12.2 · p. 52

• A two-layer tanh neural network fit 30 PDE data points with final loss approximately 0.007 for the diffusion-advection parameterization.Training used ADAM at learning rates 10−2 for 200 iterations and 10−3 for 1000 iterations, with Neumann zero-flux boundary conditions.Source: §13 · p. 52, Figure 16 · p. 54
• The backsolve adjoint was unstable and diverged, so stabilized adjoints were required; the trained neural network’s forward pass took around 0.9 seconds.The fit used a missing function wT = cos(sin(T 3)) + sin(cos(T 2)).Source: §13 · pp. 52, 54
• 13,737 seconds versus approximately 15,000x acceleration for the full Boussinesq benchmark, which used 128×128×128 spatial points and adaptive CFL-capped time stepping.The benchmark matched boundary conditions and used more spatial points than the parameterization to resolve mean statistics.Source: §13 · p. 54, Figure 16 · p. 54

Physical models and closure

• The full FENE-P formulation accounts for nonzero normal stresses that store deformation memory alongside the one-dimensional shear stress.Source: §14 · pp. 54–55
• Oldroyd-B is presented as an approximation to FENE-P, alongside an exact closure relation for the one-dimensional stress.Source: §14 · p. 55

UDE construction

• A UDE extends the viscoelastic model with one additional memory field and is trained against simulated FENE-P data using λ = 2, L = 2, and η = 4.Source: §14 · p. 55

Training procedure

• The UDE uses two single-hidden-layer, four-neuron tanh networks and minimizes stress-squared error over 100 evenly spaced times in [0, 2π] with ADAM at learning rate 0.015.The fluid is assumed initially at rest, so its initial stress is zero.Source: §14 · p. 55