Learning Neural Event Functions for Ordinary Differential Equations
Authors: Ricky T. Q. Chen
Brandon Amos
Maximilian Nickel
Venue: ICLR 2021 Poster (2021)
Keywords: differential equations
implicit differentiation
point processes

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.

Published as a conference paper at ICLR 2021
LEARNING NEURAL EVENT FUNCTIONS FOR ORDINARY DIFFERENTIAL EQUATIONS
Ricky T. Q. Chen
Brandon Amos, Maximilian Nickel Facebook Al Research {bda, 8 com
University of Toronto; Vector Institute rtqichenêcs. toronto.cdu
ABSTRACT
The existing Neural ODE formulation relies on an explicit knowledge of the mination time. We extend Neural ODEs to implicitly defined termination crite- ria modeled by neural event functions, which can be chained together and dif- ferentiated through. Neural Event ODEs are capable of modeling discrete and instantaneous changes in continuous-time system, without prior knowledge of when these changes should occur 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 based simulation- based training of point processes with applications in discrete control.
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 introduce discontinuous changes system that cannot be modeled by an ODE alone. Examples being collision in physical systems, chemical reactions, or switching dynamics (Ackerson & Fu, create discrete outputs from continuous
1970). Another part of the motivation such is the case in point processes and event-driven sampling (e.g. Steinbrecher & Shaw (2008); Pe- ters et al. (2012); Bouchard- Côté al. (2018)). In general, event function is tool for monitoring a continuous-time system and performing instantaneous interventions when events occur.
process;
The use of ordinary differential equation (ODE) solvers within deep learning frameworks has al- lowed 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., 201 8; Zhang et al., 2018; 8; 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 chysics-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, investigate the parameterization and learning of termination criteria, such that the termination time is only implicitly defined and will depend on changes in the continuous-time state. For this, 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 da [(t,,(t), ()) 0 parameters of f-with
4t with an initial state z(ta) 2g- The solution at a time value can be written as
- =(7) = 20 + L: [(1,((1),)) dt.
ODESolve(zo, f, to T.
(1)
In the context Neural ODE, f can be defined using Lipschitz-contineen neural network. However, since the state zt) is defined through infinitesimal changes, 2(t) is always continuous in
Work done while Facehook AI Research.

---

Published as a conference paper at ICLR 2021
10.0
training region
extrapolation
7.5
x(O)
2.5
0.0 0
4
8
10
t
Neural ODE
Neural Event ODE
Truth
Figure 1: Dynamics of bouncing ball can recovered by a Neural Event ODE, Meanwhile, non-linear Neural ODE has trouble modeling sudden changes and performs poorly at extrapolation.
While smooth trajectories can 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 motivating example of system with discontinuous transitions, sider modeling a bouncing ball with classical mechanics. In an environment with constant gravity, a Markov state for representing the ball is combination of position æ=t) E R and velocity v(t) E R.
dz(t) [o(t), a),
(t) = [x(t), (()),
(2)
dt
where a is scalar for acceleration, in our context gravitational constant.
To simulate this system, we need to be mindful that the ball will eventually pass through the ground-when st)) S for some 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 instantancously. However, no such ODE can model this behavior because vtt) needs change discontinuously the moment of impact. This simple bouncing ball is an example of a scenario that would be ill-suited for 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 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, =t), Φ) be an event function with Φ denoting 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 gt",, xtt),, Φ) = 0. conditioned on some initial value. We express this relationship as
t²,2(t") = oossolveEvent
0,
(3)
Note that in contrast eq. (1), there is predetermined termination time. The time of termination t" has to be solved alongside the initial value problem as it depends on the trajectory 2(t). Nev- ertheless, ODESoiveEvent strictly generalizes ODESolve lvee since the event function can simply encode an explicit termination time and is reduced back into an ODESolve, The benefits of using ODEsolveEvens lie being able to define event functions that depend on the evolving state.
Going back to the bouncing ball example, we can simply introduce event function to detect when the ball hits the ground, í.e. g(t. 2(t), .) x(t) r. We can then instantaneously modify the state so that z'("") = [I=( -(1 a))) where a is the fraction of momentum that is absorbed by the contact, and then resume solving the ODE eq. (2) with this new state z((t*).
Figure shows the bouncing ball example being fit by Neural ODE and Neural Event ODE where both f and g are neural networks. The Neural ODE model parameterizes a non-linear function for while the Neural Event ODE parameterizes f and g as linear functions of =(t). We see that the Neural Event ODE can perfectly recover the underlying physics and extrapolate seamlessly. Mcanwhile, 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.
2

---

Published as a conference paper at ICLR 2021
The Neural Event ODE, while being capable of modeling discontinuities in is continuous func- tion 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 such as =(t"), and 8. Before this, we briefly summarize how gradients can be computed through any black-box ODE solver.
BACKGROUND: DIFFERENTIATING THROUGH ODE SOLUTIONS
Consider scalar-valued loss function that depends on the output an ODE solver,
L(z(r)) = ((DDESolve f. to T, 0))
(4)
where f(t, 2) describes the dynamics. To optimize L, we require the gradients with respect to each of the inputs: 20. to. т and 0. All of these inputs influence the loss function through the intermediate states zt),, for torte [to, 7), and their gradients can be expressed in relation to the adjoint state att)
dt which contains the gradient of all intermediate states. 25(2
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) a(()) - (1))))
(5)
dt
Oz
which when combined with z(1) an ordinary differential equation that-by solving the adjoint state nackwards-in-tine similar to continuous-time chain rule-allows to compute veclor- Jacobian products of the form oT [2(r))
DE where & is any of the inputs 20, lo T, 0. For instance, dL
with v = AL
d5(7). the product AL
d=(rr effectively propagates the gradient vector from the dāa
d5o d=(T)
The ability to propagate gradients allows ODESolve to be final state, ) to the intial state, dan
AL,
used within meverse-mode automatic differentiation (Baydin et al., 2018)
We use the method in Chen al. (2018), which solves the adjoint state and parameter gradients jointly backwards-in-time alongside the state 2f). This method does not require intermediate values of (t) be stored and only invokes ODESolve once for gradient computation.
There exist other notable approaches for solving the adjoint equations with different memory- compute trade-offs, such storing all intermediate quantities (also known as discrete adjoint) (e.g. Zhang & Sandu 2014), more sophisciated methods of checkpointing (Chen et al., 2016; Gholami al., 2019; Zhuang al., 2020), the use of interpolation schemes (Hindmarsh al., 2005; Daulbaev et al., 2020), and symplectic integration (Zhuang et al., 2021; Matsubara et al., 2021). Any of these approaches can used and is tangential our contributions.
3 DIFFERENTIATING THROUGH EVENT HANDLING
In event-terminated ODE solve, the final time value 1 is not input argument but function of the other inputs. As such, for gradient-based optimization, we would need to propagate gradients from 1a to the input arguments of CDESol veEvent (eq. (3)).
Consider a loss function that depends on the outputs of ODESolveEvent,
=(())((t")) where t", =(t") ODESo vet.eslllveeeeetto.....
(6)
Without loss of generality, we can move the parameters ø inside the state and set do(l) = 0. As long as we can compute gradients w.r.t zo. these will include gradients w.r.t. Φ. This simplifies the event function to g(t,
di
Furthermore, we can interpret the event function to solving an ODE at every evaluation (as opposed to passing the event function input to ODE solver) conditioned on the input argu- ments. This simplifies the event handling procedure to finding the root of
Froo (1,, 20, lo, A g(, = ODE.ol vee lo, 1))

---

Published as a conference paper at ICLR 2021
and factorizes the ODESolveEvent procedure into two steps:
tt = arg min, 2 to such that groot (t, 20, to. 0) =
(8)
(z(()) ODESolve(Zo, f. fo, t ()).
It is obviously computationally infeasible to numerically solve ODE within the inner loop of root finding procedure, but this re-interpretation allows us to use existing tools to derive the gradients for ODESo veEvent 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 (Zo, to 8). Then the gradient satisfies
dt
[greol (t", 20, to. 8)
0z(tt)
(9)
de
dt
Oz
Though Groot requires solving ODE, the derivative of z(t") w.r.t t" is just f "(t",((t')
dgroaa (t", 20, l)
ag(t",((())
dz f"
(10)
dt
Ot
Taking into account that the loss function may directly depend on both t and 2(t"), the gradient from the loss function L w.r.t. an input & is
dL = dg
ØT. OT. 5*) )z(t")
dgma (f", 20, to 0) -1
"i(t",z(("))
=(tt))"aa(t"")
aL TT2(t")
o0
dt
Oz
0C
0z(t*)
0c
d
(11)
Re-organizing this equation, can reduce this
dL
(z((t))
TT
(12)
de
O
where e
aL 0z
T fr) )z((")
(-
Alpoa (t" .0. 0) (- dt
(t))))
AL Dz((")'
0=
All quantities in can be computed efficiently since and t* are scalar quantities. As they only require gradients from g. there no need differentiate through the ODE simulation to compute e. Finally, the vector-Jacobian product in eq. (12) can be computed with single ODESolve.
We implemented our method in the torchdiffeq (Chen, 2018) library written in the Py- Torch (Paszke et al., 2019a) framework, allowing to make use of GPU-enabled ODE solvers. We implemented event handling capabilities for all ODE solvers the library along with gradient computation for event functions.
Diffcrentiable event handling generalizes many numerical methods that often have specialized meth- ods for gradient computation, such ray tracing (Li et al., 2018), physics engines (de Avila Belbute- Peres et al., 2018; Hu et al., 2020), and spiking neural networks (Wunderlich Pehle, 2020).
3.1 NEURAL EVENT FUNCTIONS AND INSTANTANEOUS UPDATES
This allows the of ODESo veEvent differentiable modeling primitive, available for usage in general deep learning or reverse-mode automatic differentiation libraries. We can then construct Neural Event ODE model that is capable of modeling variable number of state discontinuities, by repeatedly invoking ODESol veEvent. The Neural Event ODE is parameterized by a drift function, event function, and instantaneous update function that determines how the event is updated after each event. All functions can have learnable parameters. The model is described concretely in Algorithm 1 which outputs event times and the piecewise continuous trajectory 2(t).
GENERAL DTAEE-DEP EVENT FUNCTIONS
The state-dependent event function is the most general form and is suitable for modeling systems that have discontinuities when crossing boundaries the state space. We apply this to modeling switching dynamical system and physical collisions in multi-body system.

---

Published as a conference paper at ICLR 2021
Algorithm The Neural Event ODE. In addition to ODE, we also model a variable number of possible event locations and how each event affects the system.
Input: An initial condition (fo, zo), final termination time T, drift function f. event func- tion g. and an instantaneous update function h. Parameters of f, h. g are notationally omitted.
i 0
while t2 < T do
li+1, 2(41 ODESo ve
Solve until the next event D Instantaneously update the state
2i+1 = 1(1i+1,2(41)
i=i+1 end while
Return: event times [t:] and the piecewise continuous trajectory (=:(t) for f, St S 1411}
4.1 SWITCIIING LINEAR DYNAMICAL SYSTEMS
Switching linear dynamical systems (SLDS) are hybrid discrete and continuous systems that choose which dynamics to follow based on discrete switch (see e.g. Ackerson & Fu 1970; Chang & Athans 1978; Fox et al. 2009; Linderman et al. 2016). They are of particular interest for extracting interpretable models from complex dynamics in areas such neuroscience (Linderman al., 2016) or finance (Fox et al., 2009).
Here, we consider linear switching dynamical systems in the continuous-time setting, placing them the framework of event handling, and evaluate whether a Neural Event ODE can re- cover the underlying SLDD dynamics. At any time, the state consists of continuous vari- able = and a switch variable w. Representing w as a one-hot vector, the drift function is then da(t) E Wm (A(m), + ((m),, where M is the number of switch states. The switch variable w can change instantaneously, resulting in discontinuous dynamics at the moment of switching.
Setup We adapt the lincar switching dynamical system example of a particle moving around a synthetic race track from Linderman et al. (2016) to continuous-ting and use fan-shaped track. We created a data set with 100 short trajectories for training and 25 longer trajectories as validation and test sets. This was done by sampling trajectory from the ground-truth system and adding random noise. Figure 2 shows the ground-truth dynamics and sample trajectories. Given observed samples of trajectories from the system learn recover the event and state updates by minimizing the squared error between the predicted and actual trajectories. Further details are in App. A.
Results For our Neural Event ODE, we relax the switch state to be a positive real vector that sums to one. At event loca- tions, we instantancously change the switch state w based on the event location. We evaluated our model in comparison to a non-linear Neural ODE that doesn't 't use event handling and recurrent neural network (RNN) baseline. We show the test
Table 1: Continuous-tine switch-
ing linear dynamical system.
Model
Test Loss (MSE)
RNN (LSTM)
0.261 0.078
loss in table I, and the visualization in fig. 2 demonstrates that are able recover the components and can generate real- istic trajectories with sharp transitions. We note that since the
Neural ODE
0.157 + 0.005
Neural Event ODE
0.093 + 0.037
switch variable is only set event locations, they only need to be accurate at the event locations.
4.2 MODELING PHYSICAL SYSTEMS WITH COLLISION
We next consider the use of event handling for modeling physical collisions. Though there exist physics-motivated parameterizations of Neural ODEs (e.g. Zhong et al., 2019; Greydanus et al., 2019; Finzi et al., 2020), these cannot model collision effects. Event handling allows us to capture discontinuities in the state (from collisions) that otherwise would be difficult and unnatural to model with a Neural ODE alone. Building on the bouncing ball example from sect. 1.1, we consider a multi-body 2D physics setting with collisions.
5

---

Published as a conference paper at ICLR 2021
(a) Ground truth
(b) RNN (LSTM)
(c) Neural ODE
(d) Neural Event ODE
Figure 2: We learn continuous switching linear dynamical systems with Neural Event ODE to model and predict the trajectory of particle traveling on loop with discontinuous changes in dynamics.
Table Test results on the physics prediction task. @N notates evaluating on sequence length of N discretized steps. As proxy for the complexity of ODE models, we also report the number of function evaluations (NFE) for the ODE dynamics and the event function.
Test Loss (MSE)
Complexity (NFE) (x10*)
Model
K25
@50
@ @100
0DE
100
EventFn@50
EventFn 100
RNN (LSTTM)
+
+:
0.24
Neural ODE
0.00 + 0.00
0.06 +: 0.01
0.18 +: 0.02
3.81 + 0.24
7.82 +: 0.43
Neural Event ODE
0.01 + 0.00
0.07 +: 0.00
0.19 +: 0.02
0.16 + 0.00
0.35 + 0.01
0.16 0.01
0.36 + 0.01
Setup We use Pymunk/Chipmunk (Blomqvist, 2011; Lembeke, 2007) to create a data set of trajec- tories from simulating two balls colliding in box with random initial positions and velocities. As part of our model, we learn (i) a neural event function to detect collision either between the two balls between each ball and the ground; (ii) instantaneous update the state reflect change in direction due to contact forces. Both the event function and instantaneous updates are parameterized as deep neural networks. Further details and exact hyprpaarameters are described in App. A.
Results We evaluate our model in comparison to RNN baseline and a non-linear Neural ODE that doesn't use events. The non-linear Neural ODE is parameterized such that the velocity the change in position, but the change in velocity is non-linear neural network. We report the test loss in table 2, and the complexity of the learned dynamics. The Neural Event ODE generalizes better than RNN bascline and matches the non-linear Neural ODE. Though the non-linear Neural ODE can perform very well, must use a very stiff dynamics function to model sudden changes in velocity at collisions. In contrast, the Neural Event ODE requires much fewer function evaluations solve and can output predictions much faster.
sample trajectories in fig. 3. We find that the RNN and Neural ODE baselines learn to reduce the MSE loss of short bounces by simply hovering the ball in mid-air. On the other hand, the Neural Event ODE can exhibit difficulty recovering from early mistakes which can lead to chaotic behavior for longer sequences.
We illustrate
5 THRESHOLD-BASED EVENT FUNCTIONS
We next discuss a special parameterization of event functions that are based on an integrate-and- threshold approach. This event function is coupled with the state and depends integral. That is, let A(t) E (0) be positive quantity. Then a threshold-basee event occurs when the integral over À reaches some threshold s. The event time is thus the solution to
accumulated
t" such that [ X(t) dt
(13)
which can be implemented with CDESol by tracking A(t) f A(s)ds as part of the state z(t) and using the event function g(t, =(t) = = A(1).

---

Published as a conference paper at ICLR 2021
RNW Ground Truth
Neural ODE
(a)
(b)
(c)
(d)
(e)
Figure 3: Test-set motion predictions of two bouncing balls with collisions. The dashed black circle indicates the initial positions of the balls and the solid black circle indicates the final positions. The Neural Event ODE model is able to recover realistic event and state updates that exhibit floating behavior (a-c) and straighter paths (d-e) than RNN baseline. (c) shows a failure mode of the RNN and Neural Event ODE models where the first collision is not correctly predicted.
This form appears in multiple areas such as neuronal dynamics (Abbott, 1999; Hodgkin & Huxley, 1952), inverse sampling (Steinbrecher & Shaw, 2008) and more generally temporal point processes. We focus our discussion around temporal point processes as they encompass other applications.
Temporal point processes (TPPs) The TPP framework is designed for modeling random se- quences of event times. Let H = [4: be sequence of event times, with & IR and E Z+. Additionally, let H((t) = (4, 1, <, H}, i.e. the history of events predating time t. A temporal point process is then fully characterized by conditional intensity function X X+(t) A(t H(t)). The star superscript is a common shorthand used to denote conditional dependence on the history (Daley & Vere-Jones, 2003). The only condition is that (t) > 0. The joint log likelihood of observing H. starting with an initial time value at ln is
f X*(T) dr.
log p ((4+)
Σ log (*()))
(14)
In the context of flexible TPP models parameterized with neural networks, Mei & Eisner (2017) used a Monte Carlo estimate the integral in eq. (14), Omi al. (2019) directly parameterized the integral instead of the intensity function, and Jia & Benson (2019) noted that this integral can be computed using an ODE solver. While these approaches can enable training flexible TPP models by maximizing log-like likelihooo, it is much less straightforward leam from simulations.
In the following, we discuss how the event function framework allows us to backpropagate through simulations of TPPs. This enables training TPPs with the "reverse KI." objective. Another form

---

Published as a conference paper at ICLR 2021
30
REINFORCE Reparam gradient
REINFORCE
Reparam gradient
Reverse KL
20
Roverse
5-
10
0-
200
400
600
800
1000
1200
1400
1600
200
400
600
800
1000
1200
1400
1600
Iteration
(a) Five events
(b) Ten events
10000
REINFORCE Reparam gradient
8000
Deterministic
5000
6000
Reward
Reparam gradient
2000
-5000
Detemministic
100
200
350
150 Iteration on
Iteration
(c) HIV dynamies model
(d) flodgkin-Huxldd dynamics model
Figure 4: Learning TPPs (a.b) through reverse KL and (c.d) for discrete control in continuous time. Our reparameterization gradient improves upon the REINFORCE gradient and we can even learn deterministic discrete-valued policy.
of simulation- based training appears in reinforcement learning, where a TPP policy can be used to perform instantaneous interventions on continuous-time environment. Our method can also be readily applied to extensions such as spatio-temporal point processes (Chen et al., 2021).
5.1 REPARAMETERIZATION GRADIENT FOR TEMPORAL POINT PROCESSES
Sampling single event can be done in two steps: sample s Exp(1) and (ii) solve for t" such that _e X (t)dt, which is exactly of the form eq. (13). This allows to reparameterize the sample t" as a transformation of noise variable s, thus allows us to take gradients of samples with respect to parameters of the TPP. Consider conditional intensity function >ă parameterized by 8,
V (4) [f([t+] = (s.)--ss(() [Vof([t((s;.))) 0)}}
(15)
provides the reparameterization gradient, where (s,, is the solution step (ii) above, and can be implemented as an ODESo veEvent which is differentiable.
We compare the reparameterization gradient against the REINFORCE gradient (Sutton al., 2000) in training Neural Jump SDE (Jia & Benson, 2019) with reverse KL objective
DxL (po, Pharget) EqG,iiPP(M [log pe ((:)) log Parger (ommmee((())) ((4)))
(16)
where Purget is taken be Hawkes point process. The reparameterization gradient of this objective requires simulating from the model and requires taking gradients through the event handling proce- dure. Results are shown in figs. 4a and 4b. The REINFORCE gradient is generally perceived have a higher variance and slower convergence, which is also reflected in our experiments. The reparam- eterization gradient performs well on sequences of either five or ten events, while the REINFORCE gradient exhibits slower convergence with longer event sequences.
Learning discrete control over continuous- time systems We show that the reparameterization gradient also allows us model discrete-valued control variables in continuous-time system using Neural Jump SDEs as control policies. We use a multivariate TPP, where sample from a particular dimension, or index, changes the control variable to discrete value corresponding to that index.
We test using the human immunodeficiency virus (HIV) dynamics model from Adams et al. (2004) which simulates interventions from treatment strategies with inhibitors. We use discrete control
8

---

Published as a conference paper at ICLR 2021
Action Potential
Input Current
80
60
60
20
Time
Time
Figure 5: Visualization of discrete-valued control (left) in continuous-time environment (right), using the Hodgkin-Huxiey model of neuronal dynamics.
model for determining whether each of two types of inhibitors should be used, resulting in 4 discrete states, similar the setup in Miller al. (2020). The reward function we use is the same as that of Adams et al. (2004). We additionally experiment with the Hodgkin-Huxley neuromal dynamics model (Hodgkin & Huxley, 1952), where we have an input stimulus that can switch between being on or off at any given time. The reward is high if the action potentials from the Hodgkin Huxley model match that of a target action potential. See fig. 5. All experiments were run with three seeds. Detailed setup in App. A.
Results are shown in figs. 4c and 4d, which shows that the reparameterization gradient outperforms REINFORCE in both settings. Interestingly, we can also train with a deterministic control pol- icy: instead of randomly sampling 8; we fix these threshold values. This deterministic policy can outperform both stochastic policies as the underlying continuous-time system is deterministic.
SCOPE AND LIMITATIONS
Minibatching Each trajectory minibatch can have different event locations. However, our ODE solvers do not have an explicit notation of independence between systems in minibatch. For batching, we can combine minibatch of ODEs into one system and then use a sparse aggregator (e.g. min or max operator) on the event functions. This does introduce some overhead we would restart integration for the combined system whenever one event function the minibatch triggers.
Runaway event boundaries In some cases, neural event function is never triggered. In such cases, there will be no gradient for the event function parameters and the model degenerates to Neural ODE alone. In high dimensions, the roots of a neural network are in general unpredictable and trajectory may never encounter root. We have found that initializing the parameters of the event function with a large standard deviation helps alleviate this problem for general neural event functions. This is less of problem for the threshold- always positive and there will always be an event the state is solved for long enough.
event functions, the integrand is
Numbers events discrete An objective function that relies on the number events can be discontinuous with respect to parameters of the event function, small change in parameter space can introduce a jump in the number events. Ultimately, this is a matter of the formulation which gradients are used. For instance, in the stochastic TPP setting, the reparameterizzttin gradient relies on the function being differentiable (Mohamed et al., 2019). For this reason, Shchur et al. (2020) used surrogate objectives for sampling-basse based training.
CONCLUSION
We consider parametrizing event functions with neural networks the context of solving ODEs, extending Neural ODEs to implicitly defined termination times. This enables modeling discrete events in continuous-timg systems-e, egg. g. the criteria and effects collision in physical systems- and simulation- based training of temporal point processes with applications to discrete control. No- tably, we can train deterministic policies with discrete actions, by making use of gradients that only exist continuous-ting setting.

---

Published as a conference paper at ICLR 2021
ACKNOWLEDGEMENTS
We thank David Duvenaud for helpful discussions. Additionally, we acknowledge the Python com- munity (Van Rossum & Drake Jr, 1995; Oliphant, 2007) for developing the core set of tools that en- abled this work, including PyTorch (Paszke et al., 2019b), torchdiffeg (Chen, 2018), higher (Grefen- stette al., 2019), Hydra (Yadan, 2019). Jupyter (Kluyver et al., 2016), Matplotlib (Hunter, 2007), seaborn (Waskom et al., 2018 8), numpy (Oliphant, 2006; Van Der Walt et al., 201 1), pandas (McK- inney, 2012), and SciPy (Jones al., 2014).
REFERENCES
Larry F Abbott. Lapicque's introduction of the integrate-and-firg model neuron (1907). Brain research bulletin, 50(5-6):303-304, 1999.
Ackerson and Fu. On state estimation in switching environments. IEEE transactions on auto- matic control, 15(1):10-17 1970.
Brian M Adams, Harvey T Banks, Hee-Dae Kwon, and Hien T Tran. Dynamic multidrug therapies for hiv: Optimal and sti control approaches. Mathematical Biosciences & Engineering, 1(2):223, 2004.
Atilim Gunes Baydin, Barak A. Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic differentiation machine learning: survey. Journal of Machine Learning Research, 18(153):1-44 2018. URL http: http::/ppp/.ll....or/////e//////
html.
Blomqvist, pymunk wttp:///www...ymntt....////////a
Alexandre Bouchard-CGil Sebastian J Vollmer, and Arnaud Doucet. The bouncy particle sam- pler: A nonreversible rejection-free markov chain monte carlo method. Journal of the American Statistical Association, 113(522):855-86)) 2018
Chaw-Bing Chang and Michael Athans. State estimation for discrete systems with switching pa- rameters, IEEE Transactions on Aerospace and Electronic Systems, (3):418-425, 1978.
Ricky T. Q. Chen. torchdiffeq, 2018.
URL
com/ wwtom//tticher...
torchdi i ffeq.
Ricky T. Q. Chen and David Duvenaud. Neural networks with cheap differential operators. In Advances in Neural Information Processing Systems, 2019.
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. In Advances in neural information processing systems, pp. 6571-6583, 2018.
Ricky T. Q. Chen, Brandon Amos, and Maximilian Nickel. Neural spatio-temporal point processes.
In International Conference on Learning Representations, 2021.
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv: 1614.660444444 2016.
Daryl Daley and David Vere-Jones. An introduction to the theory of point processes, volume I: Elementary theory and methods. Verlag New York Berlin Heidelberg: Springer, 2003.
Talgat Daulbaev, Alexandr Katrutsa, Larisa Markeeva, Julia Gusak, Andrzej Cichocki, and Ivan
Oseledets. Interpolated adjoint method for neural ODEs. arXiv preprint arXiv:2003.0527/ 2020.
Filipe de Avila Belbute-Peres, Kevin Smith, Kelsey Allen, Josh Tenenbaum, and Zico Kolter. End-to-end differentiable physics for learning and control. In S. Ben- gio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnet (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Asso- ciates, Inc., 2018. URL https: htttps::/ppooceeeiee weweeiinnnn.nnuns.seeeeee..e nwwrips..cc.... htt.e//p.ccc/paaerr///////////// 842424a1d05 90095055566000 6ec4 1ac10/099000eeeo0000eee665555---
pdf.
10

---

Published as a conference paper at ICLR 2021
Edward De Brouwer, Jaak Simm, Adam Arany, and Yves Moreau. GRU-ODE-bayes, Continuous modeling of sporadically-obsere time series. In Advances Neural Information Processing Systems, pp. 7379-7390, 2019.
Emilien Dupont, Arnaud Doucet, and Yee Whye Teh. Augmented neural ODEs. In Advances in Neural Information Processing Systems, pp. 3140-3150, 2019.
Mare Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolu- tional neural networks for equivariance to lie groups arbitrary continuous data. arXiv preprint arXiv:2002. 12800, 2020.
Emily Fox, Erik B Sudderth, Michael I Jordan, and Alan S Willsky. Nonparametric Bayesian learn- ing of switching linear dynamical systems. In Advances neural information processing systems, pp. 457-464, 2009.
Wulfram Gerstner, Werner M Kistler, Richard Naud, and Liam Paninski. Neuronal dynamics: From single neurons to networks and models of cognition. Cambridge University Press, 2014.
Amir Gholami, Kurt Keutzer, and George Biros. Anode: Unconditionally accurate memory-efficient gradients for neural odes. arXiv preprint arXiv: 2019.
Michael B Giles and Niles A Pierce. An introduction to the adjoint approach to design. Flow, furbulence and combustion, 65(3-4):393-455, 2000.
Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. FFJORD: Free-form continuous dynamics for scalable reversible generative models. arXiv
preprint arXiv:1810.01367 2018.
Edward Grefenstette, Brandon Amos, Denis Yarats, Phu Mon Htut, Artem Molchanov, Franziska Meier, Douwe Kiela, Kyunghyun Cho, and Soumith Chintala. Generalized inner loop meta- learning. arXiv preprint 11X1v:1910.01727, 2019.
Samuel Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian neural networks. In Advances in Neural Information Processing Systems, pp. 15379-15389, 2019.
Kunal Gupta and Manmohan Chandraker. Neural mesh flow: 3d manifold mesh generationvia diffeomorphic flows. arXiv preprint arXiv:2007.10973, 2020.
Alan C Hindmarsh, Peter Brown, Keith E Grant, Steven Lee, Radu Serban, Dan E Shumaker, and Carol S Woodward. SUNDIALS: Suite of nonlinear and differentialValgebraie equation solvers. ACM Transactions on Mathematical Software (TOMS), 31(3):633-396, 2005.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735-1780, 1997.
Allan Hodgkin and Andrew F Huxley. The components membrane conductance in the giant axon of loligo. The Journal of physiology, 1 116(4):473, 1952.
Yuanming Hu, Luke Anderson, Tzu-Mao Li, Qi Sun, Nathan Carr, Jonathan Bagan-Kelley, and Fredo Durand. Difftaichi: Differentiable programming for physical simulation. In International Conference on Learning Representations, 2020.
John D Hunter. Matplotlib: A 2d graphics environment. Computing in science & engineering, 9(3): 90, 2007.
Junteng Jia and Austin R Benson. Neural jump stochastic differential equations. In Advances in Neural Information Processing Systems, pp. 9847-9858, 2019.
Eric Jones, Travis Oliphant, and Pearu Peterson. {SciPy): Open source scientific tools for (Python). 2014.
Patrick Kidger, James Morrill, James Foster, and Terry Lyons. Neural controlled differential equa- tions for irregular time series. arXiv preprint wrXiv:2005.08926, 2020.

---

Published as a conference paper at ICLR 2021
Thomas Kluyver, Benjamin Ragan-Kelley. Fernando Pérez, Brian E Granger, Matthias Bussonnier, Jonathan Frederic, Kyle Kelley, Jessica B Hamrick, Jason Grout, Sylvain Corlay, et al. Jupyter notebooks-a publishing format for reproducible computational workflows. In ELPUB, pp. 87-90, 2016.
Steven G Krantz and Harold Parks. The implicit function theorem: history, theory, and applica- fions. Springer Science & Business Media, 2012.
Yann Le Cun. A theoretical framework for back propaaaroa..tion. 1988.
S Lembeke. Chipmunk physics httt:://chh.oouuu--ppyyiic. wtt.ssissieeeeeeeeeeeeeeeeeeeeee 2007.
Tzu-Mao Li, Miika Aittala, Frédo Durand, and Jaakko Lehtinen. Differentiable Monte Carlo ray tracing through edge sampling. ACM Transactions on Graphics (TOG), 37(6):1-11, 2018.
Scott W Linderman, Andrew C Miller, Ryan Adams, David M Blei, Liam Paninski, and Matthew Johnson. Recurrent switching linear dynamical systems. arXiv preprint arXiv:1610.08400 2016.
Takashi Matsubara, Yuto Miyatake, and Takaharu Yaguchi. Symplectic adjoint method for exact gradient neural ode with minimal memory. arXiv preprint arXiv:2102.0977 2021.
Wes McKinney. Python for data analysis: Data wrangling with Pandas, NumPy, and IPython. O'Reilly Media, Inc.", 2012.
Hongyuan Mei and Jason Fisner. The Neural Hawkes Process: A neurally self modulladug multivari- ate point process. Advances Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, 4-9 December 2017, Long Beach, CA, USA, pp. 6754-6764, 2017.
John Miller, Chloe Hsu, Jordan Troutman, Juan Perdomo, Tijana Zrnic, Lydia Liu, Yu Sun, Ludwig Schmidt, and Moritz Hardt. Whynot, 2020. URL httpp:://a...oog/1..... 3875775.
Shakir Mohamed, Mihaela Rosca, Michael Figurnov, and Andriy Mnih. Monte carlo gradient esti- mation in machine learning. arXiv preprint arXiv:1906.00 10652, 2019.
Travis E Oliphant. A guide to NumPy, volume 1. Trelgol Publishing USA, 2006.
Travis E Oliphant. Python for scientific computing. Computing in Science & Engineering, 9(3): 10-20, 2007.
Takahiro Omi, Kazuyuki Aihara, al. Fully neural network based model for general temporal point processes. In Advances in Neural Information Processing Systems, pp. 2122-2132, 2019.
Derek Onken, Samy Wu Fung, Xingjian Li, and Lars Ruthotto. OT-Flow: Fast and accurate contin- uous normalizing flows via optimal transport. arXiv preprint arXiv:2006.00104 2020.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. H. Wallach, H. Larochelle, A. Beygelzimer, F, d' d'Alché-Buc, E. Fox, and R. Gameet (eds.), Advances in Neural Information Processing Systems 32, pp. 8024 8035-8035. Curran Associates, Inc., 2019a.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. Advances in neural information processing systems, pp- 8026-8037, 2019b.
Elias AJF Peters al. Rejection-free monte carlo sampling for general potentials. Physical Review E, 85(2)30 2012.
Lev Semenovich Pontryagin, EF Mishchenko, VG Boltyanskii, and RV Gamkrelidze. The mathe- matical theory of optimal processes. 1962.
12

---

Published as a conference paper at ICLR 2021
Davis Rempe, Tolga Birdal, Yongheng Zhao, Zan Gojeic, Srinath Sridhar, and Leonidas Guibas. CaSPR: Learning canonical spatiotemporal point cloud representations. arXiv preprint arXiv:2008.02772,, 2020.
Yulia Rubanova, Ricky T. Chen, and David Duvenaud. Latent ODEs for iregularly-samppld time series. arXiv preprint arXiv:1907.088:: 1907.03907 2019.
LF Shampine and S Thompson. Event location for ordinary differential equations. Computers & Mathematics with Applications, (9(556))43-5) 2000.
Olcksandr Shehur, Nicholas Gao, Marin Biloš, and Stephan Günnemann. Fast and flexible temporal point processes with triangular maps. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
György Steinbrecher and William T Shaw. Quantile mechanics. 2008.
Richard S Sutton, David A McAllester, Satinder P Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. Advances in neural informa- lion processing systems, pp. 1057-1063, 2000.
Stefan Van Der Walt, S Chris Colbert, and Gael Varoquaux. The numpy array: structure for efficient numerical computation. Computing in Science & Engineering, 13(2)22, 2011.
Guido Van Rossum and Fred L Drake Jr. Python reference manual. Centrum voor Wiskunde en Informatica Amsterdam, 1995.
Michael Waskom, Olga Botvinník, Drew O'Kane, Paul Hobson, Joel Ostblom, Saulius Lukauskas, David C Gemperline, Tom Augspurger, Yaroslav Halchenko, John B. Cole, Jordi Warmenhoven, Julian de Ruiter, Cameron Pye, Stephan Hoyer, Jake Vanderplas, Santi Villalba, Gero Kunter, Eric Quintero, Pete Bachant, Marcel Martin, Kyle Meyer, Alistair Miles, Yoav Ram, Thomas Brunner, Tal Yarkoni, Mike Lee Williams, Constantine Evans, Clark Fitzgerald, Brian, and Adel Qalich. mwaskom/seaborn: v0.9.0 (july 2018), July 2018. URL attps://doii org/10.
5281/zenode 1313201.
Timo C Wunderlich and Christian Pehle. Eventprop: Backpropagation for exact gradients in spiking neural networks. arXiv preprint wxXiv:2009008376 2020.
Omry Yadan. Hydra a framework for elegantly configuring complex applications. Github, 2019.
URL https::pp///s /ttt://gttttt git. com/
Guandao Yang, Xun Huang. Zekun Hao, Ming- Yu Liu, Serge Belongie, and Bharath Hariharan. Pointflow: 3d point cloud generation with continuous normalizing flows. In Proceedings of the IEEE International Conference Computer Vision, pp. 4541-4550, 2019.
Han Zhang, Xi Gao, Jacob Unterman, and Tom Arodz. Approximation capabilities of neural ODEs and invertible residual networks. International Conference on Machine Learning. 2020.
Hong Zhang and Adrian Sandu. FATODE: a library for forward, adjoint, and tangent linear integra- tion of ODEs. SIAM Journal on Scientific Computing, 36(5):C504-C523 C523, 2014.
Linfeng Zhang, Lei Wang, et al. Monge-Ampére flow for generative modeling. arXiv preprint arXiv:1809. 100.188880.1888. 2018.
Yaofeng Desmond Zhong, Biswadip Dey, and Amit Chakraborty. Symplectic ODE-net: Learning hamiltonian dynamics with control. arXiv preprint 1rXiv:1109.12007 2019.
Juntang Zhuang, Nicha Dvoreck Xiaoxiao Li, Sekhar Tatikonda, Xenophon Papademetris, and James Dunean. Adaptive checkpoint adjoint method for gradient estimation in neural ODE, arXiv preprint arXiv:2006.02499, 2020.
Juntang Zhuang, Nicha C Dvomek. sekhar tatikonda, and James Duncan. MALI: A memory effi- cient and reverse accurate integrator for neural ODEs. In International Conference on Learning Representations, 2021.
13

---

Published as a conference paper at ICLR 2021
A
EXPERIMENTAL DETAILS
Continuous-time Switching Linear Dynamical Systems
DATA We constructed a fan-shaped system, similar basehall track. Let A be the rotation matrix
0
1 0
(17)
-1
Then the ground truth dynamical system followed
zA+ [o -1 -1] 1
2]
if 2
dx
if xg 2 and 21 < 2.
(18)
if xp < 0 and ø < 2
We discretized 100 sequences length 4 into 50 discrete steps for training. For each of validation and test, discretized 25 sequences of length 12 into 150 discrete steps. We added independent
Gaussian noise with standard deviation 0.05 to training sequences.
ARcurtECturE We set our dynamics function to be a weighted form of eq. (18), where the weights did not change over time and were only modified at event times. We parameterize the event function to be the product of two tanh-gated outputs of linear function.
EventFn(x)) = II ttan (wen((W [Wa x + 6)]s
(19)
And we parameterized the instantaneous change weights neural network with 2 hidden lay- ers, each with 1024 hidden units, with ReLU activation functions. The output is passed through softmax to ensure the weights
to one.
BASELINES We modeled the drift for the non-linear Neural ODE baseline MLP with 2 hidden layers with 256 hidden units each and the ReLU activation function. The LSTM baseline has a hidden state size of 128, and the hidden state is passed through MLP decoder (with 2 hidden layers of 128 units each and the ReLU activation) to predict the displacement in position between each also tested, but could not generalize well.
time step. LSTMs that directly predicts the position
TRAINING We use mean squared loss on both the position and the change in position. Let I I,..., N, be the ground truth positions and 3y the model's predicted positions. The training objective is then
N
N-1 Σ t=1
Σ
I
+
(---1-00 xx)- (Te+1 - -,))²
(20)
N
N
where we tried À E {0, 0.00. and chose À 0.01 for all models as the validation was lower. For validation and test, we only used the mean squared error on the positions I₁. For optimization, we used Adam with the default learning rate of 0.001 and a cosine learning decay. All models were trained with batch size of 1 for 25000 iterations.
Modeling Physical Systems with Collision
(Blomqvist, 2011; Lembeke, 2007) library to simulate two balls of radius 0.5 in a [0. 5]2 box. The initial position is randomly sampled and the initial velocity is We then simulated for 100 steps. We sampled 1000 initial positions for training and 25 initial positions each for validation and test.
DATA We used the Symunk/Chhppmuni
ARCHITECTURE We parameterized the event function as deep neural network. The outputs are then passed through a tanh and then multiplied together to form single scalar. The event function took as input the positions the two balls.
EventFn(())
II (tan((MLP())))
(21)
where M 8. The neural network is a multilayer perceptron (MLP) with hidden layer of 128 hidden units. We parameterized the instantaneous update function to be MLP with hidden layers
14

---

Published as a conference paper at ICLR 2021
with 512 hidden units. The instantaneous update took input the position, velocity, and the pre-tanh outputs of the neural event function. We used relative and absolute tolerances of 1E-8 for solving the ODE.
BASELINES The recurrent neural network baseline uses the LSTM architecture (Hochreiter & Schmidhuber, 1997) and outputs the difference position between time steps, which we found to be more stable than directly predicting the absolute position. The non-linear Neural ODE base- line uses the domain knowledge that velocity is the change in position. We then used a MLP with hidden layers and 256 hidden units model the instantaneous change in velocity.
TRAINING We used mean squared loss on the position of the two balls. For optimization, we used Adam with learning rate 0.0005 for the event function and 0.0001 for the instantancous update. We also clipped gradient norms at 5.0. All models were trained for 1000000 iterations, where each iteration used subsequence of 25 steps as the target.
Discrete Control over Continuous-time Systems
HIV DYNAMICS We used the human immunodeficiency virus (HIV) model of Adams et al. (2004), which describes the dynamics of infected cells depending on treatments representing reverse tran- scriptase (RT) in hibitors and protease inhibitors (PIs). We used viscre treatment control as was done by Miller et al. (2020) where RT could either be used (with a strength of 0.7) or unused, and PI could either be used (with strength of 0.3) or unused. Thus this results in discrete control variable with k = 4 states, The reward is the same as in Adams et al. (2004).
HODGKIN- HUXLEY NEURONAL DYNAMICS This model describes the propagation of action po- tentials in neurons (Hodgkin & Huxley, 1952: Gerstner al., 2014). We discrete valued input current as the control variable, with values 0 or 100. This then stimulates an action potential fol- lowing the Hodgkin- Huxley model. We use specfic action potential pattern the target A', and train a control policy that can recover this action potential. The reward is set to an integral over the infinitesimal negative squared error at each time value.
Reward 1'M (((A()-A'())) dt
(22)
where b = 600 acts a baseline for the REINFORCE gradient estimator. Figure shows the target action potential A and the input current used to generate it.
CONTROL/ / POLICY We used a multivariate temporal point process (TPP) with & dimensions where a sample from each dimension instantaneouslly changes the control variable to the discrete state corresponding to that dimension. The conditional intensity function of the TPP was parameterized as a function continous timm hidden state of width 64. The continuous-timg hidden state followed neural ODE with 3 hidden layers and sine activation functions. The hidden state is instantancously modified by a MLP with 2 hidden layers and sine activation functions. The conditional intensity function is function of the hidden state and is modeled by MLP with I hidden layer and ReLU activation. All hidden dimensions were 64.
TRAINING All models were trained using Adam with learning rate 0.0003.
15


Reviews:
[{"review_id": "zbOhuMKpDye", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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. \n\nThe paper is well written and relatively easy to follow. \nThe benefits that Neural Event ODEs provide for modeling discontinuous dynamics already becomes apparent from the formulation of its ODE solver. The simplicity of the approach is another advantage, and I can see many possible applications/use cases that can benefit from such an ODE solver.\n\nNevertheless, there were a few questions and remarks I had when reading the paper:\n\n- In experiment 1, the setup of the Neural ODE is not clear. In particular: \n - How many switch variables have you defined? In case of 3, have you tested what happens if you specify more (e.g. 4-6) to see how sensitive the model is to this parameter?\n - In the result paragraph, a classifier is mentioned that probably should refer to a classifier over $w$. Is this classifier trained by having a weighted average in $f$, i.e. $\\frac{d z(t)}{dt}=\\sum_{w=1}^{M} p(w)\\left[A^{(w)}z+b^{(w)}\\right]$, or how is the exact setup?\n\n- Table 2 evaluates the bouncing ball collision experiment on the mean squared error of the predicted trajectories. Based on those scores, the Neural Event ODE performs slightly worse than Neural ODE which is not clearly discussed in the section. Is the conclusion of Neural Event ODE being able to generalize better in this experiment based on qualitative results? If so, wouldn't have an adversarial metric reflected this result better? Besides the region close to the ground, it is hard to tell which model \\textit{generalizes} better.\n\n- In section 6, the issue of minibatching is mentioned. How does this effect the results of Neural Event ODEs shown in section 4 and 5.1? Have those models been trained with batch size 1 in contrast to the other methods, or using methods like gradient accumulation of multiple single-batches (hence only time of training increased)?\n\n- Most experimental setups in the paper clearly require the modeling of discountinuous dynamics. Have you tested Neural Event ODEs on tasks where the system to model does not have obvious discontinuities, for instance flow-based generative modeling? Would you see any potential advantages of your method there?\n\n- Experimental details such as the concrete parameterization of $f$ and hyperparameter values is missing for reproducibility of the results. No supplementary materials have been submitted in which this information could have been outlined.\n\nOverall, I think that the ideas presented in the paper are a valuable contribution to the research community and therefore, I would recommend the paper for acceptance, especially if the points mentioned above can be clarified.\n\nAdditional comments: Page 4, third line from the bottom has a double punctuation.", "extra_scores": {}, "extra_text": {}}, {"review_id": "e_7bL36pQ2", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: The paper provides a follow up approach to the Neural ODE [Chen et al., 2019] approach, where the termination time is implicitly defined by neural event functions rather than being pre-determined -- which is a promising direction for dealing with state discontinuities. The proposed method is tested on simple dynamical and two-body systems. \n\nStrengths: \n + Paper is well written \n + Promising direction for dealing with state discontinuities\n + Code will be released upon publication\n \nWeak Points: \n\n1) IMO, an explicit algorithm description should be included in the manuscript.\n\n2) I found the experimental part of the paper to be weak. Overall lots of model details are missing, comparison to baselines should include extra analysis and the approach should be tested on more challenging domains. More specifically,\n\na) When comparing to LSTMs and Neural ODEs, only sample trajectories and test MSE loss are presented. A precise description of the comparison baselines is missing (number of layers, hidden sizes, optimization setup, etc). An analysis comparing the number of parameters of each baseline would also be interesting. What is the odesolve used? Lots of missing details.\n\nb) How does the number of iterations on the ODE solve behave wtr to the errors? Do we see a similar behavior as in Neural ODE, where the number of function evaluations that the ODE solver makes increases along training? \n\nc) in the collision example, table 2 indicates neural ODE has similar (if not better) performance than the proposed approach. What happens when you have more bouncing balls?\n\nd) In the temporal point process examples from Figure 4, I think a more detailed problem overview (even if on the appendix) would help to understand the relevance of the results. \n\ne) Have you tried this approach on normalizing flows? \n\nf) In Figure 2, authors should explain the color code used. \n\nSo although I like the general research direction this paper takes, I think the experimental setup needs to be extended significantly -- which would add up to a very different submission, so at this time in my opinion it's not passing the ICLR acceptance threshold.  \n\n_______________________________________________________________________________________________________________________________________________________\nUPDATE: After reading the rebuttal I think most of my concerns have been addressed and I am updating my score accordingly.  ", "extra_scores": {}, "extra_text": {}}, {"review_id": "dlCeSAMEKVI", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary\n--\n\nThis authors extend neural ODEs to implicitly defined termination criteria modelled by 'neural event functions'. This allows neural ODEs to model abrupt changes in the dynamics (such as collisions or switching dynamics). The authors present how the even handling can be differentiated through, and include a representative set of example studies in the experiments.\n\nThe paper is well written, easy to follow, and presents relevant background and related work. The idea extends the already well-studied methodology related t neural ODE models, and even if the idea behind the contribution is small, the paper should be of interest for the audience of ICLR. The paper reads well, and the examples and illustrations help understanding.\n\nI liked this paper and found it interesting and well presented. The concerns below are reflected in my score.\n\n\nConcerns\n--\n\n1. The main reason for my low positive score is that the originality of the idea is rather low, and I view the paper mostly as an incremental addition to the toolbox of neural ODE methods. Nevertheless, I found the paper interesting and worthy of publication. The well-chosen experiments and illustrative examples further help communicate the idea, which I see as a strong point of this paper.\n\n2. In the experiments, it would have been interesting to see comparisons to not purely learning based models. This applies especially to the collision example, where a more classical method (which would most likely require more hand-tailored prior knowledge of the task) would have been a good baseline (showing how close or well these models get to a customised baseline model). I am fully aware that this would require a considerable amount of work and these comparisons are often omitted in ML papers. Maybe something to think of still.\n\n3. Minor: I recommend going through the paper and making sure especially citations and equations read well as part of the main text.", "extra_scores": {}, "extra_text": {}}, {"review_id": "DhE_7lTQwHb", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nThe paper is well written, with very few typos and amazing attention to detail (nice colors, thought-through figures, great content organization). Also, the math is well-explained. Overall, a very good read and a nice contribution. I recommend acceptance.\n\nJust a few points/questions for the authors:\n- Q1: How does your method scale as the dimension increases? I am thinking about for instance the experiment of section 4.1: if the dynamics is higher-dimensional (with still with 3 discontinuities in the vector field), is the method still able to infer the correct switching times?\n- Q2: Still about Figure 2.. how is it possible in panel (d) that the inferred discontinuity in the flow is not aligned with the discontinuity in the field? \n- Q3: I am a bit confused about how you infer successive switching times t*1, t*2 etc.. do you have to specify the number of switching times at the beginning of the inference procedure? Do you take gradients with respect to all the switching times together? Could you please add an explanation in the rebuttal and in your revised version of the paper?\n- Clarifications: I would add a few lines explaining what the “adjoint state” is before formula 5.\n- typos: Dots at the end of formulas are sometimes missing. End of page 4, “SLDS dynamics..”", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents an extension of the neural ODE approach to include discrete changes in the continuous-time dynamics. All reviewers agree the contribution made by this paper is worth publishing. Most of the reviewers' concerns have been answered in the rebuttal and I therefore recommend accepting this paper.

Decision: Accept (Poster)