NEURAL EXECUTION ENGINES
Authors: Yujun Yan
Kevin Swersky
Danai Koutra
Parthasarathy Ranganathan
Milad Hashemi
Venue:  (2020)
Keywords: neural computation
strong generalization
numerical reasoning

Abstract: 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.

Under review as a conference paper at ICLR 2020
NEURAL EXECUTION ENGINES
Anonymous authors
Paper under double-blind review
ABSTRACT
Turing complete computation and reasoning are often regarded as necessary pre- cursors to general intelligence. There has been significant body of work studying neural networks that mimic general computation, but these networks fail to general- ize 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 transformer in order to allow them to implement rudimentary functions with strong generalization. We call this model the Neural Execution Engine, and show that 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.
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 al. [0166] Kurach et al. 2016), often inspired by concepts found in conventional computer systems, like pointers (Vinyals) al. 2015). However, these neural networks still have difficulty learning complex programs from onput/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 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 's algorithm for shortest paths) and express them terms of series of subroutines, as a software engineer would. Each subroutine represents 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 supervised manner, and obtain complex behaviors through composition.
Although each subroutine represents a simple task compared to the full algorithm, this is nevertheless 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 show that while model trained on 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 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 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 conditional masking mechanism. We show that with this,

---

Under review as a conference paper at ICLR 2020
transformers 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 family of models that represent the current state-of-the-arr in sequence Iearn- ing (Vaswani et al. 2017 Devlin et al. 2018 Radford et al. 2019). In contrast to recurrent networks that process inputs sequentially, transtormers process entire sequence simultaneously using 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 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, 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 E X' and output token sequences y1, ................................ YL E E,, where X,, Yj E z+, a transformer learns a mapping X' + V. First, the tokens are individually embedded to form &, E Idd. 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 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 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 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 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 useful abstraction to examine because it is clearly defined at high level: the sorting task involves mapping list of unsorted numbers to its sorted counterpart. Generalization (performance on unseen sequences) is clearly distinct from strong generulization (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 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 strong generalization, 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.
NEURAL EXECUTION ENGINES
Simply stated, neural execution engine (NEE) is transformer that takes input binary numbers and an encoding mask, and outputs either data values, a pointer, or both. The pointer can optionally be used modify the mask the next time the NEE is invoked. This architecture is reminiscent of
2

---

Under review as a conference paper at ICLR 2020
graph attention networks (GATs) (eelickvvi et al. 2017). where the graph is used 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 apply this mechanism, before demonstrating its performance on several algorithmic tasks in subsequent sections.
3.1 METHODOL.OGY
We now describe the key components of the NEE mechanism.
Bitwise Embeddings, As imput to an NEE, we embed binary vectors using a linear projection. This is equivalent 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 v, for each bit i, for an n-bit input vector x, we would compute Σ VyV; For example, emb(1001) = U0+
Two important tokens for our purposes are start and end These commonly used in natural language data to denote the start and end of sequence. We use as input to the decoder, and to denote the end of an input sequence. As we are dealing with numerical computation, we find convenient to set 8 = O and to have the model learn an embedding vector for such that co. That is, the model will learn > x for all x o e and that e + =
Conditional Masking, The encoder of a transformer takes as input both a set of values and 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 memory state, and the mask represents set of pointers into the memory. NEE effectively learns where to focus its attention for performing computation.
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 b, = 0 would indicate that the 2th input should be considered, and hy- = would indicate that the 7t input should be ignored. We will refer to b the mask vector.
Given an input mask vector b, we need to compute the updated mask vector 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 using 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 sequence of inputs where subsequences are dependent, we concatenate the subsequences and delineate them using the token. For example, 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, O, O, to consider all tokens, or [0, 1, 1, 1] if we intend to iterate over the subsequences in order. 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 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, opposed sequence.
3

---

Under review as a conference paper at ICLR 2020
4 LEARNING TO SORT
As mentioned in Section 2 sorting frame exploration into the capability of neural networks to mimic general execution.
4.1 A SEQUENCE TO SEQUENCE VIEW
We first study how well the vanilla transformer (Section 2.1 Iearns to sort. We model it in a conven- tional seq2seq fashion (Sutskever et al. 2014) with input examples of unsorted sequences (length L) of n--it binary numbers (L=8 and n=8) and output examples of the correctly sorted sequence. The in- puts to the encoder and decoder are sequences of bitwise embedded vec- tors (Section]) 2.3). The decoder uses greedy decoding strategy. The outputs of the decoder are one-hot 256-dimensional vectors representing the unnormalized log- probabilities of the output numbers. We find that randomly generated numbers easier sort than numbers that differ in a small number(c.g.. 1 vs. 2, 53 vs. 54). Thus, we include both examples in the training distribution (70% random numbers, 30% numbers with small differences). [ The performance of this vanilla transformer, evaluated as achieving an exact content and positional match the correct output example, is shown in Figure (where the test distribution consists of 60% random numbers and 40% numbers with small differences).
Modelss
seg. m
Accuracy 0
100
all modifcaatins
90.00%
moditications we symmetnio attemiio
Accuracy (96)
98.339
60
moditication but with dot product artn
B7 7..59%6
40
-1. modifialiioss
a1. woddPPaatoooseeee
vanilia trand formor
98.67%
modefcations s0co,
one-hot embediingg
varilla transformes
93.11%
sirila ooe hot emt
20
80 100
(cotauit NLPP)
Length Test Sequences
Figure I: Seq2Seq sorting performance of transformer variants trained on sequences of up to length
To increase performance, we make number of modifications better tune the encoder process binary numbers. These modifications are driven by the unsatisfactory performance of the vanilla transformer in distinguishing close numbers (Appendix [A.2] Figure [cc). To aid in the ability of the number to learn numerical similarity in carry-out cases when many bits flip (32 vs. 31, 64 vs. 63), we use bitwise embeddings and replace the dot product self-attention with a symmetric feed-forward self-attention (Figure 11c). To help handle small differences in general, we use shared instead of independent linear projections queries, keys and values. To aid reconstructing small numerical differences, we scale up the residual connection strength by factor of 1.5. These modifications over the original transformer model are shown in the Appendix (Figure 10b).
Figure D shows the sorting performance of five models: (1) the model with all the modifications mentioned above, (2) the model with all modifications but with non-symmetric attention, (3) the model with all modifications but with dot-product attention (instead of symmetric feed-forward self- attention). (4) the vanilla transformer model with bitwise embedding, and (5) the vanilla transformer model with one-hot embedding (which is default NLP). The results of more variants and data mixes are shown in the Appendix (Figure IT). We observe that while the transformer variants all have reasonably high performance in sorting lists of numbers that are S 8 elements long from input/output examples, the models fail to generalize to longer sequences, and performance sharply drops as the sequence length increases.
To understand why performance degrades as the test sequences get longer, we plot the attention matrix of the last layer in the decoder (Figure 2h). The model accurately attends to the first few numbers in the sequence (distinct dots in the chart) but the attention distribution becomes "fuzzy" as the sequence length increases beyond numbers, often resulting in the same num- ber being repeatedly predicted. Next, seek to investigate the following question: Can we retain high attention resolution (and correspond- ing model accuracy) by restricting the model to only observe the first scenario repeatedly?
20
40
60
80
100
20
40
60
100
20
40
60
80 100
Fuzzy Attention (seq2seq)
(b) Clear Attention
(Selection Sort NEE)
Figure 2: Visualizing the decoder attention weights sorting transformers. Attention over each row. The at- tention in the seq2seq transformer saturates as the output sequence length increases, while NEE maintains sharp attention weights.
Throughout this work, the preponderance of that regenerated numbers off by small differences.
4

---

Under review as a conference paper at ICLR 2020
4.2 SUBROUTINE LEVEL TRACES: SELECTION SORT
Using this insight, we break the problem up into pieces and increase the amount of supervision provided during training. Instead of learning to perform the entire computation from only input/output examples, pick simple sorting algorithm emulate (selection sort in Figure 33 and train the transformer to learn one action that can be composed in loop to solve the task-much as a train NEE
programmer would while writing code. Our guide for selecting these subroutines on is to determine the largest source code function to learn, where the output of the function is data-dependent. That is, static parameters (like the XOR reshape in Section 4.3 are not learned. For the case of selection sort, the task is to find the minimum of the current list.
Since the transformer had dif- ficulty clearly attending to val- ues beyond the training sequence length, we restrict the locations in the unsorted sequence where the transformer could possibly attend in every iteration. This is accomplished by producing conditional mask to mask out the data elements that have al- ready been appended to the sorted_l st and feeding that mask back into the transformer (shown the right side of Fig- ure B). This modification sepa- rates the control from the com- putation (which elements should be considered) from the computation itself (find the minimum value the list). This allows the transformer to learn output logits of much larger magnitude, resulting in sharper attention, as shown in Figure 2p. Our experimental results consequently demonstrate strong generalization, perfectly sorting sequences of up to length 100, shown in Figure 12a (Appendix A.2
solection sdat(a)
sorted 1ist while (Len(date]>0)1-
Jain Indax, aii__eianett... (ind_mia((ataa)
data.deleta Index)
Neural Execution Engine
sorted 11st (mneeddeendd
Engine
aorted
find min(data)
min elenent
min index -1
for index, element enumerate(dsta)) if [element nin elenent): nin nin elanent elament
sppend()
learned
_in_index- index
[min index, min elenent]
sorted 1ist
Figure 3: Selection sort code translated to NEE architecture. The colored box (f (fina iinddd.. mn) learned from input/output examples, but potential code implementation shown.
4.3 RECURSIVE SUBROUTINES: MERGE SORT
Given the success of an NEE repeatedly performing an action (finding the minimum of list) in order implement selection sort, we explore how well the approach generalizes to other sorting algorithms. As programmer could write smaller subroutines to accomplish the sorting task in many different ways, NEEs should be able mimic these subroutines from execution level traces. We therefore select very different sorting algorithm for study, merge sort. The code for one implementation of merge-sort shown in Figure
data end-start
This code is broadly broken up into two subroutines, data decom- position (merge sort) and action (merge). Every call to me merge.sor divides the list half until there is one element left, which by definition is already sorted. Then, merq qe unrolls the recursive tree, combining every 2 elements (then every 4, 8, etc.) until the list is fully sorted. Recursive algorithms like merge sort generally consist of these two steps (the "recursive case" and the "base case").
merge_ merge_sort(dee port(ddaaa start, and): ead):
reshapa()
(start end)
Neural Execution Engine
merge_ sortidata, start, mid) mmrge sarro__sort(daaaaa mid and)
Imerge
mid, endi
XOE
partially sorted data learnad Boak
Figure 4: Merge code translated NEE architecture.
The recursive case is commonly based around data movement, and we use a reshape() operation to divide combine lists every time and emulate the start and mid pointers in the application (this a static parameter). It is possible to use the learned mask the transformer to learn the recursive decomposition serially (e.g., have the mask shift between the first pair of numbers, then the second
5

---

Under review as a conference paper at ICLR 2020
pair, etc.), but the reshape function allows us teach the transformer to swap all pairs of data for every level the tree in parallel.
More concretely, the transformer in Figure impleeents the computation described in Figure 5 The mask controls which numbers are seen by the encoder (I meaning masked). Every timestep, the model outputs the smallest number from the unmasked numbers and the position of the currently selected number. This is used to produce the next mask with a SHIFT and a XOR. We do not learn these operations since they are static and don't depend on data. The model stops when it outputs an end token. F shows the computation for one leaf of the recursive decomposition, but the reshape allows the transformer to compute on all leaf nodes a given level in the in parallel. Given this model, Figure 6 demonstrates that the NEE is able to learn merge sort with perfect strong generalization over long sequences (up to length 100). The NEE was trained to perform the computation underlying merge sort on sequences of length - 8.
Input Seq
3 1
8
8
3
B e 3
Accuracy (96)
0
XOR
0
1
XOR
10
100
Transformer
Transformer
90
75 100
Length of Test Sequences
3
Pair 0
SHIFT
0
SHIFT
0100
RETURN
Figure 6: NEE merge sort: strong generalization over long sequences. Training involved quences length S 8.
Figure 5: NEE merge sort trace.
5 COMPOSABLE SUBROUTINES, SHORTEST PATH
While both merge sort and selection sort demonstrated that an NEE can compose subroutines repeatedly sort a list with perfect accuracy, programmers often need to compose multiple different subroutines to achieve larger goals. In this section, we study whether multiple NEEs can be composed execute a more complicated algorithm.
To that end, we study a graph algorithm, Dijkstra's algorithm to find shortest paths, shown on the left of Figure 7 which mirrored in the source code on the left. The algorithm consists of four major steps: (1) Initialization: set the distance from the source node to the other nodes to infinity (end in Section 3. then append them into queue structure for processing: (2) Compute newly found path from the source node to all neighbors of the selected node; (3) Update path lengths if they are smaller than the stored lengths; (4) Select he node with the smallest distance the source node and remove from the queue. The algorithm repeats steps (2) as long there are elements in the queue.
Computing Dijkstra's algorithm requires the NEEs to learn the three corresponding subroutines, as shown on the right in Figure 7 Finding the min- imum between the possible_paths
shortest _acls__pa_______ patb(graph,
shorteat_
greph.adj chortest_paca
discs nodes [] anchor node
Neural Execution Engine
node node_13st 13se graph.get_
while list: - [possible pachs graph.adj (anchor node),
possible
batha
Execution Engine
shortest path)
Ishortest clementuise possible peths,
shortest pathll
t whortest path
and shortest.path as well as the minimum current shortest_path can be accomplished by the modified transformer trained to accomplish the same goal for sorting in Section 4.2
fanchor,
min(shortest path)
Exscution Engine
sode List. anchor node)
dists. appand append (win, d(at)
append anchor_ node)
anokor_node _ode
nin dist
append[] )
dista,
Figure Dijkstra code translated to architecture.
The new challenges are to learn addition (Section 5.I and to compose these learned models sequentially to achieve shortest path in generalizable way (Section 5.2)

---

Under review as a conference paper at ICLR 2020
5.1 LEARNING ADDITION
The addition subroutine in Dijkstra's algorithm iteratively adds all of the elements in the adjacency list of selected node to the shortest path that is currently known. To learn this function, we train an NEE to learn the sum of two numbers, then apply it to Dijkstra's algorithm.
Given the input/output examples for this subroutine, we find that addition is more challenging task than sorting. Zero and infinity (end token, Sec. 2.3 are important concepts for the model learn. Our bitwise embedding already represents as a directly learnable vector. We use a moving average on the parameters of the neural network in evaluation to reduce the noise in weights (Kushner & Yang 1995 Polyak & Juditsky 1992), and with this, our modified transformer architecture is capable of learning addition to 100% accuracy.
the zero vector. We find it useful to embed infinity
To understand the number system that the NEE has learned, visualize the structure of the learned embeddings using PCA to 3 dimensions, and compare the embeddings learned from addition and sorting, shown in Figure 88 and Figure respectively. Each node is colored based the number represents. We find that a highly structured number system has been learned, and the addition embeddings are rigorously arranged such that the sequence of numbers that increase by 1 are placed on a line (shown with arrows in Figure Sp). The sorting embeddings exhibit many small clusters and the numbers in each cluster are related. More detailed visualizations are provided in Appendix AA.3 AA.3
0.5
-0.5
200
-1.5
1.5
02
-1.5
-0.5
0.5
1.5
64
22
100
58
H64
20
60
+64
62
+64
Inf
+1 per arrow
Inf
204
208
(a) Addition: The embeddings a sequence of numbers that increase by are placed on a line.
(b) Sorting: The numbers per cluster are related according to the pattern of gray arrows (bottom). Figure 8: Visualization of learned addition and sorting embeddings.
In addition testing unseen pairs of numbers in the addition task, unseen numbers to further test the generalization of our model. As shown in Table [ even with half the data range held out of training, the model achieves high accuracy. This study corroborates prior work that shows that the model only needs to observe the bit being toggled to generalize to the numerical range represented by any given bit (Shi al., 2019). This a promising result as suggests we may be able to extend the framework to much larger bit vectors, where observing every number in training intractable.
also test completely
Table I: Performance for different hold-out % of 8-bit number ranges.
Accuracy (5))
Figure 9: NEE Dijkstra: Training on graphs with < 10 nodes generalizes per- fectly to < 40-node graphs
100
Hold-out & 6.25%
12.5%
25%
50%
Accuracy 99.96%
100%
99.2%
98.72%
30 Sizea of Test Graphs
20
5.2 RESULTS
Given that NEEs can be trained to perform the necessary subroutines underlying Dijkstra's 's algorithm, we demonstrate that they can then be composed to reliably compute the overall algorithm and generalize to much larger graphs. A step-by-step view of the operations that the NEEs perform shown in Figure [5] in Appendix A..4. The examples are Erdös-Rént Rényi random graphs (since the types of graphs do not matter). We train on graphs with up 8 nodes and test on graphs up to 40 nodes, with 100 graphs evaluated at each size. Weights are randomly assigned within the allowed 8-bit number range. We evaluate the prediction accuracy on the final output (the shortest path of all nodes to the source nodes) and achieve 100% test accuracy with graph sizes up to 40 nodes (Figure 9).
7

---

Under review as a conference paper at ICLR 2020
6 RELATED WORK
There have been a number of recent proposals for neural networks that learn to perform computation that are inspired by computer architectures (Graves et al., 2014: 2016 Kaiser & Sutskever, 2016 Kurach et al. 2016 Vinyals et al. 2015). These attempt to learn complex algorithms purely from weak supervision, L.e., input/output pairs. Through universality, these architectures are able to theoretically represent any computable function, however practically they tend to have issues generalizing to longer sequences. They are typically trained on scalar data values in limited ranges, or focus purely on pointer arithmetic. Our focus is on learning simple, composable functions on raw numerical values in way that will generalize to extremely long sequences that have not been seen during training. Furthermore, instead of developing a custom architecture, we focus on using transformers (Vaswani 2017) with very little modification, leveraging the encoding mask to allow the transformer to focus its capacity on relevant input locations.
al.
Neural program synthesis (Neelakantan et al., 2016 Reed De Freitas 2016 Parisotto et al. 2017,
Devlin et al.] [217]]
(117/Car Cai et 2017 Bunel et al. 2018) uses neural networks to generate programs
that are capable of solving tasks represented by input/output pairs, with the goal of finding a "correct"
program such that it will generalize beyond the training distribution. Related to this area, there is
body of research at the intersection of logic programming and neural networks (Evans & Grefenstette 2018 Dong ct al. 2019), that uses neural networks to learn set of rules entailing only positive examples from a set of positive and negative examples. Inputs are treated as separate objects, and specific rules are given to the system for the neural network to use. For example, in learning to sort, Dong et al. (2019) give their system all of the numbers involved as one-hot encoded vectors, along their ground-truth comparative relationships, whereas we focus on learning these relationships from binary numbers. Similar to our own work, |Reed & De Freitas (2016); (hi et al.] (2019); Cai et al. (2017) use execution traces as a form of stronger supervision. Cai et al. (2017) uses execution traces with tail recursion (where the subroutines call themselves), albeit for program synthesis, and show that this leads to improved generalization.
There have been several works on using neural networks to learn number systems for performing arithmetic. Paccanaro & Hinton (2001) directly embeds integers in the range [-10, 10) as vectors and trains these, along with matrices representing relationships between objects. They train these on triplets involving two objects and a relation, to perform modular arithmetic, and show that these embeddings and relationships are capable of learning coherent structure and generalizing to new triplets. [Sutskever & Hinton (2009) expand on this idea, modeling objects as matrices as well, so that relationships can equivalently be treated as objects, allowing the system to learn higher-orde relationships. Trask et al. (2018) explores the (poor) generalization capability of neural networks on values inputs outside of their training range, and develops new architectures that better suited to scalar arithmetic, improving extrapolation. Several papers have used neural networks to learn binary arithmetic with varying degrees of success (Joulin & Mikolov 2015 Kaiser & Sutskever 2016). Recently, (Shi et al. (2019) showed that graph neural networks are capable of learning directly from 64-bit binary memory states from execution traces of assembly code, and that this representation generalizes better when predicting program behavior outside of the training range. Finally, Wallace jet al. (2019) explored the ability of various language models to perform numerical reasoning after being trained on question/answer dataset; they found that language models exhibit high degree of numeracy, especially when using character-lovel embeddings.
7 CONCLUSION
We propose neural execution engines (NEEs), which leverage 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- Non-Neumann computers, which execute individually encoded instructions sequentially (Von Neumann 1993).
8

---

Under review as a conference paper at ICLR 2020
REFERENCES
Dzmitry Bahdanau. Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning align and translate. International Conference on Learning Representations, 2015.
Rudy Bunel, Matthew Hausknecht, Jacob Devlin, Rishabh Singh, and Pushmeet Kohli. Leveraging grammar and reinforcement learning for neural program synthesis. International Conference Learning Representations, 2018.
Jonathon Cai, Richard Shin, and Dawn Song. Making neural programming architectures generalize via recursion. International Conference on Learning Representations, 2017.
Jacob Devlin, Jonathan Uesato, Surya Bhupatiraju, Rishabh Singh, Abdel-rahman Mohamed, and Pushmeet Kohli. Robustfill: Neural program learning under noisy i/o. International Conference on Machine Learning, 2017.
Jacob Devlin, Ming- bidirectional transformers for language understanding. arXiv preprint / 2018.
Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep
Honghua Dong, Jiayuan Mao, Tian Lin, Chong Wang. Lihong Li, and Denny Zhou. Neural logic
machines. International Conference on Learning Representations, 2019.
Richard Evans and Edward Grefenstette. Learning explanatory rules from noisy data. Journal of Artificial Intelligence Research, 61:1-64, 2018.
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv preprint arXiv: 1410.5401 2014.
Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwińska, Sergio Gómez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, et al. Hybrid computing using neural network with dynamic external memory. Nature, 47776))):471, 2016.
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural networks, 2(5):359-366, 1989.
Armand Joulin and Tomas Mikolov. Inferring algorithmic patterns with stack- suuck-aggmetedd recurrent nets. In Advances in neural information processing systems, pp. 190-198, 2015.
Lukasz Kaiser and Ilya Sutskever. Neural gpus learn algorithms. International Conference on Learning Representations, 2016.
Karol Kurach, Marcin Andrychowicz, and Ilya Sutskever. Neural random-access machines. Interna- tional Conference on Learning Representations, 2016.
Harold Kushner and Jichuan Yang. Stochastic approximation with averaging and feedback: Rapidly convergent" on-line" algorithms. IEEE Transactions on Automatic Control, 40(1):24 34, 1995.
Arvind Neelakantan, Quoc Le, and Ilya Sutskever. Neural programmer: Inducing latent programs with gradient descent. International Conference on Learning Representations, 2016.
Alberto Paccanaro and Geoffrey Hinton. Learning distributed representations of concepts using linear relational embedding. IEEE Transactions on Knowledge and Data Engineering, 13(2): 232-244, 2001.
Emilio Parisotto, Abdel-rahman Mohamed, Rishabh Singh, Lihong Li, Dengyong Zhou, and Pushmeet Kohli. Neuro-symbolic program synthesis. International Conference on Learning Representations, 2017.
Boris T Polyak and Anatoli B Juditsky. Acceleration of stochastic approximation by averaging. SIAM Journal on Control and Optimization, 30(4):838-85) 855, 1992.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models unsupervised multitask learners. 2019.

---

Under review as a conference paper at ICLR 2020
Scott Reed and Nando De Freitas. Neural programmer-imterprenn International Conference Learning Representations, 2016.
Zhan Shi, Kevin Swersky, Daniel Tarlow, Parthasarathy Ranganathan, and Milad Hashemi. Learning execution through neural code fusion. arXiv preprint arXiv: /906. 07181, 2019.
llya Sutskever and Geoffrey E Hinton. Using matrices to model symbolic relationship. In Advances in neural information processing systems, pp. 1593-1600, 2009.
lya Sutskever, Oriol Vinyals, and Quoe V Le. Sequence to sequence learning with neural networks. In Advances neural information processing systems, pp. 3104-3112, 2014.
Andrew Trask, Felix Hill, Scott Reed, Jack Rae, Chris Dyer, and Phil Blunsom. Neural arithmetic logic units. In Advances in Neural Information Processing Systems, 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural information processing systems, 2017.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint arXiv: 2017.
Oriol Vinyals, Meire Fortunato, and Navdeep Jaitly. Pointer networks. In Advances in Neural Information Processing Systems, pp. 2692-2700, 2015.
John Von Neumann. First draft of a report on the edvac. IEEE Annals of the History of Computing, 15(4)227-75, 1993.
Eric Wallace, Yizhong Wang, Sujian Li, Sameer Singh, and Matt Gardner. Do nlp models know numbers? probing numeracy in embeddings. arXiv preprint arXIv:1909.07940, 2019.
Ronald J Williams and David Zipser. A learning algorithm for continually running fully recurrent neural networks. Neural computation, 1(2):270-280, 1989.
10

---

Under review as a conference paper at ICLR 2020
A APPENDIX
A.1 HYPERPARAMETERS
For sorting, we found it sufficient to use bitwise embeddings of dimension d 16. For addition, which is more difficult task, we found it necessary to increase this to d = 24. We use positional encoding for the sorting task, and single-headed attention for all tasks. In the addition task, we use exponential moving average of the parameters during learning under two time-scale stochastic approximation (Polyak Juditsky 1992) for additional stability in the learned model. The remaining transformer hyperparameters, aside from the changes described next, are set their defaults.
A.2 SORTING ABLATIONS
In this section, we show the result of our architecture modifications to performance in learning to sort. We study two scenarios: the first is a seq2seq setup, where the NEE must learn to sort the entire sequence in encoding/greedy decoding step (making it equivalent to transformer). The second when implementing the selection sort algorithm, where the NEE must output the min value and location, and where the updated mask based this is fed back into the NEE for subsequent iterations, until the e token is returned. Note that the second scenario has much more supervision than the first scenario, which must learn an entire sorting algorithm internally.
Within each scenario, we study 3 different data distributions: the first is where we train on purely random sequences with tokens in [0, 255) U {e). The second is mixed setting, where 60% of the examples are drawn randomly, and 40% are drawn from more difficult distribution, where the numbers are closer in value. The third purely difficult setting, where all sequences have numbers that are close to each other in value.
We ablate specific architectural changes in these settings. The original and modified encoder are represented visually in in Figure I0 Specifically, the architectural choices we test are:
Scaling up the strength of the residual connections by factor of 1.5.
Using an MLP-based attention module (Bahdanau et al. 2015).
Symmetrizing the MLP-based attention by flipping the order of the inputs and averaging the
resulting logit values.
Using the standard scaled dot product attention.
Using binary encoding of input values.
Using one-hot encoding of the input values.
Using a binary encoding as the input, but without any linear embedding.
Sharing the bitwise embedding projection between the query, key, and value in the attention mechanism.
Add & Norm
Add & Norm
Scaled Residual Connection
Feed-Forward
Feed-Forward
Add & Norm
Add Norm
Self Self-Attention Layer
Feed-Forward Self-attention
Positional Encoding
Positional Encoding
(optional)
Input Embedding
(optional)
Bitwise Embedding
(a) Vanilla encoder
(b) Modified encoder
Figure Bascline transformer (a) and modified transformer (b).
11

---

Under review as a conference paper at ICLR 2020
The test accuracy, measured as getting all values and their positions correct, sequences of length 8 in the mixed setting is shown in Table 2
We can see that the architectural changes help improve performance on these sequences up to near-perfect accuracy.
Table 2: Seq2Seq performance for transformer variants at training length of mixed test sets.
Models
Accuracy @ seq_len
all modifications
99.00%
all modifications..o.. wo-sstooowwwwoymmmmmm.
98.33%
all wwmoiiiicatinnseepp..rr
95.89%
modifications excp shared proj
98.44%
all modifications excp dot att
97.56%
all mwl.mmiicatiios..oo wot.sooooochooheeeeeeeee
89.56%
all www.iiiicaionssbiineee caiis_biirr__reee _emb
84.78%
vanilla
96.67%
Comilla_binaryccom _emb
77.11% %
Vanilla one hot emb
93.1 %
In Figure I1 we show the strong generalization performance of the different architectures. While some changes are able to improve performance in this regime, the performance ultimately drops steeply as the length of the test sequence increases. This is consistent across all test scenarios.
We next run the same ablation, but this time with selection sort using NEEs (of different variants) to perform the Find_min operation, and with feedback via conditional masking; the results are shown in Figure 12, In this case, most architectural choices perform well, and strongly generalize to much longer sequences than they were trained on. This shows that NEE is quite robust to the choice of architecture. The exceptions important for encoding number similarity. Note that the raw bits as inputs still work much better than it does in seq2seq setting regardless in short or long sequences.
using raw binary inputs, showing that bitwise embedding is
Here we list out some random and hard examples as well as the corresponding output (containing some errors) from the vanilla transformer with one-hot encoded input numbers (each number has an independent embedding), which is commonly used in natural language models. The symbol represents the end token.
be seen that the model makes mistakes in hard examples.
Random examples: 100 62 14 66 241 1 63 237 181 52 71 254 246 145 118 28
Output from Vanilla one hot.emb:
1 62 63 66 100 114 237 253 e 28 52 71 118 145 181 246 254
Hard examples:
132 126 131 129 127 130 128 125 e 238 239 241 240 243 237 242 244
Output from Vanilla one hot emb:
125 126 127 128 129 130 132
237 238 240 244 243 237 242
12

---

Under review as a conference paper at ICLR 2020
al modifications
al modilications. _ecp__ excp_dot_as dot_tt
vanila
all htt_moddfiiiatin__eeeel cecpeon__________ee______ swcp_shared__prg
medifications excp all modificationn. ooooon_hhh
vanilla hot_emmeb
al medifications__w wo_ vanilla vanilla_birary _mb
all
www.mmoocctiio..ceeoo hwntttion_nnnarrr..oooo emb
100
80
Accuracy (%)
60
40
20
8
20
40
60
80
100
Length of Test Sequences
(a) Mixed test sets
100
80
Accuracy (%)
60
40
20
0
8
40
100
Length of Test Sequences
(b) Random test sets
100
80
Accuracy (%)
60
40
20
8
20 Length of Test Sequences (c) Hard test sets
40
60
80
100
Figure 11: Seq2Seq strong generalization performance (without mask feedback to encoder) on mixed test consist of 60% random examples and 40% hard examples where the numbers close each other. (b) Random test sets. (c) hard test sets, where test sets consist 100% hard examples where the numbers are close to each other. All models trained on sequences < and tested up to length 100. Vanilla corresponds to the original transformer, with bitwise embeddings.
sets, where
13

---

Under review as a conference paper at ICLR 2020
100
modifaatio
mod Poation moco
Accuracy (%)
80
modoootoon wa Syft
Festtre, soco Lat
60
Ven cow. hot amb
31. mod Reation, binory, emb
po
40
20
0
20
30
40
50
60
70
80
90
100
Length of Test Sequences
(a) Mixed test sets
100
modiication
modrication
Accuracy (%)
80
al modifcationa. sym
al. modiications,
modiiccaiion,, one hot emb
vanila binary
60
vanila
vanila one hot anb
al. modiicatione, binary. emb
40
a1. mmodiiationn aww.adaardddo. proj
20
8
20
30
40
50
60
70
80
90
100
Length of Test Sequences
(b) Random test sets
100
all modifiatiooss
all modoooattionss.. exop_res
Accuracy (%)
80
al modifictions we
ml.mmoiizzzions
modiiicctinn
varila_binary_om
60
one hot,
al al_mmdiiiation birary_emb al. [_Looiiiaaions,. eep_t__eop____eee excp_shared _rared_____prr
40
20
20
30
40
50
60
70
80
90
100
Length of Test Sequences
(c) Hard test sets
Figure 12: Selection sort strong generalization performance (with mask feedback the encoder) on mixed sets, where test consist 60% random examples 40% hard examples where the numbers are close to each other. (b) Random test sets. (c) hard test sets, where test sets consist 100% hard examples where the numbers are close to each other. All models trained on sequences < and tested up to length 100. Vanilla corresponds to the original transformer, with single modification.
14

---

Under review as a conference paper at ICLR 2020
A.3 DETAILED VISUALIZATION OF LEARNED NUMBER EMBEDDINGS
In Figure 13 we show more detailed visualizations of the learned bitwise embeddings. These are 3-dimensional PCA projections of the full embedding matrix, capturing approximately 75% of the total variance. The main takeaway is that the network is able to learn a coherent number system with a great deal of structure, and that this structure varies depending on the specific task of the network. This is reminiscent of Paccanaro & Hinton (200)) where linear embeddings learned the correct structure to solve a simple modular arithmetic task. Also of note is that the network learns to embed correspond to infinity, outside of this structure.
Perhaps the most interesting feature is found in Figure 13c Here, we trained the network on the addition task, however we randomly held out 50% of the input numbers, shown in red. The network still learns a coherent embedding structure, and places the held-out numbers their correct positions. This is important feature if the network is to generalize to larger numbers. In particular, we hope to scale NEE to operate on 64-bit values, where sceing each integer at training time is impossible.
Future work will also investigate the resulting embedding from NEE that performs multiple or more complex tasks.
210
194
215
198g 208
209
182
213
82
1.5
868
66
218
149
80
70
200g
154g)
1--
2000
210
14
147
14
200
81
68
65
6 217
221
145707
204g
20'
154p
129
94
25
240g
2.28
205
85
19
78
247
20
20
165
17
90
245
43
198
114
157
220
25/
260
178
141
119
100
234
162
0.5
2520
848
1760000
215
181
25
07
249
177228
102
60
233
161
126
237
123
48 10g
34
190g
1840
620
10%
18: 184
187
17
125
121
49
36
180g
171
33
1.5
10 155
189
178
100
62
173
56
18g
61
45
-1.5
0.55
155
(a) Sorting Fmheeding
15

---

Under review as a conference paper at ICLR 2020
602
063 696
0.8
037 4:00
005
035 698
699
834
007
(161
@11
860
068
001
(225
642
#05
0.4
0.6
059
0 08
Q55
67 2230
02
668
(229
228
049 1:12
'27
648
047
1-1
045
6:09
73 G236
644
72
0107
71 1234
(236
470
(23
69 0232
850
74
@33
0.2
053 16
050
051
16
777
67
e0
(239
75 9236
067
455
954
021
055 18
0-
19
B3
G
224
099
m30 22 052
02
063
051
25
G 88
167 220
28
G 86
(249
B5 (248
09 672
0/3
247
58
G252
-0.2
3: m98
217
11
anl
.255
234
80
30 2002
888
357 2000
6300
-0.4
002
4444
45 206
45 (208
025 068
-0.6
022
006
484 Q212
4146
(211
47 210
46
C209
020 002
628
007 060
059
-0.8
511 1414
50
1113
601 084
650
093
5155 (218
G154 (217
1215
695
57
(220
0868
(221
1222
222
0.5
-0.5
-1
-1.5
0
0.5
1.5
(b) Addition Embedding
866
897
688
08
01 (225
664
#0b
1228
(227
080
230
00:
0231
668
94
433
<236
233
198
95
6338
476
110
98
02
69
12
203
844
445
1206
24
17G242
21
34
44
75
13
98
19
23
46
25/g50
1247
44
(251
250
221
802
666
114
2256
403 221
(215
70
220
200
(222
061
80
222
B7
52
625
47
90
228
60
0151
800
55
58
39
-25
06
(c) Generalization on Addition (numbers randomly held out of training colored red).
Figure 13: 3-dimensional PCA projections of learned bitwisc embeddings for da) sorting, addition, and addition with 50% of the numbers withheld from training.
16

---

Under review as a conference paper at ICLR 2020
A.4 DUKSTRA'S ALGORITHM PERFORMANCE AND EXECUTION TRACE
In Figure 15 we show an execution trace of Dijkstra's algorithm with NEE components for the toy graph shown in Figure [4] where we choose "A" be the source node. Dijkstra's algorithm the weighted adjacency matrix, and maintains in memory the current shortest path values from each node the source, shortest-path, along with the currently considered node that denote as the "current node" or anchor. NEE also maintains variable mask that marks which nodes have already been considered. Given nodes, and beginning with the source node, Dijkstra's algorithm can be broken up into 3 distinct phases (after initialization), with each one using a NEE to perform underlying computation.
1. Reference the row of the adjacency matrix indexed by the current node, dists.curret..rm _node, elementwise add this to the current shortest path values possible.paths dists.current.noc
shortest.path
2. Take the new shortest path the elementwise min between shortest.pate and cesssibleppahh..
shortest.path
foannentwisemminif http://ssmmmmsoottttptpttttpppe htt.:esss._pptt...oosssepeppptte
3. Select the new "current" node from among those not masked out, and update the mask. anchor elementwise haapeeettiiieeaaannn(aaaa argmin(shortes..5) (mask of 1])
mask (mR(ma one.hot(anchor)
4. If 3; mask, = N then terminate, otherwise go back to step 1.
Putting this together, we find that training NEE components to execute the phases on graphs of up to nodes in size generalizes perfectly to graphs of up 40 nodes (Figure 9)
Figure 14: Graph used for shortest path execution trace in Figure 15 Node "A" is used the
17

---

Under review as a conference paper at ICLR 2020
DAT
Tenskerer
Tandhann
Dutores
Curret Mari
0000
Mensrr Stse
Clametatte Surt
Gariirduia Vir
Oummet Noc
Cument Masi
Hemeyy
Tur dertr
Arjeeny Marls
P
Tepeefonoo
Toreeforee
Zbooee [ 3
Cuvett
Figure 15: An exccution trace the Dijkstra algorithm shortest paths with NEE components. Highlighted arrays and values correspond to referencing the memory inputs of the same corresponding color at that iteration.
18


Reviews:
[{"review_id": "rke3dweCtB", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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).\n\nThe authors further motivate their work by assuming that these phenomena are due to lack of prior structure that can be alleviated by further supervision during training.  The goal is for complex behaviour to emerge via composition of simple functions (which clearly follows the deep learning paradigm).Specifically, the authors propose a modification to the transformer architecture, that does not use positional encodings (the authors mention that this was detrimental to their work - it would be good to provide some more insight into that) and single-headed attention.  The main contribution seems to be adding the self-attention mask that is learned, along with execution traces that have been used in previous work.  An relatively small increase in performance is observed due to this, but it seems that the experiment is limited (no standard deviation in results, so I presume one run with one initialization).  Therefore it seems to me that the contribution of this paper is limited in terms of technical contribution.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "SJl4vu-hFB", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper investigates an interesting problem of building a program execution engine with neural networks. The authors proposed a transformer-based model to learn basic subroutines, such as comparison, find min, and addition, and apply them in several standard algorithms, such as sorting and Dijkstra’s.\n\nPros:\n1. The method achieves generalization towards longer sequences than the sequences in the training set in several algorithms.\n2. The method represents numbers in binary form and the visualization shows that it learns embeddings from fixed-range integer numbers in a well-structured manner.\n3. The learned NEE subroutines are tested in a variety of standard algorithms, such as multiple sorting algorithms and Dijkstra’s shortest path algorithm. The experiments further demonstrate that several NEE subroutines can be composed together in complex algorithms.\n\nCons:\n1. NEE mostly focuses on learning low-level subroutines such as number comparison or addition. Therefore, it has to be used along with conventional subroutines, and cannot completely replace the full execution in complex algorithms, which have sophisticated control logic, such as if/else and for-loops. When the transformer model is used alone in the sorting task (Sec. 4.1), the performance degrades substantially as the sequence length gets longer.\n2. Although the method achieved some degree of strong generalization, it lacks a formal way to verify the correctness of the learned subroutines, as opposed to prior work on program synthesis (Cai et al. 2017) that can prove the generalization of their model with recursion. \n3. The method relies on detailed execution traces for supervised learning which can be costly to obtain.\n\nQuestions:\n1. Confusing sentence: “Can we retain high attention resolution by restricting the model to only observe the first scenario repeatedly?” Can you elaborate on what you meant here?\n2. From Figure 1, it seems that the model with dot product attention generalizes better in longer sequences than the one with all modifications. What’s the reason?\n3. I would like to better understand the limitation of these learned NEE subroutines in long sequences. For instance, in Figure 8 and Figure 9, how would the model perform beyond the lengths of the sequences tested here? Would the performance maintain at 100% or decrease gradually as the sequences get even longer?\n4. I am curious to know how this method could be extended to support more complex number systems, such as float numbers, and more complex data structures beyond sequences, such as binary trees and priority queues. I’d love to hear what the authors have to say about this.\n5. I'd also like to know if the number embeddings learn in different algorithms would exhibit different structures (by examining the visualization of number embeddings learned in different tasks).\n6. As NEE focuses on learning the basic subroutines while NPI aims to learn the high-level program executions, I think that it’d be very interesting to see how these two can combine their complementary strengths to build a complete neural-based execution engine.\n\nTypos:\nSelect he node --> Select the node\n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}, {"review_id": "HygQp_7b9r", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "# Summary\n\nThis paper trains a network to mimic simple known algorithms in a way that guarantees that they generalize to\nout-of-distribution test instances. The network mimics the algorithms by running repeatedly in a loop where each\niteration of the loop runs a Transformer and outputs a mask that tells the next iteration the inputs to process. The\nsetup is tested on sorting, adding, and graph algorithms, and found to learn regular number representations that\nsupposedly aid generalization.\n\n# Review\n\nThis paper has an admirable and useful goal, but the way it is currently implemented and presented is not ready for\npublication at ICLR.\n\nMy main issue is with the training/testing setup and its presentation. The authors assume a certain structure of an\nalgorithm (for instance, the iterative structure of recursive selection sort), delegate one or more modules inside this\nstructure to be implemented by neural networks, and train them only.\nMost of the \"strong generalization\" is coming from the fact that the iterative structure is fixed. The work abstracts\nout the most complex parts of each algorithm. In Figure 3, for instance, the NN must learn to find the smallest element\namong the non-masked-out ones on the input, return it, and mask it out. This is a much simpler task than the whole\nsorting algorithm. Training the network to solve \"find_min\" != claiming that the network solves and strongly generalizes\non \"sort\".\n\nImportant training details are left unspecified. How is the data for training NEEs generated? For instance, for training\nthe network in Figure 3, do you trace the whole selection sort on a randomly generated list, and collect the\nintermediate input/output pairs for \"find_min\"? If so, it's absolutely unsurprising that the process also works for\nlonger lists -- see above. Are composable NEEs, like the three networks in Figure 7, trained jointly or separately? Do\nthey observe their own outputs that are fed into subsequent NEE networks, or are the previous outputs teacher-forced,\nor are they pre-trained? Many of these details need to be clarified precisely to make the experimental setup\nverifiable.\nSome important details are presented factually without any motivation. For example, why does Figure 5 use\nSHIFT and XOR? Why, in general, the next mask produced by a NEE is XORed with a previous one instead of replacing it?\n\nI liked the embedding visualizations, which clearly demonstrate structure in the latent space driven by (a) the\nbinary number representations, and (b) the addition task objective. In addition to regular ordering structure (needed to\nimplement addition), the latent space also clearly exhibits regularities inherent to the binary representation, such as\nthe shift by 64 in Figure 8a. While interesting, this only confirms the findings of Shi et al., albeit in a more pure\nexperimental setting.\n\nIn summary, the scope of experiments and presentation of results would need to be significantly improved in order for\nthis work to reach the quality bar of ICLR.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper investigates the problem of building a program execution engine with neural networks. While the reviewers find this paper to contain interesting ideas, the technical contributions, scope of experiments, and the presentation of results would need to be significantly improved in order for this work to reach the quality bar of ICLR.
Decision: Reject