SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Xiao Shou
xshou01@gmail.com
Tian Gao
tgao@us.ibm.com
Dharmashankar Subramanian
dharmash@us.ibm.com
Oktie Hassanzadeh
hassanzadeh@us.ibm.com
Kristin Bennett
bennek@rpi.edu
Although pairwise causal relations have been extensively studied in observational longitudinal analyses across many disciplines, incorporating knowledge of causal pairs into deep learning models for temporal event sequences remains largely unexplored. In this paper, we propose a novel approach for enhancing the performance of transformer-based models in multivariate event sequences by injecting pairwise qualitative causal knowledge such as 'event Z amplifies future occurrences of event Y'. We establish a new framework for causal inference in temporal event sequences using a transformer architecture, providing a theoretical justification for our approach, and show how to obtain unbiased estimates of the proposed measure. Experimental results demonstrate that our approach outperforms several state-of-the-art models in terms of prediction accuracy by effectively leveraging knowledge about causal pairs. We also consider a unique application where we extract knowledge around sequences of societal events by generating them from a large language model, and demonstrate how a causal knowledge graph can help with event prediction in such sequences. Overall, our framework offers a practical means of improving the performance of transformer-based models in multivariate event sequences by explicitly exploiting pairwise causal information.
Multivariate event sequences are datasets where different types of events occur sequentially, without meaningful time stamps. They can be viewed as categorical time series datasets where exactly one event from a known set of event types occurs at any position in a sequence. Applications of such datasets abound in various fields, including the social sciences, healthcare, finance, advertising, and engineering. Mining patterns from event sequences often provides meaningful insights for prediction, detection and analysis [1]; for example, analysis of such sequences is performed to study patterns of medication use or disease progression in healthcare, and to understand stock price prediction and fraud detection in finance. Modeling events is of interest in numerous fields besides data mining; for instance, there has been a surge of interest over the last several years around modeling time-stamped event streams in machine learning, e.g. [2, 3, 4], as well as in leveraging events for various tasks in natural language processing, e.g. [5, 6].
The field of machine learning has broadly experienced substantial attention around augmenting existing data with various forms of background knowledge, an area that some have referred to as informed machine learning [7]. A body of prior work has shown how incorporating various forms of domain knowledge could potentially aid deep learning models in performing better, more efficiently,
and more reliably [8], especially when the amount of available data is low. Such knowledge could also help make the models more interpretable and easier to understand, which is often important in critical real-world applications such as healthcare [9]. A commonly explored direction is one where physical rules are injected to control model prediction in physics regulated systems [10]. Knowledge augmentation is also being pursued in the setting of pretrained large language models [11]. There have been some recent related forays for machine learning of event models, such as through leveraging temporal logical rules [12] and data-dependent knowledge measures [13]. A recent study considers injecting qualitative structural statements in score-based learning systems for temporal point processes, which results in closed form and interpretable score-based loss [14].
In this paper, we address an important setting that brings together aspects of knowledge injection with modeling event sequences: how to effectively augment domain knowledge while learning multivariate event sequences using neural models. Specifically, we consider background knowledge of the form of qualitative statements such as: 'event type Z amplifies (or inhibits) event type Y'. Such statements embed the notion of pairwise causality and can be better paraphrased as 'occurrences of event type Z make future occurrences of event type Y more (or less) likely'. Such causal knowledge can oftentimes be acquired directly from causal knowledge graphs [15, 16, 17, 18] or learned from data.
Integrating pairwise causal statements into learning neural models for temporal event sequences poses at least two major technical challenges that need to be suitably tackled:
1. Causal inference usually entails examining both the factual outcome as well as the counterfactual outcome for a binary treatment. Since we consider a specific type of non-i.i.d. temporal setting where events occur sequentially, one cannot directly apply the standard potential outcome framework [19] without appropriate definitions that account for the sequential nature of the data.
2. Pairwise causal statements about event types reflect a type-to-type global property that may or may not have been inferred from instance-level data. It is relatively straightforward to use such information through additional regularization terms for loss functions to learn typical parametric models; this is because the (learned) parameters of such models are global and stationary, and thus any deviations from the statement can be penalized. However, a neural event sequence model such as an autoregressive transformer captures the dynamics of event instances, thus requiring a different approach for injecting type-level knowledge.
Contributions. In this paper, we incorporate causal knowledge augmentation into the modeling of temporal event sequences. Specifically, we propose pairwise causality guided learning in transformers that suitably addresses the afore-mentioned challenges. Our contributions are as follows: 1) We formulate the problem of incorporating pairwise causal knowledge into neural autoregressive models in the context of temporal event sequences, proposing suitable definitions for a formal causal inference investigation. 2) We provide theoretical justification for our procedure, particularly to combat the issue of time confounding and to achieve representational balance. 3) We conduct a detailed empirical evaluation demonstrating superior predictive performance as compared to state-of-the-art baselines on synthetic datasets as well as real-world benchmarks. 4) We explore a novel application involving large language model generated event sequences, illustrating both how a causal knowledge graph could aid in event prediction, as well as how one could potentially build event models through knowledge injection without relying on existing event sequences.
Our work is interdisciplinary in nature, touching upon topics such as event sequence modeling, causal inference, neural networks, and natural language processing. We briefly review closely related work across some of these topics.
Markov Models. Markov models have been widely deployed over several decades for analyzing temporal event sequences. Besides the classic  $k$ th-order Markov models where the probability of the next state depends solely on the previous  $k$  positions in the sequence, other variations include (but are not limited to) hidden Markov models [20], models incorporating linear combinations [21], variable order Markov models [22], summary Markov models [23], etc. While such models have advantages around interpretability and low data requirements, they are often unable to capture long-range and potentially complex historical dependencies, unlike neural sequence models.
Figure 1: An illustrative event sequence over labels  $\{A, B, C, D, E\}$  where event pair A and B are of interest. An example of event labels from healthcare is as follows:  $A$ : Emergency visit,  $B$ : Cardiac event,  $C$ : Medication refill,  $D$ : Allergy diagnosis,  $E$ : Anxiety diagnosis.
Neural Sequence Models. RNNs are a class of neural networks that process sequences of inputs. In an RNN, the output at each time step depends not only on the current input but also on the previous output. Variants of RNNs such as Long Short-Term Memory (LSTM) and Gated Recurrent Units (GRU) have been shown to be effective in various applications, including natural language processing, speech recognition, and time-series analysis. State-of-the-art sequence models involve the use of attention-based transformers. Current research focuses on making more reliable predictions from factual associations, for example in large language models [24], and on equipping such transformer-based models with arithmetic, commonsense or symbolic reasoning [25]. The importance of causal reasoning for natural language understanding has also recently emerged [26].
Event Models. Event sequences are a popular source of data for numerous data mining algorithms [1, 27, 28]. There is also a long line of research around modeling time-stamped streams using temporal point processes [29] that has become popular in machine learning. Graphical event models (GEMs), also known as local independence graphs, are a family of graphical models that represent temporal point processes [30, 2, 3]. GEMs capture the conditional independence structure of the continuous-time dynamics of event processes and are closely related to the notion of Granger-causality [31, 32, 33]. A recent Bayesian approach uses a score-based framework for injecting qualitative statements while learning a sub-family of parametric GEMs [14]. Our work here differs in that we model event sequences without time-stamp information, and inject pairwise causal knowledge into transformer models, which are known to be more flexible and therefore better than typical parametric models at tasks such as event prediction [34].
Sequential Treatment Effect Estimation. Substantial research has been conducted on time-varying treatment estimation for longitudinal observation data [35, 36], and many of the state-of-the-art approaches leverage neural models [37]. Some prominent examples over the last couple of years include recurrent marginal structural networks [38], counterfactual recurrent networks [39], G-Nets [40] and causal transformers [41]. In addition, some recent work has studied causal inference for temporal process data [42, 43], by leveraging time information from event stream datasets. This line of work does not consider injecting pairwise causal knowledge into neural temporal event sequence models.
Causality in NLP. Feder et al. [44] provide an excellent survey of work done at the intersection of causal inference and NLP. There is also a body of work on the extraction of causal relations from natural language descriptions [45]. Causal knowledge derived from text either through inference or extraction can be used for the injection of pairwise causal knowledge into event models.
3 Causal Event Pairs in Event Sequences
In this section, we first introduce basic notation and terminology and then motivate how one can incorporate pairwise causal knowledge in neural autoregressive models such as attention based transformers for temporal event sequences.
3.1 Event Sequences & History-Confounding Dynamics
Consider a multivariate event sequence dataset  $\mathbf{D}$  which consists of  $K$  sequences: indexed by  $k$ , each sequence takes the form of  $\mathbf{D}_{\mathbf{k}} = [l_{i}]_{i=1}^{N_{k}}$ , where event label (or type)  $l_{i}$  at position  $i$  in the
sequence is from a discrete label set,  $l_{i} \in \mathcal{L}$ . The total number of events and event types are denoted  $N$  and  $M$  respectively, i.e.  $N = \sum_{k=1}^{K} N_{k}$  and  $M = |\mathcal{L}|$ . Let  $h_{i} = \{l_{j}\}_{j=1}^{i}$  denote the history of past events up to  $i$ . Consider the following example in a typical event sequence:
Example 1 Figure 1 shows an example event sequence with 5 event labels,  $\mathcal{L} = \{A,B,C,D,E\}$ . Consider the 3 causal pairs  $(A,B),(B,C),(A,C)$  potentially provided by a domain expert.  $B$  causes the probability of future occurrences of  $A$  to increase, denoted as  $B \nearrow A$ . Another 'competing' event  $C$  causes the probability of future occurrences of  $A$  to decrease, i.e.  $C \searrow A$ . Event  $C$  also excites  $B$ , i.e.  $C \nearrow B$ . Note that it can be challenging to determine the influence of how one event affects another in an observational temporal event sequence.
Example 1 entails pairwise qualitative statements for sequential events occurring over time, and therefore requires us to establish some formal tools; we turn to causal inference and the potential outcome framework [19] in particular to address this issue. We formulate treatment, covariate and outcome models for any event sequence  $\mathbf{D}_k$  as follows. Note that we abuse notation slightly and sometimes use lower case (e.g.  $y, z$ ) for event labels to indicate that they can also be viewed as instances of random variables.
Definition 1 Formulation of Causal Event Pair: For a pair of event labels  $(z,y)$ , the binary treatment variable  $Z_{i}$  at time  $i$  is whether or not label  $z$  occurs at position  $i$ , i.e.,  $l_{i} = z$  or  $l_{i} \neq z$ . For a specific  $Z_{i}$  at time  $i$ , the corresponding covariates are historical occurrences of event labels including  $z$ , i.e.  $h_{i-1}$ . The outcome is the probability of occurrence of  $y$  (at least once) in the next window  $w$  given the treatment and covariates, i.e.,  $p(I(y)_i^w | l_i, h_{i-1}) = 1 - p(l_{i+1} \neq y, l_{i+2} \neq y, \dots, l_{i+w} \neq y | l_i, h_{i-1}) = 1 - p(l_{i+1} \neq y, l_{i+2} \neq y, \dots, l_{i+w} \neq y | h_i)$  for  $w \in \{1,2,3,\dots,W\}$  where  $W \ll N_k$ .
In the above definition,  $I(y)_i^w$  is an indicator function for occurrence of  $y$  (at least once) within the next window  $w$  from position  $i$ . Many other choices of outcome models are possible. For instance, one could choose the next arrival time of  $Y$ , however we may not observe  $Y$  at all in a finite time horizon. The benefit of our definition is that one can easily compute the probability of  $Y$  in the next window  $w$ . The simplest setting is when the window  $w = 1$ , which we will use as default. For cases where  $w > 1$ , the same principle holds and we discuss this variation in Section 4.3. We define the propensity score as follows.
Definition 2 Under Definition 1, a valid propensity score is defined as  $p(l_{i+1} = z|h_i)$ .
It is challenging to directly compute the propensity score under the raw history vector  $h_i$  due to the varying length of history. Furthermore, such history can be of high dimension for large  $N_k$ . A possible approach is to compute the propensity score in its latent representation. A neural autoregressive model with parameters  $\theta$  is trained to maximize its log likelihood and we thus obtain its propensity score  $p_{\theta^*}(l_i = z|h_{i-1})$  at each time instance  $i$  with a sample/batch size  $B$ :
$$
\theta^ {*} = \operatorname {a r g m a x} _ {\theta} \mathbb {E} _ {B \sim p \left(D _ {k}\right)} \left[ \sum_ {j = 1} ^ {| B |} \sum_ {i = 1} ^ {N _ {k} - 1} p _ {\theta} \left(l _ {i + 1} \mid h _ {i}\right) \right] \tag {1}
$$
Transformer for Event Sequences. Consider an autoregressive transformer  $G$  which takes an event (sub)sequence and outputs a sequence of categorical probability distribution:  $\mathcal{H} \rightarrow \mathcal{Y}$ . An event (sub)sequence  $H$  is first embedded into  $\mathbf{X}$  via position embedding. The two major components of a self attention layer are attention layers (Attn) and position-wise feed-forward networks (FFN) via residual connection. We show the  $h$ -head attention with upper triangular attention mask  $\mathbb{M}$  so that only future events attend to the past in our setting:
$$
\operatorname {A t t n} (\mathbf {X}) = \mathbf {X} + \sum_ {i = 1} ^ {h} \mathbf {W} _ {O} ^ {i} \mathbf {W} _ {V} ^ {i} \mathbf {X} \cdot \left(\sigma \left[ \left(\mathbf {W} _ {K} ^ {i} X\right) ^ {T} \mathbf {W} _ {Q} ^ {i} \mathbf {X} \right] \odot \mathbb {M}\right), \tag {2}
$$
$$
F F N (X) = A t t n (X) + \mathbf {W} _ {2} \cdot R e L U (\mathbf {W} _ {1} \cdot A t t n (\mathbf {X}) + \mathbf {b} _ {1} \mathbf {1} ^ {T}) + \mathbf {b} _ {2} \mathbf {1} _ {n} ^ {T},
$$
where  $\mathbf{W}_Q^i\mathbf{W}_K^i\mathbf{W}_V^i$  and  $\mathbf{W}_O^i$  are query, key, value and output matrix in the  $i^{th}$  head with appropriate dimension.  $\sigma$  is the softmax operation.
Theorem 2 Under Definition 1, an autoregressive transform  $G$  outputs a valid propensity score, i.e.,  $p(l_{i+1} = z|h_i)$ . A valid propensity score is computed by using only past history.
Proof Sketch. The only operations that are not point-wise are in the attention module. Let  $\tilde{\mathbf{X}} = \mathbf{W}_O^i\mathbf{W}_V^i\mathbf{X}$  and  $\tilde{\mathbf{A}} = \sigma[(\mathbf{W}_K^i X)^T(\mathbf{W}_Q^i\mathbf{X})] \odot \mathbb{M}$ . Thus for each head  $i$ ,  $(\tilde{\mathbf{X}}\tilde{\mathbf{A}})_{ij} = \sum_k \tilde{\mathbf{X}}_{ik}\tilde{\mathbf{A}}_{kj}$ . Notice masked attention  $\tilde{\mathbf{A}}$  is upper triangular,  $\sum_k \tilde{\mathbf{X}}_{ik}\tilde{\mathbf{A}}_{kj} = \sum_{k \leq j} \tilde{\mathbf{X}}_{ik}\tilde{\mathbf{A}}_{kj}$ . Hence only past history is involved for predicting the next treatment assignment  $l_{i+1} = z$ .
Theorem 3 Under Definition 1, an autoregressive transformer  $G$  with  $L$  blocks generates a full spectrum of balancing scores at each instance  $i = \{1, \dots, N_k - 1\}$  for each layer  $l = \{1, 2, \dots, L\}$ .
Proof Sketch: The proof follows Proposition 4 from [46] and Theorem 1 above as Transformer  $G$  with  $L$  blocks can be viewed as a sequence of compositions, i.e.  $G \coloneqq \sigma(f_{L} \circ f_{L} - 1 \circ \dots f_{1} \circ f_{0}(H))$ , where  $f_{0}$  is the position embedding,  $f_{l}$  is ATTN followed by FFN, and  $\sigma$  is the softmax operation.
We propose an approach to inject pairwise causal knowledge, referring to our model as a pairwise-causality transformer for event sequences (PC-TES). More precisely, we assume that such causal pair knowledge holds for the average case, i.e., on average, for every sequence  $k$ , event type  $Z$  decreases/increases the probability of future occurrences of event type  $Y$ . Formally, this impacts the difference  $\mathbb{E}[P(l_{i+1} = y)|l_i = z] - \mathbb{E}[P(l_{i+1} = y|l_i \neq z]$  for arbitrary position  $i$  in the sequence. We make the following assumption regarding measured confounding in temporal sequences, which is not uncommon in sequential treatment and precision medicine [39, 41]. This aspect also frequently appears in economic event sequences [23] and diabetic patient activities [47]. While such confounding terms can be trivially handled by identifying other parental states in typical parametric models (see e.g. [14]), it is not straightforward for neural models.
Assumption 1 An observational multivariate event sequence under the above framework involves time-varying confounding, i.e., for a causal pair  $(Y,Z)$ , history not only affects the occurrence of future events  $Y$  but also affects the occurrences of  $Z$ .
This assumption holds for Example 1 – the effect of  $B$  on  $A$  is not only confounded by  $C$  but also any past history; it also implies that any causal pair  $(Z, Y)$  has potential history confounding as long as  $Z$  and  $Y$  are not process independent [30]. In Example 1, occurrences of  $B$  in the past affect both the occurrences of outcome  $A$  and  $B$ . In addition, occurrences of  $C$  in history on the other hand affect the occurrences of  $A$  and  $B$ . The effect of  $C$  on the occurrences of  $A$  are thus confounded by past histories involving  $B$  or  $C$  itself and needs to be adjusted. A full qualitative statement in our case becomes 'on average, for each sequence  $k$ , event type  $Z$  decreases/increases the probability of future occurrences of event type  $Y$  while holding history constant'. The additional three standard assumptions of consistency, sequential ignorance, and sequential overlap can be made similarly to time series observational data [48, 36] to identify the potential counterfactual outcomes in our setting (see Appendix for further details).
4.1 Inverse Probability Weighting in Event Sequences
We define 2 relevant terms with respect to the statement  $Z$  reduces the occurrence of  $Y$  via inverse probability weighting. Without loss of generality, consider a window  $w = 1$ .
Definition 3 For an event sequence  $k$ , the instance-level effect (ILE) at position  $i$  is the difference between potential probabilities of  $Y$  at  $i + 1$  for occurrence and nonoccurrence of  $Z$  at  $i$ , i.e.  $p(l_{i+1} = y|h_{i-1}, l_i = z) - p(l_{i+1} = y|h_{i-1}, l_i \neq z)$ .
Definition 4 For an event sequence  $k$ , the sequence-level effect (SLE) is averaged over all positions of ILEs, i.e.  $\frac{1}{n-1} \sum_{i=1}^{n-1} (p(l_{i+1} = y|h_{i-1}, l_i = z) - p(l_{i+1} = y|h_{i-1}, l_i \neq z))$ .
For an observational temporal event sequence, we impose inverse probability weighting to adjust for time confounding:  $\hat{\tau} = \frac{1}{n - 1}\sum_{i = 1}^{n - 1}\left(\frac{p(l_{i + 1} = y|h_{i - 1},l_i = z)}{p(l_i = z|h_{i - 1})} -\frac{p(l_{i + 1} = y|h_{i - 1},l_i\neq z)}{p(l_i\neq z|h_{i - 1})}\right)$ .
Theorem 4 Under Definition 1, Assumption 1 and the consistency, sequential ignorability, and sequential overlap assumptions, given a learned transformer model  $G$ , our inverse probability weighted sequence-level effect (SLE) under  $G$  is an unbiased estimator of the true effect  $\tau = \mathbb{E}[P(l_{i+1} = y)|l_i = z] - \mathbb{E}[P(l_{i+1} = y)|l_i \neq z]$ .
Proof Sketch:  $\hat{\tau}$  can be rewritten as  $\frac{1}{n-1} \sum_{i=1}^{n-1} \frac{(l_i = z)p^*(l_{i+1} = y)}{p^*(l_i = z)} - \frac{1}{n-1} \sum_{i=1}^{n-1} \frac{(l_i \neq z)p^*(l_{i+1} = y)}{1 - p^*(l_i = z)}$  where  $*$  stands for conditioning on the historical representation obtained from  $G$ , i.e.  $H_i$  and  $H_{i-1}$  respectively for numerator and denominator. Thus we can show each  $\mathbb{E}\frac{(l_i = z)p^*(l_{i+1} = y)}{p^*(l_i = z)} = \mathbb{E}[P(l_{i+1} = y)^{Z_i = 1}]$  and  $\mathbb{E}\frac{(l_i \neq z)p^*(l_{i+1} = y)}{1 - p^*(l_i = z)} = \mathbb{E}[P(l_{i+1} = y)^{Z_i = 0}]$  using the law of iterated expectation where the subscripts  $Z_i = 0$  and  $Z_i = 1$  are for each stratum.
In practice, if the number of event types is large, causing  $p^*(l_i = z)$  to be small, we can use standardized stable inverse probability to avoid numerically instability, similar to [42] for temporal event sequences:  $\hat{\tau} = \frac{1}{n-1} \sum_{i=1}^{n-1} \left( \frac{p(z)p(l_{i+1} = y|h_{i-1}, l_i = z)}{p(l_i = z|h_{i-1})} - \frac{(1-p(z))p(l_{i+1} = y|h_{i-1}, l_i \neq z)}{p(l_i \neq z|h_{i-1})} \right)$  where  $p(z)$  is a prior belief of probability of  $z$  occurring in the sequence  $k$ .
4.2 Incompatibility Framework
Our injected statements are qualitative, and the ground truth effect  $\tau$  is usually unknown even for domain experts. Hence we only determine whether it deviates from the statement in practice. To do so, we adopt an incompatibility framework [49, 14] where qualitative knowledge is injected as a loss term. For the statement  $Z$  inhibits  $Y$ 's for a sequence  $k$ , our proposed combined loss is:
$$
L _ {t o t} = - \sum_ {i = 1} ^ {n _ {k}} p ^ {*} (l _ {i}) + \alpha \max  (\frac {1}{n - 1} \sum_ {i = 1} ^ {n - 1} (\frac {\mathbb {1} (l _ {i} = z) p ^ {*} (l _ {i + 1} = y)}{p ^ {*} (l _ {i} = z)} - \frac {\mathbb {1} (l _ {i} \neq z) p ^ {*} (l _ {i + 1} = y)}{1 - p ^ {*} (l _ {i} = z)}), 0) \tag {3}
$$
For the statement  $Z$  amplifies  $Y$ , for a sequence  $k$ , the loss is:
$$
L _ {t o t} = - \sum_ {i = 1} ^ {n _ {k}} p ^ {*} (l _ {i}) + \alpha \max  (\frac {1}{n - 1} \sum_ {i = 1} ^ {n - 1} (\frac {\mathbb {1} (l _ {i} \neq z) p ^ {*} (l _ {i + 1} = y)}{1 - p ^ {*} (l _ {i} = z)} - \frac {\mathbb {1} (l _ {i} = z) p ^ {*} (l _ {i + 1} = y)}{p ^ {*} (l _ {i} = z)}), 0) \tag {4}
$$