LEARNING NEURAL EVENT FUNCTIONS FOR ORDINARY DIFFERENTIAL EQUATIONS

Abstract

The existing Neural ODE formulation relies on an explicit knowledge of the termination time. We extend Neural ODEs to implicitly defined termination criteria modeled by neural event functions, which can be chained together and differentiated through. Neural Event ODEs are capable of modeling discrete and instantaneous changes in a continuous-time system, without prior knowledge of when these changes should occur or how many such changes should exist. We test our approach in modeling hybrid discrete-and continuous-systems such as switching dynamical systems and collision in multi-body systems, and we propose simulation-based training of point processes with applications in discrete control.

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 cannot 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-Côté 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 dz dt = f (t, z(t), θ)-where θ are parameters of f -with an initial state z(t 0 ) = z 0 . The solution at a time value τ can be written as ODESolve(z 0 , f, t 0 , τ, θ) z(τ ) = z 0 + τ 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.

