How to transfer algorithmic reasoning knowledge to learn new algorithms?

1 Introduction . Transfer learning [ 2 ] has been responsible for significant successes in multiple areas of machine learning , including Natural Language Processing ( NLP ) [ 3 , 4 ] and Computer Vision ( CV ) [ 5 ] . Pretraining and reusing the learned weights , by freezing them as feature extractors , or fine-tuning from them as an initialisation , are common approaches to transfer in these domains . This has enabled successful learning on problems where data is limited in some form . Algorithmic reasoning on graphs [ 1 , 6 ] is a fundamental problem that has been studied under the assumption that we have access to the execution traces of the algorithms we want to learn . In practice , such as many real world reasoning tasks , this may not be true . Due to the limited data , we look to transfer learning to enable us to solve algorithmic tasks without intermediate steps . Specifically , we ∗Corresponding author 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) investigate how to transfer knowledge between similar graph algorithms in the setting where we have access to the execution traces for some ( e.g . PRIM [ 7 , 8 ] ) , but not others ( e.g . DIJKSTRA [ 9 ] ) . This has not been explored before , but is important as many reasoning related fields care deeply about systematic generalisation , while only having input-output pairs to learn from . One such example is knowledge graph link prediction . There exists a reasoning process that can answer the question : what is the tail entity given the head entity and the relation . However , we do not have access to its execution trace , i.e . the step-by-step deductions of the reasoning process , but we do have input-output pair examples . We envision that the set-up studied and the direction proposed in this paper will help to learn neural networks that can solve reasoning-style problems by using other reasoning knowledge as an inductive bias . A slightly different , but related application may be when the data the graph algorithm needs to operate on is encoded in a high-level space . This is a scenario encountered in reinforcement learning and concurrent work Deac et al . [ 10 ] has already started investigating the process of pre-training an encoder with graph algorithms . Veličković et al . [ 1 ] successfully trained graph neural networks ( GNNs ) [ 11 , 12 , 13 , 14 ] to execute graph algorithms . Two key ingredients were access to the intermediate steps of the algorithm and algorithmic alignment [ 15 ] . Algorithmic alignment [ 15 ] refers to the concept of a computation structure—in our case neural networks—and the structure of an algorithm ‘ aligning ’ —in this paper graph algorithms . ‘ Alignment ’ means there exists a mapping between substructures of the computation architecture and simple substeps of the algorithm , where the substructures can ‘ easily ’ compute the corresponding substep they are mapped to—for some definition of easy , usually linear . Extending this work , we enable solving tasks without access to the intermediate steps using other similar algorithms as an inductive bias . To do this we add several new algorithms . Further , since the algorithms studied in [ 1 ] were all expressible with only linear components in the GNN , we extend their architecture with a more expressive encoder to enable algorithmic alignment with more tasks . One assumption we make is that the algorithms are similar ( e.g . both PRIM and DIJKSTRA greedily select nodes from a priority queue ) . This shared algorithmic knowledge should serve as an inductive bias to enable a neural network to systematically generalise , even when learning only on input-output pairs . We hypothesise that due to differences between feature extraction and algorithmic reasoning , successful approaches to transfer as used in CV and NLP will provide only minimal improvements when transferring from one base algorithm to a target algorithm . While features across images may be similar , the features of algorithms differs ( e.g . shortest distance in DIJKSTRA and lightest edge to the tree in PRIM ) , but are processed in a conceptually similar manner . Our intuition is that this conceptual relationship may not yield weights that are near each other in the weight space , thus making transfer difficult . We instead propose to use the base algorithms as inductive biases by training them with the intermediate steps along side training the target algorithm , for which we do not provide intermediate supervision . Our second hypothesis is that this will help systematic generalisation . We validate both hypotheses—transfer does not help , while multi-task learning helps—empirically . The contributions of this paper are : 1. presenting a new benchmark for transferring algorithmic reasoning knowledge on graphs ; 2. sampling the best trajectory to stabilise training when no execution traces are available ; 3. show that standard transfer techniques fail to significantly improve systematic generalisation ; 4. demonstrate how systematic generalisation can instead be improved on algorithmic tasks with multi-task learning . 2 Related Work . Neural Execution : Many papers have studied how to execute algorithms before [ 16 , 17 , 18 , 19 , 20 ] . With the rise of graph representation learning [ 21 , 22 , 23 ] and graph neural networks [ 11 , 12 , 13 , 14 ] , recent work has looked at executing graph algorithms [ 1 , 6 ] . This line of work assumes access to the intermediate steps of the algorithm . Further , people have looked at teaching transformers to learn explicitly denoted subroutines and then learn how to combine them [ 24 ] . The key difference between our work and prior work is that we look at the setting of learning graph algorithms where we do not have access to the execution trace and we do not denote subroutines explicitly . Instead we propose to use similar algorithms as an inductive bias and look at how to enable this transfer of algorithmic knowledge . Transfer learning : Transfer learning was originally proposed in the 1970s [ 25 , 26 ] . Today , it is a common tool in NLP and Computer Vision to use pre-trained models and fine-tune them on the target task [ 2 , 3 , 5 ] . This is helpful when the target task is similar and the available data is limited in some way . Our work differs in two ways ; firstly , we focus on algorithmic reasoning knowledge , which we hypothesis will require different techniques than reusing weights . Secondly , the data on the target task is partially missing specifically the execution trace of the algorithm is missing . 3 Background . This section gives background on the algorithms and graphs used for the experiments . To be precise we consider a graph to be a tuple G = ( V , E ) with a set of vertices V and a set of edges E ⊆ V ×V 2 . This section explains the algorithms studied in this paper in more detail . After that , we give a brief introduction into graph neural networks , the standard architecture paradigm for graph inputs . 3.1 Algorithms . We broadly study two classes of algorithmic reasoning : parallel and sequential reasoning . Specifically , the parallel algorithms ( shown in Algorithm 1 ) studied in this paper exchange messages with their neighbours until an equilibrium is reached . Sequential algorithms , presented in Algorithm 2 , greedily remove elements from a priority queue and updating the neighbouring nodes ’ keys . Algorithm 1 Parallel Input : graph G , weights w , source index i initialise_nodes ( G.vertices , i ) repeat for ( u , v ) ∈ G.edges ( ) do relax_edge ( u , v , w ) end for until none of the nodes change Algorithm 2 Sequential Input : graph G , edge weights w , source node index i initialise_nodes ( G.vertices , i ) Q← PriorityQueue ( G.vertices ) repeat u← Q.pop_min ( ) for v ∈ G.neighbours ( u ) do relax_edge ( u , v , w ) end for until Q is empty Different algorithms will implement their own relax_edge and initialise_nodes functions , however , the overall framework stays the same and can be learned by our GNN architecture ( § 4.1 ) . For the sequential algorithms each node has a state feature indicating whether it has been removed from the priority queue , a key feature for the priority queue , and a pointer to the predecessor node . We study the following algorithms : PARALLEL 1 . BREADTH-FIRST SEARCH ( BFS ) 2 . BELLMAN-FORD 3 . WIDEST PATH ( PARALLEL ) 4 . MOST RELIABLE-PATH ( PARALLEL ) SEQUENTIAL 1 . PRIMS 2 . DIJKSTRA 3 . DEPTH-FIRST SEARCH ( DFS ) 4 . WIDEST PATH ( PARALLEL ) 5 . MOST RELIABLE-PATH ( PARALLEL ) We give the pseudo-code for all algorithms in the supplementary . Note that the relax_edge function is increasingly difficult to learn for the neural network as you go down the list : PRIMS uses the edge weight as a key , thus the relax_edge function is the identity , for DIJKSTRA and DEPTH-FIRST SEARCH addition is necessary to compute the edge update , for WIDEST-PATH we need to take a maximum , which with a ReLU activation can still be done exactly , and for MOST RELIABLE-PATH the neural network needs to learn to approximate multiplication.3 2Throughout the paper we use graphs that are simple with weights . 3We note that the tasks are all linear time in the size of the graph . More challenging problems such as NP-hard problems are primarily more difficult because of their increased run-time and hence longer sequences allowing 3.2 Graph Neural Networks . A general framework to describe several graph neural network architecture is the message passing framework introduced by Gilmer et al . [ 13 ] . Such graph neural networks ( GNNs ) consist of three parts : a message function M , an update function U , and an aggregator ⊕ . M and U are arbitrary neural networks . The high-level idea is 1. each node computes a message for its neighbours using M ; 2. then each node aggregates the received messages using ⊕ ; 3. finally , each node updates its embedding using U . The message function in this paper takes as input the sending and receiving nodes ’ representation as well as an encoding of the edge feature along which the message is sent . The aggregator chosen is the max aggregator applied element-wise as it was determined to work best by Veličković et al . [ 1 ] and provides algorithmic alignment [ 15 ] between the architecture and the tasks . The update function takes as input the node ’ s previous embedding and aggregated messages . 3.3 Problem definition . We study learning graph algorithms that take in a graph G , a weight function w : E → R that assigns a weight to each edge in the graph , and node features X : V → Rk . Algorithms compute an output Y : V → F k and a predecessor P : V → V at each time step . This encompasses a large class of graph algorithms solving tasks such as reachability or shortest-path . For instance , for a sequential algorithm the ith node ’ s features Xt [ i ] would be the key value and the state variable . Given Xt , the graph G , and the edge weights , the graph algorithm at each iteration returns the node features Yt and the predecessor for each Pt ( pred variable in pseudocode § 3.1 ) . YT and PT at the final step are considered the output of the algorithm . The intermediate steps refers to Xt , Yt , Pt ∀t ∈ { 1 , . . . , T − 1 } after each iteration .

Summary: The paper studies the problem of transferring algorithmic reasoning knowledge from a task to another similar task. The authors first propose a set of benchmark tasks to evaluate transfer learning for algorithmic reasoning on graphs.   The authors then propose modifications to the Neural Extractor model by replacing the node encoder with a latent edge encoder. Along with this, they propose using two latent encoders, one with a linear layer and one with a non-linear layer to allow the model to learn more types of tasks. Also, a separate encoder is used for each task.   The authors experiment with several transfer-learning and Multi-task learning (MTL) approaches on the proposed benchmark tasks.  The authors conclude that standard methods for transfer learning are ineffective on transferring algorithmic reasoning knowledge. They further show that the MTL approach effectively transfers knowledge thus improving systematic generalization.