NEURAL EXECUTION ENGINES

1 INTRODUCTION . Neural networks are universal function approximators ( Hornik et al. , 1989 ) , meaning that provided enough data and perfect optimization they should be able to learn arbitrarily complicated functions . In recent years , there have been proposals of neural network architectures that are designed to implement general programs ( Graves et al. , 2014 ; Kaiser & Sutskever , 2016 ; Graves et al. , 2016 ; Kurach et al. , 2016 ) , often inspired by concepts found in conventional computer systems , like pointers ( Vinyals et al. , 2015 ) . However , these neural networks still have difficulty learning complex programs from input/output pairs , in the sense of strong generalization . That is , generalizing to data distributions that do not necessarily correspond to the training distribution , such as longer sequences and new values . We hypothesize that much of this difficulty stems from a lack of prior structure , and given enough structure in the form of supervision over intermediate program states , we can train networks to faithfully implement different algorithms . We take several basic algorithms ( selection sort , merge sort , Dijkstra ’ s algorithm for shortest paths ) and express them in terms of a series of subroutines , as a software engineer would . Each subroutine represents a simple function , and can be composed with others to express the algorithm . In this way , we train neural networks to perform relatively simple tasks in a supervised manner , and obtain complex behaviors through composition . Although each subroutine represents a simple task compared to the full algorithm , this is nevertheless a challenging learning domain for several reasons . First , each subroutine still requires the network to learn a function in such a way that it can strongly generalize outside of its training distribution . Next , as the goal is to learn general computation , the network will operate on raw numbers : taking as input numbers , or distributions over sets of numbers that it may not have even seen in training . Lastly , each subroutine must be performed accurately enough so that composition results in accurate inference over long runs of the program . Our main contribution is to show that while a model trained on a complex task in an end-to-end fashion may generalize well on in-distribution test data , this does not necessarily lead to strong generalization . However , the same underlying architecture can be made to strongly generalize by introducing minor modifications and more supervision . This provides a starting point for gradually reducing the amount of required supervision and increasing the sizes of the learned subroutines in order to work towards end-to-end learning of complex algorithms with neural networks . Specifically , we leverage the transformer ( Vaswani et al. , 2017 ) to learn the subroutines underlying several common yet sophisticated algorithms from input/output execution traces . Our model uses binary number representations for data values , and separates the notion of control ( which part of the input to consider ) from execution ( what to compute ) via a conditional masking mechanism . We show that with this , transformers can learn effective representations for accurately performing fundamental numeric tasks like comparison and addition , and that allowing the transformer to modulate its own mask in subsequent subroutine calls allows it to generalize to runs of the program that greatly exceed the length of the traces it was trained on , resulting in near perfect performance on larger tasks . We refer to these networks over subroutines as neural execution engines ( NEEs ) . 2 BACKGROUND . 2.1 TRANSFORMERS . Transformers are a family of models that represent the current state-of-the-art in sequence learning ( Vaswani et al. , 2017 ; Devlin et al. , 2018 ; Radford et al. , 2019 ) . In contrast to recurrent networks that process inputs sequentially , transformers process an entire sequence simultaneously using a self-attention mechanism that allows each input token to attend to every other token in the sequence when computing its transformation . This mechanism occurs within an attention block , and several blocks are stacked to form both an encoder to embed an input sequence , and a decoder to produce an output sequence . Here , we will formalize some of the components of the vanilla transformer architecture so that we can outline our modifications in subsequent sections . Given input token sequences x1 , x2 , . . . , xL1 ∈ X and output token sequences y1 , y2 , . . . , yL2 ∈ Y , where xi , yj ∈ Z+ , a transformer learns a mapping X → Y . First , the tokens are individually embedded to form x̂i , ŷj ∈ Rd . The main module of the transformer architecture is the self-attention layer , which outputs a transformation of each vector as a convex combination of values , modulated by the affinity of the vector to every other vector in the sequence . Note that for this paper , we do not use positional encodings ( which we found to hurt performance in our tasks ) , and single-headed attention . Self-attention layers are followed by a position-wise feed-forward neural network layer , forming a self-attention block . These blocks are composed to form the encoder and decoder of the transformer , with the outputs of the encoder being used as queries and keys for the decoder . More details can be found in Vaswani et al . ( 2017 ) . An important component for our purposes is the self-attention mask . This is used to prevent certain positions from propagating information to other positions . Typically this is used for decoding , to ensure that the model can only condition on past outputs during sequential generation . In our case , we will consider masking in the encoder as an explicit way for the model to condition on the part of the sequence that it needs at a given point in its computation . 2.2 THE SORTING TASK . We use sorting to frame our exploration into the capability of neural networks to mimic general execution . Sorting is a useful abstraction to examine because it is clearly defined at a high level : the sorting task involves mapping a list of unsorted numbers to its sorted counterpart . Generalization ( performance on unseen sequences ) is clearly distinct from strong generalization ( performance on sequences that are longer that the training distribution ) . Additionally , there are many different algorithms that computer scientists traditionally use to implement sorting . This provides a rich domain of interest in which we can use supervision to learn different sorting algorithms . 2.3 NUMBER REPRESENTATIONS . An essential component of general computation is manipulating numbers ( Von Neumann , 1993 ) . Given that our goal is strong generalization , it is also necessary for the number system that we use to work in large ranges and generalize outside of its training domain ( as it is not tractable to train the network on all integers ) . Neural networks generally use either categorical , one-hot , or integer number representations . Prior work has found that scalar numbers have difficulty representing large ranges ( Trask et al. , 2018 ) and that binary numbers are a useful representation that generalize outside of their training domain ( Kaiser & Sutskever , 2016 ; Shi et al. , 2019 ) , assuming each bit is toggled during training . Binary is a hierarchical format , where each bit extends the training range of the string exponentially , leading to a small bit string that can represent massive number ranges . In this paper , we focus on unsigned integers using the binary representation . 3 NEURAL EXECUTION ENGINES . Simply stated , a neural execution engine ( NEE ) is a transformer that takes as input binary numbers and an encoding mask , and outputs either data values , a pointer , or both . The pointer can optionally be used to modify the mask the next time the NEE is invoked . This architecture is reminiscent of graph attention networks ( GATs ) ( Veličković et al. , 2017 ) , where the graph is used to dictate the encoder mask . An NEE is essentially a GAT that can modify its own graph . In this section , we will describe the specifics behind how we apply this mechanism , before demonstrating its performance on several algorithmic tasks in subsequent sections . 3.1 METHODOLOGY . We now describe the key components of the NEE mechanism . Bitwise Embeddings . As input to an NEE , we embed binary vectors using a linear projection . This is equivalent to defining a learnable vector for each bit position , and then summing these vectors elementwise , modulated by the value of their corresponding bit . That is , given an embedding vector vi for each bit i , for an n-bit input vector x , we would compute x̂ = ∑n i=1 xivi . For example , emb ( 1001 ) = v0+v3 . Two important tokens for our purposes are start s and end e. These are commonly used in natural language data to denote the start and end of a sequence . We use s as input to the decoder , and e to denote the end of an input sequence . As we are dealing with numerical computation , we find it convenient to set s = 0 and to have the model learn an embedding vector for e such that e =∞ . That is , the model will learn e > x for all x 6= e and that e+ x = e. Conditional Masking . The encoder of a transformer takes as input both a set of values and a mask , which is used to force the encoder to ignore certain inputs . In the NEE mechanism , we use the output pointer of the decoder to modify the mask for a subsequent call of the encoder . In this way , the inputs represent a memory state , and the mask represents a set of pointers into the memory . NEE effectively learns where to focus its attention for performing computation . A mask M is a binary matrix , where a value of 0 indicates that the input should be considered , and 1 indicates that it should be ignored . The encoder mask is used to zero out attention weights of ignored input numbers . Given a binary vector b , a mask is obtained by broadcasting b along the sequence length dimension and other higher dimensions . The value bi = 0 would indicate that the ith input should be considered , and bi = 1 would indicate that the ith input should be ignored . We will refer to b as the mask vector . Given an input mask vector b , we need to compute the updated mask vector b̂ to serve as the next input mask . As most entries of the mask do not change with each iteration , we output only the information required to update the mask . Practically , this means a one-hot encoding o representing the position of interest to mask out . We then compute the updated mask b̂ using a bitwise XOR function , b̂ = XOR ( b , o ) . This is sufficient for the subroutines considered in this paper , but the transformer can output arbitrary masks , and future work will focus on allowing the network to make more complex updates . Using this mechanism , we can train NEE to mimic the behavior of a variety of subroutines by training it on execution traces using teacher forcing ( Williams & Zipser , 1989 ) . At inference , we simply choose the argmax of the pointer output head . Encoding Independent Inputs . We will show examples where the function must operate on a sequence of inputs , but where subsets of the inputs can be treated independently . Specifically , an elementwise sum or min between pairs of numbers . For transformers , this can be done by an explicit reshape operation that turns the sequence into a mini-batch , enabling parallel processing on the entire batch . A subsequent reshape can turn the output back into a sequence . Encoding Dependent Inputs . When we receive a sequence of inputs where subsequences are dependent , we concatenate the subsequences and delineate them using the e token . For example , if we had two subsequences [ a , b ] and [ c , d ] , we would input [ a , b , e , c , d , e ] . For the corresponding mask , we could give [ 0 , 0 , 1 , 0 , 0 , 1 ] to consider all tokens , or [ 0 , 1 , 1 , 0 , 1 , 1 ] if we intend to iterate over the subsequences in order . If required , NEE will output a pointer , and we can use this to update the mask as needed . This mechanism will be useful in merge sort , where a fundamental subroutine is merging two sorted lists into a single sorted list . Decoding . Decoding in a NEE involves giving the start token s to the decoder and using it to compute attention blocks over the encoder outputs . In this case , we only consider outputting a single value and/or pointer , as opposed to a sequence .

Summary: his paper deals with the problem of designing neural network architectures that can learn and implement general programs.  The authors are motivated by problems in such works, mainly generalization to testing distributions that do not necessarily correspond to the training distribution.  It should be made clear that the latter is a general problem in machine learning (with phenomena such as covariate shift being commonplace), the authors particularly relate this work to predicting new values and longer sequences (i.e. strong generalization).