SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yutong Xia $^{1}$ , Yuxuan Liang $^{2*}$ , Haomin Wen $^{3}$ , Xu Liu $^{1}$ , Kun Wang $^{4}$ , Zhengyang Zhou $^{4}$ , Roger Zimmermann $^{1}$
$^{1}$ National University of Singapore
<sup>2</sup>The Hong Kong University of Science and Technology (Guangzhou)
$^{3}$ Beijing Jiaotong University  $^{4}$ University of Science and Technology of China
{yutong.x,yuxliang}@outlook.com; {liuxu,rogerz}@comp.nus.edu.sg wenhaomin@bjtu.edu.cn;wk520529@mail.ustc.edu.cn;zzy0929@ustc.edu.cn
Spatio-Temporal Graph (STG) forecasting is a fundamental task in many real-world applications. Spatio-Temporal Graph Neural Networks have emerged as the most popular method for STG forecasting, but they often struggle with temporal out-of-distribution (OoD) issues and dynamic spatial causation. In this paper, we propose a novel framework called CaST to tackle these two challenges via causal treatments. Concretely, leveraging a causal lens, we first build a structural causal model to decipher the data generation process of STGs. To handle the temporal OoD issue, we employ the back-door adjustment by a novel disentanglement block to separate the temporal environments from input data. Moreover, we utilize the front-door adjustment and adopt edge-level convolution to model the ripple effect of causation. Experiments results on three real-world datasets demonstrate the effectiveness of CaST, which consistently outperforms existing methods with good interpretability. Our source code is available at https://github.com/yutong-xia/CaST.
Individuals enter a world with intrinsic structure, where components interact with one another across space and time, leading to a spatio-temporal composition. Spatio-Temporal Graph (STG) has been pivotal for incorporating this structural information into the formulation of real-world issues. Within the realm of smart cities [60], STG forecasting (e.g., traffic prediction [58, 49, 17, 52] and air quality forecasting [28, 48, 27]) has become instrumental in informed decision-making and sustainability. With recent advances in deep learning, Spatio-Temporal Graph Neural Networks (STGNNs) [50, 20] have become the leading approach for STG forecasting. They primarily use Graph Neural Networks (GNN) [23] to capture spatial correlations among nodes, and adopt Temporal Convolutional Networks (TCN) [2] or Recurrent Neural Networks (RNN) [13] to learn temporal dependencies.
However, STG data is subject to temporal dynamics and may exhibit various data generation distributions over time, also known as temporal out-of-distribution (OoD) issues or temporal distribution shift [6, 67]. As depicted in Figure 1a, the training data (periods A and B) and test data derive from different distributions, namely  $P_A(x) \neq P_B(x) \neq P_{test}(x)$ . Most prior studies [26, 61, 56, 15, 30] have overlooked this essential issue, which potentially results in suboptimal performance of STGNNs that are trained on a specific time period to accurately predict future unseen data.
Meanwhile, dynamic spatial causation is another essential nature of STG data that must be addressed for effective and unbiased representation learning in STG forecasting. While the majority of STGNNs
* Yuxuan Liang is the corresponding author of this paper. Email: yuxiang@outlook.com
rely on a distance-based adjacency matrix to perform message passing in the spatial domain [61, 26, 11], they lack adaptability to dynamic changes in the relationships between nodes. This matrix is also sometimes inaccurate, as two closely located nodes may not necessarily have causal relationships, e.g., nodes belonging to different traffic streams. As an alternative solution, the attention mechanism [14, 64, 9] calculates the dynamic spatial correlations between nodes adaptively based on their input features. However, they still fall short of capturing the ripple effects of causal relations. Similar to how node signals can propagate information across graphs over time, causal relations (perceived as edge signals) can also exhibit this effect. For example in Figure 1b, when an accident occurs between node A and B at  $t = 1$ , it directly reduces the causal relation 1. At  $t = 2$ , this effect propagates to other relations, such as weakening relation 2 and strengthening relation 3. This happens because the accident decreases the proportion of traffic flow from node B observed by node D, thus simultaneously increasing the proportion of traffic flow from node C observed by D.
In this paper, our goal is to concurrently tackle the temporal OoD issue and dynamic spatial causation via causal treatments [36]. Primarily, we present a Structural Causal Model (SCM) to gain a deeper understanding of the data generation process of STG data. Based on SCM, we subsequently propose to 1) utilize back-door adjustment to enhance the generalization capability for unseen data; 2) apply front-door adjustment along with an edge-level convolution operator to effectively capture the dynamic causation bet
Figure 1: (a) Illustration of temporal OoD. (b) Spatial causal relationship in the traffic system.
- A causal lens and treatment for STG data. We propose a causal perspective to decipher the underlying mechanisms governing the data generation process of STG-structured data. Building upon the causal treatment, we devise a novel framework termed Causal Spatio-Temporal neural networks (CaST) for more accurate and interpretable STG forecasting.
- Back-door adjustment for handling temporal OoD. We articulate that the temporal OoD arises from unobserved factors, referred to as temporal environments. Applying the back-door adjustment, we design a disentanglement block to separate the invariant part (we call it entity) and environments from input data. These environments are further discretized by vector quantization [46] which incorporates a learnable environment codebook. By assigning different weights to these environment vectors, our model can effectively generalize on OoD data from unseen environments.
- Front-door adjustment for capturing dynamic spatial causation. Adopting a distanced-based adjacency matrix to capture spatial information around a node (referred to as spatial context) could include spurious causation. However, stratifying this context is computationally intensive, making back-door adjustments impractical for spatial confounding. We thus utilize the front-door adjustment and introduce a surrogate to mimic node information filtered based on the actual causal part in the spatial context. To better model the ripple effect of causation, we propose a novel de-confounding module to generate surrogate representations via causal edge-level convolution.
- Empirical evidence. We conduct extensive experiments on three real-world datasets to validate the effectiveness and practicality of our model. The empirical results demonstrate that CaST not only outperforms existing methods consistently on these datasets, but can also be easily interpreted.
Spatio-Temporal Graph Forecasting. Recently, STG forecasting has garnered considerable attention, with numerous studies focusing on diverse aspects of this domain. Based on the foundation of GNNs [23], STGNNs [50, 18, 38, 20] have been developed to learn the spatio-temporal dependencies in STG data. By incorporating temporal components, such as TCN [2] or RNN [13], STGNNs are capable of modeling both spatial correlations and temporal dependencies in STG data. Pioneering examples include DCRNN [26], STGCN [61], and ST-MGCN [11]. Following these studies, Graph WaveNet [56] and AGCRN [1] leverage an adaptive adjacency matrix to improve the predictive performance. ASTGCN [14] and GMAN [64] utilize attention mechanisms to learn dynamic spatiotemporal dependencies within STG data. STGODE [8] and STGNCDE [5] capture the continuous spatial-temporal dynamics by using neural ordinary differential equations. ST-MetaNet [34, 35] and AutoSTG [33] exploit meta learning and AutoML for learning STGs, respectively. However, none of these approaches can simultaneously address the temporal OoD issue and dynamic spatial causation.
Causal Inference. Causal inference [36, 12] seeks to investigate causal relationships between variables, ensuring stable and robust learning and inference. Integrating deep learning techniques with causal inference has shown great promise in recent years, especially in computer vision [62, 51, 29], natural language processing [37, 44], and recommender system [65, 10]. However, in the field of STG forecasting, the application of causal inference is still in its infancy. Related methods like graph-based causal models are typically designed for graph/node classification tasks [43, 66] and link prediction [25]. For sequential data, causal inference is commonly used to address the temporal OoD issue by learning disentangled seasonal-trend representations [53] or environment-specific representations [57]. When adapting to STGs, these methods face hurdles, as graph-based models cannot tackle temporal OoD issues, while sequence-based models fail to accommodate spatial dependencies. In this study, we investigate the STG data generation process through a causal lens and employ causal techniques to mitigate confounding effects in both the temporal and spatial domains for STG forecasting.
3 Causal Interpretation for STG Data Generation
Problem Statement. We denote  $X^t \in \mathbb{R}^{N \times D}$  as the signals of  $N$  nodes at time step  $t$ , where each node has  $D$  features. Given the historical signals from the previous  $T$  steps, we aim to learn a function  $\mathcal{F}(\cdot)$  that forecasts signals over
the next  $S$  steps:  $[X^{(t - T):t}] \xrightarrow{\mathcal{F}(\cdot)} [Y^{(t + 1):(t + S)}]$ , where  $X^{(t - T):t} \in \mathbb{R}^{T \times N \times D}$ ,  $Y^{(t + 1):(t + S)} \in \mathbb{R}^{S \times N \times D'}$  and  $D'$  is the output dimension. For
(b)
Figure 2: SCMs of (a) STG generation under real-world scenarios; (b) back-door adjustment for  $E$ ; (c) front-door adjustment for  $C$ .
conciseness, we refer to  $X^{(t - T):t}$  as  $X$  and  $Y^{(t + 1):(t + S)}$  as  $Y$  in the rest of the paper.
A Causal Look on STG. From a causal standpoint, we construct a Structural Causal Model (SCM) [36] (see Figure 2a) to illustrate the causal relationships among four variables: temporal environment  $E$ , spatial context  $C$ , historical node signals  $X$ , and future signals  $Y$ . Arrows from one variable to another signify causal-effect relationships. For simplicity, we assume  $E$  and  $C$  are mutually independent. Based on the above definitions, the causal relationships in Figure 2a can be denoted as  $P(X,Y|E,C) = P(X|E,C)P(Y|X,E,C)$ . We detail these causal-effect relationships below:
-  $X \gets E \to Y$ . The temporal OoD is an inherent property of STG data, e.g.,  $X$  and  $Y$ , where  $P(X^t) \neq P(X^{t + \Delta t})$  at different time steps  $t$  and  $t + \Delta t$ . This phenomenon can arise due to changes in external variables over time, which we refer to as temporal environments  $E$  [57]. For example, external factors such as weather and events can significantly affect traffic flow observations.
-  $X \gets C \to Y$ . The historical and future data  $X$  and  $Y$  are intrinsically affected by the encompassing spatial context  $C$  surrounding a node. This influence, however, can comprise both spurious and genuine causal components. The spurious aspects may encompass nodes that exhibit either spatial distance or semantic similarity yet lack causal connections, as elucidated in Section 1.
-  $X \to Y$ . This relation is our primary goal established by the prediction model  $Y = \mathcal{F}(X)$ , which takes historical data  $X$  as input and produces predictions for future node signals  $Y$ .
Confounders and Causal Treatments. Upon examining SCM, we observe two back-door paths between  $X$  and  $Y$ , i.e.,  $X \gets E \to Y$  and  $X \gets C \to Y$ , where the temporal environment  $E$  and spatial context  $C$  act as confounding factors. This implies that some aspects of  $X$ , which are indicative of  $Y$ , are strongly impacted by  $E$  and  $C$ . To mitigate the negative effect of the two confounders, we leverage the causal tools [12, 36] and do-calculus on variable  $X$  to estimate  $P(Y|do(X))$ , where  $do(\cdot)$  denotes the do-calculus. For the temporal OoD, we employ a popular de-confounding method called back-door adjustment [57, 43] to block the back-door path from  $E$  to  $X$  (the red dashed arrow in Figure 2b), so as to effectively remove  $E$ 's confounding effect. This necessitates implicit environment stratification (see Eq. 1). Spatial confounding, however, cannot be addressed by spatial context stratification, due to the computational burden engendered by the multitude of nodes, each exhibiting a unique contextual profile. Fortunately, the front-door adjustment allows us to introduce a mediating variable  $X^*$  between  $X$  and  $Y$  to mimic a more accurate representation excluded the spurious parts in  $C$  (the red node in Figure 2c) [54]. Note that we do not use the front-door adjustment to  $E$  because this method mandates that the mediating variable is only affected by the cause variable and not by other confounding factors. While for temporal OoD scenarios, unseen future environments can be affected by time-varying factors, thus influencing the mediating variable. These two approaches effectively de-confound  $E$  and  $C$ 's confounding effects, explained as follows.
Figure 3: The pipeline of CaST. Env: Environment. Ent: Entity. Feat: Feature.
Back-door Adjustment for  $E$ . To forecast future time series  $Y$  based on historical data  $X$ , it is imperative to address the confounding effect exerted by the temporal confounder  $E$ . To achieve this, we initially envisage a streamlined SCM, where  $E$  constitutes the sole parents of  $X$  (temporarily disregarding  $C$ ) and employ back-door adjustment [36] to estimate  $P(Y|do(X))$  by stratifying  $E$  into discrete components  $E = \{e_i\}_{i=1}^{|E|}$ :
$$
P (Y | d o (X)) = \sum_ {i = 1} ^ {| E |} P (Y | X, E = e _ {i}) P (E = e _ {i}) \tag {1}
$$
where the prior probability distribution of the environment confounder  $P(E)$  is independent of  $X$  and  $Y$ , allowing us to approximate the optimal scenario by enumerating  $e_i$ .
Front-door Adjustment for  $C$ . Once we have dealt with  $E$ , our next step is to de-confound the effect of spurious spatial context  $C$  by using the front-door adjustment [12]. In the SCM depicted in Figure 2c, an instrumental variable  $X^{*}$  is introduced between  $X$  and  $Y$  to mimic the node representation conditioned on their real causal relationships with other nodes. We then estimate the causal effect of  $X$  on  $Y$  as follows:
$$
P (Y | d o (X)) = \sum_ {x ^ {*}} \sum_ {x ^ {\prime}} P (Y | X ^ {*} = x ^ {*}, X = x ^ {\prime}) P (X = x ^ {\prime}) P (X ^ {*} = x ^ {*} | X) \tag {2}
$$
By observing  $(X, X^{*})$  pairs, we can estimate  $P(Y|X^{*}, X)[54]$ . This front-door adjustment provides a reliable estimation of the impact of  $X$  on  $Y$  while circumventing the confounding associations caused by  $C$ . We put the derivations of Eq. 1 and Eq. 2 in Appendix A.
4 Model Instantiations
We implement the above causal treatments by proposing a Causal Spatio-Temporal neural network (CaST), as depicted in Figure 3. Our method takes historical observations  $X$  as inputs to predict future signals  $Y$ . We will elaborate on the pipeline and each core component in the following parts.
Back-door Adjustment (see the top half of Figure 3). In order to attain Eq. 1, two steps need to be taken: (1) separating the environment feature from the input data, and (2) discretizing the environments. To accomplish this, we introduce an Environment Disentangler block, and a learnable Environment Codebook to obtain the desired stratification of environments.
Front-door Adjustment (see the bottom half of Figure 3). Obtaining  $X^{*}$  and collecting the  $(X, X^{*})$  pairs to instantiate Eq. 2 is a non-trivial task that involves two main obstacles: (1) enumerating each spatial context, which can be computationally expensive, especially for large graphs, and (2) quantifying the causal effect of  $X$  on  $X^{*}$ . To tackle these challenges, we introduce a block called Hodge-Laplacian (HL) Deconfounder, which is a neural topologically-based de-confounding module, to capture the dynamic causal relations of nodes as well as position embeddings to learn the nodes' global location information. With these two techniques, we can approximate the surrogate  $X^{*}$ .
(a) Env Disentangler
(b) HL Deconfounder
Figure 4: (a) The structure of Env Disentangler. AvgPool: average pooling. Linear: linear projection. Attn: attention. FFT: Fast Fourier Transform. iFFT: inverse FFT. (b) The overview of HL Deconfounder. GCN: graph convolution. (c) The embedding space for the Environment Codebook. The output  $\hat{H}_e$  is projected onto the closest vector  $e_3$  and the gradient  $\nabla_z\mathcal{L}$  pushes  $H_{e}$  to change.
(c) Embedding Space
4.1 Temporal Environment Disentanglement
Overview. As shown in Figure 3, the input signals  $X$  are first mapped to latent space as  $H \in \mathbb{R}^{T \times N \times F}$  by a Backbone Encoder before entering the Env Disentangler, where  $F$  means the hidden dimension. Then, this block separates  $H$  into the environmental feature  $H_{e} \in \mathbb{R}^{N \times F}$  and entity feature  $H_{i} \in \mathbb{R}^{N \times F}$ , analogous to background and foreground objects in the computer vision field [51]. Specifically, it captures environmental and entity information using two distinct components (see Figure 4a), including 1) EnvEncoder, which consists of a series of 1D convolutions, average pooling, and a linear projection; 2) EntEncoder, which extracts features from both time and frequency domains via Fast Fourier Transform and self-attention mechanism, respectively. The intuition of the block design of EnvEncoder and EntEncoder are discussed in Appendix B. After disentangled, for  $H_{e}$  we compare it to an Environment Codebook and select the closest vector as the final representation  $\hat{H}_{e} \in \mathbb{R}^{N \times F}$ . The handling of  $H_{i}$  will be explained in Section 4.2.
Environment Codebook. To stratify the environment  $E$  in Eq. 1, we draw inspiration from [46] and develop a trainable environment codebook  $e = \{e_1, e_2, \dots, e_K\}$ , which defines a latent embedding space  $e \in \mathbb{R}^{K \times F}$ . Here,  $K$  signifies the discrete space size (i.e., the total number of environments), and  $F$  denotes the dimension of each latent vector  $e_i$ . As depicted in Figure 3, after acquiring the environment representation  $H_e$ , we use a nearest neighbor look-up method in the shared embedding space  $e$  to identify the closest latent vector for each node's environment representation. Given the environment feature of the  $i$ -th node  $H_e(i) \in \mathbb{R}^F$ , this process is calculated in the posterior categorical distribution  $q(z_{ij} = k|H_e(i))$  as follows:
$$
q \left(z _ {i j} = k \mid H _ {e} (i)\right) = \left\{ \begin{array}{l l} 1 & \text {f o r} k = \arg \min  _ {j} \left| \left| H _ {e} (i) - e _ {j} \right| \right| _ {2}, \quad j \in \{1, 2, \dots , K \}, \\ 0 & \text {o t h e r w i s e .} \end{array} \right. \tag {3}
$$
Once obtaining the latent variable  $z \in \mathbb{R}^{N \times K}$ , we derive the final environment representation  $\hat{H}_e \in \mathbb{R}^{N \times F}$  by replacing each row in  $H_e$  with its corresponding closest discrete vector in  $e$ . Note that this categorical probability in Eq. 3 is only used during the training process, whereas a soft probability is used during testing to enable generalization to unseen environments. This soft probability signifies the likelihood of environment representation for each node belonging to each environment, denoted as  $\hat{H}_e(i) = \sum_{j=1}^K q(z_{ij}|H_e(i))e_j$ , where  $q(z_{ij}|H_e(i))$  ranges from 0 to 1. More discussion on how we achieve OoD generalization is provided in Appendix G.
Representation Disentanglement. We expect the environment and entity representations to be statistically independent, where entity representations carry minimal information (MI) about the environment. To achieve this, we employ an optimization objective inspired by Mutual Information Neural Estimation [3]. MI measures the information shared between  $H_{e}$  and  $H_{i}$ , which is calculated using the Kullback-Leibler (KL) divergence between the joint probability  $P(H_{e};H_{i})$  and the product of marginal distributions  $P(H_{e})P(H_{i})$ :
$$
\mathcal {I} \left(H _ {e}, H _ {i}\right) = D _ {K L} \left[ P \left(H _ {e}, H _ {i}\right) \right] \left| \left| P \left(H _ {e}\right) P \left(H _ {i}\right) \right. \right]. \tag {4}
$$
By minimizing the mutual information, the overlap between  $H_{e}$  and  $H_{i}$  decreases. When it approaches zero, each representation is ensured to possess only self-contained information. This approach transforms disentanglement into an optimization issue, which will be introduced in Section 4.3.
4.2 Spatial Context Filtering
Overview. Until now, we finished separating environment-entity and stratifying the environment  $E$ . We then shift our focus to the entity  $H_{i}$ . Our goal is to derive a surrogate  $\hat{H}_{i}$  (i.e., the latent variable of  $X^{*}$  in Figure 2c) that emulates a node representation containing only information propagated based on genuine causation within their spatial context. As emphasized in Section 1, it is essential to account for the ripple effects of causal relationships to accurately learn the surrogate. Thus the challenge is how can we effectively model the ripple effect of dynamic causal relationships? Since nodes' causal relations can naturally be regarded as edge features, an intuitive solution for this challenge is to execute convolution operations on edges. Inspired by [16], we build a higher-order graph over edges and use an edge-level spectral filter, i.e., the Hodge-Laplacian operator, to represent the propagation of causal relations. This forms the core of HL Deconfounder block (see Figure 4b).
Moreover, the locational information of nodes incorporates a more global spatial perspective, which can be seamlessly implemented using a position embedding [31]. We showcase the effects of it in Appendix F. Ultimately, the HL Deconfounder block ingests several inputs: the entity variable  $H_{i}$ , the edge signal  $X_{ed} \in \mathbb{R}^{M \times F'}$ , the boundary operator  $\partial[24]$ , and the position embedding  $P \in \mathbb{R}^{N \times D_p}$  where  $M$  and  $F'$  mean the number of nodes and the dimension of their features in the built higher-order graph, respectively, and  $D_{p}$  denote the embedding dimension. These inputs are processed to yield  $\hat{H}_{i} \in \mathbb{R}^{N \times F}$ . We then provide an exhaustive breakdown of this process, complemented with associated formulations.
Edge Graph Construction. We first construct a higher-order graph over edges by employing the boundary operator [24], which is a mathematical tool used in graph topology to connect different graph elements. Specifically, the first-order boundary operator  $\partial_1$  maps pairs of nodes to edges, while the second-order boundary operator  $\partial_2$  maps pairs of edges to triangles. Here we use  $\partial_1$  and  $\partial_2$  to establish a higher-order graph on edges, which facilitates subsequent convolution operations.