Programming with a Differentiable Forth Interpreter

Introduction

A central goal of Artificial Intelligence is the creation of machines that learn as effectively from human instruction as they do from data. A recent and important step towards this goal is the invention of neural architectures that learn to perform algorithms akin to traditional computers, using primitives such as memory access and stack manipulation (Graves et al. 2014; Joulin & Mikolov 2015; Grefenstette et al. 2015; Kaiser & Sutskever 2015; Kurach et al. 2016; Graves et al. 2016). These architectures can be trained through standard gradient descent methods, and enable machines to learn complex behaviour from input-output pairs or program traces. In this context, the role of the human programmer is often limited to providing training data. However, training data is a scarce resource for many tasks. In these cases, the programmer may have partial procedural background knowledge: one may know the rough structure of the program, or how to implement several subroutines that are likely necessary to solve the task. For example, in programming by demonstration (Lau et al. 2001) or query language programming (Neelakantan et al. 2015a) a user establishes a larger set of conditions on the data, and the model needs to set out the details. In all these scenarios, the question then becomes how to exploit various types of prior knowledge when learning algorithms.

To address the above question we present an approach that enables programmers to inject their procedural background knowledge into a neural network. In this approach, the programmer specifies a program sketch (Solar-Lezama et al. 2005) in a traditional programming language. This sketch defines one part of the neural network behaviour. The other part is learned using training data. The core insight that enables this approach is the fact that most programming languages can be formulated in terms of an abstract machine that executes the commands of the language. We implement these machines as neural networks, constraining parts of the networks to follow the sketched behaviour. The resulting neural programs are consistent with our prior knowledge and optimised with respect to the training data.

In this paper, we focus on the programming language Forth (Brodie 1980), a simple yet powerful stack-based language that facilitates factoring and abstraction. Underlying Forth’s semantics is a simple abstract machine. We introduce $\partial$4, an implementation of this machine that is differentiable with respect to the transition it executes at each time step, as well as distributed input representations. Sketches that users write define underspecified behaviour which can then be trained with backpropagation.

For two neural programming tasks introduced in previous work (Reed & de Freitas 2015) we present Forth sketches that capture different degrees of prior knowledge. For example, we define only the general recursive structure of a sorting problem. We show that given only input-output pairs, $\partial$4 can learn to fill the sketch and generalise well to problems of unseen size. In addition, we apply $\partial$4 to the task of solving word algebra problems. We show that when provided with basic algorithmic scaffolding and trained jointly with an upstream LSTM (Hochreiter & Schmidhuber 1997), $\partial$4 is able to learn to read natural language narratives, extract important numerical quantities, and reason with these, ultimately answering corresponding mathematical questions without the need for explicit intermediate representations used in previous work.

The contributions of our work are as follows: i) We present a neural implementation of a dual stack machine underlying Forth, ii) we introduce Forth sketches for programming with partial procedural background knowledge, iii) we apply Forth sketches as a procedural prior on learning algorithms from data, iv) we introduce program code optimisations based on symbolic execution that can speed up neural execution, and v) using Forth sketches we obtain state-of-the-art for end-to-end reasoning about quantities expressed in natural language narratives.

The Forth Abstract Machine

Forth is a simple Turing-complete stack-based programming language (ANSI 1994; Brodie 1980). We chose Forth as the host language of our work because i) it is an established, general-purpose high-level language relatively close to machine code, ii) it promotes highly modular programs through use of branching, loops and function calls, thus bringing out a good balance between assembly and higher level languages, and importantly iii) its abstract machine is simple enough for a straightforward creation of its continuous approximation. Forth’s underlying abstract machine is represented by a state $S=(D,R,H,c)$, which contains two stacks: a data evaluation pushdown stack $D$ (data stack) holds values for manipulation, and a return address pushdown stack $R$ (return stack) assists with return pointers and subroutine calls. These are accompanied by a heap or random memory access buffer $H$, and a program counter $c$.

A Forth program $P$ is a sequence Forth is a concatenative language. of Forth words (i.e. commands) $P=w_{1}\ldots w_{n}$. The role of a word varies, encompassing language keywords, primitives, and user-defined subroutines (e.g. DROP discards the top element of the data stack, or DUP duplicates the top element of the data stack). In this work, we restrict ourselves to a subset of all Forth words, detailed in Appendix A. Each word $w_{i}$ defines a transition function between machine states $w_{i}:S\rightarrow S$. Therefore, a program $P$ itself defines a transition function by simply applying the word at the current program counter to the current state. Although usually considered as a part of the heap $H$, we consider Forth programs $P$ separately to ease the analysis.

An example of a Bubble sort algorithm implemented in Forth is shown in Listing 1 in everything except lines 3b-4c. The execution starts from line 12 where literals are pushed on the data stack and the SORT is called. Line 10 executes the main loop over the sequence. Lines 2-7 denote the BUBBLE procedure – comparison of top two stack numbers (line 3a), and the recursive call to itself (line 4a). A detailed description of how this program is executed by the Forth abstract machine is provided in Appendix B. Notice that while Forth provides common control structures such as looping and branching, these can always be reduced to low-level code that uses jumps and conditional jumps (using the words BRANCH and BRANCH0, respectively). Likewise, we can think of sub-routine definitions as labelled code blocks, and their invocation amounts to jumping to the code block with the respective label.

∂\partial4: Differentiable Abstract Machine

When a programmer writes a Forth program, they define a sequence of Forth words, i.e., a sequence of known state transition functions. In other words, the programmer knows exactly how computation should proceed. To accommodate for cases when the developer’s procedural background knowledge is incomplete, we extend Forth to support the definition of a program sketch. As is the case with Forth programs, sketches are sequences of transition functions. However, a sketch may contain transition functions whose behaviour is learned from data.

To learn the behaviour of transition functions within a program we would like the machine output to be differentiable with respect to these functions (and possibly representations of inputs to the program). This enables us to choose parameterised transition functions such as neural networks.

To this end, we introduce $\partial$4, a TensorFlow (Abadi et al. 2015) implementation of a differentiable abstract machine with continuous state representations, differentiable words and sketches. Program execution in $\partial$4 is modelled by a recurrent neural network (RNN), parameterised by the transition functions at each time step.

where $\mathbf{p}\in\{\mathbf{d},\mathbf{r}\}$, $\text{inc}(\mathbf{p})=\mathbf{p}^{T}\mathbf{R^{1}+}$, $\text{dec}(\mathbf{p})=\mathbf{p}^{T}\mathbf{R^{-}}$, and $\mathbf{R^{1}+}$ and $\mathbf{R^{1}-}$ are increment and decrement matrices (left and right circular shift matrices).

Popping is realized by multiplying the TOS pointer and the memory buffer, and decreasing the TOS pointer:

It is straightforward to convert Forth words, defined as functions on discrete machine states, to functions operating on the continuous space $\mathcal{S}$. For example, consider the word DUP, which duplicates the top of the data stack. A differentiable version of DUP first calculates the value $\mathbf{e}$ on the TOS address of $\mathbf{D}$, as $\mathbf{e}=\mathbf{d}^{T}\mathbf{D}$. It then shifts the stack pointer via $\mathbf{d}\leftarrow\text{inc}(\mathbf{d})$, and writes $\mathbf{e}$ to $\mathbf{D}$ using $\text{write}_{\mathbf{D}}(\mathbf{e},\mathbf{d})$. The complete description of implemented Forth Words and their differentiable counterparts can be found in Appendix 4.

2 Forth Sketches

We define a Forth sketch $\mathbf{P_{\theta}}$ as a sequence of continuous transition functions $\mathbf{P}=\mathbf{w}_{1}\ldots\mathbf{w}_{n}$. Here, $\mathbf{w}_{i}\in\mathcal{S}\rightarrow\mathcal{S}$ either corresponds to a neural Forth word or a trainable transition function (neural networks in our case). We will call these trainable functions slots, as they correspond to underspecified “slots” in the program code that need to be filled by learned behaviour.

We allow users to define a slot $\mathbf{w}$ by specifying a pair of a state encoder $\mathbf{w}_{\text{enc}}$ and a decoder $\mathbf{w}_{\text{dec}}$. The encoder produces a latent representation $\mathbf{h}$ of the current machine state using a multi-layer perceptron, and the decoder consumes this representation to produce the next machine state. We hence have $\mathbf{w}=\mathbf{w}_{\text{dec}}\circ\mathbf{w}_{\text{enc}}$. To use slots within Forth program code, we introduce a notation that reflects this decomposition. In particular, slots are defined by the syntax { encoder -> decoder } where encoder and decoder are specifications of the corresponding slot parts as described below.

We provide the following options for encoders:

produces a static representation, independent of the actual machine state.

$\,e_{1}\ldots e_{m}$: concatenates the elements $e_{1}\ldots e_{m}$ of the machine state. An element can be a stack item Di at relative index $i$, a return stack item Ri, etc.

represent chained trans-formations, which enable the multilayer perceptron architecture. Linear N projects to N dimensions, and sigmoid and tanh apply same-named functions elementwise.

Users can specify the following decoders:

$\,w_{1}\ldots w_{m}$: chooses from the Forth words $w_{1}\ldots w_{m}$. Takes an input vector $\mathbf{h}$ of length $m$ to produce a weighted combination of machine states $\sum_{i}^{m}h_{i}\mathbf{w_{i}}(\mathbf{S})$.

$\,e_{1}\ldots e_{m}$: directly manipulates the machine state elements $e_{1}\ldots e_{m}$ by writing the appropriately reshaped and softmaxed output of the encoder over the machine state elements with $\text{write}_{\mathbf{M}}$.

$\,e_{1}\ldots e_{m}$: permutes the machine state elements $e_{1}\ldots e_{m}$ via a linear combination of $m!$ state vectors.

3 The Execution RNN

We model execution using an RNN which produces a state $\mathbf{S}_{n+1}$ conditioned on a previous state $\mathbf{S}_{n}$. It does so by first passing the current state to each function $\mathbf{w}_{i}$ in the program, and then weighing each of the produced next states by the component of the program counter vector $\mathbf{c}_{i}$ that corresponds to program index $i$, effectively using $\mathbf{c}$ as an attention vector over code. Formally we have:

Clearly, this recursion, and its final state, are differentiable with respect to the program code $\mathbf{P}_{\theta}$, and its inputs. Furthermore, for differentiable Forth programs the final state of this RNN will correspond to the final state of a symbolic execution (when no slots are present, and one-hot values are used).

4 Program Code Optimisations

The $\partial$4 RNN requires one-time step per transition. After each time step, the program counter is either incremented, decremented, explicitly set or popped from the stack. In turn, a new machine state is calculated by executing all words in the program and then weighting the result states by the program counter. As this is expensive, it is advisable to avoid full RNN steps wherever possible. We use two strategies to avoid full RNN steps and significantly speed-up $\partial$4: symbolic execution and interpolation of if-branches.

Whenever we have a sequence of Forth words that contains no branch entry or exit points, we can collapse this sequence into a single transition instead of naively interpreting words one-by-one. We symbolically execute (King 1976) a sequence of Forth words to calculate a new machine state. We then use the difference between the new and the initial state to derive the transition function of the sequence. For example, the sequence R> SWAP >R that swaps top elements of the data and the return stack yields the symbolic state $D=r_{1}d_{2}\ldots d_{l}$. and $R=d_{1}r_{2}\ldots r_{l}$. Comparing it to the initial state, we derive a single neural transition that only needs to swap the top elements of $\mathbf{D}$ and $\mathbf{R}$.

We cannot apply symbolic execution to code with branching points as the branching behaviour depends on the current machine state, and we cannot resolve it symbolically. However, we can still collapse if-branches that involve no function calls or loops by executing both branches in parallel and weighing their output states by the value of the condition. If the if-branch does contain function calls or loops, we simply fall back to execution of all words weighted by the program counter.

5 Training

Our training procedure assumes input-output pairs of machine start and end states $(\mathbf{x}_{i},\mathbf{y}_{i})$ only. The output $\mathbf{y}_{i}$ defines a target memory $\mathbf{Y}^{D}_{i}$ and a target pointer $\mathbf{y}^{d}_{i}$ on the data stack $\mathbf{D}$. Additionally, we have a mask $\mathbf{K}_{i}$ that indicates which components of the stack should be included in the loss (e.g. we do not care about values above the stack depth). We use $\mathbf{D}_{T}(\theta,\mathbf{x}_{i})$ and $\mathbf{d}_{T}(\theta,\mathbf{x}_{i})$ to denote the final state of $\mathbf{D}$ and $\mathbf{d}$ after $T$ steps of execution RNN and using an initial state $\mathbf{x}_{i}$. We define the loss function as

where $\mathcal{H}(\mathbf{x},\mathbf{y})=-\mathbf{x}\log{\mathbf{y}}$ is the cross-entropy loss, and $\theta$ are parameters of slots in the program $P$. We can use backpropagation and any variant of gradient descent to optimise this loss function. Note that at this point it would be possible to include supervision of the intermediate states (trace-level), as done by the Neural Program Interpreter (Reed & de Freitas 2015).

Experiments

We evaluate $\partial$4 on three tasks. Two of these are simple transduction tasks, sorting and addition as presented in (Reed & de Freitas 2015), with varying levels of program structure. For each problem, we introduce two sketches.

We also test $\partial$4 on the more difficult task of answering word algebra problems. We show that not only can $\partial$4 act as a standalone solver for such problems, bypassing the intermediary task of producing formula templates which must then be executed, but it can also outperform previous work when trained on the same data.

Specific to the transduction tasks, we discretise memory elements during testing. This effectively allows the trained model to generalise to any sequence length if the correct sketch behaviour has been learned. We also compare against a Seq2Seq (Sutskever et al. 2014) baseline. Full details of the experimental setup can be found in Appendix E.

2 Sorting

Sorting sequences of digits is a hard task for RNNs, as they fail to generalise to sequences even marginally longer than the ones they have been trained on (Reed & de Freitas 2015). We investigate several strong priors based on Bubble sort for this transduction task and present two $\partial$4 sketches in Listing 1 that enable us to learn sorting from only a few hundred training examples (see Appendix C.1 for more detail):

Permute. A sketch specifying that the top two elements of the stack, and the top of the return stack must be permuted based on the values of the former (line 3b). Both the value comparison and the permutation behaviour must be learned. The core of this sketch is depicted in Listing 1 (b lines), and the sketch is explained in detail in Appendix D.

Compare. This sketch provides additional prior procedural knowledge to the model. In contrast to Permute, only the comparison between the top two elements on the stack must be learned (line 3c). The core of this sketch is depicted in Listing 1 (c lines).

In both sketches, the outer loop can be specified in $\partial$4 (Listing 1, line 10), which repeatedly calls a function BUBBLE. In doing so, it defines sufficient structure so that the behaviour of the network is invariant to the input sequence length.

A quantitative comparison of our models on the Bubble sort task is provided in Table 1. For a given test sequence length, we vary the training set lengths to illustrate the model’s ability to generalise to sequences longer than those it observed during training. We find that $\partial$4 quickly learns the correct sketch behaviour, and it is able to generalise perfectly to sort sequences of $64$ elements after observing only sequences of length two and three during training. In comparison, the Seq2Seq baseline falters when attempting similar generalisations, and performs close to chance when tested on longer sequences. Both $\partial$4 sketches perform flawlessly when trained on short sequence lengths, but under-perform when trained on sequences of length 4 due to arising computational difficulties (Compare sketch performs better due to more structure it imposes). We discuss this issue further in Section 5.

3 Addition

Next, we applied $\partial$4 to the problem of learning to add two n-digit numbers. We rely on the standard elementary school addition algorithm, where the goal is to iterate over pairs of aligned digits, calculating the sum of each to yield the resulting sum. The key complication arises when two digits sum to a two-digit number, requiring that the correct extra digit (a carry) be carried over to the subsequent column.

We assume aligned pairs of digits as input, with a carry for the least significant digit (potentially $0$), and the length of the respective numbers. The sketches define the high-level operations through recursion, leaving the core addition to be learned from data.

The specified high-level behaviour includes the recursive call template and the halting condition of the recursion (no remaining digits, line 2-3). The underspecified addition operation must take three digits from the previous call, the two digits to sum and a previous carry, and produce a single digit (the sum) and the resultant carry (lines 6a, 6b and 7a, 7b). We introduce two sketches for inducing this behaviour:

Manipulate. This sketch provides little prior procedural knowledge as it directly manipulates the $\partial$4 machine state, filling in a carry and the result digits, based on the top three elements on the data stack (two digits and the carry). Depicted in Listing 2 in green.

Choose. Incorporating additional prior information, Choose exactly specifies the results of the computation, namely the output of the first slot (line 6b) is the carry, and the output of the second one (line 7b) is the result digit, both conditioned on the two digits and the carry on the data stack. Depicted in Listing 2 in blue.

The rest of the sketch code reduces the problem size by one and returns the solution by popping it from the return stack.

In a set of experiments analogous to those in our evaluation on Bubble sort, we demonstrate the performance of $\partial$4 on the addition task by examining test set sequence lengths of 8 and 64 while varying the lengths of the training set instances (Table 2). The Seq2Seq model again fails to generalise to longer sequences than those observed during training. In comparison, both the Choose sketch and the less structured Manipulate sketch learn the correct sketch behaviour and generalise to all test sequence lengths (with an exception of Manipulate which required more data to train perfectly). In additional experiments, we were able to successfully train both the Choose and the Manipulate sketches from sequences of input length $24$, and we tested them up to the sequence length of $128$, confirming their perfect training and generalisation capabilities.

4 Word Algebra Problems

Word algebra problems (WAPs) are often used to assess the numerical reasoning abilities of schoolchildren. Questions are short narratives which focus on numerical quantities, culminating with a question. For example:

A florist had 50 roses. If she sold 15 of them and then later picked 21 more, how many roses would she have?

Answering such questions requires both the understanding of language and of algebra — one must know which numeric operations correspond to which phrase and how to execute these operations.

Previous work cast WAPs as a transduction task by mapping a question to a template of a mathematical formula, thus requiring manuall labelled formulas. For instance, one formula that can be used to correctly answer the question in the example above is (50 - 15) + 21 = 56. In previous work, local classifiers (Roy & Roth 2015; Roy et al. 2015), hand-crafted grammars (Koncel-Kedziorski et al. 2015), and recurrent neural models (Bouchard et al. 2016) have been used to perform this task. Predicted formula templates may be marginalised during training (Kushman et al. 2014), or evaluated directly to produce an answer.

In contrast to these approaches, $\partial$4 is able to learn both, a soft mapping from text to algebraic operations and their execution, without the need for manually labelled equations and no explicit symbolic representation of a formula.

Our model is a fully end-to-end differentiable structure, consisting of a $\partial$4 interpreter, a sketch, and a Bidirectional LSTM (BiLSTM) reader.

The BiLSTM reader reads the text of the problem and produces a vector representation (word vectors) for each word, concatenated from the forward and the backward pass of the BiLSTM network. We use the resulting word vectors corresponding only to numbers in the text, numerical values of those numbers (encoded as one-hot vectors), and a vector representation of the whole problem (concatenation of the last and the first vector of the opposite passes) to initialise the $\partial$4 heap $\mathbf{H}$. This is done in an end-to-end fashion, enabling gradient propagation through the BiLSTM to the vector representations. The process is depicted in Figure 1.

The sketch, depicted in Listing 3 dictates the differentiable computation. Due to space constraints, we present the core of the sketch here. For the full sketch, please refer to Listing 4 in the Appendix. First, it copies values from the heap $\mathbf{H}$ – word vectors to the return stack $\mathbf{R}$, and numbers (as one-hot vectors) on the data stack $\mathbf{D}$. Second, it contains four slots that define the space of all possible operations of four operators on three operands, all conditioned on the vector representations on the return stack. These slots are i) permutation of the elements on the data stack, ii) choosing the first operator, iii) possibly swapping the intermediate result and the last operand, and iv) the choice of the second operator. The final set of commands simply empties out the return stack $\mathbf{R}$. These slots define the space of possible operations, however, the model needs to learn how to utilise these operations in order to calculate the correct result.

We evaluate the model on the Common Core (CC) dataset, introduced by Roy & Roth 2015. CC is notable for having the most diverse set of equation patterns, consisting of four operators (+, -, $\times$, $\div$), with up to three operands.

We compare against three baseline systems: (1) a local classifier with hand-crafted features (Roy & Roth 2015), (2) a Seq2Seq baseline, and (3) the same model with a data generation component (GeNeRe) Bouchard et al. 2016. All baselines are trained to predict the best equation, which is executed outside of the model to obtain the answer. In contrast, $\partial$4 is trained end-to-end from input-output pairs and predicts the answer directly without the need for an intermediate symbolic representation of a formula.

Results are shown in Table 3. All RNN-based methods (bottom three) outperform the classifier-based approach. Our method slightly outperforms a Seq2Seq baseline, achieving the highest reported result on this dataset without data augmentation.

Discussion

$\partial$4 bridges the gap between a traditional programming language and a modern machine learning architecture. However, as we have seen in our evaluation experiments, faithfully simulating the underlying abstract machine architecture introduces its own unique set of challenges.

One such challenge is the additional complexity of performing even simple tasks when they are viewed in terms of operations on the underlying machine state. As illustrated in Table 1, $\partial$4 sketches can be effectively trained from small training sets (see Appendix C.1), and generalise perfectly to sequences of any length. However, difficulty arises when training from sequences of modest lengths. Even when dealing with relatively short training length sequences, and with the program code optimisations employed, the underlying machine can unroll into a problematically large number states. For problems whose machine execution is quadratic, like the sorting task (which at input sequences of length 4 has 120 machine states), we observe significant instabilities during training from backpropagating through such long RNN sequences, and consequent failures to train. In comparison, the addition problem was easier to train due to a comparatively shorter underlying execution RNNs.

The higher degree of prior knowledge provided played an important role in successful learning. For example, the Compare sketch, which provides more structure, achieves higher accuracies when trained on longer sequences. Similarly, employing softmax on the directly manipulated memory elements enabled perfect training for the Manipulate sketch for addition. Furthermore, it is encouraging to see that $\partial$4 can be trained jointly with an upstream LSTM to provide strong procedural prior knowledge for solving a real-world NLP task.

Related Work

The idea of program synthesis is as old as Artificial Intelligence, and has a long history in computer science (Manna & Waldinger 1971). Whereas a large body of work has focused on using genetic programming (Koza 1992) to induce programs from the given input-output specification (Nordin 1997), there are also various Inductive Programming approaches (Kitzelmann 2009) aimed at inducing programs from incomplete specifications of the code to be implemented (Albarghouthi et al. 2013; Solar-Lezama et al. 2006). We tackle the same problem of sketching, but in our case, we fill the sketches with neural networks able to learn the slot behaviour.

Our work is closely related to probabilistic programming languages such as Church (Goodman et al. 2008). They allow users to inject random choice primitives into programs as a way to define generative distributions over possible execution traces. In a sense, the random choice primitives in such languages correspond to the slots in our sketches. A core difference lies in the way we train the behaviour of slots: instead of calculating their posteriors using probabilistic inference, we estimate their parameters using backpropagation and gradient descent. This is similar in-kind to TerpreT’s FMGD algorithm (Gaunt et al. 2016), which is employed for code induction via backpropagation. In comparison, our model which optimises slots of neural networks surrounded by continuous approximations of code, enables the combination of procedural behaviour and neural networks. In addition, the underlying programming and probabilistic paradigm in these programming languages is often functional and declarative, whereas our approach focuses on a procedural and discriminative view. By using an end-to-end differentiable architecture, it is easy to seamlessly connect our sketches to further neural input and output modules, such as an LSTM that feeds into the machine heap.

Recently, there has been a surge of research in program synthesis, and execution in deep learning, with increasingly elaborate deep models. Many of these models were based on differentiable versions of abstract data structures (Joulin & Mikolov 2015; Grefenstette et al. 2015; Kurach et al. 2016), and a few abstract machines, such as the NTM (Graves et al. 2014), Differentiable Neural Computers (Graves et al. 2016), and Neural GPUs (Kaiser & Sutskever 2015). All these models are able to induce algorithmic behaviour from training data. Our work differs in that our differentiable abstract machine allows us to seemingly integrate code and neural networks, and train the neural networks specified by slots via backpropagation. Related to our efforts is also the Autograd (Maclaurin et al. 2015), which enables automatic gradient computation in pure Python code, but does not define nor use differentiable access to its underlying abstract machine.

The work in neural approximations to abstract structures and machines naturally leads to more elaborate machinery able to induce and call code or code-like behaviour. Neelakantan et al. 2015a learned simple SQL-like behaviour–—querying tables from the natural language with simple arithmetic operations. Although sharing similarities on a high level, the primary goal of our model is not induction of (fully expressive) code but its injection. (Andreas et al. 2016) learn to compose neural modules to produce the desired behaviour for a visual QA task. Neural Programmer-Interpreters (Reed & de Freitas 2015) learn to represent and execute programs, operating on different modes of an environment, and are able to incorporate decisions better captured in a neural network than in many lines of code (e.g. using an image as an input). Users inject prior procedural knowledge by training on program traces and hence require full procedural knowledge. In contrast, we enable users to use their partial knowledge in sketches.

Neural approaches to language compilation have also been researched, from compiling a language into neural networks (Siegelmann 1994), over building neural compilers (Gruau et al. 1995) to adaptive compilation (Bunel et al. 2016). However, that line of research did not perceive neural interpreters and compilers as a means of injecting procedural knowledge as we did. To the best of our knowledge, $\partial$4 is the first working neural implementation of an abstract machine for an actual programming language, and this enables us to inject such priors in a straightforward manner.

Conclusion and Future Work

We have presented $\partial$4, a differentiable abstract machine for the Forth programming language, and showed how it can be used to complement programmers’ prior knowledge through the learning of unspecified behaviour in Forth sketches. The $\partial$4 RNN successfully learns to sort and add, and solve word algebra problems, using only program sketches and program input-output pairs. We believe $\partial$4, and the larger paradigm it helps establish, will be useful for addressing complex problems where low-level representations of the input are necessary, but higher-level reasoning is difficult to learn and potentially easier to specify.

In future work, we plan to apply $\partial$4 to other problems in the NLP domain, like machine reading and knowledge base inference. In the long-term, we see the integration of non-differentiable transitions (such as those arising when interacting with a real environment), as an exciting future direction which sits at the intersection of reinforcement learning and probabilistic programming.

We thank Guillaume Bouchard, Danny Tarlow, Dirk Weissenborn, Johannes Welbl and the anonymous reviewers for fruitful discussions and helpful comments on previous drafts of this paper. This work was supported by a Microsoft Research PhD Scholarship, an Allen Distinguished Investigator Award, and a Marie Curie Career Integration Award.

References

Appendix

Appendix A Forth Words and their implementation

We implemented a small subset of available Forth words in $\partial$4. The table of these words, together with their descriptions is given in Table 4, and their implementation is given in Table 5. The commands are roughly divided into 7 groups. These groups, line-separated in the table, are:

{num}, 1+, 1-, DUP, SWAP, OVER, DROP, +, -, *, /

IF..ELSE..THEN, BEGIN..WHILE..REPEAT, DO..LOOP

Appendix B Bubble sort algorithm description

An example of a Forth program that implements the Bubble sort algorithm is shown in Listing 1. We provide a description of how the first iteration of this algorithm is executed by the Forth abstract machine:

The program begins at line 11, putting the sequence on the data stack $D$, followed by the sequence length 4. Note that Forth uses Reverse Polish Notation and that the top of the data stack is 4 in this example. It then calls the SORT word.

For a sequence of length $4$, SORT performs a do-loop in line 9 that calls the BUBBLE sub-routine $3$ times. It does so by decrementing the top of $D$ with the 1- word to $3$. Subsequently, $3$ is duplicated on $D$ by using DUP, and $0$ is pushed onto $D$.

DO consumes the top two stack elements $3$ and $0$ as the limit and starting point of the loop, leaving the stack $D$ to be . We use the return stack $R$ as a temporary variable buffer and push $3$ onto it using the word >R. This drops $3$ from $D$, which we copy from $R$ with R@

Next, we call BUBBLE to perform one iteration of the bubble pass, (calling BUBBLE $3$ times internally), and consuming $3$. Notice that this call puts the current program counter onto $R$, to be used for the program counter $c$ when exiting BUBBLE.

Inside the BUBBLE subroutine, DUP duplicates $3$ on $R$. IF consumes the duplicated $3$ and interprets is as TRUE. >R puts $3$ on $R$.

Calling OVER twice duplicates the top two elements of the stack, to test them with <, which tests whether $2<7$. IF tests if the result is TRUE ($0$), which it is, so it executes SWAP.

To prepare for the next call to BUBBLE we move $3$ back from the return stack $R$ to the data stack $D$ via R>, SWAP it with the next element, put it back to $R$ with >R, decrease the TOS with 1- and invoke BUBBLE again. Notice that $R$ will accumulate the analysed part of the sequence, which will be recursively taken back.

When we reach the loop limit we drop the length of the sequence and exit SORT using the ; word, which takes the return address from $R$. At the final point, the stack should contain the ordered sequence .

Appendix C Learning and Run Time Efficiency

When measuring the performance of the model as the number of training instances varies, we can observe the benefit of additional prior knowledge to the optimisation process. We find that when stronger prior knowledge is provided (Compare), the model quickly maximises the training accuracy. Providing less structure (Permute) results in lower testing accuracy initially, however, both sketches learn the correct behaviour and generalise equally well after seeing $256$ training instances. Additionally, it is worth noting that the Permute sketch was not always able to converge into a result of the correct length, and both sketches are not trivial to train.

In comparison, Seq2Seq baseline is able to generalise only to the sequence it was trained on (Seq2Seq trained and tested on sequence length 3). When training it on sequence length 3, and testing it on a much longer sequence length of 8, Seq2Seq baseline is not able to achieve more than $45\%$ accuracy.

We tested the models to train on datasets of increasing size on the addition task. The results, depicted in Table 4 show that both the choose and the manipulate sketch are able to perfectly generalise from $256$ examples, trained on sequence lengths of $8$, tested on $16$. In comparison, the Seq2Seq baseline achieves $98\%$ when trained on $16384$ examples, but only when tested on the input of the same length, $8$. If we test Seq2Seq as we tested the sketches, it is unable to achieve more $19.7\%$.

C.2 Program Code Optimisations

We measure the runtime of Bubble sort on sequences of varying length with and without the optimisations described in Section 3.4. The results of ten repeated runs are shown in Figure 5 and demonstrate large relative improvements for symbolic execution and interpolation of if-branches compared to non-optimised $\partial$4 code.

Appendix D ∂\partial4 execution of a Bubble sort sketch

Listing 1 (lines 3b and 4b – in blue) defines the BUBBLE word as a sketch capturing several types of prior knowledge. In this section, we describe the Permute sketch. In it, we assume BUBBLE involves a recursive call, that terminates at length 1, and that the next BUBBLE call takes as input some function of the current length and the top two stack elements.

The input to this sketch are the sequence to be sorted and its length decremented by one, $n-1$ (line 1). These inputs are expected on the data stack. After the length ($n-1$) is duplicated for further use with DUP, the machine tests whether it is non-zero (using IF, which consumes the TOS during the check). If $n-1>0$, it is stored on the $R$ stack for future use (line 2).

At this point (line 3b) the programmer only knows that a decision must be made based on the top two data stack elements D0 and D-1 (comparison elements), and the top return stack, R0 (length decremented by 1). Here the precise nature of this decision is unknown but is limited to variants of permutation of these elements, the output of which produce the input state to the decrement -1 and the recursive BUBBLE call (line 4b). At the culmination of the call, R0, the output of the learned slot behaviour, is moved onto the data stack using R>, and execution proceeds to the next step.

Figure 2 illustrates how portions of this sketch are executed on the $\partial$4 RNN. The program counter initially resides at >R (line 3 in $\mathbf{P}$), as indicated by the vector $\mathbf{c}$, next to program $\mathbf{P}$. Both data and return stacks are partially filled ($\mathcal{R}$ has 1 element, $\mathcal{D}$ has 4), and we show the content both through horizontal one-hot vectors and their corresponding integer values (colour coded). The vectors $\mathbf{d}$ and $\mathbf{r}$ point to the top of both stacks, and are in a one-hot state as well. In this execution trace, the slot at line 4 is already showing optimal behaviour: it remembers the element on the return stack (4) is larger and executes BUBBLE on the remaining sequence with the counter $n$ subtracted by one, to 1.

Appendix E Experimental details

The parameters of each sketch are trained using Adam (Kingma & Ba 2015), with gradient clipping (set to 1.0) and gradient noise (Neelakantan et al. 2015b). We tuned the learning rate, batch size, and the parameters of the gradient noise in a random search on a development variant of each task.

The Seq2Seq baseline models are single-layer networks with LSTM cells of $50$ dimensions.

The training procedure for these models consists of 500 epochs of Adam optimisation, with a batch size of 128, a learning rate of 0.01, and gradient clipping when the L2 norm of the model parameters exceeded 5.0. We vary the size of training and test data (Fig. 3), but observe no indication of the models failing to reach convergence under these training conditions.

E.2 Sorting

The Permute and Compare sketches in Table 1 were trained on a randomly generated train, development and test set containing 256, 32 and 32 instances, respectively. Note that the low number of dev and test instances was due to the computational complexity of the sketch.

The batch size was set to a value between $64$ and $16$, depending on the problem size, and we used an initial learning rate of $1.0$.

E.3 Addition

We trained the addition Choose and Manipulate sketches presented in Table 2 on a randomly generated train, development and test sets of sizes 512, 256, and 1024 respectively. The batch size was set to $16$, and we used an initial learning rate of $0.05$

E.4 Word Algebra Problem

The Common Core (CC) dataset (Roy & Roth 2015) is partitioned into a train, dev, and test set containing 300, 100, and 200 questions, respectively. The batch size was set to $50$, and we used an initial learning rate of $0.02$. The BiLSTM word vectors were initialised randomly to vectors of length $75$. The stack width was set to $150$ and the stack size to $5$.

Appendix F Qualitative Analysis on BubbleSort of PC traces

In Figure 6 we visualise the program counter traces. The trace follows a single example from start, to middle, and the end of the training process. In the beginning of training, the program counter starts to deviate from the one-hot representation in the first $20$ steps (not observed in the figure due to unobservable changes), and after two iterations of SORT, $\partial$4 fails to correctly determine the next word. After a few training epochs $\partial$4 learns better permutations which enable the algorithm to take crisp decisions and halt in the correct state.

Appendix G The complete Word Algebra Problem sketch

The Word Algebra Problem (WAP) sketch described in Listing 3 is the core of the model that we use for WAP problems. However, there were additional words before and after the core which took care of copying the data from the heap to data and return stacks, and finally emptying out the return stack.

The full WAP sketch is given in Listing 4. We define a QUESTION variable which will denote the address of the question vector on the heap. Lines 4 and 5 create REPR_BUFFER and NUM_BUFFER variables and denote that they will occupy four sequential memory slots on the heap, where we will store the representation vectors and numbers, respectively. Lines 7 and 8 create variables REPR and NUM which will denote addresses to current representations and numbers on the heap. Lines 10 and 11 store REPR_BUFFER to REPR and NUM_BUFFER to NUM, essentially setting the values of variables REPR and NUM to starting addresses allotted in lines 4 and 5. Lines 14-16 and 19-20 create macro functions STEP_NUM and STEP_REPR which increment the NUM and REPR values on call. These macro functions will be used to iterate through the heap space. Lines 24-25 define macro functions CURRENT_NUM for fetching the current number, and CURRENT_REPR for fetching representation values. Lines 28-32 essentially copy values of numbers from the heap to the data stack by using the CURRENT_NUM and STEP_NUM macros. After that line 35 pushes the question vector, and lines 36-40 push the word representations of numbers on the return stack.

Following that, we define the core operations of the sketch. Line 43 permutes the elements on the data stack (numbers) as a function of the elements on the return stack (vector representations of the question and numbers). Line 45 chooses an operator to execute over the TOS and NOS elements of the data stack (again, conditioned on elements on the return stack). Line 47 executes a possible swap of the two elements on the data stack (the intermediate result and the last operand) conditioned on the return stack. Finally, line 49 chooses the last operator to execute on the data stack, conditioned on the return stack.

The sketch ends with lines 52-55 which empty out the return stack.