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: This paper presents an interesting perspective that connects ordinary differential equations with diffeomorphisms. More specifically, possibly complicated and thus difficult-to-learn ODE dynamics are expressed as a morphed version of a simple "base" ODE system. The authors propose to learn the base dynamics alongside with the corresponding diffeomorphism. As experimentally shown, this approach leads to faster training, somewhat increases accuracy and allows for the long-term behaviour analysis of certain types of base ODE systems.