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: The derivative function of an ODE defines a vector field over the input space (augmented with time). Neural ODEs learn such vector fields when solving initial value problems. The authors propose a method that directly learns the flow of such a vector field, by learning a function that respects the constraints of a proper flow. This method avoids the need to perform a costly integration operation like in NODEs while also performing similarly or better.