Latent Execution for Neural Program Synthesis

At a high level, the program decoder (Fig. 2(b)) takes a latent vector $h_{t-1}$ that represents the generated partial program, the previous (generated) program token $p_{t-1}$, and outputs the latent vector $h_{t}$ and the next program token $p_{t}$ to be generated at time step $t$:

Here the recurrent model is conditioned on the IO pair $IO_{t-1}$. When $IO_{t}=IO:=(I,O)$ for every $t$, i.e., $IO_{t}$ remains constant over the entire recurrent generation process, Eqn. 1 represents the standard recurrent architecture used in most autoregressive natural language models [24, 49], and is also used in prior works on program synthesis from input-output examples [17, 9].

For program decoding, the decoder first takes two attention vectors $s_{t}^{I}$ and $s_{t}^{O}$ computed from IO pairs and latent vector $h_{t-1}$ via double attention [17], and utilizes a max pooling layer to compute an aggregated embedding $m_{t}$ for all IO pairs (Fig. 2(b)):

Here the superscript $(j)$ indicates that the representation is for the $j$-th IO pair, $[a;b]$ is vector concatenation of $a$ and $b$, and $W$ is a trainable matrix. To facilitate the prediction of long programs, we compute an attention vector $d_{t}$ over previously generated program tokens using the standard attention mechanism [5, 30]:

Finally, the next token $p_{t}$ is sampled from $\mathbb{P}[p_{t}]=\mathrm{Softmax}(Vd_{t})_{p_{t}}$ where $V$ is a trainable matrix.

As shown in our experiments (Sec. 5), the standard program decoder architecture may not be able to achieve strong performance in program synthesis when the program complexity increases. One main reason is that the standard program decoder only takes the initial IO pairs as the input without considering the program execution, thus the learned representation for the partial program does not effectively guide the synthesis process. Motivated by prior works that utilize execution traces for Karel program synthesis [12, 44, 47], in this paper, we introduce latent executor (Fig. 2(c)) which maintains a second representation $\hat{I}_{t}$ during program decoding. Intuitively, $\hat{I}_{t-1}$ models the hypothetical input of the partial program $p_{t\ldots T}$ so that its output becomes $O$. Given the estimated input $\hat{I}_{t-1}$ and the latent vector $h_{t}$, the latent executor returns $\hat{I}_{t}$ at the next time step $t$:

The collection of $\{\hat{I}_{t}\}_{t=0}^{T}$ is the latent execution trace (LaET). With the help of latent executor, we now use the IO pairs $IO_{t-1}:=(\hat{I}_{t-1},O)$ instead of $(I,O)$ for the program decoder (Eqn. 1).

We train our model with supervised learning, by minimizing the sum of token prediction loss $\mathcal{L}_{Prog}$, and the latent executor loss $\mathcal{L}_{Exec}$:

Specifically, $\mathcal{L}_{Prog}:=\sum_{t=1}^{T}\mathrm{Loss}(p_{t},p^{\star}_{t})$ is the step-by-step cross-entropy loss between the predicted programs $p_{1\ldots T}$ and the ground truth programs $p^{\star}_{1\ldots T}$.

For latent executor, since the semantics of partial programs (e.g., partial C programs) are not always well-defined, there is no step-by-step training supervision. However, the output of the executor should be consistent with the program specification after taking the annotated ground truth program as the input. Therefore, we set $\hat{I}_{0}=I$ (true input) and minimize the distance between $\hat{I}_{T}$ and $O$ (true output) after the program finishes:

Note that $\mathcal{L}_{Exec}$ does not rely on any assumptions of the partial program semantics, and thus is applicable to both domain-specific languages and general-purpose programming languages such as C. In our evaluation, equipping with the latent executor significantly improves the program prediction performance, where each program could include up to 256 tokens.

Interestingly, once our model is trained on the initial random generated programs $\mathcal{D}_{0}$, the predicted program becomes more concise and resembles human-written code. While the exact token match accuracy is low even on the training set, the model still satisfies the IO pairs for many problems. We leverage such a phenomenon to construct a new dataset $\mathcal{D}_{1}$ with higher-quality programs from $\mathcal{D}_{0}$. Specifically, we run beam search on the trained model to predict program $p_{0\ldots T}$ given input-output pairs in the training set. If model prediction $p_{0\ldots T}$ satisfies all the input-output examples and held-out cases, we replace the original program $p^{\star}_{0\ldots T}$ with $p_{0\ldots T}$ in $\mathcal{D}_{1}$, and keep $p^{\star}_{0\ldots T}$ otherwise. Afterward, we re-train the model on $\mathcal{D}_{1}$. In Sec. 5, we will demonstrate that the retraining process further improves the model performance, especially with smaller training datasets.

Collecting large-scale high-quality datasets for program synthesis requires a lot of human efforts, and we aim to reduce the manual work for dataset construction.

Our data generator is built upon Csmith [54], a random C code generation tool originally designed for finding bugs in compilers. Following the common practice of generating input-output pairs, for each program, we randomly sample 5 numerical lists as the program inputs, and execute the program to obtain the corresponding output lists. This is similar to existing works on PBE problems that sample programs based on a probabilistic context-free grammar, randomly generate valid inputs for the programs and obtain the outputs [40, 15, 6]. This creates infinite samples for synthesizing programs in domain-specific languages. While the programs sampled in this way differ from human-written code, Sec. 3.4 shows that they can be converted to be more concise and human-like.

The subset of language features used. Our generated program has variable declaration, variable assignment, and expressions with addition or subtraction operations. The programs also have non-sequential statements, including If statements, For loops, Continue and Break statements. Except for the input argument which is a list, all variables declared are integers, and all program statements are integer manipulation. Each expression has at most 2 mathematical operations, and chaining the full C program could perform multi-step numerical calculation (e.g., p0 = p0 - p1 + p2; p0 = p0 - 1;). Looping statements other than For (i.e., While or Do-While loops) are not supported. Note that we only constrain the final program length ($\leq 256$ tokens) and the program can have nested for-loops and complicated if-conditions.

Post-processing. We perform a few post-processing steps to obtain our final programs from programs generated by Csmith (see Fig. 1 for an example). We resample different components of the program, so that (1) each constant numerical value lies in $[-4,4]$, (2) mathematical operators only contain addition and subtraction, and (3) upper/lower limits of For loops are positive and within the length of the list. Programs are discarded if they are trivial (e.g., constant or identity mappings), or the input-output examples include values out of the range $[-4,4]$.

Final dataset. We reweight the program distribution so that at least half of them include For loops. Our full dataset includes $500K$ samples in the training set, $1K$ samples in the validation set, and $1K$ samples in the test set. As shown in Fig. 1, the randomly sampled program may contain redundant statements, which can be easily avoided by human programmers. We compare our restricted C domain to prior datasets of programming by example in Table 1.

For program decoder, predicting the next program token $p_{t}$ is non-trivial, especially when mathematical reasoning is required [43, 29]. To improve the program synthesis performance for domains involving numerical calculation, such as our restricted C domain, we design an associative memory structure named operation predictor (Fig. 2(d)), based on the following intuition: given the input $I=2$ and output $O=4$, human would infer that “$O=I+2$” might be the desired operation and write down the code accordingly. To materialize such an intuition, we create a pre-computed table that covers all possible integer addition and subtraction operations for valid input and output list values. We defer the details of the model architecture to Appendix B.2. The program decoding process remains similar to the one described in Sec. 3, and we highlight the key differences as follows.

The operation predictor takes two attention vectors $s_{t}^{I}$ and $s_{t}^{O}$ as the representations of input-output examples, and yields an operator embedding $\hat{op}_{t}$. To compute the aggregated embedding vector for all input-output examples, we modify Eqn. 2 to also take $\hat{op}_{t}$ as an input of the max pooling layer:

To train the operation predictor, we add an additional loss $\mathcal{L}_{Op}$:

$\mathcal{L}_{Op}$ is designed to ensure that the operation predictor predicts operations related to IO pairs, and we defer the details to Appendix B.2.

Limitations. In our current implementation of the operation predictor, the operation table is only able to enumerate the arithmetic operations over a pre-defined constant set, thus it requires that the set of possible numerical values in input-output pairs is finite. One way of extending our operation predictor to support potentially unbounded numerical calculation is to combine it with the subword tokenizer, which has been commonly used in recent language models [16, 11, 4]. We consider designing general-purpose number representation for better mathematical reasoning as future work.

In Fig. 3, we show the effectiveness of retraining on decoded Karel programs (Sec. 3.4). We observe that retraining for one iteration is sufficient, and it significantly improves the generalization accuracy by over $3\%$. To understand the differences between predicted programs and randomly generated programs, we demonstrate the changes of dataset distributions after each retraining iteration in Fig. 4(a) and 4(b). We observe that the model learns to predict more concise programs than the ground truth for a large proportion of input-output examples, and considerably alters the dataset distribution so that it becomes more concentrated on short programs with simplified control flow structures. Specifically, from Fig. 4(a), although the initial Karel dataset seems to include a large proportion of complicated programs with different control flow constructs, our model synthesizes straight-line programs for nearly half of the samples, which means that many loops and branches in the annotated ground truth programs are unnecessary. This distribution shift also explains the gap between the exact match and generalization accuracies. The program distribution after the second retraining iteration is largely similar to the first iteration, thus retraining for more iterations does not considerably improve the performance. Note that in the second iteration, the synthesizer tends to generate slightly more complicated programs than the first iteration, in order to deal with the cases when the input-output examples oversimplify the intended program functionality. For example, sometimes the input-output examples do not cover the edge cases that the robot may encounter, thus adding additional If branches could avoid the crashes when testing on held-out cases.

Fig. 7(a) presents the results of retraining on decoded C programs. Similarly, retraining improves the prediction accuracy, especially when the training set is small. From Fig. 7(b) and 8(a), we again observe that the model tends to predict shorter programs than the random code, and it eliminates unnecessary control flows to simplify the programs. We present more examples in Appendix D.