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 work provides an extension to the neural ODEs framework to include discrete changes (i.e. switching) in continuous-time dynamics. The authors provide a few examples of such systems (bouncing balls, collisions of particles, discrete control systems) and derive formally the gradients with respect to the unknown switching time (which is a solution to the so-called event function), where a discontinuity (the switch) happens. The authors implement their method in the torchdiffeq library of Chen et al. 2018, and provide an extensive experimental evaluation in this manuscript.