Universal Differential Equations for Scientific Machine Learning

Introduction

Recent advances in machine learning have been dominated by deep learning which uses readily available “big data” to solve previously difficult problems such as image recognition and natural language processing . While some areas of science have begun to generate the large amounts of data required to train deep learning models, notably bioinformatics , in many areas the expense of scientific experiments has prohibited the effectiveness of these ground breaking techniques. In these domains, mechanistic models are still predominantly deployed due to the inaccuracy of deep learning techniques with small training datasets. While these mechanistic models are constrained to be predictive by uses prior structural knowledge (i.e. a form of inductive bias), the data-driven approach of machine learning can be more flexible and allows one to drop the simplifying assumptions required to derive theoretical models. The purpose of this work is to give a mathematical framework and accompanying software tool which bridges the gap by merging the best of both methodologies while mitigating the deficiencies.

This work falls into the burgeoning field of “scientific machine learning” which seeks to integrate machine learning derived idea into traditional engineering-related methods in order to utilize domain knowledge and known physical information into the learning process . It has recently been shown to be advantageous to merge differential equations with machine learning. Physics-Informed Neural Networks (PINNs) use partial differential equations in the cost functions of neural networks to incorporate prior scientific knowledge . While this has been shown to be a form of data-efficient machine learning for some scientific applications, PINNs frame the solution process as a large optimization. The data efficiency of PINNs has been shown to further improve when encoding structural assumptions like energy conservation by directly modeling the Hamiltonian . While this formalism incorporates the knowledge of physical systems into machine learning, it does not incorporate the numerical techniques which have led to stable and efficient solvers the large majority of scientific models. Treating the training process as fully implicit is compute-intensive which recent results have shown is increasingly difficult as the stiffness of the model increases . While work has demonstrated that efficiency is greatly improved when incorporating classical discretization techniques into the training process (as “discrete physics-informed neural networks”, such as in multi-step neural networks ), the formalism provided by PINNs is not conducive to the full use of classical model simulation software and thus every major PINN software framework, such as deepxde and SimNet , does not have the ability to automatically combine scientific machine learning training with the highly efficient differential equation solvers and adjoint techniques developed over the last century.

Thus as a basis for scientific machine learning that incorporates the efficient numerical solver and adjoint techniques, we developed a formalism which we denote the universal differential equation (UDE). UDEs are differential equations which are defined in full or part by a universal approximator. A universal approximator is a parameterized object capable of representing any possible function in some parameter size limit. Common universal approximators in low dimensions include Fourier or Chebyshev expansions, while common universal approximators in high dimensions include neural networks and other models used through machine learning. Mathematically in its most general form, the UDE is a forced stochastic delay partial differential equation (PDE) defined with embedded universal approximators:

where $\alpha(t)$ is a delay function and $W(t)$ is the Wiener process.

In this manuscript we describe how the tools of the SciML software ecosystem give rise to efficient training and analysis of the wide variety of UDEs which show up throughout the scientific literature. An astute reader may recognize a neural ordinary differential equation $u^{\prime}=\text{NN}_{\theta}(u,t)$ as the specific case of a one-dimensional UDE defined in full by a neural network . The previously mentioned discrete physics-informed neural networks also arise when one uses the appropriate fixed time step method as the solver for the UDE. In addition, neural network based optimal control , i.e. finding a neural network $c(u,t)$ which minimizes a loss of a differential equation solution $u^{\prime}=f(u,c(u,t),t)$, and model augmentation can similarly be phrased within this framework. Therefore, this architecture has an expansive reach in its utility throughout SciML and control applications.

It is thus in this context that we demonstrate the SciML ecosystem as a set of tools for handling the wide array of possible UDEs with solvers and adjoints handling all of the cases of adaptivity, stiffness, stochasticity, delays, and more. Figure 1 gives a general overview of how the tools of the SciML ecosystem connect to give rise to a modular and composable toolkit for scientific machine learning. The examples of this paper are constructed as the interplay of over 200 dependent libraries constructed by and for the SciML ecosystem, from high-level symbolic computing tools to low-level customized BLAS implementations to improve the performance of internal matrix factorizations over commonly used tools like OpenBLAS or MKL. To simplify the discussion, we will focus on how the composition of the numerical libraries and compiler-based code generation mechanisms gives rise to an efficient computing stack for SciML applications.

In Section 2.1 we showcase how the DifferentialEquations.jl solvers, the DiffEqSensitivity.jl adjoint methods, and the DiffEqFlux.jl helper functions combine to give rise to efficient and stable training framework for UDEs. In the sections following, we focus on use cases which go beyond the obvious neural ODE and optimal control applications in order to demonstrate the generality of the framework and formalism. In Section 2.3 we highlight how the DataDrivenDiffEq.jl symbolic regression tooling can be used in conjunction with the UDE framework in order to augment models with known physical information and decrease the data requirements for symbolic model discovery. In Section 2.4 we showcase how recent methods in solving high-dimensional parabolic partial differential equations can be phrased as a universal stochastic differential equation, and thus these methodologies can be extended to high order, implicit, and adaptive methods through this formulation on the SciML tools. In Section 2.5 we demonstrate how the discovery of non-local operators for reduced order modeling, known as parameterizations in the climate modeling literature, can similarly be phrased as a UDE training problem. Together this shows how the SciML ecosystem and its UDE formalism substantially advances the ability for scientists and engineers to combine all scientific knowledge with the latest techniques of machine learning and numerical analysis to arrive at a highly efficient and flexible set of automated training techniques.

Results

Training a UDE amounts to minimizing a cost function $C(\theta)$ defined on $u_{\theta}(t)$, the current solution to the differential equation with respect to the choice of parameters $\theta$. One choice of cost function is the Euclidean distance $C(\theta)=\sum_{i}\left\|u_{\theta}(t_{i})-d_{i}\right\|$ at discrete data points $(t_{i},d_{i})$. When optimized with local derivative-based methods, such as stochastic gradient decent, ADAM , or L-BFGS , this requires the calculation of $\frac{dC}{d\theta}$ which by the chain rule amounts to calculating $\frac{du}{d\theta}$. Thus the problem of efficiently training a UDE reduces to calculating gradients of the differential equation solution with respect to parameters.

Efficient methods for calculating these gradients are commonly called adjoints . The asymptotic computational cost of these methods does not grow multiplicatively with the number of state variables and parameters like numerical or forward sensitivity approaches, and thus it has been shown empirically that adjoint methods are more efficient on large parameter models . The DiffEqSensitivity.jl module uses the composable multiple dispatch based modular architecture of the DifferentialEquations.jl in order to extend all 300+ solvers for a wide variety of equations, including ordinary differential equations, stochastic differential equations, delay differential equations, differential-algebraic equations, stochastic delay differential equations, and hybrid equations which incorporate jumps and Levy processes. The full set of adjoint options available in this new software, which includes continuous adjoint methods and pure reverse-mode AD approaches, is described in Supplement S6.

To the authors’ knowledge, this is the first differential equation library which includes choices from all of the demonstrated categories of forward and adjoint sensitivity analysis methods. The importance of this fact is because each of these methods offers a substantial trade-off and thus all of these adjoints have different contexts in which they are the most applicable. Methods via solving ODEs and SDEs in reverse are the common adjoint utilized in neural ODE software such as torchdiffeq and are O(1) in memory (and the neural SDE can utilize the virtual Brownian tree for O(1) in memory ), but are known to be unstable under certain conditions such as on stiff equations . Checkpointed interpolation adjoints are available which do not require stable reversibility of the ODEs while retaining a relatively low-memory implementation via checkpointing. Another stabilized adjoint technique is the continuous quadrature adjoint which trades off increased memory use to reduce the computational complexity with respect to parameters from cubic to linear. These methods are unique to the SciML ecosystem and have recently demonstrated around three orders of magnitude performance improvements for large stiff partial differential equations . Section 2.2 and Section 2.5.1 are noted cases which is not stable under the reversed adjoint but stable under the checkpointing and quadrature adjoint approach, where in the former it is demonstrated that the adjoint techniques of alternative packages diverge due to this numerical issue. All of the aforementioned adjoint methods fall under the continuous optimize-then-discretize approach. Through the integration with automatic differentiation, discrete adjoint sensitivity analysis is implemented through both tape-based reverse-mode and source-to-source translation , with computational trade-offs between the two approaches. The former can be faster on scalarized heterogeneous differential equations while the latter is more optimized for homogeneous vectorized functions calls like are demonstrated in neural networks and discretizations of partial differential equations. Supplement S8 describes the literature behind continuous vs discrete adjoint approaches and showcases it as a general trade-off between performance and stability. Additionally, continuous and discrete forward mode sensitivity analysis approaches are also provided and optimized for equations with smaller numbers of parameters. Given the large number of variables involved in choosing a correct derivative calculation method, Supplement S7 describes a decision tree to guide users towards an appropriate choice. A compiler analysis of the user’s differential equation function is provided by default to automatically choose an efficient adjoint choice.

As described in Supplement S6.2, these adjoints use reverse-mode automatic differentiation for vector-transposed Jacobian products within the adjoint definitions to reduce the computational complexity. This has been shown to contribute up to two orders of magnitude in performance increases for large stiff differential equations over purely numerical vector-Jacobian product approaches which are common in other stiff ODE solver libraries with adjoints like Sundials CVODES and PETSc TS . In addition, the module DiffEqFlux.jl handles compatibility with the Flux.jl neural network library so that common deep architectures, such as convolutional neural networks and recurrent neural networks, automatically uses all efficient backpropagation kernels whenever encountered in the derivative calculation of differential equations. Thus together these three tools, DifferentialEquations.jl, DiffEqSensitivity.jl, and DiffEqFlux.jl, combine to give a UDE training framework which covers the vast set of combinations to allow efficiently training each model.

2 Features and Performance of the SciML Ecosystem

We assessed the viability of alternative differential equation libraries for universal differential equation workflows by comparing the features and performance of the given libraries. Table 1 demonstrates that the SciML ecosystem is the only differential equation solver library with deep learning integration that supports stiff ODEs, DAEs, DDEs, stabilized adjoints, distributed and multithreaded computation. We note the importance of the stabilized adjoints in Section 2.5.1 as many PDE discretizations with upwinding exhibit unconditional instability when reversed, and thus this is a crucial feature when training embedded neural networks in many PDE applications. Table 2 demonstrates that the SciML ecosystem exhibits more than an order of magnitude performance when solving ODEs against torchdiffeq of up to systems of 1 million equations. Because the adjoint calculation itself is a differential equation, this also corresponds to increased training times on scientific models. We note that torchdiffeq’s adjoint calculation diverges on all but the first two examples due to the aforementioned stability problems.

To reinforce this measure of performance, Supplement S10 demonstrates a 100x performance difference over torchdiffeq when training the spiral neural ODE from . Additionally we note that the author of the tfdiffeq library of TensorFlow has previous concluded “speed is almost the same as the PyTorch (torchdiffeq) code base ($\pm 2\%$)”https://github.com/titu1994/tfdiffeq/tree/v0.0.1-pre0#caveats. In addition, Supplement S10 demonstrates a 1,600x performance advantage for the SciML ecosystem over torchsde using the geometric Brownian motion example from the torchsde documentation . Given the computational burden, the mix of stiffness, and non-reversibility of the examples which follow in this paper, these results demonstrate that the SciML ecosystem is the first deep learning integrated differential equation software ecosystem that can train all of the equations necessary for the results of this paper. Note that this does not infer that our solvers will demonstrate more than an order of magnitude performance difference or more on all types of equations. For example, large non-stiff neural ODEs used in image classification tend to be dominated by large dense matrix multiplications and thus are in a different performance regime. However, these results demonstrate that on the equations generally derived from scientific models (ODEs derived from PDE semi-discretizations, heterogeneous differential equation systems, and neural networks in sufficiently small systems) that an order of magnitude or more performance difference exists over a wide set of cases.

In the following sections we will demonstrate the various applications which are enabled by this combination of efficient libraries.

3 Knowledge-Enhanced Symbolic Regression via UDEs

Automatic reconstruction of models from observable data has been extensively studied. Many methods produce non-symbolic representations by learning functional representations or through dynamic mode decomposition (DMD, eDMD) . Symbolic reconstruction of equations has utilized symbolic regressions which require a pre-chosen basis , or evolutionary methods to grow a basis . However, a common thread throughout much of the literature is that added domain knowledge constrains the problem to allow for more data-efficient reconstruction . Here we detail how using UDEs with symbolic regression can improve the data efficiency and applicability of the techniques.

As a motivating example, take the Lotka-Volterra system:

Assume that a scientist has a time series of measurements dense in the prey $x$ and sparse in the predator $y$ from this system but knows the birth rate $\alpha$ of $x$ and assumes a linear decay of the population $y$. With only this information, a scientist can propose the knowledge-based UODE as:

To then learn the missing portions of the equation, we can directly apply symbolic regression to the trained UAs in order to arrive at symbolic equations suggesting the missing equations. In contrast, the popular SINDy method normally approximates derivatives using a spline over the data points and subsequently applies symbolic regression to the full equation. The use of such interpolating spline techniques for derivative estimates implies dense enough data for accurate derivative calculations, which is not required in the UODE framework as the trained neural network can be sampled continuously and information can be spread throughout different data sources (such as structural inductive bias) as long as a causal and identifiable relation within the system is present. Figure 2 shows the UODE-based symbolic regression, Figure 3 showcases its ability to fully and accurately recover the correct symbolic equations from the limited data, while the same symbolic regression method with the same data used in the SINDy approach with derivative smoothing is unable to recover the equations. Supplement S11.1 further demonstrates additional examples using this approach and showcases how the UDE improves the performance in the presense of noise. Together this highlights how incorporating prior structural knowledge through the UDE framework can improve the performance of symbolic regression techniques.

3.2 Generalizing Sparse Regression of Non-ODEs via UDEs

The use of UDEs for extending sparse regression training methods also applies to extending the types of problems which sparse regression can be applied to. For example, one might wish to encode prior knowledge of the conservation equation $1=\sum_{i}u_{i}(t)$ to the evolution of the states of a discovered model. In this case, a universal differential-algebraic equation (DAE) of the form:

can be utilized to encode this prior knowledge as a DAE $Mu^{\prime}=U_{\theta}(u)$ where $M$ is singular. One the UA is trained, symbolic regression on the $U_{\theta}(u)$ subsequently discovers the dynamical equations. Similarly, one may know the general evolution of the system but be unaware of the stochastic features, leading to a universal stochastic differential equation of the form:

Tutorials in DiffEqFlux.jl showcase how to train the embedded UAs within these objects, which can then subsequently be symbolically regressed upon as in Section 2.3.

As a concrete example, we show how UDEs can turn equation discovery of partial differential equations into a low-dimensional sparse regression problem. To demonstrate discovery of spatio-temporal equations directly from data, we consider data generated from the one-dimensional Fisher-KPP (Kolmogorov–Petrovsky–Piskunov) PDE :

with periodic boundary conditions. Such reaction-diffusion equations appear in diverse physical, chemical and biological problems . To learn from the generated data, we define the UPDE:

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

It is impractical to solve high dimensional PDEs with mesh-based techniques since the number of mesh points scales exponentially with the number of dimensions. Given this difficulty, mesh-free methods based on universal approximators such as neural networks have been constructed to allow for direct solving of high dimensional PDEs . Recently, methods based on transforming partial differential equations into alternative forms, such as backwards stochastic differential equations (BSDEs), which are then approximated by neural networks have been shown to be highly efficient on important equations such as the nonlinear Black-Scholes and Hamilton-Jacobi-Bellman (HJB) equations . Here we will showcase how one of these methods, a deep BSDE method for semilinear parabolic equations , can be reinterpreted as a universal stochastic differential equation (USDE) to generalize the method and allow for enhancements like adaptivity, higher order integration for increased efficiency, and handling of stiff driving equations through the SciML software.

with a terminal condition $u(T,x)=g(x)$. Supplement S12 describes how this PDE can be solved by approximating by approximating the FBSDE:

where $U^{1}_{\theta_{1}}$ and $U^{2}_{\theta_{2}}$ are UAs and the loss function is given by the requiring that the terminating condition $g(X_{T})=u(X_{T},W_{T})$ is satisfied.

A fixed time step Euler-Maryumana discretization of this USDE gives rise to the deep BSDE method . However, this form as a USDE generalizes the approach in a way that makes all of the methodologies of our USDE training library readily available, such as higher order methods, adaptivity, and implicit methods for stiff SDEs. As a motivating example, consider the classical linear-quadratic Gaussian (LQG) control problem in 100 dimensions:

We solve the PDE by training the USDE using an adaptive Euler-Maruyama method as described in Supplement S12. Supplementary Figure 15 showcases that this methodology accurately solves the equations, effectively extending recent algorithmic advancements to adaptive forms simply be reinterpreting the equation as a USDE. While classical methods would require an amount of memory that is exponential in the number of dimensions making classical adaptively approaches infeasible, this approach is the first the authors are aware of to generalize the high order, adaptive, highly stable software tooling to the high-dimensional PDE setting.

5 Accelerated Scientific Simulation with Automatically Constructed Closure Relations

As an example of directly accelerating existing scientific workflows, we focus on the Boussinesq equations . The Boussinesq equations are a system of 3+1-dimensional partial differential equations acquired through simplifying assumptions on the incompressible Navier-Stokes equations, represented by the system:

where ${\bf u}=(u,v,w)$ is the fluid velocity, $p$ is the kinematic pressure, $\nu$ is the kinematic viscosity, $\kappa$ is the thermal diffusivity, $T$ is the temperature, and $b$ is the fluid buoyancy. We assume that density and temperature are related by a linear equation of state so that the buoyancy $b$ is only a function $b=\alpha gT$ where $\alpha$ is the thermal expansion coefficient and $g$ is the acceleration due to gravity.

This system is commonly used in climate modeling, especially as the voxels for modeling the ocean in a multi-scale model that approximates these equations by averaging out the horizontal dynamics $\overline{T}(z,t)=\iint T(x,y,z,t)\,dx\,dy$ in individual boxes. The resulting approximation is a local advection-diffusion equation describing the evolution of the horizontally-averaged temperature $\overline{T}$:

This one-dimensional approximating system is not closed since $\overline{wT}$ (the horizontal average temperature flux in the vertical direction) is unknown. Common practice closes the system by manually determining an approximating $\overline{wT}$ from ad-hoc models, physical reasoning, and scaling laws. However, we can utilize a UDE-automated approach to learn such an approximation from data. Let

where $P$ are the physical parameters of the Boussinesq equation at different regimes of the ocean, such as the amount of surface heating or the strength of the surface winds . Using data from average temperatures $\overline{T}$ and known physical parameters $P$, the non-locality of the convection term may be captured by training a universal diffusion-advection partial differential equation. Supplementary Figure 16 demonstrates the accuracy of the approach using a deep UPDE with high order stabilized-explicit Runge-Kutta (ROCK) methods where the fitting is described in Supplement S13. To contrast the trained UPDE, we directly simulated the 3D Boussinesq equations under similar physical conditions and demonstrated that the neural parameterization results in around a 15,000x acceleration. This demonstrates that physical-dependent parameterizations for acceleration can be directly learned from data utilizing the previous knowledge of the averaging approximation and mixed with a data-driven discovery approach.

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

All continuum materials satisfy conservation equations for mass and momentum. The difference between an elastic solid and a viscous fluid comes down to the constitutive law relating the stresses and strains. In a one-dimensional system, an elastic solid satisfies $\sigma=G\gamma$, with stress $\sigma$, strain $\gamma$, and elastic modulus $G$, whereas a viscous fluid satisfies $\sigma=\eta\dot{\gamma}$, with viscosity $\eta$ and strain rate $\dot{\gamma}$. Non-Newtonian fluids have more complex constitutive laws, for instance when stress depends on the history of deformation,

alternatively expressed in the instantaneous form :

where the history is stored in $\phi_{i}$. To become computationally feasible, the expansion is truncated, often in an ad-hoc manner, e.g. $\phi_{n}=\phi_{n+1}=\cdots=0$, for some $n$. Only with a simple choice of $G(t)$ does an exact closure condition exist, e.g. the Oldroyd-B model. For a fully nonlinear approximation, we train a UODE according to the details in Supplement S14 to learn a closure relation:

from the numerical solution of the FENE-P equations, a fully non-linear constitutive law requiring a truncation condition . Figure 5 compares the neural network approach to a linear, Oldroyd-B like, model for $\sigma$ and showcases that the nonlinear approximation improves the accuracy by more than 50x. We note that the neural network approximation accelerates the solution by 2x over the original 6-state DAE, demonstrating that the universal differential equation approach to model acceleration is not just applicable to large-scale dynamical systems like PDEs but also can be effectively employed to accelerate small scale systems.

Discussion

While many attribute the success of deep learning to its blackbox nature, the key advances in deep learning applications have come from incorporating inductive bias into architectures. Deep convolutional neural networks for image processing directly utilized the local spatial structure of images by modeling convolution stencil operations. Similarly, recurrent neural networks encode a forward time progression into a deep learning model and have excelled in natural language processing and time series prediction. Here we present a software designed for generating such inductive biased architectures for a wide variety of scientific domains. Our results show that by building these hybrid mechanistic models with machine learning, we can arrive at similar efficiency advancements by utilizing all known prior knowledge of the underlying problem’s structure. While we demonstrate the utility of UDEs in equation discovery, we have also demonstrated that these methods are capable of solving many other problems such as high dimensional partial differential equations. Many methods of recent interest, such as discrete physics-informed neural networks, can additionally be written as discretizations of UDEs and thus efficiently implemented using the SciML tools.

Our software implementation is the first deep learning integrated differential equation library to include the full spectrum of adjoint sensitivity analysis methods that is required to both efficiently and accurately handle the range of training problems that can arise from universal differential equations. We have demonstrated orders of magnitude performance advantages over previous machine learning enhanced adjoint sensitivity ODE software in a variety of scientific models and demonstrated generalizations to stiff equations, DAEs, SDEs, and more. While the results of this paper span many scientific disciplines and incorporate many different modeling approaches, together all of the examples shown in this manuscript can be implemented using the SciML software ecosystem in just hundreds of lines of code each, with none of the examples taking more than half an hour to train on a standard laptop. This both demonstrates the efficiency of the software and its methodologies, along with the potential to scale to much larger applications.

Code and Data Availability

The code for reproducing the computational experiments can be found at:

The following are the core libraries of the SciML ecosystem which together implement the functionality described in the paper:

All of the data for the experiments are simulated in the example codes.

Acknowledgements

We thank Jesse Bettencourt, Mike Innes, and Lyndon White for being instrumental in the early development of the DiffEqFlux.jl library, Tim Besard and Valentin Churavy for the help with the GPU tooling, and David Widmann and Kanav Gupta for their fundamental work across DifferentialEquations.jl. Special thanks to Viral Shah and Steven Johnson who have been helpful in the refinement of these ideas. We thank Charlie Strauss, Steven Johnson, Nathan Urban, and Adam Gerlach for enlightening discussions and remarks on our manuscript and software. We thank Stuart Rogers for his careful read and corrections. We thank David Duvenaud for extended discussions on this work. We thank the author of the torchsde library, Xuechen Li, for optimizing the SDE benchmark code.

References

DiffEqFlux.jl Pullback Construction

The DiffEqFlux.jl pullback construction is not based on just one method but instead has a dispatch-based mechanism for choosing between different adjoint implementations. At a high level, the library defines the pullback on the differential equation solve function, and thus using a differential equation inside of a larger program leads to this chunk as being a single differentiable primitive that is inserted into the back pass of Flux.jl when encountered by overloading the Zygote.jl and ChainRules.jl rule sets. For any ChainRules.jl-compliant reverse-mode AD package in the Julia language, when a differential equation solve is encountered in any Julia library during the backwards pass, the adjoint method is automatically swapped in to be used for the backpropagation of the solver. The choice of the adjoint is chosen by the type of the sensealg keyword argument which are fully described below.

Given a function $f(x)=y$, the pullback at $x$ is the function:

where $f^{\prime}(x)$ is the Jacobian $J$. We note that $B_{f}^{x}(1)=\left(\nabla f\right)^{T}$ for a function $f$ producing a scalar output, meaning the pullback of a cost function computes the gradient. A general computer program can be written as the composition of discrete steps:

and thus the vector-Jacobian product can be decomposed:

which allows for recursively decomposing a the pullback to a primitively known set of $\mathcal{B}_{f^{i}}^{x}$:

where $x_{i}=\left(f^{i}\circ f^{i-1}\circ\ldots\circ f^{1}\right)(x)$. Implementations of code generation for the backwards pass of an arbitrary program in a dynamic programming language can vary. For example, building a list of function compositions (a tape) is provided by libraries such as Tracker.jl and PyTorch , while other libraries perform direct generation of backward pass source code such as Zygote.jl , TAF , and Tapenade .

2 Backpropagation-Accelerated DAE Adjoints for Index-1 DAEs with Constraint Equations

Before describing the modes, we first describe the adjoint of the differential equation with constraints. The following derivation is based on but modified to specialize on index 1 DAEs with linear mass matrices. The constrained ordinary differential equation:

We wish to solve for some cost function $G(u,p)$ evaluated throughout the differential equation, i.e.:

To derive this adjoint, introduce the Lagrange multiplier $\lambda$ to form:

Since $u^{\prime}=f(u,p,t)$, we have that:

for $s_{i}$ being the sensitivity of the ith variable. After applying integration by parts to $\lambda^{\ast}Ms^{\prime}$, we require that:

If $G$ is discrete, then it can be represented via the Dirac delta:

at the data points $(t_{i},d_{i})$. Therefore, the derivative of an ODE solution with respect to a cost function is given by solving for $\lambda^{\ast}$ using an ODE for $\lambda^{T}$ in reverse time, and then using that to calculate $\frac{dG}{dp}$. At each time point where discrete data is found, $\lambda$ is then changed using a callback (discrete event handling) by the amount $g_{u}$ to represent the Dirac delta portion of the integral. Lastly, we note that $\frac{dG}{du_{0}}=-\lambda(0)$ in this formulation.

We have to take care of consistent initialization in the case of semi-explicit index-1 DAEs. We need to satisfy the system of equations

where d and a denote differential and algebraic variables, and $f$ and $g$ denote differential and algebraic equations respectively. Combining the above two equations, we know that we need to increment the differential part of $\lambda$ by

with $\mu(T)=0$ can be appended to the system of equations to perform the quadrature for $\frac{dG}{dp}$. We note that this formulation allows for a single linear solve to generate a guaranteed consistent initialization via a linear solve. This is a major improvement over the previous technique when applied to index 1 DAEs in mass matrix form since the alternative requires a full consistent initialization of the DAE, a problem which is known to have many common failure modes .

3 Current Adjoint Calculation Methods

From this setup we have the following 8 possible modes for calculating the adjoint, with their pros and cons.

QuadratureAdjoint: a quadrature-based approach. This utilizes interpolation of the forward solution provided by DifferentialEquations.jl to calculate $u(t)$ at arbitrary time points for doing the calculations with respect to $\frac{df}{du}$ in the reverse ODE of $\lambda$. From this a continuous interpolatable $\lambda(t)$ is generated, and the integral formula for $\frac{dG}{dp}$ is calculated using the QuadGK.jl implementation of Gauss-Kronrod quadrature. While this approach is memory heavy due to requiring the interpolation of the forward and reverse passes, it can be the fastest version for cases where the number of ODE/DAE states is small and the number of parameters is large since the QuadGK quadrature can converge faster than ODE/DAE-based versions of quadrature. This method requires an ODE or a DAE.

InterpolatingAdjoint: a checkpointed interpolation approach. This approach solves the $\lambda(t)$ ODE in reverse using an interpolation of $u(t)$, but appends the equations for $\mu(t)$ and thus does not require saving the timeseries trajectory of $\lambda(t)$. For checkpointing, a scheme similar to that found in SUNDIALS is used. Points $(u_{k},t_{k})$ from the forward solution are chosen as the interval points. Whenever the backwards pass enters a new interval, the ODE is re-solved on $t\in[t_{k-1},t_{k}]$ with a continuous interpolation provided by DifferentialEquations.jl. For the reverse pass, the tstops argument is set for each $t_{k}$, ensuring that no backwards integration step lies in two checkpointing intervals. This requires at most a total of two forward solutions of the ODE and the memory required to hold the interpolation of the solution between two consecutive checkpoints. Note that making the checkpoints at the start and end of the integration interval makes this equivalent to a non-checkpointed interpolated approach which replaces the quadrature with an ODE/SDE/DAE solve for memory efficiency. This method tends to be both stable and require a minimal amount of memory, and is thus the default. This method requires an ODE, SDE, or a DAE.

BacksolveAdjoint: a checkpointed backwards solution approach. Following , after a forward solution, this approach solves the $u(t)$ equation in reverse along with the $\lambda(t)$ and $\mu(t)$ ODEs. Thus, since no interpolations are required, it requires $\mathcal{O}(1)$ memory. Unfortunately, many theoretical results show that backwards solution of ODEs is not guaranteed to be stable, and testing this adjoint on the universal partial differential equations like the diffusion-advection example of this paper showcases that it can be divergent and is thus not universally applicable, especially in cases of stiffness. Thus for stability we modify this approach by allowing checkpoints $(u_{k},t_{k})$ at which the reverse integration is reset, i.e. $u(t_{k})=u_{k}$, and the backsolve is then continued. The tstops argument is set in the integrator to require that each checkpoint is hit exactly for this resetting to occur. By doing so, the resulting method utilizes $\mathcal{O}(1)$ memory + the number of checkpoints required for stability, making it take the least memory approach. However, the potential divergence does lead to small errors in the gradient, and thus for highly stiff equations we have found that this is only applicable to a certain degree of tolerance like $10^{-6}$ given reasonable numbers of checkpoints. When applicable this can be the most efficient method for large memory problems. This method requires an ODE, SDE, or a DAE.

ForwardSensitivity: a forward sensitivity approach. From $u^{\prime}=f(u,p,t)$, the chain rule gives $\frac{d}{dt}\frac{du}{dp}=\frac{df}{du}\frac{du}{dp}+\frac{df}{dp}$ which can be appended to the original equations to give $\frac{du}{dp}$ as a time series, which can then be used to compute $\frac{dG}{dp}$. While the computational cost of the adjoint methods scales like $\mathcal{O}(N+P)$ for $N$ differential equations and $P$ parameters, this approach scales like $\mathcal{O}(NP)$ and is thus only applicable to models with small numbers of parameters (thus excluding neural networks). However, when the universal approximator has small numbers of parameters, this can be the most efficient approach. This method requires an ODE or a DAE.

ForwardDiffSensitivity: a forward-mode automatic differentiation approach, using ForwardDiff.jl to calculate the forward sensitivity equations, i.e. an AD-generated implementation of forward-mode “discretize-then-optimize”. Because it utilizes a forward-mode approach, the scaling matches that of the forward sensitivity approach and it tends to have similar performance characteristics. This method applies to any Julia-based differential equation solver.

TrackerAdjoint: a Tracker-driven taped-based reverse-mode discrete adjoint sensitivity, i.e. an AD-generated implementation of reverse-mode “discretize-then-optimize”. This is done by using the TrackedArray constructs of Tracker.jl to build a Wengert list (or tape) of the forward execution of the ODE solver which is then reversed. This method applies to any Julia-based differential equation solver.

ZygoteAdjoint: a Zygote-driven source-to-source reverse-mode discrete adjoint sensitivity, i.e. an AD-generated implementation of reverse-mode “discretize-then-optimize”. This utilizes the Zygote.jl system directly on the differential equation solvers to generate a source code for the reverse pass of the solver itself. Currently this is only directly applicable to a few differential equation solvers, but is under heavy development.

ReverseDiffAdjoint: A ReverseDiff.jl taped-based reverse-mode discrete adjoint sensitivity, i.e. an AD-generated implementation of reverse-mode “discretize-then-optimize”. In contrast to TrackerAdjoint, this methodology can be substantially faster due to its ability to precompile the tape but only supports calculations on the CPU.

For each of the non-AD approaches, there are the following choices for how the Jacobian-vector products $Jv$ (jvp) of the forward sensitivity equations and the vector-Jacobian products $v^{\prime}J$ (vjp) of the adjoint sensitivity equations are computed:

Automatic differentiation for the jvp and vjp. In this approach, automatic differentiation is utilized for directly calculating the jvps and vjps. ForwardDiff.jl with a single dual dimension is applied at $f(u+\lambda\epsilon)$ to calculate $\frac{df}{du}\lambda$ where $\epsilon$ is a dual dimensional signifier. For the vector-Jacobian products, a forward pass at $f(u)$ is utilized and the backwards pass is seeded at $\lambda$ to compute the $\lambda^{\prime}\frac{df}{du}$ (and similarly for $\frac{df}{dp}$). Note that if $f$ is a neural network, this implies that this product is computed by starting the backpropagation of the neural network with $\lambda$ and the vjp is the resulting return. Four methods are allowed to be chosen for performing the internal vjp calculations:

Zygote.jl source-to-source transformation based vjps. Note that only non-mutating differential equation function definitions are supported in this mode. This mode is the most efficient in the presence of neural networks.

Enzyme.jl source-to-source transformation basd vjps. This is the fastest vjp choice in the presence of heavy scalar operations like in chemical reaction networks, but is currently not compatible with garbage collection and thus requires non-allocating $f$ functions.

ReverseDiff.jl tape-based vjps. This allows for JIT-compilation of the tape for accelerated computation. This is a the fast vjp choice in the presence of heavy scalar operations like in chemical reaction networks but more general in application than Enzyme. It is not compatible with GPU acceleration.

Tracker.jl with arrays of tracked real values is utilized on mutating functions.

The internal calculation of the vjp on a general UDE recurses down to primitives and embeds optimized backpropagations of the internal neural networks (and other universal approximators) for the calculation of this product when this option is used.

Numerical differentiation for the jvp and vjp. In this approach, finite differences is utilized for directly calculating the jvps and vjps. For a small but finite $\epsilon$, $\left(f(u+\lambda\epsilon)-f(u)\right)/\epsilon$ is used to approximate $\frac{df}{du}\lambda$. For vjps, a finite difference gradient of $\lambda^{\prime}f(u)$ is used.

Automatic differentiation for Jacobian construction. In this approach, (sparse) forward-mode automatic differentiation is utilized by a combination of ForwardDiff.jl with SparseDiffTools.jl for color-vector based sparse Jacobian construction. After forming the Jacobian, the jvp or vjp is calculated.

Numerical differentiation for Jacobian construction. In this approach, (sparse) numerical differentiation is utilized by a combination of DiffEqDiffTools.jl with SparseDiffTools.jl for color-vector based sparse Jacobian construction. After forming the Jacobian, the jvp or vjp is calculated.

In total this gives 48 different adjoint method approaches, each with different performance characteristics and limitations. A full performance analysis which measures the optimal adjoint approach for various UDEs has been omitted from this paper, since the combinatorial nature of the options requires a considerable amount of space to showcase the performance advantages and generality disadvantages between each of the approaches. A follow-up study focusing on accurate performance measurements of the adjoint choice combinations on families of UDEs is planned.

Sensitivity Algorithm Decision Tree

If no enealg is provided by the user, DiffEqSensitivity.jl will automatically choose an adjoint technique from the list. By default, a stable adjoint with an auto-adapting vjp choice is used. The following decision tree is similar to that within the defaults (but is continuously updating due to various performance changes).

If there are 50 parameters+states or less, consider using forward-mode sensititivites. If the f function is not ForwardDiff-compatible, use ForwardSensitivty, otherwise use ForwardDiffSensitivty as its more efficient.

For larger equations, give BacksolveAdjoint and InterpolatingAdjoint a try. If the gradient of BacksolveAdjoint is correct, many times it’s the faster choice so choose that (but it’s not always faster!). If your equation is stiff or a DAE, skip this step as BacksolveAdjoint is almost certainly unstable.

If your equation does not use much memory and you’re using a stiff solver, consider using QuadratureAdjoint as it is asymptotically more computationally efficient by trading off memory cost.

If the other methods are all unstable (check the gradients against each other!), then ReverseDiffAdjoint is a good fallback on CPU, while TrackerAdjoint is a good fallback on GPUs.

After choosing a general sensealg, if the choice is InterpolatingAdjoint, QuadratureAdjoint, or BacksolveAdjoint, then optimize the choice of vjp calculation next:

If your function has no branching (no if statements) and is heavily scalarized, use ReverseDiffVJP(true).

If your calculations are on the CPU and your function is very scalarized in operations but has branches, choose ReverseDiffVJP().

If your calculations are on the CPU or GPU and your function is very vectorized, choose ZygoteVJP().

Else fallback to TrackerVJP() if Zygote does not support the function.

If none of the reverse-mode AD based vjps work on your function, fallback to autojacvec=true (for forward-mode AD via ForwardDiff) or false for numerical Jacobians.

Continuous vs Discrete Adjoints

Previous research has shown that the discrete adjoint approach is more stable than continuous adjoints in some cases while continuous adjoints have been demonstrated to be more stable in others and can reduce spurious oscillations . This trade-off between discrete and continuous adjoint approaches has been demonstrated on some equations as a trade-off between stability and computational efficiency . Care has to be taken as the stability of an adjoint approach can be dependent on the chosen discretization method , and our software contribution helps researchers switch between all of these optimization approaches in combination with hundreds of differential equation solver methods with a single line of code change.

Integration with Existing Code

The open-source differential equation solvers of DifferentialEquations.jl were developed in a manner such that all steps of the programs have a well-defined pullback when using a Julia-based backwards pass generation system. Our software allows for automatic differentiation to be utilized over differential equation solves without any modification to the user code. This enables the simulation software already written with DifferentialEquations.jl, including large software infrastructures such as the MIT-CalTech CLiMA climate modeling system and the QuantumOptics.jl simulation framework , to be compatible with all of the techniques mentioned in the rest of the paper. Thus while we detail our results in isolation from these larger simulation frameworks, the UDE methodology can be readily used in full-scale simulation packages which are already built on top of the Julia SciML ecosystem.

Benchmarks

The three ODE benchmarks utilized the Lorenz equations (LRNZ) weather prediction model from and the standard ODE IVP Testset :

The 28 ODE benchmarks utilized the Pleiades equation (PLEI) celestial mechanics simulation from and the standard ODE IVP Testset :

written in the form $u=[x_{i},y_{i},x^{\prime}_{i},y^{\prime}_{i}]$.

The rest of the benchmarks were derived from a discretization a two-dimensional reaction diffusion equation, representing systems biology, combustion mechanics, spatial ecology, spatial epidemiology, and more generally physical PDE equations:

where $\alpha_{A}(x)=1$ if $x>80$ and 0 otherwise on the domain $x\in$, $y\in$, and $t\in$ with zero-flux boundary conditions. For the purpose of parameter gradient tests, we treated calculated the derivative of the solution with respect to the entires of $\alpha_{A}(x)$. The diffusion constant $D$ was chosen as $100$ and all other parameters were left at $1.0$. In this parameter regime the ODE was non-stiff as indicated by solves with implicit methods not yielding performance advantages. The diffusion operator was discretized using the second order finite difference stencil on an $N\times N$ grid, where $N$ was chosen to be 16, 32, 64, 128, 256, and 512. To ensure fairness, the torchdiffeq functions were compiled using torchscript which we varified improved performance. The code for reproducing the benchmark can be found at:

We omit the the gradient performance benchmarks for this case from the main manuscript since the backsolve adjoint method of torchdiffeq is unstable on all of the partial differential equation examples. We note that similarly when using BacksolveAdjoint with the SciML tools we similarly see a divergence, while other adjoints such as InterpolatingAdjoint do not diverge, reinforcing the point that this is due to lack of stability in the algorithm. In the cases where torchdiffeq does not diverge, we see torchdiffeq 12,000x slower and 1,200x slower on Lorenz and Pleiades respectively. However, we note that this is because at the size of those equations forward sensitivity analysis is more efficient which is not available in torchdiffeq, and thus this overestimates the general performance difference in the derivative calculations.

2 Neural ODE Training Benchmark

The spiral neural ODE from was used as the benchmark for the training of neural ODEs. The data was generated according from the form:

where $A=[-0.1,2.0;-2.0,-0.1]$ on $t\in[0,1.5]$ where data was taken at 30 evenly spaced points. Each of the software packages trained the neural ODE for 500 iterations using ADAM with a learning rate of 0.05. The defaults using the SciML software resulted in a final loss of 4.895287e-02 in 7.4 seconds, the optimized version (choosing BacksolveAdjoint with compiled ReverseDiff vector-jacobian products) resulted in a final loss of 2.761669e-02 in 2.7 seconds, while torchdiffeq achieved a final loss of 0.0596 in 289 seconds. To ensure fairness, the torchdiffeq functions were compiled using torchscript. Code to reproduce the benchmark can be found at:

3 SDE Solve Benchmark

The torchsde benchmarks were created using the geometric Brownian motion example from the torchsde README. The SDE was a 4 independent geometric Brownian motions:

where $\mu=0.5$ and $\sigma=1.0$. Both software solved the SDE 100 times using the SRI method with fixed time step chosen to give 20 evenly spaced steps for $t\in$. The SciML ecosystem solvers solved the equation 100 times in 0.00115 seconds, while torchsde v0.1 took 1.86 seconds. We contacted the author who rewrote the Brownian motion portions into C++ and linked it to torchsde as v0.1.1 and this improved the timing to roughly 5 seconds, resulting in a final performance difference of approximately 1,600x. The code to reproduce the benchmarks and the torchsde author’s optimization notes can be found at:

Sparse Identification of Missing Model Terms via Universal Differential Equations

The SINDy algorithm enables data-driven discovery of governing equations from data. Notice that to use this method, derivative data $\dot{X}$ is required. While in most publications on the subject this information is assumed. However, for our studies we assume that only the time series information is available. Here we modify the algorithm to apply to only subsets of the equation in order to perform equation discovery specifically on the trained neural network, and in our modification the $\dot{X}$ term is replaced with $U_{\theta}(t)$, the output of the universal approximator, and thus is directly computable from any trained UDE.

After training the UDE, choose a set of state variables:

and compute a the action of the universal approximator on the chosen states:

Then evaluate the observations in a basis $\Theta(X)$. For example:

where $X^{P_{i}}$ stands for all $P_{i}$th order polynomial terms such as

Using these matrices, find this sparse basis $\mathbf{\Xi}$ over a given candidate library $\mathbf{\Theta}$ by solving the sparse regression problem $\dot{X}=\mathbf{\Theta}\mathbf{\Xi}$ with $L_{1}$ regularization, i.e. minimizing the objective function $\left\|\mathbf{\dot{X}}-\mathbf{\Theta}\mathbf{\Xi}\right\|_{2}+\lambda\left\|\mathbf{\Xi}\right\|_{1}$. This method and other variants of SINDy applied to UDEs, along with specialized optimizers for the LASSO $L_{1}$ optimization problem, have been implemented by the authors and collaborators as the DataDrivenDiffEq.jl library on top of the ModelingToolkit.jl and Symbolics.jl computer algebra systems.

On the Lotka-Volterra equations, we trained a UDE model in two different scenarios, starting from $x_{0}=0.44249296,~{}y_{0}=4.6280594$ with parameters are chosen to be $\alpha=1.3,~{}\beta=0.9,~{}\gamma=0.8,~{}\delta=1.8$. The neural network consists of an input layer, two hidden layers with 5 neurons and a linear output layer modeling the polynomial interaction terms. The input and hidden layer have gaussian radial basis activation functions. We trained for 200 iterations with ADAM with a learning rate $\gamma=10^{-1}$. We then switched to BFGS with an initial stepnorm of $\gamma=10^{-2}$ setting the maximum iterations to 10000. Typically the training converged after 400-600 iterations in total.

Scenario 1) consists of a trajectory with 31 points measured in $x(t),y(t)$ with a constant step size $\Delta t=0.1$. We assumed perfect knowledge about the linear terms of the equations and their corresponding parameters $\alpha,~{}\delta$. The trajectory has been perturbed with additive noise drawn from a normal distribution scaled by $5\%$ of its mean. The loss was chosen was the L2 loss $\mathcal{L}=\sum_{i}(u_{\theta}(t_{i})-d_{i})^{2}$.

Scenario 2) consists of a trajectory with 61 points measured in $x(t)$ with a constant step size $\Delta t=0.1$ and 6 points measured in $y(t)$ with a constant step size $\Delta t=1.2$. The trajectory has been perturbed with additive noise drawn from a normal distribution scaled by $1\%$ of its mean.We assumed perfect knowledge about the linear terms of the first differential equation and their corresponding parameters $\alpha$. In addition to the UDE a linear decay rate with unknown parameter was added in the second differential equation governing the predator dynamics. The parameter was included in the training process. The data was divided into $j=5$ different datasets $d$ containing $i=13$ measurements in $x,~{}t$ and 2 measurements in $y$ for the initial and boundary condition. The loss was chosen was similar to shooting like techniques plus a regularization term $\mathcal{L}=\sum_{j}(\sum_{i}(u_{x,\theta(t_{i},j)}-d_{x,i,j})^{2}+\|u_{y,\theta(t_{1}3,j)}-d_{y,2,j}\|)+\lambda\|\theta\|^{2}$.

From the trained neural network, data was sampled over the original trajectory and fitted using the SR3 method with varying threshold with $\lambda=exp10.(-7:0.1:3)$. A pareto optimal solution was selected via the L1 norm of the coefficients and the L2 norm of the corresponding error between the differential data and estimate. The knowledge-enhanced neural network returned zero for all terms except non-zeros on the quadratic terms with nearly the correct coefficients that were then fixed using an additional sparse regression over the quadratic terms. The resulting parameters extracted are $\beta\approx 0.9239$ and $\gamma\approx 0.8145$. Performing a sparse identification on the ideal, full solution and numerical derivatives computed via an interpolating spline resulted in a failure. After determining the results of the symbolic regression, the learned ODE model was then trained to refit the parameters before extrapolating. In comparison, interpolations of the sparse predator measurements have been performed using linear, quadratic and cubic spline and a polynomial interpolation of order 5.

Using SINDy with same library of candidate functions over the interpolated data and their numerical derivatives leads to the the results listed in 3. Given the sparsity of the data, quality of the resulting fit and its derivative, none of the attempts were able to recover the true equation describing the predator dynamics.

The UDE training procedure was evaluated 500 times in total for Scenario 1) given the initial trajectory with varying noise levels $0.1,~{}0.5,~{}1.0,~{}2.5,~{}5\%$ in terms of the mean of the trajectory (100 repetitions each). The results of the successful recovery of the missing equations are shown in Fig. 8, the training loss is shown in Fig. 9. Overall, the recovery rate is $(50.4\pm 25.7)\%$ for all 498 error-free runs . Two runs failed due to numerical instabilities of the combination of trained parameters and numerical integrator. Given that the study has conducted in a brute-force manor, e.g. no prior data cleaning, no hyperparameter optimization, no (automatic) pre-selection of the networks architecture or candidate selection has been performed, the sensitivity of the networks initial parameters is also captured in this study. However, since the scope of this work is to highlight the usefulness of UDEs as means to extract specific information from time series we limited ourselves to this naive approach. Further research can highlight additional methods to improve noise robustness.

In Fig. 10 some examples for failed recovery can be seen. It is worth noting that the mere visual quality of the fit points in most cases, except for Sample 450, would indicate a success. In some cases, e.g. Sample 500, some discontinuities can be seen in the solution, possibly due to overfitting. However, adding a regularization penalty to the overall loss function has shown little effect and was hence neglected.

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

Additionally, a full recovery based on a dataset of the Hudson Bay Data of Hares and Lynx between 1900 and 1920, taken from and originally published in has been performed. The data has been normalized to the interval $(0,1]$ and the assumed model is given as

Incorporating a linear birth rate of the prey and a linear decay rate of the predator. The neural network consists of an input layer, two hidden layers with 5 neurons and a linear output layer. Again, radial basis activations have been used, except for the last hidden layer which used a tanh activation. As in Scenario 2) we started by using a shooting loss with L2 regularization of the parameter, intentionally widthholding data of the predators to smoothly optimize the parameters using ADAM with a learning rate of $0.1$ for 100 iterations and BFGS with an initial stepnorm of $0.01$ until converged ( maximum 200 iterations). After this initial fit, we switched onto a an L2-Norm loss with regularization and trained until convergence using BFGS. The results are shown in Figure 11.

Afterwards we performed a sparse regression with a candidate library of multivariate polynomials up to degree 5 and sinusoidal signals in the states on the recovered signal of the UDE using sequentially thresholded regression. The UDE has been subsampled to an interval of 0.5 years, efficcently augmenting the limited data. The mixed terms were recovered successful. The resulting symbolic model has been post-fitted with an L2-Norm loss to decrease its resulting error, as can be seen in Figure 12. A long term estimation can seen in Figure 13, using the parameters $\alpha=0.557$, $\delta=0.826$, $\beta=-1.70$ and $\gamma=2.04$.

1.3 Application to the Reconstruction of the Fisher-KPP Equations

To generate training data for the 1D Fisher-KPP equation 7, we take the growth rate and the diffusion coefficient to be $r=1$ and $D=0.01$ respectively. The equation is numerically solved in the domain $x\in$ and $t\in[0,T]$ using a 2nd order central difference scheme for the spatial derivatives and the time-integration is done using the Tsitouras 5/4 Runge-Kutta method. We implement periodic boundary condition $\rho(x=0,t)=\rho(x=1,t)$ and initial condition $\rho(x,t=0)=\rho_{0}(x)$ is taken to be a localized function given by

with $\Delta=0.2$ which represents the width of the region where $\rho\simeq 1$. The data are saved at evenly spaced points with $\Delta x=0.04$ and $\Delta t=0.5$.

In the UPDE 8, the growth neural network $\textrm{NN}_{\theta}(\rho)$ has 4 densely connected layers with 10, 20, 20 and 10 neurons each and $\tanh$ activation functions. The diffusion operator is represented by a CNN that operates on an input vector of arbitrary size. It has 1 hidden layer with a $3\times 1$ filter $[w_{1},w_{2},w_{3}]$ without any bias. To implement periodic boundary conditions for the UPDE at each time step, the vector of values at different spatial locations $[\rho_{1},\rho_{2},\dots,\rho_{N_{x}}]$ is padded with $\rho_{N_{x}}$ to the left and $\rho_{1}$ to the right. This also ensures that the output of the CNN is of the same size as the input. The weights of both the neural networks and the diffusion coefficient are simultaneously trained to minimize the loss function

where $\lambda$ is taken to be $10^{2}$ (note that one could also structurally enforce $w_{3}=-(w_{1}+w_{2}))$. The second term in the loss function enforces that the differential operator that is learned is conservative—that is, the weights sum to zero. The training is done using the ADAM optimizer with learning rate $10^{-3}$.

Similar to the Lotka-Volterra experiments, symbolic regression over monomials up to degree 10 has been performed to recover the nonlinear term of the Fisher-KPP Equations $U_{\theta}$ from Section 11.1.4 in Scenario 3). The system has been simulated with $r=1$ for $t=s$ with a time resolution of $\Delta t=0.5s$ and 25 measurements in $x$. Normally distributed noise with $2.5\%$ of the mean has been added. Figure 14 shows the recovered dynamics nonlinear term via the UDE approach, using a network similar to Scenario 1). We trained the network using a loss function similar to Section 11.1.4 with a regularization of $1$ using ADAM with a learning rate of $0.1$ for 200 iterations and then switching to BFGS with an initial stepnorm of $0.01$. The recovered equation is $U_{\theta}(\rho)=1.0\rho-1.0\rho^{2}$. The symbolic regression has been performed using sequentially thresholded least squares over thresholds $\lambda=[10^{-3},10^{2}]$ logarithmic evenly distributed with a step size of $0.01$.

1.4 Analysis of Alternative Function Approximators in Fisher-KPP

Additionally, we analyzed the performance of alternative universal approximators for this spatiotemporal PDE discovery task. To do this, we analyzed both the parameter efficiency and the computational efficiency. For the parameter efficiency, we attempted to quantify the minimum numbers of parameters that could robustly identify the terms of the PDEs. It is predicted by the lottery ticket hypothesis that low parameter neural networks can generally produce accurate fits, but have a low probability of being discovered. Thus we defined robustness as the ability to perform 5 optimizations with random initializations mixed with stochastic optimizers and recover a fitted solution to a loss of 0.01. All optimizations used the same optimization process which was 400 iterations of ADAM to find a local minima and then BFGS to complete the optimization (which automatically exited when the solution achieves a local maximum according to the default detection criteria of the Optim.jl library ).

For the neural network we choose 1 hidden layer with a $\tanh$ activation function and scaled the size of the hidden layer from 2-4, corresponding to 4, 7, and 15 parameters. At 4 parameters, all 5 optimization runs failed to produce a loss of 0.01, while at 7 and 15 all optimizations passed. To test against a non-neural network function approximator we chose to use the Fourier basis. Tests with 3, 5, 7, and 15 parameters all give fits to a loss of 0.01 on all 5 optimization runs, demonstrating that the smaller parameter Fourier basis model was more robust than the smallest parameter neural network model. The least parameter neural network, the 7 parameter version, took approximately 2,500 seconds with a standard deviation of approximately 1,000 seconds, a minimum time of approximately 1,300 seconds and a maximum time of approximately 3,300 seconds. The Fourier basis with 7 parameters took approximately 250 seconds to train, with a standard deviation of approximately 5 seconds, a minimum time of approximately 242 seconds and a maximum time of approximately 255 seconds. This demonstrates that the training time with the Fourier basis was on average an order of magnitude faster than the neural network when comparing the same parameter size between the models.

All of the output losses and times for the experiments are stored as comments in the code for the experiments FisherKPPCNNSmall.jl and FisherKPPCNNFourier.jl in the example repository. The implementations of the classical basis functions can be found in the DiffEqFlux.jl repository https://diffeqflux.sciml.ai/dev/layers/BasisLayers/ with an associated tutorial on classical basis functions and TensorLayer https://diffeqflux.sciml.ai/dev/examples/tensor_layer/.

2 Discovery of Robertson’s Equations with Prior Conservation Laws

On Robertson’s equations, we trained a UDAE model against a trajectory of 10 points on the timespan $t\in[0.0,1.0]$ starting from $y_{1}=1.0$, $y_{2}=0.0$, and $y_{3}=0.0$. The parameters of the generating equation were $k_{1}=0.04$, $k_{2}=3e7$, and $k_{3}=1e4$. The universal approximator was a neural network with one hidden layers of size 64. The equation was trained using the BFGS optimizer to a loss of $9e-6$.

Adaptive Solving for the 100 Dimensional Hamilton-Jacobi-Bellman Equation

with a terminal condition $u(T,x)=g(x)$. In this equation, tr is the trace of a matrix, $\sigma^{T}$ is the transpose of $\sigma$, $\nabla u$ is the gradient of $u$, and $\text{Hess}_{x}u$ is the Hessian of $u$ with respect to $x$. Furthermore, $\mu$ is a vector-valued function, $\sigma$ is a $d\times d$ matrix-valued function and $f$ is a nonlinear function. We assume that $\mu$, $\sigma$, and $f$ are known. We wish to find the solution at initial time, $t=0$, at some starting point, $x=\zeta$.

Let $W_{t}$ be a Brownian motion and take $X_{t}$ to be the solution to the stochastic differential equation

with a terminal condition $u(T,x)=g(x)$. With initial condition $X(0)=\zeta$ has shown that the solution to 9 satisfies the following forward-backward SDE (FBSDE) :

with terminating condition $g(X_{T})=u(X_{T},W_{T})$. Notice that we can combine 60 and 12.1 into a system of $d+1$ SDEs:

where $U_{t}=u(t,X_{t})$. Since $X_{0}$, $\mu$, $\sigma$, and $f$ are known from the choice of model, the remaining unknown portions are the functional $\sigma^{T}(t,X_{t})\nabla u(t,X_{t})$ and initial condition $U(0)=u(0,\zeta)$, the latter being the point estimate solution to the PDE.

To solve this problem, we approximate both unknown quantities by universal approximators:

Therefore we can rewrite 62 as a stochastic UDE of the form:

with initial condition $(X_{0},U_{0})=(X_{0},U^{2}_{\theta_{2}}(X_{0}))$.

To be a solution of the PDE, the approximation must satisfy the terminating condition, and thus we define our loss to be the expected difference between the approximating solution and the required terminating condition:

Finding the parameters $(\theta_{1},\theta_{2})$ which minimize this loss function thus give rise to a BSDE which solves the PDE, and thus $U^{2}_{\theta_{2}}(X_{0})$ is the solution to the PDE once trained.

2 The LQG Control Problem

This PDE can be rewritten into the canonical form by setting:

where $\overline{\sigma}=\sqrt{2}$, T = 1 and $X_{0}=(0,...,0)\in R^{100}$. The universal stochastic differential equation was then supplemented with a neural network as the approximator. The initial condition neural network was had 1 hidden layer of size 110, and the $\sigma^{T}(t,X_{t})\nabla u(t,X_{t})$ neural network had two layers both of size 110. For the example we chose $\lambda=1$. This was trained with the LambaEM method of DifferentialEquations.jl with relative and absolute tolerances set at $1e-4$ using 500 training iterations and using a loss of 100 trajectories per epoch.

On this problem, for an arbitrary $g$, one can show with Itô’s formula that:

which was used to calculate the error from the true solution.

Reduction of the Boussinesq Equations

As a test for the diffusion-advection equation parameterization approach, data was generated from the diffusion-advection equations using the missing function $\overline{wT}=\cos(\sin(T^{3}))+\sin(\cos(T^{2}))$ with $N$ spatial points discretized by a finite difference method with $t\in[0,1.5]$ with Neumann zero-flux boundary conditions. A neural network with two hidden layers of size 8 and $\tanh$ activation functions was trained against 30 data points sampled from the true PDE. The UPDE was fit by using the ADAM optimizer with learning rate $10^{-2}$ for 200 iterations and then ADAM with a learning rate of $10^{-3}$ for 1000 iterations. The resulting fit is shown in 16 which resulted in a final loss of approximately $0.007$. We note that the stabilized adjoints were required for this equation, i.e. the backsolve adjoint method was unstable and results in divergence and thus cannot be used on this type of equation. The trained neural network had a forward pass that took around 0.9 seconds.

For the benchmark against the full Bossinesq equations, we utilized Oceananigans.jl . It was set to utilize adaptive time stepping to maximize the time step according to the CFL condition number (capped at $\text{CFL}\leq 0.3$) and matched the boundary conditions, along with setting periodic boundary conditions in the horizontal dimension. The Bossinesq simulation used $128\times 128\times 128$ spatial points, a larger number than the parameterization, in order to accurately resolve the mean statistics of the 3-dimensional dynamics as is commonly required in practice . The resulting simulation took 13,737 seconds on the same computer used for the neural diffusion-advection approach, demonstrating the approximate 15,000x acceleration.

Automated Derivation of Closure Relations for Viscoelastic Fluids

is the upper convected derivative, and $L$, $\eta$, $\lambda$ are parameters . For a one dimensional strain rate, $\dot{\gamma}=\dot{\gamma}_{12}=\dot{\gamma}_{21}\neq 0$, $\dot{\gamma}_{ij}=0$ else, the one dimensional stress required is $\sigma=\sigma_{12}$. However, $\sigma_{11}$ and $\sigma_{22}$ are both non-zero and store memory of the deformation (normal stresses). The Oldroyd-B model is the approximation:

As an arbitrary nonlinear extension, train a UDE model using a single additional memory field against simulated FENE-P data with parameters $\lambda=2$, $L=2$, $\eta=4$. The UDE model is of the form,

where $U_{0},U_{1}$ are neural networks each with a single hidden layer containing 4 neurons. The hidden layer has a tanh activation function. The loss was taken as $\mathcal{L}=\sum_{i}(\sigma(t_{i})-\sigma_{\text{FENE-P}}(t_{i}))^{2}$ for 100 evenly spaced time points in $t_{i}\in[0,2\pi]$, and the system was trained using an ADAM iterator with learning rate 0.015. The fluid is assumed to be at rest before $t=0$, making the initial stress also zero.