Dynamically Pruned Message Passing Networks for Large-scale Knowledge Graph Reasoning

1 INTRODUCTION . Modern deep learning systems should bring in explicit reasoning modeling to complement their black-box models , where reasoning takes a step-by-step form about organizing facts to yield new knowledge and finally draw a conclusion . Particularly , we rely on graph-structured representation to model reasoning by manipulating nodes and edges where semantic entities or relations can be explicitly represented ( Battaglia et al. , 2018 ) . Here , we choose knowledge graph scenarios to study reasoning where semantics have been defined on nodes and edges . For example , in knowledge base completion tasks , each edge is represented by a triple 〈head , rel , tail〉 that contains two entities and their relation . The goal is to predict which entity might be a tail given query 〈head , rel , ? 〉 . Existing models can be categorized into embedding-based and path-based model families . The embedding-based ( Bordes et al. , 2013 ; Sun et al. , 2018 ; Lacroix et al. , 2018 ) often achieves a high score by fitting data using various neural network techniques but lacks interpretability . The pathbased ( Xiong et al. , 2017 ; Das et al. , 2018 ; Shen et al. , 2018 ; Wang , 2018 ) attempts to construct an explanatory path to model an iterative decision-making process using reinforcement learning and recurrent networks . A question is : can we construct structured explanations other than a path to better explain reasoning in graph context . To this end , we propose to learn a dynamically induced subgraph which starts with a head node and ends with a predicted tail node as shown in Figure 1 . Graph reasoning can be powered by Graph Neural Networks . Graph reasoning needs to learn about entities , relations , and their composing rules to manipulate structured knowledge and produce structured explanations . Graph Neural Networks ( GNNs ) provide such structured computation and also inherit powerful data-fitting capacity from deep neural networks ( Scarselli et al. , 2009 ; Battaglia et al. , 2018 ) . Specifically , GNNs follow a neighborhood aggregation scheme to recursively aggregate information from neighbors to update node states . After T iterations , each node can carry structure information from its T -hop neighborhood ( Gilmer et al. , 2017 ; Xu et al. , 2018a ) . GNNs need graphical attention expression to interpret . Neighborhood attention operation is a popular way to implement attention mechanism on graphs ( Velickovic et al. , 2018 ; Hoshen , 2017 ) by focusing on specific interactions with neighbors . Here , we propose a new graphical attention mechanism not only for computation but also for interpretation . We present three considerations when constructing attention-induced subgraphs : ( 1 ) given a subgraph , we first attend within it to select a few nodes and then attend over those nodes ’ neighborhood for next expansion ; ( 2 ) we propagate attention across steps to capture long-term dependency ; ( 3 ) our attention mechanism models reasoning process explicitly through pipeline disentangled from underlying representation computing . GNNs need input-dependent pruning to scale . GNNs are notorious for their poor scalability . Consider one message passing iteration on a graph with |V | nodes and |E| edges . Even if the graph is sparse , the complexity of O ( |E| ) is still problematic on large graphs with millions of nodes and edges . Besides , mini-batch based training with batch size B and high dimensions D would lead to O ( BD|E| ) making things worse . However , we can avoid this situation by learning input-dependent pruning to run computation on dynamical graphs , as an input query often triggers a small fraction of the entire graph so that it is wasteful to perform computation over the full graph for each input . Cognitive intuition of the consciousness prior . Bengio ( 2017 ) brought the notion of attentive awareness from cognitive science into deep learning in his consciousness prior proposal . He pointed out a process of disentangling high-level factors from full underlying representation to form a lowdimensional combination through attention mechanism . He proposed to use two recurrent neural networks ( RNNs ) to encode two types of state : unconscious state represented by a high-dimensional vector before attention and conscious state by a derived low-dimensional vector after attention . We use two GNNs instead to encode such states on nodes . We construct input-dependent subgraphs to run message passing efficiently , and also run full message passing over the entire graph to acquire features beyond a local view constrained by subgraphs . We apply attention mechanism between the two GNNs , where the bottom runs before attention , called Inattentive GNN ( IGNN ) , and the above runs on each attention-induced subgraph , called Attentive GNN ( AGNN ) . IGNN provides representation computed on the full graph for AGNN . AGNN reinforces representation within a cohesive group of nodes to produce sharp semantics . Experimental results show that our model attains very competitive scores on HITS @ 1,3 and the mean reciprocal rank ( MRR ) compared to the best embedding-based method so far . More importantly , we provide explanations while they do not . 2 ADDRESSING THE SCALE-UP PROBLEM . Notation . We denote training data by { ( xi , yi ) } Ni=1 . We denote a full graph by G = 〈V , E〉 with relations R and an input-dependent subgraph by G ( x ) = 〈VG ( x ) , EG ( x ) 〉 which is an induced subgraph of G. We denote boundary of a graph by ∂G where V∂G = N ( VG ) − VG and N ( VG ) means neighbors of nodes in VG . We also denote high-order boundaries such as ∂2G where V∂2G = N ( N ( VG ) ) ∪ N ( VG ) − VG . Trainable parameters include node embeddings { ev } v∈V , relation embeddings { er } r∈R , and weights used in two GNNs and an attention module . When performing full or pruned message passing , node and relation embeddings will be indexed according to the operated graph , denoted by θG or θG ( x ) . For IGNN , we use Ht of size |V| ×D to denote node hidden states at step t ; for AGNN , we use Ht ( x ) of size |VG ( x ) | × D to denote . The objective is written as ∑N i=1 l ( xi , yi ; θG ( xi ) , θG ) , where G ( xi ) is dynamically constructed . The scale-up problem in GNNs . First , we write the full message passing in IGNN as Ht = fIGNN ( Ht−1 ; θG ) , ( 1 ) where fIGNN represents all involved operations in one message passing iteration over G , including : ( 1 ) computing messages along each edge with the complexity1 of O ( BD|E| ) , ( 2 ) aggregating messages received at each node with O ( BD|E| ) , and ( 3 ) updating node states with O ( BD|V| ) . For T -step propagation , the per-batch complexity is O ( BDT ( |E|+ |V| ) ) . Considering that backpropagation requires intermediate computation results to be saved during one pass , this complexity counts for both time and space . However , since IGNN is input-agnostic , node representations can be shared across inputs in one batch so that we can remove B to get O ( DT ( |E| + |V| ) ) . If we use a sampled edge set Ê from E such that |Ê | ≈ k|V| , the complexity can be further reduced to O ( DT |V| ) . The pruned message passing in AGNN can be written as Ht ( x ) = fAGNN ( H t−1 ( x ) , Ht ; θG ( x ) ) . ( 2 ) Its complexity can be computed similarly as above . However , we can not remove B. Fortunately , subgraph G ( x ) is not G. If we let x be a node v , G ( x ) grows from a single node , i.e. , G0 ( x ) = { v } , and expands itself each step , leading to a sequence of ( G0 ( x ) , G1 ( x ) , . . . , GT ( x ) ) . Here , we describe the expansion behavior as consecutive expansion , which means no jumping across neighborhood allowed , so that we can ensure that Gt ( x ) ⊆ Gt−1 ( x ) ∪ ∂Gt−1 ( x ) ⊆ Gt−2 ( x ) ∪ ∂2Gt−2 ( x ) . ( 3 ) Many real-world graphs follow the small-world pattern , and the six degrees of separation implies G0 ( x ) ∪ ∂6G0 ( x ) ≈ G. The upper bound of Gt ( x ) can grow exponentially in t , and there is no guarantee that Gt ( x ) will not explode . Proposition . Given a graph G ( undirected or directed in both directions ) , we assume the probability of the degree of an arbitrary node being less than or equal to d is larger than p , i.e. , P ( deg ( v ) ≤ d ) > p , ∀v ∈ V . Considering a sequence of consecutively expanding subgraphs ( G0 , G1 , . . . , GT ) , starting with G0 = { v } , for all t ≥ 1 , we can ensure P ( |VGt | ≤ d ( d− 1 ) t − 2 d− 2 ) > p d ( d−1 ) t−1−2 d−2 . ( 4 ) The proposition implies the guarantee of upper-bounding |VGt ( x ) | becomes exponentially looser and weaker as t gets larger even if the given assumption has a small d and a large p ( close to 1 ) . We define graph increment at step t as ∆Gt ( x ) such that Gt ( x ) = Gt−1 ( x ) ∪ ∆Gt ( x ) . To prevent Gt ( x ) from explosion , we need to constrain ∆Gt ( x ) . Sampling strategies . A simple but effective way to handle the large scale is to do sampling . 1 . ∆Gt ( x ) = ∂̂Gt−1 ( x ) , where we sample nodes from the boundary of Gt−1 ( x ) . 2 . ∆Gt ( x ) = ∂Ĝt−1 ( x ) , where we take the boundary of sampled nodes from Gt−1 ( x ) . 3 . ∆Gt ( x ) = ∂̂Ĝt−1 ( x ) , where we sample nodes twice from Gt−1 ( x ) and from ∂Ĝt−1 ( x ) . 4 . ∆Gt ( x ) = ∂̂Ĝt−1 ( x ) , where we sample nodes three times with the last from ∂̂Ĝt−1 ( x ) . Obviously , we have ∂̂Ĝt−1 ( x ) ⊆ ∂̂Ĝt−1 ( x ) ⊆ ∂Ĝt−1 ( x ) and Gt−1 ( x ) ∪ ∂Ĝt−1 ( x ) ⊆ Gt−1 ( x ) ∪ ∂Gt−1 ( x ) . Further , we let N1 and N3 be the maximum number of sampled nodes in ∂Ĝt−1 ( x ) and the last sampling of ∂̂Ĝt−1 ( x ) respectively and let N2 be per-node maximum sampled neighbors in ∂̂Gt−1 ( x ) , and then we can obtain much tighter guarantee as follow : 1We assume per-example per-edge per-dimension time cost as a unit time . 1 . P ( |V∆Gt ( x ) | ≤ N1 ( d− 1 ) ) > pN1 for ∂Ĝt−1 ( x ) . 2 . P ( |V∆Gt ( x ) | ≤ N1N2 ) = 1 and P ( |V∆Gt ( x ) | ≤ N1 ·min ( d−1 , N2 ) ) > pN1 for ∂̂Ĝt−1 ( x ) . 3 . P ( |V∆Gt ( x ) | ≤ min ( N1N2 , N3 ) ) = 1 for ∂̂Ĝt−1 ( x ) . Attention strategies . Although we guarantee |VGT ( x ) | ≤ 1 + T min ( N1N2 , N3 ) by ∂̂Ĝt−1 ( x ) and constrain the growth of Gt−1 ( x ) by decreasing either N1N2 or N3 , smaller sampling size means less area explored and less chance to hit target nodes . To make efficient selection rather than random sampling , we apply attention mechanism to do the top-K selection where K can be small . We change ∂̂Ĝt−1 ( x ) to ∂̂Gt−1 ( x ) ̃ where ∼ represents the operation of attending over nodes and picking the top-K . There are two types of attention operations , one applied to Gt−1 ( x ) and the other applied to ∂̂Gt−1 ( x ) ̃ . Note that the size of ∂̂Gt−1 ( x ) ̃ might be much larger if we intend to sample more nodes with larger N2 to sufficiently explore the boundary . Nevertheless , we can address this problem by using smaller dimensions to compute attention , since attention on each node is a scalar requiring a smaller capacity compared to node representation vectors computed in message passing .

Summary: The authors propose to use sampling methods in order to apply graph neural networks to large scale knowledge graphs for semantic reasoning. To this end induced subgraphs are constructed in a data-dependent way using an attention mechanism. This improves the efficiency and leads to interpretable results. The experiments show some improvements over path- and embedding-based methods.