Neural Execution of Graph Algorithms

1 INTRODUCTION . A multitude of important real-world tasks can be formulated as tasks over graph-structured inputs , such as navigation , web search , protein folding , and game-playing . Theoretical computer science has successfully discovered effective and highly influential algorithms for many of these tasks . But many problems are still considered intractable from this perspective . Machine learning approaches have been applied to many of these classic tasks , from tasks with known polynomial time algorithms such as shortest paths ( Graves et al. , 2016 ; Xu et al. , 2019 ) and sorting ( Reed & De Freitas , 2015 ) , to intractable tasks such as travelling salesman ( Vinyals et al. , 2015 ; Bello et al. , 2016 ; Kool et al. , 2018 ) , boolean satisfiability ( Selsam et al. , 2018 ; Selsam & Bjørner , 2019 ) , and even probabilistic inference ( Yoon et al. , 2018 ) . Recently , this work often relies on advancements in graph representation learning ( Bronstein et al. , 2017 ; Hamilton et al. , 2017 ; Battaglia et al. , 2018 ) with graph neural networks ( GNNs ) ( Li et al. , 2015 ; Kipf & Welling , 2016 ; Gilmer et al. , 2017 ; Veličković et al. , 2018 ) . In almost all cases so far , ground-truth solutions are used to drive learning , giving the model complete freedom to find a mapping from raw inputs to such solution1 . Many classical algorithms share related subroutines : for example , shortest path computation ( via the Bellman-Ford ( Bellman , 1958 ) algorithm ) and breadth-first search both must enumerate sets of edges adjacent to a particular node . Inspired by previous work on the more general tasks of program synthesis and learning to execute ( Zaremba & Sutskever , 2014 ; Kaiser & Sutskever , 2015 ; Kurach et al. , 2015 ; Reed & De Freitas , 2015 ; Santoro et al. , 2018 ) , we show that by learning several algorithms simultaneously and providing a supervision signal , our neural network is able to demonstrate positive knowledge transfer between learning different algorithms . The supervision signal is driven by how a known classical algorithm would process such inputs ( including any relevant intermediate outputs ) , providing explicit ( and reusable ) guidance on how to tackle graph-structured problems . We call this approach neural graph algorithm execution . ∗Work performed while the author was at DeepMind . 1We note that there exist good reasons for choosing this approach , e.g . ease of optimisation . Given that the majority of popular algorithms requires making discrete decisions over neighbourhoods ( e.g . “ which edge should be taken ? ” ) , we suggest that a highly suitable architecture for this task is a message-passing neural network ( Gilmer et al. , 2017 ) with a maximisation aggregator—a claim we verify , demonstrating clear performance benefits for simultaneously learning breadth-first search for reachability with the Bellman-Ford algorithm for shortest paths . We also verify its applicability to sequential reasoning , through learning Prim ’ s algorithm ( Prim , 1957 ) for minimum spanning trees . Note that our approach complements Reed & De Freitas ( 2015 ) : we show that a relatively simple graph neural network architecture is able to learn and algorithmically transfer among different tasks , do not require explicitly denoting subroutines , and tackle tasks with superlinear time complexity . 2 PROBLEM SETUP . 2.1 GRAPH ALGORITHMS . We consider graphs G = ( V , E ) where V is the set of nodes ( or vertices ) and E is the set of edges ( pairs of nodes ) . We will consider graphs for two purposes : 1 ) as part of the task to be solved ( e.g. , the graph provided as input to breadth first search ) , 2 ) as the input to a graph neural network . A graph neural network receives a sequence of T ∈ N graph-structured inputs . For each element of the sequence , we will use a fixed G and vary meta-data associated with the nodes and edges of the graph in the input . In particular , to provide a graph G = ( V , E ) as input to a graph neural network , each node i ∈ V has associated node features ~x ( t ) i ∈ RNx where t ∈ { 1 , . . . , T } denotes the index in the input sequence and Nx is the dimensionality of the node features . Similarly , each edge ( i , j ) ∈ E has associated edge features ~e ( t ) ij ∈ RNe where Ne is the dimensionality of the edge features . At each step , the algorithm produces node-level outputs ~y ( t ) i ∈ RNy . Some of these outputs may then be reused as inputs on the next step ; i.e. , ~x ( t+1 ) i may contain some elements of ~y ( t ) i . 2.2 LEARNING TO EXECUTE GRAPH ALGORITHMS . We are interested in learning a graph neural network that can execute one or more of several potential algorithms . The specific algorithm to be executed , denoted A , is provided as an input to the network . The structure of the graph neural network follows the encode-process-decode paradigm ( Hamrick et al. , 2018 ) . First we will describe the encode-process-decode architecture and then describe the specific parameterisations that are used for each sub-network . The relation between the variables used by the graph algorithm and our neural algorithm executor setup is further illustrated in Figure 1 . For each algorithm A we define an encoder network fA . It is applied to the current input features and previous latent features ~h ( t−1 ) i ( with ~h ( 0 ) i = ~0 ) to produce encoded inputs , ~z ( t ) i , as such : ~z ( t ) i = fA ( ~x ( t ) i , ~h ( t−1 ) i ) ( 1 ) The encoded inputs are then processed using the processor network P . The processor network shares its parameters among all algorithms being learnt . The processor network takes as input the encoded inputs Z ( t ) = { ~zti } i∈V and edge features E ( t ) = { ~e ( t ) ij } e∈E and produces as output latent node features , H ( t ) = { ~hti ∈ RK } i∈V : H ( t ) = P ( Z ( t ) , E ( t ) ) ( 2 ) The node and algorithm specific outputs are then calculated by the decoder-network , gA : ~y ( t ) i = gA ( ~z ( t ) i , ~h ( t ) i ) ( 3 ) Note that the processor network also needs to make a decision on whether to terminate the algorithm . This is performed by an algorithm-specific termination network , TA , which provides the probability of termination τ ( t ) —after applying the logistic sigmoid activation σ—as follows : τ ( t ) = σ ( TA ( H ( t ) , H ( t ) ) ) ( 4 ) where H ( t ) = 1|V | ∑ i∈V ~h ( t ) i is the average node embedding . If the algorithm hasn ’ t terminated ( e.g . τ ( t ) > 0.5 ) the computation of Eqns . 1–4 is repeated—with parts of ~y ( t ) i potentially reused in ~x ( t+1 ) i . In our experiments , all algorithm-dependent networks , fA , gA and TA , are all linear projections , placing the majority of the representational power of our method in the processor network P . As we would like the processor network to be mindful of the structural properties of the input , we employ a graph neural network ( GNN ) layer capable of exploiting edge features as P . Specifically , we compare graph attention networks ( GATs ) ( Veličković et al. , 2018 ) . ( Equation 5 , left ) against message-passing neural networks ( MPNNs ) ( Gilmer et al. , 2017 ) ( Equation 5 , right ) : ~h ( t ) i = ReLU ∑ ( j , i ) ∈E a ( ~z ( t ) i , ~z ( t ) j , ~e ( t ) ij ) W~z ( t ) j ~h ( t ) i = U ~z ( t ) i , ⊕ ( j , i ) ∈E M ( ~z ( t ) i , ~z ( t ) j , ~e ( t ) ij ) ( 5 ) where W is a learnable projection matrix , a is an attention mechanism producing scalar coefficients , while M and U are neural networks producing vector messages . ⊕ is an elementwise aggregation operator , such as maximisation , summation or averaging . We use linear projections for M and U . Note that the processor network P is algorithm-agnostic , and can hence be used to execute several algorithms simultaneously . Lastly , we note that the setup also easily allows for including edge-level outputs , and graph-level inputs and outputs—however , these were not required in our experiments . 3 EXPERIMENTAL SETUP . Graph generation To provide our learner with a wide variety of input graph structure types , we follow prior work ( You et al. , 2018 ; 2019 ) and generate undirected graphs from seven categories : • Ladder graphs ; • 2D grid graphs ; • Trees , uniformly randomly generated from the Prüfer sequence ; • Erdős-Rényi ( Erdős & Rényi , 1960 ) graphs , with edge probability min ( log2 |V | |V | , 0.5 ) ; • Barabási-Albert ( Albert & Barabási , 2002 ) graphs , attaching either four or five edges to every incoming node ; • 4-Community graphs—first generating four disjoint Erdős-Rényi graphs with edge probability 0.7 , followed by interconnecting their nodes with edge probability 0.01 ; • 4-Caveman ( Watts , 1999 ) graphs , having each of their intra-clique edges removed with probability 0.7 , followed by inserting 0.025|V | additional shortcut edges between cliques . We additionally insert a self-edge to every node in the graphs , in order to support easier retention of self-information through message passing . Finally , we attach a real-valued weight to every edge , drawn uniformly from the range [ 0.2 , 1 ] . These weight values serve as the sole edge features , e ( t ) ij , for all steps t. Note that sampling edge weights in this manner essentially guarantees the uniqueness of the recovered solution , simplifying downstream evaluation . We also ignore corner-case inputs ( such as negative weight cycles ) , leaving their handling to future work . We aim to study the algorithm execution task from a “ programmer ” perspective : human experts may manually inspect only relatively small graphs , and any algorithms derived from this should apply to arbitrarily large graphs . For each category , we generate 100 training and 5 validation graphs of only 20 nodes . For testing , 5 additional graphs of 20 , 50 and 100 nodes are generated per category . Parallel algorithms We consider two classical algorithms : breadth-first search for reachability , and the Bellman-Ford algorithm ( Bellman , 1958 ) for shortest paths . The former maintains a single-bit value in each node , determining whether said node is reachable from a source node , and the latter maintains a scalar value in each node , representing its distance from the source node . In both cases , the algorithm is initialised by randomly selecting the source node , s. As the initial input to the algorithms , x ( 1 ) i , we have : BFS : x ( 1 ) i = { 1 i = s 0 i 6= s Bellman-Ford : x ( 1 ) i = { 0 i = s +∞ i 6= s ( 6 ) This information is then propagated according to the chosen algorithm : a node becomes reachable from s if any of its neighbours are reachable from s , and we may update the distance to a given node as the minimal way to reach any of its neighbours , then taking the connecting edge : BFS : x ( t+1 ) i = 1 x ( t ) i = 1 1 ∃j . ( j , i ) ∈ E ∧ x ( t ) j = 1 0 otherwise B-F : x ( t+1 ) i = min ( ~x ( t ) i , min ( j , i ) ∈E x ( t ) j + e ( t ) ji ) ( 7 ) For breadth-first search , no additional information is being computed , hence y ( t ) i = x ( t+1 ) i . Additionally , at each step the Bellman-Ford algorithm may compute , for each node , the “ predecessor ” node , p ( t ) i in the shortest path ( indicating which edge should be taken to reach this node ) . This information is ultimately used to reconstruct shortest paths , and hence represents the crucial output : Bellman-Ford : p ( t ) i = i i = s argmin j ; ( j , i ) ∈E x ( t ) j + e ( t ) ji i 6= s ( 8 ) Hence , for Bellman-Ford , ~y ( t ) i = p ( t ) i ‖x ( t+1 ) i , where ‖ is concatenation . To provide a numerically stable value for +∞ , we set all such entries to the length of the longest shortest path in the graph + 1 . We learn to execute these two algorithms simultaneously—at each step , concatenating the relevant ~x ( t ) i and ~y ( t ) i values for them . As both of the algorithms considered here ( and most others ) rely on discrete decisions over neighbourhoods , learning to execute them should be naturally suited for the MPNN with the max-aggregator—a claim which we directly verify in the remainder of this section . Sequential algorithms Unlike the previous two algorithms , single iterations of many classical graph algorithms will specifically focus on one node at a time—very often the case with constructive tasks . We seek to demonstrate that our neural graph algorithm execution paradigm aligns well with this setting too , and in this context we study Prim ’ s algorithm ( Prim , 1957 ) for minimum spanning trees . Prim ’ s algorithm maintains a partially constructed minimum spanning tree ( MST ) —initially , it is a singleton tree consisting of only a source node , s. At each step , Prim ’ s algorithm searches for a new node to connect to the MST—chosen so that the edge attaching it to the tree is the lightest possible : Prim : x ( 1 ) i = { 1 i = s 0 i 6= s Prim : x ( t+1 ) i = 1 x ( t ) i = 1 1 i = argmin j s.t . x ( t ) j =0 min k s.t . x ( t ) k =1 e ( t ) jk 0 otherwise ( 9 ) Once the new node is selected , the algorithm attaches it to the MST via this edge—similarly to Bellman-Ford , we can keep track of predecessor nodes , p ( t ) i : Prim : p ( t ) i = i i = s p ( t−1 ) i i 6= s ∧ x ( t ) i = 1 argmin j s.t . x ( t ) j =1 e ( t ) ij x ( t ) i = 0 ∧ x ( t+1 ) i = 1 ⊥ ( undefined ) otherwise ( 10 ) In Equations 9–10 , the boxed updates are the only modifications to the state of the algorithm at step t— centering only on the selected node to attach to the MST . Once again , the algorithm requires discrete decisions based on neighbourhood edge weights , hence we expect outperformance of MPNN-max . For a visualisation of the expected alignment between a graph algorithm and our neural graph executors , refer to Figure 2 . We provide an overview of all of the inputs and supervision signals used here for the executor in Appendix A. Neural network architectures To assess the comparative benefits of different architectures for the neural algorithm execution task , we consider many candidate networks executing the computation of Equations 1–5 , especially the processor network P : For the MPNN update rule , we consider maximisation , mean and summation aggregators . For the GAT update rule , we consider the originally proposed attention mechanism of Veličković et al . ( 2018 ) , as well as Transformer attention ( Vaswani et al. , 2017 ) ; Additionally for GAT , we consider also attending over the full graph—adding a second attention head , only acting on the non-edges of the graph ( and hence not accepting any edge features ) . The two heads ’ features are then concatenated and passed through another linear layer . Analogously to our expectation that the best-performing MPNN rule will perform maximisation , we attempt to force the attentional coefficients of GAT to be as sharp as possible—applying either an entropy penalty to them ( as in Ying et al . ( 2018 ) ) or the Gumbel softmax trick ( Jang et al. , 2016 ) . We perform an additional sanity check to ensure that a GNN-like architecture is necessary in this case . Prior work ( Xu et al. , 2019 ) has already demonstrated the unsuitability of MLPs for reasoning tasks like these , and they will not support variable amounts of nodes . Here , instead , we consider an LSTM ( Hochreiter & Schmidhuber , 1997 ) architecture into which serialised graphs are fed ( we use an edge list , in a setup similar to ( Graves et al. , 2016 ) ) . In all cases , the neural networks compute a latent dimension of K = 32 features , and are optimised using the Adam SGD optimiser ( Kingma & Ba , 2014 ) on the binary cross-entropy for the reachability predictions , mean squared error for the distance predictions , categorical cross-entropy for the predecessor node predictions , and binary cross-entropy for predicting termination ( all applied simultaneously ) . We use an initial learning rate of 0.0005 , and perform early stopping on the validation accuracy for the predecessor node ( with a patience of 10 epochs ) . If the termination network TA does not terminate the neural network computation within |V | steps , it is assumed terminated at that point . For Prim ’ s algorithm , as only one node at a time is updated , we optimise the categorical cross-entropy of predicting the next node , masked across all the nodes not added to the MST yet . It should be noted that , when learning to execute Bellman-Ford and Prim ’ s algorithms , the prediction of p ( t ) i is performed by scoring each node-pair using an edge-wise scoring network ( a neural network predicting a scalar score from ~h ( t ) i ‖~h ( t ) j ‖~e ( t ) ij ) , followed by a softmax over all neighbours of i .

Summary: The paper suggests training neural network to imitate graph algorithms in a more fine-grained way than done before: by learning primitives and subroutines rather than the final output. The paper makes quite a strong case for the advantage of this approach, citing the fact that many graph algorithms share subroutines, which could simplify learning and support joint training and transfer learning. The experiments are detailed an elaborate.