GraphFormers: GNN-nested Transformers for Representation Learning on Textual Graph

1 Introduction . The textual graph is a widely existed data format , where each node is annotated with its textual feature . The representation learning on textual graph is to generate low-dimensional node embeddings based on the individual textual features and the information from the neighbourhood . In recent years , the breakthroughs in pretrained language models and graph neural networks contribute to the development of corresponding techniques . Particularly , with pretrained language models , such as BERT ( Devlin et al. , 2018 ) and RoBERTa ( Liu et al. , 2019a ) , the underlying semantics of texts can be captured more precisely ; at the same time , with graph neural networks , like GraphSage ( Hamilton et al. , 2017a ) and GAT ( Veličković et al. , 2018 ) , neighbours can be effectively aggregated for more informative node embeddings . It is necessary to combine both techniques for better textual graph representation . As ∗Work was done by 2021.01 during Junhan Yang and Shitao Xiao ’ s internship in MSRA 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . suggested by GraphSage ( Hamilton et al. , 2017a ) and PinSage ( Ying et al. , 2018 ) , the textual feature can be independently modeled by text encoders and further aggregated by rear-mounted GNNs for the final node embeddings . Such a representation paradigm has been widely adopted by subsequent works on various scenarios ( Zhu et al. , 2021 ; Li et al. , 2021 ; Hu et al. , 2020 ; Liu et al. , 2019b ; Zhou et al. , 2019 ) , where GNNs are combined with powerful PLM-based text encoders . The above way of combination is called the “ Cascaded Transformers-GNN ” architecture ( Figure 1 A ) , as the language models ( built upon Transformers ) are deployed ahead of the GNN component . With the above architecture , the text encoding and the graph aggregation are performed in two consecutive steps , where there is no information exchange between the nodes when text embeddings are generated . However , the above workflow is defective considering that the linked nodes are correlated , whose underlying semantics can be mutually enhanced . For example , given a node “ notes on transformers ” and its neighbour “ tutorials on machine translation ” ; by making reference to the whole context , the “ transformers ” here can be interpreted as a machine learning model , rather than an electric device . Our Work . We propose “ GNN-nested Transformers ” ( GraphFormers ) , which are highlighted for the fusion of GNNs and language models ( Figure 1 B ) . In GraphFormers , the GNN components are nested alongside the transformer layers ( TRM ) of language models , where the text encoding and graph aggregation are fused as an iteratively workflow . In each iteration , the linked nodes will exchange information with each other in the layerwise GNN component ; thus , each node will be augmented by its neighbourhood information . The transformer component will work on the augmented node features , where increasingly informative node representations can be generated for the next iteration . Compared with the cascaded architecture , GraphFormers achieve more sufficient utilization of the cross-node information on graph , which significantly benefit the representation quality . Given that the layerwise GNN components merely involve simple and effective multi-head attention , GraphFormers preserve comparable running costs as the existing cascaded Transformers-GNN models . On top of the proposed model architecture , we further improve GraphFormers ’ representation quality and practicability as follows . Firstly , the training of GraphFormers is likely to be shortcut : in many cases , the center node itself can be “ sufficiently informative ” , where the training tasks can be accomplished without leveraging the neighbourhood information . As such , GraphFormers may end up with insufficiently trained GNNs . Inspired by recent success of curriculum learning ( Bengio et al. , 2009 ) , we propose to train the model progressively : the first round of training is performed with manipulated data , where the nodes are randomly polluted ; thus , it becomes harder to make prediction merely rely on the center nodes , and the model will be forced to leverage the whole input nodes . The second round of training gets back to the unpolluted data , where the model will be fit into the targeted distribution . Another concern about GraphFormers is that all the linked nodes are mutually dependent in the representation process : once a new node is presented , all the neighbours , regardless of whether they have been processed before , need to be encoded from scratch . As a result , a great deal of unnecessary computations will be incurred . We introduce unidirectional graph attention to alleviate this problem : only the center node is required to make reference to the neighbours , while the neighbour nodes remain independently encoded . By this means , the existing neighbours ’ encoding results can be cached and reused , which significantly saves the computation cost . Extensive evaluations are conducted with three million-scale textual graph datasets : DBLP , Wiki and Product , where the representation quality is measured by the link prediction accuracy . According to our experiment results , GraphFormers significantly outperform the SOTA cascaded TransformersGNN baselines with comparable running efficiency . 2 Related Work . The textual graph representation is an important research topic in multiple areas , such as natural language processing , information retrieval and graph learning ( Yang et al. , 2015 ; Wang et al. , 2016b , a ; Yasunaga et al. , 2017 ; Wang et al. , 2019a ; Xu et al. , 2019 ) . To learn high-quality representation for textual graph , techniques on natural language understanding and graph representation need to be jointly leveraged . In recent years , breakthroughs on pretrained language models ( PLM ) and graph neural networks ( GNN ) significantly advance the development of corresponding techniques . PLM . The PLMs are proposed to learn universal language models with neural networks trained on large-scale corpus . The early works were based on shallow networks , e.g , word embeddings learned by Skip-Gram ( Mikolov et al. , 2013 ) and GloVe ( Pennington et al. , 2014 ) . In recent years , the backbone networks are being quickly scaled up : from EMLo ( Peters et al. , 2018 ) , GPT ( Radford et al. , 2018 ) , to BERT ( Devlin et al. , 2018 ) , XLNet ( Yang et al. , 2019 ) , T5 ( Raffel et al. , 2019 ) , GPT-3 ( Brown et al. , 2020 ) . The large-scale models , which get fully trained with massive data , demonstrate superior performances on general NLP tasks . One of the most critical usages of PLMs is text representation , where the underlying semantics of texts are captured by low-dimensional embeddings . Such embeddings achieve competitive results on downstream tasks , like text retrieval and classification ( Reimers and Gurevych , 2019 ; Luan et al. , 2020 ; Gao et al. , 2021 ; Su et al. , 2021 ) . GNN . Graph neural networks are recognized as powerful tools of modeling graph data ( Hamilton et al. , 2017b ; Zhou et al. , 2020 ) . Such methods ( e.g. , GCN ( Kipf and Welling , 2016 ) , GAT ( Veličković et al. , 2018 ) , GraphSage ( Hamilton et al. , 2017a ) ) learn effective message passing mechanisms such that information between the nodes can get aggregated for expressive graph representations . Graph neural networks may also incorporate node attributes , like texts ; and it ’ s quite straightforward to leverage GNNs and PLMs for textual graph representation following the “ cascaded architecture ” suggested by GraphSage ( Hamilton et al. , 2017a ) : the node features are independently encoded at first ; then , the node embeddings are aggregated via GNNs to generate the final representations . Such a representation paradigm is widely adopted by subsequent works ( Zhu et al. , 2021 ; Li et al. , 2021 ; Hu et al. , 2020 ; Liu et al. , 2019b ; Zhou et al. , 2019 ) . However , the above approaches treat the text encoding and graph aggregation as two consecutive steps , where the node-level features are independently processed . Our work is different from these approaches as the text encoding and graph aggregation are fused as an iterative workflow based on the “ GNN-nested Transformers ” . 3 GraphFormers . In this work , we deal with textual graph data , where each node x is a text . The node x together with its neighbours Nx are denoted as Gx . Our model learns the embedding for node x based on its own textual feature and the information of its neighbourhood Nx . The generated embeddings are expected to capture the relationship between the nodes , i.e. , to accurately predict whether two nodes xq and xk are connected based on the embedding similarity . 3.1 GNN-nested Transformers . The encoding process of GraphFormers is indicated as follows . The input nodes ( the center node and its neighbours ) are tokenized into sequences of tokens , with special tokens [ CLS ] padded in the front , whose states are used for node representation . The input sequences are mapped into the initial embedding sequences { H0g } G based on the summation of word embeddings and position embeddings . The embedding sequences are encoded by multiple layers of GNN-nested Transformers ( shown as Figure 2 ) , where the graph aggregation and text encoding are iteratively performed . • Graph Aggregation in GNN . Each node is enhanced by its neighbourhood information based on the layerwise graph aggregation . For each node in the l-th layer , the first token-level embedding ( corresponding to [ CLS ] ) is taken as the node-level embedding : zlg ← Hlg [ 0 ] . The node-level embeddings are gathered from all the nodes and passed to the layerwise GNN for graph aggregation . We leverage Multi-Head Attention ( MHA ) to encode the node-level embeddings ZlG ( { zlg } G ) , similar as GAT ( Veličković et al. , 2018 ) . For each attention head , the scaled dot-product attention is performed as : ẐlG = MHA ( Z l G ) ; MHA ( ZlG ) = Concat ( head1 , ... , headh ) ; headj = softmax ( QKT√ d +B ) V ; Q = ZlGW Q j ; K = Z l GW K j ; V = Z l GW V j ; ( 1 ) In the above equations , WQj , W K j , and W V j are the projection matrices of MHA , corresponding to the j-th attention head . A learnable position bias B is added to the dot-product result ; the positions differentiate the relationship between the nodes ; i.e. , “ center-to-center ” ( x to x ) , “ center-to-neighbour ” ( x to Nx ) , and “ neighbour-to-neighbour ” ( Nx to Nx ) , respectively . Each of the embeddings ẑlg ( ẑ l g ∈ ẐlG ) is dispatched to its original node and concatenated ( ⊕ ) with the token-level embeddings , which gives rise to the graph-augmented token-level embeddings : Ĥlg ← Concat ( ẑlg , Hlg ) . ( 2 ) In this place , the GNN-processed node-level embeddings ẐlG can be interpreted as “ messagers ” , with which the neighbourhood information can be introduced to each of the nodes . • Text Encoding in Transformer . The graph-augmented token-level embeddings Ĥlg are processed by the transformer component ( Vaswani et al. , 2017 ) , where the following computations are performed : Ĥlg = LN ( H l g +MHA asy ( Ĥlg ) ) ; Hl+1g = LN ( Ĥ l g +MLP ( Ĥ l g ) ) . ( 3 ) In the above equations , MLP is the Multi-Layer Projection unit , and LN is the Layer-Norm unit . We use asymmetric Multi-Head Attention ( MHAasy ) , where Q , K , V are computed as : Q = HlgW Q j ; K = Ĥ l gW K j ; V = Ĥ l gW V j . ( 4 ) Therefore , the output sequence Hl+1g will be of the same length as the input sequence H l g. The encoding result will be used as the input token-level embeddings for the next layer . The node-level embedding at the last layer zLx ( i.e. , H L g [ 0 ] ) will be used as the final node representation . •Workflow . We summarize GraphFormers ’ encoding workflow as Algorithm 1 . The initial tokenlevel embeddings { H0g } G are independently encoded by the first Transformer layer TRM 0 . For a L-layer GraphFormers , the graph aggregation and text encoding are iteratively performed for the subsequent L-1 steps ( from 1 to L− 1 ) . In each step , the node-level embeddings ZlG are gathered and Algorithm 1 : GraphFormers ’ Workflow Input : The input graphs G ( consist of the center node x and its neighbours ) . Output : The embedding for the center node hx . 1 begin 2 for each text g ∈ G do 3 H1g ← TRM 0 ( H0g ) ; // Get the initial token-level embeddings 4 for l = 1 , ... , L− 1 do 5 ZlG ← { zlg|g∈G } ; // Gather node-level embeddings to GNN 6 ẐlG ← GNN ( ZlG ) ; // Graph aggregation in GNN 7 for each text g ∈ G do 8 Ĥlg ← Concat ( ẑlg , Hlg ) ; // Get graph-augmented token-level embeddings 9 Hl+1g ← TRM l ( Ĥlg ) ; // Text encoding in Transformer 10 Return hx ← zLx ; processed by the layerwise GNN component . The output node-level embeddings ẐlG are dispatched to their original nodes , which generates the graph-augmented token-level embeddings Ĥlg . The graphaugmented token-level embedding are further processed by the Transformer component . Finally , The node-level embedding ( for the center node x ) in the last layer zLx is taken as our representation result . • Encoding Complexity . Given an input of M nodes , each one has P tokens ; the time complexity of each layer ’ s encoding operation is O ( M2 +MP 2 ) : the graph aggregation takes O ( M2 ) , because M node-level embeddings are gathered for multi-head attention ; the text encoding takes O ( MP 2 ) , as each of the M node calls for the multi-head attention of P tokens . Compared with Transformers , the GNN ’ s computation cost is much smaller , mainly because of two reasons : 1 ) M2 MP 2 in general , 2 ) operations like MLP are not needed in graph aggregation . Therefore , the working efficiency of GraphFormers is close to the cascaded GNN-Transformers as the extra computation cost of layerwise graph aggregation is relatively small . Such a property is also empirically verified in our experiment .

Summary: This paper proposes GraphFormers for text graph tasks where GNNs are nested in each layer of the transformer-based pretrained language model. In this way, each node can appropriate graph structure to build semantic representations more effectively. Two-stage progressive training and unidirectional graph aggregation are further proposed to improve representation quality and avoid unnecessary computation. They use three real-world million-scale textual graph datasets (DBMLP, Wikidata, Product Graph) and demonstrate consistent improvements compared with traditional cascaded transformers-GNN models.