Mamba Hawkes Process

Hawkes processes are widely used for temporal prediction in various fields. To enhance their performance, many deep learning approaches have been applied. [7] introduced the Recurrent Marked Temporal Point Process (RMTPP) model, which uses recurrent neural networks (RNNs) to learn the influence of event history on the intensity function. [31] employed two RNNs to model event sequences: one for background intensity and another for the impact of historical events, enabling effective end-to-end training. Similarly, [24] proposed a continuous-time LSTM model to capture the self-modulating nature of Hawkes processes, addressing the inhibiting and exciting effects of prior events.

With the development of self-attention mechanisms, self-attention-based neural Hawkes processes were proposed. [36] utilized self-attention to enhance these processes, while [39] used the transformer encoder to convert sequence data into continuous conditional intensity functions. UTHP [34] incorporated RNNs and CNNs to address issues in THP, such as parallel processing and recursive learning. TAA-THP [35] improved attention structures by incorporating temporal encoding. Lastly, [32] proposed Hawkesformer, linking hierarchical attention mechanisms to Hawkes processes for information cascade prediction.

State space models (SSMs) were initially developed as mathematical tools to describe dynamic systems in modern control theory. With the introduction of HiPPO [12] initialization, the Linear State-Space Layer (LSSL) [14] demonstrated the capability to handle long-range dependencies in sequential data. However, LSSL’s computational and memory overhead make it impractical for large-scale applications. Structured state space models (S4) [13] address these issues by using reparameterization to enhance computational efficiency, providing an effective alternative to traditional attention mechanisms.

Several recent variants of S4 have been proposed to achieve linear time attention, including H3 [9], Gated State Space [23], Hyena [25], and RWKV [28]. Mamba [11] introduces a data-dependent selection mechanism to S4, improving the capture of long-range context as sequence length increases. Mamba not only achieves linear time efficiency in long-sequence modeling but also outperforms Transformers across various benchmarks. Recently, Jamba [22] has been introduced as a novel hybrid model that combines Transformer and Mamba layers in a mixture-of-experts (MoE) architecture. Jamba interleaves blocks of Transformer and Mamba layers, harnessing the strengths of both model families.

A Temporal Point Process (TPP) [6, 5] is a stochastic process that defines a probability distribution over event sequences. In this process, the number of points (events) $K$ and their locations (arrival times) $t_{i}$ are random. The realization of a TPP can be represented as a sequence of discrete events with $\{t_{i}\}\in\mathbb{R}^{+}$ and $i\in\mathbb{Z}^{+}$ abstracted as points on a timeline. We can represent a TPP realization by a counting measure $N(t)=\sum_{i}^{n}\mathbb{I}(t_{i}<t)$, for $t\in[0,T]$. The intensity characterzing a TPP can be interprested as the expected number of events per unit of time and is defined as:

where $\mathcal{H}_{t}=\{t_{i}:t_{i}<t\}$ is the event history until time $t$, which acts as a filtration to the process.

The Hawkes process [15, 19] is a typical temporal point process, and it models past events and predicts the timestamp of the following event by its conditional intensity function, which is defined as:

where $\mu\geq 0$, named base intensity, is an exogenous component of the intensity function independent of the history, while $\psi(t)>0$ is an endogenous component dependent on the history. Besides, $\psi(t)$ is a triggering kernel containing the peer influence of past events. Due to the self-exciting charities, the Hawkes process has recently received wide attention in event sequence modeling.

The structured state space sequence models (S4) [13] is defined by the simple equation 3, which maps a 1-dimensional function or sequence $x(t)\in\mathbb{R}\mapsto y(t)\in\mathbb{R}$ through an implicit latent state $h(t)\in\mathbb{R}^{N}$:

where $\mathbf{A}\in\mathbb{R}^{N\times N},\mathbf{B}\in\mathbb{R}^{N\times 1},\mathbf{C}\in\mathbb{R}^{1\times N}$ are parameters of neural networks in deep learning. To deal with the discrete input sequence $\boldsymbol{x}=(x_{0},x_{1},...)\in\mathbb{R}^{L}$, S4 discretizes these parameters in Eq. (3) using a step size $\Delta$, where the continuous parameters $\mathbf{A},\mathbf{B}$ are converted into discrete parameters $\overline{\mathbf{A}}=f_{A}(\Delta,\mathbf{A}),\overline{\mathbf{B}}=f_{B}(\Delta,\mathbf{B})$, where the pair $(f_{A},f_{B})$ is called a discretization rule [13]. Various rules can be used such as the zero-order hold (ZOH) defined as follows:

Then, this recurrent rule can be expanded as:

where $L$ denotes the length of the input sequence $\boldsymbol{x}$ and $\overline{\mathbf{K}}\in\mathbb{R}^{L}$ is the convolution kernel.

A recent development in state space layers is selective SSMs [11] (S6). These models utilize time-variant SSMs, namely, where the discrete matrices $\bar{A},\bar{B},$ and $C$ of each channel are modified over $L$ time steps depending on the input sequence. Unlike traditional state-space layers, which operate individually on each channel, selective state-space layers compute the SSM matrices $\bar{A}_{i},\bar{B}_{i},C_{i}$ for all $i\leq L$ based on all the channels, and then apply the time-variant recurrent rule individually for each channel.

Thus, we denote the entire input sequence by $\hat{x}:=(\hat{x}_{1},\cdots,\hat{x}_{L})\in\mathbb{R}^{L\times D}$ where $\hat{x}_{i}\in\mathbb{R}^{D}$. The per-time discrete matrices $\bar{A_{i}},\bar{B_{i}},$ and $C_{i}$ are defined as follows:

where $f_{A},f_{B}$ represents the discretization rule, $S_{B},S_{C},S_{\Delta}$ are linear projection layers, and softplus is an elementwise function that is a smooth approximation of ReLU.

Denote the sequence $\mathcal{S}=\{(t_{1},k_{1}),(t_{2},k_{2}),...,(t_{n},k_{n})\}$ as the temporal event sequence, where $K$ is the number of events. Let $\Delta_{i}=t_{i}-t_{i-1}$ represent the temporal differences, with $\Delta_{1}=t_{1}$ for convention, hence the temporal sequence is represented by

Additionally, let $\mathbf{k}_{i}$ be the one-hot vector of the events. Inspired by the discretization

the hidden states deffer by time gap $\Delta t$ are related by the formula, we put the temporal differences into the equation directly to construct our Mamba Hawkes Process (MHP) structure. Now we state our construction.

Let $W^{e}$ be the event embedding matrix with dimensions $D\times K$, where $D$ is the dimension of the hidden layers of Mamba blocks. The event embedding is defined as $x_{t_{i}}=\mathbf{k}_{i}(W^{e})^{T}$

In the Mamba architecture, the matrices $\Delta,\mathbf{B}$ and $\mathbf{C}$ are time-dependent and are obtained by linear projection from $x_{t}$. However, for the Hawkes Process, the approach is different, as it requires the use of temporal features. Specifically, we make $\mathbf{B}$ and $\mathbf{C}$ time-dependent, and $\Delta=(\Delta_{1},\Delta_{2},...,\Delta_{n})$ is defined by the temporal differences as above. Following Equation 9, we have $\Delta_{i}=t_{i}-t_{i-1}$, thus we replace $t$ and $t+\Delta t$ by $t_{i-1}$ and $t_{i}$. We define

are obtained by a linear transformation of the vector $x_{t_{i}}$. We have the transition formulas for the Mamba Hawkes process:

where

is the temporal-dependent coefficients. Hence, the temporal information is incorporated into our recurrence process.

The final output state is denoted by $\mathbf{O}=(\mathbf{o}_{1},\mathbf{o}_{2},...,\mathbf{o}_{n})$, we thus have

where $W_{i},b_{i}$ are the parameters for the two layers MLP. We have $\mathbf{h}(t_{j})=\mathbf{H}(j,:)$ in our case.

The intensity function of Neural Hawkes process is given by

where $\lambda_{k}$ is the intensity function of the $k$-th event, and

where $t$ is defined on interval $t\in\left[t_{j},t_{j+1}\right)$, and $f_{k}(x)=\beta_{k}\log\left(1+\exp\left(x/\beta_{k}\right)\right)$ is the Softplus function. For the log-likelihood, it is given by

Denote the log-likelihood of the event sequence $\mathcal{S}$ as $\mathcal{L}$.

For the prediction of next event type and timestamp, we train two linear layers $P^{e},P^{t}$

For the sequence $\mathcal{S}=\{(t_{1},k_{1}),(t_{2},k_{2}),...,(t_{n},k_{n})\}$, we define

where ${t}_{j}$ is the true timestamp of event $j$. $\mathcal{L}_{event}$ is the cross-entropy loss that measures the accuracy of the event type prediction, and $\mathcal{L}_{time}$ is the MSE loss that measures the accuracy of the time prediction. The training loss can then be defined as

where $\beta,\gamma$ are hyper-parameters to control the range of event and time losses.

Next we provide a proposition to show that, with specific choices of parameters, the Mamba Hawkes Process recurrence can degenerate to an RNN architecture similar to RMTPP as in [7].

Proof: We consider that if a given input $x_{t}$ should be completely ignored (as necessary in the synthetic tasks), all $D$ channels should ignore it, and so we project the input down to 1 dimension before repeating/broadcasting with $\Delta$.

In Mamba Hawkes Process, we set $\Delta_{i}=t_{i}-t_{i-1}$, when $N=1,\boldsymbol{A}=-1,\boldsymbol{B}=1$. By applying the zero-order hold $(\mathrm{ZOH})$ discretization formulas:

Denote that $g_{t_{i}}=\exp(t_{i-1}-t_{i})$, thus the final discrete recurrence is

as desired. We finish the proof.

From the construction we can see that the temporal differences have a canonical way as the time scale variables. This architecture can inherent the temporal information naturally.

In [22], the authors combine the Mamba layer and Transformer layer to create a Jamba structure, demonstrating impressive capabilities in large language models. Inspired by this, we propose combining the Mamba structure and Transformer structure to develop a new model architecture, which we call the Mamba Hawkes Process Extension (MHP-E).

Explicitly, Let MambaBlock be the Mamba structure defined above, TransformerBlock be the Transformer blocks. Given the temporal sequence $\mathcal{S}=\{(t_{1},k_{1}),(t_{2},k_{2}),...,(t_{n},k_{n})\}$ and the corresponding event one-hot vector $\mathbf{k}_{i}$, we first apply the embedding layer:

and we let the encoding vector pass through the Mamba layers and Transformer layers

and then we can use $h$ as the hidden layer to compute the log-likelihood $\mathcal{L}$, event cross entropy loss $\mathcal{L}_{event}$ and the mean square temporal loss $\mathcal{L}_{time}$.

We need to analyze the construction of the above architecture to understand its implicit ideas. In this architecture, we did not use temporal encoding for the Transformer, neither absolute nor relative. Instead, we used the Mamba layer as an encoder to simultaneously encode the temporal and event features. Consequently, the hidden layer obtained from the Mamba layer inherently incorporates this information.

We design our MHP models and MHP-E models as follows: For MPH, we set the dimension in our construction of the embedding as describing in the following table:

We set all other architecture hyper-parameters as suggested in [11]:

For the MHP-E, the Mamba part we use the same construction as MHP except we choose n_layers to be 2 since we only need it to encode the temporal and event features. For the Transformer blocks, we apply the same architecture as the Transformer Hawkes Process (THP) as described in [39]. Similarly, we follow the code from [39] and use $\beta=1$ and $\gamma=$1e-4 for the loss function. To avoid NaN values during our training, we apply Softplus function and Clamp function to the temporal difference $\Delta$.

All experiments are performed using GPU RTX A6000 with 48GB memory, and spend less than a minute for the training process each epoch.