Learning Simple Algorithms from Examples

by   Wojciech Zaremba, et al.

We present an approach for learning simple algorithms such as copying, multi-digit addition and single digit multiplication directly from examples. Our framework consists of a set of interfaces, accessed by a controller. Typical interfaces are 1-D tapes or 2-D grids that hold the input and output data. For the controller, we explore a range of neural network-based models which vary in their ability to abstract the underlying algorithm from training instances and generalize to test examples with many thousands of digits. The controller is trained using Q-learning with several enhancements and we show that the bottleneck is in the capabilities of the controller rather than in the search incurred by Q-learning.



page 6


Towards Modular Algorithm Induction

We present a modular neural network architecture Main that learns algori...

Hybrid LQG-Neural Controller for Inverted Pendulum System

The paper presents a hybrid system controller, incorporating a neural an...

Learning Controllers for Unstable Linear Quadratic Regulators from a Single Trajectory

We present the first approach for learning – from a single trajectory – ...

Reinforcement Learning Neural Turing Machines - Revised

The Neural Turing Machine (NTM) is more expressive than all previously c...

A neural network based heading and position control system of a ship

Heading and position control system of ships has remained a challenging ...

Robustness of Control Design via Bayesian Learning

In the realm of supervised learning, Bayesian learning has shown robust ...

Neural Network Predictive Controller for Grid-Connected Virtual Synchronous Generator

In this paper, a neural network predictive controller is proposed to reg...

Code Repositories


We present an approach for learning simple algorithms such as copying, multi-digit addition and single digit multiplication directly from examples. Our framework consists of a set of interfaces, accessed by a controller. Typical interfaces are 1-D tapes or 2-D grids that hold the input and output data. For the controller, we explore a range of neural network-based models which vary in their ability to abstract the underlying algorithm from training instances and generalize to test examples with

view repo


A Python implementation of Zaremba's "Learning Simple Algorithms from Examples" work.

view repo
This week in AI

Get the week's most popular data science and artificial intelligence research sent straight to your inbox every Saturday.

1 Introduction

Many every day tasks require a multi-step interaction with the world. For example, picking an apple from a tree requires visual localization of the apple; extending the arm and then fine muscle control, guided by visual feedback, to pluck it from the tree. While each individual procedure is not complex, the task nevertheless requires careful sequencing of operations across both visual and motor systems.

This paper explores how machines can learn algorithms involving a similar compositional structure. Since our emphasis is on learning the correct sequence of operations, we consider the domain of arithmetic where the operations themselves are very simple. For example, although learning to add two digits is straightforward, solving addition of two multi-digit numbers requires precise coordination of this operation with movement over the sequence and recording of the carry. We explore a variety of algorithms in this domain, including complex tasks involving addition and multiplication.

Our approach formalizes the notion of a central controller that interacts with the world via a set of interfaces, appropriate to the task at hand. The controller is a neural network model which must learn to control the interfaces, via a set of discrete actions (e.g. “move input tape left”, “read”, “write symbol to output tape”, “write nothing this time step” ) to produce the correct output for given input patterns. Specifically, we train the controller from large sets of examples of input and output patterns using reinforcement learning. Our reward signal is sparse, only being received when the model emits the correct symbol on the output tape.

We consider two separate settings. In the first, we provide supervision in the form of ground truth actions. In the second, we train only with input-output pairs (i.e. no supervision over actions). While we are able to solve all the tasks in the latter case, the supervised setting provides insights about the model limitations and an upper bound on the performance. We evaluate our model on sequences far longer than those present during training. Surprisingly, we find that controllers with even modest capacity to recall previous states can easily overfit the short training sequences and not generalize to the test examples, even if the correct actions are provided. Even with an appropriate controller, off-the-shelf -learning fails on the majority of our tasks. We therefore introduce a series of modifications that dramatically improve performance. These include: (i) a novel dynamic discount term that makes the reward invariant to the sequence length; (ii) an extra penalty that aids generalization and (iii) the deployment of Watkins Q-lambda [Sutton & Barto (1998)].

We would like to direct the reader to the video accompanying this paper (https://youtu.be/GVe6kfJnRAw). This gives a concise overview of our approach and complements the following explanations. Full source code for this work can be found at https://github.com/wojzaremba/algorithm-learning.

2 Model

Our model consists of an RNN-based controller that accesses the environment through a series of pre-defined interfaces. Each interface has a specific structure and set of actions it can perform. The interfaces are manually selected according to the task (see Section 3). The controller is the only part of the system that learns and has no prior knowledge of how the interfaces operate. Thus the controller must learn the sequence of actions over the various interfaces that allow it to solve a task. We make use of three different interfaces:

Input Tape: This provides access to the input data symbols stored on an “infinite” 1-D tape. A read head accesses a single character at a time through the read action. The head can be moved via the left and right actions.

Input Grid: This is a 2D version of the input tape where the read head can now be moved by actions up, down, left and right.

Output Tape: This is similar to the input tape, except that the head now writes a single symbol at a time to the tape, as provided the controller. The vocabulary includes a no-operation symbol (NOP) enabling the controller to defer output if it desires. During training, the written and target symbols are compared using a cross-entropy loss. This provides a differentiable learning signal that is used in addition to the sparse reward signal provided by the -learning.

Fig. 1(a) shows examples of the input tape and grid interfaces. Fig. 1(b) gives an overview of our controller–interface abstraction and Fig. 1(c) shows an example of this on the addition task (for two time steps).

(a) ControllerController InputController Output
Input Interface

Output Interface

Memory Interface
Input Interface

Output Interface

Memory Interface

Past StateFuture State
Figure 1: (a): The input tape and grid interfaces. Both have a single head (gray box) that reads one character at a time, in response to a read action from the controller. It can also move the location of the head with the left and right (and up, down) actions. (b) An overview of the model, showing the abstraction of controller and a set of interfaces (in our experiments the memory interface is not used). (c) An example of the model applied to the addition task. At time step , the controller, a form of RNN, reads the symbol from the input grid and outputs a no-operation symbol () on the output tape and a down action on the input interface, as well as passing the hidden state to the next timestep.

For the controller, we explore several recurrent neural network architectures: two different sizes of 1-layer LSTM [

Hochreiter & Schmidhuber (1997)

], a gated-recurrent unit (GRU)[

Cho et al. (2014)] and a vanilla feed-forward network. Note that RNN-based models are able to remember previous network state, unlike the the feed-forward network. This is important because some tasks explicitly require some form of memory, e.g. the carry in addition.

The simple algorithms we consider (see Section 3) have deterministic solutions that can be expressed as a finite state automata. Thus during training we hope the controller will implicitly learn the correct automata from the training samples, since this would ensure generalization to sequences of arbitrary length. On some tasks like reverse, we observe a higher-order form of over-fitting: the model learns to solve the training tasks correctly and generalizes successfully to test sequences of the same length (thus is not over-fitting in the standard sense). However, when presented with longer test sequences the model fails completely. This suggests that the model has converged to an incorrect local minima, one corresponding to an alternate automata which have an implicit awareness of the sequence length of which they were trained. See Fig. 4 for an example of this on the reverse task. Note that this behavior results from the controller, not the learning scheme, since it is present in both the supervised (Section 5) and -learning settings (Section 6). These experiments show the need to carefully adjust the controller capacity to prevent it learning any dependencies on the length of training sequences, yet ensuring it has enough state to implement the algorithm in question.

As illustrated in Fig. 1

(c), the controller passes two signals to the output tape: a discrete action (move left, move right, write something) and a symbol from the vocabulary. This symbol is produced by taking the max from the softmax output on the top of the controller. In training, two different signals are computed from this: (i) a cross-entropy loss is used to compare the softmax output to the target symbol and (ii) a discrete 1/0 reward if the symbol is correct/incorrect. The first signal gives a continuous gradient to update the controller parameters via backpropagation. Leveraging the reward requires reinforcement learning, since many actions might occur before a symbol is written to the output tape. Thus the action output of the controller is trained with reinforcement learning and the symbol output is trained by backpropagation.

3 Tasks

We consider six different tasks: copy, reverse, walk, multi-digit addition, 3 number addition and single digit multiplication. The input interface for copy and reverse is an input tape, but an input grid for the others. All tasks use an output tape interface. Unless otherwise stated, all arithmetic operations use base 10. Examples of the six tasks are shown in Fig. 2.

Copy: This task involves copying the symbols from the input tape to the output tape. Although simple, the model still has to learn the correspondence between input and output symbols, as well as executing the move right action on the input tape.

Reverse: Here the goal is to reverse a sequence of symbols on the input tape. We provide a special character “r” to indicate the end of the sequence. The model must learn to move right multiple times until it hits the “r” symbol, then move to the left, copying the symbols to the output tape.

Walk: The goal is to copy symbols, according to the directions given by an arrow symbol. The controller starts by moving to the right (suppressing prediction) until reaching one of the symbols . Then it should change it’s direction accordingly, and copy all symbols encountered to the output tape.

Addition: The goal is to add two multi-digit sequences, provided on an input grid. The sequences are provided in two adjacent rows, with their right edges aligned. The initial position of the read head is the last digit of the top number (i.e. upper-right corner). The model has to: (i) memorize an addition table for pairs of digits; (ii) learn how to move over the input grid and (iii) discover the concept of a carry.

3 Number Addition: As for the addition task, but now three numbers are to be added. This is more challenging as the reward signal is less frequent (since more correct actions must be completed before a correct output digit can be produced). Also the carry now can take on three states (0, 1 and 2), compared with two for the 2 number addition task.

Single Digit Multiplication: This involves multiplying a single digit with a long multi-digit number. It is of similar complexity to the 2 number addition task, except that the carry can take on more values .

Copy Reverse Walk Addition 3 number addition Single digit multiplication

Figure 2: Examples of the six tasks, presented in their initial state. The yellow box indicates the starting position of the read head on the Input Interface. The gray characters on the Output Tape are target symbols used in training.

4 Related Work

A variety of recent work has explored the learning of simple algorithms. Many of them are different embodiments of the controller-interface abstraction formalized in our model. The Neural Turing Machine (NTM) [

Graves et al. (2014)] uses a modified LSTM [Hochreiter & Schmidhuber (1997)] as the controller, and has three inferences: sequential input, delayed output and a differentiable memory. The model is able to learn simple algorithms including copying and sorting. The Stack RNN [Joulin & Mikolov (2015)] has an RNN controller and three interfaces: sequential input, a stack memory and sequential output. The learning of simple binary patterns and regular expressions is demonstrated. A closely related work to this is [Das et al. (1992)], which was recently extended in the Neural DeQue [Grefenstette et al. (2015)] to use a list instead. End-to-End Memory Networks [Sukhbaatar et al. (2015)] use a feed-forward network as the controller and interfaces consisting of a soft-attention input, plus a delayed output (by a fixed number of “hops”). The model is applied to simple Q&A tasks, some of which involve logical reasoning. In contrast, our model automatically determines when to produce output and uses more general interfaces.

However, most of these approaches use continuous interfaces that permit training via back-propagation of gradients. Our approach differs in that it uses discrete interfaces thus is more challenging to train since as we must rely on reinforcement learning instead. A notable exception is the Reinforcement Learning Neural Turing Machine (RLNTM) [Zaremba & Sutskever (2015)] which is a version of the NTM with discrete interfaces. The Stack-RNN [Joulin & Mikolov (2015)] also uses a discrete search procedure for its interfaces but it is unclear how this would scale to larger problems.

The problem of learning algorithms has its origins in the field of program induction [Nordin (1997); Liang et al. (2013); Wineberg & Oppacher (1994); Solomonoff (1964)]. In this domain, the model has to infer the source code of a program that solves a given problem. This is a similar goal to ours, but in quite a different setting. I.e. we do not produce a computer program, but rather a neural net that can operate with interfaces such as tapes and so implements the program without being human-readable. A more relevant work is [Schmidhuber (2004)

] which learns an algorithms for the Hanoi tower problem, using a simple form of program induction and incremental learning components. Genetic algorithms [

Holland (1992); Goldberg (1989)] also can be considered a form of program induction, but are mostly based on a random search strategy rather than a learned one.

Similar to [Mnih et al. (2013)], we train the controller to approximate the -function. However, we introduce several modifications on top of the classical -learning. First, we use Watkins [Watkins (1989); Sutton & Barto (1998)]. This helps to overcome a non-stationary environment. We are unaware of any prior work that uses Watkins for this purpose. Second, we reparametrized function, to become invariant to the sequence length. Finally, we penalize , which might help to remove positive bias [Hasselt (2010)].

5 Supervised Experiments

To understand the behavior of our model and to provide an upper bound on performance, we train our model in a supervised setting, i.e. where the ground truth actions are provided. Note that the controller must still learn which symbol to output. But this now can be done purely with backpropagation since the actions are known.

To facilitate comparisons of difficulty between tasks, we use a common measure of complexity, corresponding to the number of time steps required to solve each task (using the ground truth actions***In practice, multiple solutions can exist (see Appendix A), thus the measure is approximate.). For instance, a reserve task involving a sequence of length requires time-steps ( steps to move to the “r” and steps to move back to the start). The conversion factors between sequence lengths and complexity are as follows: copy=1; reverse=2; walk=1; addition=2; 3 row addition=3 and single digit multiplication=1.

For each task, we train a separate model, starting with sequences of complexity 6 and incrementing by 4 once it achieves 100% accuracy on held-out examples of the current length. Training stops once the model successfully generalizes to examples of complexity 1000. Three different cores for the controllers are explored: (i) a 200 unit, 1-layer LSTM; (iii) a 200 unit, 1-layer GRU model and (iii) a 200 unit, 1-layer feed-forward network. An additional linear layer is placed on top of these model that maps the hidden state to either action for a given interface, or the target symbol.

In Fig. 3 we show the accuracy of the different controllers on the six tasks for test instances of increasing complexity, up to time-steps. The simple feed-forward controller generalizes perfectly on the copy, reverse and walk tasks but completely fails on the remaining ones, due to a lack of required memoryAmmending the interfaces to allow both reading and writing on the same interface would provide a mechanism for long-term memory, even with a feed-forward controller. But then the same lack of generalization issues (encountered with more powerful controllers) would become an issue.. The RNN-based controllers succeed to varying degrees, although some variability in performance is observed.

Further insight can be obtained by examining the internal state of the controller. To do this, we compute the autocorrelation matrixLet be the controller state at time , then the autocorrelation between time-steps and is given by where . is the number of time steps (i.e. complexity). of the network state over time when the model is processing a reverse task example of length , having been trained on sequences of length or shorter. For this problem there should be two distinct states: move right until “r” is reached and then move left to the start. Fig. 1 plots for models with three different controllers. The larger the controller capacity, the less similar the states are within the two phases of execution, showing how it has not captured the correct algorithm. The figure also shows the confidence in the two actions over time. In the case of the high capacity models, the initial confidence in the move left action is high, but this drops off after moving along the sequence. This is because the controller has learned during training that it should change direction after at most steps. Consequently, the unexpectedly long test sequence makes it unsure of what the correct action is. By contrast, the simple feed-forward controller does not show this behavior since it is stateless, thus has no capacity to know where it is within a sequence. The equivalent automata is shown in Fig. 4(a), while Fig. 4(b) shows the incorrect time-dependent automata learned by the over-expressive RNN-based controllers. We note that this argument is empirically supported by our results in Table 2, as well as related work such as [Graves et al. (2014)] and [Joulin & Mikolov (2015)] which found limited capacity controllers to be most effective. For example, in the latter case, the counting and memorization tasks used controllers with just and units respectively.

Figure 3: Test accuracy for all tasks with supervised actions over runs for feed-forward (green), GRU (red) and LSTM (yellow) controllers. In this setting the optimal policy is provided. Complexity is the number of time steps required to compute the solution. Every task has slightly different conversion factor between complexity and the sequence length: a complexity of for copy and walk would mean input symbols; for reverse would correspond to input symbols; for addition would involve two long numbers; for 3 row addition would involve three long numbers and for single digit multiplication would involve a single long number.

Table 1: Three models with different controllers (feed-forward, 200 unit LSTM and 400 unit LSTM) trained on the reverse task and applied to a 20 digit test example. The top row shows confidence values for the two actions on the input tape: move left (green) and move right (red) as a function of time. The correct model should be equivalent to a two-state automata (Fig. 4), thus we expect to see the controller hidden state occupy two distinct values. The autocorrelation matrices (whose axes are also time) show this to be the case for the feed-forward model – two distinct blocks of high correlation. However, for the LSTM controllers, this structure is only loosely present in the matrix, indicating that they have failed to learn the correct algorithm.









Figure 4:

(a): The automata describing the correct solution to the reverse problem. The model first has to go to the right while suppressing prediction. Then, it has to go to the left and predict what it sees at the given moment (this figure illustrates only actions over the Input Tape). (b) Another automata that solves the reverse problem for short sequences, but does not generalize to arbitrary length sequences, unlike (a). Expressive models like LSTMs tend to learn such incorrect automata.

6 Q-Learning

In the previous section, we assumed that the optimal controller actions were given during training. This meant only the output symbols need to be predicted and these could be learned via backpropagation. We now consider the setting where the actions are also learned, to test the true capabilities of the models to learn simple algorithms from pairs of input and output sequences.

We use -learning, a standard reinforcement learning algorithm to learn a sequence of discrete actions that solves a problem. A function

, the estimated sum of future rewards, is updated during training according to:


Taking the action in state causes a transition to state , which in our case is deterministic. is the reward experienced in the state . The discount factor is and is the learning rate. The another commonly considered quantity is . is called the value function, and is the expected sum of future rewards starting from the state . Moreover, and are function values for the optimal policy.

Our controller receives a reward of every time it correctly predicts a digit (and otherwise). Since the overall solution to the task requires all digits to be correct, we terminate a training episode as soon as an incorrect prediction is made. This learning environment is non-stationary, since even if the model initially picks the right actions, the symbol prediction is unlikely to be correct, so the model receives no reward. But further on in training, when the symbol prediction is more reliable, the correct action will be rewarded§§§If we were to use reinforcement to train the symbol output as well as the actions, then the environment would be stationary. However, this would mean ignoring the reliable signal available from direct backpropagation of the symbol output.. This is important because reinforcement learning algorithms assume stationarity of the environment, which is not true in our case. Learning in non-stationary environments is not well understood and there are no definitive methods to deal with it. However, empirically we find that this non-stationarity can be partially addressed by the use of Watkins [Watkins (1989)], as detailed in Section 6.2.

6.1 Dynamic Discount

The purpose of the reinforcement learning is to learn a policy that yields the highest sum of the future rewards. -learning does it indirectly by learning a -function. The optimal policy can be extracted by taking over . Note that shifting or scaling induces the same policy. We propose to dynamically rescale so (i) it is independent of the length of the episode and (ii) is within a small range, making it easier to predict.

We define to be our reparametrization. should be roughly in range , and it should correspond to how close we are to . could be decomposed multiplicatively as . However, in practice, we do not have access to , thus instead we use an estimate of future rewards based on the total number of digits left in the sequence. Since every correct prediction yields a reward of , the optimal policy should achieve sum of future rewards equal to the number of remaining symbols to predict. The number of remaining symbols to predict is known and we denote it by . Note that this is a form of supervision, albeit a weak one.

Therefore, we normalize the -function by the remaining sum of rewards left in the task:

We assume that transitions to , and we re-write the -learning update equations:

Note that , with equality if no digit was predicted at the current time-step. As the episode progresses, the discount factor decreases, making the model greedier. At the end of the sequence, the discount drops to .

6.2 Watkins

The update to in Eqn. 1 comes from two parts: the observed reward and the estimated future reward . In our setting, there are two factors that make the former far more reliable than the latter: (i) rewards are deterministic and (ii) the non-stationarity (induced by the ongoing learning of the symbol output by backpropagation) means that estimates of are unreliable as environment evolves. Consequently, the single action recurrence used in Eqn. 1 can be improved upon when on-policy actions are chosen. More precisely, let be consecutive actions induced by :

Then the optimal follows the following recursive equation:

and the update rule corresponding to Eqn. 1 becomes:

This is a special form of Watkins [Watkins (1989)] where . The classical applications of Watkins suggest choosing a small , which trades-off estimates based on various numbers of future rewards. rolls back to the classical -learning. Due to reliability of our rewards, we found to be better than , however this needs further study.

Note that this unrolling of rewards can only take place until a non-greedy action is taken. When using an -greedy policy, this means we would expect to be able to unroll steps, on average. For the value of used in our experiments, this corresponds to steps on average.

6.3 Penalty on -function

After reparameterizing the -function to (Section 6.1), the optimal should be 1 for the correct action and zero otherwise. To encourage our estimate to converge to this, we introduce a penalty that “pushes down” on incorrect actions: . This has the effect of introducing a margin between correct and incorrect actions, greatly improving generalization. We commence training with and make it non-zero once good accuracy is reached on short samples (introducing it from the outset hurts learning).

6.4 Reinforcement Learning Experiments

We apply our enhancements to the six tasks in a series of experiments designed to examine the contribution of each of them. Unless otherwise specified, the controller is a 1-layer GRU model with 200 units. This was selected on the basis of its mean performance across the six tasks in the supervised setting (see Section 5). As the performance of reinforcement learning methods tend to be highly stochastic, we repeat each experiment times with a different random seed. Each model is trained using characters which takes hrs. A model is considered to have successfully solved the task if it able to give a perfect answer to test instances, each digits in length. The GRU model is trained with a batch size of , a learning rate of , using the same initialization as [Glorot & Bengio (2010)] but multiplied by 2. All tasks are trained with the same curriculum used in the supervised experiments (and in [Joulin & Mikolov (2015)]), whereby the sequences are initially of complexity (corresponding to 2 or 3 digits, depending on the task) and once 100% accuracy is achieved, increased by until the model is able to solve validation sequences of length .

For 3-row addition, a more elaborate curriculum was needed which started with examples that did not involve a carry and contained many zero. The test distribution was unaffected. Some examples:   






We show results for various combinations of terms in Table 2. The experiments demonstrate that standard -learning fails on most of our tasks (first six columns). Each of our additions (dynamic discount, Watkins and penalty term) give significant improvements. When all three are used our model is able to succeed at all tasks, providing the appropriate curriculum and controller are used. For the reverse and walk tasks, the default GRU controller failed completely. However, using a feed-forward controller instead enabled the model to succeed, when dynamic discount and Watkins was used. As noted above, the 3-row addition required a more careful curriculum before the model was able to learn successfully. Increasing the capacity of the controller (columns 2-4) hurts performance, echoing Fig. 1. The last two columns of Table 2 show results on test sequences of length 1000. Except for multiplication, the models still generalized successfully.

Test length 100 100 100 100 100 100 100 100 1000 1000
#Units 600 400 200 200 200 200 200 200 200 200
Discount 1 1 1 0.99 0.95 D D D D D
Task Penalty


Copying 30% 60% 90% 50% 70% 90% 100% 100% 100% 100%


Reverse 0% 0% 0% 0% 0% 0% 0% 0% 0% 0%
Reverse (FF controller) 0% 0% 0% 0% 0% 0% 100% 90% 100% 90%


Walk 0% 0% 0% 0% 0% 0% 10% 90% 10% 80%
Walk (FF controller) 0% 0% 0% 0% 0% 0% 100% 100% 100% 100%


2-row Addition 10% 70% 70% 70% 80% 60% 60% 100% 40% 100%


3-row Addition 0% 0% 0% 0% 0% 0% 0% 0% 0% 0%
3-row Addition (extra curriculum) 0% 50% 80% 40% 50% 50% 80% 80% 10% 60%


Single Digit Multiplication 0% 0% 0% 0% 0% 100% 100% 100% 0% 0%
Table 2: Success rates for classical -learning (columns 2-5) versus our enhanced -learning. A GRU-based controller is used on all tasks, except reverse and walk which use a feed-forward network. Curriculum learning was also used for the 3-row addition task (see text for details). When dynamic discount (D), Watkins and the penalty term are all used the model consistently succeeds on all tasks. The model still performs well on test sequences of length 1000, apart from the multiplication task. Increasing the capacity of the controller results in worse performance (columns 2-4).

Fig. 5 shows accuracy as a function of test example complexity for standard -learning and our enhanced version. The difference is performance is clear. At very high complexity, corresponding to ’s of digits, the accuracy starts to drop on the more complicated tasks. We note that these trends are essentially the same as those observed in the supervised setting (Fig. 3), suggesting that -learning is not to blame. Instead, the inability of the controller to learn an automata seems to be the cause. Potential solutions to this might include (i) noise injection, (ii) discretization of state, (iii) a state error correction mechanism or (iv) regularizing the learned automata using MDL principles. However, this issue, the inability of RNN to perfectly represent an automata can be examined separately from the setting where actions have to be learnt (i.e. in the supervised domain).

Figure 5: Test accuracy as a function of task complexity ( runs) for standard -learning (blue) and our enhanced version (dynamic discount, Watkins and penalty term). Accuracy corresponds to the fraction of correct test cases (all digits must be predicted correctly for the instance to be considered correct).

Further results can be found in the appendices. For the addition task, our model was able to discover multiple correct solutions, each with a different movement pattern over the input tape (see Appendix A). Table 3 in Appendix B sheds light on the trade-off between errors in actions and errors in symbol prediction by varying the base used in the arithmetic operations and hence the size of the target vocabulary. Appendix C explores the use of non-integer rewards. Surprisingly, this slows down training, relative to the 0/1 reward structure.

7 Discussion

We have explored the ability of neural network models to learn algorithms for simple arithmetic operations. Through experiments with supervision and reinforcement learning, we have shown that they are able to do this successfully, albeit with caveats. -learning was shown to work as well as the supervised case. But, disappointingly, we were not able to find a single controller that could solve all tasks. We found that for some tasks, generalization ability was sensitive to the memory capacity of the controller: too little and it would be unable to solve more complex tasks that rely on carrying state across time; too much and the resulting model would overfit the length of the training sequences. Finding automatic methods to control model capacity would seem to be important in developing robust models for this type of learning problem.


We wish to thank Jason Weston, Marc’Aurelio Ranzato and Przemysław Mazur for useful discussions, and comments. We also thank Christopher Olah for LSTM figures that have been used in the paper and the accompanying video.


Appendix A: Different Solutions to Addition Task

On examination of the models learned on the addition task, we notice that three different solutions were discovered. While they all give the correct answer, they differ in their actions over the input grid, as shown in Fig. 6.

Figure 6: Our model found three different solutions to the addition task, all of which give the correct answer. The arrows show the trajectory of the read head over the input grid.

Appendix B:Reward Frequency vs Reward Reliability

We explore how learning time varies as the size of the target vocabulary is varied. This trades off reward frequency and reliability. For small vocabularies, the reward occurs more often but is less reliable since the chance of the wrong action sequence yielding the correct result is relatively high (and vice-versa for for larger vocabularies). For copying and reverse tasks, altering the vocabulary size just alters the variety of symbols on the tape. However, for the arithmetic operations this involves a change of base, which influences the task in a more complex way. For instance, addition in base requires the memorization of digit-to-digit addition table of size instead of for the base . Table 3 shows the median training time as a function of vocabulary size. The results suggest that an infrequent but reliable reward is preferred to a frequent but noisy one.

TaskVocabulary Size 2 3 4 5 6 7 8 9 10
Coping 1.4 0.6 0.6 0.6 0.6 0.5 0.6 0.6 0.6
Reverse (FF controller) 6.5 23.8 3.1 3.6 3.8 2.5 2.8 2.0 3.1
Walk (FF controller) 8.7 6.9 6.8 4.0 6.2 5.3 4.4 3.9 11.1
Addition 250.0 30.9 14.5 26.1 21.8 21.9 25.0 23.4 21.1
3-number Addition (extra curriculum) 250.0 61.5 250.0 250.0 112.2 178.2 93.8 79.1 81.9
Single Digit Multiplication invalid 6.2 17.8 20.9 21.4 21.5 22.3 23.3 24.7
Table 3: Median training time (minutes) over runs as we vary the base used (hence vocabulary size) on different problems. Training stops when the model successfully generalizes to test sequences of length . The results show the relative importance of reward frequency versus reliability, with the latter being more important.

Appendix C: Reward Structure

Reward in reinforcement learning systems drives the learning process. In our setting we control the rewards, deciding when, and how much to give. We now examine various kinds of rewards and their influence on the learning time of our system.

Our vanilla setting gives a reward of for every correct prediction, and reward

for every incorrect one. We refer to this setting as “0/1 reward”. We consider two other settings in addition to this, both of which rely on the probabilities of the correct prediction. Let

be the target symbol and be the probability of predicting label .

In setting “Discretized reward”, we sort . That gives us an order on indices , i.e. . “Discretized reward” yields reward iff , reward iff , and reward iff . Otherwise, environment gives a reward . In the “Continuous reward” setting, a reward of is given for every prediction. One could also consider reward , however this quantity is unbounded, and further processing might be necessary to make it work.

Table 4 gives results for the three different reward structures, showing training time for the five tasks (training is stopped once the model generalizes to test sequences of length ). One might expect that a continuous reward would convey more information than a discrete one, thus result in faster training. However, the results do not support this hypothesis, as training seems harder with continuous reward than a discrete one. We hypothesize, that the continuous reward makes environment less stationary, which might make -learning less efficient, although this needs further verification.

TaskReward Type 0/1 reward Discretized reward Continuous reward
Coping 0.6 0.6 0.8
Reverse (FF controller) 3.1 3.1 59.7
Walk (FF controller) 11.1 9.5 250.0
Addition 21.1 21.6 24.2
3-number Addition (extra curriculum) 81.9 77.9 131.9
Single Digit Multiplication 24.7 26.5 26.6
Table 4: Median training time (minutes) for the five tasks for the three different reward structures. “0/1 reward”: the model gets a reward of for every correct prediction, and otherwise. “Discretized reward” provides a few more values of reward prediction, if sufficiently close to the correct one. “Continuous reward” gives a probability of correct answer as the reward. See text for details.