Heterogeneous Graph Neural Networks for End-to-End Traffic Assignment
  and Traffic Flow Learning
The traffic assignment problem is one of the significant components of
traffic flow analysis for which various solution approaches have been proposed.
However, deploying these approaches for large-scale networks poses significant
challenges. In this paper, we leverage the power of heterogeneous graph neural
networks to propose a novel data-driven approach for end-to-end traffic
assignment and traffic flow learning. Our model integrates an adaptive graph
attention mechanism with auxiliary "virtual" links connecting
origin-destination node pairs, This integration enables the model to capture
spatial traffic patterns across different links, By incorporating the
node-based flow conservation law into the overall loss function, the model
ensures the prediction results in compliance with flow conservation principles,
resulting in highly accurate predictions for both link flow and flow-capacity
ratios. We present numerical experiments on urban transportation networks and
show that the proposed heterogeneous graph neural network model outperforms
other conventional neural network models in terms of convergence rate and
prediction accuracy. Notably, by introducing two different training strategies,
the proposed heterogeneous graph neural network model can also be generalized
to different network topologies. This approach offers a promising solution for
complex traffic flow analysis and prediction, enhancing our understanding and
management of a wide range of transportation systems.
Tong Liu, Hadi Meidani
# Heterogeneous Graph Neural Networks for Data-driven Traffic Assignment

###### Abstract

The traffic assignment problem is one of the significant components of traffic flow analysis for which various solution approaches have been proposed. However, deploying these approaches for large-scale networks poses significant challenges. In this paper, we leverage the power of heterogeneous graph neural networks to propose a novel data-driven approach for traffic assignment and traffic flow learning. The proposed model is capable of capturing spatial traffic patterns across different links, yielding highly accurate results. We present numerical experiments on urban transportation networks and show that the proposed heterogeneous graph neural network model outperforms other conventional neural network models in terms of convergence rate, training loss, and prediction accuracy. Notably, the proposed heterogeneous graph neural network model can also be generalized to different network topologies. This approach offers a promising solution for complex traffic flow analysis and prediction, enhancing our understanding and management of a wide range of transportation systems.

keywords: traffic assignment, graph neural network, traffic flow prediction, flow conservation +
Footnote †: journal: Transportation Research Part C: Emerging Technologies

## 1 Introduction

The traffic assignment problem, as a significant component of transportation network performance analysis, helps understand spatial and temporal traffic flow patterns for effective transportation network management and provides insights into bottlenecks, and traffic congestion (Nie et al., 2004). Congestion is a major concern in urban areas, leading to increased travel times, fuel consumption, and environmental pollution. By identifying the bottlenecks in transportation networks, city planners can make a strategic plan to increase road capacities to mitigate traffic congestion.

To this end, the objective of the traffic assignment problem is to determine the traffic flow distribution and identify traffic bottlenecks on a given road network. Numerous mathematical models have been proposed to tackle traffic assignment problems. The user equilibrium (UE) principle and system optimum (SO) principles are two fundamental but effective solutions to traffic assignment problems. The UE principle, also known as Wardrop's first principle, states that drivers between each origin-destination (OD) pair cannot unilaterally reduce travel costs by unilaterally shifting to another route (Kuang et al., 2021). The system optimum assignment, which is based on Wardrop's second principle, emphasizes driver cooperation to minimize the total system travel time (Seliverstov et al., 2017). These models have been used extensively in practice to estimate traffic flow patterns and optimize transportation infrastructure investments. However, these approaches are computationally expensive for a large-scale network. Substantial efforts have been made to find efficient algorithms to solve this problem (Nie et al., 2004; Di Lorenzo et al., 2015). Furthermore, researchers have proposed various extensions to the UE and SO models, such as the stochastic user equilibrium (Damberg et al., 1996), multi-modal traffic assignment (Pi et al., 2019), and dynamic traffic assignment (Ben-Akiva et al., 2012). These models incorporate additional factors, such as travel time variability, driver behaviors, and network capacity constraints, to provide more realistic traffic assignment solutions. However, there is still a pressing need for more effective and accurate methods for traffic assignment problems, particularly for traffic planning and optimization.
To perform transportation performance analysis, OD demand needs first to be obtained, which is a significant component in traffic assignment problems. OD demand estimation has been widely studied in the transportation engineering literature. The statistical model is one of the common approaches for OD estimation. It leverages multiple data sources to estimate vehicle flow between different locations in a network (Hazelton, 2008; Jin et al., 2014). Furthermore, the increasing availability of crowdsourced data, including GPS data and social media, enables us to extract OD demand and traffic flow data at a large scale and over extended periods (Toole et al., 2015; Zhang et al., 2020). Besides, neural network approaches (Tang et al., 2021; Zhang et al., 2021) have also been applied to estimate OD estimation, which allows for the extraction of complex patterns from high-dimensional and noisy transportation data.

In recent years, data-driven approaches have leveraged the availability of large-scale transportation data and computational resources in various transportation problems. For instance, the recurrent neural network (RNN), including the long short-term memory (LSTM) and gated recurrent unit (GRU) networks, have been applied to different prediction tasks in transportation (Zhaowei et al., 2020; Fang et al., 2022). However, standalone RNNs lack the ability to learn spatial correlation. To overcome this weakness and capture the spatio-temporal relationship of traffic flows at different locations, Ye et al. (2019) combines convolutional neural networks (CNN) with LSTM to predict multiple transportation demands using real-world taxi and sharing bike demand data. Recently, the graph neural network (GNN) has emerged as a powerful tool in transportation performance analysis because of its capability to effectively model and analyze data represented in the form of a graph. Compared with RNNs and CNNs, a key advantage of GNNs is their ability to capture spatial and relational information inherent in transportation networks.

Graph neural networks have been shown to be effective in network modeling and management of transportation systems (Rahman and Hasan, 2022; Liu and Meidani, 2023a). However, these models encountered a few limitations. Firstly, the proposed GNN models can only be trained on a fixed topology, and need to be re-trained whenever the topology of the road network changes. Additionally, these studies didn't adequately consider the dynamic nature of traffic flow, e.g., when link capacities alter due to traffic accidents or damages to roadways. These drawbacks limit the adaptability and applicability of these models in real-world scenarios.

To address this challenge, in this paper, we propose a novel GNN-based approach for static traffic assignment. Specifically, we focus on predicting the flow-capacity ratio of different scenarios in the context of traffic assignment. In particular, to address the long-range effects caused by OD pairs that are multiple-hops apart, we create a heterogeneous GNN that consists of additional "virtual" links. The major contributions of this work are as follows: (1) a heterogeneous GNN model is proposed with "virtual" OD-links and the "real" roadway links, which incorporates the attention-based mechanism and message-passing mechanism; (2) the proposed model also integrates a physics-based loss based on the flow conservation law into the total loss function to accelerate the learning process; (3) the proposed model is generalizable to different road network topologies, link characteristics, and OD demands.

The remainder of this article is structured as follows. General backgrounds on the traffic assignment problem, the neural network and graph neural network models are presented in Section 2. Section 3 includes the explanation of the proposed heterogeneous graph neural network for traffic flow learning. Furthermore, the experiments with urban road networks and generalized synthetic networks are presented to demonstrate the accuracy and generalization capability of the proposed framework in Section 4. Finally, the conclusion and discussion of the proposed framework are presented in Section 5.

## 2 Technical Background

### Traffic Assignment Problem

The traffic assignment problem involves assigning traffic volumes or flows to each edge in the network. Given a transportation network represented as a graph \(\mathcal{G}=(\mathcal{V},\mathcal{E})\), where nodes \(\mathcal{V}\) represent intersections of roads and edges \(\mathcal{E}\) represent roads or links connecting these locations, The general form of traffic assignment problem can be considered as an optimization task:

\[\min_{f}:\quad\sum_{e\in\mathcal{E}}Z_{e}(f_{e}), \tag{1}\]
where \(f_{e}\) and \(Z_{e}(f_{e})\) are the total flow and the link cost function on link \(e\), respectively. The link cost function can be expressed in terms of travel time, travel distance, or other relevant factors. Therefore, the objective function is to minimize the summation of link cost over the graph. The traffic assignment problem can have different formulations depending on the specific objectives and constraints considered in the optimization process. For instance, under user equilibrium formulation, the traffic assignment problem could be solved by optimizing Beckmann's formulation (Beckmann et al., 1956):

\[\begin{split}\min:& z(x)=\sum_{e\in\mathcal{E}} \int_{0}^{x_{e}}t_{e}(\omega)\mathrm{d}\omega\\ \text{s.t.}&\sum_{k}f_{k}^{rs}=q_{rs},\ \forall r,s\in \mathcal{V},\\ & f_{k}^{rs}\geq 0,\ \forall k,r,s\in\mathcal{V},\\ & x_{e}=\sum_{rs}\sum_{k}f_{k}^{rs}\zeta_{e,k}^{rs},\forall e\in \mathcal{E},\end{split} \tag{2}\]

where \(t_{e}(f_{e})\) is the link travel time function, \(q_{rs}\) is the total demand from source \(r\) to destination \(s\), \(f_{k}^{rs}\) represents the flow on \(k^{\text{th}}\) path from \(r\) to s. \(\zeta_{e,k}^{rs}\) is the binary value, which equals to 1 when link \(e\) is on \(k^{\text{th}}\) connecting \(r\) and \(s\).

### Neural Network

Without loss of generality, we will start by considering a neural network with only one layer. Given a \(p\)-dimensional input vector \(h^{k}\in\mathbb{R}^{p}\), \(q\)-dimensional the output \(h^{k+1}\in\mathbb{R}^{q}\) of single layer neural network, The single-layer neural network with index \(k\) can be expressed as:

\[\mathbf{h}^{k+1}=\sigma(\mathbf{h}^{k}\mathbf{W}_{k}+\mathbf{b}_{k}), \tag{3}\]

where \(\mathbf{W}_{k}\in\mathbb{R}^{p\times q}\) and \(\mathbf{b}_{k}\in\mathbb{R}^{1\times q}\) represent the weight and bias term, respectively. The non-linear activation function \(\sigma(\cdot)\) is utilized in the neural network. Theoretically, a single-layer neural network with an infinite number of neurons can approximate any continuous function to arbitrary accuracy, given a sufficiently large dataset (Hornik et al., 1989). However, due to limitations in network width, dataset size, and the challenge of tuning parameters, a single-layer network is not optimal for achieving top performance, which leads to overfitting and poor generalization performance. To alleviate the limitation, multiple neural network layers are stacked together to enhance its expressibility and capture complex hierarchical features.

### Graph Neural Network

Neural networks have shown remarkable performance in various applications. In most neural network applications, input data structures are normally fixed, which is also called Euclidean data. However, non-Euclidean data structure such as graph-structured data is pervasive in different applications. The complexity and variability of the graph structure data make it difficult to model with conventional neural network architectures. To address this challenge, GNNs are specifically designed to handle graph-structured data. It operates on the node features and edge features and learns to extract embedding from nodes and edges, aiming to capture the underlying graph structure.

There are different types of graph neural network formulation. One of the popular approaches is the spectral approach (Wang and Zhang, 2022). Spectral graph convolution is a type of convolution operation on graph signals that uses the graph Fourier transform. It operates in the frequency domain and utilizes the eigenvalues and eigenvectors of the graph Laplacian to filter the node features. Given a graph \(\mathcal{G}=(\mathcal{V},\mathcal{E})\) with adjacency matrix \(\mathbf{A}\) and diagonal degree matrix \(\mathbf{D}=\text{diag}(\mathbf{A}\mathbf{\bar{1}})\), the Laplacian matrix and normalized Laplacian matrix of the graph is defined as \(\mathbf{L}=\mathbf{D}-\mathbf{A}\) and \(\mathbf{L}_{\text{norm}}=\mathbf{D}^{-\frac{1}{2}}\mathbf{L}\mathbf{D}^{-\frac{1}{2}}\), respectively. The spectral graph convolution is defined mathematically as:

\[g_{\theta}*\mathbf{x}=\mathbf{U}g_{\theta}(\mathbf{U}^{T}\mathbf{x}), \tag{4}\]

where \(g_{\theta}\) is a filter with learnable parameters \(\theta\), \(\mathbf{x}\in\mathbb{R}^{|\mathcal{V}|\times N_{F}}\) is the input features with \(|\mathcal{V}|\) nodes and \(N_{F}\) features per node, and \(\mathbf{U}\) is the eigenvectors of \(\mathbf{L}_{\text{norm}}\). The input signal is first transformed into the spectral
domain. The features are passed through the learnable filter and transformed back into the spatial domain. The graph spectral operator can be applied to graphs of varying sizes. As a different approach to modeling graph data, the graph attention network (GAT) learns the graph feature by computing attention scores for each node based on its features and the features of its neighbors (Velickovic et al., 2017). The graph attention network computes the new node representation \(\mathbf{x}_{i}^{\prime}\) for each node \(i\) as follows:

\[\mathbf{x}_{i}^{\prime}=\sigma\left(\sum_{j=1}^{N}\alpha_{ij}\mathbf{W}_{x}\mathbf{x}_{j} \right), \tag{5}\]

where \(\sigma\) is an activation function, \(\mathbf{W}_{x}\) is a learnable weight matrix, and \(\alpha_{ij}\) is the attention weight assigned to the node \(j\) related to its neighbour node \(i\). The attention weights are computed as follows:

\[\alpha_{ij}=\frac{\exp(\sigma(\mathbf{a}^{T}[\mathbf{W}_{x}\mathbf{x}_{i}\oplus\mathbf{W}_{x} \mathbf{x}_{j}]))}{\sum_{k\in\mathcal{N}(i)}\exp(\sigma(\mathbf{a}^{T}[\mathbf{W}_{x}\mathbf{ x}_{i}\oplus\mathbf{W}_{x}\mathbf{x}_{j}]))}, \tag{6}\]

where \(\mathbf{a}\) is a learnable weight vector, \(\mathcal{N}(i)\) is the set of neighboring nodes of node \(i\), \(\oplus\) denotes concatenation function. The graph attention mechanism can be stacked into multiple layers, with each layer learning increasingly complex representations of the graph. The attention mechanism allows the network to learn the different importance of different nodes within a neighborhood, which can improve model performance.

The aforementioned formulation is valid for homogeneous graphs, where all nodes and edges have the same semantic meaning. However, it is noted that real-world graphs are not always homogeneous. For instance, in the literature citation graph, nodes can represent various entities such as papers, authors, and journals, while edges may denote different semantic relationships. When the graph contains different types of nodes or edges, it is considered as a heterogeneous graph. Utilizing GNNs on heterogeneous graphs offers notable advantages over homogeneous counterparts, particularly in the ability to learn type-specific representations for each node and edge type (Wang et al., 2019; Fu et al., 2020). This allows for more accurate and targeted modeling of each entity and relationship, leading to improved performance on downstream tasks (Zhao et al., 2021). In the following sections, we will leverage the expressiveness of the heterogeneous graph neural network to estimate the traffic flow performance under different OD demand settings.

## 3 Traffic Flow Learning using Heterogeneous Graph Neural Networks

In this section, we elaborate on the proposed architecture of the heterogeneous graph neural network for traffic flow learning. The detail of the proposed model is shown in Fig. 1. It consists of the embedding preprocessing block, the heterogeneous graph attention block, and the prediction block. The details of each module are described as follows.

### Graph Construction & Feature Preprocessing

The heterogeneous graph \(\mathcal{G}=(\mathcal{V},\mathcal{E}_{r},\mathcal{E}_{v})\) for traffic flow learning consists of one type of node representing the intersections of road segments and two edge types: real links and virtual links. The real links represent the road segments in the road network, while the virtual links represent the auxiliary link between the origin and destination nodes. The auxiliary links can be considered as the edge augmentation technique to enhance the feature updating. The node features \(\mathbf{X}_{n}\in\mathbb{R}^{|\mathcal{V}|\times|\mathcal{V}|+2}\) are the origin-destination demand of each node and the node geocordinates. The real link features \(\mathbf{X}_{e,r}\in\mathbb{R}^{|\mathcal{E}_{r}|\times 2}\) include the free-flow travel time and the link capacity, while the virtual link features \(\mathbf{X}_{e,v}\in\mathbb{R}^{|\mathcal{E}_{v}|\times 1}\) are binary values indicating the existence of OD demand between origin and destination nodes.

Furthermore, the original node features are often sparse and non-normalized. To address this issue, we employ a preprocessing step to encode the raw features into a lower-dimensional representation that captures the essential characteristics of the data while preserving the semantic information. The generated node feature embedding size is \(\mathbf{X}_{n}^{0}\in\mathbb{R}^{|\mathcal{V}|\times N_{v}}\), where \(N_{v}\) is the embedding size. Similarly, the edge features are also normalized before being propagated in the message passing. The normalized edge feature of real edges and virtual edges are denoted as \(\mathbf{X}_{e,r}^{0}\in\mathbb{R}^{|\mathcal{E}_{r}|\times 2}\) and \(\mathbf{X}_{e,v}^{0}\in\mathbb{R}^{|\mathcal{E}_{r}|\times 1}\), respectively. Additionally, it should be noted that there is an overlap between the real links and virtual links in the heterogeneous graph because it is possible that a road and an OD demand both exist between the same node pair.
### Graph Spatial Features Extraction

As discussed in Section 2.3, effectively aggregating information from nodes and edges of different types is a major challenge for heterogeneous graphs. In order to address this challenge, we propose a novel approach that leverages attention mechanisms on both real links and virtual links. This allows our model to selectively focus on the most relevant information from each edge type, improving the overall performance of the model.

As the first step of spatial feature extraction, to enhance the modeling capability and capture more fine-grained relationships, we employ a key-query-value mechanism on node features in addition to the weighted sum operation. For single-head attention, the node features are transformed into key, query, and value matrices through linear transformations:

\[\mathbf{K},\mathbf{Q},\mathbf{V}=[\mathbf{W}_{Q}\mathbf{X}_{n}^{0},\mathbf{W}_{K}\mathbf{X}_{n}^{0},\mathbf{W}_ {V}\mathbf{X}_{n}^{0}], \tag{7}\]

where \(\mathbf{W}_{K}\in\mathbb{R}^{N_{v}\times N_{h}}\), \(\mathbf{W}_{Q}\in\mathbb{R}^{N_{v}\times N_{h}}\), and \(\mathbf{W}_{V}\in\mathbb{R}^{N_{v}\times N_{h}}\) are learnable weight matrices. \(N_{h}\) is the output size of the linear transformation. The value matrix captures the learned representations and importance scores for each node. The query matrix represents the target node's characteristics, while the key matrix contains the features of neighboring nodes. The attention scores between the query vector and each key are calculated with the dot product, normalized by a softmax function to ensure that the weights sum up to one:

\[s_{ji}=\frac{\exp(\mathbf{q}_{j}^{T}\mathbf{k}_{i})}{\sum_{k\in\mathcal{N}(i)}\exp( \mathbf{q}_{k}^{T}\mathbf{k}_{i})}, \tag{8}\]

where \(\mathbf{k}_{i}\) and \(\mathbf{q}_{j}\) denote the \(i^{\text{th}}\) and \(j^{\text{th}}\) column of matrix \(\mathbf{K}\) and \(\mathbf{Q}\), respectively. The attention scores are then used in the weighted sum of the corresponding values, allowing the model to focus on more informative features. When the edge consists of multiple features, the output of the attention layer \(\mathbf{R}\in\mathbb{R}^{|\mathcal{V}|\times(N_{h}\times h)}\) is calculated using the attention score from Equation 8, value vector, and edge features from the neighbor of node \(i\):

\[\mathbf{R}=\bigoplus_{i=1}^{N}\left(\sum_{l=1}^{L}\sum_{k\in\mathcal{N}(i)}s_{ik} \mathbf{v}_{k}x_{ik,l}\right), \tag{9}\]

where \(\mathbf{v}_{k}\) denote the \(k^{\text{th}}\) column of value matrix \(\mathbf{V}\), \(x_{ik,l}\) represent the \(l^{\text{th}}\) edge feature on link \((i,k)\). This allows the model to focus on the most relevant output vectors for a given input vector. For multi-head

Figure 1: The illustration of the heterogeneous graph neural network for traffic flow learning. The original node feature is transformed using multi-layer perception. Then, the node feature and edge feature are first passed through the virtual links and then passed through the real links. The flow-capacity ratio of each link is calculated using the source node feature, destination node feature, and edge feature.
attention, the process of Equation 7, 8, and 9 is repeated \(h\) times to get multiple head attention output \([\mathbf{R}^{(0)},\mathbf{R}^{(1)},,\mathbf{R}^{(h-1)}]\), where \(h\) is the number of heads. The learned embedding of original node features is then passed into a linear transformation layer. We concatenate the output tensors \(\mathbf{R}^{(0)},\mathbf{R}^{(1)},\ldots,\mathbf{R}^{(h-1)}\) along the feature dimension and multiply a learnable weight matrix to produce the final output tensor \(\mathbf{O}\):

\[\mathbf{O}=\bigoplus_{i=0}^{h-1}\left(\mathbf{W}_{o}\mathbf{R}^{(i)}\right). \tag{10}\]

where \(\mathbf{W}_{o}\in\mathbb{R}^{(N_{h}\times h)\times N_{o}}\), \(N_{o}\) is the output dimension of the final node embedding of a single layer GNN.

As shown in Figure 1, multiple GNN layers are stacked to extract spatial features. The first GNN layer aggregates the node features along the virtual links. While the following GNN layers aggregate the node features along the real link. Compared with the homogeneous graph constructed with only real links, the link augmentation and message passing through virtual links can be considered as a dimension-reduction technique to reduce the number of hops required for distant nodes to gather messages. As a result, it requires fewer GNN layers for effective feature aggregation and following edge prediction in the heterogeneous graph.

### Graph Edge Prediction

To predict the transportation performance at the edge level, the node embedding of the source node and destination node, and the real edge feature are concatenated and passed into the multi-layer perceptron to have the flow prediction \(\tilde{f}_{e}\) of real link \(e\):

\[\tilde{f}_{e}=\texttt{MLP}([\mathbf{O}_{e,src}\oplus\mathbf{O}_{e,dst}\oplus\mathbf{X}_{ e,r}];\mathbf{W}_{e},\mathbf{b}_{e}), \tag{11}\]

where \(\mathbf{O}_{e,src}\) and \(\mathbf{O}_{e,dst}\) represent the source node embedding and destination node embedding, respectively. \(\mathbf{W}_{e}\) and \(\mathbf{b}_{e}\) is the learnable parameters associated with the multilayer perception. After generating the final prediction, selecting an appropriate loss function becomes crucial to ensure the model's effective convergence. The proposed model utilizes a twofold loss function. The first part is the supervised loss, which measures the discrepancy between prediction and ground truth:

\[L_{s}=\frac{1}{|\mathcal{E}_{r}|}\sum_{e\in\mathcal{E}_{r}}\|f_{e}-\tilde{f}_ {e}\|, \tag{12}\]

where the \(f_{e}\) and \(\tilde{f}_{e}\) represent the ground truth and prediction of flow on link \(e\in\mathcal{E}\). The second part of the loss function is from the node-based flow conservation law, where the total flow of traffic entering a node equals the total flow of traffic exiting that node. The node-based flow conservation law can be represented mathematically:

\[\sum_{k}f_{ki}-\sum_{j}f_{ij}=\Delta f_{i}=\begin{cases}&\sum_{v\in\mathcal{V }}O_{v,i}-\sum_{v\in\mathcal{V}}O_{i,v},\quad\text{ if }i\in\mathcal{V}_{OD},\\ &0\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \text{ otherwise },\end{cases} \tag{13}\]

where \(f_{ki}\) denotes the flow on the link \((k,i)\), \(\Delta f_{i}\) represents the difference between flow receiving and sending at node \(i\), \(O_{v,i}\) represents the number of OD demand from \(v\) to \(i\). \(\mathcal{V}\) and \(\mathcal{V}_{OD}\) denote the entire node set and the origin-destination node set, respectively. The node-based flow conservation law can be considered as a normalization loss or auxiliary loss to ensure the predicted traffic flow at each node satisfies the flow conservation principle. One common way to incorporate conservation law into the loss function is to define a residual loss function:

\[L_{f}=\sum_{i}\ |\sum_{k\in\mathcal{N}_{i}(i)}\tilde{f}_{ki}-\sum_{j\in \mathcal{N}_{o}(i)}\tilde{f}_{ij}-\Delta f_{i}|, \tag{14}\]

where \(\mathcal{N}_{i}(i)\) and \(\mathcal{N}_{o}(i)\) represent the incoming edges and outcoming edges of node \(i\), respectively. The normalization loss \(L_{f}\) measures how the flow prediction satisfies the flow conservation law. Minimizing this loss function during training will encourage the model to learn traffic flow patterns that satisfy the conservation law. The total loss for the flow prediction \(L_{total}\) is the weighted summation of the supervised loss and the conservation loss:
\[L_{total}=w_{s}L_{s}+w_{f}L_{f}, \tag{15}\]

where the \(w_{s}\) and \(w_{f}\) represent the normalized weight for supervised loss and the conservation loss, respectively.

## 4 Numerical Experiments

To evaluate the accuracy, efficiency, and generalization capability of the proposed graph neural network, two numerical experiments are conducted. The first experiment is on urban transportation networks with synthetic traffic data. The second experiment is on multiple synthetic graphs with different topologies. Furthermore, we also synthesized variations in the OD demand and link capacity, as will be explained later.

### Experiments on Urban Transportation Networks

As case studies, three urban transportation networks are selected: Sioux Falls network, East Massachusetts Network (EMA), and Anaheim network. The information about the network topology, link capacity, and the default OD demand of these networks are obtained from (Bar-Gera et al., 2023). The details of the network topologies are shown in Table 1. To create demand variation, we scaled the demand by a scaling factor according to

\[\tilde{O}_{s,t}=\delta^{o}_{s,t}\ O_{s,t}, \tag{16}\]

where \(O_{s,t}\) is the default OD demand between source \(s\) and destination \(t\) and \(\delta^{o}_{s,t}\sim U(0.5,1.5)\) is the uniformly distributed random scaling factor for the OD pair (\(s\), \(t\)). Additionally, to account for variations in network properties, variable link capacities are created according to

\[\tilde{c}_{a}=\delta^{c}_{a}\ c_{a}, \tag{17}\]

where \(c_{a}\) is the original link capacity for link \(a\), and \(\delta^{c}_{a}\) is the scaling factor for link \(a\). Capacity variations are considered to be due to traffic accidents, road construction/damage, and adverse weather conditions, which reduce the link capacity. In this work, three levels of capacity reduction are considered: (L): light disruption with \(\delta^{c}_{a}\sim U(0.8,1.0)\); (M) moderate disruption with \(\delta^{c}_{a}\sim U(0.5,1.0)\); (H) high disruption with \(\delta^{c}_{a}\sim U(0.2,1.0)\). The size of the dataset for each network at each disruption scenario is 5000, which is split into the training set and the testing set with a ratio of 80% and 20%, respectively.

\begin{table}
\begin{tabular}{c c c c c} \hline \hline Network Name & \(|\mathcal{V}|\) & \(|\mathcal{E}|\) & Average Degree & OD Demand \\ \hline Sioux Falls & 24 & 76 & 3.17 & 188,960 \\ EMA & 74 & 258 & 3.49 & 132,106 \\ Anaheim & 416 & 914 & 3.05 & 226,279 \\ \hline \hline \end{tabular}
\end{table}
Table 1: The detail of urban transportation network. Three networks, including Sioux Falls, East Massachusetts, and Anaheim, are considered.

Figure 2: The illustrations of urban transportation networks, including Sioux Falls, EMA, and Anaheim. The link color represents the link capacity of each link.
The ground truth data used in training and testing steps is obtained from the user equilibrium (UE)-based traffic assignment. Specifically, the traffic flow and flow-capacity ratio (v/c ratio) of each link are the quantities of interest. The GNN model is implemented using PyTorch (Paszke et al., 2019), and DGL (Wang et al., 2019). The preprocessing layer consists of a three-layer fully connected neural network with an embedding size of 64. The number of GNN layers in the proposed model is 3. The first layer is applied on virtual links, and the following layers are applied on real links. The number of heads in the attention block is 8. For hyper-parameter selection, the hidden layer size is chosen as 64, which is common in neural network implementation (Liu and Meidani, 2023). Rectified Linear Unit (ReLU) activation (Agarap, 2018) is chosen to introduce nonlinearity. Mini-batch stochastic gradient descent was implemented in the training process. More specifically, the adaptive moment estimation (Adam) optimizer (Kingma and Ba, 2014) is adopted with a learning rate of 0.001. The batch size of training is 128. The weights for supervised loss and the conservation loss in equation 15 are chosen as 1.0 and 0.001, respectively, in order to keep both loss terms in the same order of magnitude. We evaluated the performance of our proposed heterogeneous GNN model (referred to by HetGAT) and compared it with three benchmark models: a fully connected neural network (FCNN), a homogeneous graph attention network (GAT), and a homogeneous graph convolution network (GCN). The FCNN consisted of five fully connected layers with an embedding size of 64. The GAT and GCN both have three layers of graph message passing layer, followed by three layers of FCNN with an embedding size of 64. The metrics to evaluate performance include the mean absolute error (MAE) and the normalized conversation loss \(L_{f}\) given by

\[\text{MAE}=\frac{1}{N}\sum_{i=1}^{N}|y_{i}-\tilde{y}_{i}|, \tag{18}\]

\[L_{f}=\frac{\sum_{i}|\sum_{k\in\mathcal{N}_{i}(i)}\tilde{f}_{ki}-\sum_{j\in \mathcal{N}_{o}(i)}\tilde{f}_{ij}-\Delta f_{i}|}{\sum_{s}\sum_{t}\tilde{O}_{s,t}}, \tag{19}\]

where \(y\) and \(\tilde{y}\) respectively represent the ground truth and predicted values for the flow-capacity ratio. Furthermore, we use the coefficient of determination, denoted by \(R^{2}\) to measure the goodness-of-fit between the GNN prediction and the ground truth. The training histories of the studied models are shown in Figure 3. The results indicate that the FCNN model performed poorly during training compared to GNN-based models, as shown by the high training loss and early stagnation. In contrast, GCN and GAT models exhibited similar convergence rates. Our proposed model outperformed both GCN and GAT in terms of training loss. Especially when used for larger networks, the proposed model demonstrated superior convergence performance compared to GCN and GAT; for the Anaheim network, the training loss of the proposed model is almost 1/3 of that of GAT in the first 25 iterations. This is because GCN and GAT only consider homogeneous edges, which limits the message passing to adjacent nodes. In contrast, the proposed GNN model uses virtual links and provides augmented connectivity to long-hop node pairs, which makes the node feature updating in HetGAT more efficiently.

After the training is finished, the model performance is evaluated on the testing set. The experiments under the urban road network are conducted in three different settings. The first setting, referred to as

Figure 3: Training loss history under urban transportation network. Three benchmarks, including FCNN, GCN, and GAT, are compared with HetGAT.
LMH-LMH, involved using all levels of disruption (flow reduction scaling levels) in both the training and testing sets. The second setting, namely L-M, involves training the model on light disruption data and testing it on medium disruption data. The third case, which is labeled as M-H, involves training on medium disruption data and testing on high disruption data. The L-M and M-H scenarios will therefore involve unseen cases that don't exist in the training data. Figure 4 plotted the predicted value and ground truth of the flow-capacity ratio on multiple samples in the Anaheim network under the LMH-LMH setting. In total, the flow-capacity ratio on 10,000 edges is predicted using HetGAT, GAT, and GCN, respectively. Figure 4 indicates that HetGAT has a relatively higher correlation coefficient, which outperforms GAT and GCN. Table 2 summarizes the prediction performance of all methods under different settings, and shows that HetGAT, compared to other models, offers better performance. When the graph size increases, the proposed model can maintain a relatively low MAE compared to GCN and GAT. For instance, in the Anaheim network, HetGAT offers MAE values that are 39%, 19%, and 31% lower than GAT values, in LMH-LMH, L-M, and M-H settings, respectively. Also, the conservation errors of HetGAT are 52%, 29%, and 55% smaller in LMH-LMH, L-M, and M-H settings, respectively. This shows that the inclusion of virtual links can assist GNN models in better learning the traffic flow patterns.

As an additional experiment, we consider a realistic scenario where the demand values are not known for the entire OD pairs. To evaluate the accuracy and robustness of the proposed model under incomplete OD demand scenarios, we introduce a random mask to the original OD demand, to simulate these missing

\begin{table}
\begin{tabular}{c c c c c|c c c c|c c} \hline \hline \multirow{2}{*}{Network} & \multirow{2}{*}{Model} & \multicolumn{3}{c}{LMH-LMH} & \multicolumn{3}{c}{L-M} & \multicolumn{3}{c}{M-H} \\ \cline{3-11}  & & MAE & \(L_{f}\) & \(R^{2}\) & MAE & \(L_{f}\) & \(R^{2}\) & MAE & \(L_{f}\) & \(R^{2}\) \\ \hline \multirow{4}{*}{Sioux Falls} & FCNN & 0.158 & 0.324 & 0.728 & 0.139 & 0.062 & 0.812 & 0.215 & 0.238 & 0.643 \\  & GAT & 0.046 & 0.092 & 0.956 & **0.022** & 0.018 & **0.991** & 0.085 & 0.057 & 0.889 \\  & GCN & 0.046 & 0.096 & 0.955 & 0.023 & 0.019 & 0.990 & 0.086 & 0.060 & 0.889 \\  & HetGAT & **0.028** & **0.054** & **0.985** & 0.026 & **0.018** & 0.986 & **0.062** & **0.039** & **0.937** \\ \hline \multirow{4}{*}{EMA} & FCNN & 0.238 & 0.422 & 0.621 & 0.225 & 0.144 & 0.795 & 0.263 & 0.446 & 0.771 \\  & GAT & 0.063 & 0.131 & 0.952 & 0.046 & 0.033 & 0.991 & 0.132 & 0.090 & 0.842 \\  & GCN & 0.071 & 0.167 & 0.944 & 0.034 & 0.043 & 0.986 & 0.161 & 0.135 & 0.814 \\  & HetGAT & **0.038** & **0.090** & **0.982** & **0.034** & **0.027** & **0.989** & **0.090** & **0.055** & **0.921** \\ \hline \multirow{4}{*}{ Anaheim} & FCNN & 0.111 & 0.721 & 0.822 & 0.199 & 0.237 & 0.801 & 0.186 & 0.386 & 0.840 \\  & GAT & 0.051 & 0.263 & 0.951 & 0.042 & 0.081 & 0.958 & 0.090 & 0.184 & 0.881 \\ \cline{1-1}  & GCN & 0.072 & 0.402 & 0.903 & 0.062 & 0.109 & 0.911 & 0.116 & 0.245 & 0.827 \\ \cline{1-1}  & HetGAT & **0.031** & **0.126** & **0.981** & **0.034** & **0.059** & **0.979** & **0.062** & **0.082** & **0.939** \\ \hline \hline \end{tabular}
\end{table}
Table 2: Performance comparison of HetGAT with benchmark methods. Three different settings are included: LMH-LMH, L-M, M-H. The mean absolute error, normalized conservation loss, and correlation coefficient are used to evaluate the prediction performance on the testing set.

Figure 4: Comparison of predicted flow-capacity ratio between ground truth and surrogate modeling in the Anaheim network under LMH-LMH setting.
demand scenarios. Specifically, given a specific missing ratio, we randomly select a number of OD pairs and mask their corresponding OD demand values as zeros in the input node feature. This way, the model is expected to learn the inherent patterns and structures of the traffic network, even when some of the demand information is missing. The effectiveness of the proposed model under incomplete OD demand scenarios will be evaluated by comparing the predicted traffic flows against the ground truth data obtained from the complete OD demand. Three missing ratios are considered in the experiment: 20%, 30%, and 40%. The training setting and the hyperparameters remain the same as those in the aforementioned experiments. Table 3 summarizes the results of prediction performance under different missing rate scenarios for the LMH-LMH setting. The FCNN model is not considered in these experiments because of their very poor performance in the previous experiments under full OD demand. As expected, overall the performances under incomplete OD demand drop, compared to that in the complete demand scenario, reported in Table 2. However, HetGAT still outperforms GAT and GCN under different networks and different missing ratios. In particular, the MAE values are about 50% smaller to those offered by GAT and GCN under various missing ratios. Also, the flow predictions by HetGAT have better compliance with the flow conservation law compared with GAT and GCN.

### Experiments on Generalized Synthetic Networks

In this set of experiments, we study the performance of the proposed model when generalized to networks with different topologies. To do so, we generate synthetic networks by starting with a grid graph and adding links between randomly selected nodes in the grid. Then, a few nodes and edges are randomly removed to emulate the real road network and increase the variability. It should be noted that the proposed HetGAT involves a node feature whose dimension is equal to the total number of nodes in the graph. Therefore, our model can be generalized to different graphs as long as the total number of nodes remain the same. Therefore, the various topologies in synthetic road networks all have the same number of nodes. Two sizes of the synthetic generalized dataset are included in the experiments: 100 and 300. For each graph size, 100 different graph topologies are generated. Three of these randomly generated graphs are shown in Figure 5. The OD demand and the link capacity are also randomly generated using the scaling factor according to Equation 16 and 17, respectively. Furthermore, the UE-based traffic assignment is used to generate the ground truth data, similarly to Section 4.1.

It should be noted that in fully connected neural networks (FCNN), the output dimension and the order of the entries should be fixed. This is while in our synthetic networks, even though the number of nodes are fixed, the number of edges are not necessarily the same in different topologies. Because of this, FCNN is not suitable for experiments on generalized synthetic networks. Therefore, GAT and GCN are the only two models that are compared with the proposed HetGAT model.

In the training, the same hyperparameters used in the previous experiments are used here, as well. Figure 6 presents the training history for the three models, and shows that GCN and GAT models struggle during the training process. In contrast, the HetGAT model can successfully learn the flow dynamics across multiple graphs, leading to a decrease in training loss over time. For a graph size of 100, GCN and GAT had a training

\begin{table}
\begin{tabular}{c c c c c|c c c c|c c c} \hline \hline \multirow{3}{*}{Network} & \multirow{3}{*}{Model} & \multicolumn{6}{c}{Missing Ratio} \\ \cline{3-11}  & & \multicolumn{3}{c}{20\%} & \multicolumn{3}{c}{30\%} & \multicolumn{3}{c}{40\%} \\ \cline{3-11}  & & MAE & \(L_{f}\) & \(R^{2}\) & MAE & \(L_{f}\) & \(R^{2}\) & MAE & \(L_{f}\) & \(R^{2}\) \\ \hline \multirow{4}{*}{Sioux Falls} & GAT & 0.048 & 0.098 & 0.954 & 0.049 & 0.102 & 0.953 & 0.049 & 0.099 & 0.952 \\  & GCN & 0.049 & 0.102 & 0.951 & 0.049 & 0.101 & 0.953 & 0.049 & 0.101 & 0.952 \\  & HetGAT & **0.033** & **0.068** & **0.981** & **0.034** & **0.069** & **0.979** & **0.033** & **0.068** & **0.982** \\ \hline \multirow{4}{*}{EMA} & GAT & 0.068 & 0.145 & 0.938 & 0.070 & 0.149 & 0.931 & 0.070 & 0.149 & 0.933 \\  & GCN & 0.079 & 0.200 & 0.936 & 0.079 & 0.204 & 0.935 & 0.077 & 0.188 & 0.939 \\  & HetGAT & **0.046** & **0.111** & **0.970** & **0.048** & **0.115** & **0.970** & **0.048** & **0.105** & **0.972** \\ \hline \multirow{4}{*}{Anaheim} & GAT & 0.067 & 0.152 & 0.924 & 0.061 & 0.135 & 0.933 & 0.071 & 0.156 & 0.914 \\  & GCN & 0.105 & 0.225 & 0.835 & 0.109 & 0.237 & 0.802 & 0.106 & 0.239 & 0.827 \\ \cline{1-1}  & HetGAT & **0.035** & **0.060** & **0.975** & **0.034** & **0.069** & **0.969** & **0.034** & **0.058** & **0.977** \\ \hline \hline \end{tabular}
\end{table}
Table 3: Comparison of the performance of HetGAT with that of GAT and GCN under incomplete OD demand. Three missing ratios are considered in the experiments: 20%, 30%, and 40%. The mean absolute error (MAE), normalized conservation loss (\(L_{f}\)), and the coefficient of determination, \(R^{2}\), are used to evaluate the prediction performance on the testing set.
loss of 0.114 and 0.124, respectively, whereas HetGAT achieved a training loss of 0.032, representing a 70% improvement. The prediction performance metrics on testing sets are compared in Table 4, where the mean absolute error of HetGAT outperformed GAT and GCN for graph sizes of 100 and 300, with values of 0.033 and 0.051, respectively.

## 5 Conclusion and Discussion

In this paper, we proposed a novel approach for traffic assignment and traffic flow learning using heterogeneous graph neural networks. We conducted extensive experiments on three real-world traffic networks to evaluate the performance of the proposed heterogeneous GNN model and compared it with the state-of-the-art models. The results show that the proposed model outperforms other models in terms of

\begin{table}
\begin{tabular}{c c c c c} \hline \hline \multirow{2}{*}{Graph Size} & \multirow{2}{*}{Metric} & \multicolumn{3}{c}{Model} \\ \cline{3-5}  & & GAT & GCN & HetGAT \\ \hline \multirow{3}{*}{100} & MAE & 0.143 & 0.179 & **0.033** \\  & \(L_{f}\) & 0.016 & 0.014 & **0.013** \\  & \(R^{2}\) & 0.602 & 0.372 & **0.979** \\ \hline \multirow{3}{*}{300} & MAE & 0.291 & 0.316 & **0.051** \\  & \(L_{f}\) & 0.018 & 0.016 & **0.013** \\ \cline{1-1}  & \(R^{2}\) & 0.412 & 0.316 & **0.977** \\ \hline \hline \end{tabular}
\end{table}
Table 4: Comparison of the performance of HetGAT with that of GAT and GCN on generalized synthetic networks. Two graph sizes of 100 and 300 are considered in the experiments. The mean absolute error (MAE), normalized conservation loss (\(L_{f}\)), and coefficient of determination (\(R^{2}\)) are used to evaluate the prediction performance on the testing set.

Figure 5: The illustrations of sampled generalized synthetic networks. The network size is 100. The link color represents the link capacity of each link.

Figure 6: Training loss history under generalized synthetic networks. Two network sizes of 100 and 300 are selected for training and testing. Two benchmarks including GCN and GAT are compared with HetGAT.
convergence rate, training loss, and prediction accuracy. The proposed model can also be generalized to different network topologies, with different topologies or network parameters, which demonstrates its potential in real-world applications. Currently, the proposed HetGAT model only learns and predicts the static traffic flow patterns. As a potential extension of this work, the proposed framework can be extended to learn the dynamic traffic flow patterns. Furthermore, the current proposed GNN model uses training data collected from conventional solvers of static traffic assignment. In future work, we will study how these models can be trained on traffic data collected from sensors, such as loop detectors, cameras, or GPS devices.

## 6 Acknowledgment

This work was supported in part by the National Science Foundation under Grant CMMI-1752302.
