Turing complete computation and reasoning are often regarded as necessary pre- cursors to general intelligence. There has been a significant body of work studying neural networks that mimic general computation, but these networks fail to generalize to data distributions that are outside of their training set. We study this problem through the lens of fundamental computer science problems: sorting and graph processing. We modify the masking mechanism of a transformer in order to allow them to implement rudimentary functions with strong generalization. We call this model the Neural Execution Engine, and show that it learns, through supervision, to numerically compute the basic subroutines comprising these algorithms with near perfect accuracy. Moreover, it retains this level of accuracy while generalizing to unseen data and long sequences outside of the training distribution. 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). We propose neural execution engines (NEEs), which leverage a learned mask and supervised execution traces to mimic the functionality of subroutines. We demonstrate that while algorithms like sorting are challenging to generalizably learn from input/output examples, we can identify smaller, simpler subroutines that transformers can learn with near-perfect strong generalization. While the functionality of these subroutines is currently limited, future work can expand the complexity of each subroutine that an NEE learns, getting closer to end-to-end learning of complex algorithms with neural networks. There are many natural extensions, including: extending the usage of binary to the transformer outputs, teaching learned masks more complex pointer primitives, using reinforcement learning to replace the supervision provided by teacher forcing, linking the generation of these models to source code, and exploring the link between the learned subroutines of NEE-like architectures and conventional von-Neumann computers, which execute individually encoded instructions sequentially (Von Neumann, 1993) .
We propose neural execution engines (NEEs), which leverage a learned mask and supervised execution traces to mimic the functionality of subroutines and demonstrate strong generalization. This paper investigates a problem of building a program execution engine with neural networks and proposes a transformer-based model to learn basic subroutines and applies them in several standard algorithms. This paper deals with the problem of designing neural network architectures that can learn and implement general programs.