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 a means to encode text feature for nodes in textual graphs, where a layerwise aggregation module (implemented as a multi-head self-attention) is appended after each Transformer encoder block that aggregates the hidden representations of (sampled) neighbor nodes. The produced contextual representations are dispatched to each node for the computation of next Transformer encoder block. To mitigate the possible model degradation issue that the aggregation module might be insufficiently trained, the authors also propose a two-stage training schedule where the model is first trained on nodes with corrupted inputs until convergence and then trained on normal data.