SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhonghang Li $^{1,2}$
Lianghao Xia²
Yong Xu $^{1,3,4}$
Chao Huang*
$^{1}$  South China University of Technology,  $^{2}$  University of Hong Kong,
<sup>3</sup> PAZHOU LAB, <sup>4</sup> Guangdong Key Lab of Communication and Computer Network
{bjdwh.zzh,chaohuang75}@gmail.com aka_xia@foxmail.com yxu@scut.edu.cn
In recent years, there has been a rapid development of spatio-temporal prediction techniques in response to the increasing demands of traffic management and travel planning. While advanced end-to-end models have achieved notable success in improving predictive performance, their integration and expansion pose significant challenges. This work aims to address these challenges by introducing a spatio-temporal pre-training framework that seamlessly integrates with downstream baselines and enhances their performance. The framework is built upon two key designs: (i) We propose a spatio-temporal mask autoencoder as a pre-training model for learning spatio-temporal dependencies. The model incorporates customized parameter learners and hierarchical spatial pattern encoding networks. These modules are specifically designed to capture spatio-temporal customized representations and intra- and inter-cluster region semantic relationships, which have often been neglected in existing approaches. (ii) We introduce an adaptive mask strategy as part of the pre-training mechanism. This strategy guides the mask autoencoder in learning robust spatio-temporal representations and facilitates the modeling of different relationships, ranging from intra-cluster to inter-cluster, in an easy-to-hard training manner. Extensive experiments conducted on representative benchmarks demonstrate the effectiveness of our proposed method. We have made our model implementation publicly available at https://github.com/HKUDS/GPT-ST.
Intelligent Transportation Systems (ITS) is a rapidly growing research field driven by advancements in sensors and communications technology [1, 12]. Spatio-Temporal (ST) prediction, including traffic flow and ride demand, aims to forecast future trends, facilitating decision-making in traffic management and risk response [43], and enhancing people's lives with future travel plans [6]. Deep learning methods, such as RNNs and CNNs, have been widely adopted to model temporal and spatial correlations [5, 51, 50, 45]. Recently, GNNs have gained attention for spatial and temporal prediction tasks, achieving remarkable results. Initially, researchers built adjacency matrices based on region attributes, such as distance [48, 54]. Subsequent studies introduced learnable graph structures to capture spatial correlations [40, 39], while others focused on dynamic relationships between regions, proposing methods for constructing dynamic graphs [14, 40]. These advancements have significantly contributed to the development of spatio-temporal prediction techniques.
Despite the remarkable results achieved by existing methods in spatio-temporal prediction, there are several issues that remain inadequately addressed. i) Lack of customized representations of specific spatio-temporal patterns. Customization can be categorized into two key aspects: time-dynamic and node-specific properties in both the time and spatial domains. Time-dynamic patterns exhibit variations across different time periods, such as contrasting patterns between weekends and
weekdays within the same region. Moreover, correlations between regions dynamically evolve over time, surpassing the capabilities of static graph representations. Node-specific patterns highlight the distinct time series observed in different regions, rather than a shared pattern. Additionally, it is important to ensure that different regions retain their individual node characteristics even after message aggregation, to prevent interference from prominent nodes in the spatial domain. We believe that encoding all of these customized characteristics is essential to guarantee the model robustness. However, existing works often fall short in providing a comprehensive consideration of these factors. ii) Insufficient consideration of spatial dependencies at different levels. Most approaches primarily focus on pairwise associations between regions when modeling spatial dependency, but they overlook the semantic relevance at different spatial levels. In the real world, regions with similar functions tend to exhibit similar spatio-temporal patterns. By performing clustering analysis on different regions, the model can explore common characteristics among similar regions, thereby facilitating improved spatial representation learning. Moreover, there is a lack of adequate modeling of high-level regional relationships across time in current studies. Spatio-temporal patterns between different types of high-level regions may exhibit dynamic transfer relationships. For instance, during work hours, there is a notable movement of people from residential areas to work areas, as illustrated in the right part of Figure 1. In this context, changes in people flow within residential areas can provide valuable auxiliary signals for predicting people flow in work areas. It highlights the importance of incorporating both fine-grained and coarse-grained correlations among different levels of regions in order to enhance the predictive capabilities of spatio-temporal prediction models.
An intuitive approach to address the aforementioned challenges is to develop an end-to-end model. However, the current models face difficulties in benefiting from such an approach, as each module in the state-of-the-art (SOTA) model is intricately refined, and any disassembly and integration may lead to a degradation in prediction performance. So, is there a strategy that can allow existing spatiotemporal methods to leverage the advantages of an end-to-end model? The answer is affirmative. Recent pre-training frameworks, such as ChatGPT [3, 27] in the field of natural language processing (NLP) and MAE [15] in computer vision (CV), have been proposed by pioneers and widely studied in their respective domains. These pre-training architectures construct unsupervised training tasks, involving masking or other techniques, to learn better representations and improve
downstream task performance. However, the application of such scalable pre-training models in the field of spatio-temporal prediction has been limited thus far. To address the aforementioned challenges, we propose a novel framework called  $\underline{\mathrm{G}}$  enerative Pre-Training for Spatio-Temporal prediction (GPT-ST). This framework is designed to seamlessly integrate into existing spatio-temporal prediction models, enhancing their performance. Our contributions are summarized as follows:
- We present GPT-ST, a novel pre-training framework designed for spatio-temporal prediction. This framework seamlessly integrates into existing spatio-temporal neural networks, resulting in performance enhancements. Our approach combines a model parameter customization scheme with a self-supervised masking autoencoding, enabling effective spatio-temporal pre-training.
- GPT-ST ingeniously utilizes a hierarchical hypergraph structure to capture spatial dependencies at different levels from a global perspective. Through collaboration with a thoughtfully designed adaptive mask strategy, the model acquires the capability to model both intra-cluster and intercluster spatial relations among regions, thereby generating robust spatio-temporal representations.
- We conduct extensive experiments on real-world datasets, and the demonstrated improvement in the performance of diverse downstream baselines showcases the superior performance of GPT-ST.
Figure 1: The motivations behind GPT-ST. The upper left figure demonstrates that traffic patterns in the same region exhibit variations across different time periods. Similarly, the lower left figure depicts how this variation can occur between different regions during the same time period. In contrast, existing works typically construct static graphs, failing to capture the dynamic nature of region relationships over time, as shown in the middle figures (lower), where region relationships change dynamically. Lastly, the right figure illustrates the commuting phenomenon between regions belonging to different categories.
Deep Spatio-temporal Prediction Techniques. Numerous studies have focused on designing deep neural networks for spatio-temporal data prediction. To capture temporal dependencies, D-LSTM [49] and GWN [40] have developed predictive frameworks based on Recurrent Neural Networks (RNNs) and Temporal Convolutional Networks (TCNs), respectively. GMAN [55] and STGNN [36] have utilized temporal self-attention networks to facilitate long-range temporal learning. In spatial relation learning, DMVST-Net [44] adopted convolutional network to model region associations. STGCN [48] and DCRNN [23] exploited Graph Neural Networks (GNNs) for message passing between regions. STAN [26] consider the interrelationships between regions based on Graph Attention Networks (GAT). Temporal graph networks serve as a similar research baseline aimed at reasoning about dynamic graph structures, such as CAWs [37] and PINT [33]. More recently, some works have explored the combination of ST models with advanced methods, such as Neural-ODE-based [10, 18] and self-supervised learning-based [24], for better spatial and temporal dependencies modelling. Compared to these works, our proposed GPT-ST framework aims at empower diverse deep spati-temporal models with the pre-training stage.
Pre-training methods. Pre-training models have made significant advancements in recent years, with notable examples including BERT [7, 28], vision transformers [9, 8], masked autoencoders [15, 20], and language models [3, 27, 34]. The success of the autoencoder-based approach, exemplified by MAE, has demonstrated the powerful expressive capability of the mask-reconstruction pre-training framework. Building upon this, STEP [31] proposed a pre-training model specifically for time series prediction. In other domains, CMSF [41] introduced a master-slave framework for identifying urban villages. The master model pre-trains region representations, while the slave model fine-tunes specific regions for accurate identification. Researchers have also recognized the need for more efficient mask mechanisms than random masking. AttMask [19] and SemMAE [21] leverage attention mechanisms and semantic information, respectively, to guide the mask operation. AdaMAE [2] proposes adaptively selecting visible signals with greater reconstruction loss, as they contain more information, and using a higher proportion of masking strategy accordingly. Despite the significant progress of pre-training models in the fields of natural language processing (NLP) and computer vision (CV), their application to spatio-temporal prediction remains relatively limited. Additionally, computational efficiency has been a major bottleneck for transformer architectures in general.
Self-supervised learning methods for GNNs. In recent years, self-supervised learning methods for graph data have gained significant attention [53, 38]. Contrastive learning-based GNNs generate different views of the original graph through data augmentation techniques [42, 29]. A loss function is then employed to maximize the consistency of positive sample pairs while minimizing the consistency of negative sample pairs across views. For example, GraphCL [47] generates two views of the graph by applying node dropping and edge shuffling, and performs contrastive learning between them. Another research direction focuses on generative graph neural networks, where the graph data itself serves as a natural supervisory signal for representation learning through reconstruction. GPT-GNN [17] conducts pre-training by reconstructing graph features and edges, while GraphMAE [16] utilizes node feature masking in both the graph encoder and decoder to reconstruct features and learn graph representations. However, spatio-temporal forecasting tasks require simultaneous consideration of complex temporal evolution patterns and spatial correlation mechanisms. Pre-training paradigms specifically designed for such tasks are still an area of exploration and research.
Spatio-Temporal Data  $\mathbf{X}$ . To capture the ST information, we represent it as a three-way tensor  $\mathbf{X} \in \mathbb{R}^{R \times T \times F}$ , where  $R$ ,  $T$ , and  $F$  correspond to the numbers of regions, time slots, and feature dimensions, respectively. Each entry  $\mathbf{X}_{r,t,f}$  is the  $f$ -th feature of the  $r$ -th region in the  $t$ -th time slot.
Spatio-Temporal Hypergraphs. We employ ST hypergraphs for ST modeling. A hypergraph  $\mathcal{H} = \{\mathcal{V},\mathcal{E},\mathbf{H}\}$  is composed of three parts: i) Vertices  $\mathcal{V} = \{v_{r,t}:r\in R,t\in T\}$ , each of which represents a region  $r$  in a specific time slot  $t$ . ii)  $H$  hyperedges  $\mathcal{E} = \{e_1,\dots,e_H\}$ , each of which connects multiple vertices to reflect the multipartite region-wise relations (e.g., all residential regions can be connected by one hyperedge). iii) The vertex-hyperedge connections  $\mathbf{H}\in \mathbb{R}^{N\times H}$ , where  $N$  denotes the number of vertices. To fully excavate the potential of hypergraphs in region-wise relation learning, we adopt a learnable hypergraph scheme where  $\mathbf{H}$  is derived from trainable parameters.
Spatio-Temporal Pre-training Paradigm. Our GPT-ST framework aims at developing a pre-trained ST representation method that facilitates the accuracy on downstream ST prediction tasks such as traffic volume forecasting. As depicted in Figure 2, the workflow of GPT-ST can be divided into the pre-training stage and the downstream task stage. We formulate the two stages as follows:
i) The pre-training stage of our GPT-ST framework adopts the masked autoencoding (MAE) task as the train-
Figure 2: Overall workflow of GPT-ST.
ing objective. The goal of this MAE task is to reconstruct masked information of the spatio-temporal data according to the unmasked ones by learning a ST representation function  $f$ . For the training ST data during the period of  $[t_{K - L + 1}, t_K]$ , we aim to minimize the following objective:
$$
\mathcal {L} \left(\left(1 - \mathbf {M}\right) \odot \mathbf {X} _ {t _ {K - L + 1}: t _ {K}}, \mathbf {W} \cdot f \left(\mathbf {M} \odot \mathbf {X} _ {t _ {K - L + 1}: t _ {K}}\right)\right) \tag {1}
$$
where  $\mathcal{L}$  denotes the measure of prediction deviation,  $\mathbf{M} \in \{0,1\}$  represents the mask tensor, and  $\odot$  denotes the element-wise multiplication operation.  $\mathbf{W}$  denotes a predictive linear layer.
ii) After the pre-training stage, the results of our GPT-ST are then utilized in the downstream task stage. The pre-trained model yields high-quality ST representations to facilitate the downstream prediction tasks such as traffic forecasting. Specifically, the downstream stage is formulated as:
$$
\boldsymbol {\zeta} = f \left(\mathbf {X} _ {t _ {K - L + 1}: t _ {K}}\right); \quad \hat {\mathbf {X}} _ {t _ {K + 1}: t _ {K + P}} = g \left(\boldsymbol {\zeta}, \mathbf {X} _ {t _ {K - L + 1}: t _ {K}}\right) \tag {2}
$$
where  $\zeta$  refers to the representation generated by  $f$  according to the historical spatio-temporal data from the previous  $L$  time slots before the  $K$ -th time slot. The output is the predictions for the next  $P$  time slots. Various existing spatio-temporal neural networks can serve as the prediction function  $g$ .
This section introduces the technical details of the proposed ST pre-training framework GPT-ST. It develops a customized temporal encoder and a hierarchical spatial encoder to obtain large modeling capacity for the pre-trained model. An adaptive mask strategy is proposed to facilitate effective MAE training. Figure 3 demonstrates the model design of our pre-trained spatio-temporal encoding model.
4.1 Customized Temporal Pattern Encoding
Initial Embedding Layer. We begin by constructing an embedding layer to initialize the representation of the spatio-temporal data  $\mathbf{X}$ . The original data undergoes normalization using the Z-Score function [1, 39] and is then masked using the mask operation. Subsequently, a linear transformation is applied to augment the representation as  $\mathbf{E}_{r,t} = \mathbf{M}_{r,t} \odot \bar{\mathbf{X}}_{r,t} \cdot \mathbf{E}_0$ , where  $\mathbf{E}_{r,t} \in \mathbb{R}^d$ ,  $\mathbf{M}_{r,t} \in \mathbb{R}^F$ , and  $\bar{\mathbf{X}}_{r,t} \in \mathbb{R}^F$  represent the representation, mask operation, and normalized spatio-temporal data for the  $r$ -th region in the  $t$ -th time slot. The variable  $d$  denotes the number of hidden units. Additionally,  $\mathbf{E}_0 \in \mathbb{R}^{F \times d}$  represents the learnable embedding vectors for the  $F$  feature dimensions.
Temporal Hypergraph Neural Network. To facilitate global relation learning, we employ the hypergraph neural network for the temporal pattern encoding, specifically as follows:
$$
\mathbf {\Gamma} _ {t} = \sigma \left(\bar {\mathbf {E}} _ {t} \cdot \mathbf {W} _ {t} + \mathbf {b} _ {t} + \mathbf {E} _ {t}\right); \quad \bar {\mathbf {E}} _ {r} = \text {H y p e r P r o p a g a t e} \left(\mathbf {E} _ {r}\right) = \sigma \left(\mathbf {H} _ {r} ^ {\top} \cdot \sigma \left(\mathbf {H} _ {r} \cdot \mathbf {E} _ {r}\right)\right) \tag {3}
$$
where  $\Gamma_t, \bar{\mathbf{E}}_t, \mathbf{E}_t \in \mathbb{R}^{R \times d}$  denote the result, intermediate, and initial region embeddings for the  $t$ -th time slot, respectively.  $\mathbf{W}_t \in \mathbb{R}^{d \times d}, \mathbf{b}_t \in \mathbb{R}^d$  represent the  $t$ -th time slot-specific parameters.  $\sigma(\cdot)$  denotes the LeakyReLU activation. The intermediate embeddings  $\bar{\mathbf{E}} \in \mathbb{R}^{R \times T \times d}$  are calculated by the hypergraph information propagation. It employs the region-specific hypergraph  $\mathbf{H}_r \in \mathbb{R}^{H_T \times T}$  to propagate the temporal embeddings  $\mathbf{E}_r \in \mathbb{R}^{T \times d}$  for the  $r$ -th region along the connections between time slots and  $H_T$  hyperedges, so as to capture the multipartite relations among time periods.
Customized Parameter Learner. To characterize the diversity of temporal patterns, our temporal encoder conducts model parameter customization for both different regions and different time periods. Specifically, the aforementioned time-specific parameters  $\mathbf{W}_t$ ,  $\mathbf{b}_t$  and the region-specific hypergraph
Figure 3: The detailed framework of the proposed GPT-ST.
parameters  $\mathbf{H}_r$  are generated through a learnable process, rather than directly utilizing independent parameters. Specifically, the customized parameters are learned as follows:
$$
\mathbf {H} _ {r} = \mathbf {c} _ {r} ^ {\top} \bar {\mathbf {H}}; \quad \mathbf {W} _ {t} = \mathbf {d} _ {t} ^ {\top} \bar {\mathbf {W}}; \quad \mathbf {b} _ {t} = \mathbf {d} _ {t} ^ {\top} \bar {\mathbf {b}}; \quad \mathbf {d} _ {t} = \operatorname {M L P} \left(\bar {\mathbf {z}} _ {t} ^ {(d)} \mathbf {e} _ {1} + \bar {\mathbf {z}} _ {t} ^ {(w)} \mathbf {e} _ {2}\right) \tag {4}
$$
where  $\bar{\mathbf{H}}\in \mathbb{R}^{d'\times H_T\times T}$ ,  $\bar{\mathbf{W}}\in \mathbb{R}^{d'\times d\times d}$ ,  $\bar{\mathbf{b}}\in \mathbb{R}^{d'\times d}$  are independent parameters for the three generated parameters, respectively.  $\mathbf{c}_r,\mathbf{d}_t\in \mathbb{R}^{d'}$  denote representations for the  $r$ -th region and the  $t$ -th time slot, respectively.  $\mathbf{c}_r$  for  $r\in R$  is a free-form parameter, while  $\mathbf{d}_t$  for  $t\in T$  is calculated from the normalized time-of-day features  $\bar{\mathbf{z}}_t^{(d)}$  and day-of-week features  $\bar{\mathbf{z}}_t^{(w)}$ .  $\mathbf{e}_1,\mathbf{e}_2\in \mathbb{R}^{d'}$  are their corresponding learnable embeddings. This parameter learner realizes spatial and temporal customization by generating parameters according to the features of specific time slots and regions.
4.2 Hierarchical Spatial Pattern Encoding
4.2.1 Hypergraph Capsule Clustering Network
Current spatial encoders primarily focus on capturing local adjacent region-wise relations while disregarding the widespread similarities that exist between distant regions. For example, business districts that are geographically apart can still exhibit similar ST patterns. In light of this, our GPT-ST introduces a hypergraph capsule clustering network to capture global region-wise similarities. This network explicitly learns multiple cluster centers as hyperedges, characterizing the global region-wise similarities. To further enhance the hypergraph structure learning, we incorporate the dynamic routing mechanism of capsule network. This mechanism iteratively updates hyperedge representations and region-hyperedge connections based on semantic similarities. As a result, it improves the clustering capability of hyperedges and facilitates the global modeling of dependencies among regions.
Specifically, we firstly obtain a normalized region embedding  $\mathbf{v}_{r,t} \in \mathbb{R}^d$  for each region  $r$  in time slot  $t$ , using the previous embedding  $\Gamma_{r,t}$  and the squash function [30]. This embedding is then used to calculate the transferred information  $\bar{\mathbf{v}}_{i|r,t} \in \mathbb{R}^d$  from each region  $r$  to each cluster center (hyperedge)  $i$ , within the  $t$ -th time slot. Formally, these two variables are calculated by:
$$
\mathbf {v} _ {r, t} = \operatorname {s q u a s h} \left(\mathbf {V} \boldsymbol {\Gamma} _ {r, t} + \mathbf {c}\right); \bar {\mathbf {v}} _ {i | r, t} = \operatorname {s q u a s h} \left(\mathbf {H} _ {i, r, t} ^ {\prime} \mathbf {v} _ {r, t}\right) \odot \mathbf {v} _ {r, t}; \operatorname {s q u a s h} (\mathbf {x}) = \frac {\left\| \mathbf {x} \right\| ^ {2}}{1 + \left\| \mathbf {x} \right\| ^ {2}} \frac {\mathbf {x}}{\left\| \mathbf {x} \right\|} \tag {5}
$$
where  $\mathbf{V} \in \mathbb{R}^{d \times d}$  and  $\mathbf{c} \in \mathbb{R}^d$  are free-form learnable parameters. The hypergraph connectivity matrix  $\mathbf{H}_t' \in \mathbb{R}^{H_S \times R}$  records the relations between  $R$  regions and  $H_S$  hyperedges as the cluster centroids. It is tailored to the  $t$ -th time slot using the aforementioned customized parameter learner, specifically by  $\mathbf{H}_t' = \text{softmax}(\mathbf{d}_t'^{\top} \bar{\mathbf{H}}')$ . Here,  $\mathbf{d}_t'$  and  $\bar{\mathbf{H}}'$  are temporal features and hypergraph embeddings.
Iterative Hypergraph Structure Learning. With the initialized region embeddings  $\mathbf{v}_{r,t}$  and hypergraph connection embeddings  $\bar{\mathbf{v}}_{i|r,t}$ , we follow the dynamic routing mechanism of capsule network to enhance the clustering effect for the hyperedges. The  $j$ -th iteration is described as follows:
$$
\mathbf {s} _ {i, t} ^ {j} = \sum_ {r = 1} ^ {R} c _ {i, r, t} ^ {j} \bar {\mathbf {v}} _ {i | r, t}; \quad c _ {i, r, t} ^ {j} = \frac {\exp \left(b _ {i , r , t} ^ {j}\right)}{\sum_ {r} \exp \left(b _ {i , r , t} ^ {j}\right)}; b _ {i, r, t} ^ {j} = b _ {i, r, t} ^ {j} + \mathbf {v} _ {r, t} ^ {\top} \operatorname {s q u a s h} \left(\mathbf {s} _ {i, t} ^ {j - 1}\right) \tag {6}
$$
where  $\mathbf{s}_{i,t} \in \mathbb{R}^d$  represents the iterative hyperedge embedding. It is calculated utilizing the iterative hyperedge-region weights  $c_{i,r,t} \in \mathbb{R}$ . The weight  $c_{i,r,t}$  is in turn calculated from the hyperedge embeddings  $\mathbf{s}_{i,t}$  from the last iteration. With this iterative process, the relation scores and the hyperedge representations are mutually adjusted by each other, to better reflect the semantic similarities between regions and the spatial cluster centroids represented by hyperedges.