Representing Long-Range Context for Graph Neural Networks with Global Attention

Graph neural networks are powerful architectures for structured datasets . However , current methods struggle to represent long-range dependencies . Scaling the depth or width of GNNs is insufficient to broaden receptive fields as larger GNNs encounter optimization instabilities such as vanishing gradients and representation oversmoothing , while pooling-based approaches have yet to become as universally useful as in computer vision . In this work , we propose the use of Transformer-based self-attention to learn long-range pairwise relationships , with a novel “ readout ” mechanism to obtain a global graph embedding . Inspired by recent computer vision results that find position-invariant attention performant in learning long-range relationships , our method , which we call GraphTrans , applies a permutation-invariant Transformer module after a standard GNN module . This simple architecture leads to state-of-the-art results on several graph classification tasks , outperforming methods that explicitly encode graph structure . Our results suggest that purely-learning-based approaches without graph structure may be suitable for learning high-level , long-range relationships on graphs . Code for GraphTrans is available at https : //github.com/ucbrise/graphtrans . 1 Introduction . Graph neural networks ( GNNs ) enable deep networks to process structured inputs such as molecules or social networks . GNNs learn mappings that compute representations at graph nodes and/or edges from the structure of and features in their neighborhoods . This neighborhood-local aggregation leverages the relational inductive bias encoded by the graph ’ s connectivity [ 3 ] . Similar to convolutional neural networks ( CNNs ) , GNNs can aggregate information from beyond local neighborhoods by stacking layers , effectively broadening the GNN receptive field . However , GNN performance drops dramatically when its depth increases [ 21 ] . This limitation has hurt the performance of GNNs on whole-graph classification and regression tasks , where we want to predict a target value describing the whole graph that may rely on long-range dependencies that may not be captured by a GNN with a limited receptive field [ 35 ] . Consider for example a large graph where node A must attend to a distant node B which is K-hops away . If our GNN layer aggregates only over a node ’ s one-hop neighborhood , then a K-layer GNN is required . However , the width of the receptive field of this GNN will grow exponentially , diluting the signal from node B . That is , simply expanding the receptive field to a K-hop neighborhood may not capture these long-range dependencies either [ 40 ] . Often , “ too deep ” GNNs lead to node representations that collapse to be equivalent over the entire graph , a phenomenon sometimes called oversmoothing or oversquashing [ 21 , 5 , 2 ] . Therefore , the maximum context size for common GNN architectures is effectively limited . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Several proposed methods combat the oversmoothing problem via intermediate pooling operations similar to those found in today ’ s CNNs . Graph pooling operations gradually coarsen the graph in progressive GNN layers , usually by collapsing neighborhoods into single nodes [ 9 , 37 , 20 , etc. ] . In theory , hierarchical coarsening should allow better long-range learning , both by reducing the distance information has to travel and by filtering out unimportant nodes . However , no graph pooling operation has been found that is as universally applicable as CNN pooling . State-of-the-art results are often obtained with models using no intermediate graph coarsening [ 27 ] , and some results suggest neighborhood-local coarsening may be unnecessary or counterproductive [ 23 ] . In this work , we take a different approach at graph pooling and learning long-range dependencies in GNNs . Like hierarchical pooling , our method is also inspired by methods for computer vision : we replace some of the atomic operations that explicitly encode relevant relational inductive biases ( i.e. , convolutions or spatial pooling in CNNs , neighborhood coarsening in GNNs ) with purely learned operations like attention [ 11 , 4 , 7 ] . Our method , which we call Graph Transformer ( GraphTrans , see Fig . 1 ) , adds a Transformer subnetwork on top of a standard GNN layer stack . This Transformer subnetwork explicitly computes all pairwise node interactions in a position-agnostic fashion . This approach is intuitive as it retains the GNN as a specialized architecture to learn local representations of the structure of a node ’ s immediate neighborhood while leveraging the Transformer as a powerful global reasoning module . This parallels recent computer vision architectures , where authors have found hard relational inductive biases important for learning short-range patterns but less useful or even counterproductive in modeling long-range dependencies [ 25 ] . As the Transformer without a positional encoding is permutationinvariant , we find it is a natural fit for graphs . Moreover , GraphTrans does not require any specialized modules or architectures and can be implemented in any framework atop any existing GNN backbone . We evaluate GraphTrans on a variety of popular graph classification datasets . We find significant improvements in accuracy on OpenGraphBenchmark [ 15 ] where we achieve state-of-the-art results on two graph classification tasks . Moreover , we find substantial improvements on the molecular dataset NCI1 . Surprisingly , we find our simple model outperforms complex baselines for long-range modeling in graphs via hierarchical clustering such as self-attention pooling [ 20 ] . Our contributions are as follows : • We show that long-range reasoning via Transformers improve graph neural network ( GNN ) accuracy . Our results suggest that modelling all pairwise node-node interactions in the graph is particularly important for large graph classification tasks . • We introduce a novel GNN “ readout module. ” Inspired by text-classification applications of Transformers , we use a special “ < CLS > ” token whose output embedding aggregates all pairwise interactions into a single classification vector . We find that this approach outperforms both non-learned readout methods like global pooling as well as learned aggregation methods like graph-specific pooling methods [ 37 , 20 ] and “ virtual node ” approaches . • Using our novel architecture GraphTrans , we obtain state-of-the-art results on several OpenGraphBenchmark [ 15 ] datasets and the NCI biomolecular datasets [ 30 ] . 2 Related Work . Graph Classification . Graph classification is an important task in real-world applications . Though GNNs encode the structured data into the node representations , aggregation of the representations to a single graph embedding for graph classification is still a problem . Similar to CNNs , pooling in GNNs can be either global , reducing a set of node and/or edge encodings to a single graph encoding , or local , collapsing subsets of nodes and/or edges to create a coarser graph . Paralleling the use of intermediate pooling within CNNs , several authors have proposed local pooling operations meant to be used within the GNN layer stack , progressively coarsening the graph . Methods proposed include both learned pooling schemes [ 37 , 20 , 14 , 16 , 1 , etc . ] and non-learned pooling methods based on classic graph coarsening schemes [ 10 , 9 , etc. ] . However , the effectiveness or necessity of hierarchical , coarsening-based pooling in GNNs is unclear [ 23 ] . On the other hand , the most common global , whole-graph pooling methods , are i ) non-learned mean or max-pooling over nodes and ii ) the “ virtual node ” approach , where a final GNN layer outputs an embedding for a single virtual node that is connected to every “ real ” node in the graph . A notable work related to graph pooling is the DAGNN ( Directed Acyclic Graph Neural Network ) of Thost and Chen [ 27 ] , which had obtained the previous state-of-the-art accuracy on OGBG-Code2 . The DAGNN layer aggregates over the entire graph within each layer via an RNN that traverses the DAG , unlike most GNN layers that only aggregate over a node ’ s neighborhood . While they did not characterize this method as a pooling operation , it is similar to GraphTrans in that it acts as a learned global pooling ( in that it aggregates the embeddings of every node in a DAG into the sink nodes ) that can model long-range dependencies . Note that GraphTrans is also complementary to DAGNN because their final graph-level pooling operation is a global max-pooling over the sink nodes rather than a learned operation . Transformers on Graphs . Several authors have investigated applications of Transformer architectures to graphs . Recent works such as Zhang et al . [ 38 ] , Rong et al . [ 24 ] , and Dwivedi and Bresson [ 12 ] propose GNN layers that let nodes attend to other nodes in some surrounding neighborhood via Transformer-style attention , whereas we use self attention for a permutation-invariant , graph-level pooling or “ readout ” operation that collapses node encodings to a single graph encoding . Of these , Zhang et al . [ 38 ] and Rong et al . [ 24 ] tackle the problem of learning long-range dependencies without over smoothing by allowing nodes to attend to more than just the one-hop neighborhood : Zhang et al . [ 38 ] take the attended neighborhood radius as a tuning parameter and Rong et al . [ 24 ] attend to neighborhoods of random size during training and inference . In contrast , we use whole-graph self-attention to allow for learning of long-range dependencies . While Zhang et al . [ 38 ] do not consider whole-graph prediction problems , in the case of Dwivedi and Bresson [ 12 ] , when a graph-wide embedding was needed for graph classification or regression , they used global average pooling over the nodes , while Rong et al . [ 24 ] take a weighted sum over nodes with the weights computed bypassing the hLv ’ s to a two-layer MLP . Note also that prior works consider graph-specific versions of a Transformer ’ s positional encoding , while we omit positional encodings to ensure permutation invariance . Efficient Transformers . Transformer [ 28 ] has been widely used in sequence modeling . Recently , modifications of the transformer architecture emerge to further improve the efficiency [ 34 , 19 , 6 ] . The LiteTransformer [ 34 ] with less FLOPs , Reformer [ 19 ] with complexity , and Performer [ 6 ] with both less computation and memory complexity . Neural architecture search ( NAS ) was also applied to Transformer to fulfill the resource constraints for the edge devices [ 32 ] . These off-the-shelf architectures are orthogonal to our GraphTrans and can be adopted to improve the scalability . 3 Motivation : Modeling Long-Range Pairwise Interactions . To summarize , attempting long-range learning on graphs via stacking GNN layers or hierarchical pooling have not yet led to performance increases , and while some works have shown some success in expanding the receptive field of a single GNN layer beyond a one-hop neighborhood [ 38 , 24 , 40 ] , it remains to be seen how this approach will scale to very large graphs with thousands of nodes . An inspiration for an alternative approach can be found in the recent computer vision literature . In the last few years , researchers have found that attention mechanisms can act as drop-in replacements for traditional CNN convolutions [ 4 , 7 ] : attention layers can learn to reproduce the strong relational inductive biases induced by local convolutions . More recently , state-of-the-art approaches to several computer vision tasks use an attention-style submodule on top of a traditional CNN backbone [ 2 , 33 , etc. ] . These results suggest that while strong relational inductive biases are helpful for learning local , short-range correlations , for long-range correlations less structured modules may be preferred [ 2 ] . We leverage this insight to the graph learning domain with our GraphTrans model , which uses a traditional GNN subnetwork as a backbone , but leaves learning long-range dependencies to a Transformer subnetwork with no graph spatial priors . As mentioned , our Transformer application lets every node attend to every other node ( unlike other approaches of applying Transformers to graphs that only allow attention to neighborhoods ) , which incentivizes the Transformer to learn the most important node-node relationships , instead of favoring nearby nodes ( the latter task having been offloaded to the preceding GNN module ) . Qualitatively , this scheme provides evidence that long-range relationships are indeed important . An example application of GraphTrans on the OGB Code2 dataset is depicted in Figure 2 . In this task , we take in the Abstract Sentence Tree obtained by parsing a Python method and need to predict the tokens that form the method name . The attention map exhibits similar patterns to those found in NLP applications of Transformers : some nodes receive significant weighting from many other nodes , regardless of the distance between them . Note that node 17 assigns significant importance to node 8 , despite these two nodes being five hops away . Also , in Figure 2 ’ s attention map , index 18 refers to the embedding corresponding to the special < CLS > token we use as a readout mechanism , described in more detail below . We allow this embedding to be learnable , so the many nodes attending to it ( represented by the many dark cells in column 18 ) may suggest these nodes are obtaining some graphgeneral memory from the learned embedding . This qualitative visualization , along with our new state-of-the-art results , suggest that removing spatial priors when learning long-range dependencies may be necessary for effective graph summarization . 4 Learning Global Information with GraphTrans Referring back to Figure 1 , GraphTrans consists of two primary modules : a GNN subnetwork followed by a Transformer subnetwork . We discuss these in detail next . GNN module . We consider graph property prediction , i.e. , for each graph G = ( V , E ) we have a graph-specific prediction target yG . We suppose that each node v ∈ V has an initial feature vector h0v ∈ Rd0 . As GraphTrans is a generally-applicable framework that can be used in concert with a variety of GNNs , we make very few assumptions on the GNN layers that feed into the Transformer subnetwork . A generic GNN layer stack can be expressed as h ` v = f ` ( h ` −1v , { h ` −1u |u ∈ N ( v ) } ) , ` = 1 , . . . , LGNN ( 1 ) where LGNN is the total number of GNN layers , N ( v ) ⊆ V is some neighborhood of v , and f ` ( · ) is some function parameterized by a neural network . Note that many GNN layers admit edge features , but to avoid notational clutter we omit discussion of them here . Transformer module . Once we have the final per-node GNN encodings hLGNNv , we pass these to GraphTrans ’ s Transformer subnetwork . The Transformer subnetwork operates as follows . We first perform a linear projection of the hLGNNv ’ s to the Transformer dimension and a Layer Normalization to normalize the embedding : h0v = LayerNorm ( W ProjhLGNNv ) ( 2 ) where W Proj ∈ RdTF×dLGNN is a learnable weight matrix , and dTF and dLGNN are the Transformer dimension and the dimension of the final GNN embedding , respectively . The projected node embeddings h0v are then fed into a standard Transformer layer stack , with no additive positional embeddings , as we expect the GNN to have already encoded the structural information into the node embeddings : a ` v , u = ( W Q ` h ` −1 v ) > ( WK ` h ` −1 u ) / √ dTF α ` v , u = softmax w∈V ( a ` v , w ) h′ ` v = ∑ w∈V α ` v , wW V ` h ` −1 w ( 3 ) where WQ ` , W K ` , W V ` ∈ RdTF/nhead×dTF/nhead are the learned query , key , and value matrices , respectively , for a single attention head in layer ` . As is standard , we run nhead parallel attention heads and concatenate the resulting per-head encodings h′ ` v . Concatenated encodings are then passed to a Transformer fully-connected subnetwork , consisting of the standard Dropout→ Layer Norm→ FC → nonlinearity→ Dropout→ FC→ Dropout→ Layer Norm sequence , with residual connections from h ` −1v to after the first dropout , and from before the first fully-connected sublayer to after the dropout immediately following the second fully-connected sublayer . < CLS > embedding as a GNN “ readout ” method . As mentioned , for whole-graph classification we require a single embedding vector that describes the whole graph . In the GNN literature , this module that collapses embeddings for every node and/or edge to a single embedding is called the “ readout ” module , and the most common readout modules are simple mean or max pooling , or a single “ virtual node ” that is connected to every other node in the network . In this work , we propose a special-token readout module similar to those used in other applications of Transformers . In text classification tasks with Transformers , a common practice is to append a special < CLS > token to the input sequence before passing it into the network , then to take the output embedding corresponding to this token ’ s position as the representation of the whole sentence . In that way , the Transformer will be trained to aggregate information of the sentence to that embedding , by calculating the one-to-one relationships between the < CLS > token and each other tokens in the sentence with the attention module . Our application of special-token readout is similar to this . Concretely , when feeding the transformed per-node embeddings h̄0v , we append an additional learnable embedding h < CLS > to the sequence , and take the first embedding h̄ < CLS > ∈ RdTF from the transformer output as the representation of the whole graph ( note that since we do not include positional encodings , placing the special token at the “ beginning ” of the sentence has no special computational meaning ; the location is chosen by convention ) . Finally , we apply a linear projection followed by a softmax to generate the prediction : y = softmax ( W outhLTF < CLS > ) . ( 4 ) where LTF is the number of Transformer layers . This special-token readout mechanism may be viewed as a generalization or a “ deep ” version of a virtual node readout . While a virtual node method requires every node in the graph to send its information to the virtual node and does not allow for learning pairwise relationships between graph nodes except within the virtual node ’ s embedding ( possibly creating an information bottleneck ) , a Transformer-style special-token readout method lets the network learn long-range node-to-node relationships in earlier layers before needing to distill them in the later layers .

Summary: Current GNN methods fail at utilizing long-range dependencies in graphs due to over smoothing. However, in some cases, long-range information might be useful. Inspired by the success of transformers in computer vision tasks, this paper tackles the challenge of learning long-range dependencies in GNNs by simply augmenting a common GNN architecture with a transformer sub-network. They also borrow a “readout” token from NLP which improves the final learned representation of the graph.