TANGNN: a Concise, Scalable and Effective Graph Neural Networks with Top-m Attention Mechanism for Graph Representation Learning

Algorithm 2 demonstrates the forward propagation process of neighbor aggregation components. In this process, at the $i$th layer, the model firstly gathers a fixed set of neighbors, $\mathcal{N}(v)$, for the node $v$. Then it aggregates the neighbor information $h_{\mathcal{N}(v)}^{i}\leftarrow AGGEREGATE(h_{u}^{i-1},u\in\mathcal{N}(v))$ (See line 2 of the code), here, the aggregation function used by the model is the Mean aggregation function. Finally, it updates the vector representation of node $v$ $h_{v}^{i}\leftarrow\sigma(W_{i}\cdot CONCAT(h_{v}^{i-1},h_{\mathcal{N}(v)}^{i}))$ (line 3 of the code), where Sigmod function $\sigma$ is a nonlinear function: $\frac{1}{1+e^{(-x)}}$, it is capable of extracting nonlinear relationships.

Algorithm 3 demonstrates the forward propagation process of the Top-$m$ attention mechanism aggregation component. During this process, apply the Top-$m$ strategy ( (Details are shown in section 4.3) ) to the output of the previous layer of each node $v$, which is $g_{v}^{i-1}{{}^{\prime}}$, to select the $m$ nodes with the highest similarity and obtain the representation $H_{\text{Top-}m}^{i}$ (see line 5 of the code). As shown in Eq. (1) to (3), the self-attention mechanism first maps it to the query ($q$), key ($k$), and value ($\nu$) vectors:

Here, $W_{q_{m}}^{i},W_{k_{m}}^{i},W_{\nu_{m}}^{i}$ are the learnable weight matrices for the $i$th layer. The query, key, and value matrices corresponding to each head typically have smaller dimensions, denoted as $d_{q},d_{k},d_{\nu}$ espectively. Next, compute the dot product of the query vector with the key vectors, scaled by a factor of $\sqrt{d_{k}}$, and then apply softmax to obtain the attention weights, as shown in Eq. (4):

Here, $A_{vu}^{i}$ represents the attention weight from node $v$ to node $u$ at the $i$th layer, where $m$ is the number of nodes obtained from the Top-$m$ attention mechanism component. Then we use Eq. (5) to calculate the output of the head:

The output vectors linearly transformed by another weight matrix $W^{o}$, as shown in Eq. (6).

Subsequently, after layer normalization and a feed-forward neural network, the output is merged $h_{v}^{i}{{}^{\prime\prime}}$ $\leftarrow$ $LayerNorm$ $(g_{v}^{i-1},\tilde{h}_{v}^{i}{{}^{\prime}})$ (line 5 of the code), and then it is fed into a feed-forward neural network (line 6 of the code). Finally, the output of the current layer is completed through a second layer normalization (line 7 of the code), obtaining the vector representation $h_{v}^{i}{{}^{\prime}}$ of node $v$ at the $i$th layer.

Algorithm 4 shows the backpropagation process of TANGNN. It firstly calls Algorithm 1 to compute the node vector representation $g_{v}$ (see line 4 of the code). Lines 5 to 7 involve the calculation of the loss value, starting with a softmax operation on each node’s vector representation to convert it into a probability distribution $p$, then converting the node’s label $y$ into one_hot encoding, and finally computing the loss value through the loss function. Lines 8 to 9 involve calculating gradients to implement the update of weight parameters.

$W$ is the training parameter for the neighbor aggregation component of Algorithm 1, and $W_{q_{m}},W_{k_{m}},W_{\nu_{m}}$ are the weight parameters for the Top-$m$ attention mechanism aggregation component, all of which are obtained through training with the following objective function:

Where the probability distribution $p$ is the desired output, and the probability distribution $q$ is the actual output. This objective function primarily measures the distance between the actual output (probability) and the desired output (probability). A smaller value indicates that the two probability distributions are closer to each other.

As shown in Figure 4, unlike TANGNN, TANGNN-FLC concatenates the node information processed by the two components only in the final layer. The neighborhood aggregation component produces the feature vector representation $h_{v}^{L}$ of the target node at the $L$th layer, and the Top-$m$ attention mechanism aggregation component then computes based on the output $h_{v}^{L-1}{{}^{\prime}}$ from the previous layer’s , produces the feature vector representation $h_{v}^{L}{{}^{\prime}}$ of the target node at the $L$th layer. These two feature vectors are concatenated and then processed through a multilayer perceptron to form the final vector representation $g_{v}$ of the target node.

TANGNN-FLC cannot effectively retain the information of each layer, and as the number of layers increases, the problem of oversmoothing may occur. To address this, we designed TANGNN-NAI based on the TANGNN framework. $h_{v}^{i}$ is the vector representation of the neighborhood aggregation component at the $i$th layer, and $h_{v}^{i}{{}^{\prime}}$ is the vector representation of the Top-$m$ attention mechanism aggregation component at the $i$th layer. These two vector representations are concatenated to form the vector representation [$h_{v}^{i}$, $h_{v}^{i}{{}^{\prime}}$], which serves as the input for the next layer of the neighborhood aggregation component, while $h_{v}^{i}{{}^{\prime}}$ serves as the input for the next layer of the Top-$m$ attention mechanism aggregation component, continuing until the final layer and then passing through a multilayer perceptron to obtain the final vector representation $g_{v}^{L}$. This method effectively captures the structural information of the graph.

The only difference between TANGNN-TAI and TANGNN-NAI is that after obtaining the concatenated vector representation at the $i$th layer, different vector representations are input into different components. The concatenated vector representation serves as the input for the next layer of the Top-$m$ attention mechanism aggregation component, while the vector representation $h_{v}^{i}$ from the $i$th layer of the neighborhood aggregation component serves as the input for the next layer of the neighborhood aggregation component.

We conducted evaluations using eight well-established datasets frequently employed for vertex embeddings. The details of these datasets are shown in Table 1 and 2, including a variety of types such as social networks, protein datasets and citation networks. Both $Cora$ (Sen et al., 2008) and $Pubmed$ (Sen et al., 2008) are recognized as a benchmark in citation network analysis, where each vertex is represented by a bag-of-words from the paper and labeled by its academic subject. $Amazon$ (Sen et al., 2008) models the interactions between users and products, representing users and products as nodes and their transactions as graph edges. $Citeseer$ (Sen et al., 2008) represents a citation network where each node holds specific characteristic information pertaining to the literature. $Reddit$ (Hamilton et al., 2017) is extensively utilized for research in social network analysis and natural language processing fields. It contains post data from the Reddit platform, organized by community (Subreddit). In this dataset, each node represents a post, and the edges in the graph represent direct links between posts or explicit interactions between users, its labels correspond to the community to which the post belongs. The $ZINC$ (Irwin & Shoichet, 2005) is a large chemical database containing millions of drug compounds’ 3D structures and attribute information. In graph neural network research, ZINC is commonly used to test molecular property prediction models, such as solubility, toxicity, or biological activity of molecules. The $QM9$ (Ramakrishnan et al., 2014) includes nearly 134k stable small molecules’ quantum chemical properties, with each molecule composed of no more than 9 atoms (C, H, O, N, F). It provides computed chemical properties of molecules, such as molecular orbital energies, dipole moments, free energies, etc. The $ArXivNet$ is derived from unarXive (Saier et al., 2023), where papers are extracted to construct its own citation network. The nodes in this dataset represent ArXiv academic papers, directed edges represent citations between documents. We use the pre-trained model SPECTER (Cohan et al., 2020) to extract semantic representations as node features based on the titles and abstracts of papers. And we utilize the DictSentiBERT (Yu & Hua, 2023) model to annotate the sentiment polarity (positive, neutral, negative) of citations. The relevant processing code and data have been published on the aforementioned Github website.

We use the following representative algorithms in comparative experiments. For all models, the learning rate is set to 0.001 and the batch size is set to 128. The number of layers (hops) for the GNN models GCN, GraphSAGE, SGCN, GAT, JK-Net, GIN, TransGNN, SAGEFormer, and NAGphormer is uniformly set to 2.

$GCN$ (Kipf & Welling, 2017) is one of the initial graph convolutional network models proposed, utilizes adjacency and node feature matrices to update and propagate node features. At each layer, the feature representation of a node is updated by aggregating features from its neighboring nodes.

$SGCN$ (Wu et al., 2019) represents a streamlined version of the graph convolutional neural network. It simplifies the graph convolution process to a single linear operation by eliminating nonlinear transformations and polynomial fitting. This reduction in complexity and computational demand enhances the model’s interpretability and efficiency.

$GraphSAGE$ (Hamilton et al., 2017) is a graph convolutional neural network that uses an aggregation function to collect neighbor information around vertices. Through training, the algorithm updates this information, allowing vertices to access higher-order neighbor information as the number of iterations increases.

$GAT$ (Veličković et al., 2018) is a neural network model based on graph attention mechanisms. Its core concept revolves around learning the relationships between nodes and the importance of features between them, utilizing attention mechanisms to dynamically calculate the influence weights of each node on its neighboring nodes.

$JK$-$Net$ (Xu et al., 2018) is a framework designed to enhance the performance of graph neural networks (GNNs) on tasks like node classification and graph classification. The fundamental concept of it involves leveraging jumping connections to amalgamate information from various levels of graph convolutional layers. This integration helps the model to more effectively learn and represent graph data.

$GIN$ (Xu et al., 2019): In the task of graph isomorphism detection, the objective is to ascertain if two graphs are isomorphic, meaning there exists a one-to-one correspondence between their nodes that preserves the connectivity of edges. GIN learns graph representations by integrating the features of nodes with those of their neighbors and applying a multilayer perceptron to introduce non-linear transformations to these combined features.

$Graphormer$ (Ying et al., 2021) utilizes the standard Transformer architecture, introducing key structural encodings such as centrality encoding, spatial encoding, and edge encoding to optimize the performance of graph representation learning. This enhances its ability to process graph structure data and improves its expressiveness.

$TransGNN$ (Zhang et al., 2023) employs an attention sampling strategy based on semantic and structural information, selecting nodes most relevant to the central node for information aggregation. It alternates between Transformer layers and GNN layers within the model, using Transformer layers to expand the receptive field of GNNs, and GNN layers to aid the Transformer in better understanding the graph structure.

$NAGphormer$ (Chen et al., 2023) introduces a neighborhood aggregation module that can aggregate features from multiple hops and converts each hop’s features into a token. Each token represents the central node’s neighborhood information at different hop distances, enabling the Transformer to be effectively applied to node classification tasks in large-scale data.

$SAT$ (Chen et al., 2022) enhances the effectiveness of graph node representations by incorporating structural information into the self-attention mechanism of the Transformer. Compared to traditional methods that focus only on node features, SAT uses extracted local subgraph structure information for attention computation, enhancing the structural awareness of node representations.

$DeepGraph$ (Zhao et al., 2023) incorporates additional substructure tokens into the model and uses local attention to deal with nodes related to specific substructures, addressing the performance bottlenecks faced by graph transformers when increasing the number of layers.

$SGFormer$ (Wu et al., 2024) uses a simple single layer attention mechanism, which, by eliminating the need for positional encoding and complex graph preprocessing steps, simplifies the model architecture. This makes information propagation more efficient and reduces computational overhead.

The objective of vertex classification tasks is to accurately predict the labels for each node in the network.

Table 3 to 5 shows our method and several other algorithms performing classification tasks on three real datasets. The experimental results show that TANGNN-LC achieved good results on all datasets. On the Cora dataset, TANGNN-LC improved performance by 12% compared to GraphSAGE (when the training dataset is 50%). On the PubMed dataset, TANGNN-LC improved by 2% compared to TANGNN (when the training dataset is 10%). Overall, compared to other algorithms, TANGNN-LC and TANGNN both have achieved better performance in classification tasks.

Figure 5 illustrates how varying the parameter $K$ affects the accuracy of different algorithms when they are tasked with classification on the Cora and PubMed datasets. The findings indicate that TANGNN maintains stability without showing a significant decline in performance with increasing $K$ values. We can see that TANGNN-LC have achieved the highest accuracy on the pubmed dataset when the value of $K$ is 5. As the depth increases, TransGNN, SGFormer, and NAGphormer show a downward trend due to the influence of GNN entities. JK-Net combines information from various levels of graph convolutional layers using skip connections, thereby retaining more node information. Therefore, the effect of JK-Net on the performance is not significant when the value of $K$ increases. GCN, SGCN and other GCN variants show a significant downward trend in Acc values after the value of $K$ reaches 4.

The objective of the link prediction task is to forecast potential connections or edges within a network using existing structural information about the network.

Figure 6 shows the results of the link prediction task on the Amazon, ArXivNet, Cora, PubMed datasets. In the Amazon dataset, As the size of the training set increased, TANGNN demonstrated improved performance. In the ArXivNet dataset, TANGNN-TAI outperformed Graphormer by 6% (When the training set constitutes 70% of the total dataset), outperformed TransGNN by about 4% (When the training set constitutes 70% of the total dataset), and TANGNN outperformed GraphSAGE by about 9% (When the training set constitutes 30% of the total dataset), while TANGNN-LC outperformed GraphSAGE by about 12% (When the training set constitutes 30% of the total dataset). In both the Cora and PubMed datasets, TANGNN-LC achieved better performance.

Figure 7 shows our analysis of the runtime for the link prediction task on the ArXivNet and Reddit datasets. In this experiment, the runtime is calculated based on the model reaching a predefined convergence condition rather than a fixed number of iterations, i.e., the training loss of the model reached a stable state and no longer changed significantly. From the figure, it is evident that Graphormer and SAT encountered memory overflow issues on the Reddit dataset, and on the ArXiNet dataset, Graphormer had a longer runtime compared to other algorithms. The experimental results indicate that TANGNN takes significantly less time on large graph datasets compared to TransGNN and NAGphormer, and TANGNN-LC can reach convergence faster than TANGNN, demonstrating its efficiency advantage in processing large-scale graph data.

The purpose of the citation sentiment prediction task in this article is to classify the sentiment polarity of edges in the citation network based on the citation relationships between literature. From the table 6, it can be seen that the performance of TANGNN-LC is superior to all other algorithms. Compared to GraphSAGE, TANGNN has improved performance by 8% (when the training dataset is 50%), while TANGNN-LC has improved performance by about 8% (when the training dataset is 70%).

The purpose of graph regression task is to predict certain numerical attributes in the graph based on known network structure information, such as importance indicators of nodes, physical and chemical properties of molecules, and other continuous variables. The evaluation metric used in the graph regression task is MAE (Mean Absolute Error), and the smaller the value of MAE, the closer the model’s prediction is to the true value.

As shown in Tables 7 and 8, it can be seen that TANGNN-FLC, TANGNN-NAI, and TANGNN perform very well on both datasets. This indicates that TANGNN and its variants have strong generalization ability and stability in processing graph regression tasks, and can more accurately predict the importance indicators of nodes and the physical and chemical properties of molecules. In contrast, other models show significant fluctuations in performance under different dataset sizes, with lower stability and accuracy compared to the TANGNN model.

We investigated the training process of GraphSAGE and TANGNN on the Cora dataset. As shown in Figure 8, in terms of training loss, TANGNN’s loss decreases faster and the final loss is lower, indicating that TANGNN has better fitting effect on the training set and faster convergence speed. In terms of validation accuracy, TANGNN’s validation accuracy rapidly increases to a higher level compared to GraphSAGE in the initial stage and maintains a relatively stable state throughout the entire training process.

Firstly, we generate representation vectors using various models, which are then visualized in a two-dimensional space using t-SNE (van der Maaten & Hinton, 2008). In this visualization, vertices from different classes are distinguished by unique colors. Ideally, vertices belonging to the same class should cluster closely together.

The visualization results, displayed in Figure 9, utilize the PubMed dataset. The results clearly show effective clustering within small, distinct regions. Our algorithm demonstrates a more distinct separation of points compared to other algorithms, highlighting its effectiveness in distinguishing between classes.

In this study, we evaluated the model’s sensitivity to the parameter by varying the number of samples $m$, as shown in Figure 10. Specifically, we chose a range of sample numbers from 1 to 50, incrementing by 5 each time (for $m$ $\geq$ 5 ), to observe the impact of this variation on the model’s accuracy. Our results indicated that on the Cora dataset, the model’s accuracy began to decline when the number of samples exceeded 35; on the PubMed dataset, the performance also started to decrease when the number of samples exceeded 40. In conclusion, based on considerations of performance and computational efficiency, we recommend setting the number of samples to 30 to balance model performance and computational costs.