Learning Neural Event Functions for Ordinary Differential Equations

1 INTRODUCTION . Event handling in the context of solving ordinary differential equations ( Shampine & Thompson , 2000 ) allows the user to specify a termination criteria using an event function . Part of the reason is to introduce discontinuous changes to a system that can not be modeled by an ODE alone . Examples being collision in physical systems , chemical reactions , or switching dynamics ( Ackerson & Fu , 1970 ) . Another part of the motivation is to create discrete outputs from a continuous-time process ; such is the case in point processes and event-driven sampling ( e.g . Steinbrecher & Shaw ( 2008 ) ; Peters et al . ( 2012 ) ; Bouchard-Côté et al . ( 2018 ) ) . In general , an event function is a tool for monitoring a continuous-time system and performing instantaneous interventions when events occur . The use of ordinary differential equation ( ODE ) solvers within deep learning frameworks has allowed end-to-end training of Neural ODEs ( Chen et al. , 2018 ) in a variety of settings . Examples include graphics ( Yang et al. , 2019 ; Rempe et al. , 2020 ; Gupta & Chandraker , 2020 ) , generative modeling ( Grathwohl et al. , 2018 ; Zhang et al. , 2018 ; Chen & Duvenaud , 2019 ; Onken et al. , 2020 ) , time series modeling ( Rubanova et al. , 2019 ; De Brouwer et al. , 2019 ; Jia & Benson , 2019 ; Kidger et al. , 2020 ) , and physics-based models ( Zhong et al. , 2019 ; Greydanus et al. , 2019 ) . However , these existing models are defined with a fixed termination time . To further expand the applications of Neural ODEs , we investigate the parameterization and learning of a termination criteria , such that the termination time is only implicitly defined and will depend on changes in the continuous-time state . For this , we make use of event handling in ODE solvers and derive the gradients necessarily for training event functions that are parameterized with neural networks . By introducing differentiable termination criteria in Neural ODEs , our approach allows the model to efficiently and automatically handle state discontinuities . 1.1 EVENT HANDLING . Suppose we have a continuous-time state z ( t ) that follows an ODE dzdt = f ( t , z ( t ) , θ ) —where θ are parameters of f—with an initial state z ( t0 ) = z0 . The solution at a time value τ can be written as ODESolve ( z0 , f , t0 , τ , θ ) , z ( τ ) = z0 + ∫ τ t0 f ( t , z ( t ) , θ ) dt . ( 1 ) In the context of a Neural ODE , f can be defined using a Lipschitz-continuous neural network . However , since the state z ( t ) is defined through infinitesimal changes , z ( t ) is always continuous in ∗Work done while at Facebook AI Research . t. While smooth trajectories can be a desirable property in some settings , trajectories modeled by an ODE can have limited representation capabilities ( Dupont et al. , 2019 ; Zhang et al. , 2020 ) and in some applications , it is desirable to model discontinuities in the state . Bouncing ball example As a motivating example of a system with discontinuous transitions , consider modeling a bouncing ball with classical mechanics . In an environment with constant gravity , a Markov state for representing the ball is a combination of position x ( t ) ∈ R and velocity v ( t ) ∈ R , z ( t ) = [ x ( t ) , v ( t ) ] , dz ( t ) dt = [ v ( t ) , a ] , ( 2 ) where a is a scalar for acceleration , in our context a gravitational constant . To simulate this system , we need to be mindful that the ball will eventually pass through the ground—when x ( t ) ≤ r for some r that is the radius of the ball—but when it hits the ground , it bounces back up . At the moment of impact , the sign of the velocity is changed instantaneously . However , no such ODE can model this behavior because v ( t ) needs to change discontinuously at the moment of impact . This simple bouncing ball is an example of a scenario that would be ill-suited for a Neural ODE alone to model . In order to model this discontinuity in the state , we can make use of event functions . Event functions allow the ODE to be terminated when a criteria is satisfied , at which point we can instantaneously modify the state and then resume solving the ODE with this new state . Concretely , let g ( t , z ( t ) , φ ) be an event function with φ denoting a set of parameters . An ODE solver with event handling capabilities will terminate at the first occurrence when the event function crosses zero , i.e . time t∗ such that g ( t∗ , z ( t∗ ) , φ ) = 0 , conditioned on some initial value . We express this relationship as t∗ , z ( t∗ ) = ODESolveEvent ( z0 , f , g , t0 , θ , φ ) . ( 3 ) Note that in contrast to eq . ( 1 ) , there is no predetermined termination time . The time of termination t∗ has to be solved alongside the initial value problem as it depends on the trajectory z ( t ) . Nevertheless , ODESolveEvent strictly generalizes ODESolve since the event function can simply encode an explicit termination time and is reduced back into an ODESolve . The benefits of using ODESolveEvent lie in being able to define event functions that depend on the evolving state . Going back to the bouncing ball example , we can simply introduce an event function to detect when the ball hits the ground , i.e . g ( t , z ( t ) , φ ) = x ( t ) − r. We can then instantaneously modify the state so that z′ ( t∗ ) = [ x ( t∗ ) , − ( 1 − α ) v ( t∗ ) ] , where α is the fraction of momentum that is absorbed by the contact , and then resume solving the ODE in eq . ( 2 ) with this new state z′ ( t∗ ) . Figure 1 shows the bouncing ball example being fit by a Neural ODE and a Neural Event ODE where both f and g are neural networks . The Neural ODE model parameterizes a non-linear function for f while the Neural Event ODE parameterizes f and g as linear functions of z ( t ) . We see that the Neural Event ODE can perfectly recover the underlying physics and extrapolate seamlessly . Meanwhile , the Neural ODE has trouble fitting to the sudden changes in dynamics when the ball bounces off the ground , and furthermore , does not generalize because the true model requires the trajectory to be discontinuous . The Neural Event ODE , while being capable of modeling discontinuities in t , is a continuous function of the parameters and hence can be trained with gradient descent . Going forwards , we will discuss how to differentiate t∗ w.r.t . the variables that depend on it , such as z ( t∗ ) , φ and θ . Before this , we briefly summarize how gradients can be computed through any black-box ODE solver . 2 BACKGROUND : DIFFERENTIATING THROUGH ODE SOLUTIONS . Consider a scalar-valued loss function L that depends on the output of an ODE solver , L ( z ( τ ) ) = L ( ODESolve ( z0 , f , t0 , τ , θ ) ) ( 4 ) where f ( t , z , θ ) describes the dynamics . To optimize L , we require the gradients with respect to each of the inputs : z0 , t0 , τ and θ . All of these inputs influence the loss function through the intermediate states z ( t ) , for t ∈ [ t0 , τ ] , and their gradients can be expressed in relation to the adjoint state a ( t ) , dLdz ( t ) which contains the gradient of all intermediate states . The adjoint method ( see e.g . Pontryagin et al. , 1962 ; Le Cun , 1988 ; Giles & Pierce , 2000 ; Chen et al. , 2018 ) provides an identity that quantifies the instantaneous change in the adjoint state : da ( t ) dt = −a ( t ) T ∂f ( t , z ( t ) , θ ) ∂z ( 5 ) which when combined with z ( t ) is an ordinary differential equation that—by solving the adjoint state backwards-in-time , similar to a continuous-time chain rule—allows us to compute vectorJacobian products of the form vT [ ∂z ( τ ) ∂ξ ] , where ξ is any of the inputs z0 , t0 , τ , θ . For instance , with v = dLdz ( τ ) , the product dL dz0 = dLdz ( τ ) T dz ( τ ) dz0 effectively propagates the gradient vector from the final state , dLdz ( τ ) , to the intial state , dL dz0 . The ability to propagate gradients allows ODESolve to be used within reverse-mode automatic differentiation ( Baydin et al. , 2018 ) . We use the method in Chen et al . ( 2018 ) , which solves the adjoint state and parameter gradients jointly backwards-in-time alongside the state z ( t ) . This method does not require intermediate values of z ( t ) to be stored and only invokes ODESolve once for gradient computation . There exist other notable approaches for solving the adjoint equations with different memorycompute trade-offs , such as storing all intermediate quantities ( also known as discrete adjoint ) ( e.g . Zhang & Sandu 2014 ) , more sophisciated methods of checkpointing ( Chen et al. , 2016 ; Gholami et al. , 2019 ; Zhuang et al. , 2020 ) , the use of interpolation schemes ( Hindmarsh et al. , 2005 ; Daulbaev et al. , 2020 ) , and symplectic integration ( Zhuang et al. , 2021 ; Matsubara et al. , 2021 ) . Any of these approaches can be used and is tangential to our contributions . 3 DIFFERENTIATING THROUGH EVENT HANDLING . In an event-terminated ODE solve , the final time value t∗ is not an input argument but a function of the other inputs . As such , for gradient-based optimization , we would need to propagate gradients from t∗ to the input arguments of ODESolveEvent ( eq . ( 3 ) ) . Consider a loss function L that depends on the outputs of ODESolveEvent , L ( t∗ , z ( t∗ ) ) where t∗ , z ( t∗ ) = ODESolveEvent ( z0 , f , g , t0 , θ , φ ) . ( 6 ) Without loss of generality , we can move the parameters φ inside the state z0 and set dφ ( t ) dt = 0 . As long as we can compute gradients w.r.t z0 , these will include gradients w.r.t . φ . This simplifies the event function to g ( t , z ) . Furthermore , we can interpret the event function to be solving an ODE at every evaluation ( as opposed to passing the event function as an input to an ODE solver ) conditioned on the input arguments . This simplifies the event handling procedure to finding the root of groot ( t , z0 , t0 , θ ) , g ( t , z = ODESolve ( z0 , f , t0 , t , θ ) ) ( 7 ) and factorizes the ODESolveEvent procedure into two steps : { t∗ = arg mint t ≥ t0 such that groot ( t , z0 , t0 , θ ) = 0 z ( t∗ ) = ODESolve ( z0 , f , t0 , t∗ , θ ) . ( 8 ) It is obviously computationally infeasible to numerically solve an ODE within the inner loop of a root finding procedure , but this re-interpretation allows us to use existing tools to derive the gradients for ODESolveEvent which can be simplified later to just solving one ODE backwards-in-time . First , the implicit function theorem ( Krantz & Parks , 2012 ) gives us the derivative from t∗ through the root finding procedure . Let ξ denote any of the inputs ( z0 , t0 , θ ) . Then the gradient satisfies dt∗ dξ = − ( dgroot ( t ∗ , z0 , t0 , θ ) dt ) −1 [ ∂g ( t∗ , z ( t∗ ) ) ∂z ∂z ( t∗ ) ∂ξ ] . ( 9 ) Though groot requires solving an ODE , the derivative of z ( t∗ ) w.r.t . t∗ is just f∗ , f ( t∗ , z ( t∗ ) ) , so dgroot ( t ∗ , z0 , t0 , θ ) dt = ∂g ( t∗ , z ( t∗ ) ) ∂t + ∂g ( t∗ , z ( t∗ ) ) ∂z T f∗ . ( 10 ) Taking into account that the loss function may directly depend on both t∗ and z ( t∗ ) , the gradient from the loss function L w.r.t . an input ξ is dL dξ = ( ∂L ∂t∗ + ∂L ∂z ( t∗ ) T f∗ ) ︸ ︷︷ ︸ dL dt∗ ( −dgroot ( t ∗ , z0 , t0 , θ ) dt −1 [ ∂g ( t∗ , z ( t∗ ) ) ∂z T ∂z ( t∗ ) ∂ξ ] ) ︸ ︷︷ ︸ dt∗ dξ + ∂L ∂z ( t∗ ) T [ ∂z ( t∗ ) ∂ξ ] . ( 11 ) Re-organizing this equation , we can reduce this to dL dξ = vT [ ∂z ( t∗ ) ∂ξ ] ( 12 ) where v = ( ∂L ∂t∗ + ∂L ∂z ( t∗ ) T f∗ ) ( −dgroot ( t ∗ , z0 , t0 , θ ) dt −1 ∂g ( t∗ , z ( t∗ ) ) ∂z ) + ∂L∂z ( t∗ ) . All quantities in v can be computed efficiently since g and t∗ are scalar quantities . As they only require gradients from g , there is no need to differentiate through the ODE simulation to compute v. Finally , the vector-Jacobian product in eq . ( 12 ) can be computed with a single ODESolve . We implemented our method in the torchdiffeq ( Chen , 2018 ) library written in the PyTorch ( Paszke et al. , 2019a ) framework , allowing us to make use of GPU-enabled ODE solvers . We implemented event handling capabilities for all ODE solvers in the library along with gradient computation for event functions . Differentiable event handling generalizes many numerical methods that often have specialized methods for gradient computation , such as ray tracing ( Li et al. , 2018 ) , physics engines ( de Avila BelbutePeres et al. , 2018 ; Hu et al. , 2020 ) , and spiking neural networks ( Wunderlich & Pehle , 2020 ) .

Summary: This paper presents Neural Event ODEs, a method to extend Neural ODEs for modeling discontinuous dynamics in a continuous-time system.  Neural Event ODEs allow to learn termination criteria dependent on the system's state while being fully differentiable. Experiments on time series and temporal point processes validate the benefits of Neural Event ODEs on discountinuous dynamics.