On Second Order Behaviour in Augmented Neural ODEs: A Short Summary
Abstract
In Norcliffe et al. [13], we discussed and systematically analysed how Neural ODEs (NODEs) can learn higher-order order dynamics.
In particular, we focused on second-order dynamic behaviour and analysed Augmented NODEs (ANODEs), showing that they can learn second-order dynamics with only a few augmented dimensions, but are unable to correctly model the velocity (ﬁrst derivative).
In response, we proposed Second Order NODEs (SONODEs), that build on top of
ANODEs, but explicitly take into account the second-order physics-based inductive biases.
These biases, besides making them more efﬁcient and noise-robust when modelling second-order dynamics, make them more interpretable than ANODEs, therefore more suitable in many real-world scientiﬁc modelling applications.
Introduction
From a "simple" swinging pendulum, to complex environmental/climate processes, to celestial-body motion – dynamical systems are fundamental means of studying complex phenomena in nature and society.
Ordinary Differential Equations (ODEs) are a common formalism for describing dynamical behaviour, which express the rate of change of systems’ states (typically) through time.
Neural
Ordinary Differential Equations (NODEs) [2] refer to continuous-depth neural networks that can represent a structure of an ODE, and as such allow for learning and modelling unknown dynamics from (partially-)observed data, which often cannot be described analytically.
Neural ODEs (NODEs) can be seen as a continuous variant of ResNet models [7], whose hidden state evolves continuously according to a differential equation
[Equation]
whose velocity is described by a neural network f ( v ) with parameters θ f and initial position given by the points of a dataset X 0 .
As shown by Chen et al.
[2], the gradients can be computed through an abstract adjoint state r ( t ) , once its dynamics are known.
Since the introduction of NODEs, there have been many studies which address different challenges and classes of modelling problems [4, 8, 11, 15, 16, 19, 22].
Our work in Norcliffe et al. [13] presented a clear in-depth study on second-order behaviour in Neural ODEs, that can be used for modelling many dynamical systems governed by second-order laws, such as oscillators, Newton’s equations etc.
In particular, we introduced Second Order Neural ODEs (SONODEs), an instance of Augmented Neural ODEs (ANODEs) [4] that explicitly takes into account physics-based inductive biases.
We studied the optimisation of SONODEs by generalising the adjoint sensitivity method to second-order models.
We showed that some useful properties of ANODEs also extend to SONODEs, which can also lead to simpler and more intuitive solutions to several benchmark problems.
Finally, through a series of experiments on synthetic and real-world dynamical systems, often with known analytic solutions, we illustrated the beneﬁts of SONODEs that relate to efﬁciency, noise-robustness and interpretability when modelling second-order dynamics.
This paper highlights and brieﬂy discusses these ﬁndings.
The code is available online at https: //github.com/a-norcliffe/sonode .
Second Order Neural Ordinary Differential Equations
Second Order Neural ODEs (SONODEs) [11, 22], with an initial position x ( t 0 ) , initial velocity ˙ x ( t 0 ) , and acceleration ¨ x are given by
[Equation]
where f ( a ) is a neural network with parameters θ f .
Alternatively, SONODEs can be seen as a system of coupled ﬁrst-order neural ODEs with state z ( t ) =
[ x ( t ) , a ( t )] :
[Equation]
This formulation shows that SONODEs are a type of ANODE [4], which append states a ( t ) to the
ODE:
[Equation]
and as such, ANODEs inherit many useful properties, such as greater function representation than
NODEs and being able to learn second order dynamics with fewer dimensions than SONODEs (see
Section 4 in [13]). In contrast to ANODEs, SONODEs constrain the structure of f ( v ) , and offer a way to adapt NODE’s ﬁrst-order adjoint method for training [2] to a second order adjoint method.
Proposition 2.1.
The adjoint state r ( t ) of SONODEs follows the second-order ODE
[Equation]
The proof and boundary conditions are given in Appendix B. Since the dynamics of the abstract adjoint vector are known, its state at all times t can be used to train the parameters θ f using the integral
[Equation]
where L denotes the loss function and t n is the timestamp of interest.
The gradient with respect to the parameters of the initial velocity network, θ g , can be found in Appendix B. To further evaluate how the second-order adjoint sensitivity method compares with the ﬁrst-order adjoint-based optimisation, we compare this gradient to the one obtained through the adjoint of the ﬁrst-order coupled ODE from
Equation (3).
Proposition 2.2.
The gradient of θ f computed through the adjoint of the coupled ODE from (3) and the gradient from (6) are equivalent.
However, the latter requires at least as many matrix multiplications as the former.
This result motivates the use of the ﬁrst-order coupled ODE, as it presents computational advantages.
The proof in Appendix B shows that this is due to the dynamics of the adjoint from the coupled
ODE, which contain entangled representations of the adjoint.
This is in contrast to the disentangled representation in Equation (5), where the adjoint state and velocity are separated.
It is the entangled representation, a typical phenomenon in ANODEs, that permits fast computation of the gradients for the coupled ODE. The entangled representation also introduces further properties when ANODEs are used as a modelling choice for second-order dynamical systems as described below.
Even though SONODEs build on ANODEs, they learn second-order dynamics in different ways (see Section 5.1. in [13] or Appendix E).
ANODEs learn an abstract function F that at t 0 is equal to the [Figure] initial velocity, and another function G that couples to F giving it the right acceleration.
ANODEs learn an entangled representation of the dynamics in this way, allowing them to model dynamics with fewer dimensions (see Section 4 in [13]).
In contrast, SONODEs are constrained to learn the acceleration and initial velocity directly.
This leads to several useful properties, for instance when investigating the dynamics of physical systems.
In the case of SONODEs, modelling this behaviour is straightforward, directly learning the acceleration as a function of position, velocity and time.
However, in the case of ANODEs, this is learnt through an abstract alternative ODE where the state and augmented dimensions are entangled, impairing their interpretability.
To illustrate this, in Figure 1 we show the trajectories of an ANODE(2) 2 and a SONODE, trained to model a two-dimensional second-order system (see Equation (13) in [13]), using the correct initial velocity.
Even though ANODE(2) is able to learn the true trajectory in real space, the augmented trajectories change each run and are substantially different from the true velocity of the underlying ODE.
In contrast, SONODE learns the correct velocity for both runs.
This clearly shows that ANODEs might not be a suitable investigative tool for scientiﬁc applications, where the physical interpretability of the results is important.
SONODEs also leverage their inductive biases to converge faster than the other models.
Figure 2a shows the training loss of NODE, ANODE(1) and SONODE on a task of modelling a damped harmonic oscillator 3 .
While NODEs and ANODEs learnt a general ˙ z , SONODEs were given ˙
z
=
[ v, f ( a ) ] and only learnt f ( a ) , therefore converging in fewer iterations.
Moreover, SONODEs are more robust to noise.
Figure 2b illustrates the ability of ANODE(1) and SONODE to learn a sine curve ( x = sin( t ) ) in varying noise regimes.
SONODEs extrapolate better than ANODEs, because they are "forced" to learn second-order dynamics, and therefore are less likely to overﬁt to the training data.
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
(a) NODE, ANODE(1) and SONODE train- ing error on harmonic oscillators.
SONODEs, with built-in second-order behaviour in the architectural choice, converge faster.
[Figure]
SONODEs are more stable than ANODEs when learning a sine curve in different noise regimes.
1000 Iterations 0 .0 .0 .0 .0 .1 .1 .
M S E
Silverbox Training MSE
SONODEANODE(1)
5000 t −−−
V
Silverbox V 2 True V
SONODEANODE(1)    4000 t 0 .0 .0 .
R unn i n g
A v e r ag e R M S
E
Silverbox Running Error
SONODEANODE(1)
Finally, we compare SONODEs and ANODEs on a real-world modeling task.
In particular, we consider modelling the Silverbox system [18], an electronic circuit resembling a Dufﬁng Oscillator, with input voltage V 1 ( t ) and measured output V 2 ( t ) . The non-linear model Silverbox represents is ¨ V 2 = a ˙ V 2 + bV 2 + cV 32 + dV 1 .
To account for this, all models included a V 32 term.
On this second-order system, shown in Figure 3, SONODEs extrapolate better than ANODEs and are able to capture the increase in the amplitude of the signal exceptionally well.
Other comparisons on real-world datasets are given in Appendix D.
Discussion
SONODEs are a special case of ANODEs, whose phase-space dynamics are restricted to model second-order behaviour.
We believe that for tasks where the trajectory is unimportant, and performance depends only on the endpoints (such as classiﬁcation), ANODEs might perform better because they are unconstrained in how they use their capacity (see Appendix D.5).
In contrast, we expect
SONODEs to outperform ANODEs both in terms of accuracy and convergence rate on time-series data whose underlying dynamics is assumed (or known) to be second-order.
In this setting, SONODEs have a unique functional solution and fewer local minima compared to ANODEs.
At the same time, they have higher parameter efﬁciency since ˙ x = v requires no parameters, so all parameters are in the acceleration.
Finally, we expect SONODEs to be more appropriate for applications in the natural sciences, where second-order dynamics are common and it is useful to recover the force equation.
Concurrent to our work [13], SONODEs have also been brieﬂy evaluated on MNIST by Massaroli et al.
[11] as part of a wider study on Neural ODEs.
In contrast, our study is focused on the theoretical understanding of second-order behaviour.
At the same time, our investigations are largely based on learning the dynamics of physical systems rather than classiﬁcation tasks.
Second-order models have also been considered in Graph Differential Equations [15] and ODE 2 VAE [22].
In the same way SONODEs assert Newtonian mechanics, other models have been made to use physical laws, guaranteeing physically plausible results, in discrete and continuous cases.
Lutter et al.
[10] apply
Lagrangian mechanics to cyber-physical systems, while Greydanus et al.
[6] and Zhong et al.
[20] use Hamiltonian mechanics to learn dynamical data.
More recently, further priors have been used on top of second order, to improve performance.
Modular Neural ODEs [21] split up the acceleration into a sum of physically meaningful terms, ¨ x = −∇ φ ( x )+ v × B ( x ) − D ( v ) v , making up a scalar potential, magnetic ﬁeld and drag force.
This separation was seen to greatly improve performance, however it requires the form of the force to be known.
More generally, one can use Privileged Information [17] to improve the learning of dynamics.
LUPI Neural ODE Processes (learning using privileged information) [3], extend the Neural ODE
Process model (NDP) [14], so that during training higher level knowledge is supplied to the model, such as conserved quantities.
Then at test time the NDP is able to infer the privileged information from observations only.
Conclusion
In our work we showed that Second Order Neural ODEs (SONODEs) are a special case of Augmented
Neural ODEs (ANODEs) with restricted dynamics.
We showed that ANODEs can learn second order dynamics using fewer dimensions than SONODEs, making them more efﬁcient.
However, their dimensions are entangled, making their solutions less interpretable than those learnt by SONODEs.
Finally, we demonstrated with synthetic and real-world datasets, that when the dynamics are known to be second order, SONODEs are both more robust and more accurate than ANODEs
A Phase Space Trajectory Proofs
Here we consider a k -th order initial value problem and its phase space properties.
Lemma A.1.
For a k-th order IVP, where the k-th derivative is Lipschitz continuous, a solution cannot have discontinuities in the time derivative of its phase space trajectory.
Proof.
Consider the phase space trajectory z ( t ) =
[Equation]
[Equation] [Equation]
k
[Equation]
1 x dt
[Equation]
Let f be the k-th time derivative of x ( t ) .
Then the time derivative of z ( t ) is
[Equation]

If for one set of ﬁnite arguments, z 1 , f ( z 1 ) is also ﬁnite, then because the gradients of f are all bounded (due to Lipschitz continuity), for any other ﬁnite arguments, z n , f ( z n ) will remain ﬁnite.
Now consider d k − 1 x dt k − 1 , its time derivative is f ( z ( t )) , which is ﬁnite for all ﬁnite z .
Therefore, d k − 1 x dt k − 1 , can’t have discontinuities with a ﬁnite derivative, and also must be ﬁnite for ﬁnite z .
Now consider d k − 2 x dt k − 2 , its time derivative is ﬁnite for all ﬁnite z , and therefore it can’t have discontinuities and also must be ﬁnite for all ﬁnite z .
This line of argument continues up to x .
The state x and all of its time derivatives up to the k -th have no discontinuities and are ﬁnite.
Therefore as long as the initial conditions z ( t 0 ) are ﬁnite, there can be no discontinuities in the time derivative of the phase space trajectory at ﬁnite time.
Proposition A.2.
For a k-th order IVP, if the k-th derivative of x is Lipschitz continuous and has no explicit time dependence, then unique phase space trajectories cannot intersect at an angle.
Similarly, a single phase space trajectory cannot intersect itself at an angle.
Proof.
Consider two trajectories z 1 ( t ) and z 2 ( t ) that have different initial conditions z 1 ( t 0 )
=
h 1 and z 2 ( t 0 )
=
h 2 .
Assume the trajectories cross at a point in phase space at an angle, z 1 ( t 1 ) =
z 2 ( t 2 ) =
˜h .
If they intersect at an angle, then evolving the two states by a small time δt << 1 , and using the
Lipschitz continuity of f , meaning that the trajectories cannot have kinks in them (as shown in
Lemma A.1), z 1 ( t 1 + δt ) (cid:54) =
z 2 ( t 2 + δt ) .
However, if they are at the same point in phase space, then they must have the same k-th order derivative, f .
All other derivatives are equal, so by evolving the states by the same small time δt << 1 , z 1 ( t 1 + δt ) =
z 2 ( t 2 + δt ) .
There is a contradiction and therefore the assumption is wrong, unique trajectories cannot cross at an angle in phase space when f is Lipschitz continuous and has no t dependence.
Now consider the single trajectory z ( t ) .
Assume it intersects itself at an angle, at t 1 and t 2 .
Now consider two particles on this trajectory, starting at t 1 − τ and t 2 − τ such that t 2 − τ > t 1 .
These two particles have different initial conditions and cross at an angle.
However, the above shows that cannot happen.
Therefore, the assumption that z ( t ) can intersect itself at an angle must be wrong.
Trajectories cannot intersect themselves in phase space at an angle.
Trajectories can, however, feed into each other representing the same particle path at different times.
Single phase space trajectories can feed into themselves representing periodic motion.
This requires a
Lipschitz continuous f , and for there to be no explicit time dependence.
If there was time dependence then two trajectories can cross at different times, and a trajectory can self intersect.
Effectively an additional dimension is added to phase space, which is time.
The propositions above would still hold because dt dt = 1 which is Lipschitz continuous.
Therefore, with time included as a phase space dimension, intersections in space are only forbidden if they occur at the same time.
B Adjoint Sensitivity Method
We present a proof to both the ﬁrst and second order Adjoint method, using a Lagrangian style approach [1, 5].
We also prove that when the underlying ODE is second order, using the ﬁrst order method on a concatenated state, z
=
[ x , v ] , produces the same results as the second order method but does so more efﬁciently.
All parameters, θ , are time-independent (so dθdt =
dtdθ = 0 ).
B.1 First Order Adjoint Method
Let L denote a scalar loss function, L = L ( x ( t n )) , the gradient with respect to a parameter θ is
[Equation]
dθ (7)
The vector ∂L ∂ x ( t n ) T is found using backpropagation.
For dynamical data the loss will depend on multiple time stamps, there is also a sum over timestamps, t n .
Therefore d x ( t n ) dθ is needed.
x ( t n ) follows
[Equation] [Equation]
) subject to
[Equation]
where X 0 is the data going into the network and is constant.
The functions f ( v ) and s describe the
ODE and the initial conditions.
Here we allow X 0 to ﬁrst go through the transformation, s ( X 0 , θ s ) .
This maintains generality and allows NODEs to be used as a component of a larger model.
For example, X 0 could go through a ResNet before the NODE, and then through a softmax classiﬁer at the end (which is accounted for in the term ∂L ∂ x
[Equation]
t n
[Equation]
Introduce the new variable F
[Equation]
[Equation]
These are equivalent because ( ˙ x − f ( v ) ) and ( x ( t 0 ) − s ) are both zero.
This means the matrices, A ( t ) and B , can be chosen freely (as long as they are well behaved, ﬁnite etc.), to make the computation easier.
The gradients of x ( t n ) with respect to the parameters are
[Equation]
Differentiating F with respect to a general parameter θ d F
[Equation]
[Equation]
Integrating by parts
[Equation] [Equation]
Substituting this in and using (cid:82) t n t 0 d ˙ x dθ dt =
[ d x dθ ] t n t 0 , gives
[Equation]
Using the freedom of choice of A ( t )
[Equation]
let it follow the ODE
[Equation] [Equation]
Where I is the identity matrix.
Then the ﬁrst term and second integral in Equation (14) become zero, yielding
[Equation]
Now using the freedom of choice of B
[Equation]
let it obey the equation
[Equation]
[Equation]
This makes the ﬁrst term in Equation (16) zero.
This gives the ﬁnal form of d F dθ
[Equation]
Subbing into Equation (11) and using the fact that f ( v ) has no θ s dependence and s has no θ f dependence
[Equation]
This leads to the gradients of the loss dL
[Equation]
Subject to the ODE for A ( t )
[Equation]
Now introduce the adjoint state r ( t )
[Equation]
Using the fact that ∂L ∂ x ( t n ) is constant with respect to time, the adjoint equations are obtained by applying the deﬁnition of the adjoint in Equation (22), to the gradients in Equation (20), and multiplying the ODE in Equation (21) by the constant − ∂L ∂ x ( t n )
[Equation]
Where the adjoint a ( t ) follows the ODE
[Equation]
The gradients are found by integrating the adjoint state, r , and the real state, x , backwards in time, which requires no intermediate values to be stored, using constant memory, a major beneﬁt over traditional backpropagation.
These are the same equations that were derived by Chen et al.
[2], however this includes the addition of letting x ( t 0 )
=
s ( X 0 , θ s ) giving the corresponding gradient, dL dθ s .
Additionally, the derivation used by Chen et al.
[2] is simpler but does not present an obvious way to extend the adjoint method to second order ODEs, which this derivation method can do, as shown next.
B.2 Second Order Adjoint
Using the same derivation method, but with a second order differential equation, a second order adjoint method is derived, according to the proposition from the main text:
Proposition 2.1.
The adjoint state r ( t ) of SONODEs follows the second order ODE
[Equation]
and the gradients of the loss with respect to the parameters of the acceleration, θ f are
[Equation] [Equation]
Proof.
In general, the loss function, L , depends on x and ˙ x
[Equation]
The gradients from the positional part and the velocity part are found separately and added.
Firstly the position
[Equation]
Subject to the second order ODE
[Equation]
Following the same procedure as in ﬁrst order, but including the initial condition for the velocity as well
[Equation] [Equation]
As before, the vectors, (¨ x − f ( a ) ) , ( ˙ x ( t 0 ) − g ) and ( x ( t 0 ) − s ) are zero, which gives freedom to choose the matrices A ( t ) , B and C to make the calculation easier.
The gradients of x ( t n ) with respect to the parameters θ are
[Equation]
Differentiating F from equation 30 with respect to a general parameter
[Equation]
(cid:90) t n
[Equation] [Equation]
(33) (cid:90) t n t
A ( t ) ∂f ( a ) ∂ ˙ x T d ˙ x dθ dt =
(cid:20)
A ( t ) ∂f ( a ) ∂ ˙ x T d x dθ (cid:21) t n t 0 − (cid:90) t n td dt (cid:18)
A ( t ) ∂f ( a ) ∂ ˙ x T (cid:19)
d x dθ dt (34)
Subbing these into Equation (32)
[Equation]
Using the freedom to choose A ( t )
[Equation]
let it follow the second order ODE
[Equation]
This makes the ﬁrst term and ﬁrst integral in Equation (35) zero
[Equation]
yielding
[Equation]
Now using the freedom of choice in B and C
[Equation]
This makes the second and third terms in Equation (37) zero, yielding
[Equation]
These give the ﬁnal gradients of x ( t n ) with respect to the parameters, by subbing the results for B ,
C and d F dθ above into Equation (31)
[Equation]
using the fact that f ( a ) , g and s only depend on the parameters θ f , θ g
and θ s respectively
[Equation]
As before, introduce the adjoint state r x ( t ) :
[Equation]
Using the fact that ∂L ∂ x ( t n ) is constant with respect to time, all the results above, and the ODE and initial conditions for A ( t ) in Equation (36) can be multiplied by − ∂L
[Equation]
) T , to get the gradients dL dθ in terms of r x ( t ) dL
[Equation]
[Equation]
Subject to the second order ODE for r x ( t )
[Equation]
Where after differentiating with the product rule the ODE in Equation (43) becomes
[Equation]
Where doing the full time derivative gives
[Equation]
Where the fact that ¨ x = f ( a ) has been used.
This is only when the loss depends on the position.
The same method is used to look at the velocity part in Equation (27) dL
[Equation] [Equation]
Where
[Equation]
The general method is to take this expression and add zeros, in the form of A ( t ) , B and C multiplied by the ODE and initial conditions, (¨ x − f ( a ) ) , ( ˙ x ( t 0 ) − g ) and ( x ( t 0 ) − s ) .
Then differentiate with respect to a general parameter θ and integrate by parts to get any integrals containing d ˙ x dθ or d ¨ x dθ in terms of d x dθ .
Then choose the ODE for A ( t ) to remove any d x dθ terms in the integral, and the initial conditions of A ( t n ) to remove the boundary terms at t n .
Then B and C are chosen to remove the boundary terms at t 0 .
After doing this the gradients of ˙ x with respect to the parameters are
[Equation]
Subject to the second order ODE for A ( t )
[Equation]
Now introduce the state r v ( t )
[Equation]
Which allows the gradients of the loss with respect to the parameters to be written as dL
[Equation]
dL
[Equation]
Where r v follows the second order ODE and initial conditions
[Equation]
Now adding the gradients from the x dependence and the ˙ x dependence together.
It can be seen that the gradients are the same in Equations (42) and (51), but just swapping r x and r v .
Additionally, it can be seen from the ODEs for r x and r v in Equations (44) and (52), that they are governed by the same, linear, second order ODE, with different initial conditions.
Therefore the gradients, dL dθ , can be written in terms of a new adjoint state, r = r x + r v
[Equation]
Where a follows the second order ODE with initial conditions
¨
[Equation]
The full derivative, d t ( ∂ ˙ x f ( a ) ) , is given by Equation (45).
The ODE can also be written compactly as
[Equation]
Just as in the ﬁrst order method, a sum over times stamps t n may be required.
This matches and extends on the gradients and ODE given by proposition 2.1.
B.3 Equivalence between the two Adjoint methods
When acting on a concatenated state, z ( t ) =
[ x ( t ) , v ( t )] , the ﬁrst order adjoint method will produce the same gradients as the second order adjoint method.
However, it is more computationally efﬁcient to use the ﬁrst order method.
This is also given in the main text as the following proposition:
Proposition 2.2.
The gradient of θ f computed through the adjoint of the coupled ODE from (3) and the gradient from (6) are equivalent.
However, the latter requires at least as many matrix multiplications as the former.
Intuitively, the ﬁrst order method will produce the same gradients because second order dynamics can be thought of as two coupled ﬁrst order ODEs, where the ﬁrst order dynamics happen in phase space.
However, this provides no information about computational efﬁciency.
We prove the equivalence and compare the computational efﬁciencies below.
Proof.
The ﬁrst order formulation of second order dynamics can be written as
[Equation]
When using index notation, x i and v i are concatenated to make z i. .
For x i and v i , the index, i, ranges from 1 to d , whereas for z
i it ranges from 1 to 2 d accounting for the concatenation.
This is represented below
[Equation]
It also extends to ˙
z i and z i ( t 0 ) , where f ( a ) i , s i and g
i also have the index range from 1 to d , but
the index of ˙ z i goes from 1 to 2 d just like for z i .
[Equation]
Using the ﬁrst order adjoint method, Equations (23) and (24)
[Equation]
and using index notation with repeated indices summed over, the gradients are
[Equation]
Where the adjoint follows the ODE
[Equation]
Where just like in z i , the index, i, ranges from 1 to 2 d in the adjoint
r i ( t ) .
When writing the sum over the index explicitly
[Equation]
Now split up the adjoint state, r , into two equally sized vectors, r A and r B , where their indices only range from 1 to d , like x , v , f ( a ) , g and s .
[Equation]
Using Equations (57), (58), (59) and (63), and subbing them into Equation (62), the derivative can be written as
[Equation]
Relabelling the indices in the second sum ( j − d ) −→
[Equation] [Equation]
Looking at speciﬁc values of i:
[Equation]
d
[Equation]
i
[Equation]
(cid:88)
[Equation]
Relabelling the ﬁrst index (
[Equation]
[Equation]
Noting that, ∂v j ∂v i = δ ij , the time derivatives can be written in vector matrix notation as
[Equation]
[Equation]
Differentiating Equation (70), and using Equation (69) for ˙
r A ( t )
[Equation]
This matches the ODE for the second order method in Equation (55).
Now applying the initial conditions, using index notation again r i ( t n
[Equation]
∂L
[Equation]
For i ≤ d
[Equation]
For i ≥ d +
[Equation]
Applying these initial conditions in r A and r B to Equation (70)
[Equation]
By looking at the ODE and initial conditions, it is clear r B is equivalent to the second order adjoint, in Equation (54).
Now looking at the gradients, and including an explicit sum over the index dL
[Equation]
The ﬁrst term is zero because v has no explicit θ dependence.
The second term, after relabelling and using summation convention becomes
[Equation]
Where ˜ θ
f
= θ f has been used, as they are both the parameters for the acceleration.
This matches the result for gradients of parameters in the acceleration term θ f , when using the second order adjoint method, because r B is the adjoint.
Looking at the gradients related to the initial conditions
[Equation]
[Equation]
After going through the previous process of separating out the sums from 1 −→ d and d + 1 −→ 2 d , then relabelling the indices on r B , this becomes
[Equation]
Using the expression for r A by rearranging Equation (70), this can be written as dL
[Equation]
The parameters ˜ θ s contain both θ s and θ g .
Looking at θ g ﬁrst, where s ( X 0 , θ s ) has no dependence dL
[Equation]
where dg dθ g can be written as a partial derivative, because X 0 and θ s have no dependence on θ g .
This expression is equivalent to dL dθ g found using the second order adjoint method.
Now looking at the parameters θ s , these parameters are in s ( X 0 , θ s ) explicitly and g ( s, θ g ) , implicitly through s .
Subbing ˜ θ s = θ
s into Equation (80) gives dL
[Equation]
Using the fact
that x ( t 0 ) =
s , this is the same result for dL dθ s found using the second order adjoint method: dL
[Equation]
All of the gradients match, so the ﬁrst order adjoint method acting on z ( t ) =
[ x ( t ) , v ( t )] will produce the same gradients as the second order adjoint method acting on x ( t ) .
Given by Equation (53).
Looking at the efﬁciencies of each method and how they would be implemented.
Both methods would integrate the state z
= [ x , v ] forward in time, with ˙ z
= [ v , f ( a ) ] .
Both methods then integrate z and the adjoint backwards, in the same way.
The difference is how the adjoint is represented.
In ﬁrst order it is represented as [ r A , r B ] where r B is the adjoint, in second order it is represented as [ r , ˙ r ] where r is the adjoint.
The time derivatives and initial conditions for the ﬁrst order adjoint representation are
[Equation]
The time derivatives and intial conditions for the second order adjoint representation are
[Equation]
Where
[Equation]
Looking at Equations (84) and (85), the second order method has the additional term, r · d t ( ∂ v ( f ( a ) )) , in the ODE, and the additional term, ( ∂ v L ) · ( ∂ v f ( a ) ) in the initial conditions.
The ﬁrst order method acting on the concatenated state, [ x , v ] , requires equal or fewer matrix multiplications than the second order method acting on x , to ﬁnd the gradients at each step and the initial conditions.
This is in the general case, but also for all speciﬁc cases, it is as efﬁcient or more efﬁcient.
The same is also true for calculating the ﬁnal gradients.
The reason for the difference in efﬁciencies is the state, r B , is the adjoint, and the state, r A , contains a lot of the complex information about the adjoint.
It is an entangled representation of the adjoint, contrasting with the disentangled second order representation [ r , ˙ r ] .
This is similar to how ANODEs can learn an entangled representation of second order ODEs and SONODEs learn the disentangled representation, seen in Section
E.
However, entangled representations are more useful here, because they do not need to be interpretable, they just need to produce the gradients, and the entangled representation can do this more efﬁciently.
This analysis provides useful information on the inner workings of the adjoint method.
It shows a second order speciﬁc method does exist, but the ﬁrst order method acting on a state z =
[ x , v ] will produce the same gradients more efﬁciently, due to how it represents the complexity.
This was speciﬁc to second order ODEs, however, the ﬁrst order adjoint will work on any system of ODEs, because any motion can be thought of as being ﬁrst order motion in phase space.
Additionally, the ﬁrst order method may be the most efﬁcient adjoint method.
The complexity going from the ﬁrst order to the second order was seen based on the calculation, so this is only likely to get worse as the system of ODEs becomes more complicated.
C Properties of SONODEs
In this section, we analyse certain properties of SONODEs and illustrate them with toy examples.
C.1 Generalised parity problem
It is known that unique trajectories in NODEs cannot cross at the same time [4, 11].
We extend this to higher order Initial Value Problems
(IVP).
Proofs are presented in Appendix A.
Proposition C.1.
For a k-th order IVP, if the k-th derivative of x is Lipschitz continuous and has no explicit time dependence, then unique phase space trajectories cannot intersect at an angle.
Similarly, a single phase space trajectory cannot intersect itself at an angle.
While this shows SONODE trajectories cannot cross in phase space, they can cross in real space if they have different velocities.
To illustrate this, we introduce a generalised parity problem, an extension to D dimensions of the g 1 d function from Dupont et al.
[4], which maps x → − x .
We remark that SONODEs should be able to learn a parity ﬂip in any number of dimensions, with a trivial solution
[Equation]
This is equivalent to all points moving in straight lines through the origin to − x ( t 0 ) .
10 Dimension −−−−−−− l n ( L )
[Figure]
For higher dimensions, we ﬁrst remark that NODEs are able to produce parity ﬂips for even dimensions by pairing off the dimensions and performing a 180 ◦ rotation in each pair.
This solution does not apply to odd-dimensional cases because there is always an unpaired dimension that is not rotated.
In addition to the dimensional-parity effect, as volume increases exponentially with the dimensionality, the density exponentially decreases (given the number of points in the dataset remains constant).
This makes it easier to manipulate the points without trajectories crossing, and so, it is expected that the problem will become easier for NODEs as dimensionality increases.
In Figure 4, we investigate parity ﬂips in higher dimensions, using 50 training points and 10 test points, randomly generated between -1 and 1 in each dimension.
For
NODEs, as predicted, the loss oscillates over dimensions and, for odd dimensions, the loss decreases with the number of dimensions.
ANODEs perform better than NODES, especially in odd dimensions, where it can rotate the unpaired dimension through the additional space.
SONODEs have the lowest loss in every generalisation, which can be associated with the existence of the trivial solution in any number of dimensions, given by Equation (87).
C.2 Nested n-spheres
Dupont et al.
[4] prove that a transformation under NODEs has to be a homeomorphism, preserving the topology of the input space, and as such, they cannot learn certain transformations.
Similarly to
ANODEs, SONODEs avoid this problem.
Proposition C.2.
SONODEs are not restricted to homeomorphic transformations in real space.
The proof can be found in Appendix C.3.
To illustrate this, we perform an experiment on the nested n-spheres problem [4], (the name is taken from [11], originally called g function [4]), where the elements of the blue class are surrounded by the elements of the red class (Figure 5) such that a homeomorphic transformation in that space cannot linearly separate the two classes.
As expected, only ANODEs and SONODEs can learn a mapping.
[Figure]
C.3 Second Order ODEs are not Homeomorphisms
One of the conditions for a transformation to be a homeomorphism is for the transformation to be bijective (one-to-one and onto).
In real space, a transformation that evolves according to a second order ODE does not have to be one-to-one.
This is demonstrated using a one-dimensional counter-example
[Equation]
If t 0 = 0 and t N =
[Equation]
So the transformation in real space is not always one-to-one, and therefore, not always a homeomorphism.
D Experimental Setup and Additional Results
We anticipate two main uses for SONODEs.
One is using an experiment in a controlled environment, where the aim is to ﬁnd values such as the coefﬁcient of friction.
The other use is when data is observed, and the aim is to extrapolate in time, but the experiment is not controlled, for example, observing weather.
We would expect for the former, a simple model with only a single linear layer would be useful, to ﬁnd those coefﬁcients, and for the latter, a deeper model may be more appropriate.
Additionally, Neural ODEs may be used in classiﬁcation or other tasks that only involve the start and endpoints of the ﬂow.
For all of these tasks we used t 0 = 0 and t 1 = 1 , and accelerations that were not time-dependent.
For tasks depending on the start and endpoint only, a deeper neural network is more useful for the acceleration.
For all experiments, except the MNIST experiment, we optimise using Adam with a learning rate of 0.01.
We also train on the complete datasets and do not minibatch.
All the experiments were repeated 3 times to obtain a mean and standard deviation.
Depending on the task at hand, we used two different architectures for NODEs, ANODEs and SONODEs.
The ﬁrst is a simple linear model, one weight matrix and bias without activations.
This architecture, in the case of NODEs, ANODEs and SONODEs, was used on Silverbox, Airplane Vibrations and Van-Der-Pol Oscillator, with the aim of extracting coefﬁcients from the models, for these tasks we also allowed ANODEs to learn the initial augmented position.
The second architecture is a fully connected network with two hidden layers of size 20, it uses ELU activations in ˙ z and tanh activations in the initial conditions.
ELU and tanh were used because they allow for negative values in the ODE [11].
When considering ANODEs, they are in a higher-dimensional space than the problem, and the result must be projected down to the lower dimensions.
This projection was not learnt as a linear layer, instead, the components were directly selected, using an identity for the real dimensions, and zero for the augmented dimensions.
This was done because a ﬁnal (or initial) learnt linear layer would hide the advantages of certain models.
For example, the parity problem can be solved easily if NODEs are given a ﬁnal linear layer, do not move the points and then multiply by -1.
For this reason, no models used a linear layer at the end of the ﬂow.
Equally, they do not initialise with a linear layer as they again hide advantages.
For example, the nested n-spheres problem, NODEs can solve this with an initial linear layer, if they were to go into a higher-dimensional space the points may already be linearly separated, as shown by Massaroli et al.
D.1 Airplane vibrations
The dataset [12] concerns real vibrations measurements of an airplane.
A shaker was attached underneath the right wing, producing an acceleration a 1 .
Additional accelerations at different points were measured including a 2 , which was examined in this experiment, the acceleration on the right wing, next to a non-linear interface of interest.
This is a higher order system, therefore it pertains to be a challenging modelling task.
The results presented in Figure 6 show that while both methods can model the dynamics reasonably well, ANODEs perform marginally better.
We conjecture that this is due to ANODEs not being restricted to second order behaviour, allowing them to partially access higher order dynamics.
We test this conjecture in Appendix D.
[Figure]
D.2 Van Der Pol Oscillator
ANODEs and SONODEs were tested on a forced Van Der Pol (VDP) Oscillator that exhibits chaotic behaviour.
More speciﬁcally, the parameters and equations of the particular VDP oscillator are:
[Equation]
[Figure]
As shown in Figure 7, while ANODEs achieve a lower training loss than SONODEs, their test loss is much greater.
We conjecture that, in the case of ANODEs, this is a case of overﬁtting.
SONODEs, on the other hand, can better approximate the dynamics, therefore they exhibit better predictive performance.
Note that, neither model can learn the VDP oscillator particularly well, which may be attributed to chaotic behaviour of the system at hand.
D.3 Third Order NODEs on Airplane Vibrations
We test Third Order Neural ODEs (TONODEs) on the Airplane Vibrations task from Appendix D.1.
The results are in Figure 8.
We see that TONODEs vastly underperform compared to ANODEs and SONODEs.
In each of the 3 repetitions of the experiment, the different initialisation found the best solution to be at zero.
Therefore, whilst the loss stays constant, the error remains large.
We hypothesise that despite theoretically being able to perform at least as well as SONODEs, TONODEs avoid exponentially growing at any point by exponentially decaying towards zero.
It is likely that by rescaling the time to be between 0 and 1, TONODE would approach a more accurate solution.
[Figure]
D.4 First Order Dynamics and Interpolation
SONODEs contain NODEs as a subset of models.
Consider ﬁrst order dynamics that is approximated by the NODE
[Equation]
[Equation]
Carrying out the full time derivative of Equation (89):
[Equation]
Which yields the SONODE equivalent of the learnt dynamics:
[Equation]
, ˜ θ f ) (91)
Additionally, it was shown in Equation (3) that SONODEs are a speciﬁc case of ANODEs that learn the initial augmented position.
Therefore, anything that NODEs can learn, SONODEs should also be able to learn, and anything SONODEs can learn, ANODEs should be able to learn.
To demonstrate that SONODEs and ANODEs can also learn ﬁrst order dynamics, we task them with learning an exponential with no noise, x ( t ) =
exp (0 .
1667 t ) .
All models, as expected, are able to learn the function, as shown in Figure 9.
[Figure]
D.5 Performance on MNIST
NODEs, SONODEs and ANODEs were tested on MNIST [9] to investigate their ability on classiﬁcation tasks.
The networks used convolutional layers, which in the case of SONODEs were used for both the acceleration and the initial velocity.
ANODEs were augmented with one additional channel as is suggested by Dupont et al.
The models used a training batch size of 128 and test batch size of 1000, as well as group normalisation.
SGD optimiser was used with a learning rate of 0.1 and
[Table]
momentum 0.9.
The cross-entropy loss was used.
The experiment was repeated 3 times with random initialisations to obtain a mean and standard deviation.
The results are given in table 1 and Figure 10.
In terms of test accuracy, SONODEs and ANODEs perform marginally better than NODEs.
ANODEs can achieve the same accuracy with fewer parameters than SONODEs because the dynamics are not limited to second order and it is only the ﬁnal state that is of concern in classiﬁcation.
However,
SONODEs are able to achieve the same accuracy with a lower number of function evaluations (NFE).
NFE denotes how many function evaluations are made by the ODE solver, and represents the complexity of the learnt solution.
It is a continuous analogue of the depth of a discrete layered network.
In the case of NODEs and ANODEs, the NFE gradually increases meaning that the complexity of the ﬂow also increases.
However, in the case of SONODEs, the NFE stays constant, suggesting that the initial velocity was associated with larger gradients (otherwise we would expect
NFE to increase for SONODEs with training).
[Figure]
s t
[Figure]
c
[Figure]
[Figure]
Epoch
[Figure]
N F E
MNIST
[Figure]
E ANODEs learning 2nd Order
Here we demonstrate the different ways that ANODEs and SONODEs learn second order dynamics, and how this affects their solutions.
Proposition E.1.
The general form ANODEs learn second order behaviour is given by:
[Equation]
Proof.
Let z ( t ) be the state vector
[ x ( t ) , a ( t )] .
The time derivatives can be written as
[Equation]
Let x ( t ) follow the second order ODE, ¨ x = ˙ F = f
( a ) ( x , ˙ x , t, θ f ) .
Differentiating F with respect to time
[Equation]
Using ˙ x = F and ˙ a =
G
[Equation]
Rearranging for G
[Equation]
In order for the solution of G to exist, the matrix ∂F ∂ a T must be invertible.
Either the dimension of a matches F , x and f ( a ) , so that ∂F ∂ a T is square, or ∂F ∂ a T has a left inverse.
Crucially, F must have explicit a dependence, or the inverse does not exist.
Intuitively, in order for real space to couple to augmented space, there must be explicit dependence.
Using the left inverse, we do not require ∂F ∂ a T to be square, so a does not necessarily need to be the same dimensionality as F and x .
ANODEs can learn second order dynamics with fewer dimensions than SONODEs.
Using the equation for G ( x , a , t, θ G ) , there is a gauge symmetry in the system, which proves proposition E.2.
Proposition E.2.
ANODEs can learn an inﬁnity of (non-trivial) functional forms to learn the true dynamics of a second order ODE in real space.
Proof.
Assume a solution for F ( x , a , t, θ F ) and G ( x , a , t, θ G ) has been found such that, ˙ F = f ( a ) and F ( x 0 , a 0 , t 0 , θ F ) = ˙ x 0 .
If an arbitrary function of x , φ ( x ) , is added to F , where φ ( x
[Equation]
=
[Equation]
The initial velocity is still the same.
The dynamics are preserved if there is a corresponding change in
G ˜ G ( x , a , t, θ
[Equation]
The proof can end here, however this can be simpliﬁed.
φ ( x ) has no explicit a or t dependence, so this equation simpliﬁes to
[Equation]
The term f ( a ) ( x , F + φ, t, θ f ) can be Taylor expanded (assuming convergence)
[Equation]
Which gives the corresponding change in G
[Equation]
+
[Equation]
This demonstrates that there are inﬁnite functional forms that ANODEs can learn.
This only considered perturbing functions φ ( x ) .
More complex functions can be added that have a or t dependence, which lead to a more complex change in G .
By contrast, we now show SONODEs have a unique functional form.
Proposition E.3.
SONODEs learn to approximate a unique functional form to learn the true dynamics of a second order ODE in real space.
Proof.
Consider a dynamical system
[Equation]
For these problems we let the loss only depend on the position, if it depends on position and velocity there would be more restrictions.
So if it is true when loss only depends on the position, it is also true when it depends on both position and velocity.
Assume that there is another system, that has the same position as a function of time
[Equation]
Where f ( x , v , t ) (cid:54) =
˜ f
(˜ x , ˜ v , t ) .
Because the initial conditions are given, the position and velocity are deﬁned at all times, and therefore position, velocity and acceleration can all be written as explicit functions of time.
x ≡ x ( t ) , v ≡ v ( t ) .
This allows for the acceleration to be written as a function of t only, f ( x , v , t ) =
f τ ( t ) for all t .
The same applies for the second system, ˜ x ≡ ˜ x
( t ) , ˜
v ≡ ˜ v ( t ) and ˜ f
(˜ x , ˜ v , t )
=
˜ f
τ ( t ) .
For all t , x ( t ) =
˜ x ( t ) , therefore, for any time increment, δt , x ( t + δt ) =
˜ x
( t + δt ) .
Taking the full time derivative of x and ˜ x ( t )
[Equation]
Using these two equations and the fact that x ( t ) =
˜ x ( t ) , it is inferred that v ( t ) =
˜ v ( t ) for all t .
Taking the full time derivative of v ( t ) and ˜ v ( t )
[Equation]
Using these two equation and the fact that v ( t )
=
˜ v ( t ) for all t , it is also inferred that f τ ( t ) =
˜ f
τ ( t ) for all t .
Using these three facts, x ( t ) =
˜ x ( t ) , v ( t ) =
˜ v ( t ) and f τ ( t ) =
˜ f
τ ( t ) .
It must also be true that f ( x ( t ) , v ( t ) , t ) =
˜ f
(˜ x ( t ) , ˜ v ( t ) , t ) −→ f ( x , v , t ) =
˜ f
( x , v , t ) .
Therefore the assumption that f ( x , v , t ) (cid:54) =
˜ f
(˜ x , ˜ v , t ) is incorrect, there can only be one functional form for f ( x , v , t ) .
Additionally, using v ( t ) =
˜ v ( t ) for all t , the initial velocities must also be the same.