Differential Equations as a Neural Network Layers | by Kevin Hannay | Towards Data Science

The main idea of artificial neural networks (ANN) is to build up representations for complicated functions using compositions of relatively simple functions called layers.
A deep neural network is one that has many layers, or many functions composed together.
Although layers are typically simple functions( e.g. relu(Wx+b)) in general they could be any differentiable functions.
The layer is specified by some finite vector of parameters θ ∈ Rp. To be practically useful we need to be able to fit this layer (or layers) to data. This involves defining a cost or loss function which measures the closeness of the model prediction to the data.
If our layers are differentiable then we can find the gradient of this cost function ∇C(θ) and use this to find a local minimum of the cost in an efficient manner.
Here I am considering differential equations models. These systems describe the time evolution of the state of a system (x) in time, using an expression which involves the derivative of the current state. In general, they take the form:
A differential equation fits into our neural network framework, as it takes in some parameters and produces the solution as output and it is differentiable.
Thus, we can use a differential equation as a layer in a neural network. This is really neat for a few reasons:
Differential equations are the fundamental language of all physical laws.
Outside of physics and chemistry differential equations are an important tool in describing the behavior of complex systems. Using differential equations models in our neural networks allows these models to be combined with neural networks approaches.
Building effective neural networks involves choosing a good underlying structure for the network. It is often easier to think of describing how a function changes in time, then it is to write down a function for the phenomena. Relatively simple rate laws can turn into very complex behavior (see the Lorenz system below!).
For simplicity, in this article, I am going to focus on neural networks with a single differential equations based layer in this article. However, these layers could easily be embedded as one layer in a deep learning project. The combination of deep learning with domain knowledge in the form of a differential equation is a game-changer in many fields.
To demonstrate how you can build your own differential equations layers into neural networks I am going to make use of the Julia Flux, DiffEqFlux and DifferentialEquations libraries. In the article, I will keep the code to small snippets to allow for developing intuition, but the full code for these examples has been posted in this Github repo.
To begin we will consider fitting the parameters of a classic model in mathematical biology. The Lotka-Volterra predator prey model describes the oscillations in a population that can be observed in simple ecological communities.
A famous example of this can be observed in the Hudson bay trading company fur trade data from the 1800s.
The equations for the model are given by
The predator species oscillations will tend to lag behind the peak in the prey species population.
Now let’s use this model as a layer in a neural network, with a set of six parameters. Two for the initial populations of the prey x0 and predators y0 and the four rate parameters α, β, δ, γ.
In Julia, we can define this system and a set of default parameters.
Now we define the differential equation layer as a function of the input parameters. The parameter array has the first two entries as the initial conditions.
The next thing to do is generate some fake data to fit from (add some noise) and then define a loss function. I used a simple mean square error loss function for our data, although we certainly could add some regularization of the parameters to the loss function easily.
Finally, we define a function to actually perform the model fitting. Here is where the Julia ecosystem really pays off as this system can easily use gradient methods for fitting the parameters.
As you can see I use 3000 Adams iterations followed by a call to BFGS to fit minimize the cost. This is a good combination, you could just jump straight to BFGS but the convergence will take much longer. It is better to let the Adams iterations get you in the neighborhood and then finish with BFGS.
The recovered parameters are given by:
The parameters are [0.9987532818277052, 0.9809616721237869, 1.5075095360807464, 1.009470157647888, 2.974936372236448, 0.993477599859459] with final loss value 2.1229151208653736The actual parameters should be: [1,1,1.5,1,3,1]
Here is a plot of our data against the simulated time-series data. As you can see the minimization finds and excellent fit to the data.
Of course, the real test for a neural network is how it generalizes to validation data. In this case, we can easily test the system for an unobserved initial condition to see if our neural network has captured the dynamics.
Great it works!
Now let’s consider the problem of recovering the parameters from a dynamical system that often comes up in my research field. The Stuart-Landau equation describes the dynamics of a nonlinear system which is near a Hopf bifurcation (common way oscillations emerge in nonlinear systems). I’m going to consider a variation of this which comes up in the study of synchrony in populations of coupled oscillators.
Part of the motivation for using this model is due to my familiarity with this model. However, I have another deeper motivation. The Lotka-Volterra predator-prey model is a great model to start with, but it is kind of a weird model for mathematical biology. It is a conservative system, much like is observed in physical system (without friction). This means instead of having an isolated attracting oscillation (called a limit cycle) the Lotka-Volterra system has a whole class of concentric cycles.
This has the potential to skew our understanding of fitting differential equation models, because conservative systems retain a memory of initial conditions. For limit cycle oscillations we lose information about the initial conditions over time (transient dynamics decay, arrow of time). Since many initial conditions end up at the same final state this could make the question of parameter recovery more difficult.
The equations for the system are given by:
these are written in polar coordinates. Thus R describes the amplitude and ψ the angle in radians. The dynamics take the qualitative form of
In Julia it is a simple matter to update our predator prey model to allow for this type of system.
Now we define a ode layer generate some data and create a loss function.
Finally, the following function creates runs an optimization to recover parameters.
The recovered parameters are given by:
The parameters are [0.02293619193779294, 0.08918144968755219, 0.08304049248335474, 6.275224626122101, 0.11589347566722544, 0.9019465183126854] with final loss value 0.7300967336849258The actual parameters are:[0.0,0.08,0.1, 6.28, 0.1, 1.0]
The fit for the training data is shown below:
We can test our model further by looking at how well it generalizes for new initial values.
Now time for a fun application. Let’s consider a parameter recovery using time-series data from a chaotic system. For this, we will use the famous lorenz system.
It is simple enough to define this system in Julia:
As a matter of demonstration let’s take a look at some time series data for the X variable generated from slightly different initial condition values.
This plot demonstrates a sensitive dependence on initial conditions which is a hallmark of chaos.
As usual, we define a differential equation layer which depends on initial conditions and parameters.
Generate some data and add some noise to the time-series and define our usual error function.
Now define the training function using the DiffEqFlux library.
Running the optimization gives a good fit to the training data, which generalizes well when we integrate for times beyond the training sample.
We can also look at the training model fit in a time-series plot.
Since the model was only trained on t ∈ [0,10], we can see that the model generalizes pretty well for t ∈ [10,20] but quickly diverges from the true value after that time.
The parameters are [0.9727986684296102, 0.034212291835795355, -0.001627150535904438, 9.99887680497206, 28.00267472723392, 2.666573010932292] with final loss value 1.7015586033442565The actual parameters are:[1.0,0.0,0.0,10.0,28.0,2.6666666666666665]
So we have found a pretty close fit.
It is relatively straightforward to include differential equation models into neural networks using the Julia ecosystem (scientific machine learning sciml). This allows us to include whole branches of knowledge through classical dynamical systems models into our neural network models for time-series data. The next step is to combine these differential equations layers into our deep learning models.
On the subject of the interface of dynamical systems and artificial neural networks. Three possible generalizations of what we have begun here are:
Neural differential equations is a term that is used to describe using an artificial neural network function as the right-hand side of a dynamical system. Since these systems make use of a general ANN function they can show poor convergence in modeling time-series.
Universal differential equations is a term that has been used for mixing a mathematical model and a ANN function. The idea here is to specify as much as possible of the dynamics and then let the ANN fill in the unmodeled dynamics, control function, etc. These systems can show much faster convergence and allow us to combine the power of neural networks to describe high dimensional functions with domain knowledge.
Hybrid dynamical systems are used to combine differential equation models (or more generally domain knowledge of some kind) with a separate machine learning approach. Many variations of these types exist but I would distinguish this class that it separates the machine learning model from the domain model. They can then be combined or fed into one another as an input.
Also, I should mention that all of the models we discussed in this article can easily be moved to the GPU for training. One simply has to declare the initial conditions to be on the GPU.
Reminder all of the code for this article can be found here.