Augmenting Neural Differential Equations to Model Unknown Dynamical Systems with Incomplete State Information

The evolution of a system can be computed when encoded as an ordinary differential equation. Unfortunately, it is common to not know the functional form of a system of interest. By appeal to the universal approximation theorem[3], a neural network can be trained to approximate the ODE governing the system. Since the neural ODE[1, 2] can be trained from data, a neural ODE can learn the rules governing a physical system with an unknown ODE functional form.

However, an ODE cannot be evolved when the state of the initial state of the system is incompletely known. For example, the rate of change of populations in the Lotka-Volterra model of a predator-prey system depends only on the predator and prey populations; but if the number of predators is not known then the rate of change of the prey cannot be computed and so the system cannot be evolved in time.

A complication occurs in training a neural ODE when the state trajectory is missing for some of the state variables. For example, if we know only the prey population time series but not the predator population then the training loss function won’t demand the predator population be predicted correctly.

Intuitively, there is some reason not to despair. If we knew the exact functional form of the right-hand side $F(.)$, but not its parameters, then, in favorable cases, we might still be able to infer a varying predator population just by observing the prey population with time. That is, a falling prey population requires there exist a predator, and it falls faster the more numerous predators. Thus perhaps we might infer the proportional fluctuation of the predators from just the prey population dynamics even without knowing the initial condition.

A worse situation occurs when the reason we lack the data on a variable is that we don’t even know the variable is a vital part of the state equations. For example, if we were not aware bears existed, a two-species ODE for rabbit and wolf population mechanics would be missing dependencies on an evolving bear predator population. This doesn’t simply make the problem harder; it may entirely prevent a solution by a two-variable ODE. The feasible trajectory space of a two-variable ODE (rabbits and wolves) does not include every trajectory available to those same two variables when embedded in a three-variable ODE (rabbits, wolves, and bears).

This paper examines that hard challenge: whether the system can still be solved when both the entire functional form of the right-hand side is unknown and, simultaneously, all of the training data for one state variable’s trajectory is unknown. Furthermore, we restrict the training to short time series but require that it extrapolate to long-duration time series in testing. In each case study, we find the neural ODE successfully learns from just short time-series data then successfully predicts much longer time spans even for initial conditions absent in the training set. The usefulness of neural ODEs for modeling unknown dynamical systems is thus demonstrated successfully.

Our principal work on this was completed in October 2019 and was subsequently published and presented in multiple, peer-reviewed, venues [8, 9, 10, 11]. In those we described our two novel Neural ODE architectures as "recurrent external channel" neural ODEs to distinguish them from the conventional recurrent neural network (internal hidden nodes). As it happens, one of our two architectures is similar to one presented at NeuroIPS 2019 by Dupont et al, but neither of us was aware of the other’s contemporaneous developments and so until now have not cited each other.[5, 12] In the interest of not fragmenting public terminology we now adopt Dupont et al’s catchy acronym "ANODE" in this summary publication. The "recurrent external channels" described in our previous work are logically the same as the augmentation variables "$\hat{a}(t)$" in Dupont et al.

Dupont et al addressed a different problem than the one we posed. Moreover, that work set the undefined initial condition of the augmentation variables to a constant (zero) which was sufficient for their posed problem. But in general, this doesn’t entirely resolve the non-crossing rule of ODE phase space trajectories that the ANODE (a.k.a. recurrent external channel) architecture was intended to address in the first place. In particular, it did not provide feasible solutions in the problems we posed for unknown systems (case-2 or case-3 below). To fully resolve this in case-3, we added a second neural net to compute unique initial conditions for the recurrent external channels that were conditioned on the data itself. Furthermore, this required that we trained both of these neural nets simultaneously. (See Methods section.)

For our case-1 (validation) experiment, we deliberately chose an identical ODE functional form and neural architecture to that of Rackauckas et al. to conveniently verify our implementation. Surprisingly, we observed a different, significantly improved, outcome arising from a modest change in the training protocol. The changes (discussed below) appear to make the solution process more robust for somewhat stiff oscillatory equations and that protocol advancement was critical in the more challenging case-2 and -3 presented below.

The neural ODE is trained on population time series data with nothing withheld. This control will validate the neural ODE is sufficiently deep to reproduce the ODE dynamics and the training protocol is sufficient on short-time spans. Testing will compare the extrapolation accuracy on long time-span test-sets.

In this case, the data for the bear population is withheld while training the neural network, except we supply the true bear population at t=0 as the initial condition. The training loss function contains just the wolf and rabbit populations, ignoring the bear population. If you like, you may imagine it as a partially known variable – you know bears are in a forest affecting the system, but they are only measured once at the start because it is too dangerous to the graduate students to repeatedly count bears in the woods.

One can see that this architecture has the capacity to correctly represent the differential equation since it has three inputs and three outputs. That is, if the neural network just happened to begin emitting the correct bear rate of change given the correct number of bear inputs, that’s self-consistent and thus it has to work just as well as the previous model. But one might doubt it can be successfully trained to achieve that. First, since there’s no constraint on the "bear" channel in the loss function it does not have to emit the correct bear population to achieve low loss. It could, for example, emit any monotonic transformation of the bear population (e.g. the cube root of the tangent of the bear population), because there are infinitely many possible ODE functional forms that are self-consistent with the wolf and rabbit population trajectories. Furthermore, this ambiguity might make it too hard to successfully learn any model at all without the helpful constraint of the withheld bear population evolution in training data.

The ODE neural net has three inputs. Two of these are the wolf and rabbit population. Intuitively, we would like to think of the third one as "bears". But as we shall see it does not turn out to be the bear population. It’s just an additional recurrent external channel. There are three outputs which provide the rate of change for the Rabbit and Wolf and recurrent external channel to the integrator.

This is almost the same as the previous test, except the external channel is not initialized to the bear population at t=0. We devised another approach to initializing it. Simply supplying a constant, like zero, would not be likely to generalize to a correct result on testing data. This can be seen by imagining two test data sets generated starting with the same initial number of rabbits and wolves but different numbers of bears. The different number of bear predators means the trajectories must be different, but every trained ANODE will have a single deterministic output for the same three inputs (wolves and rabbits and zero). Thus it could not match both data sets. This is a special case of the general rule that ODEs cannot produce crossing trajectories.

Thus more than just the ANODE architecture is needed to resolve this. We implemented a separate neural network outside of the neural ODE which receives not just the t=0 wolf and rabbit populations but the first ten time points of the wolf and rabbit population. It outputs a single scalar to initiate (t=0) the recurrent external channel. Thus the initial conditions delivered to the neural ODE itself are the ordinary t=0 values of the rabbit and wolf, as well as this derived initial condition for the recurrent channel.
Note that both of the initial-condition-predictor and the neural net inside the ODE must be trained simultaneously. We are not separately training the external net to predict some known bear population; we would not have that initial bear information since we are pretending we don’t know there are bears at all. This means we also back-propagate through both neural nets including through all iterations of the solver and neural ODE to further back-propagate through the initial condition determining network.

Through training, the prediction from integrating the neural ODE goes from uniformed and inaccurate to matching the true solution of the Lotka-Volterra equations. The gradual progression of training to match the Lotka-Volterra equations is shown in figure 1(d). To test if the neural network actually learned to generalize rather than memorize, the model is applied to an initial condition it hasn’t trained on and, even more stringently, tested over a longer time span, so it must extrapolate using the $F(.)$ it has learned. The plots in figure 2(c) are the results from these tests in each experiment.

To test what it learned, the neural ODE and Lotka-Volterra equations are solved with an initial condition not seen in the training set and over a longer period of time than in training. Figure 2(a) shows the prediction from the neural ODE matches up with the true solution to the Lotka-Volterra equations in these testing conditions very well. This shows the neural net was successfully trained through back-propagation through the ODE solver. Given it extrapolates to the held-out testing data establishes the neural network learned the physics of the system rather than memorizing the test data. This demonstrates neural ODEs are able to learn the physics of unknown systems if no data is withheld.

Again, to test what it learned, the neural ODE and Lotka-Volterra equations are solved with an initial condition not seen in the training set and over a longer period of time than in training. Figure 2(b) shows the prediction from integrating the neural ODE matches up with the true solution to the Lotka-Volterra equations in only two of the three species. This is confusing at first. How can the neural ODE model the rabbits and wolves without knowing the bear population?

We were initially surprised that the bear population did not match the recurrent external state variable. It seemed like that since we gave it the true bear population as the initial condition that the simplest outcome would be to have reproduced the original ODE. But it did not. And it didn’t have to because the loss function did not impose this. So it doesn’t need to predict the bear population, the recurrent channel only needs to transmit and update some information that aids it in predicting the other two populations. However, it is hard to imagine a general class of transformations that also would be self-consistent with every initial value matching the bear population at t=0. Thus we expected it would be biased to find the true bear population, and we were surprised when it found some alternative solution.

Figure 2(c) shows the prediction from integrating the neural ODE matches up with the true solution to the Lotka-Volterra equations in two of the three species, and has different behavior in the third. Again, only the rabbit and wolf populations match, and they match very well. Notice that, unlike the previous case, it now does not match the initial value of the bears. It does seem to have oscillatory dynamics like the true bear population but magnitudes and phases evolve differently over time.

Our new neural network architecture was designed specifically to address the central problem arising from incomplete state specifications. It preserves the core property of all stationary ODEs: trajectories don’t cross. The crossing of trajectories does not occur in the fully specified ODE of known systems simply because the next integration step along the time derivative $\frac{\partial\hat{x}(t)}{\partial t}$ depends only on the current position $F[\hat{x}(t)]$. Two ramifications occur for unknown systems. First, if one is unaware of an additional member variable in an N-dimensional $\hat{x}$, then the reduced state space of N-1 dimensions can have a position that maps to many states in the N-dimensional space with different gradients, and so paths cross. Second, even if one knows the problem is N-dimensional if the initial condition isn’t specified for all the elements of $\hat{x}$, the same problem arises.

The first problem is solved by augmenting an N-1 dimensional ODE with an additional idler variable meant to make this mapping unique again. We may not know what this variable is, but we just add it on. This need not be done blindly. We need only need to add as many variables as required to remove trajectory crossings in the training data.

Secondly, how do we initialize these new variables or any other variables for which we lack that initial data? We address this by training a second neural net which takes in a time series of the variables whose trajectory we do know, and outputs a unique initial condition. Thus the trajectory ambiguity of the unknown initial condition is removed and again paths don’t cross. We note that this degeneracy could be broken by any arbitrary mapping of known data to a unique initial condition (e.g. a hash function) but such would be unlikely to have nice continuity properties and thus requires a deeper neural net to handle the arbitrary discontinuity. Therefore, we used a neural net to more smoothly estimate the initial condition.

We implemented the augmented neural ODEs and differential equation solvers in the Julia programming language [6] within an existing neural ODE framework[7]. This was all completed by writing code in the Julia programming language version 1.0. The truth ODE function for the Lotka-Volterra equations was written and solved in each case for truth data. Neural networks were created with the Flux package [16]. Differential equations were solved with the DifferentialEquations package [17], and back-propagation through a neural ODE used the DiffEqFlux package [4].
The neural ODE network structure was intentionally kept the same as ones used for solving the canonical Lotka-Volterra problem with 3 dense hidden layers, 32 nodes wide, using swish activations. These had 2 (case 1) or 3 (cases 2&3) inputs/outputs. The structure of our initial condition neural net was 20 inputs (10 each for r and w), one output (the initial condition for the external channel), with 3 hidden dense layers 32 nodes wide with swish activation. Training used the Flux.jl package [16] with ADAM optimization. The code and documentation written for this project can be found online at https://github.com/robertstrauss/nnDiffEq.