ReGrAt: Regularization in Graphs using Attention to handle class imbalance

Handling class imbalance is a long-studied problem in the context of Machine Learning, and much work has been done previously on it. As mentioned earlier, the methods are broadly classified into three types: data-level, model-level, and hybrid approaches.

Data level methods usually resample the data in a way to make the distribution more balanced than before. SMOTE (Chawla et al. 2002) is a widely used state-of-the-art method that handles data imbalance by adding synthetic data to make data balanced. In SMOTE, artificial samples are produced by interpolating between existing minority samples and their nearest minority neighbors. SMOTE has been tried in the context of graphs to tackle imbalance where new nodes are created from minority classes. (Zhao, Zhang, and Wang 2021). An edge generator is used to find out the links between new nodes and the existing ones. An embedding space is constructed to encode the similarity among the nodes and generate new samples. In addition, an edge generator is trained simultaneously to model the relational information and provide it for those new samples.

Model-level methods used till now use a cost-sensitive approach which increases the priority of minority classes. The loss function is reweighted to achieve this. Hybrid approaches train multiple models using the methods mentioned above and then create an Ensemble of the models to get the final results (Shi et al. 2021).

Graph Neural Networks (GNNs) have received much attention and developing rapidly because of the need to work on non-linear structured data. Graph Convolutional Networks (GCNs) (Kipf and Welling 2016) perform similar operations like CNNs. The node in the model learns features by inspecting the neighboring nodes. The major difference between CNNs and GNNs is that CNNs are specially built to operate on structured data, whereas GNNs are the generalized version of CNNs where the numbers of node connections vary and the nodes are unordered. CNNs are used in Euclidean data and GNNs are used in non-Euclidean data. Current GNNs follow a message-passing framework, which is composed of pattern extraction and interaction modeling within each layer. Graph Attention Networks (Veličković et al. 2017) overcome the shortcoming of GCNs by having dynamic weights instead of having static ones. This is done by using masked self-attention layers.

We use two types of base loss functions which are used specifically when the data is imbalanced. Weighted cross entropy loss and Focal loss (Lin et al. 2017). Weighted cross-entropy loss assigns every class different weights based on the number of samples in the class. The weights are assigned as follows:

$N(c)$ is the number of samples in class $c$. This ensures that classes with less samples that is, the minority classes get more weightage than the majority class while calculating the loss.

Lin et al. (2017) proposed an algorithm that was helping solve the problem of extreme class imbalance in object detection problems by suggesting Focal loss. We use focal loss in the context of graphs. Focal loss reshapes the cross-entropy loss in such a way that it reduces the impact of easily classified examples and stresses more on the hard classified examples. This is done by multiplying the cross-entropy loss by a modulating factor, $(1-pt)^{\gamma}$. The hyperparameter $\gamma$ acts as a scaling factor. It adjusts the rate at which easy examples are down-weighted. For easily classified examples, where pt tends towards 1, it causes the modulating factor to approach 0 and helps in reducing the sample’s effect on the loss. Focal loss is calculated according to the following equation.

As the value of $\gamma$ increases the model focuses more on the misclassified examples. Experiments have shown that using 0.6 as the $\gamma$ value gives optimal results.

We propose the following regularizer $KLDiv_{att}$. Out of all the attention heads used, we make one of the attention heads focus solely on the minority nodes. We find out $Adj_{minority}\subset Adj$ by taking rows of only the minority nodes from the adjacency matrix. Similarly, we find out $a_{minority}\subset a$ by taking rows of only minority nodes from the attention weights. We calculate the attention regularizer term by taking KLDivergence of $Adj_{minority}$ and $a_{minority}$. KLDivergence is used to measure the difference between two distributions. The $Adj_{minority}$ matrix consists of zeros and ones. Its value is one when the node has a minority node as its neighbor and is zero otherwise. $a_{minority}$ consists of the weights of minority nodes which are between zero and one. We want to maximise the weights and bring them numerically close to one of those nodes whose neighbours are minority nodes. This is similar to the distribution in $Adj_{minority}$. Therefore, we calculate KLDivergence as follows:

We suggest a new loss function by adding $KLDiv_{att}$ to the cross entropy loss function already used.

The model is then trained on this custom loss function. $\lambda$ is a hyperparameter which can be tuned between 0 to 1. Both the loss functions that are used aid in handling the imbalance along with the regularizer that we are using.

The datasets used in these experiments are Cora and CiteSeer. These are widely used benchmark datasets (Wu et al. 2020). Both are citation network datasets. The number of edges, nodes, node features and the number of classes are shown in Table 2.

The distribution of classes is shown in Table 3. We observe that the distribution of classes is uneven and imbalanced. We assume L1, L3, L5 and L6 as minority classes for Cora and the rest as majority classes. Similarly, we assume L3 and L4 as our minority classes in Citeseer.

Our experimental settings are as follows: A semi-supervised setup is used for node classification. Each class has 20 examples. The validation set has 500 nodes, and the test set has 1000 nodes. Learning rate is set to 5e-3 and weight decay is set to 5e-4 for all experiments. The test set has the same minority nodes as the training set as used by standard graph methods. We use a total of two attention heads while using Graph Attention networks. We vary the number of epochs varies between 100-1000. We found that training the model for 300 epochs had the best results while using the regularizer. We report the macro F1 score and the individual F1 scores of the minority nodes.

Table 1 shows the results of node classification for Cora and CiteSeer datasets. The methods that we use are: GCNs with weighted CE, GCNs with FL, existing GraphSMOTE method, Regularizer with weighted CE and Regularizer with FL. The metrics that we report are accuracy, AUC-ROC score and the macro F1 score. The macro-F1 score is the most suitable and important metric for node classification.

Comparison with GCNs: We run the experiments on Graph Convolutional Networks as our baseline using both weighted CE and FL as our loss. We see that all three performance metrics: Accuracy, AUC-ROC Score, and F1 Score, increase with the addition of Regularizer as shown in table 1. The individual F1 scores of minority nodes also increase, as shown in Table 4 for Cora and in Table 5 for CiteSeer dataset.

Comparison with GraphSMOTE: We compare our results with the existing state-of-the-art method to handle imbalance in Graph networks. We see a considerable rise in the F1 scores for both Cora and CiteSeer datasets, as shown in Table 1.

The method $Reg+CE$ works the best as it has the highest F1 Score. There is a rise in the individual minority F1 scores also apart from L5, as shown in Table 4 for the Cora dataset. The rise in the individual F1 Scores in minority classes, as shown in Tables 4, 5 indicates that the algorithm is focusing more on the minority nodes than before.