Original: N EURAL S PATIO -T EMPORAL P OINT P ROCESSES
Revised: N EURAL S PATIO -T EMPORAL P OINT P ROCESSES

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: We propose a new class of parameterizations for spatio-temporal point processes which leverage Neural ODEs as a computational method and enable ﬂexible, highﬁdelity models of discrete events that are localized in continuous time and space.
Revised: We propose a new class of parameterizations for spatio-temporal point processes which leverage Neural ODEs as a computational method and enable ﬂexible, highﬁdelity models of discrete events that are localized in continuous time and space.

Original: Central to our approach is a combination of recurrent continuous-time neural networks with two novel neural architectures, i.e., Jump and Attentive Continuoustime Normalizing Flows.
Revised: Central to our approach is a combination of recurrent continuous-time neural networks with two novel neural architectures, i.e. , Jump and Attentive Continuoustime Normalizing Flows.

Original: This approach allows us to learn complex distributions for both the spatial and temporal domain and to condition non-trivially on the observed event history.
Revised: This approach allows us to learn complex distributions for both the spatial and temporal domain and to condition non-trivially on the observed event history.

Original: We validate our models on data sets from a wide variety of contexts such as seismology, epidemiology, urban mobility, and neuroscience.
Revised: We validate our models on data sets from a wide variety of contexts such as seismology, epidemiology, urban mobility, and neuroscience.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Modeling discrete events that are localized in continuous time and space is an important task in many scientiﬁc ﬁelds and applications.
Revised: Modeling discrete events that are localized in continuous time and space is an important task in many scientiﬁc ﬁelds and applications.

Original: Spatio-temporal point processes (STPPs) are a versatile and principled framework for modeling such event data and have, consequently, found many applications in a diverse range of ﬁelds.
Revised: Spatio-temporal point processes (STPPs) are a versatile and principled framework for modeling such event data and have, consequently, found many applications in a diverse range of ﬁelds.

Original: This includes, for instance, modeling earthquakes and aftershocks (Ogata, 1988; 1998), the occurrence and propagation of wildﬁres (Hering et al., 2009), epidemics and infectious diseases (Meyer et al., 2012; Schoenberg et al., 2019), urban mobility (Du et al., 2016), the spread of invasive species (Balderama et al., 2012), and brain activity (Tagliazucchi et al., 2012). 
Revised: This includes, for instance, modeling earthquakes and aftershocks (Ogata, 1988; 1998), the occurrence and propagation of wildﬁres (Hering et al., 2009), epidemics and infectious diseases (Meyer et al., 2012; Schoenberg et al., 2019), urban mobility (Du et al., 2016), the spread of invasive species (Balderama et al., 2012), and brain activity (Tagliazucchi et al., 2012). x t * ( t )

Original: It is of great interest in all of these areas to learn high-ﬁdelity models which can jointly capture spatial and temporal dependencies and their propagation effects.
Revised: It is of great interest in all of these areas to learn high-ﬁdelity models which can jointly capture spatial and temporal dependencies and their propagation effects.

Original: However, existing parameterizations of STPPs are strongly restricted in this regard due to computational considerations:
Revised: However, existing parameterizations of STPPs are strongly restricted in this regard due to computational considerations:

Original: In its general form,
Revised: In its general form,

Original: STPPs require to solve multivariate integrals for computing likelihood values and thus have primarily been studied within the context of different approximations and model restrictions.
Revised: STPPs require solving multivariate integrals for computing likelihood values and thus have primarily been studied within the context of different approximations and model restrictions.

Original: This includes, for instance, restricting the model class to parameterizations with known closed-form solutions (e.g., exponential Hawkes processes (Ozaki, 1979)), to restrict dependencies between the spatial and temporal domain (e.g., independent and unpredictable marks (Daley & Vere-Jones, 2003)), or to discretize continuous time and space (Ogata, 1998).
Revised: This includes, for instance, restricting the model class to parameterizations with known closed-form solutions ( e.g. , exponential Hawkes processes (Ozaki, 1979)), to restrict dependencies between the spatial and temporal domain ( e.g. , independent and unpredictable marks (Daley & Vere-Jones, 2003)), or to discretize continuous time and space (Ogata, 1998).

Original: These restrictions and approximations—which can lead to mis-speciﬁed models and loss of information—motivatedrecently the development of neural temporal point processes such as Neural Hawkes Processes (Mei & Eisner, 2017) and Neural
Revised: These restrictions and approximations—which can lead to mis-speciﬁed models and loss of information—motivated the development of neural temporal point processes such as Neural Hawkes Processes (Mei & Eisner, 2017) and Neural

Original: Jump SDEs (Jia & Benson, 2019).
Revised: Jump SDEs (Jia & Benson, 2019).

Original: While these methods are more ﬂexible, they can still require approximations such as Monte-Carlo sampling of the likelihood (Mei & Eisner, 2017; Nickel & Le, 2020) and, most importantly, only model restricted spatial distributions (Jia & Benson, 2019).
Revised: While these methods are more ﬂexible, they can still require approximations such as Monte-Carlo sampling of the likelihood (Mei & Eisner, 2017; Nickel & Le, 2020) and, most importantly, only model restricted spatial distributions (Jia & Benson, 2019).

Original: To overcome these issues, we propose a new class of parameterizations for spatio-temporal point processes which leverage Neural ODEs as a computational method and allow to deﬁne ﬂexible, high-ﬁdelity models for spatio-temporal event data.
Revised: To overcome these issues, we propose a new class of parameterizations for spatio-temporal point processes which leverage Neural ODEs as a computational method and allow to deﬁne ﬂexible, high-ﬁdelity models for spatio-temporal event data.

Original: We build upon ideas of Neural Jump SDEs (Jia & Benson, 2019) and Continuous-time Normalizing Flows (CNFs; Chen et al. 2018; Grathwohl et al. 2019; Mathieu & Nickel 2020) to learn parametric models of spatial (or mark 1 ) distributions that are deﬁned continuously in time. 
Revised: We build upon ideas of Neural Jump SDEs (Jia & Benson, 2019) and Continuous-time Normalizing Flows (CNFs; Chen et al. 2018; Grathwohl et al. 2019; Mathieu & Nickel 2020) to learn parametric models of spatial (or mark 1 ) distributions that are deﬁned continuously in time. Normalizing ﬂows are known to be ﬂexible universal density estimators ( e.g. Huang et al. 2018; 2020; Teshima et al. 2020;

Original: Our approach allowsto compute the  exact likelihood values even for highly complex spatio-temporal distributions and creates smoothly changing spatial distributions that naturally beneﬁts spatio-temporal modeling.
Revised: Kong & Chaudhuri 2020) while retaining computational tractability. As such, our approach allows the computation of exact likelihood values even for highly complex spatio-temporal distributions, and our models create smoothly changing spatial distributions that naturally beneﬁts spatio-temporal modeling.

Original: Central to our approach, are two novel neural architectures based on CNFs—using either discontinuous jumps in distribution or self-attention—to condition spatial distributions on the event history.
Revised: Central to our approach, are two novel neural architectures based on CNFs—using either discontinuous jumps in distribution or self-attention—to condition spatial distributions on the event history.

Original: To the best of our knowledge, this is the ﬁrst method that combines the ﬂexibility or neural TPPs with the ability to learn high-ﬁdelity models of continuous marks that can have complex dependencies on the event history.
Revised: To the best of our knowledge, this is the ﬁrst method that combines the ﬂexibility of neural TPPs with the ability to learn high-ﬁdelity models of continuous marks that can have complex dependencies on the event history.

Original: In addition to our modeling contributions, we also construct ﬁve new pre-processed data sets for benchmarking spatio-temporal event models.
Revised: In addition to our modeling contributions, we also construct ﬁve new pre-processed data sets for benchmarking spatio-temporal event models.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: In the following, we give a brief overview of two core frameworks which our method builds upon, i.e., spatio-temporal point processes and continuous-time normalizing ﬂows.
Revised: In the following, we give a brief overview of two core frameworks which our method builds upon, i.e. , spatio-temporal point processes and continuous-time normalizing ﬂows.

Original: Event Modeling with Point Processes Spatio-temporal point processes are concerned with modeling sequences of random events in continuous space and time (Moller & Waagepetersen, 2003;
Revised: Event Modeling with Point Processes Spatio-temporal point processes are concerned with modeling sequences of random events in continuous space and time (Moller & Waagepetersen, 2003;

Original: Baddeley et al., 2007).
Revised: Baddeley et al., 2007).

Original: Let H =
Revised: Let H =

Original: { ( t i , x i ) } ni =1 denote the sequence of event times t i ∈ R and their associated locations x i ∈ R d , the number of events n being also random.
Revised: { ( t i , x i ) } ni =1 denote the sequence of event times t i ∈ R and their associated locations x i ∈ R d , the number of events n being also random.

Original: Additionally, let
Revised: Additionally, let

Original: H t = { ( t i , x i )
Revised: H t = { ( t i , x i )

Original: | t i < t, t i ∈ H} denote the history of events predating time t .
Revised: | t i < t, t i ∈ H} denote the history of events predating time t .

Original: A spatio-temporal point process is then fully characterized by its conditional intensity function
Revised: A spatio-temporal point process is then fully characterized by its conditional intensity function

Original: [Equation]
Revised: [Equation]

Original: where B ( x , ∆ x ) denotes a ball centered at x ∈ R d and with radius ∆ x .
Revised: where B ( x , ∆ x ) denotes a ball centered at x ∈ R d and with radius ∆ x .

Original: The only condition is that λ ( t, x | H t ) ≥ 0 and need not be normalized.
Revised: The only condition is that λ ( t, x | H t ) ≥ 0 and need not be normalized.

Original: Given i − 1 previous events, the conditional intensity function describes therefore the instantaneous probability of the i -th event occurring at t and location x .
Revised: Given i − 1 previous events, the conditional intensity function describes therefore the instantaneous probability of the i -th event occurring at t and location x .

Original: In the following, we will use the common star superscript shorthand λ ∗ ( t, x ) =
Revised: In the following, we will use the common star superscript shorthand λ ∗ ( t, x ) =

Original: λ ( t, x | H t ) to denote conditional dependence on the history.
Revised: λ ( t, x | H t ) to denote conditional dependence on the history.

Original: The joint log-likelihood of observing H within a time interval of [0 , T ] is then given by (Daley & Vere-Jones, 2003, Proposition 7.3.III)
Revised: The joint log-likelihood of observing H within a time interval of [0 , T ] is then given by (Daley & Vere-Jones, 2003, Proposition 7.3.III)

Original: [Equation]
Revised: [Equation]

Original: Training general STPPs with maximum likelihood is difﬁcult as eq.
Revised: Training general STPPs with maximum likelihood is difﬁcult as eq.

Original: (2) requires solving a multivariate integral.
Revised: (2) requires solving a multivariate integral.

Original: This need to compute integrals has driven research to focus around the use of kernel density estimators (KDE) with exponential kernels that have known anti-derivatives (Reinhart et al., 2018).
Revised: This need to compute integrals has driven research to focus around the use of kernel density estimators (KDE) with exponential kernels that have known anti-derivatives (Reinhart et al., 2018).

Original: Continuous-Time Normalizing Flows Normalizing ﬂows (Dinh et al., 2014; 2016; Rezende &
Revised: Continuous-time Normalizing Flows Normalizing ﬂows (Dinh et al., 2014; 2016; Rezende &

Original: Mohamed, 2015) is a class of density models that describe ﬂexible distributions by parameterizing an invertible transformation from a simpler base distribution, which enables exact computation of the probability of the transformed distribution, without any unknown normalization constants.
Revised: Mohamed, 2015) is a class of density models that describe ﬂexible distributions by parameterizing an invertible transformation from a simpler base distribution, which enables exact computation of the probability of the transformed distribution, without any unknown normalization constants.

Original: Given a random variable x 0 with known distribution p ( x 0 ) and an invertible transformation T ( x ) , the transformed variable T ( x 0 ) is a random variable with a probability distribution function that satisﬁes
Revised: Given a random variable x 0 with known distribution p ( x 0 ) and an invertible transformation F ( x ) , the transformed variable F ( x 0 ) is a random variable with a probability distribution function that satisﬁes

Original: [Equation]
Revised: [Equation]

Original: There have been many advances in parameterizing T with ﬂexible neural networks that also allow for cheap evaluations of eq.
Revised: There have been many advances in parameterizing F with ﬂexible neural networks that also allow for cheap evaluations of eq.

Original: We focus our attention on Continuous-time Normalizing Flows (CNFs), which parameterizes this transformation with a Neural ODE (Chen et al., 2018).
Revised: We focus our attention on Continuous-time Normalizing Flows (CNFs), which parameterizes this transformation with a Neural ODE (Chen et al., 2018).

Original: CNFs deﬁne an inﬁnite set of distributions on the real line that vary smoothly across time, and will be our core component for modeling events in the spatial domain.
Revised: CNFs deﬁne an inﬁnite set of distributions on the real line that vary smoothly across time, and will be our core component for modeling events in the spatial domain.

Original: Let p ( x 0 ) be the base distribution 2 .
Revised: Let p ( x 0 ) be the base distribution 2 .

Original: We then parameterize an instantaneous change in the form of an ordinary differential equation (ODE),
Revised: We then parameterize an instantaneous change in the form of an ordinary differential equation (ODE),

Original: d x t dt = f
Revised: d x t dt = f

Original: ( t, x t ) , where the subscript denotes dependence on t .
Revised: ( t, x t ) , where the subscript denotes dependence on t .

Original: This function can be parameterized using any Lipschitz-continuous neural network.
Revised: This function can be parameterized using any Lipschitz-continuous neural network.

Original: Conditioned on a sample x 0 from the base distribution, let x t be the solution of the initial value problem 3 at time t , i.e. it is from a trajectory that passes through x 0 at time 0 and satisﬁes the ODE d x t / dt = f .
Revised: Conditioned on a sample x 0 from the base distribution, let x t be the solution of the initial value problem 3 at time t , i.e. it is from a trajectory that passes through x 0 at time 0 and satisﬁes the ODE d x t / dt = f .

Original: We can express the value of the solution at time t as
Revised: We can express the value of the solution at time t as

Original: [Equation]
Revised: [Equation]

Original: The distribution of x t then also continuously changes in t through the following equation,
Revised: The distribution of x t then also continuously changes in t through the following equation,

Original: [Equation]
Revised: [Equation]

Original: In practice, eq.
Revised: In practice, eq.

Original: (4) and eq.
Revised: (4) and eq.

Original: (5) are solved together from 0 to t , as eq.
Revised: (5) are solved together from 0 to t , as eq.

Original: (5) alone is not an ordinary differential equation but the combination of x t and log p ( x t ) is.
Revised: (5) alone is not an ordinary differential equation but the combination of x t and log p ( x t ) is.

Original: The trace of the Jacobian ∂f∂x ( τ, x τ ) can be estimated using a Monte Carlo estimate of the identity (Skilling, 1989; Hutchinson, 1990), tr ( A ) =
Revised: The trace of the Jacobian ∂f∂x ( τ, x τ ) can be estimated using a Monte Carlo estimate of the identity (Skilling, 1989; Hutchinson, 1990), tr ( A ) =

Original: E v ∼N (0 , 1) [ v T Av ] .
Revised: E v ∼N (0 , 1) [ v T Av ] .

Original: This estimator relies only on a vector-Jacobian product, which can be efﬁciently computed in modern automatic differentiation and deep learning frameworks.
Revised: This estimator relies only on a vector-Jacobian product, which can be efﬁciently computed in modern automatic differentiation and deep learning frameworks.

Original: This has been used (Grathwohl et al., 2019) to scale CNFs to higher dimensions using a Monte Carlo estimate of the log likelihood objective,
Revised: This has been used (Grathwohl et al., 2019) to scale CNFs to higher dimensions using a Monte Carlo estimate of the log likelihood objective,

Original: [Equation]
Revised: [Equation]

Original: which, even if only one sample of v is used, is still amenable to training with stochastic gradient descent.
Revised: which, even if only one sample of v is used, is still amenable to training with stochastic gradient descent.

Original: Gradients with respect to any parameters in f can be computed with constant memory by solving an adjoint ODE in reverse-time as described in Chen et al.
Revised: Gradients with respect to any parameters in f can be computed with constant memory by solving an adjoint ODE in reverse-time as described in Chen et al.

Original:  N EURAL S PATIO -T EMPORAL P OINT P ROCESSES
Revised:  N EURAL S PATIO -T EMPORAL P OINT P ROCESSES

Original: We are interested in modeling high-ﬁdelity distributions in continuous time and space that can be updated based on new event information.
Revised: We are interested in modeling high-ﬁdelity distributions in continuous time and space that can be updated based on new event information.

Original: For this purpose, we use the Neural ODE framework to parameterize a STPP by combining ideas from Neural Jump SDEs and Continuous Normalizing
Revised: For this purpose, we use the Neural ODE framework to parameterize a STPP by combining ideas from Neural Jump SDEs and Continuous Normalizing

Original: Flows to create highly ﬂexible models that still allow exact likelihood computation.
Revised: Flows to create highly ﬂexible models that still allow exact likelihood computation.

Original: We ﬁrst (re-)introduce necessary notation.
Revised: We ﬁrst (re-)introduce necessary notation.

Original: Let H =
Revised: Let H =

Original: { ( t i , x ( i ) t i ) } denote a sequence of event times t i ∈ [0 , T ] and locations x ( i ) t i ∈ R d .
Revised: { ( t i , x ( i ) t i ) } denote a sequence of event times t i ∈ [0 , T ] and locations x ( i ) t i ∈ R d .

Original: The superscript indicates an association with the i -th event, and the use of subscripting with t i will be useful later in the continuous-time modeling framework.
Revised: The superscript indicates an association with the i -th event, and the use of subscripting with t i will be useful later in the continuous-time modeling framework.

Original: Following Daley & Vere-Jones (2003), we decompose the conditional intensity function as 
Revised: Following Daley & Vere-Jones (2003), we decompose the conditional intensity function as [Equation]

Original: [Equation]
Revised: [Equation]

Original: where λ ∗ ( t ) is the ground intensity of the temporal process and where p ∗ ( x | t ) is the conditional density of a mark x at t given H t .
Revised: where λ ∗ ( t ) is the ground intensity of the temporal process and where p ∗ ( x | t ) is the conditional density of a mark x at t given H t .

Original: The star superscript is used as again shorthand to denote dependence on the history.
Revised: The star superscript is used as again shorthand to denote dependence on the history.

Original: Since (cid:82) R d p ∗ ( x | t ) = 1 , eq.
Revised: Since (cid:82) R d p ∗ ( x | t ) = 1 , eq.

Original: (7) allows us now to simplify the log-likelihood function of the joint process from eq.
Revised: (7) allows us now to simplify the log-likelihood function of the joint process from eq.

Original: (2), such that
Revised: (2), such that

Original: [Equation]
Revised: [Equation]

Original: Furthermore, based on eq.
Revised: Furthermore, based on eq.

Original: (7), we can derive separate models for the ground intensity and conditional mark density whichwe will condition both on a continuous-time hidden state with jumps.
Revised: (7), we can derive separate models for the ground intensity and conditional mark density which will be jointly conditioned on a continuous-time hidden state with jumps.

Original: In the following, we will ﬁrst describe how we construct a latent dynamics model, which we use to compute  the ground intensity λ ∗ ( t ) .
Revised: In the following, we will ﬁrst describe how we construct a latent dynamics model, which we use to compute This initial value problem has a unique solution if f is Lipschitz-continuous along the trajectory to x t , by the Picard–Lindel¨of theorem.

Original:  We will then propose three novel CNF-based approaches for modeling the conditional mark density p ∗ ( x | t ) .
Revised: the ground intensity λ ∗ ( t ) . We will then propose three novel CNF-based approaches for modeling the conditional mark density p ∗ ( x | t ) .

Original: We will ﬁrst describe an unconditional model, which is already a strong baseline when spatial event distributions only follow temporal patterns and there is little to no correlation between the spatial observations.
Revised: We will ﬁrst describe an unconditional model, which is already a strong baseline when spatial event distributions only follow temporal patterns and there is little to no correlation between the spatial observations.

Original: We then devise two new methods of conditioning on the event history H ; one explicitly modeling instantaneous changes in distribution, and another that uses an attention mechanism which is more amenable to parallelism.
Revised: We then devise two new methods of conditioning on the event history H : one explicitly modeling instantaneous changes in distribution, and another that uses an attention mechanism which is more amenable to parallelism.

Original: Latent Dynamics and Ground Intensity For the temporal variables { t i } ,we follow the work of Jia & Benson (2019) and parameterize the intensity function using hidden state dynamics with jumps.
Revised: Latent Dynamics and Ground Intensity For the temporal variables { t i } , parameterize the intensity function using hidden state dynamics with jumps, similar to the work of Jia & Benson (2019).

Original: Speciﬁcally, we evolve a continuous-time hidden state h and set
Revised: Speciﬁcally, we evolve a continuous-time hidden state h and set

Original: [Equation]
Revised: [Equation]

Original: where g λ is a neural network with a positive output.
Revised: where g λ is a neural network with a softplus nonlinearity applied to the output, to ensure the intensity is positive.

Original: We then capture conditional dependencies through the use of a continuously changing state h t with instantaneous updates when conditioned on an event.
Revised: We then capture conditional dependencies through the use of a continuously changing state h t with instantaneous updates when conditioned on an event.

Original: The architecture is analogous to a recurrent neural network with a continuous-time hidden state (Rubanova et al., 2019).
Revised: The architecture is analogous to a recurrent neural network with a continuous-time hidden state (Mei & Eisner, 2017; Che et al., 2018; Rubanova et al., 2019) modeled by a Neural ODE.

Original: This provides us with a vector representation h t at every time value t that acts as both a summary of the history of events and as a predictor of future behavior.
Revised: This provides us with a vector representation h t at every time value t that acts as both a summary of the history of events and as a predictor of future behavior.

Original: Instantaneous updates to h t allow to incorporate abrupt changes to the hidden state that are triggered by observed events.
Revised: Instantaneous updates to h t allow to incorporate abrupt changes to the hidden state that are triggered by observed events.

Original: This mechanism is important for modeling TPPs and allows past events to inﬂuence future dynamics in a discontinuous way (e.g., modeling immediate shocks to a system).
Revised: This mechanism is important for modeling point processes and allows past events to inﬂuence future dynamics in a discontinuous way ( e.g. , modeling immediate shocks to a system).

Original: We use f h to model the continuous change in the form of an ODE and g h to model instantaneous changes based on an observed event. 
Revised: We use f h to model the continuous change in the form of an ODE and g h to model instantaneous changes based on an observed event. [Equation] (An initial hidden state [Equation]

Original: [Equation]
Revised: [Equation]

Original: between event times
Revised: between event times

Original: [Equation]
Revised: [Equation]

Original:  The use of ε is to portray that h t is a c`agl`ad function, i.e. left-continuous with right limits, with a discontinuous jump modeled by g h .
Revised: Continuous evolution [Equation] [Equation] The use of ε is to portray that h t is a c`agl`ad function, i.e. left-continuous with right limits, with a discontinuous jump modeled by g h .

Original: The parameterization of continuous-time hidden states in the form of eqs.
Revised: The parameterization of continuous-time hidden states in the form of eqs.

Original: (10) to (12) has been used for time series modeling (Rubanova et al., 2019;
Revised: (10) to (12) has been used for time series modeling (Rubanova et al., 2019;

Original: De Brouwer et al., 2019) as well as TPPs (Jia &
Revised: De Brouwer et al., 2019) as well as TPPs (Jia &

Original: Benson, 2019).
Revised: Benson, 2019).

Original: We parameterize f h as a standard multi-layer fully connected neural network, and use the GRU update (Cho et al., 2014) to parameterize g h , as was done in Rubanova et al.
Revised: We parameterize f h as a standard multi-layer fully connected neural network, and use the GRU update (Cho et al., 2014) to parameterize g h , as was done in Rubanova et al.

Original: Time-Varying CNF
Revised: Time-varying CNF

Original: The ﬁrst model we consider is a straightforward application of the CNF to time-variable observations.
Revised: The ﬁrst model we consider is a straightforward application of the CNF to time-variable observations.

Original: Assuming that the spatial distribution is independent of prior events,we have
Revised: Assuming that the spatial distribution is independent of prior events,

Original: [Equation]
Revised: [Equation]

Original: where x ( i ) τ is the solution of the ODE f with initial value x ( i ) t i  at τ = t i .
Revised: where x ( i ) τ is the solution of the ODE f with initial value x ( i ) t i , the observed event location, at τ = t i , the observed event time.

Original: Thus, the spatial distribution of an event  changes with respect to the time it occurs.
Revised: The spatial distribution of an event modeled by a Time-varying CNF changes with respect to the time it occurs.

Original: Some spatio-temporal data sets exhibit mostly temporal patterns and little to no dependence on previous events in the spatial domain, which would make a time-varying CNF a good ﬁt.
Revised: Some spatio-temporal data sets exhibit mostly temporal patterns and little to no dependence on previous events in the spatial domain, which would make a time-varying CNF a good ﬁt.

Original: Nevertheless, this model lacks the ability to capture spatial propagation effects, as it does not condition on previous event observations.
Revised: Nevertheless, this model lacks the ability to capture spatial propagation effects, as it does not condition on previous event observations.

Original: A major beneﬁt of this model is the ability to evaluate the joint log-likelihood fully in parallel across events, since there are no dependencies between events.
Revised: A major beneﬁt of this model is the ability to evaluate the joint log-likelihood fully in parallel across events, since there are no dependencies between events.

Original: Most modern ODE solvers that we are aware of only allow a scalar terminal time.
Revised: Most modern ODE solvers that we are aware of only allow a scalar terminal time.

Original: Thus, to solve all n integrals in eq.
Revised: Thus, to solve all n integrals in eq.

Original: (13) with a single call to an
Revised: (13) with a single call to an

Original: ODE solver, we can simply reparameterize all integrals with a consistent dummy variable and track the terminal time in the state.
Revised: ODE solver, we can simply reparameterize all integrals with a consistent dummy variable and track the terminal time in the state (see Appendix F for detailed explanation).

Revised: Intuitively, the idea is that we can reparameterize ODEs that are on t ∈ [0 , t i ] into an ODE on s ∈ [0 , 1] using the change of variables s = t / t i (or t = st i ) and scaling the output of f by t i .

Original: The joint ODEwe use after the change of variables is 
Revised: The joint ODE is then d

Original: [Equation]
Revised: [Equation]

Original: which gives
Revised: which gives

Original: [Equation]
Revised: [Equation]

Original: Thus the full trajectories between 0 to t i for all events can be computed in parallel using this augmented ODE by simply integrating once from s = 0 to s = 1 .
Revised: Thus the full trajectories between 0 to t i for all events can be computed in parallel using this augmented ODE by simply integrating once from s = 0 to s = 1 .

Original: Jump CNF
Revised: Jump CNF

Original: For the second model, we condition the dynamics deﬁning the continuous normalizing ﬂow on the hidden state h , allowing the normalizing ﬂow to update its distribution based on changes in H .
Revised: For the second model, we condition the dynamics deﬁning the continuous normalizing ﬂow on the hidden state h , allowing the normalizing ﬂow to update its distribution based on changes in H .

Original: For this purpose, we deﬁne continuous-time spatial distributions by making again use of two components: (i) a continuous-time normalizing ﬂow that evolves the distribution continuously, and (ii) a standard (discrete-time) ﬂow model that changes the distribution instantaneously after conditioning on new events.
Revised: For this purpose, we deﬁne continuous-time spatial distributions by making again use of two components: (i) a continuous-time normalizing ﬂow that evolves the distribution continuously, and (ii) a standard (discrete-time) ﬂow model that changes the distribution instantaneously after conditioning on new events.

Original: As normalizing ﬂows parameterize distributions through transformations of the samples, these continuous- and discrete-time transformations are composable in a straightforward manner and are end-to-end differentiable.
Revised: As normalizing ﬂows parameterize distributions through transformations of the samples, these continuous- and discrete-time transformations are composable in a straightforward manner and are end-to-end differentiable.

Original: The generative process of a single event in a Jump CNF is given by:
Revised: The generative process of a single event in a Jump CNF is given by:

Original: [Equation]
Revised: [Equation]

Original: between eventtimes
Revised: between event

Original: [Equation]
Revised: [Equation]

Original:  at event [Equation]
Revised: The initial distribution can be parameterized by a normalizing ﬂow. In practice, we set a base distribution at a negative time value and model p ( x 0 ) using the same CNF parameterized by f x .

Original: The instantaneous updates (or jumps) describe conditional updates in distribution after each new event has been observed.
Revised: The instantaneous updates (or jumps) describe conditional updates in distribution after each new event has been observed.

Original: This conditioning on h t i is required for the continuous and instantaneous updates to depend on the history of observations.
Revised: This conditioning on h t i is required for the continuous and instantaneous updates to depend on the history of observations.

Original: Otherwise, a Jump CNF would only be able to model the marginal distribution and behave similarly to a time-varying CNF.
Revised: Otherwise, a Jump CNF would only be able to model the marginal distribution and behave similarly to a time-varying CNF.

Original:  The ﬁnal probability of an event x t at some t > t n after observing n events is given by the sum of changes according to the continuous- and discrete-time normalizing ﬂows.
Revised: We solve for h t alongside x t . The ﬁnal probability of an event x t at some t > t n after observing n events is given by the sum of changes according to the continuous- and discrete-time normalizing ﬂows.

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Revised: event

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]

Original: [Equation]
Revised: [Equation]

Original: As the instantaneous updates must be applied sequentially in a Jump CNF, we can only compute the integrals in eq.
Revised: As the instantaneous updates must be applied sequentially in a Jump CNF, we can only compute the integrals in eq.

Original: (18) one at a time.
Revised: (18) one at a time.

Original: As such, the number of initial value problems scales linearly with the number of events in the history because the ODE solver must be restarted between each instantaneous update to account for the discontinuous change to state.
Revised: As such, the number of initial value problems scales linearly with the number of events in the history because the ODE solver must be restarted between each instantaneous update to account for the discontinuous change to state.

Original: This incurs a substantial cost when the number of events is large.
Revised: This incurs a substantial cost when the number of events is large.

Original: Attentive CNF To design a spatial model with conditional dependencies that alleviates the computational issues of Jump CNFs and can be computed in parallel, we make use of efﬁcient attention mechanisms based on the Transformer architecture (Vaswani et al., 2017).
Revised: Attentive CNF To design a spatial model with conditional dependencies that alleviates the computational issues of Jump CNFs and can be computed in parallel, we make use of efﬁcient attention mechanisms based on the Transformer architecture (Vaswani et al., 2017).

Original: Denoting only the spatial variables for simplicity, each conditional distribution log p ( x t i | H t i ) can be modeled by a CNF that depends on the sample path of prior events.
Revised: Denoting only the spatial variables for simplicity, each conditional distribution log p ( x t i | H t i ) can be modeled by a CNF that depends on the sample path of prior events.

Original: Speciﬁcally, we take the dummy-variable reparameterization of eq.
Revised: Speciﬁcally, we take the dummy-variable reparameterization of eq.

Original: (14) and modify it so that the i -th event depends on all previous events using a Transformer architecture for f ,
Revised: (14) and modify it so that the i -th event depends on all previous events using a Transformer [Figure] ) t [Figure]

Revised: [Figure]

Revised: ) t t 4 t 3 t 2 t

Original: [Equation]
Revised: T [Figure] ( 4 ) t

Revised: Jump CNF x ( 1 ) t

Revised: Sequence 1 Sequencex ( 2 ) t x ( 3 ) t tttt

Revised: T x ( 4 ) t

Revised: Attentive

Revised: [Figure]

Original:  With this formulation, the trajectory of x ( i ) τ depends continuously on the trajectory of x ( j ) τ for all j <
Revised: architecture for f , [Equation] [Equation] With this formulation, the trajectory of x ( i ) τ depends continuously on the trajectory of x ( j ) τ for all j <

Original: i and the hidden state h .
Revised: i and the hidden states h prior to the i -th event.

Original: Similar to eq.
Revised: Similar to eq.

Original: (14), an Attention CNF can now solve for the trajectories of all events in parallel but simultaneously depend non-trivially on H .
Revised: (14), an Attention CNF can now solve for the trajectories of all events in parallel while simultaneously depending non-trivially on H .

Original: To parameterize f Attn , we use an embedding layer followed by two multihead attention (MHA) blocks and an output layer to map back into the input space.
Revised: To parameterize f Attn , we use an embedding layer followed by two multihead attention (MHA) blocks and an output layer to map back into the input space.

Original: We use the Lipschitz-continuous multihead attention from Kim et al.
Revised: We use the Lipschitz-continuous multihead attention from Kim et al.

Original: (2020) as they recently showed that the dot product multihead attention (Vaswani et al., 2017) is not Lipschitz-continuous  and thus may be ill-suited for parameterizing ODEs. 
Revised: (2020) as they recently showed that the dot product multihead attention (Vaswani et al., 2017) is not Lipschitz-continuous in and thus may be ill-suited for parameterizing ODEs. [Figure]

Original: Low-variance Log-likelihood Estimation
Revised: Low-variance Log-likelihood Estimation

Original: The variance of the  stochastic trace estimator in eq.
Revised: The variance of the Hutchinson stochastic trace estimator in eq.

Original: (6) grows with the squared Frobenius norm of the Jacobian, (cid:80) ij [ ∂f / ∂x ] 2 ij (Hutchinson, 1990).
Revised: (6) grows with the squared Frobenius norm of the Jacobian, (cid:80) ij [ ∂f / ∂x ] 2 ij (Hutchinson, 1990).

Original: For attentive CNFs, we can remove some of the non-diagonal elements of the Jacobian and achieve a lower variance estimator.
Revised: For attentive CNFs, we can remove some of the non-diagonal elements of the Jacobian and achieve a lower variance estimator.

Original: The attention mechanism creates a block-triangular Jacobian, where each block corresponds to one event, but the elements outside of the block-diagonal are solely due to the multihead attention.
Revised: The attention mechanism creates a blocktriangular Jacobian, where each block corresponds to one event, but the elements outside of the blockdiagonal are solely due to the multihead attention.

Original: By detaching the gradient connections between different events in the MHA blocks, we can create a surrogate Jacobian matrix that do not contain cross-event partial derivatives.
Revised: By detaching the gradient connections between different events in the MHA blocks, we can create a surrogate Jacobian matrix that do not contain cross-event partial derivatives.

Original: This effectively allows us to apply the stochastic trace estimator on a matrix that has the same diagonal elements as the Jacobian ∂f / ∂x —and thus has the same expected value—but has zeros outside of the block-diagonal, leading to a lower variance trace estimator.
Revised: This effectively allows us to apply Hutchinson’s estimator on a matrix that has the same diagonal elements as the Jacobian ∂f / ∂x —and thus has the same expected value—but has zeros outside of the block-diagonal, leading to a lower variance trace estimator.

Original: The procedure consists of selectively removing partial derivatives and is straightforward, but notationally cumbersome; the interested reader can ﬁnd the details in Appendix D.
Revised: The procedure consists of selectively removing partial derivatives and is straightforward but notationally cumbersome; the interested reader can ﬁnd the details in Appendix E.

Revised: An ablation experiment is shown in Figure 3 for training on the P INWHEEL data set, where the lower variance estimates (and gradients) ultimate led to faster convergence and better converged models.

Original:  R ELATED WORK
Revised:  R ELATED W

Revised: ORK

Original: Neural Temporal Point Processes Modeling real-world data using restricted models such as
Revised: Neural Temporal Point Processes Modeling real-world data using restricted models such as

Original: Exponential Hawkes Processes (Ozaki, 1979) may lead to poor results due to model mis-speciﬁcation.
Revised: Exponential Hawkes Processes (Ozaki, 1979) may lead to poor results due to model mis-speciﬁcation.

Original: This motivated a variety of recent works  to explore neural networks for the parameterization of TPPs.
Revised: While this has led to many works on improving the Hawkes process ( e.g. Linderman & Adams 2014; Li & Zha 2014; Zhao et al. 2015; Farajtabar et al. 2017; Li & Ke 2020), recent works have begun to explore neural network parameterizations of TPPs.

Original: A common approach is to use recurrent neural networks to accumulate the event history in a latent state from which the intensity value can then be derived.
Revised: A common approach is to use recurrent neural networks to accumulate the event history in a latent state from which the intensity value can then be derived.

Original: Models of this form include, for instance, 
Revised: Models of this form include, for instance, Recurrent Marked Temporal Point

Original: Recurrent Marked Temporal Point Processes (RMTPPs; Du et al. 2016) and Neural Hawkes Processes (NHPs; Mei & Eisner 2017).
Revised: Processes (RMTPPs; Du et al. 2016) and Neural Hawkes Processes (NHPs; Mei & Eisner 2017).

Original: In contrast to our approach, these methods can not compute the exact likelihood of the model and have to resort to Monte-Carlo sampling for its approximation.
Revised: In contrast to our approach, these methods can not compute the exact likelihood of the model and have to resort to Monte-Carlo sampling for its approximation.

Original: However, this approach is especially problematic for commonly occurring clustered and bursty event sequences as it either requires a very high sampling rate or ignores important temporal dependencies (Nickel & Le, 2020).
Revised: However, this approach is especially problematic for commonly occurring clustered and bursty event sequences as it either requires a very high sampling rate or ignores important temporal dependencies (Nickel & Le, 2020).

Original: To overcome this issue, Jia & Benson (2019) proposed recently Neural Jump SDEs which extend the Neural ODE framework and allow to compute the exact likelihood for neural TPPs.
Revised: To overcome this issue, Jia & Benson (2019) proposed recently Neural Jump SDEs which extend the Neural ODE framework and allow to compute the exact likelihood for neural TPPs, up to numerical errors.

Original: This method is closely related to our approach and we build on its ideas to compute the ground intensity of the STPP.
Revised: This method is closely related to our approach and we build on its ideas to compute the ground intensity of the STPP.

Original: However, current Neural Jump SDEs —as well as NHPs and RMTPPs—are not well-suited for modeling complex continuous mark distributions as they are restricted to methods such as Gaussian mixture models in the spatial domain.
Revised: However, current Neural Jump SDEs —as well as NHPs and RMTPPs—are not well-suited for modeling complex continuous mark distributions as they are restricted to methods such as Gaussian mixture models in the spatial domain.

Original: Finally, Shchur et al.
Revised: Finally, Shchur et al.

Original: ; Mehrasa et al.
Revised: (2019) and Mehrasa et al.

Original: (2019) considered to combine TPPs and CNFs, however for different purposes as in our case, i.e., for intensity-free learning of TPPs.
Revised: (2019) considered combining TPPs with ﬂexible likelihood-based models, however for different purposes as in our case, i.e. , for intensity-free learning of only temporal point processes.

Original: Continuous Normalizing Flows The ability to describe an inﬁnite number of distributions with a Continuous Normalizing Flow has been used by a few recent works.
Revised: Continuous Normalizing Flows The ability to describe an inﬁnite number of distributions with a Continuous Normalizing Flow has been used by a few recent works.

Original: Some works in computer graphics have used the interpolation effect of CNFs to model transformations of point clouds (Yang et al., 2019;
Revised: Some works in computer graphics have used the interpolation effect of CNFs to model transformations of point clouds (Yang et al., 2019;

Original: Rempe et al., 2020; Li et al., 2020).
Revised: Rempe et al., 2020; Li et al., 2020).

Original: CNFs have also been used in sequential latent variable models (Deng et al., 2020;
Revised: CNFs have also been used in sequential latent variable models (Deng et al., 2020;

Original: Rempe et al., 2020).
Revised: Rempe et al., 2020).

Original: However, such works do not align the “time” axis of the CNF with the temporal axis of observations, and do not train on observations at more than one value of “time” in the CNF.
Revised: However, such works do not align the “time” axis of the CNF with the temporal axis of observations, and do not train on observations at more than one value of “time” in the CNF.

Original: In contrast, we align the time axis of the CNF with the time of the observations, directly using its ability to model distributions on a real-valued axis.
Revised: In contrast, we align the time axis of the CNF with the time of the observations, directly using its ability to model distributions on a real-valued axis.

Original: A closely related application of CNFs to spatio-temporal data was done by Tong et al.
Revised: A closely related application of CNFs to spatio-temporal data was done by Tong et al.

Original: (2020), who modeled the distribution of cells in a developing human embryo system at ﬁve ﬁxed time values.
Revised: (2020), who modeled the distribution of cells in a developing human embryo system at ﬁve ﬁxed time values.

Original: In contrast to this, we extend to applications where observations are made at arbitrary time values and further jointly model such distributions with a temporal time process.
Revised: In contrast to this, we extend to applications where observations are made at arbitrary time values, jointly modeling space and time within the spatio-temporal point process framework.

Original: Furthermore, Mathieu & Nickel (2020); Lou et al.
Revised: Furthermore, Mathieu & Nickel (2020); Lou et al.

Original: (2020) recently proposed extensions of CNFs to Riemannian manifolds.
Revised: (2020) recently proposed extensions of CNFs to Riemannian manifolds.

Original: For our proposed approach, this is especially interesting in the context of earth and climate science, as it allows us to model STPPs on the sphere simply by replacing the CNF with its Riemannian equivalent.
Revised: For our proposed approach, this is especially interesting in the context of earth and climate science, as it allows us to model STPPs on the sphere simply by replacing the CNF with its Riemannian equivalent.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: Data Sets As the world naturally moves along a temporal axis, many data sets can be represented as spatio-temporal events.
Revised: Data Sets As the world naturally moves along a temporal axis, many data sets can be represented as spatio-temporal events.

Original: We pre-process data from open sources and make them suitable for spatiotemporal event modeling.
Revised: We pre-process data from open sources and make them suitable for spatio-temporal event modeling.

Original: Varying across a wide range of domains, the data sets we consider are: earthquakes, pandemic spread, consumer demand for a bike sharing app, and high-amplitude brain signals from fMRI scans.
Revised: Varying across a wide range of domains, the data sets we consider are: earthquakes, pandemic spread, consumer demand for a bike sharing app, and high-amplitude brain signals from fMRI scans.

Original: We brieﬂy describe these data sets here; further details and pre-processing steps can be found in Appendix C.
Revised: We brieﬂy describe these data sets here; further details, pre-processing steps, and data set diagnostics can be found in Appendix C.

Original: P INWHEEL This is a synthetic data set with strong dependencies between spatial samples (see ﬁg. 2).
Revised: P INWHEEL This is a synthetic data set with multimodal and non-Gaussian spatial distributions designed to test the ability to capture drastic changes due to event history (see ﬁg. 5).

Original: The data set consists of 10 clusters which form a pinwheel structure.
Revised: The data set consists of 10 clusters which form a pinwheel structure.

Original: Events are sampled from a multivariate Hawkes process such that events from one cluster willdrastically increase the probability of  events in the next cluster in a clock-wise manner.
Revised: Events are sampled from a multivariate Hawkes process such that events from one cluster will increase the probability of observing events in the next cluster in a clock-wise rotation. Number of events per sequences ranges between 4 to 108.

Revised: E ARTHQUAKES For modeling earthquakes and aftershocks, we gathered location and time of all earthquakes in Japan from 1990 to 2020 with magnitude of at least 2.5 from the U.S. Geological

Revised: Survey (2020).

Revised: Number of sequences

Revised: 18 to 543.

Original:  [Figure]
Revised: C OVID -19 C ASES We use data released publicly by The New York Times (2020) on daily COVIDcases in New Jersey state. The data is aggregated at the county level, which we dequantize uniformly across the county. Number of events per sequences ranges between 3 to 323. [Figure]

Original: J u m p
Revised: J u m p

Original: C N F
Revised: C N F

Original: C ond iti on a l
Revised: C ond iti on a l

Original: KD E
Revised: KD E

Original: [Figure]
Revised: [Figure]

Original: E ARTHQUAKES

Original: For modeling earthquakes and aftershocks, we gathered location and time of all earthquakes in Japan from 1990 to 2020 with magnitude of at least 2.5 from the U.S. Geological

Original: Survey (2020).

Original: We split the data into individual sequences using sliding windows of length 30 days.

Original: C OVID -19 C ASES We use data released publicly by The New York Times (2020) on daily COVIDcases in New Jersey state, from March to July of 2020.

Original: The data is aggregated at the county level, which we dequantize uniformly across the county.

Original: We also dequantize the temporal axis by assigning new cases uniformly within the day.

Original: We split the data into sliding windows of 7 days.

Original: B OLD [-
-] Thisdata consists of fMRI scansof 4 participants as they are given visual stimuli (Chang et al., 2019).
Revised: B OLD [-
-]5000 This consists of fMRI scans as participants are given visual stimuli (Chang et al., 2019).

Original: Weuse the sessions of a single patient and convert brain responses into spatio-temporal events following the z-score thresholding approach in Tagliazucchi et al.
Revised: We convert brain responses into spatio-temporal events following the z-score thresholding approach in Tagliazucchi et al.

Original:  In addition to these datasets, we also report in Appendix A results for C ITIBIKE ,which is a data set consisting of rental events in a bike sharing service in New York City.
Revised: Number of events per sequences ranges between 6 to 1741. In addition to these datasets, we also report in Appendix A results for C ITIBIKE , a data set consisting of rental events from a bike sharing service in New York City.

Original: Results To evaluate the capability of our proposed models, we compare against commonly-used baselines and state-of-the models.
Revised: Baselines To evaluate the capability of our proposed models, we compare against commonly-used baselines and state-of-the-art models.

Original: In some settings, ground intensity and conditional mark density are independent of each other and we can freely combine different baselines for the temporal and spatial domains.
Revised: In some settings, ground intensity and conditional mark density are independent of each other and we can freely combine different baselines for the temporal and spatial domains.

Original: As temporal baselines, we use a homogeneous Poisson process, a self-correction process,and a Hawkes process.
Revised: As temporal baselines, we use a homogeneous Poisson process, a self-correction process, a Hawkes process, and the Neural Hawkes Process, which were trained using their ofﬁcially released code.

Original: As spatial baselines, we use a conditional kernel density estimator (KDE), where p ( x | t ) is essentially modeled as a history-dependent Gaussian mixture model (see
Revised: As spatial baselines, we use a conditional kernel density estimator (KDE) with learned parameters, where p ( x | t ) is essentially modeled as a history-dependent Gaussian mixture model (see

Original: Appendix B), as well as the Time-varying CNF (see section 3).
Revised: Appendix B), as well as the Time-varying CNF.

Original: In addition, we also compare to our implementation of Neural Jump SDEs (Jia & Benson, 2019) where the spatial distribution is a 
Revised: In addition, we also compare to our implementation of Neural Jump SDEs (Jia & Benson, 2019) where the spatial distribution is a Gaussian mixture model.

Original: [Table]

Original: Gaussian mixture model, but we improve the  architecture to use our GRU-based continuous-time hidden states for fair comparison, as we found the simpler parameterization in Jia & Benson (2019) to be numerically unstable for long sequences.
Revised: We use the same architecture as our GRU-based continuous-time hidden states for fair comparison, as we found the simpler parameterization in Jia & Benson (2019) to be numerically unstable for large number of events. Range of hyperparameter values are outlined in Appendix D.

Original:  The results of our evaluation are shown in table 1.
Revised: Results The results of our evaluation are shown in table 1.

Original: It can be seen that our proposed neural STPPs achieve strong results and typically outperform the comparison models by a large margin.
Revised: We highlight all results where the intervals containing one standard deviation away from the mean overlap. [Table]

Original: Across all data sets, the Time-varying CNF outperforms the conditional KDE baseline despite not being conditional on history.
Revised: Across all data sets, the Time-varying CNF outperforms the conditional KDE baseline despite not being conditional on history.

Original: This suggests that the overall spatial distribution is rather complex and cannot be modeled with simple Gaussian clusters.
Revised: This suggests that the overall spatial distribution is rather complex. We also see from Figure 7 that Gaussian clusters tend to compensate for far-reaching events by learning a larger band-width whereas a ﬂexible CNF can easily model multi-modal event propagation.

Original: On P INWHEEL and E ARTHQUAKES , the historydependent Jump and Attentive CNF models achievesubstantially better log-likelihoods in both the temporal and spatial domain.
Revised: The Jump and Attentive CNF models achieve better log-likelihoods than the Time-varying CNF, suggesting prediction in these data sets beneﬁt from modeling dependence on event history.

Original: For C OVID -19, the self-exciting Hawkes process is a strong baseline— which aligns with similar results for other infectious diseases (Park et al., 2019) — but the neural
Revised: For C OVID -19, the self-exciting Hawkes process is a strong baseline which aligns with similar results for other infectious diseases (Park et al., 2019), but Neural STPPs can achieve substantially better spatial likelihoods.

Original: STPPs achieve again substantially better spatial likelihoods. When compared to Neural Jump SDEs, our proposed models show similar strong performance. First, it can be seen that on all datasets our modelsshow far better spatial results which illustrates the beneﬁts from using expressive CNFs for the  spatial domain.
Revised: Overall, NHP is competitive with the Neural Jump SDE; however, it tends to fall short of the Attentive CNF which jointly models spatial and temporal variables. In a closer comparison to the temporal likelihood of Neural Jump SDEs (Jia & Benson, 2019), we ﬁnd that overly-restricted spatial models can negatively affect the temporal model since both domains are tightly coupled.

Original: Second, since our realization of Neural Jump SDEs and our STPPs use the same  architecture to model the temporal domain, their temporal likelihoods are often close.
Revised: Since our realization of Neural Jump SDEs and our STPPs use the same underlying architecture to model the temporal domain, the temporal likelihood values are often close.

Original: However, this notably not the case for P INWHEEL and B OLD 5000, which seems to indicate that too restricted spatial models can negatively affect the temporal model since both domains are tightly coupled.
Revised: However, there is still a statistically signiﬁcant difference between our Neural STPP models and Neural Jump SDEs even for the temporal log-likelihood on all data sets.

Original: Finally, we note that the results of the Jump and Attentive CNFs are typically close.
Revised: Finally, we note that the results of the Jump and Attentive CNFs are typically close.

Original: On E ARTHQUAKES , C OVID -19, and B OLD 5000 the attentive model  achieves even the best temporal or spatial likelihood, while being substantially faster to compute (see Appendix A for a runtime analysis).
Revised: The attentive model generally achieves better temporal log-likelihoods while maintaining competitive spatial log-likelihoods. This difference is likely due to the Attentive CNF’s ability to attend to all previous events, while the Jump CNF has to compress all history information inside the hidden state at the time of event. The Attentive CNF also enjoys substantially faster computations (see Appendix A).

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: To learn high-ﬁdelity models of stochastic events occurring in continuous space and time, we have proposed a new class of parameterizations for spatio-temporal point processes.
Revised: To learn high-ﬁdelity models of stochastic events occurring in continuous space and time, we have proposed a new class of parameterizations for spatio-temporal point processes.

Original: Our approach combines ideas of Neural Jump SDEs with Continuous Normalizing Flows and allows to retain the ﬂexibility of neural temporal point processes while enabling highly expressive models of continuous marks.
Revised: Our approach combines ideas of Neural Jump SDEs with Continuous Normalizing Flows and allows to retain the ﬂexibility of neural temporal point processes while enabling highly expressive models of continuous marks.

Original: We leverage Neural ODEs as a computational method what allows  to compute theexact likelihood of the joint model and  show  on spatio-temporal datasets  from a wide range of domains that our approach achieves state-of-the-art performance.
Revised: We leverage Neural ODEs as a computational method that allows computing, up to negligible numerical error, the likelihood of the joint model, and we show that our approach achieves state-ofthe-art performance on spatio-temporal datasets collected from a wide range of domains.

Original: One limitation of our method is currently that calls to ODE solvers cancomputationally be  expensive and preventthus the modeling of long sequences.
Revised: One limitation of our history-dependent models is the reliance on sequentially computing a continuoustime hidden states, which can be computationally expensive and prevent the modeling of extremely large number of events.

Original: Improvements to the scalability of Neural ODEs and RNN hybrids would beneﬁt our approach substantially and are a very promising direction for future work.
Revised: Improvements to the scalability of Neural ODE and RNN hybrids would beneﬁt our approach substantially and are a very promising direction for future work.

Original: Another promising area for future work are applications of our method in earth and climate science which often are concerned with modeling highly complex spatio-temporal data.
Revised: Another promising area for future work are applications of our method in earth and climate science which often are concerned with modeling highly complex spatio-temporal data.

Original: In this context, the use of Riemannian CNFs (Mathieu & Nickel, 2020; Lou et al.,  2020) isalso especially interesting as it allows us to model Neural STPPs on  the sphere simply by  replacing the CNF  with its Riemannian counterpart.
Revised: In this context, the use of Riemannian CNFs (Mathieu & Nickel, 2020; Lou et al., 2020; Falorsi & Forr´e, 2020) is especially interesting as it allows us to model Neural STPPs on manifolds ( e.g. the earth’s surface) by simply replacing the CNF in our models with a Riemannian counterpart.

Original: A A DDITIONAL R ESULTS 
Revised: A A DDITIONAL R ESULTS AND F IGURES

Original: [Figure] 
Revised: [Figure] [Figure]

Original: Citibike NY
Revised: [Table] [Figure] [Table] [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Original: [Figure] 
Revised: [Figure] [Figure]

Revised: x (1)0 x (2)0 x (3)0 x (4) x (1)0 x (2) x (3)0 x (4) t = t =

Revised: t i x (1)0 x (2)0 x (3)0 x (4) x (1)0 x (2) x (3)0 x (4) t

Revised: = t

Revised: =

Revised: t i x (1)0 x (2)0 x (3)0 x (4) x (1)0 x (2) x (3)0 x (4) t

Revised: = t

Revised: =

Original:  B B ASELINE
Revised: t i B B ASELINE

Original: Our self-excitation baseline uses a Hawkes process to model the temporal variable, then uses a
Revised: Our self-excitation baseline uses a Hawkes process to model the temporal variable, then uses a

Original: Gaussian mixture model to describe the spatial distribution conditioned on history of events.
Revised: Gaussian mixture model to describe the spatial distribution conditioned on history of events.

Original: This corresponds to the following likelihood decomposition
Revised: This corresponds to the following likelihood decomposition

Original: [Equation]
Revised: [Equation]

Original:  This dependence structure allows the usage of simple temporal point processes to model t i , e.g. a Hawkes process, since temporal variables do not depend on the spatial information.
Revised: Note that t i does not depend on the spatial variables associated with previous events. This dependence structure allows the usage of simple temporal point processes to model t i , e.g. a Hawkes process, since temporal variables do not depend on the spatial information.

Original: The spatial distribution conditions all past events as well as the current time of occurance.
Revised: The spatial distribution conditions all past events as well as the current time of occurance.

Original: Our baseline model assumes a simple Gaussian conditional model, that new events are likely to appear near previous events.
Revised: Our baseline model assumes a simple Gaussian conditional model, that new events are likely to appear near previous events.

Original: [Equation]
Revised: [Equation]

Original: This parameteric model has two learnable parameters: σ 2 and τ , which control the rate of decay in the spatial and temporal domains, respectively.
Revised: This parameteric model has two learnable parameters: σ 2 and τ , which control the rate of decay in the spatial and temporal domains, respectively.

Original: However, this Gaussian spatial model assumes events are propagated in all directions equally and can only model local self-excitation behavior.
Revised: However, this Gaussian spatial model assumes events are propagated in all directions equally and can only model local self-excitation behavior.

Original: These assumptions are often suitable for simplicity but are generally incorrect for many spatio-temporal data.
Revised: These assumptions are often used for simpliﬁcations but are generally incorrect for many spatio-temporal data.

Original: To name a few, earthquakes occur more frequently along boundaries of tectonic plates, epidemics propagate along trafﬁc routes, taxi demands saturate locally and change as customers move around.
Revised: To name a few, earthquakes occur more frequently along boundaries of tectonic plates, epidemics propagate along trafﬁc routes, taxi demands saturate locally and change as customers move around.

Original: C P RE - PROCESSING STEPS FOR EACH DATA SET
Revised: C P RE - PROCESSING STEPS FOR EACH DATA SET

Revised: P INWHEEL We sample from a multivariate Hawkes process with 10 dimensions.

Revised: We turn this into continuous spatial variables by assigning each dimension to a cluster from a “pinwheel” distribution, and sample from the corresponding cluster for each event.

Original:  E ARTHQUAKES For modeling earthquakes and aftershocks, we gathered location and time of all earthquakes in Japan from 1990 to 2020 with magnitude of at least 2.5 from the U.S. Geological
Revised: Number of events per sequences ranges between 4 to 108. E ARTHQUAKES For modeling earthquakes and aftershocks, we gathered location and time of all earthquakes in Japan from 1990 to 2020 with magnitude of at least 2.5 from the U.S. Geological

Original: Survey (2020).
Revised: Survey (2020).

Original: We removed earthquakes from 2010 November to 2011 December, as these sequences were too long and only served as outliers in the data.

Original: Starting from January 01, 1990, we created sequences with a gap of 7 days.
Revised: Starting from January 01, 1990, we created sequences with a gap of 7 days.

Original: Each sequence was of length 30 days.
Revised: Each sequence was of length 30 days.

Original: We ensured there was no contamination between train/val/test sets by removing intermediate sequences.
Revised: We ensured there was no contamination between train/val/test sets by removing intermediate sequences.

Revised: We removed earthquakes from 2010 November to 2011 December, as these sequences were too long and only served as outliers in the data.

Original: This resulted intraining sequences, 50 validation sequences, and 50 test sequences.
Revised: This resulted in 950 training sequences, 50 validation sequences, and 50 test sequences.

Original:  C OVID -19 C ASES We use data released publicly by The New York Times (2020) on daily COVIDcases in the New Jersey state, from March to July of 2020.
Revised: Number of events per sequence ranges between 18 to 543. C OVID -19 C ASES We use data released publicly by The New York Times (2020) on daily COVIDcases in the New Jersey state, from March to July of 2020.

Original: The data is aggregated at the county level, which we dequantize uniformly across the county.
Revised: The data is aggregated at the county level, which we dequantize uniformly across the county.

Original: We also dequantize the temporal axis by assigning new cases uniformly within the day.
Revised: We also dequantize the temporal axis by assigning new cases uniformly within the day.

Original: Starting at March 15, and every 3 days, we took a 7 day length sequence.
Revised: Starting at March 15, and every 3 days, we took a 7 day length sequence.

Original: For each sequence, we sampled each event with a probability of 0 .
Revised: For each sequence, we sampled each event with a probability of 0 .

Original: This was donetimes per sequence.
Revised: This was donetimes per sequence.

Original: We ensured there was no contamination between train/val/test sets by removing intermediate sequences.
Revised: We ensured there was no contamination between train/val/test sets by removing intermediate sequences.

Original: This resulted in 1450 training sequences, 100 validation sequences, andtest sequences. 
Revised: This resulted in 1450 training sequences, 100 validation sequences, andtest sequences. Number of events per sequence ranges between 3 to 323.

Original: C ITIBIKE Citibike is a bike sharing service in New York City.
Revised: C ITIBIKE Citibike is a bike sharing service in New York City.

Original: We treat the start of each trip as an event, and use the data from April to August of 2019.
Revised: We treat the start of each trip as an event, and use the data from April to August of 2019.

Original: We split into sequences of length 1 day starting at 5:00am of each day.
Revised: We split into sequences of length 1 day starting at 5:00am of each day.

Original: For each sequence, we subsampled with a probability of 0 .
Revised: For each sequence, we subsampled with a probability of 0 .

Original: 005 per event,times.
Revised: 005 per event,times.

Original: This resulted in 2440 training sequences, 300 validation sequences, and 320 test sequences.
Revised: This resulted in 2440 training sequences, 300 validation sequences, [Table] [Table]

Original:  B OLD 5000 This consists of fMRI scans of four participants as they are given visual stimuli (Chang et al., 2019).
Revised: [Figure] e n [Figure] B OLD 5000 This consists of fMRI scans of four participants as they are given visual stimuli (Chang et al., 2019).

Original: We use the sessions of a single patient and for each run, we split into 3 equal length sequences.
Revised: We use the sessions of a single patient and for each run, we split into 3 sequences, treated individually.

Original: We converted brain responses into spatio-temporal events following the z-score thresholding approach in Tagliazucchi et al. 2016, Equation (2).
Revised: We converted brain responses into spatio-temporal events following the z-score thresholding approach in Tagliazucchi et al. 2016, Equation (2).

Original: We used a threshold of γ = 6 . 0 .
Revised: We used a threshold of γ = 6 . 0 .

Original: We split the data into 1050 training sequences, 150 val sequences, 220 test sequences. 
Revised: We split the data into 1050 training sequences, 150 val sequences, 220 test sequences. Number of events per sequence ranges between 6 to 1741.

Revised: Each data set contains sequences with highly variable number of events, with varying degrees of dependence between events, making them difﬁcult to model with traditional point process models.

Revised: We plot histograms showing the number of events per sequence in Figure 11.

Original: D R EMOVING CROSS - EVENT PARTIAL DERIVATIVES
Revised: D H YPERPARAMETERS C HOSEN AND T ESTED

Revised: For the time-varying, jump, and attentive CNF models, we parameterized the CNF drift as a multilayer perceptron (MLP) with dimensions [ d − 64 − 64 − 64 − d ] , where d is the number of spatial variables. We swept over activation functions between using softplus or a time-dependent Swish (Ramachandran et al., 2017). TimeDependentSwish

Revised: [Equation]

Revised: (22) where σ is the logistic sigmoid function, (cid:12) is the Hadamard (element-wise) product, and β : R → R d z is a MLP with widths [1 − 64 − d z ] where d z is the dimension of z , using the softplus activation function.

Revised: We ultimately decided on using the time-dependent Swish for all experiments.

Revised: We swept over the MLP for deﬁning f h for the continuous-time hidden state in eq.

Revised: (11) using hidden widths of [8 − 20] , [32 − 32] , [64 − 64] , [32 − 32 − 32] , and [64 − 64 − 64] .

Revised: The majority of models used 32 − 32 as it provided enough ﬂexibility while remaining easy to solve.

Revised: We used the softplus activation function.

Revised: We tried MLP for parameterizing the instantaneous change in eq.

Revised: (12); however, it was too unstable for long sequences.

Revised: We therefore switched to the GRU parameterization, which takes an input (new event), the hidden state at the time of event, and outputs a new hidden state.

Revised: We regularized the L 2 norm of the hidden state drift with a strength of 1e-4, chosen from { 0, 1e-4, 1e-3, 1e-2 } .

Revised: We optionally used optimal transport-inspired regularization from Finlay et al.

Revised: (2020), which adds a Frobenius norm regularization to the gradient of the drift in addition to the L 2 norm regularization, to the CNF models with a strength of { 0, 1e-4, 1e-3, 1e-2 } .

Revised: The Time-varying and

Revised: Attentive CNF models did not require regularization and were mostly kept at 0, but the Jump CNF models beneﬁted from some amount of regularization to avoid numerical instability.

Revised: To model a non-trivial spatial distribution for the entire data interval, we shift the data interval to start at t = 2 for all CNF models.

Revised: Thus the interval used for parameterizing the CNF is [2 , T + 2] .

Revised: Generally, the “time” variable is a dummy one; we can place the base distribution at any time, and we can choose any interval on the real line to be the data interval; this does not limit the model in any way.

Revised: For the Jump CNF, we used a composition of 4 radial ﬂows (Rezende & Mohamed, 2015) to parameterize the instantaneous updates in eq.

Revised: All parameters of the radial ﬂows were parameterized to be the output of a MLP that takes as input the hidden state at the time of the event (before the hidden state is updated based on the current event).

Revised: The radial ﬂows were initialized in such a way that the log determinant is near zero.

Revised: For the Attentive CNF, the drift function consists of

Revised: Time-dependent MLP( d − 64 − 64 )

Revised: → 2 × MultiheadAttention → Time-dependent MLP( 64 − 64 − d ) where the Time-dependent MLPs make use of the TimeDependentSwish.

Revised: As was done in Vaswani et al.

Revised: (2017), the multihead attention is used within a residual branch, except we swapped LayerNorm

Revised: (Ba et al., 2016) with ActNorm (Kingma & Dhariwal, 2018) as LayerNorm has an unbounded Lipschitz and can be ill-suited for use in ODEs.

Revised: We tested both standard multihead attention (Vaswani et al., 2017) and the Lipschitz multihead attention (Kim et al., 2020).

Revised: The Lipschitz multihead attention typically produced similar validation NLL as the standard multihead attention but were more stable on multiple occassions.

Revised: We therefore kept the Lipschitz multihead attention for all experiments.

Revised: We additionally, use an auxiliary (non-attentive, simply with the two multihead attention layers removed)

Revised: CNF to map from 0 ( i.e. the time of base distribution) to the beginning of the data interval.

Revised: We initialized all Neural ODEs (for the hidden state and CNFs) with zero drift by initializing the weights and biases of the ﬁnal layer to zero.

Revised: The log-likelihood values reported are after the spatial variables have been standardized using the empirical mean and standard deviation from the training set.

Revised: We train and test on log-likelihood (in nats) per event, which normalizes eq.

Revised: (8) of each sequence by the number of events.

Revised: All integrals were solved using Chen (2018) to within a relative and absolute tolerance of 1E-4 or 1E-6, chosen based on preliminary testing for convergence and stability.

Revised: Our implementation of the Neural Jump SDE shares the same continuous-time hidden state parameterization but uses a mixture of Gaussians as the spatial model.

Revised: We used 5 mixtures, and a MLP that maps from the hidden state to the parameters of this mixture of Gaussians (the means, log standard deviations, and mixture coefﬁcients).

Original: This results in a lower-variance gradient estimator for training, and allows parallel computation of conditional log probabilities at test time.
Revised: This results in a lower-variance gradient estimator for training, and allows parallel computation of conditional log probabilities at test time.

Original: We ﬁrst summarily describe the attention mechanism.
Revised: We ﬁrst summarily describe the attention mechanism.

Original: For an input X ∈ R n × d representing the sequence of n variables { x (0) s , . . .
Revised: For an input X ∈ R n × d representing the sequence of n variables { x (0) s , . . .

Original: , x ( n ) s } at some value of s , this attention mechanism creates logits
Revised: , x ( n ) s } at some value of s , this attention mechanism creates logits

Original: P ∈ R n × n and values V ∈ R n × d , both dependent on X , such that the output is 
Revised: P ∈ R n × n and values V ∈ R n × d , both dependent on X , such that the output is O = softmax ( P ) [Equation] V

Original: [Equation]
Revised: [Equation]

Original:  where the softmax is taken over each row of P .
Revised: (23) where the softmax is taken over each row of P .

Original: The output is then added to X as a residual connection.
Revised: The output is then added to X as a residual connection.

Original: The multihead attention computes P in a way such that P ij depends on X i and X j , and V i depends on X i .
Revised: The multihead attention computes P in a way such that P ij depends on X i and X j , and V i depends on X i .

Original: This is true for both the vanilla MHA (Vaswani et al., 2017) and the L2 MHA
Revised: This is true for both the vanilla MHA (Vaswani et al., 2017) and the L2 MHA

Original: (Kim et al., 2020).
Revised: (Kim et al., 2020).

Original: For our use case, P ij is set to − inf for j > i as we don’t want to attend to future events.
Revised: For our use case, P ij is set to − inf for j > i as we don’t want to attend to future events.

Original: We retain only the block-diagonal gradients.
Revised: We retain only the block-diagonal gradients where each block contains variables corresponding to one event.

Revised: This is equivalent to removing all the cross-event dependencies.

Original: [Equation]
Revised: [Equation]

Revised: F P ARALLEL S OLVING OF M ULTIPLE ODE S WITH V ARYING I NTERVALS

Revised: Our numerical ODE solvers integrate a single ODE system dxdt =

Revised: f ( t, x ) , where x ∈ R d and f : R 1+ d → R d , on a single ﬁxed interval [ t start , t end ] .

Revised: We can express the inputs and outputs of an

Revised: ODE solver with

Revised: [Equation]

Revised: t end t start

Revised: [Equation]

Revised: where x 0 is a vector containing the initial state at the initial time t 0 .

Revised: Multiple ODEs Now suppose we have a set of M systems ( i.e. dx m dt = f m for m = 1 , . . .

Revised: , M ) that we would like to solve.

Revised: If all systems had the same initial time and desired output time of t start and t end , we can readily create a joint system

Revised: [Equation]

Revised: that follows

Revised: [Equation]

Revised: Solving this joint system can be done in parallel with a single call to ODESolve :

Revised: [Equation]

Revised: which computes x m ( t 1 ) for all m = 1 , . .

Revised: , M .

Revised: This is the standard method used for solving a batch of Neural ODEs.

Revised: Adding dependencies is straightforward Note that extending this further, the ODE systems do not need to be independent.

Revised: They can depend on other variables at the same concurrent time value because the joint system is still an ordinary differential equation.

Revised: For instance, we can have

Revised: [Equation]

Revised: (28) where f m : R 1+ Md → R d .

Revised: Each system is now a partial differential equation, but the joint system x joint is still an ODE and can be solved with one call to ODESolve .

Revised: Varied time intervals Now suppose each system has a different time interval that we want to solve.

Revised: Different initial times and different end times.

Revised: Let’s denote the start and end time for the m -th system as t ( m ) start and t ( m ) end respectively.

Revised: We can construct a dummy variable that always integrates from 0 to 1, and perform a change of variables (reparameterization) to transform every system to use this dummy variable.

Revised: As a concrete example of this reparameterzation procedure, consider just one system x ( t ) with drift function f ( t, x ) that we want to integrate from t start to t end with the initial value x 0 .

Revised: We can transform x ( t ) using the relation s = t − t start t end − t start , or equivalently

Revised: [Equation]

Revised: [Equation]

Revised: into a solution ˜ x ( s ) on the unit interval [0 , 1] such that

Revised: [Equation]

Revised: The drift function for ˜ x then follows as

Revised: [Equation]

Revised: Now since ˜ x (0) =

Revised: x ( t start ) and ˜ x

Revised: (1) =

Revised: x ( t end ) , the following are equivalent

Revised: [Equation]

Revised: Putting it all together Let ˜ x m ( s ) be the reparameterized solution for x m ( t ) such that

Revised: [Equation]

Revised: [Equation]

Revised: We can then solve for all M systems, with different varying time intervals, using

Revised: [Equation]

Revised: joint

Revised: [Equation]

Revised: that follows

Revised: [Equation]

Revised: Solving this system to s = 1 yields ˜ x m (

Revised: [Equation]

Revised: = x

Revised: [Equation]

Revised: [Equation]

Revised: Assuming t ( m ) start = 0 for all m = 1 , . .

Revised: , M in order to reduce notational complexity, we can write this joint system in terms of the original systems as

Revised: [Equation]

Revised: This is the joint system written in equation 14.

Revised: The joint system in equation 19 adds dependence between the M systems but can still be solved with a single ODESolve .