Source-linked AI summary

Learning Neural Event Functions for Ordinary Differential Equations

Ricky T. Q. Chen, Brandon Amos, Maximilian Nickel

arXiv:2011.03902v4cs.LGstat.ML

TL;DR

Neural ODEs require a fixed termination time, limiting their fit to systems with state discontinuities. This paper learns differentiable event functions to define termination implicitly, and reports results on switching systems and collisions, including faster predictions with fewer function evaluations for collision modeling.

  • Problem

    Neural ODEs use fixed termination times, and ODE trajectories cannot represent the instantaneous state changes involved in events such as bouncing-ball collisions.

    Source: §1 · p. 1, §1.1 · p. 2
  • Method

    The paper parameterizes event functions with neural networks, differentiates through their implicitly determined termination times, and applies them to event-driven systems.

    Source: §1 · p. 1, §4 · p. 4
  • Results

    For collision modeling, Neural Event ODE generalizes better than the RNN baseline, matches the nonlinear Neural ODE, and uses fewer function evaluations for faster predictions.

    Source: §4.2 · p. 6
  • Takeaways & Limitations

    The approach models discrete events in continuous-time systems and supports simulation-based temporal point-process training with applications to discrete control.

    Source: §7 · p. 9

    A neural event function may never trigger, leaving it without a gradient and reducing the model to a Neural ODE alone.

    Source: §6 · p. 9

Abstract

from arXiv · show

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

Neural ODEs use fixed termination times and cannot represent some state discontinuities, motivating learned, differentiable event criteria. Event handling lets a solver stop at a state-dependent event, modify the state, and resume integration.

1 INTRODUCTION

  • Existing Neural ODEs use fixed termination times; learned, differentiable event criteria let termination depend on the evolving state and support state discontinuities.The approach derives gradients for neural event functions, allowing the model to handle discontinuities automatically.Source: §1 · p. 1

1.1 EVENT HANDLING

  • An event function terminates ODE solving when it crosses zero, after which the state can be changed instantaneously and integration resumed.The event time depends on the trajectory rather than being specified in advance.Source: §1.1 · p. 2
  • For the bouncing-ball example, Neural Event ODE recovers the underlying physics and extrapolates, while Neural ODE struggles with bounce discontinuities.The Neural Event ODE uses linear functions for both dynamics and event detection; the Neural ODE uses a nonlinear dynamics function.Source: §1.1 · p. 2, Figure 1 · p. 2
  • Although its trajectories can be discontinuous in time, Neural Event ODE is continuous in its parameters and can therefore be trained with gradient descent.The paper discusses differentiating the event time with respect to dependent variables, including the event state and model parameters.Source: §1.1 · p. 3

2 BACKGROUND: DIFFERENTIATING THROUGH ODE SOLUTIONS

The background explains how the adjoint method computes gradients through ODE solutions. It describes backward integration for vector-Jacobian products and notes alternative approaches with different memory-compute trade-offs.

2 BACKGROUND: DIFFERENTIATING THROUGH ODE SOLUTIONS

  • The adjoint state tracks how the loss changes with intermediate ODE states, and its backward evolution relates gradients of solver inputs to the loss.The relevant inputs include the initial state, initial time, solution time, and dynamics parameters.Source: §2 · p. 3
  • Solving the adjoint state backward alongside the trajectory yields vector-Jacobian products for ODE inputs and enables reverse-mode automatic differentiation.The described implementation jointly solves state, adjoint, and parameter gradients backward without storing intermediate states, invoking the solver once for gradient computation.Source: §2 · p. 3
  • Other adjoint approaches trade memory and computation differently, using stored intermediate quantities, checkpointing, interpolation, or symplectic integration.The paper treats these alternatives as compatible with its contribution and tangential to it.Source: §2 · p. 3

3 DIFFERENTIATING THROUGH EVENT HANDLING

The section derives gradients through event-determined termination times and uses differentiable event handling as a modeling primitive. Repeated event solves, state updates, and learnable functions support trajectories with a variable number of discontinuities.

3 DIFFERENTIATING THROUGH EVENT HANDLING

  • The implicit function theorem provides gradients through the root-finding procedure that determines the event time.The loss may depend directly on both the event time and the state at that time.Source: §3 · p. 4
  • The event-time gradient computation avoids differentiating through the ODE simulation for its scalar quantities, and one ODESolve computes the vector-Jacobian product.Source: §3 · p. 4
  • Unlike the cited specialized applications with event functions known in advance, this approach learns the event function for machine-learning applications.Source: §3 · p. 4

3.1 NEURAL EVENT FUNCTIONS AND INSTANTANEOUS UPDATES

  • Repeated ODESolveEvent calls combine learnable drift, event, and instantaneous-update functions to produce event times and a piecewise continuous trajectory with variable numbers of discontinuities.Source: §3.1 · p. 4

4 GENERAL STATE-DEPENDENT EVENT FUNCTIONS

Neural Event ODEs use state-dependent events and instantaneous updates to model hybrid continuous and discrete dynamics. Experiments apply this approach to switching systems and collisions, reporting sharp transitions and faster predictions, with long-sequence recovery vulnerable to early errors.

4 GENERAL STATE-DEPENDENT EVENT FUNCTIONS

  • State-dependent event functions suit discontinuities at state-space boundaries; repeated solves and instantaneous updates return event times and piecewise continuous trajectories.Source: §4 · pp. 4–5

4.1 SWITCHING LINEAR DYNAMICAL SYSTEMS

  • On a noisy fan-track task with 100 training and 25 longer validation and test trajectories, the model recovers components and generates realistic trajectories with sharp transitions.Source: §4.1 · p. 5, Figure 2 · p. 6

4.2 MODELING PHYSICAL SYSTEMS WITH COLLISION

  • On physics prediction, Neural Event ODE generalizes better than the RNN and matches the nonlinear Neural ODE, while using fewer function evaluations and predicting faster.Source: Table 2 · p. 6, §4.2 · p. 6
  • For collisions between two balls and a box, the model learns neural event detectors and instantaneous state updates that reflect contact-force changes in direction.Source: §4.2 · p. 6
  • On longer sequences, the Neural Event ODE can struggle to recover from early mistakes, leading to chaotic behavior.Source: §4.2 · p. 6

5 THRESHOLD-BASED EVENT FUNCTIONS

Threshold-based event functions turn an accumulated intensity into an event time, supporting simulation-based learning of temporal point processes and discrete control in continuous-time systems. Reparameterization gradients outperform REINFORCE in the reported point-process and control experiments.

5 THRESHOLD-BASED EVENT FUNCTIONS

  • A threshold-based event occurs when the integral of a positive intensity reaches a threshold, so its event time is obtained by solving for that crossing.The threshold can be represented as part of the state, with an event function defined by the threshold minus the accumulated intensity.Source: §5 · p. 6

5.1 REPARAMETERIZATION GRADIENT FOR TEMPORAL POINT PROCESSES

  • The reparameterization gradient outperforms REINFORCE in both discrete-control settings, and a deterministic policy can outperform both stochastic policies in the deterministic system.Figure 4 covers reverse-KL learning in panels (a,b) and continuous-time discrete control in (c,d).Source: Figure 4 · p. 8, §5.1 · p. 9
  • Sampling an event time uses exponential noise and solves for the time at which integrated intensity reaches that sample, making the sample differentiable through ODESolveEvent.The resulting reparameterization gradient is used to train temporal point processes through simulation.Source: §5.1 · p. 8
  • For Hawkes-process reverse-KL training, reparameterization performs well on five- and ten-event sequences, while REINFORCE converges more slowly on longer sequences.The reparameterization objective requires simulating from the model and differentiating through event handling.Source: §5.1 · p. 8
  • Neural Jump SDE control policies use multivariate point-process event indices to select discrete control values in continuous-time systems.Experiments include HIV treatment with two inhibitor choices and Hodgkin–Huxley dynamics with an input stimulus that switches on or off.Source: §5.1 · pp. 8–9

6 SCOPE AND LIMITATIONS

The approach has practical constraints in minibatching, event triggering, and objectives that depend on event counts. These include integration overhead, absent event gradients when boundaries are never reached, and discontinuous objectives when event counts change.

Minibatching

  • Combining minibatched trajectories requires restarting integration whenever any event function triggers, introducing overhead.A sparse min or max aggregator can combine the event functions.Source: §6 · p. 9

Runaway event boundaries

  • If a neural event function never triggers, its parameters receive no gradient and the model reduces to a Neural ODE.In high dimensions, neural-network roots can be unpredictable; large-standard-deviation initialization helps alleviate this, while threshold-based functions are less affected because their integrand is always positive.Source: §6 · p. 9

Numbers of events is discrete

  • Objectives depending on event counts can be discontinuous because small parameter changes may cause jumps in the number of events.For stochastic temporal point processes, reparameterization gradients rely on differentiability; Shchur et al. (2020) used surrogate objectives for sampling-based training.Source: §6 · p. 9

7 CONCLUSION

The paper extends Neural ODEs with neural event functions for implicitly defined termination times. It applies this approach to discrete events in continuous-time systems and simulation-based point-process training for discrete control.

Conclusion

  • Neural event functions extend Neural ODEs to implicitly defined termination times.Source: §7 · p. 9
  • The approach models discrete events in continuous-time systems, including collision criteria and effects in physical systems.Source: §7 · p. 9
  • Simulation-based temporal point-process training supports deterministic policies with discrete actions using gradients available in a continuous-time setting.Source: §7 · p. 9

A EXPERIMENTAL DETAILS

The experiments describe setups for switching dynamics, physical collisions, and discrete control. They specify the data, architectures, baselines, and training procedures used across these settings.

Continuous-time Switching Linear Dynamical Systems

  • The switching-dynamics experiment uses a fan-shaped system, trains on 100 sequences of length 4, and tests on 25 sequences of length 12.Sequences are discretized into 50 training steps and 150 validation and test steps; training data receive Gaussian noise with standard deviation 0.05. The model uses event-time weight updates, with an event function formed from two tanh-gated linear outputs.Source: A EXPERIMENTAL DETAILS · p. 14

Modeling Physical Systems with Collision

  • The collision experiment simulates two radius-0.5 balls in a [0, 5]^2 box, using 1,000 training initial positions and 25 each for validation and test.The event function takes the balls' positions as input; the instantaneous update uses a neural network. The LSTM baseline predicts position differences, while the nonlinear Neural ODE models velocity as the change in position.Source: Modeling Physical Systems with Collision · pp. 14–15

Discrete Control over Continuous-time Systems

  • Discrete control is evaluated on HIV dynamics with four treatment states and Hodgkin–Huxley dynamics with input current values of 0 or 100.A multivariate temporal point process changes the control state at each event; its intensity depends on a continuous-time neural ODE hidden state, which is instantaneously modified by an MLP.Source: Modeling Physical Systems with Collision · p. 15
Loading 2011.03902v4…