TUNEUP: A TRAINING STRATEGY FOR IMPROVING GENERALIZATION OF GRAPH NEURAL NETWORKS Anonymous authors Paper under double-blind review

Abstract

Despite many advances in Graph Neural Networks (GNNs), their training strategies simply focus on minimizing a loss over nodes in a graph. However, such simplistic training strategies may be sub-optimal as they neglect that certain nodes are much harder to make accurate predictions on than others. Here we present TUNEUP, a curriculum learning strategy for better training GNNs. Crucially, TUNEUP trains a GNN in two stages. The first stage aims to produce a strong base GNN. Such base GNNs tend to perform well on head nodes (nodes with large degrees) but less so on tail nodes (nodes with small degrees). So, the second stage of TUNEUP specifically focuses on improving prediction on tail nodes. Concretely, TUNEUP synthesizes many additional supervised tail node data by dropping edges from head nodes and reusing the supervision on the original head nodes. TUNEUP then minimizes the loss over the synthetic tail nodes to finetune the base GNN. TUNEUP is a general training strategy that can be used with any GNN architecture and any loss, making TUNEUP applicable to a wide range of prediction tasks. Extensive evaluation of TUNEUP on five diverse GNN architectures, three types of prediction tasks, and both inductive and transductive settings shows that TUNEUP significantly improves the performance of the base GNN on tail nodes, while often even improving the performance on head nodes, which together leads up to 58.5% relative improvement in GNN predictive performance. Moreover, TUNEUP significantly outperforms its variants without the two-stage curriculum learning, existing graph data augmentation techniques, as well as other specialized methods for tail nodes.

1. INTRODUCTION

Graph Neural Networks (GNNs) are one of the most successful and widely used paradigms for representation learning on graphs, achieving state-of-the-art performance in a variety of prediction tasks, such as semi-supervised node classification (Kipf & Welling, 2017; Velickovic et al., 2018) , link prediction (Hamilton et al., 2017; Kipf & Welling, 2016) , and recommender systems (Ying et al., 2018; He et al., 2020) . There has been a surge of work on improving GNN model architectures (Velickovic et al., 2018; Xu et al., 2019; 2018; Shi et al., 2020; Klicpera et al., 2019; Wu et al., 2019; Zhao & Akoglu, 2019; Li et al., 2019; Chen et al., 2020; Li et al., 2021) and task-specific losses (Kipf & Welling, 2016; Rendle et al., 2012; Verma et al., 2021; Huang et al., 2021) . Despite all these advances, strategies for training a GNN on a given supervised loss remain largely simplistic. Existing work has focused on simply minimizing the given loss over nodes in a graph. While such a simplistic default strategy already gives a strong performance, the strategy may still be sub-optimal as it neglects that some nodes are much harder to make accurate predictions on than others. Consequently, a GNN trained with the default strategy may significantly under-perform on those hard nodes, resulting in overall sub-optimal predictive performance. Here we present TUNEUP to better train a GNN on a given supervised loss. The key motivation behind TUNEUP is that GNNs tend to under-perform on tail nodes, i.e., nodes with a small number of neighbors (Liu et al., 2021) . In practice, performing well on tail nodes is important since they are prevalent in real-world scale-free graphs (Clauset et al., 2009) and newly-arriving cold-start nodes (Lika et al., 2014) . To better train a GNN on those hard-to-predict tail nodes, the key idea of TUNEUP is to use a curriculum learning strategy (Bengio et al., 2009) ; TUNEUP first trains a GNN

Semi-sup node classification

Link prediction Recommender systems Figure 1 : Degree-specific generalization performance of the base GNN and TUNEUP in the transductive setting. The x-axis represents the node degrees in the training graph, and the y-axis is the generalization performance averaged over nodes with the specific degrees. We see from the dotted blue curves that the base GNN tends to perform poorly on tail nodes, i.e., nodes with small degrees. Our TUNEUP (denoted by the solid orange curves) improves or at least maintains the base GNN performance on almost all node degrees. The improvement is more significant on tail nodes. to perform well on relatively easy head nodes, i.e., nodes with a large number of neighbors. It then proceeds to improve the performance on the hard tail nodes. Specifically, TUNEUP uses the two-stage strategy to train a GNN. In the first stage, TUNEUP employs the default training strategy, i.e., simply minimizing the given supervised loss, to produce a strong base GNN to start with. The base GNN tends to perform well on head nodes, but poorly on tail nodes (see the dotted blue curves in Figure 1 ). To mitigate this issue, the second stage of TUNEUP focuses on improving the performance on the tail nodes. Specifically, TUNEUP synthesizes many additional tail node inputs by dropping edges from head nodes. TUNEUP then adds target supervision (e.g., class labels for node classification, edges for link prediction) on the synthetic tail nodes by reusing the supervision on the original head nodes (before dropping edges). Finally, TUNEUP finetunes the base GNN by minimizing the loss over the increased supervised tail node data. The dedicated training on the synthetic tail nodes allows the resulting GNN to perform much better on the real tail nodes, while often even improving the performance on head nodes. TUNEUP is simple to implement on top of the default training pipeline of GNNs, as shown in Algorithm 1. Moreover, TUNEUP can be used to train any GNN model with any supervised loss, making it generally applicable to a broad range of node and edge-level prediction tasks. We extensively evaluate TUNEUP on a wide range of settings. We consider five diverse GNN architectures, three types of key prediction tasks (semi-supervised node classification, link prediction, and recommender systems) with a total of eight datasets, as well as both transductive (i.e., prediction on nodes seen during training) and inductive (i.e., prediction on new nodes never seen during training) settings. For the inductive setting, we additionally consider the challenging cold-start scenario (i.e., limited edge connectivity from new nodes) by randomly removing certain portions of edges from new nodes. Across all settings, TUNEUP produces consistent improvement on the generalization performance of GNNs. In the transductive setting, TUNEUP significantly improves the performance of base GNNs on tail nodes, while oftentimes even improving the performance on head nodes (see Figure 1 ). Moreover, our ablation study shows that the two-stage curriculum training strategy of TUNEUP is critical and gives significantly improved performance over its variant strategy without curriculum learning. Finally, we extensively compare our TUNEUP against recent graph augmentation techniques (Rong et al., 2020; Liu et al., 2022) and specialized methods for tail nodes (Liu et al., 2021; Zheng et al., 2022; Zhang et al., 2022; Kang et al., 2022) . Our TUNEUP outperforms all these methods in all settings, while being simpler and more general. Overall, our work demonstrates that training strategies can play an important role in improving generalization performance of GNNs.

2. GENERAL SETUP AND TUNEUP

TUNEUP is a curriculum learning strategy to train any GNN model with any supervised loss to solve node or edge-level prediction tasks over graphs. We first provide a general task setup for machine learning on graphs and review the default training strategy of GNNs to solve the task. We then present

