L EARNING E FFICIENT AND R OBUST O RDINARY D IFFERENTIAL E QUATIONS VIA D IFFEOMORPHISMS
Anonymous authors
Paper under double-blind review
A BSTRACT
Advances in differentiable numerical integrators have enabled the use of gradient descent techniques to learn ordinary differential equations (ODEs), where a ﬂexible function approximator (often a neural network) is used to estimate the system dynamics, given as a time derivative.
However, these integrators can be unsatisfactorily slow and unstable when learning systems of ODEs from long sequences.
We propose to learn an ODE of interest from data by viewing its dynamics as a vector ﬁeld related to another base vector ﬁeld via a diffeomorphism (i e., a differentiable bijection).
By learning both the diffeomorphism and the dynamics of the base
ODE, we provide an avenue to ofﬂoad some of the complexity in modelling the dynamics directly on to learning the diffeomorphism.
Consequently, by restricting the base ODE to be amenable to integration, we can speed up and improve the robustness of integrating trajectories from the learned system.
We demonstrate the efﬁcacy of our method in training and evaluating benchmark ODE systems, as well as within continuous-depth neural networks models.
We show that our approach attains speed-ups of up to two orders of magnitude when integrating learned ODEs.
I NTRODUCTION
The problem of ﬁtting an ordinary differential equation (ODE) to observed data is ubiquitous throughout many disciplines of the natural sciences and engineering (Perko, 1991).
Although traditional approaches have focused on ﬁxed-form system dynamics and inferring their parameters, here we consider the more general problem of learning ODEs, when their dynamics are completely unknown.
This problem arises, for example, in robot learning where ODEs are often used to parameterize learned motion (Sindhwani et al., 2018; Singh et al., 2020).
In deep learning, this problem appears within the context of Neural ODEs (Chen et al., 2018), a family of continuous-depth models where the evolution of hidden states is an ODE.
Recent developments in learning ODEs have allowed the use of differentiable adaptive step-size numerical integrators to train neural network dynamics via the adjoint method (Chen et al., 2018).
The learned ODE system allows us to integrate continuous trajectories at different initial conditions.
To roll out long and complicated trajectories with a numerical integrator, the neural network dynamics model is queried sequentially at each step.
This can be unsatisfactorily slow for time critical applications, such as those in robot control, and is known to suffer from numerical instabilities (Gholami et al., 2019; Choromanski et al., 2020).
With these challenges in mind, we propose an alternative approach to learning ODEs: we view the desired target ODE dynamics as a vector ﬁeld that is a “morphed” version of an alternative base vector ﬁeld via a diffeomorphism , i.e., a bijective mapping where both the mapping and its inverse are differentiable.
Thus, instead of directly modelling the time derivative of the desired ODE with a neural network, we use an invertible neural network to learn the diffeomorphism that relates the target ODE to a base ODE.
Crucially, by restricting the base ODE to be less complex and more amenable to integration, we can obtain a solution to the (more complex) target ODE by integrating the simpler base ODE and passing its solution through the bijection.
Figure 1 shows an example of related vector ﬁelds: the simple integral curves in the left ﬁgure are passed through a diffeomorphism to give corresponding complicated curves in the right.
Unlike rolling out a trajectory, evaluating the diffeomorphism needs not to be sequential, allowing for efﬁcient GPU computation.
We investigate restricting the ﬂexibility of the base ODE to improve integration efﬁciency, and ofﬂoading more of the representation burden to the diffeomorphism.
Speciﬁcally, we assume the dynamics of the base ODE to be: (1) linear or (2) modelled by a neural network.
Restricting the base
ODE to be linear allows the computation of closed-forms solutions, providing major speed beneﬁts.
In this set-up we can achieve signiﬁcant speed-ups of up to two orders of magnitude when employing
GPUs, as compared against differentiable integrators with standard settings.
We can also restrict the learned ODE to be provably asymptotically stable by adding simple constraints to the linear base
ODE.
Alternatively, when additional ﬂexibility is required, we remove the restrictive assumption of a linear base ODE and model the dynamics using a neural network.
We show that in this setting we can improve the performance of learning challenging ODEs compared to existing differentiable integrators, even when we use a simpler neural network for the base ODE.
In summary, our main contributions are:
a novel paradigm to learn ODEs from data: invertible neural networks are trained to “morph” the target ODE to an alternative related base ODE, which can be more tractably integrated; 2.
analysis of the base as (i) a linear ODE and (ii) a non-linear ODE with neural network dynamics.
In the linear case, we demonstrate how by restricting the ﬂexibility of the base ODE, we can obtain closed-form integrals, providing signiﬁcant speed-ups to integration.
In the non-linear case, we demonstrate that by learning ODEs as related vector ﬁelds, we can ﬂexibly learn challenging ODEs with simpler networks; 3.
a principled method to enforce asymptotic stability of learned ODEs, by adding restrictions to the base ODE.
Proofs and additional details can be found in the appendices.
R ELATED W ORK
Learning of ODEs and Neural ODEs: Dynamical systems governed by ODEs can be found throughout many disciplines of science and engineering.
Earlier work on approximating free-form dynamics of ODEs include gradient matching (Ramsay et al., 2006) and using Gaussian processes (Heinonen et al., 2018).
Most recent work on this problem model the unknown dynamics with a neural network and leverage differentiable numerical integrators, which use the adjoint method (Pontryagin et al., 1962) to train in a memory-usage tractable manner (see, e.g., Chen et al., 2018).
A particular usage of ODE learning is within Neural ODEs, which are neural network models that model the hidden state as continuous ODEs rather than discrete layers (Chen et al., 2018; Massaroli et al., 2020).
Other continuous neural network models which incorporate an ODE, such as latent
ODEs (Rubanova et al., 2019) have found application in time-series tasks.
Subsequent strategies have been introduced to improve the training of these models, including augmenting the ODE state-space (Dupont et al., 2019), hyper-network extensions (Choromanski et al., 2020), regularisation techniques (Finlay et al., 2020;
Pal et al., 2021), and investigating integrator step-sizes (Ott et al., 2021).
At the core of all neural ODE models is the differentiable integrator used to learn the underlying ODE.
Our proposed approach improves the learning of the underlying ODE, and is compatible with models that incorporate learnable ODEs.
We note that the term “neural ODE” has typically been used in the literature to refer to neural networks that incorporate ODEs, including the original work in Chen et al.
However, “neural ODE” has occasionally been used to refer to an ODE with dynamics parameterized by a neural network (Norcliffe et al., 2021).
To disambiguate, throughout our paper, we refrained from referring to the latter model as “neural ODEs”, but rather as “ODEs with dynamics parameterized by a neural network”.
Invertible neural networks and Normalizing Flows: Invertible neural networks (INNs) are a class of function approximators that learn bijections where the forward and inverse mapping and their
Jacobians can be efﬁciently computed (Ardizzone et al., 2019).
INNs are typically constructed by invertible building blocks, such as those introduced in Kingma et al.
; Dinh et al.
; Durkan et al.
Advances in INNs are largely motivated by normalizing ﬂows (Rezende & Mohamed, 2015; Papamakarios et al., 2021), an approach to construct a ﬂexible probability distribution by ﬁnding a differentiable bijection between the target distribution and a base distribution.
Our approach is similar in spirit to normalizing ﬂows, as we analogously aim to learn a diffeomorphism that relates the vector ﬁelds of the target ODE and some base ODE.
However, unlike normalizing ﬂows, we do not require the burdensome computation of Jacobian determinants (Karami et al., 2019) to obtain trajectories.
A separate line of work, broadly characterized as continuous normalizing ﬂows , use ODEs to build invertible approximators (Grathwohl et al., 2019; Chen et al., 2018).
Our work proposes the opposite where invertible approximators are used to learn ODEs.
P RELIMINARIES
In this section we introduce the problem of learning ODE dynamics with neural networks.
We then describe the notions of tangent spaces and pushforwards, which will be used in section 4 to develop our method.
3. L EARNING ODE S WITH NEURAL NETWORKS
Consider a dynamical system given by ordinary differential equations of the form:
[Equation]
where t is time, y ( t ) are the states at time t , and f provides the dynamics.
Unlike traditional approaches where f is assumed known with only a few parameters to estimate from data, here we consider the more general problem where the dynamics are completely unknown.
Thus, we can use a ﬂexible mapping f ω as given by a neural network with parameters ω .
Henceforth, we will drop the explicit dependence on time, and consider the autonomous ODEs given by y (cid:48) ( t ) =
f ( y ( t )) .
Nonautonomous ODEs, which explicitly depend on time, can be equivalently expressed as autonomous
ODEs by adding a dimension to the states y (Davis et al., 2020).
For an initial condition y t 0 at start t 0 , and some end time t e , a solution of the ODE can be evaluated by a numerical integrator (ODESolve), such as Runge-Kutta methods (Butcher, 1987):
[Equation]
The learning problem involves estimating, with f ω , the dynamics of the ODE, provided n t observations y obst 1 . . .
y obst nt at speciﬁed times.
We can learn the ODE by optimising the parameters ω to minimize a loss between the observations at the given times and the integrated ODE, (cid:96) ( ω ) =
Loss( { y obst i } n t
i
=1 , { y ( t i ) } n t i =1 ) , where { y ( t i ) } n t i =1 are obtained by solving eq.
Advances in the neural ODE literature have introduced differentiable numerical integrators, which allow gradient-based techniques to be used to optimize l ( ω ) .
Furthermore, by using the adjoint sensitivity method as outlined in Chen et al.
(2018), the gradients of adaptive integrators can be obtained in a memory tractable manner, without differentiating through the integrator operations.
Nevertheless, the ﬂexibility of neural network dynamics for ODE learning comes at the expense of a high-computational cost and potential numerical instabilities, especially when considering long trajectories.
We will develop in section 4 an alternative method that transforms this problem into that of learning a simpler ODE along with a diffeomorphism, by treating the dynamics of the original (target) ODE and simpler (base) ODE as related vector ﬁelds.
3. T ANGENT S PACES AND P USHFORWARDS
As mentioned above, we shall be analysing the system dynamics of ODEs as vector ﬁelds.
Here we brieﬂy introduce the differential geometry notions of tangent spaces and pushforwards, which will be used to deﬁne related vector ﬁelds , a core concept underpinning our methodology.
Tangent Spaces: A manifold is a space that locally resembles Euclidean space.
Throughout this paper, all manifolds will be assumed to be differentiable, with deﬁned tangent spaces .
For an
T M
T N
M
N
DF
X
F
F −
Y
[Figure] [Figure]
n -dimensional manifold M , at a point p ∈ M , the tangent space T p M is an n -dimensional real vector space, where each element passes p tangentially and is referred to as a tangent vector .
The tangent space provides a higher-dimensional analogue of a tangent plane at a point on a surface.
The collection of tangent spaces for all points on M is known as the tangent bundle denoted by T M .
Pushfoward: For a mapping F : M → N between two manifolds, M and N , the pushforward by F is a linear mapping between the tangent spaces of the manifolds, D p F : T p M → T F ( p ) N .
Tangent vectors at p in the domain M can be mapped to tangent vectors at the corresponding point
F ( p ) in the codomain N via the pushforward.
This is computed by the matrix product of the Jacobian of F at p , J F ( p ) , and a tangent vector at p .
M ETHODOLOGY
We study the dynamics of ODEs, f , as vector ﬁelds, and solutions as their integral curves.
We model the desired ODE dynamics as a target vector ﬁeld that is related to another base vector ﬁeld.
First, we introduce the concept of related vector ﬁelds, outline how they can be learned, and elaborate on the beneﬁts of learning them.
Then, we describe possible choices of base vector ﬁeld models.
4. R ELATED V ECTOR F IELDS FOR ODE L EARNING
A vector ﬁeld X deﬁned on manifold M is a function that assigns a tangent vector X p ∈ T p M to each point p ∈ M .
Intuitively, our aim is to construct a mapping F which shapes the manifold where a base vector ﬁeld X is deﬁned, such that the pushforward of X by F extrinsically appears “morphed” to match the data.
What are the requirements of these mappings, for the pushforward of vector ﬁelds to be valid?
Provided a mapping between manifolds F : M → N , we can push a single vector, X p ∈ T p M , to the tangent space of N at F ( p ) , T F ( p ) N , via the pushforward , D p F ( X p ) .
Figure 2a sketches out an example of the pushforward of a tangent vector between tangent spaces.
However, this notion does not extend in general to vector ﬁelds.
If F is injective and non-surjective, the pushforward of
X outside the image of F is not deﬁned.
If F is surjective and non-injective, there may be multiple differing pushforwards for a point.
In special cases when the pushforward by F deﬁnes a valid vector ﬁeld on the codomain N , the vector ﬁeld and its pushforward are known to be F -related.
Latent
ODEs, as described in (Rubanova et al., 2019), which use an auto-encoder to map input data to a latent space, where an ODE is learned is similar in spirit to our method.
Latent ODEs, however, do not deﬁne a valid vector ﬁeld in the input space.
Deﬁnition 4.1 (Related vector ﬁelds) .
Let F : M → N be a smooth mapping of manifolds.
A vector ﬁeld X on M and a vector ﬁeld Y on N are related by F , or F -related , if for all p ∈ M ,
[Equation]
F ( p ) .
Related vector ﬁelds arise in particular when F is a diffeomorphism, i.e. a bijective mapping, where both the mapping itself and its inverse are differentiable.
Proposition 4.1 (Proposition 8.19 in Lee (2012)) .
Suppose F : M → N is a diffeomorphism between smooth manifolds M , N .
For every vector ﬁeld X on M , there is a unique vector ﬁeld Y on N that is F -related to X .
By considering the F -related properties of vector ﬁelds, we have a pathway to deﬁne unknown vector ﬁelds using the pushforward of F , as shown in ﬁg. 2b.
If vector ﬁeld X on M is F -related to some vector ﬁeld Y on N , instead of directly evaluating the vector ﬁeld Y , we can instead obtain tangent values for any q ∈ N , via N F − 1 −−−→ M X −→ T M DF −−→ T N .
Therefore, the vector attached for each q ∈ N is, Y q =
D F − 1 ( q ) F ( X F − 1 ( q ) ) =
J F ( F − 1 ( q ))
X F − 1 ( q ) , where J F is the Jacobian of F .
4. D IFFEOMORPHISM LEARNING VIA I NVERTIBLE N EURAL N ETWORKS
The machinery to learn invertible mappings has seen extensive development with the progress of normalizing ﬂows for estimating probability distributions.
Invertible neural networks (INNs,
Ardizzone et al., 2019) are function approximators which learn differentiable bijections.
INNs can be trained on a forward mapping, and get the inverse with no additional work, by the deﬁnition of their architecture.
In this paper, we use INNs of the type described in Dinh et al.
The basic unit is a reversible block, where inputs are split into two halves, u 1 and u 2 .
The outputs v 1 and v 2 are:
[Equation]
where (cid:12) indicates element-wise multiplication, and t 1 , t 2 and s 1 , s 2 are functions modelled by fully-connected neural networks with non-linear activations.
These expressions are clearly invertible:
[Equation]
Note that the functions t 1 , t 2 and s 1 , s 2 themselves are not required to be invertible.
4. ODE S OLUTIONS AS I NTEGRAL C URVES OF R ELATED V ECTOR F IELDS
Why would it be beneﬁcial to construct a desired vector ﬁeld Y indirectly, by way of a related X ?
We shall answer this by considering integral curves on Y , which represent solutions to the ODE associated with Y .
An integral curve of Y on N is a differentiable curve y : R → N , whose velocity at each point is equal to the value of Y at that point, i.e. y (cid:48)
( t ) =
Y y ( t ) ∈
T y ( t ) N , for all t ∈ R .
The integral curves of F -related vector ﬁelds are also linked by F : integral curves on one vector ﬁeld are uniquely mapped to the other via a single pass through F , and the Jacobian J F is not required.
Proposition 4.2 (Proposition 9.6 in Lee (2012)) .
Suppose X and Y are vector ﬁelds on manifolds
M and N respectively.
X and Y are related by mapping F : M → N if and only if for each integral curve x : R → M , y = F ( x ) is an integral curve of Y .
In the ODE learning problem outlined in section 3.1, we denote Y to be the vector ﬁeld associated with the target ODE.
During both training and inference, we need to obtain integral curves y of Y either by numerical integration, or by y = F ( x ) , where x denotes the corresponding integral curve of
X , related to Y via the diffeomorphism F .
Critically, the Jacobian of F does not need to be evaluated when we are working with the integral curves.
If integral curves of X can be found in a more efﬁcient, or less error-prone manner, than by numerically integrating curves of Y , we can leverage the relationship y
= F ( x ) for ODE learning.
This can be done by an INN, F θ , with parameters θ .
We denote the base ODE as x (cid:48) ( t ) =
g ϕ
( x ( t )) , where ϕ are parameters.
We can then use the target ODE within some learning problem, minimising a loss over target ODE trajectories and observations:
[Equation]
where there are n t data time points, y 0 is an initial condition for the system, and y obst i are observed data points at times t i .
After training, the dynamics of the target ODE is given by y (cid:48)
[Equation]
) =
J F θ ( F − 1 θ ( y ( t )))
g ϕ ( F − 1 θ ( y ( t ))) and the ODE solutions (integrals) are obtained with:
[Equation]
In practice, we are often required to evaluate an entire trajectory, i.e., y ( t ) at multiple times t 1 , . . .
, t end with one initial y 0 , as outlined in Algorithm 1.
This allows us to batch up the pass through F θ , which makes this highly efﬁcient when executed on a GPU.
The beneﬁts of our method are apparent when it is advantageous to integrate the base ODE and then pass the solution through the diffeomorphism, F θ , rather than numerically integrate the target ODE.
[Table]
The integral is easier to solve.
[Equation]
[Table]
through INN can be efficiently computed on GPUs.
Next, we investigate restricting the ﬂexibility of the base ODE, so that it is more amenable to integration, ofﬂoading the complexity of learning to the diffeomorphism.
We consider two choices of base
ODEs: (1) Linear ODE; (2) ODE with neural network dynamics.
4. L INEAR ODE AS B ASE : F AST I NTEGRATION AND S TRAIGHTFORWARD A SYMPTOTIC S TABILITY
We can speed-up integration signiﬁcantly by modelling the base as a Linear ODE, of the form x (cid:48) ( t ) =
A x ( t ) , where x ( t ) ∈ R n are n -dimensional variables, and A ∈ R n × n .
Linear ODEs can be solved very efﬁciently as they admit closed-form solutions.
Provided an initial solution x 0 , the solution of x ( t ) and that of the target
ODE y ( t ) are: x ( t ) =
(cid:80) nk
=1 ( l k · x 0 ) r k exp( λ k t ) , and y ( t ) =
F θ ( x ( t )) , where l k , r k and λ k are the corresponding left, right eigenvectors and eigenvalues of matrix A , respectively.
We learn the eigenvalues and eigenvectors of matrix
A jointly with diffeomorphism F θ .
Linear ODEs are also interesting because their long-term behavior, which is determined by their eigenvalues, is easy to analyse.
We shall see how this property allows us to craft the long-term behavior of the target ODE.
In particular, in many applications, consideration is given to the asymptotic properties of ODEs, namely what happens to the solutions after a long period of time.
Will they converge to equilibrium points, periodic orbits, or diverge and ﬂy off?
Our method provides a straightforward way to restrict the learned ODE to be asymptotically stable.
In robot motion generation problems, such as that in Sindhwani et al.
(2018), we aim to learn an asymptotically stable
ODE.
We begin by deﬁning equilibrium points and asymptotic stability of ﬁrst order ODEs.
Deﬁnition 4.2 (Equilibrium point) .
An equilibrium point y ∗ of an ODE y (cid:48) ( t ) =
f ( y ( t )) , is a point where f ( y ∗ )
= 0 .
Deﬁnition 4.3 (Asymptotic stability) .
An ODE y (cid:48) ( t ) =
f ( y ( t ) , t ) is asymptotically stable if for every solution y ( t ) , there exists a δ > 0 , such that whenever || y ( t 0 ) − y ∗ || <
δ , then y ( t ) → y ∗ as t → ∞ , where y ∗ is some equilibrium point.
Intuitively, asymptotically stable systems of ODEs will always settle at some equilibrium points after a long period of time.
In the context of vector ﬁelds related by a diffeomorphism, the asymptotic stability properties of the ODEs are shared.
Theorem 4.1.
Suppose two ODEs x (cid:48) ( t ) =
g ( x ( t )) , y (cid:48) ( t ) =
f ( y ( t )) are related via y ( t )
=
F ( x ( t )) , where F is a diffeomorphism.
If the former ODE is asymptotically stable with n e equilibrium points x ∗ 1 , . .
, x ∗ n e , then the latter is also asymptotically stable, with equilibrium points F ( x ∗ 1 ) , . . .
, F ( x ∗ n e ) .
Therefore, if we can restrict the base ODE to be asymptotically stable, then the target ODE learned by our method is also asymptotically stable.
When the base is an n dimensional linear ODE, we can restrict it to be asymptotically stable by directly learning the eigenvalues, λ i for i = 1 , . . . , n , and constraining them to be negative.
This can be done by setting λ i = − ( s λ i ) 2 − ε , where ε is a small positive constant, and learning s λ i instead of learning the eigenvalues.
4. N EURAL N ETWORK ODE AS B ASE : I MPROVED R OBUSTNESS FOR ‘D IFFICULT ’ ODE S
Using linear systems as base ODEs provides a dramatic increase in speed at the cost of ﬂexibility.
We observe that the computation overhead of a single backward pass F − 1 θ and a batched single forward pass F θ is minimal when compared with numerical integration which requires sequential querying.
When the ODE is difﬁcult to learn, we can also parameterize the dynamics of the base
ODE using a neural network.
This is particularly appealing for ODEs that are considered stiff, with rapid varying of the solution in time in one dimension, while the other dimensions remain largely unchanged.
Existing differentiable integrators often struggle to directly learn these ODEs.
Directly learning the target ODE will require exceedingly small step-sizes (Hairer et al., 1993).
Although by simply scaling the data before training, we can lessen the stiffness of the ODE to learn, the relatively rapid changes isolated in a single dimension can still result in the ODE being hard to learn.
To tackle this, we learn a neural network dynamics of the base ODE jointly with the diffeomorphism.
The burden of accurately representing the stiff dynamics is shared by F θ , providing additional ﬂexibility.
The diffeomorphism can learn to relate the target ODE to a base that is amenable to integration.
This allows us to more accurately learn challenging ODEs, with a much smaller neural network model of the base ODE, than that used to directly learn the target ODE.
As passes of the INN are inexpensive, we gain a speed-up during integration.
E XPERIMENTAL R ESULTS
We empirically evaluate the ability of our method to speed up the integration of learned ODEs, along with the robustness of integration when learning difﬁcult ODE systems.
Throughout this section, we compare the error and integration times of our method against a variety of solvers.
We include ﬁxed step-size solvers: Euler’s, midpoint, and Runge-Kutta 45 (RK4), and the adaptive step-size solvers Dormand–Prince 5 (DOPRI5) and Dormand–Prince 8 (DOPRI8).
Unless speciﬁed otherwise, for ﬁxed step-size solvers, we set the step-size equal to the smallest time increment of outputs.
For adaptive step-size solvers, we set absolute and relative tolerances to 10 − 5 .
We augment the ODE states for all the systems in accordance to Dupont et al.
(2019) in all of the ODEs trained during the experiments, except when recreating latent ODE results from Rubanova et al.
(2019), where we use the original implementation provided by the authors.
The differentiable solvers are implemented in the torchdiffeq library.
Details on experiments available in the appendices.
5. I NTEGRATION S PEED - UPS BY L EARNING WITH A L INEAR ODE B ASE
We test our hypothesis that the availability of a closed-form expression for the integral, when using a linear base ODE, can provide signiﬁcant integration speed-ups.
We evaluate on learning synthetic
ODE systems, real-world robot demonstrations, and within continuous deep learning models.
Here, we report test error/accuracy and integration times.
Training times can be found in the appendices.
Learning 3D Lotka-Volterra: We train and evaluate models on data from the 3D Lokta-Volterra system, which models the dynamics of predator-prey populations.
The data is corrupted by white noise with standard deviation of 0 .
We train our model using a linear base ODE, and assess the capability of our model in interpolating the data points at 10 x the data resolution, and generalizing to 16 hidden test initial conditions to integrate trajectories, also at 10 x the data resolution.
We report the integration time for generalization.
Figure 3 shows interpolation results and newly generated trajectories, where we see that our model is able to capture the dynamics of the system.
Furthermore, table 1 provides a quantitative evaluation, where we see that our method is signiﬁcantly faster
[Table]
than competing approaches using numerical integrators with speed-ups of more than two orders of magnitude, while achieving comparable or better accuracy.
[Figure]
A Time Critical Application—Robot Motion Generation as Stable ODE Learning: The ability to quickly roll-out trajectories is crucial in motion robotics settings.
We consider the application of generating robot manipulator motion trajectories from provided demonstrations.
In particular, Sindhwani et al.
(2018) showed that modelling the motion as trajectories of a stable ODE is critical for generalizing and being robust to perturbations in initial conditions.
The goal is to learn an ODE system where trajectories integrated at different starting points mimic the shown demonstrations.
We use three sets of real-world data from (Khansari-Zadeh & Billard, 2011) of trajectories: drawing “S” shapes; placing a cube on a shelf; drawing out large “C” shapes.
We use 70% of the data for training, and test our generalization on the remaining demonstrations.
In these datasets, the motions are known to converge to equilibrium points.
Hence, we constrain the learned ODE to be asymptotically stable.
We report the performance and run-times of generalizing to new starting points in table 1.
We see that our approach is competitive in the quality of generalized trajectories, while achieving speed-ups of more than two-orders of magnitude.
ODE Learning for Continuous
Deep Learning Models: We evaluate our method as a component of Latent
ODEs (Rubanova et al., 2019), a continuous deep learning model.
Latent
ODEs embed the time series observations as hidden states via an encoderdecoder.
An ODE, with dynamics parameterized by a neural network, is ﬁt on the hidden states which allows for irregularly sampled series.
In these experiments, our method is applied with a linear ODE base to learn the dynamics governing the hidden states.
We report results for the periodic curves and the human activity classiﬁcation problem used in the original latent ODE paper (Rubanova et al., 2019), as well as for ECG classiﬁcation and image sequence generation.
The periodic curves problem requires us to reconstruct trajectories at different resolutions, with 100 and 1000 time-steps.
The human activity classiﬁcation dataset contains readings over the body of subjects over time, and the problem has a sequence-to-sequence setup, requiring us to predict the human activity at each time point.
The ECG problem is a sequence classiﬁcation problem, where the sequences are
ECG signals.
In the image sequence generation problem, we train on 100 sequences of rotating “3” characters, similar to that in Yildiz et al.
(2019), learning a high-dimensional ODE which governs a latent representation of the image rotation over time.
We test on 100 unseen “3” characters and report results at 10 x training resolution.
The performance and times spent on integrating the hidden state dynamics are reported in table 2.
We observe that by leveraging the closed-form expression of integrals of the base ODEs, we achieve ODE integration times that are hundreds of times faster, while achieving competitive performances against compared differentiable integrators.
We note that the main cost of the integral in our method is the pass through the invertible neural network.
GPUs allow us to batch the pass at practically constant cost, whereas the sequential nature of integrators give linear increases in run-time.
A qualitative evaluation of our method on the image sequence problem is shown in ﬁg. 4.
We observe the structure of the rotated “3” character remains consistent with the test initial image.
Our approach, when restricted to a linear base ODE, is able to learn high-dimensional
ODEs which can be integrated signiﬁcantly faster, without compromising performance.
5. R OBUST I NTEGRATION BY L EARNING WITH A N ON -L INEAR N EURAL N ETWORK B ASE
We test our hypothesis that using a neural network base ODE allows us to learn ODEs that are difﬁcult to integrate.
We learn and evaluate models trained on the chaotic Lorenz system, and the stiff
[Table]
when using latent ODEs on the tasks of periodic curve reconstruction using 100 and 1000 time-steps, the classiﬁcation problems of human activity and ECG
[Table]
and the image sequence generation with our method using a linear base ODE and competing numerical integrators.
[Figure]
Robertson’s system (ROBER, Robertson, 1966).
We use a DOPRI5 adaptive step-size integrator to learn our base ODE, generate trajectories at 10 x data resolution.
Figure 5 (left) illustrates a generated trajectory from the Lorenz system, which closely resembles the data points (blue).
We can see that both the initial large and small variations, indicated by the small dense spiral at the end, are captured.
Figure 5 (center) illustrates the particular rapidly changing second dimension of the ROBER system.
We see that the added ﬂexibility of the diffeomorphism allows us to better capture the rapid variations over time, while the directly learned ODE struggles to handle the sudden increase.
Following the equation rescaling described in Kim et al.
, before we train on our method and comparisons, we rescale our data by the maximum training value in each dimension and operate in logscale time.
Figure 5 (right) provides the performance and integration times of learning with our method and baseline integrators, where we see that our method is more accurate than competing approaches while also requiring shorter integration times, due to having a simpler network modelling the dynamics.
C ONCLUSIONS
We have proposed a novel approach to learning ODEs with unknown dynamics, which uses invertible neural networks to learn a diffeomorphism relating a desired target ODE to a base ODE that is often easier to integrate.
We have investigated using a base ODE that is linear or parameterized by a neural network.
By leveraging the closed form solution of linear ODEs, our method provides signiﬁcant speed-ups and allows for asymptotic property constraints on the learned ODEs.
Alternatively, by using a base ODE parameterized by a neural network, our approach can learn “difﬁcult” ODEs, with simpler networks modeling their dynamics.
We have validated our method by learning ODEs on synthetic and real-world data, on robotic learning problems and within continuous-depth neural network models.
Future work could explore more on how to balance ofﬂoading the burden of learning to the diffeomorphism and the base ODE.
R EPRODUCIBILITY S TATEMENT
It is extremely important that the work published in ICLR is reproducible.
To this end, we have included source code for our experiments, including code to generate the benchmark dynamical systems, as supplementary materials to the submission.
Furthermore, additional details of experiments and data, along with clear explanations of any assumptions and a complete proof of the theoretical claims are included in the appendix.
A A DDITIONAL R ESULTS : R EPRESENTATION L OAD OF THE INN
In the setup of using a non-linear neural network base ODE, both the INN and neural network dynamics are trained jointly.
We empirically investigate the representation burden on the invertible neural network, when using a non-linear neural network base ODE.
We train, for 5000 epochs, augmented ODEs with neural network dynamics models of different parameter sizes, without the diffeomorphism, on the Lorentz system and report the MAE values of integrating with the dopri integrator with the same setup as the Lorentz experiments in Section 5.2.
We then compare take the neural network dynamics model with the smallest number of parameters (with a parameter count of 2256), and add INNs with an increasing number of invertible layers.
The results of the ODE with neural network parameters only, with increasing parameter counts, are tabulated in table 3, while the results of the neural networks with an increasing number of invertible layers are tabulated in table 4.
Each invertible layer contains 15090 parameters.
We note here that the querying the INN (and thereby having a large INN) adds negligible integration time, as during the entire integration we only need one forward and one inverse pass of the INN, while the system dynamics needs to be queried sequentially.
We see that, in general, adding parameters to the system dynamics of a learned ODE and to the diffeomorphism adds to the representation power of our ODE models.
However, we clearly see that the addition of an INN improves performance, and additional layers to the INN, up to around 5 invertible layers, improves performance.
We see that the model with 5 invertible layers and a base
ODE of 2256 parameters greatly outperforms the ODE model with only a neural network dynamics of 83006, while the two models are similar in parameter count.
This indicates that the invertible neural network plays a large role in representing the learned ODE.
[Table] [Table]
B A DDITIONAL R ESULTS : T RAINING TIMES
We present the training times for directly learning ODEs with our method, using a linear base ODE.
These include the training times on the 3D Lotka-Volterra, and the robot imitation datasets, outlined in Sections 5.1.1 and 5.1.2 of the main paper.
We run training for 1000 iterations, where in each iteration the batch includes the entire training set.
We see that, by leveraging the closed-form solution of linear ODEs, our method is able to also drastically speed up training.
Additionally the parallel nature of passing through the invertible neural network allows more consistent training times across datasets.
[Table]
C A DDITIONAL R ESULTS : A BLATION S TUDY
We study the effects of the number of layers in the invertible neural network and number of parameters in the sub-network, which are the main hyper-parameters of the invertible neural networks used.
To this end, we conduct ablation studies of the speed and performance of our method on the real-world datasets outlined in section 5.1.2 of the paper.
Our basic model uses an invertible neural network with 5 layers and sub-networks in the invertible network had 1500 hidden dimension size.
We alter the number of layers to be: 2 , 3 , 4 , 5 , 6 , 7 ,
[Table]
and hidden dimensions of the sub-networks within the invertible network to be: 500 , 1000 , 1500 , 2000 , 2500 .
The results are presented below:
[Table]
We see that as we increase the number of invertible network layers, the integration times increase, while the hidden dimension size of the sub-networks within the invertible network does not visibly affect the integration times.
Overall, the generalisation performance improves as the number of invertible layers are used, up to some number of layers.
Beyond this number of layers, adding layers does not vary performance signiﬁcantly.
Additionally, the hidden dimension sizes, for the values tested do not greatly vary the generalisation performance.
D P ROOFS
Proofs for Propositions 4.1 and 4.2 can be found in Lee (2012) as Propositions 8.19 and 9.6.
Theorem 4.1.
Suppose two ODEs x (cid:48) ( t ) =
g ( x ( t )) , y (cid:48) ( t ) =
f ( y ( t )) are related via y ( t )
=
F ( x ( t )) , where F is a diffeomorphism.
If the former ODE is asymptotically stable with n e equilibrium points x ∗ 1 , . .
, x ∗ n e , then the latter ODE is also asymptotically stable, with equilibrium points
F ( x ∗ 1 ) , . .
, F ( x ∗ n e ) .
Proof.
First we show F ( x ∗ 1 ) , . . .
, F ( x ∗ n e ) are equilibrium points of ODE y (cid:48) ( t ) =
f ( y ( t )) .
By y ( t ) =
F ( x ( t )) , we can write the time derivatives y (cid:48) at F ( x ) as
[Equation]
where J F ( x ( t )) is the Jacobian of F .
F is a diffeomorphism and hence invertible over its domain.
By the inverse function theorem (Dontchev & Rockafellar, 2009), the Jacobian J F ( x ( t )) is invertible, and furthermore, by the invertible matrix theorem (Horn & Johnson, 2012), it has a null-space containing only the zero vector.
Therefore, y (cid:48) ( t ) =
f ( F ( x ( t )))
=
J F ( x ( t )) g ( x ( t ))
=
if and only if g ( x ( t )) =
As g ( x ∗ ( t )) = 0 for x ∗ ∈ { x ∗ 1 . . .
x ∗ n e } , then we also have f ( F ( x ∗ ( t )))
= 0 , hence y ∗ ∈ { F ( x ∗ 1 ) , . . .
, F ( x ∗ n e ) } gives equilibrium points for y (cid:48) ( t ) =
f ( y ( t )) .
We now show asymptotically stability of y (cid:48)
( t ) =
f ( y ( t )) , by the existence of a Lyapunov function (Lefschetz & Alverson, 1962), V y :
R n → R , where n is the dimension of y , such that ∂V y
( y ) ∂t <for all y ∈ R n \ { F ( x ∗ 1 ) , . .
, F ( x ∗ n e ) } , and ∂V y ( y ∗ ) ∂t = 0 for y ∗ ∈ { F ( x ∗ 1 ) , . . .
, F ( x ∗ n e ) } .
The existence of such a Lyapunov function is a necessary and sufﬁcient condition for stability.
We assume the candidate function to be V y
=
V x ( F − 1 ( y )) , where V x is a valid Lyapunov function of the asymptotically stable x (cid:48) ( t ) =
g ( x ( t )) , with ∂V x ( x ) ∂t < 0 for x ∈ R b \ { x ∗ 1 , . . .
, x ∗ n e } and ∂V x ( x ∗ ) ∂t = 0 for x ∗ ∈ { x ∗ 1 , . . .
, x ∗ n e } .
Consider the time derivative of the candidate function:
[Equation]
Equation (11) by the inverse function theorem (Dontchev & Rockafellar, 2009).
Therefore, our candidate V y is a valid Lyapunov function for y (cid:48) ( t ) =
f ( y ( t )) .
Thus, the system y (cid:48) ( t ) =
f ( y ( t )) is asymptotically stable.
E A DDITIONAL I MPLEMENTATION DETAILS
We run all of our experiments on a machine with an Intel i7-3770k 3.50GHz processor, 32GB
RAM and an NVIDIA GTX
GPU, with 8GB vRAM.
For all of our experiments, we use the optimizer ADAM with step-size 10 − 4 , except for the experiments in the Latent ODE, which where we use the standard set-up from the Latent-ODE repository (Rubanova et al., 2019).
The dynamics models of compared ODEs have the architecture: Input->dense(Input dimensions, 150)->tanh()->dense(150,150)->tanh()->dense(150,150)->tanh()->dense(150,150)->tanh()- >dense(150,150)->tanh()->dense(150,output dimensions)->output.
Except for the Latent-ODE comparisons where settings from the original repository (Rubanova et al., 2019) is used.
For all of the experiments, except latent ODE experiments where we follow the original set-up, we train foriterations in total.
For all the experiments where we directly learn a dynamical system, we use an invertible neural network with 5 invertible layers, and sub-networks with one hidden layer of 1500 units.
For non-linear base ODEs parameterized with a simple neural networks, we use the architecture:
Input->dense(Input dimensions,30)->tanh()->dense(30,30)->tanh()->dense(30,30)->tanh()- >dense(30,Output dimensions)->Outputs.
Additionally, all learned dynamics, both with ours and compared methods, excepted when adhering to the original Latent-ODE set-up, were augmented with the same number of additional zeros as original state dimensions, for example 3 dimensional systems were augmented to 6 dimensions.
An exception to this is the high-dimensional image rotation problem, where we found that adding half as many augmented states as the original state dimensions was sufﬁcient.
The Lotka-Volterra system used has the dynamics:
[Equation] [Equation]
[Equation]
for t ∈ [0 , 7] with initial conditions { (5 , 5 , 1) , (2 , 6 , 6) , (3 , 1 , 4) , (7 , 1 , 2) , (6 , 2 , 4) , (3 , 3 , 1) , (2 , 2 , 2) , (4 , 4 , 3) , (3 , 3 , 4) , (1 , 1 , 5) } .
The Lorenz system used has the dynamics:
[Equation] [Equation]
for t ∈ [0 , 2] with the initial conditions (0 . 15 , 0 . 15 , 0 . 15) .
The Robertson’s system used has the dynamics:
[Equation]
for t ∈ [0 , 120] with the initial conditions (1 , 0 , 0) .
During training and evaluation, we rescale the data dimensions, and roll out the ODE in log-space.
In the latent-ODE problem setup, an observable time-series is assumed to have latent variables which follow some ODE dynamics, and uses an Encoder -> ODE -> Decoder architecture where an ODE is used to model the hidden latent dimensions between the Encoder and Decoder.
Note that a valid ODE is not guaranteed in the space of observable data, but only in the latent dimensions.
Our set-up follows the repository given by Rubanova et al.
(2019), with the training settings for the
Encoder and Decoder architecture as below:
Periodic 100 : We train the entire model for 500 epochs with Adamax optimizer and an initial learning rate of 10 − 2 .
We sub-sample 5% of the original time points and the size of the latent state is 10.
The noise weight is set as 0.01 and the total number of time points is 100.
For the Neural ODE architectures, there is one layer in the recognition ODE and one layer in the generative ODE, andunit per layers.
For the GRU unit there exists 100 units per layer for the GRU update network.
All the above settings are exactly the same as the conﬁguration given in repository (Rubanova et al., 2019).
Periodic 1000 : Settings are the same as Periodic 100 , except that the total number of time points is set as 1000 to predict for ﬁner time steps.
Human Activity : The model is trained for 200 epochs, with a dimensionality of 15 in the latent state.
There are 4 layers in the recognition ODE and 2 layers in the generative ODE, and 500 units per layer.
The GRU unit has exists 50 units per layer.
ECG : Settings are the same as the classiﬁcation task of Human Activity , except that we use the ECG
Heartbeat data available at Dataset (2018), removing the ‘0’ class.
Rotating Image Sequence : We use the ﬁrst 100 MNIST “3” characters as training and the next “3” characters as test.
We create a sequence of 44 images for each initial character until we rotate the initial image by 180 degrees.
During testing, we integrate to obtain a sequence of 440, at 10 x data resolution.
We obtain latent representations of each image by training a convolutional autoencoder to obtain a 64 dimension latent vector.
We ﬁt the ODE on these latent dimensions, with an ADAM optimizer with learning rate 5 × 10 − 4 .
ODE dimensions and sequence lengths : The following table contains details on the dimensions of the ODE models.
For latent ODE models, these are the dimensions of the latent state.
We also provide the lengths of the integrated trajectories of the ODEs during inference.
F A DDITIONAL F IGURES
We provide ﬁgures for learning an additional Lorenz system for t ∈ [0 , 5] , with trajectory at initial condition ( − 3 . 1 , − 1 . 15 , 8 . 15) .
We see that our method, with a base ODE parameterized by a neural network, can generate trajectories that closely match the ground truth:
[Figure] [Figure]
We provide the change in coordinates over time, for the trajectory shown in ﬁgure 5(a) in the paper:
[Figure]
We provide additional plots of trajectories, at different start points, from a learned Lotka-Volterra system.
The ground truth data is in blue, while generated trajectory, of 10 x data resolution, is in red.
We provide an additional ﬁgure for trajectories generated at unseen starting points after being trained on the “imitation C” training data.
The four generated trajectories are in red, while the ground truths are in blue.
Our generated trajectories match the ground truth, and accurately capture the motion of drawing a “C” character.
G L ICENSES FOR P ACKAGES
Common scientiﬁc packages used in our code include: (i) Numpy (Harris et al., 2020)
(BSD license), for general linear algebra and miscellaneous math operations (ii) Matplotlib (Hunter, 2007)
(BSD compatible custom license), for plotting ﬁgures.
More specialized packages used include (i) FrEIA (Ardizzone et al., 2019)
(MIT license), for invertible neural networks; (ii) TorchDiffEq (Chen et al., 2018)
(MIT license), for differentiable numerical integrators; (iii) Pytorch (Paszke et al., 2019)
(BSD license), for optimisation and automatic differentiation; (iv) Latent-ODE (Rubanova et al., 2019)
(MIT license), for latent ODE implementation.