Neural Flows: Efficient Alternative to Neural ODEs

1 Introduction . Ordinary differential equations ( ODEs ) are among the most important tools for modeling complex systems , both in natural and social sciences . They describe the instantaneous change in the system , which is often an easier way to model physical phenomena than specifying the whole system itself . For example , the change of the pendulum angle or the change in population can be naturally expressed in the differential form . Similarly , Chen et al . [ 11 ] introduce neural ODEs that describe how some quantity of interest represented as a vector x , changes with time : ẋ = f ( t , x ( t ) ) , where f is now a neural network . Starting at some initial value x ( t0 ) we can find the result of this dynamic at any t1 : x ( t1 ) = x ( t0 ) + ∫ t1 t0 f ( t , x ( t ) ) dt = ODESolve ( x ( t0 ) , f , t0 , t1 ) . ( 1 ) It is sufficient for f to be continuous in t and Lipschitz continuous in x to have a unique solution , by the Picard–Lindelöf theorem [ 14 ] . This mild condition is already satisfied by a large family of neural networks . In most practically relevant scenarios , the integral in Equation 1 has to be solved numerically , requiring a trade-off between computation cost and numerical precision . Much of the follow up work to [ 11 ] focused on retaining expressive dynamics while requiring fewer solver evaluations [ 22 , 37 ] . In the machine learning context we are given a set of initial conditions ( often at t0 = 0 ) and a loss function for the solution evaluated at time t1 . One example is modeling time series where the latent state is evolved in continuous time and is used to predict the observed measurements [ 16 ] . Here , unlike in physics for example , the function f is completely unknown and needs to be learned from data . Thus , [ 11 ] used neural networks to model it , for their ability to capture complex dynamics . However , note that this comes at the cost of the ODE being non-interpretable . ∗Work partially done during an internship at Amazon Research . Correspondence to : bilos @ in.tum.de . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Since solving an ODE is expensive , we want to find a way to keep the desired properties of neural ODEs at a much smaller computation cost . If we take a step back , we see that neural ODEs take initial values as inputs and return non-intersecting solution curves ( Figure 1 ) . In this paper we propose to model the solution curves directly , with a neural network , instead of specifying the derivative . That is , given an initial condition we return the solution with a single forward pass through our network . Straight away , this leads to improvements in computation performance because we avoid using ODE solvers altogether . We show how our method can be used as a faster alternative to ODEs in existing models [ 9 , 16 , 34 , 69 ] , while improving the modeling performance . In the following , we derive the conditions that our method needs to satisfy and propose different architectures that implement them . 2 Neural flows . In this section , we present our method , neural flows , that directly models the solution curve of an ODE with a neural network . For simplicity , let us briefly assume that the initial condition x0 = x ( t0 ) is specified at t0 = 0 . We handle the general case shortly . Then , Equation 1 can be written as x ( t ) = F ( t , x0 ) , where F is the solution to the initial value problem , ẋ = f ( t , x ( t ) ) , x0 = x ( 0 ) . We will model F with a neural network . For this , we first list the conditions that F must satisfy so that it is a solution to some ODE . Let F : [ 0 , T ] × Rd → Rd be a smooth function satisfying : i ) F ( 0 , x0 ) = x0 , ( initial condition ) ii ) F ( t , · ) is invertible , ∀t . ( uniqueness of the solution given the initial value x0 ; i.e. , the curves specified by F corresponding to different initial values should not intersect for any t ) There is an exact correspondence between a function F with the above properties and an ODE defined with f such that the derivative ddtF ( t , x0 ) matches f ( t , x ( t ) ) everywhere , given x0 = x ( 0 ) [ 47 , Theorem 9.12 ] . In general , we can say that f defines a vector field and F defines a family of integral curves , also known as the flow in mathematics ( not to be confused with normalizing flow ) . As F will be parameterized with a neural network , condition i ) requires that its parameters must depend on t such that we have the identity map at t = 0 . Note that by providing x0 we define a smooth trajectory F ( · , x0 ) — the solution to some ODE with the initial condition at t0 = 0 . If we relax the restriction t0 = 0 and allow x0 to be specified at an arbitrary t0 ∈ R , the solution can be obtained with a simple procedure . We first go back to the case t = 0 where we obtain the corresponding “ initial ” value x̂0 : = x ( 0 ) = F−1 ( t0 , x0 ) . This then gives us the required solution F ( · , x̂0 ) to the original initial value problem . Thus , we often prefer functions with an analytical inverse . Finally , we tackle implementing F . The second property instructs us that the function F ( t , · ) is a diffeomorphism on Rd . We can satisfy this by drawing inspiration from existing works on normalizing flows and invertible neural networks [ e.g. , 17 , 2 ] . In our case , the parameters must be conditioned on time , with identity at t = 0 . As a starting example , consider a linear ODE f ( t , x ( t ) ) = Ax ( t ) , with x ( 0 ) = x0 . Its solution can be expressed as F ( t , x0 ) = exp ( At ) x0 , where exp is the matrix exponential . Here , the learnable parametersA are simply multiplied by t to ensure condition i ) ; and given fixed t , the network behaves as an invertible linear transformation . In the following we propose other , more expressive functions suitable for applications such as time series modeling . 2.1 Proposed flow architectures . ResNet flow . A single residual layer xt+1 = xt + g ( xt ) [ 30 ] bears a resemblance to Equation 1 and can be seen as a discretized version of a continuous transformation which inspired the development of neural ODEs . Although plain ResNets are not invertible , one could use spectral normalization [ 26 ] to enforce a small Lipschitz constant of the network , which guarantees invertibility [ 2 , Theorem 1 ] . Thus , ResNets become a natural choice for modeling the solution curve F resulting in the following extension — ResNet flow : F ( t , x ) = x+ ϕ ( t ) g ( t , x ) , ( 2 ) where ϕ : R→ Rd . This satisfies properties i ) and ii ) from above when ϕ ( 0 ) = 0 and |ϕ ( t ) i| < 1 ; and g : Rd+1 → Rd is an arbitrary contractive neural network ( Lip ( g ) < 1 ) . One simple choice for ϕ is a tanh function . The inverse of F can be found via fixed point iteration similar to [ 2 ] . GRU flow . Time series data is traditionally modeled with recurrent neural networks , e.g. , with a GRU [ 12 ] , such that the hidden state ht−1 is updated at fixed intervals with the new observation xt : ht = GRUCell ( ht−1 , xt ) = zt ht−1 + ( 1− zt ) ct , ( 3 ) where zt and ct are functions of the previous state ht−1 and the new input xt . De Brouwer et al . [ 16 ] derived the continuous equivalent of this architecture called GRU-ODE ( see Appendix A.1 ) . Given the initial condition h0 = h ( t0 ) , they evolve the hidden state h ( t ) with an ODE , until they observe new xt1 at time t1 , when they use Equation 3 to update it : h̄t1 = ODESolve ( h0 , GRU-ODE , t0 , t1 ) , ht1 = GRUCell ( h̄t1 , xt1 ) . ( 4 ) Here , we will derive the flow version of GRU-ODE . If we rewrite Equation 3 by regrouping terms : ht = ht−1 + ( 1− zt ) ( ct − ht−1 ) , we see that GRU update acts as a single ResNet layer . Definition 1 . Let fz , fr , fc : Rd+1 → Rd be any arbitrary neural networks and let z ( t , h ) = α · σ ( fz ( t , h ) ) , r ( t , h ) = β · σ ( fr ( t , h ) ) , c ( t , h ) = tanh ( fc ( t , r ( t , h ) h ) ) , where α , β ∈ R and σ is a sigmoid function . Further , let ϕ : R → Rd be a continuous function with ϕ ( 0 ) = 0 and |ϕ ( t ) i| < 1 . Then the evolution of GRU state in continuous time is defined as : F ( t , h ) = h+ ϕ ( t ) ( 1− z ( t , h ) ) ( c ( t , h ) − h ) . ( 5 ) Theorem 1 . A neural network defined by Equation 5 specifies a flow when the functions fz , fr and fc are contractive maps , i.e. , Lip ( f· ) < 1 , and α = 25 , β = 4 5 . We prove Theorem 1 in Appendix A.3 by showing that the second summand on the right hand side in Equation 5 satisfies Lipschitz constraint making the whole network invertible . We also show that the GRU flow has the same desired properties as GRU-ODE , namely , bounding the hidden state in ( −1 , 1 ) and having the Lipschitz constant of 2 . Note that GRU flow ( Equation 5 ) acts as a replacement to ODESolve in Equation 4 . Alternatively , we can append xt to the input of fz , fr and fc , which would give us a continuous-in-time version of GRU . Coupling flow . The disadvantage of both ResNet flow and GRU flow is the missing analytical inverse . To this end , we propose a continuous-in-time version of an invertible transformation based on splitting the input dimensions into two disjoint sets A and B , A∪B = { 1 , 2 , . . . , d } [ 17 ] . We copy the values indexed by B and transform the rest conditioned on xB which gives us the coupling flow : F ( t , x ) A = xA exp ( u ( t , xB ) ϕu ( t ) ) + v ( t , xB ) ϕv ( t ) , ( 6 ) where u , v are arbitrary neural networks and ϕu ( 0 ) = ϕv ( 0 ) = 0 . We can easily see that this satisfies condition i ) , and it is invertible by design regardless of t [ 17 ] . Since some values stay constant in a single layer , we apply multiple consecutive transformations , choosing different partitions A and B . For all three models we can stack multiple layers F = F1 ◦ · · · ◦ Fn and still define a proper flow since the composition of invertible functions is invertible , and consecutive identities give an identity . We can think of ϕ ( including ϕu , ϕv ) as a time embedding function that has to be zero at t = 0 . Since it is a function of a single variable , we would like to keep the complexity low and avoid using general neural networks in favor of interpretable and expressive basis functions . A simple example is linear dependence on time ϕ ( t ) = αt , or tanh ( αt ) for ResNet flow . We use these in the experiments . An alternative , more powerful embedding consists of Fourier features ϕ ( t ) i = ∑ k αik sin ( βikt ) .

Summary: This work proposes directly modelling the solution trajectories of an ODE, circumventing the need for costly numerical ODE solvers. Specifically, they model the function which returns a solution to the initial value problem, and derive conditions this function must satisfy in order to implicitly induce a valid ODE.  Crucially, the authors explain that Neural ODE models are not very interpretable, unlike more traditional applications of ODEs. Since we aren’t able to use the dynamics for interpretability, it makes sense to not modely them explicitly, and instead model its integral (via IVP) directly. (This is more my interpretation than what the authors stated explicitly.)  The authors propose several architectures, including ResNet flow and GRU flow. The issue with these models is that these models can only be inverted approximately and require expensive iterative schemes. A third approach is to design architectures inspired by normalizing flows, which trades off some expressivity for greater efficiency. The authors also consider the approximation capabilities of these models.  Section 3 introduces specific models for time-series and density estimation, and Section 4 compares these models to baselines on the relevant tasks.