Learning Efficient and Robust Ordinary Differential Equations via Diffeomorphisms

1 INTRODUCTION . The problem of fitting 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 fixed-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 field that is a “ morphed ” version of an alternative base vector field 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 fields : the simple integral curves in the left figure 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 efficient GPU computation . We investigate restricting the flexibility of the base ODE to improve integration efficiency , and offloading more of the representation burden to the diffeomorphism . Specifically , 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 benefits . In this set-up we can achieve significant 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 flexibility 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 : 1. 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 flexibility of the base ODE , we can obtain closed-form integrals , providing significant speed-ups to integration . In the non-linear case , we demonstrate that by learning ODEs as related vector fields , we can flexibly 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 . 2 RELATED WORK . 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 . ( 2018 ) . 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 efficiently computed ( Ardizzone et al. , 2019 ) . INNs are typically constructed by invertible building blocks , such as those introduced in Kingma et al . ( 2016 ) ; Dinh et al . ( 2017 ) ; Durkan et al . ( 2019 ) . Advances in INNs are largely motivated by normalizing flows ( Rezende & Mohamed , 2015 ; Papamakarios et al. , 2021 ) , an approach to construct a flexible probability distribution by finding a differentiable bijection between the target distribution and a base distribution . Our approach is similar in spirit to normalizing flows , as we analogously aim to learn a diffeomorphism that relates the vector fields of the target ODE and some base ODE . However , unlike normalizing flows , 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 flows , 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 . 3 PRELIMINARIES . 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.1 LEARNING ODES WITH NEURAL NETWORKS . Consider a dynamical system given by ordinary differential equations of the form : y′ ( t ) = f ( y ( t ) , t ) , y ( 0 ) = y0 , ( 1 ) 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 flexible 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′ ( 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 yt0 at start t0 , and some end time te , a solution of the ODE can be evaluated by a numerical integrator ( ODESolve ) , such as Runge-Kutta methods ( Butcher , 1987 ) : y ( te ) = yt0 + ∫ te t0 fω ( y ( t ) ) dt = ODESolve ( fω , yt0 , te − t0 ) . ( 2 ) The learning problem involves estimating , with fω , the dynamics of the ODE , provided nt observations yobst1 . . .y obs tnt at specified 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 , ` ( ω ) = Loss ( { yobsti } nt i=1 , { y ( ti ) } nt i=1 ) , where { y ( ti ) } nt i=1 are obtained by solving eq . ( 2 ) . 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 flexibility 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 fields . 3.2 TANGENT SPACES AND PUSHFORWARDS . As mentioned above , we shall be analysing the system dynamics of ODEs as vector fields . Here we briefly introduce the differential geometry notions of tangent spaces and pushforwards , which will be used to define related vector fields , 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 defined tangent spaces . For an n-dimensional manifoldM , at a point p ∈ M , the tangent space TpM 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 onM is known as the tangent bundle denoted by TM . 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 , DpF : TpM → TF ( p ) N . Tangent vectors at p in the domainM can be mapped to tangent vectors at the corresponding point F ( p ) in the codomainN via the pushforward . This is computed by the matrix product of the Jacobian of F at p , JF ( p ) , and a tangent vector at p .

Summary: In order to better learn differential equations that produce discrete data, this paper proposes to learn a coordinate transform (i.e. the `diffeomorphism’ in the title) in addition to learning the vector field. It is argued that, after this coordinate transformation, one can either learn a linear ODE or learn a neural network parameterized ODE, and obtain improved efficiency and robustness.