Two-Stage Training of Graph Neural Networks for Graph Classification

Abstract

Graph neural networks (GNNs) have received massive attention in the field of machine learning on graphs. Inspired by the success of neural networks, a line of research has been conducted to train GNNs to deal with various tasks, such as node classification, graph classification, and link prediction. In this work, our task of interest is graph classification. Several GNN models have been proposed and shown great accuracy in this task. However, the question is whether usual training methods fully realize the capacity of the GNN models. In this work, we propose a two-stage training framework based on triplet loss. In the first stage, GNN is trained to map each graph to a Euclidean-space vector so that graphs of the same class are close while those of different classes are mapped far apart. Once graphs are well-separated based on labels, a classifier is trained to distinguish between different classes. This method is generic in the sense that it is compatible with any GNN model. By adapting five GNN models to our method, we demonstrate the consistent improvement in accuracy and utilization of each GNN’s allocated capacity over the original training method of each model up to $5.4\%$ points in 12 datasets.

Appendix: Effect of Classification Layers

Table 12 Classification accuracy when the classifier consists of one layer and when the classifier is tuned using up to three layers on the benchmark datasets. 2STG $=1$ and 2STG $\le 3$ represent 2STG with one classification layer and up to three classification layers (as detailed in Sect. 4.1.4), respectively.

Table 13 Classification accuracy when the classifier consists of one layer and when the classifier is tuned using up to three layers on the New York City Taxi datasets. 2STG $=1$ and 2STG $\le 3$ represent 2STG with one classification layer and up to three classification layers (as detailed in Sect. 4.1.4), respectively.

We report in Tables 12 and 13 the classification accuracy when the classifier only consists of one layer and when the classifier is tuned using up to three layers (as described in Sect. 4.1.4). In some cases, the 1-layer classifier could achieve an accuracy that is close or slightly higher than that of the classifier using up to three layers. However, tuning the model using up to three classification layers allows us to achieve better accuracy than using only one layer for the classifier in most cases. The few exceptions are highlighted in bold italic in Tables 12 and 13. These results indicate that using a strong classifier is generally helpful in enhancing the classification accuracy. We also visualize in Fig. 7 the cases in which the differences between using one classification layer and using up to three classification layers are the highest. In particular, we highlight the t-SNE visualization of the embeddings generated by GraphSage 2STG+ for the datasets PTC-FM and JAN. G. when using one classification layer and up to three classification layers, respectively.

Fig. 7

Graph embeddings generated by GraphSAGE for the datasets PTC-FM and JAN. G. in the two cases of 2STG+ with one classification layer (left) and 2STG+ with up to three classification layers (right), respectively. These are the graphs embeddings (i.e., ${\hat{f}}(g_i)$ of each graph $g_i$), which are the input of the classification part in Fig. 2. T-SNE is applied to reduce the embedding dimension from 64 to 2 for visualization. Due to the capacity difference of the classifiers (one layer vs. up to three layers), there is a gap in the final classification accuracies of the two cases reported in Tables 12 and 13