Learning Mixture of Neural Temporal Point Processes for Multi-dimensional Event Sequence Clustering

1 Introduction

Many real-world processes can be characterized by timestampedmulti-typed 1 event sequences, e.g. e-Commerce [Xu et al., 2014] and electronic health records [Enguehard et al., 2020]. Given a set of event sequences, a natural task is to cluster them into groups such that sequences generated by similar underlying dynamics are grouped. Multi-dimensional event sequence clustering is important to downstream applications. For instance, clustering user's behaviors in eCommerce helps to promote recommendation system.

∗Correspondence author: Junchi Yan. The work was in part supported by Shanghai Municipal Science and Technology Major Project (2021SHZDZX0102), National NSFC (U19B2035, 61972250, 72061127003), and funding from Ant Group.

1 We use the term 'multi-dimensional' & 'typed' interchangeably.

Despite the importance, research on event sequence clustering (especially for multi-dimensional events i.e. typed events) is still limited. [Xu and Zha, 2017] propose a Dirichlet Mixture model of Hawkes Processes (DMHP) where event sequences are characterized via Hawkes processes [Hawkes, 1971] with different parameters. However, as a traditional Temporal Point Process (TPP), it is hard for Hawkes process to model complex real-world dynamics due to its limited parametric assumption. The recent work [Wu et al., 2022] takes a deep reinforcement learning view of event sequence clustering. But it only works for uni-dimensional event sequences and its generative adversarial learning framework can be inefficient and less stable.

With the development of deep learning, many Neural Temporal Point Processes (NTPPs) have been developed. Recurrent Neural Networks (RNN) [Du et al., 2016; Mei and Eisner, 2016], Neural Ordinary Differential Equations [Jia and Benson, 2019; Chen et al., 2021] and Transformers [Zhang et al., 2020; Zuo et al., 2020] are used to model the conditional intensity function of event sequences. [Shchur et al., 2020; Sharma et al., 2021; Omi et al., 2019] propose intensityfree NTPPs to directly model the conditional distribution of inter-event times or cumulative hazard function. Many Maximum Likelihood Estimation free (MLE-free) NTPPs are also developed [Xiao et al., 2017; Upadhyay et al., 2018; Yan et al., 2018; Li et al., 2018]. Although these NTPPs are more flexible than Hawkes process, they can not perform clustering because they assume all observed sequences are generated by similar dynamics to which they strike to fit.

We propose mixture of neural temporal point processes (NTPP-MIX), a general framework that can incorporate various MLE-based NTPPs for multi-dimensional event sequence clustering. Specifically, the prior of the mixing coefficients is modeled by a Dirichlet distribution. Then the coefficients are used to model the prior of cluster assignment. For each sequence, its conditional probability on given assignment is modeled via a series of NTPPs. Our goal is to maximize the marginal probability and get the posterior of cluster assignment. Therefore we introduce variational distribution to approximate the posterior and combine variational EM and Stochastic Gradient Descent (SGD) to efficiently train the framework. In E-step, we fix parameters for NTPPs and use variational distributions to approximate the true posterior. While in M-step, we fix variational distributions and update parameters for NTPPs using SGD. The contributions are:

1) We develop a general framework NTPP-MIX for multidimensional event sequence clustering using neural networks. To our best knowledge, it is the first neural-network-based framework for multi-dimensional event sequence clustering and can utilize various previous MLE-based NTPPs as its inside module. To put it in the context, the work [Xu and Zha, 2017] is limited to the parametric model, neural model [Wu et al., 2022] can only handle uni-dimensional events.

2) We propose a training algorithm based on variational EM and SGD to efficiently train the framework and give a theoretical proof of its correctness and convergence.

3) Extensive experimental results on both synthetic and real-world benchmarks show the effectiveness of our NTPPMIX framework. Source code will be made public available.

2 Preliminaries and RelatedWorks

Traditional temporal point processes Temporal Point Process (TPP) [Daley and David, 2007] is a useful tool to model a set of similar event sequences. We denote a multidimensional event sequence as S = {(ti, vi)}Li=1 , where ti ∈ [0, T) is timestamp and vi ∈ {1, 2, . . . , C} is event type. A TPP first defines the conditional intensity function:

where Nu(t) denotes the counting process which counts the number of events of type u until time t, Ht = {(ti, vi) : ti < t} represents the history up to t. The log likelihood of S is:

Most traditional TPP models design a parametric form for intensity, which is often trained via MLE [Yan et al., 2016].

Neural temporal point processes Neural Temporal Point Processes (NTPPs) have been recently proposed to capture complex dynamics. [Du et al., 2016] use RNNs to give a flexible form of intensity. [Mei and Eisner, 2016] further improve NTPP by constructing continuous-time RNNs. [Jia and Benson, 2019; Chen et al., 2018] use the continuous-depth Neural Ordinary Differential Equations [Chen et al., 2018] to achieve more flexible continuous-time state evolution. Transformer or Attention mechanism is used in [Zhang et al., 2020; Zuo et al., 2020] to capture long-term dependency.

The above NTPPs model the conditional intensity, while there are intensity-free NTPPs. [Omi et al., 2019] model the integral of the intensity. [Shchur et al., 2020; Sharma et al., 2021] directly model the distribution of inter-event times.

Both intensity-based and intensity-free NTPPs are trained via MLE. Besides MLE, there are also MLE-free [Xiao et al., 2017; Yan et al., 2018; Li et al., 2018] NTPPs.

Although these NTPPs are more flexible than traditional TPPs, they can not directly perform clustering as they assume observed sequences are generated by similar dynamics.

Event sequence clustering methods. Although event sequence modeling is well-studied, there are few studies about event sequence clustering. [Wu et al., 2022] propose a reinforcement learning model but it only works for unidimensional event sequences. To our best knowledge, the most related work to ours is DMHP [Xu and Zha, 2017] which mixes Hawkes processes for clustering. Clusters are defined as Hawkes processes with different parameters. However, real-world data may be complex and can not be well modeled by Hawkes processes, e.g. a new event can inhibit other events instead of the triggering assumption of Hawkes.

3 NTPP-MIX: Mixture of NTPPs for Multi-dimensional Event Clustering

In this section, we propose mixture of neural temporal point processes (NTPP-MIX) framework. It combines NTPP and variational inference for multi-dimensional event clustering.

3.1 Model Framework

Given multi-dimensional event sequences S = {Sn}Nn=1 , where Sn = {(ti, vi)}Ln i=1 is a sequence defined in Sec. 2. Our goal is to divide these N sequences into K groups such that the sequences generated by similar dynamics are grouped. For each Sn, there is a corresponding one-hot vector zn ∈ {0, 1}K denoting the group it belongs to. Denote the hidden variables as Z = {z1, . . . , zN} and our goal is to give a probabilistic estimation of Z. We combine NTPPs and mixture model to enable a probabilistic framework for this clustering problem. Our NTPP-MIX is general and can utilize various NTPPs as its basic component as shown in Fig. 1.

We first assume the prior of mixing coefficients π follows:

which is a Dirichlet distribution, where α0 = [α0, . . . , α0] ∈ RK . Setting a prior for π results in a variational model that helps to select K [Bishop and Nasrabadi, 2007].

Given π, the conditional distribution of Z is formed as: p(Z|π) = ∏Nn=1 ∏Kk=1 πznk k where πk is the k-th dimension of π. This equals to p(znk = 1|π) = πk. Given Z, we use K NTPPs to model the conditional probability of S:

where θ = {θk}Kk=1 and each θk denotes parameters of the k-th NTPP model. For an event sequence Sn, we input it into the k-th NTPP, and get the conditional probability p(Sn|θk). Note that NTPP here can be any MLE-based NTPP such as RMTPP [Du et al., 2016], NHP [Mei and Eisner, 2016], SAHP [Zhang et al., 2020], THP [Zuo et al., 2020], LogNormMix [Shchur et al., 2020], etc.

Summarizing above equations, we can get the joint distribution p(S,Z,π|θ). Our goal is to maximize marginal distribution p(S|θ) w.r.t. θ and get the posterior p(Z|S, θopt):

Note Z is discrete and we use integration for notation brevity.

3.2 Training via Variational EMand SGD

It is hard to optimize marginal probability p(S|θ) and get the posterior p(Z|S, θopt) directly as Eq. 5 contains complex integral terms and parameters for networks. To address this problem, we combine variational EM and SGD for training.

We introduce a variational distribution q(Z, π) to approximate p(Z, π|S, θ). Using q(Z, π), we can get:

where we define:

Therefore, maximizing L(q, θ) with respect to q and θ is equivalent to maximizing log p(S|θ) and minimizing KL(q‖p). Maximizing log p(S|θ) makes our model characterize observed sequences well. Minimizing KL(q‖p) approximates the posterior. Our goal becomes maxθ,q L(q, θ).

In variational inference, L(q, θ) is called evidence lower bound (ELBO) [Hoffman et al., 2013]. Using mean field approximation [Opper and Saad, 2001], we assume q(Z, π) = q(Z)q(π) and maximize L(q, θ) with respect to q(Z), q(π), θ iteratively according to the following process.

Update q(Z) (E-step) Fixing q(π) and θ, viewing terms irrelevant to q(Z) as constant, we can derive:

where we define a new distribution with log ˜ p(S,Z|θ) = Eq(π)[log p(S,Z,π|θ)], the KL divergence is minimized by setting q∗ (Z) = ˜ p(S, Z|θ):

Normalizing the above formulation, we obtain:

Note that q(Z) is an approximation of p(Z|S, θ), when the model is well-trained, p(znk = 1|Sn, θopt) ≈ rnk, i.e. the posterior probability that Sn in group k is rnk.

Update q(π) (E-step). Fixing q(Z) and θ, viewing terms irrelevant to q(π) as constant, we can derive:

Similarly, log ˜ p(S, π|θ) = Eq(Z)[log p(S, Z, π|θ)], the KL divergence is minimized by setting q∗ (π) = ˜ p(S, π|θ):

Normalizing the above formulation, we recognize q∗ (π) is a Dirichlet distribution:

Update θ (M-step). Fixing q(Z) and q(π), viewing terms irrelevant to θ as constant, we can derive:

where L(θ) denotes the term in L(q, θ) that is only related to θ. As θ is the parameter for NTPPs, it is hard to obtain the optimal value of θ like q(Z), q(π). Instead, we use SGD with

-L(θ) as objective function for optimization. Considering the problem is non-convex, we perform M > 1 steps of SGD:

Although L(q, θ) can not reach the maximum w.r.t. θ like the update of q(Z), q(π), it is still reasonable to assume that L(q, θ) gets greater after M steps of SGD.

As each of the three steps w.r.t. q(Z), q(π), θ makes L(q, θ) greater, performing them one after another, we get:

where q(t) (Z), q(t) (π), θ(t) denote distribution/parameters after iteration t. This suggests that after each iteration, L(q, θ) gets greater than that of the last iteration. Performing these three steps iteratively, we can get the optimal solution.

Mini-batch training. Mini-batch training is often used to train NTPPs efficiently. We also incorporate it into our training algorithm. It is natural to update q(Z) and θ with a batch of samples, but not for q(π) as it is a global distribution for all samples. Inspired by the momentum technique [Sutskever et al., 2013], we update αk (parameter for q(π)) by:

where S′ is a batch of sequences sampled from the whole dataset. αS′ k is the parameter estimated on this batch. η is the update rate that conrtols how much to update with this batch.

Update the cluster number K. Thanks to the variational framework, NTPP-MIX has a natural tendency to set some mixture weights close to zero and choose a suitable number of effective components automatically. When K is unknown, we initialize K as a large number and remove redundant clusters during training. After each iteration of E-step and Mstep, we can get q∗ (π) and its parameter α. As q(π) is an approximation of p(π|S, θ), a small αk indicates the effective number of sequences corresponding to cluster k is small, so we remove it. The removement also helps to speed up the training process as the number of NTPPs to fit is decreasing.

Detailed training algorithm is shown in Algorithm 1.

4 Experiments

4.1 Protocols

We conduct experiments on 4 synthetic datasets and 3 realworld datasets to evaluate the performance of NTPP-MIX.

Synthetic datasets We generate event sequenceswith 5 different processes: Homo Homogeneous Possion process. Inhomo In-homogeneous Possion process. Inhibit Process in which events of different types inhibit each other. Excite Process in which events excite each other (Hawkes). In&Ex Process in which relation is either inhibition or excitation. Each process generates 4,000 event sequences, and each sequence contains Ln = 50 events with number of event types

Algorithm 1: Training algorithm for NTPP-MIX

Input: Sequences S = {Sn}Nn=1 . Initial number of clusters K (to be automatically updated). Mini-batch size B. Training epochs T . Number of SGD steps to perform in each M-step M. q(π) update rate η. Minimum cluster size ϵ. Weights for prior p(π): α0. Output: Cluster assignment q∗ (Z) ≈ p(Z|S, θopt). 1 Initialize {θk}Kk=1 for K NTPPs. Initialize {αk = α0+N K }Kk=1 for q(π); 2 Pre-train each NTPP on S to prevent bad initialization; 3 for epoch = 1, · · · , T do 4 for iter = 1, · · · , ⌈ NB ⌉ do 5 Sample a batch of sequences S′ from S; // E-step: update q(Z′ ) 6 Eq(π)[log πk] = ψ(αk) - ψ(∑Kk=1 αk) 7 log p(Sn|θk) = log[NTPP(Sn|θk)] Sn ∈ S′ 8 rnk = exp(Eq(π)[logπk]+log p(Sn|θk)) ∑Kk=1 exp (Eq(π)[log πk]+log p(Sn|θk)) // E-step: update q(π) 9 αk = (1 - η)αk + η(α0 + N |S′ | ∑Sn∈S′ rnk) // M-step: update θ 10 for step = 1, · · · ,M do 11 ˜ L(θ) = ∑Sn∈S′ ∑Kk=1 rnk log p(Sn|θk); 12 Update θ using SGD with - ˜ L(θ) as loss; // Remove redundant clusters 13 ∀k, remove θk and distribute αk to other clusters equally if αk < ϵ; 14 Compute rnk, n ∈ {1, . . . , N}, k ∈ {1, . . . , K} using the trained model (line 6 to 8); 15 Return q∗ (znk = 1) = rnk;

C = 5.We merge above sequences to generate 4 datasets with different ground truth number of clusters (KGT ):

KGT = 2: Homo + In-homo; KGT = 3: Homo + In-homo + Inhibit; KGT = 4: Homo + In-homo + Inhibit + Excite; KGT = 5: Homo + In-homo + Inhibit + Excite + In&Ex.

Real-world datasets We use 3 real-world datasets:

1. Stock 2 contains daily stock prices of 31 companies. We extract three types of event: 'up', 'down', 'unchanged' from each stock. We further partition sequences by every season and obtain 1,488 sequences with average length of 49.
2. eCommerce 3 contains users' behavior in a cosmetics online store. Each user's behaviors are categorized into four types: 'view', 'cart', 'remove-from-cart', 'purchase'. This dataset contains 6,200 sequences with average length of 64.

3) Neuron 4 contains spike record of 219 M1 neurons. Ev-

• ery time a neuron spikes is recorded as an event, these events form a uni-dimensional event sequence. This dataset contains

2 www.kaggle.com/szrlee/stock-time-series-20050101-to-20171231

3 www.kaggle.com/mkechinov/ecommerce-events-history-in-cosmetics-shop

4 www.kordinglab.com/spykes/index.html

3,718 sequences with average length of 265.

Models for comparison. We compare the following clustering methods on both synthetic and real-world datasets:

1. HKS+BGM. Learn a specific Hawkes process for each event sequence and apply Bayesian Gaussian Mixture (BGM) model to learned parameters for clustering.

2) ADM4+BGM. Learn a specific ADM4 [Zhou et al., 2013] for each sequence and apply BGM to the learned parameters. 3) NPHC+BGM. Learn Non Parametric Hawkes Cumulant (NPHC) [Xu et al., 2016] for each sequence and apply BGM. 4) DIS+SC. Define distance for event sequence [Iwayama et al., 2017], apply Spectral Clustering to the distance matrix. 5) DMHP [Xu and Zha, 2017]. This is the most related method to our work, which mixes several Hawkes processes to generate a Dirichlet Mixture Model of Hawkes Processes. 6) RMTPP-MIX. Our NTPP-MIX with RMTPP [Du et al., 2016] as components. RMTPP uses RNN to model intensity. 7) SAHP-MIX. Our NTPP-MIX with SAHP [Zhang et al., 2020] as components. It uses Transformer to model intensity. 8) LNM-MIX. Our NTPP-MIX with LNM [Shchur et al., 2020] as components. LNM is intensity-free and MLE-based.

The first four methods are two-step pipelines which extract features and use them for clustering. DMHP and three NTPPMIXs are joint models which perform clustering directly.

We implement all three versions of NTPP-MIX using PyTorch. On all datasets, hyper-parameters in Algorithm 1 are set as: B = 128, T = 100,M = 5, α0 = 1K , pre-train epoch is set to 5. On synthetic datasets, η = 0.01, ϵ = N 5∗K , Adam with lr = 0.005 is used for SGD; On real-world datasets, η = B N , ϵ = N 10∗K , Adam with lr = 0.0005 is used.

Clustering metrics. We use three metrics: the first two for synthetic datasets following [Xu and Zha, 2017; Wu et al., 2022], the third for real-world datasets where the ground truth assignment is unknown following [Zhang and Yan, 2021].

1. Clustering Purity (CP). Ratio of correct assignments.
2. Adjusted Rand index (ARI). Similarity of learned and ground truth assignments [Nguyen et al., 2010].

3) Expected Log Likelihood (ELL). The expected log likelihood w.r.t. cluster assignment on testing set, i.e. we randomly split the data set into training and testing sets, train the model on training set, fix q(π), θ and estimate the cluster assignments q(Ztest) for testing set, then report Eq(Ztest)[log p(Stest|Ztest, θ)]. A model can get higher ELL on testing set by recovering the real distribution.

4.2 Experiments on Synthetic Datasets

We first assume the ground truth cluster number KGT is given and assign K = KGT for each method. Table 1 shows the CP and ARI. In general, joint models (DMHP and NTPPMIX) outperform two-step pipelines. As feature extraction and clustering of two-step models are relatively independent, the extracted feature may not be suitable for clustering. Performance of DIS+SC is satisfactory on KGT = 2, 3 but drops sharply on KGT = 4, 5. This shows it is hard to define a universal distance for typed event sequences.

For joint models, our NTPP-MIXs constantly outperform DMHP, no matter what NTPP we use. The mixture components in DMHP are Hawkes processes, i.e. different clusters are defined as Hawkes processes with different parameters. This limits model capacity as dynamic behind data can be much more complex. Moreover, when we add sequences generated by Hawkes processes in KGT = 4, performance of DMHP drops suddenly and the clustering result collapses to one cluster. The mixture components of our NTPP-MIXs are neural-network-based NTPPs, which are more general and flexible. The standard deviation also shows our NTPP-MIXs are more stable than DMHP. We further visualize similarity matrices of some models on KGT = 5 in Fig. 2. It can be

90

seen that result of LNM-MIX is the closest to ground truth.

We further compare NTPP-MIXs' ability to choose the appropriate number of clusters with DMHP. We set the initial cluster number K = 2KGT , and run 10 trials with different random initialization. The selected number of clusters is shown in Table 2. Note that number of clusters selected by our NTPP-MIX is closer to ground truth and more stable.

4.3 Experiments on Real-world Datasets

For real-world datasets, we randomly sample 60% sequences for training and use the other 40% for testing. Table 3 reports the ELL on testing set, which shows that NTPP-MIXs outperform baselines by a large margin. Two-step models perform poorly and are not as interpretable as joint models: it is hard to explain what we get by directly clustering parameters of Hawkes processes. DMHP also performs not very well due to the strong parametric assumption. Moreover, DMHP collapses on Neuron and assigns all sequences to one cluster.

Fig. 3 shows clustered sequences and corresponding intensities of Neuron using SAHP-MIX: in Cluster #1, when an event occurs, the intensity increases first then decreases like traditional Hawkes process. In Cluster #2, the intensity has a small constant base and spikes sharply when an event occurs. In Cluster #3, intensity decreases when no event occurs, while a new event will pull the intensity back to the baseline.

4.4 Parameter Sensitivity Study

We perform sensitivity analysis of our training algorithm on Dataset KGT = 5. Number of SGD steps in M-step (M in Eq. 15): In Fig. 4(a), when M = 1 all three NTPP-MIXs perform not so well, as we have stated in Eq. 15 that it needs multiple SGD steps. When M ≥ 3, all models converge and get close performance. q(π) update rate (η in Eq. 16): η controls how much update is made to q(π) with a batch of data. We range η from 0.001 to 0.1 and get Fig. 4(b). RMTPPMIX performs not so well when η is too large or small, while SAHP-MIX and LNM-MIX are stable w.r.t η.

5 Conclusion

We have proposed NTPP-MIX, a framework that can utilize existing NTPPs for multi-dimensional event sequence clustering. NTPP-MIX can be efficiently trained by our devised combination of variational EM algorithm and SGD. This is the first neural model for multi-dimensional events clustering, and it is more flexible than traditional parametric models. Experiments show the effectiveness of NTPP-MIX in terms of fitting complex dynamics of event sequence, against state-ofthe-art methods on public datasets.