Empowering Graph Representation Learning with Paired Training and Graph Co-Attention

1 INTRODUCTION AND RELATED WORK . We study the task of graph-level representation learning : i.e. , computing representations of entire input graphs , for the purposes of downstream tasks ( such as graph classification or regression ) . This is typically a step-up in complexity compared to node classification or link prediction , given that the learning algorithm must aggregate useful structural information across the graph into a single prediction—relying on only this global supervision signal ( as opposed to having feedback from every node/edge of the graph ) . Perhaps the highest challenge this kind of architecture must face is inductivity and generalisation across structures . Specifically , an inductive model must be readily applicable across several graph structures—including ones unseen during training . Additionally , the model is tasked with discovering interesting structural “ motifs ” across the entire dataset of graphs , whose presence or absence may help determine the overall predictions . However , even enabling inductivity is not a traditionally simple task in graph representation learning , as many prior approaches ( Bruna et al. , 2013 ; Perozzi et al. , 2014 ; Defferrard et al. , 2016 ) are not inductive by design . Furthermore , even the models that are currently used for graph-level representation learning ; e.g . Gilmer et al . ( 2017 ) ; Ying et al . ( 2018 ) ; Xu et al . ( 2018 ) ; Lu et al . ( 2019 ) , operate over only a single graph at a time—making it challenging for them to reason about common substructures across graphs from a graph-level supervision signal alone . In this manuscript , we propose the approach of paired training—i.e. , learning representations over pairs of input graphs at once . Intuitively , as long as we allow for dataflow between the representations of the two graphs within a pair , this allows the graph neural network to directly observe related ( sub ) structures from other inputs , to solidify its decision making . We note that in the context of graph-structured inputs this may be particularly useful as , unlike simpler inputs such as images or text , there are no guarantees that different graphs within a dataset will have equal or even similar overall structure . To facilitate this dataflow , we propose the usage of graph co-attention for exchanging representations of nodes across the two graphs . Intuitively , this operator performs attention ( Bahdanau et al. , 2014 ; Vaswani et al. , 2017 ) over the fully-connected bipartite graph , with one part corresponding to all nodes in one graph . This allows every node of the first graph to detect and reuse useful patch representations in the second graph ( in a form of hierarchical graph matching ) , and vice-versa . Initially , we validate our model performance on a pairwise graph classification task—classifying drug pairs for side effects caused by drug-drug interactions ( DDI ) ( Jin et al. , 2017 ; Zitnik et al. , 2018 ) . In this setting , a pairwise approach is natural , as we inherently have to classify pairs of graphs . We demonstrate that learning a joint representation using graph co-attention provides substantial benefits to predictive power , setting the state-of-the-art result on this task . From there , we demonstrate the applicability of our approach to arbitrary multi-graph datasets ; for this , we leverage the QM9 dataset for predicting quantum chemistry properties of small molecules ( Ramakrishnan et al. , 2014 ) . As such , it represents a challenging graph regression problem . We propose using paired training to perform regression on two molecules at once , demonstrating clear benefits to doing so . In a similar vein , we execute variants of our model on standard graph kernel classification benchmarks ( Kersting et al. , 2016 ) , showing advantages to generic graph classification . Our approach paves the way to a promising direction for graph-level prediction tasks , that is in principle applicable to any kind of multi-graph dataset , especially under availability of large quantities of labelled examples . Our model builds up on a large existing body of work in graph convolutional networks ( Bruna et al. , 2013 ; Defferrard et al. , 2016 ; Kipf & Welling , 2016a ; Gilmer et al. , 2017 ; Veličković et al. , 2018 ) , that have substantially advanced the state-of-the-art in many tasks requiring graph-structured input processing ( such as the chemical representation ( Gilmer et al. , 2017 ; De Cao & Kipf , 2018 ; You et al. , 2018 ) of the drugs leveraged here ) . Furthermore , we build up on work proposing co-attention ( Lu et al. , 2016 ; Deac et al. , 2018 ) as a mechanism to allow for individual set-structured datasets ( such as nodes in multimodal graphs ) to interact . Specifically , such mechanisms have already been used for explicit matching of graph structure motifs ( Li et al. , 2019 ) , and therefore represent a natural methodology for our purposes . Overall , these ( and related ) techniques lie within the domain of graph representation learning , one of the latest major challenges of machine learning ( Bronstein et al. , 2017 ; Hamilton et al. , 2017 ; Battaglia et al. , 2018 ) , with transformative potential across a wide spectrum of potential applications , extending outside the biochemical domain . 2 ARCHITECTURE . In this section , we will present the main building blocks used within our architecture for paired graph representation learning . This will span a discussion of the way the input to the model is encoded , followed by an overview of the individual computational steps of the model . As the tasks considered in this manuscript operate over molecules , we will treat all inputs as molecular graphs . 2.1 INPUTS . The molecules , dx , are represented as graphs consisting of atoms , a ( dx ) i as nodes , and bonds between those atoms ( a ( dx ) i , a ( dx ) j ) as edges . For each atom , the following input features are recorded : the atom number , the number of hydrogen atoms attached to this atom , and the atomic charge . For each bond , a discrete bond type ( e.g . single , double etc . ) is encoded as a learnable input edge vector , e ( dx ) ij . 2.2 MESSAGE PASSING . Within each of the two molecules separately , our model applies a series of message passing ( Gilmer et al. , 2017 ) layers . Herein , nodes are allowed to send arbitrary vector messages to each other along the edges of the graph , and each node then aggregates all the messages sent to it . Let ( dx ) hti denote the features of atom i of molecule x , at time step t. Initially these are set to projected input features , i.e . : ( dx ) h0i = fi ( a ( dx ) i ) ( 1 ) where fi is a small multilayer perceptron ( MLP ) neural network . In all our experiments , this MLP ( and all subsequently referenced MLPs and projections ) projects its input to 50 features . Note that this effectively encodes one-hot inputs such as atom type into a learnable embedding . Considering atoms i and j of molecule x , connected by a bond with edge vector e ( dx ) ij , we start by computing the message , ( dx ) mtij , sent along the edge j → i . The message takes into account both the features of node j and the features of the edge between i and j : ( dx ) mtij = f t e ( e ( dx ) ij ) f tv ( ( dx ) ht−1j ) ( 2 ) where f te and f t v are small MLPs , and is elementwise vector-vector multiplication . Afterwards , every atom i aggregates the messages sent to it via summation . This results in an internal message for atom i , ( dx ) mti : ( dx ) mti = ∑ ∀j∈N ( i ) ( dx ) mtij ( 3 ) where the neighbourhood , N ( i ) , defines the set of atoms linked to i by an edge . 2.3 CO-ATTENTION . Message passing provides a robust mechanism for encoding within-molecule representations of atoms . For learning an appropriate joint molecule-molecule representation , however , we allow atoms to interact across molecule boundaries via a co-attentional mechanism ( Deac et al. , 2018 ) . Consider two atoms , i and j , of molecules x and y , respectively . Let their features at time step t be ( dx ) hti and ( dy ) htj , just as before . For every such pair , we compute an attentional coefficient , αtij using a simplified version of the Transformer ( Vaswani et al. , 2017 ) attention mechanism : αtij = softmaxj ( 〈 Wtk ( dx ) ht−1i , W t k ( dy ) ht−1j 〉 ) ( 4 ) where Wtk is a learnable projection matrix , 〈· , ·〉 is the inner product , and the softmax is taken across all nodes j from the second molecule . The coefficients αtij may be interpreted as the importance of atom j ’ s features to atom i . These coefficients are then used to compute an outer message for atom i of molecule x , ( dx ) nti , expressed as a linear combination ( weighted by α t ij ) of all ( projected ) atom features from molecule y : ( dx ) nti = ∑ ∀j∈dy αtij ·Wtv ( dy ) ht−1j ( 5 ) where Wtv is a learnable projection matrix . It was shown recently that multi-head attention can stabilise the learning process , as well as learn information at different conceptual levels ( Veličković et al. , 2018 ; Qiu et al. , 2018 ) . As such , the mechanism of Equation 5 is independently replicated across K attention heads ( we chose K = 8 ) , and the resulting vectors are concatenated and MLP-transformed to provide the final outer message for each atom . ( dx ) nti = f t o K ‖ k=1 ∑ ∀j∈dy ( k ) αtij · ( k ) Wtv ( dy ) ht−1j ( 6 ) where f to is a small MLP and ‖ is featurewise vector concatenation . Drug x Drug y j l ( dx ) mti k ( dx ) ntu ( dx ) mtil ( dx ) mtij ( dx ) mtik v αtuv Figure 1 : The illustration of a single step of message passing ( computing the inner message , ( dx ) mti ) , and co-attention ( computing the outer message ( dx ) ntu ) on two nodes ( i and u ) of molecule x . The computation of Equations 4–6 is replicated analogously for outer messages from molecule x to atoms of molecule y . We have also attempted to use other popular attention mechanisms for computing the αtij values—such as the original Transformer ( Vaswani et al. , 2017 ) , GAT-like ( Veličković et al. , 2018 ) and tanh ( Bahdanau et al. , 2014 ) attention—finding them to yield weaker performance than the approach outlined here . A single step of message passing and co-attention , as used by our model , is illustrated by Figure 1 . 2.4 UPDATE FUNCTION . Once the inner messages , ( dx ) mti ( obtained through message passing ) , as well as the outer messages , ( dx ) nti ( obtained through co-attention ) are computed for every atom i of each of the two molecules ( dx/dy ) , we use them to derive the next-level features , ( dx ) hti , for each atom . At each step , this is done by aggregating ( via summation ) the previous features ( representing a skip connection ( He et al. , 2016 ) ) , the inner messages and outer messages , followed by layer normalisation ( Ba et al. , 2016 ) : ( dx ) hti = LayerNorm ( ( dx ) ht−1i + ( dx ) mti + ( dx ) nti ) ( 7 ) The operations of Equations 2–7 are then repeated for T propagation steps—here , we set T = 3 . Refer to Figure 2 for a complete visualisation of our architecture . As will be demonstrated by our results , using co-attention to enable the model to propagate the information between two molecules from the beginning—thus learning a joint representation of the molecules—valuably contributes to the predictive power of the model .

Summary: In this paper, the authors proposed a method to extend graph-based learning with a co-attentional layer. Feeding graphs pairwisely into the model allows nodes to easily exchange information with nodes in other graphs as well as within the graph. This method outperforms other previous ones on a pairwise graph classification task (drug-drug interaction prediction).