Spatio-temporal Prototype-based Hierarchical Learning for OD Demand Prediction

1 Introduction

With rapid urbanization and the ever-growing demand for transportation, urban mobility systems face significant challenges, including traffic congestion [Gong et al., 2020b; Ji et al., 2023], resource inefficiency [Gross et al., 2006], and prolonged passenger wait times [Gong et al., 2020a]. Understanding passenger movement patterns is a critical pathway to optimizing urban transportation systems. Origin-Destination (OD) demand prediction, which focuses on forecasting the number of passengers traveling from specific origins to destinations, has emerged as a critical research area in urban computing [Rong et al., 2024]. OD demand prediction can identify the underlying spatio-temporal patterns of passenger movements and forecast their mobility intentions, thereby offering essential insights into the dynamics of urban mobility.

∗Corresponding author.

Due to its significance in practical applications, OD demand prediction has recently attracted extensive attention in both academic and industrial communities [Wang et al., 2019; Han et al., 2022; Ye et al., 2024; Yu et al., 2025]. Some studies transform passenger flow data into grid-based representations and utilize CNNs to capture spatial dependencies [Wang et al., 2019; Jiang, 2023]. To model more complex spatial relationships, other approaches represent passenger flow data as graphs and employ GNNs to extract intricate spatio-temporal patterns [Liu et al., 2022; Ye et al., 2024; Liu et al., 2024]. However, despite these advancements, existing methods still face several limitations. One significant issue is the failure to adequately capture the distinct characteristics of a station when it functions as an origin (O) compared to when it serves as a destination (D). As illustrated in Figure 1, the demand patterns of Node c as an origin can differ substantially from its traffic demand distribution patterns as a destination. A clear real-world example of this phenomenon is the contrasting traffic patterns observed in workplace districts during morning and evening peak hours. Previous studies often model origin and destination features jointly, which overlooks their inherent semantic differences.

The second challenge lies in capturing the implicit spatiotemporal patterns from a global perspective. As illustrated in Figure 1, different nodes may exhibit similar spatio-temporal evolution patterns regardless of their geographical proximity. We refer to these shared patterns, which are applicable across multiple nodes, as spatio-temporal prototypes. The prototypes often represent the traffic demand evolution features of a specific category of nodes, providing a higher-level understanding of urban mobility dynamics. Effectively capturing the spatio-temporal prototypes and modeling their intricate correlations is a crucial yet unexplored step in OD prediction.

To address these challenges, we propose STPro, a spatiotemporal prototype based dual-branch hierarchical framework for OD demand prediction. Our framework works at two levels to comprehensively capture spatio-temporal dynamics. At the macro level, we focus on learning spatiotemporal prototypes to capture the shared mobility patterns across nodes. These prototypes are constructed using dynamic clustering, which adaptively groups nodes with similar spatio-temporal evolution trends. At the micro level, we extract the fine-grained spatio-temporal features for individual nodes with a hypervariate graph that fully connects all nodes. To address the inherent semantic differences between origin and destination roles, we design a dual-branch architecture that separately models the characteristics of nodes as origins (O-branch) and destinations (D-branch). At each branch, spatio-temproal features learned from the micro and macro levels are fused to generate the unified embeddings. To enable effective interaction between the two branches, we employ a cross-attention mechanism to integrate mutual information between the O-branch and D-branch, which is then used to make final OD demand predictions. Extensive experiments conducted on two real-world datasets demonstrate our method, as it consistently outperforms baseline methods. In summary, our work makes the following key contributions:

• We propose STPro, a hierarchical framework for OD demand prediction. We utilize dynamic clustering to construct spatio-temporal prototypes at the macro level and leverage the fine-grained node features at the micro level. This hierarchical integration enriches spatiotemporal semantics and enhances prediction accuracy.
• We design a dual-branch architecture to separately model the semantic characteristics of nodes as origins (O-branch) and destinations (D-branch). To model the dynamic interactions between origin and destination roles, a cross-attention mechanism is employed to extract and integrate mutual information between the two branches.
• Extensive experiments are conducted on two widely used real-world datasets to evaluate the STPro. The experimental results demonstrate that our approach can build effective spatio-temporal prototypes and achieve state-of-the-art performance in both single-step and multi-step OD demand prediction.

2 RelatedWorks

2.1 Origin-Destination Demand Prediction

Origin-destination demand prediction is a critical task in urban mobility analysis. It aims to forecast traffic demand between two locations and has attracted significant research attention in recent years. With the rapid development of deep learning techniques, neural network-based methods have emerged as the dominant approach for addressing this task. Shi et al. [2020] used LSTM units to extract temporal characteristics for each pair of OD and separately learned the spatial dependencies of the origins and destinations. Ke et al. [2021] represented OD pairs with multiple OD graphs and developed a spatio-temporal encoder-decoder residual framework to model the spatial dependencies between different OD pairs and the temporal dependencies of OD pairs themselves. Liu et al. [2022] proposed a heterogeneous information aggregation mechanism, which fully utilizes incomplete OD matrices from historical data to jointly learn the evolution patterns of OD and DO rides. Liu et al. [2024] introduced ODMixer, a method for learning traffic evolution by analyzing all OD pairs. They proposed a fine-grained spatio-temporal MLP architecture specifically designed for metro OD prediction.

However, previous studies do not explore the hierarchical spatio-temporal information inherent in OD demand, which is crucial for capturing both global trends and local variations. Furthermore, most of them predict the OD matrix without considering the distinct semantic attributes and intrinsic correlations between origins and destinations

2.2 Prototype Learning

Prototypes have been extensively utilized in various machine learning paradigms, including transfer learning [Quattoni et al., 2008], multi-task learning [Kang et al., 2011], and few-shot learning [Snell et al., 2017; Liu et al., 2019; Li et al., 2021a]. Traditionally, a prototype is defined as the mean feature vector of samples within the same class [Wieting et al., 2015; Babenko and Lempitsky, 2015]. This definition leverages the central tendency of feature distributions to represent class-specific information, which has proven effective in tasks requiring generalization across similar instances. For instance, in distributed machine learning systems, prototypes have been employed to capture taskagnostic information, enabling the fusion of multi-task models for new tasks [Hoang et al., 2020]. Furthermore, the generalization capability of prototypes has been harnessed in federated learning, where they assist local training by aggregating semantic knowledge from distributed data sources [Michieli and Ozay, 2021; Tan et al., 2021; Li et al., 2021b; Mu et al., 2023].

In this paper, the spatio-temporal prototypes are defined as dynamic clustering centers that adaptively capture shared mobility patterns across nodes. Unlike static mean prototypes, our approach leverages dynamic clustering to iteratively update prototypes based on evolving spatio-temporal characteristics. This dynamic nature allows the prototypes to reflect temporal variations and spatial dependencies, which are critical for accurate OD demand prediction.

3 Problem Statement

In this section, we provide the definition of origin-destination demand prediction.

The primary objective of this work is to predict future OD demands based on past observed demand data. Given the passenger flow between stations (or regions) at a specific time, we formulate this information into an OD matrix Xt ∈ RN×N , where N denotes the number of stations. We treat these stations as nodes in the urban network in the following context. Here, we represent the node features from two perspectives and generate the input for the O-branch and Dbranch. The i-th row of Xt represents the flow from the node i to all other nodes. From the origin perspective, we consider the origin feature of the node vOi t ∈ RN and the set of all nodes is VOt = {vO1 t , ..., vOi t , ..., vON t } ∈ RN×N . Similarly, the j-th column of the Xt represents the flow from all nodes to node j. The destination feature of the node is vDj t , and the set of all nodes is VDt = {vD1 t , ..., vDi t , ..., vDN t } ∈ RN×N . In this work, we utilize the passenger flow from previous T time steps to forecast the future OD demand matrix. The origin feature of T time steps VO will be the input for the Obranch and the destination feature VD will be the input for the O-branch. Finally, the prediction task is formulated as:

4 ProposedMethod

In this paper, we propose a dual-branch hierarchical model capable of learning the multi-level information of OD to effectively model the spatio-temporal distribution of OD. As shown in Figure. 2, our model consists of an O-branch hierarchical module, a D-branch hierarchical module, and a crossattention mechanism for O-D interaction modeling. Specifically, at the micro level of each branch, an OD hypervariate graph is constructed to learn spatio-temporal unified node information. Meanwhile, dynamic clustering is applied to generate OD prototypes at the macro level, exploring typical spatio-temporal patterns in the city. Subsequently, the dualbranch information is interacted to enhance the modeling. We develop a hierarchical OD framework to jointly predict multilevel complete OD demand for the future.

4.1 Hierarchical ODModeling

We propose a hierarchical model for OD demand forecasting that integrates macro and micro levels to capture both global and local spatio-temporal patterns. The micro level captures spatio-temporal local features via fully connected nodes, while the macro level extracts global OD flow patterns by dynamic clustering. Hierarchical interaction combines two levels, providing multi-level representations in the OD system. However, the origin and destination OD flow of a node is affected by different factors that we should consider separately. Therefore, we model the hierarchical structure in both branches (O-branch and D-branch).

OD Hypervariate Graph

Urban OD flow is essentially a spatio-temporal traffic system with a non-Euclidean topology. Recently, GCN has been proven to be effective for spatio-temporal non-Euclidean data embedding [Yi et al., 2024]. However, the unified spatiotemporal graph structure is usually unknown in OD flow scenarios [Zhang et al., 2022; Han et al., 2022; Liu et al., 2022]. Inspired by these works, we construct spatio-temporal OD hypervariate graphs with two branches and utilize graph convolutional units to learn the spatio-temporal representations for OD prediction.

To capture long-term spatio-temporal information, we choose NT nodes from T time steps for the OD hypervariate graph. In the O-branch, we define the set of nodes for T time steps as UOt = {VOt-T+1 , ..., VOt } ∈ RNT×N , and construct an OD hypervariate graph GOt = {UOt , AOt }. Since the prior graph structure is unknown, and nodes are spatially and temporally correlated with each other because of the time lag effect [Wei, 1990], we assume all nodes in the graph GOt are fully connected. Therefore, AOt ∈ RNT×NT is the adjacency matrix initialized to make GOt as a spatio-temporal fully-connected graph. In the same way, in the D-branch, we construct the graph GDt = {UDt , ADt }.

Based on the fully connected graphs, we utilize the standard graph convolution [Kipf and Welling, 2016] to learn spatio-temporal representations HO/D t ∈ RNT×d . In detail, for each branch, we update its representation by aggregating the state embedding vectors of all its fully connected neighbors including itself. In formulation, the representation of HO/D t follows:

where WO/D t ∈ RN×d is a weight matrix; the superscript O and D denotes the spatio-temporal representation generation processes of the O-branch and D-branch, respectively.

OD Prototypes Embedding

The traffic system naturally exhibits a spatio-temporal hierarchical structure [Guo et al., 2021]. OD system not only includes the basic micro layer of spatio-temporal nodes but also has the macro layer of prototypes. While focusing on the local spatio-temporal behavior of individual OD pairs, we also aim to capture the global spatio-temporal OD flow patterns.

To explore the dynamic hierarchical relationships in the OD network, we use a learnable matrix ΘO/D to dynamically cluster spatio-temporal nodes into different spatio-temporal prototypes, where ΘO/D ∈ Rd×I and I is the number of the prototypes. In detail, the incidence probability matrix ΛO/D t ∈ RNT×I between spatio-temporal nodes and prototypes is formalized as:

In OD systems, there are complex dependencies between prototypes. Therefore, we introduce a learnable matrix PO/D ∈ RI×I to capture the connection relationships of the prototypes. Here, PO/D is regarded as the spatio-temporal prototype graph used for graph convolution. Then, we utilize the graph convolution [Zhao et al., 2023] to learn spatiotemporal prototype embeddings, which have spatio-temporal global information from the hierarchical OD demand network. Specifically, we generate each prototype's embedding EO/D t ∈ RI×d by aggregating the embeddings of all spatiotemporal nodes and the spatio-temporal prototype graph:

where ϕ is a non-linear activation function.

OD Hierarchical Fusion

To realize the interaction between the macro and micro layers, we propose an OD hierarchical fusion to integrate local spatio-temporal node features and global spatio-temporal prototype features. First, we define the transformation matrix TranO/D t ∈ RNT×I based on the incidence probability matrices in the previous section:

We transform the obtained local node features and the global prototype features into the same dimension using the transformation matrix:

where EO/D tran ∈ RNT×d and each variable in EO/D tran contains information about all nodes of the same prototype. To integrate features from different levels, we design the hierarchical fusion function as follows:

where the readout is a summation operation.

4.2 Dual ODAttention

After obtaining clear multi-level representations from the two branches, we explicitly model the correlation between the two branches to enhance the overall representation of OD demand. Inspired by Liu et al. [2022], we introduce a dual OD attention mechanism to model the O and D distributions jointly by propagating their mutual information in a dual manner. Following this interaction, the multi-level O and D features become informative for OD demand flow prediction.

Our dual OD attention is implemented with cross-branch cross-attention, where the bidirectional transformer propagates information from the O-branch to the D-branch, as well as from the D-branch to the O-branch. Specifically, FOt and FDt are first respectively fed into three linear layers for query, key, and value embedding:

where all convolution kernel sizes are 1 ∗ 1 with individual parameters, and the biases for the query, key, and value have dimensions of NT × d . Same as FOt and FDt , these query, key, value features also have a dimension of NT × d.

Based on attention mechanisms, the cross-branch information propagation is performed with the following formulations:

where softmax(QOt (KDt )⊤) and softmax(QDt (KOt )⊤) are two propagation coefficients that dynamically determine the amount of information propagated between the features of the O-branch and the D-branch. Through this process, the representations of the O-branch and the D-branch not only reinforce each other but also better capture the intrinsic interaction between the origin and destination attributes of nodes in the OD system.

4.3 ODDemand Prediction

Finally, the extracted features are fused and used for OD demand prediction. The model is capable of predicting both single-step and multi-step OD demands. For multi-step OD demand prediction, the model directly outputs predictions for multiple future time steps to avoid error accumulation. The model uses Mask OD Loss to measure prediction errors and continuously improve its performance.

To predict the OD demand at timestamp t + 1, the model utilizes the interaction features from previous timestamps, including O2D features and D2O features. First, we combine the two features and multiply them by the transformation matrix Wz resulting in a matrix with dimensions NT ×N. Then, we reshape the matrix into T ×N×N. Finally, we sum along the temporal dimension to obtain a combined matrix with dimension of 1 ×N ×N, aggregating the information across T time steps. The combination feature can be expressed as follows:

After combining spatio-temporal features, a convolution layer is employed to make the prediction. The OD demand prediction result at the time t +m can be expressed as follows:

where ˆ Y t+m is the prediction results for future m time steps and Ker denotes the convolution kernel, with the number of kernels corresponding to the predicted the next m time steps, and the size of each kernel being 1 ∗ 1.

5 Experiments

5.1 Experimental Setup

Datasets

The two OD datasets used in our experiments are NYCTOD2018 and NYC-TOD2019 [Yu et al., 2025]. New York City taxi datasets were provided by the New York City Taxi and Limousine Commission (TLC). We selected the yellow taxi data for this study. According to the division of TLC, Manhattan is divided into 69 nodes. The NYC-TOD2018 dataset collected passenger demand data for 69 areas in Manhattan from January 1,2018 to December 31, 2018, and the NYC-TOD2019 dataset collected passenger demand data for 69 areas in Manhattan from January 1, 2019, to December 31, 2019. The time interval for both datasets is 1 hour.

Evaluation Setting and Metrics

All models take 3 hours (3 time steps) of historical data as input and make predictions for the future 1/2/3 hours (1/2/3 time steps) of OD demand. For both datasets, we use 60% of the data for training, 20% for validation, and the remaining 20% for testing the model's performance. We choose widely used Mean Absolute Error (MAE) and Root Mean Square Error (RMSE) as the evaluation metrics.

Baselines

We conduct a comprehensive comparison of the proposed STPro with a wide range of baselines, spanning from traditional statistical methods to state-of-the-art graph-based approaches. The baselines include: (1) the statistical-based method HA; (2) RNN-based machine learning methods, such as GRU [Cho et al., 2014] and LSTM [Hochreiter and Schmidhuber, 1997]; (3) recent graph-based traffic prediction methods, including STGCN [Yu et al., 2017], GraphWaveNet [Wu et al., 2019], and DCRNN [Li et al., 2017]; and (4) state-of-the-art graph-based OD prediction methods, such as GEML [Wang et al., 2019], CMOD [Han et al.,

2022], HIAM [Liu et al., 2022], C-AHGCSP [Ye et al., 2024], and ODMixer [Liu et al., 2024]. Notably, STGCN, GraphWaveNet, and DCRNN are representative models originally designed for traffic forecasting. To adapt them for OD demand prediction, we re-implemented these models based on their official code.

5.2 Comparison with AdvancedMethods

The performance summary of single-step and multi-step OD demand predictions for all comparison methods is presented in Table 1. From the results, we observe that traditional methods exhibit relatively higher MAE and RMSE values across different time intervals, indicating their limitations in capturing complex spatio-temporal patterns. In contrast, graphbased traffic prediction methods, such as STGCN, GraphWaveNet, and DCRNN, demonstrate improved performance over traditional baselines and RNN-based models by leveraging graph convolution to model spatial and temporal dependencies. Among these, GraphWaveNet achieves notable results, particularly in multi-step predictions, due to its use of dilated causal convolution, which effectively captures longterm temporal dependencies. However, this mechanism appears less effective in capturing fine-grained details critical for single-step predictions, which may explain its relatively weaker performance in such scenarios. While graph-based methods have made significant advancements in modeling traffic dynamics, they primarily focus on establishing relationships between individual traffic points, often overlooking the unified spatio-temporal prototypes.

Our proposed method introduces a novel hierarchical spatio-temporal framework that effectively captures mobility patterns at multiple scales, addressing the limitations of existing approaches. By integrating micro-level fine-grained details and macro-level global trends, our model achieves a balanced representation of OD demand dynamics, enabling accurate predictions for both single-step and multi-step forecasting tasks. The experimental results demonstrate the robustness and effectiveness of our approach. On the NYCTOD2018 dataset, our model achieves a 5.58% reduction in RMSE for 2-hour predictions, showcasing its ability to capture short-term mobility patterns. Similarly, on the NYCTOD2019 dataset, our model improves RMSE by 1.55% for 2-hour predictions, further validating its generalizability across different datasets. These consistent improvements highlight the superiority of STPro in modeling the hierarchical spatio-temporal dynamics of OD systems.

5.3 Ablation Studies

In this section, we conduct comprehensive ablation studies to evaluate the contribution of each component in the proposed STPro. Table 2 summarizes the average performance of all model variants for the 3-hour prediction task. The key components under investigation include the dynamic clustering module (Prototypes), the O-branch for modeling nodes as origins, and the D-branch for modeling nodes as destinations.

Effectiveness of Prototypes

From Table 2, it can be observed that the model with prototypes significantly outperforms other variants in OD prediction, indicating that the dynamic clustering module effectively enhances OD prediction performance. Specifically, the dynamic clustering module generates spatio-temporal prototypes by aggregating features of similar samples into dynamic centers. These prototypes enable the model to better learn the intrinsic patterns of OD demand. For instance, on the NYC-TOD2018 dataset, the STPro model with prototypes achieves a 1.57% reduction in MAE and a 4.61% reduction in RMSE compared to the variant without prototypes. Similar performance improvements were also observed on the NYCTOD2019 dataset. These experimental results further validate the effectiveness of the dynamic clustering module in capturing spatio-temporal correlations.

Effectiveness of Dual OD Branch

The results in Table 2 also highlight the importance of the dual OD branch structure, which explicitly models the interactions between origins and destinations. The model incorporating OD branch interaction significantly outperforms the variant without interaction, demonstrating its ability to capture the complex relationships between origins and destinations. These correlations are often overlooked or inadequately modeled in traditional approaches, and the OD branch interaction module addresses this limitation.

6 Model Analysis

6.1 Impact of the Number of Prototypes

In this section, we analyze the impact of the number of prototypes I on the prediction performance. As described in the methodology, we cluster N × T nodes to form I prototypes, which serve as high-level representations of shared spatiotemporal patterns. To determine the value of I, we first conducted a wide-range grid search and identified that models with I around twenty consistently achieved superior performance. Based on this observation, we performed a detailed study of I around twenty to evaluate the average performance for both single-step and multi-step predictions.

As shown in Figure 3, both MAE and RMSE reach the minimum values for single-step and multi-step predictions when I equals 23. When the number of prototypes I deviates from 23-either by being too small or too large-the prediction errors increase. We infer that the number of prototypes should align with the real-world number of macro spatio-temporal patterns inherent in the dataset. The prototypes can reflect the true diversity of real-world traffic patterns captured in the data. An appropriate value of I enables the model to better capture the underlying spatio-temporal correlations in the OD demand data. While the choice of I has an implicit impact on prediction performance, our hierarchical approach consistently achieves competitive results compared to state-of-theart methods.

6.2 Exploration of Prototype GenerationMethods

In this section, we conduct a detailed exploration of various spatio-temporal prototype generation methods to validate the effectiveness of our proposed approach. Specifically, we compare our method with the following approaches: K-Means clustering (K C), Agglomerative clustering (A C), Temporal clustering (T C), and Spatial clustering (S C). The experimental results, presented in Figure 4, clearly demonstrate that our proposed clustering method ST C significantly outperforms other approaches. This superior performance stems from its ability to jointly model spatial and temporal information, effectively capturing the complex dependencies inherent in spatio-temporal data. In addition to their suboptimal performance in prediction accuracy, K C and A C suffer from high computational complexity and low efficiency. On the other hand, S C and T C are limited by their focus on either spatial or temporal features. Our proposed method effectively captures the complex spatial and temporal relationships in OD demand data and ensures a more accurate prediction.

6.3 Experiments on the Stability of the STPro

To further validate the stability of our model, we analyze the prediction performance of STPro across different time periods. This analysis focuses on the model's ability to handle varying levels of OD flow complexity during weekdays and weekends. We provide the comparison with the six bestperforming baselines in Figure 5. On weekdays, all methods often exhibit higher prediction errors during the morning and evening periods, primarily due to the increased complexity of traffic patterns during peak hours. Our method not only achieves optimal prediction results across all time periods but also demonstrates exceptional performance during peak hours, outperforming some methods even during off-peak hours. Even though the traffic patterns on weekends differ from those on weekdays, our method consistently achieves the best prediction performance across both scenarios. By maintaining stable and accurate predictions under diverse conditions, our method ensures reliable performance across a wide range of real-world applications, from weekday commutes to weekend leisure travel.

6.4 Model Applicability Analysis

To validate the applicability of STPro, we conducted additional experiments on datasets from a new city and extended the model to OD flow prediction. Experimental results presented in Table 3 demonstrate that STPro maintains outstanding prediction performance on the HZMetro dataset [Yu et al., 2025], collected from Hangzhou, China. For OD flow prediction, it focuses on real-time flow estimation based on incomplete trip observations (e.g., subway station entries without exits) [Yu et al., 2025]. The results on HZMOD [Liu et al., 2022] highlight STPro's robustness against sparsity issues in OD flow prediction.

7 Conclusion

In this paper, we present STPro, a dual-branch hierarchical model designed for OD demand prediction. STPro captures multi-level spatio-temporal information, combining fine-grained details at the micro level with city-wide patterns at the macro level. STPro effectively models urban mobility dynamics by leveraging dynamic clustering to identify spatiotemporal prototypes and employing a dual-branch structure to model origin-destination interactions. Experiments on two real-world datasets show that STPro achieves state-of-the-art performance in both single-step and multi-step OD demand predictions. This work advances OD demand prediction and provides a robust framework for spatio-temporal modeling in intelligent transportation systems. Future research will explore extending the framework with additional data sources and applying it to other spatio-temporal tasks.

Acknowledgments

This work was supported in part by the National Natural Science Foundation of China (62476154, 62202270), and Major Basic Research Project of Shandong Provincial Natural Science Foundation (ZR2024ZD03) and the Shandong Excellent Young Scientists Fund (Oversea) (2022HWYQ044), and the Taishan Scholar Project of Shandong Province (tsqn202306066), and Shandong University Qilu Young Scholars Program.

Contribution Statement

Shilu Yuan and Xiaoyu Li have made equal contributions to this work.