Masked Label Prediction: Unified Message Passing Model for Semi-Supervised Classification

Graph neural network ( GNN ) and label propagation algorithm ( LPA ) are both message passing algorithms , which have achieved superior performance in semisupervised classification . GNN performs feature propagation by a neural network to make predictions , while LPA uses label propagation across graph adjacency matrix to get results . However , there is still no good way to combine these two kinds of algorithms . In this paper , we proposed a new Unified Message Passaging Model ( UniMP ) that can incorporate feature propagation and label propagation with a shared message passing network , providing a better performance in semisupervised classification . First , we adopt a Graph Transformer jointly label embedding to propagate both the feature and label information . Second , to train UniMP without overfitting in self-loop label information , we propose a masked label prediction strategy , in which some percentage of training labels are simply masked at random , and then predicted . UniMP conceptually unifies feature propagation and label propagation and be empirically powerful . It obtains new state-of-the-art semi-supervised classification results in Open Graph Benchmark ( OGB ) . 1 INTRODUCTION . There are various scenarios in the world , e.g. , recommending related news and products , discovering new drugs , or predicting social relations , which can be described as graph structures . Many methods have been proposed to optimize these graph-based problems and achieved significant success in many related domains such as predicting the properties of nodes ( Yang et al. , 2016 ; Kipf & Welling , 2016 ) , links ( Grover & Leskovec , 2016 ; Battaglia et al. , 2018 ) , and graphs ( Duvenaud et al. , 2015 ; Niepert et al. , 2016 ; Bojchevski et al. , 2018 ) . In the task of semi-supervised node classification , we are required to learn with labeled examples and then make predictions for those unlabeled ones . To better classify the nodes ’ labels in the graph , based on the Laplacian smoothing assumption ( Li et al. , 2018 ; Xu et al. , 2018b ) , the message passing models were proposed to aggregate the information from its connected neighbors in the graph , acquiring enough facts to produce a more robust prediction for unlabeled nodes . Generally , there are two kinds of practical methods to implement message passing model , the Graph Neural Networks ( GNNs ) ( Kipf & Welling , 2016 ; Hamilton et al. , 2017 ; Xu et al. , 2018b ; Liao et al. , 2019 ; Xu et al. , 2018a ; Qu et al. , 2019 ) and the Label Propagation Algorithms ( LPAs ) ( Zhu , 2005 ; Zhu et al. , 2003 ; Zhang & Lee , 2007 ; Wang & Zhang , 2007 ; Karasuyama & Mamitsuka , 2013 ; Gong et al. , 2016 ; Liu et al. , 2019 ) . GNNs combine graph structures by propagating and aggregating nodes features through several neural layers , which get predictions from feature propagation . While LPAs make predictions for unlabeled instances by label propagation iteratively . Since GNN and LPA are based on the same assumption , making semi-supervised classifications by information propagation , there is an intuition that incorporating them together for boosting performance . Some superior studies have proposed their graph models based on it . For example , APPNP ( Klicpera et al. , 2019 ) and TPN ( Liu et al. , 2019 ) integrate GNN and LPA by concatenating them together , and GCN-LPA ( Wang & Leskovec , 2019 ) uses LPA to regularize their GCN model . How- ever , as shown in Tabel 1 , aforementioned methods still can not incorporate GNN and LPA within a message passing model , propagating feature and label in both training and prediction procedure . To unify the feature and label propagation , there are mainly two issues needed to be addressed : Aggregating feature and label information . Since node feature is represented by embeddings , while node label is a one-hot vector . They are not in the same vector space . In addition , there are different between their message passing ways , GNNs can propagate the information by diverse neural structures likes GraphSAGE ( Hamilton et al. , 2017 ) , GCN ( Kipf & Welling , 2016 ) and GAT ( Veličković et al. , 2017 ) . But LPAs can only pass the label message by graph adjacency matrix . Supervised training . Supervised training a model with feature and label propagation will overfit in self-loop label information inevitably , which makes the label leakage in training time and causes poor performance in prediction . In this work , inspired by several advantages developments ( Vaswani et al. , 2017 ; Wang et al. , 2018 ; Devlin et al. , 2018 ) in Natural Language Processing ( NLP ) , we propose a new Unified Message Passing model ( UniMP ) with masked label prediction that can settle the aforementioned issues . UniMP is a multi-layer Graph Transformer , jointly using label embedding to transform nodes labels into the same vector space as nodes features . It propagates nodes features like the previous attention based GNNs ( Veličković et al. , 2017 ; Zhang et al. , 2018 ) . Meanwhile , its multi-head attentions are used as the transition matrix for propagating labels vectors . Therefore , each node can aggregate both features and labels information from its neighbors . To supervised training UniMP without overfitting in self-loop label information , we draw lessons from masked word prediction in BERT ( Devlin et al. , 2018 ) and propose a masked label prediction strategy , which randomly masks some training instances ’ label embedding vectors and then predicts them . This training method perfectly simulates the procedure of transducing labels information from labeled to unlabeled examples in the graph . We conduct experiments on three semi-supervised classification datasets in the Open Graph Benchmark ( OGB ) , where our new methods achieve novel state-of-the-art results in all tasks , gaining 82.56 % ACC in ogbn-products , 86.42 % ROC-AUC in ogbn-proteins and 73.11 % ACC in ogbnarxiv . We also conduct the ablation studies for the models with different inputs to prove the effectiveness of our unified method . In addition , we make the most thorough analysis of how the label propagation boosts our model ’ s performance . 2 METHOD . We first introduce our notation about graph . We denote a graph as G = ( V , E ) , where V denotes the nodes in the graph with |V | = n and E denotes edges with |E| = m. The nodes are described by the feature matrix X ∈ Rn×f , which usually are dense vectors with f dimension , and the target class matrix Y ∈ Rn×c , with the number of classes c. The adjacency matrix A = [ ai , j ] ∈ Rn×n is used to describe graph G , and the diagonal degree matrix is denoted by D = diag ( d1 , d2 , ... , dn ) , where di = ∑ j aij is the degree of node i . A normalized adjacency matrix is defined as D −1A or D− 1 2AD− 1 2 , and we adopt the first definition in this paper . 2.1 FEATURE PROPAGATION AND LABEL PROPAGATION . In semi-supervised node classification , based on the Laplacian smoothing assumption , the GNN transforms and propagates nodes features X across the graph by several layers , including linear layers and nonlinear activation to build the approximation of the mapping : X → Y . The feature propagation scheme of GNN in layer l is : H ( l+1 ) = σ ( D−1AH ( l ) W ( l ) ) Y = fout ( H ( L ) ) ( 1 ) where the σ is an activation function , W ( l ) is the trainable weight in the l-th layer , and the H ( l ) is the l-th layer representations of nodes . H ( 0 ) is equal to node input featuresX . Finally , a fout output layer is applied on the final representation to make prediction for Y . As for LPA , it also assumes the labels between connected nodes are smoothing and propagates the labels iteratively across the graph . Given an initial label matrix Ŷ ( 0 ) , which consists of one-hot label indicator vectors ŷ0i for the labeled nodes or zeros vectors for the unlabeled . A simple iteration equation of LPA is formulated as following : Ŷ ( l+1 ) = D−1AŶ ( l ) ( 2 ) Labels are propagated from each other nodes through a normalized adjacency matrix D−1A . 2.2 UNIFIED MESSAGE PASSING MODEL . As shown in Figure 1 , we employ a Graph Transformer , jointly using label embedding to construct our unified message passing model for combining the aforementioned feature and label propagation together . Graph Transformer . Since Transformer ( Vaswani et al. , 2017 ) has been proved being powerful in NLP , we employ its vanilla multi-head attention into graph learning with taking into account the case of edge features . Specifically , given nodes features H ( l ) = { h ( l ) 1 , h ( l ) 2 , ... , h ( l ) n } , we calculate multi-head attention for each edge from j to i as following : q ( l ) c , i = W ( l ) c , qh ( l ) i + b ( l ) c , q k ( l ) c , j = W ( l ) c , kh ( l ) j + b ( l ) c , k ec , ij = Wc , eeij + bc , e α ( l ) c , ij = 〈q ( l ) c , i , k ( l ) c , j + ec , ij〉∑ u∈N ( i ) 〈q ( l ) c , i , k ( l ) c , u + ec , iu〉 ( 3 ) where 〈q , k〉 = exp ( qT k√ d ) is exponential scale dot-product function and d is the hidden size of each head . For the c-th head attention , we firstly transform the source feature h ( l ) i and distant feature h ( l ) j into query vector q ( l ) c , i ∈ Rdand key vector k ( l ) c , j ∈ Rd respectively using different trainable parameters W ( l ) c , q , W ( l ) c , k , b ( l ) c , q , b ( l ) c , k . The provided edge features eij will be encoded and added into key vector as additional information for each layer . After getting the graph multi-head attention , we make a message aggregation from the distant j to the source i : v ( l ) c , j = W ( l ) c , vh ( l ) j + b ( l ) c , v ĥ ( l ) i = ∥∥∥C c=1 [ ∑ j∈N ( i ) α ( l ) c , ij ( v ( l ) c , j + ec , ij ) ] r ( l ) i = W ( l ) r h ( l ) i + b ( l ) r β ( l ) i = sigmoid ( W ( l ) g [ ĥ ( l ) i ; r ( l ) i ; ĥ ( l ) i − r ( l ) i ] ) h ( l+1 ) i = ReLU ( LayerNorm ( ( 1− β ( l ) i ) ĥ ( l ) i + β ( l ) i r ( l ) i ) ) ( 4 ) where the ‖ is the concatenation operation for C head attention . Comparing with the Equation 1 , multi-head attention matrix replaces the original normalized adjacency matrix as transition matrix for message passing . The distant feature hj is transformed to vc , j ∈ Rd for weighted sum . In addition , inspired by ( Li et al. , 2019 ; Chen et al. , 2020 ) to prevent oversmoothing , we propose a gated residual connections between layers by ri ∈ Rd and β ( l ) i ∈ R1 . Specially , similar to GAT , if we apply the Graph Transformer on the output layer , we will employ averaging for multi-head output as following : ĥ ( l ) i = 1 C C∑ c=1 [ ∑ j∈N ( i ) α ( l ) c , ij ( v ( l ) c , j + e ( l ) c , ij ) ] h ( l+1 ) i = ( 1− β ( l ) i ) ĥ ( l ) i + β ( l ) i r ( l ) i ( 5 ) Label Embedding and Propagation . We propose to embed the partially observed labels information into the same space as nodes features : Ŷ ∈ Rn×c → Ŷe ∈ Rn×f , which consist of the label embedding vector for labeled nodes and zeros vectors for the unlabeled . And then , we combine the label propagation into Graph Transformer by simply adding the nodes features and labels features together as propagation features ( H0 = X + Ŷe ) ∈ Rn×f . We can prove that by mapping partiallylabeled Ŷ and nodes features X into the same space and adding them up , our model is unifying both label propagation and feature propagation within a shared message passing framework . Let ’ s take Ŷe = Ŷ We and A∗ to be normalized adjacency matrix D−1A or the attention matrix from our Graph Transformer like Equation 3 . Then we can find that : H ( 0 ) = X + Ŷ Wc H ( l+1 ) = σ ( ( ( 1− β ) A∗ + βI ) H ( l ) W ( l ) ) ( 6 ) where β can be the gated function like Equation 4 or a pre-defined hyper-parameters like APPNP ( Klicpera et al. , 2019 ) . For simplification , we let σ function as identity function , then we can get : H ( l ) = ( ( 1− β ) A∗ + βI ) l ( X + Ŷ Wc ) W ( 1 ) W ( 2 ) . . .W ( l ) = ( ( 1− β ) A∗ + βI ) lXW + ( ( 1− β ) A∗ + βI ) lŶ WcW ( 7 ) where W = W ( 1 ) W ( 2 ) . . .W ( l ) . Then we can find that our model can be approximately decomposed into feature propagation ( ( 1 − β ) A∗ + βI ) lXW and label propagation ( ( 1 − β ) A∗ + βI ) lŶ WcW .

Summary: The authors proposed a unified message passing model to make a graph neural network to be able to incorporate both label propagation and feature propagation. Compared to previous work, the proposed model can also make use of partial label information in both training and inference stages. Experiments on three OGBN datasets show that the proposed methods achieve promising performance.