SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Value Module We maintain a set of explored values  $S$  which contains variable tokens for constructing lambdas, lambda expressions, and non-lambda expressions including constants and input variables. The Value Module embeds each element of  $S$  forming a matrix of embeddings  $E_{S} \in \mathbb{R}^{|S| \times d}$ . Elements of  $S$  are embedded as follows. A variable token is embedded as a vector of trainable parameters. Note that the set of such variable tokens is fixed and determined by the DSL. $^3$  A lambda expression is embedded by  $\mathbf{s} + \mathbf{z}$ , where  $\mathbf{s}$  is the property signature of this lambda function encoded by an MLP, and  $\mathbf{z}$  is an embedding of the weight of this value. Non-lambda expressions are embedded like lambda expressions except using a different MLP to encode their property signatures.
**Argument Selector Module** Given an operator  $op$ , we use an operator-specific  $\mathrm{LSTM}_{op}$  to select the arguments using a pointer mechanism [35] from the encoded value matrix  $E_{S}$ , in an autoregressive way. In addition to selecting  $\mathrm{arity}(op)$  arguments, for an argument  $a_{k}$  that is a lambda expression, we also need to predict the variables  $i_{k}$  as required for MERGE, where  $i_{k}$  is a tuple of  $\mathrm{arity}(a_{k})$  variable tokens. All of the  $a_{k}$  and  $i_{k}$  predictions are done as a single autoregressive sequence. Furthermore,
3The higher-order functions in our DSL expect lambdas with at most 2 variables. Thus, it is unnecessary to create lambdas with  $3+$  variables, so the only variables needed for MERGE are  $v_{1}, v_{2}, u_{1}$ , and  $u_{2}$ .
for convenience we predict all of the  $a_{k}$  arguments first, followed by the  $i_{k}$  variable tuples which are constrained (via masking and padding) to include exactly arity  $(a_{k})$  variable tokens.
Search with Restarts We also change the inference time search procedure. Recall that CROSS-BEAM uses random sampling during evaluation, making the search nondeterministic. In LAMBDA-BEAM, instead of performing one synthesis search until timeout, we restart the search from scratch whenever the search has run for a certain amount of time without finding a solution. Even though this discards work done in previous searches, in practice this helps LAMBDA-BEAM solve more tasks because it may be otherwise difficult to recover from exploring the wrong search direction.
In this section, we experimentally evaluate the effectiveness of LAMBDABEAM, comparing to prior neural and symbolic approaches in an integer list manipulation domain.
4.1 Integer List Manipulation DSL
To measure a synthesizer's ability to create and use lambda expressions, we create a domain-specific language (DSL) that emphasizes lambda expressions and higher-order functions. Specifically, the DSL from DeepCoder [3] includes higher-order functions and has been used in subsequent work [37, 34]. However, DeepCoder's DSL only contains a hardcoded set of lambda functions and is not expressive enough to fully exercise LAMBDABEAM's ability to create arbitrary lambda expressions. Therefore, we extend DeepCoder's DSL by allowing lambda expressions to include arbitrary compositions of DSL operations, and replacing the hardcoded lambda functions with DSL operations and literal constants that enable a superset of the original functionality. For example, DeepCoder's original DSL includes hardcoded lambdas such as  $(\lambda x.x - 1)$ ,  $(\lambda x.y.x - y)$ , and  $(\lambda x.y.\max \{x,y\})$ . By introducing first-order functions including subtract and max, and constant literals including 0 and 1, we can create the hardcoded lambdas as well as lambdas like  $(\lambda x.\max \{x,0 - x\})$  that were not possible in the original DeepCoder DSL. Additionally, we add a new if-then-else operation, which further enriches our space of possible programs. The full DSL contains 23 first-order operations, 5 higher-order operations, and 6 integer literals, described fully in Appendix C. In our DSL, all integers are in the range  $[-256, 255]$  as in DeepCoder, and lists have lengths in the range [0, 10].
4.2 Experimental Setup
Training Data Similar to previous works including CROSSBEAM, we create synthetic training data by generating random tasks within our DSL. This is done by performing exhaustive bottom-up searches starting from random inputs and enumerating programs in order of increasing weight, and then sampling a subset of the resulting programs to serve as training tasks. Each task has between 2 and 5 I/O examples and between 1 and 3 input variables, and we sample tasks such that approximately  $80\%$  of them use lambdas in the ground-truth program. We used a time limit of 1 hour per data generation search (reaching programs of weight at most 12), sampling up to 1600 tasks per search, and performing enough searches parallelized across cloud CPU workers such that training uses less than 1 epoch over the dataset. We furthermore removed from the training dataset all programs that would solve any of our 200 evaluation tasks, described below.
Evaluation Tasks For evaluation, we use 100 handwritten evaluation tasks plus 100 synthetically generated tasks, with a time limit of 10 minutes per task. The handwritten tasks include all 9 "example programs" from Appendix A of the DeepCoder paper [3], plus other tasks that we created from scratch by brainstorming many natural but varied tasks that we could solve using our DSL (near the end of this process, it became quite difficult to come up with new tasks that were not merely slight variations of existing ones). Handwritten tasks include between 1 and 3 input variables, 3 I/O examples if the output is a list or 5 I/O examples if the output is an integer, and a handwritten ground-truth solution that has minimal weight to our knowledge. When creating I/O examples, we aimed to make the examples informative but reasonably succinct with lists of length at most 10. Every DSL operation is used in the solution for at least 4 handwritten tasks, and every higher-order operation is used in at least 10. For the 100 synthetic tasks, we sampled distinct random programs using the same procedure as for generating training data, except also enforcing that we have exactly
10 tasks of each weight between 3 and 12 inclusive. Appendix D contains example handwritten and synthetic tasks, along with LAMBDABEAM's solutions for them.
Approaches Our experiments compare several approaches:
1. LAMBDBEAM with random restarts: We trained the LAMBDBEAM model using on-policy training as in CROSSBEAM. The model has about 13 million trainable parameters and was trained on about 6.5 million tasks, which took about a week of training using 8 V100 GPUs. See Appendix E for more details on the model architecture and training. During evaluation, we use only 1 V100 GPU and perform random restarts every 6 seconds on the handwritten tasks, or every 30 seconds on the synthetic tasks, both chosen from a coarse search over restart frequencies. We run this approach for 5 trials using different randomness for the UniqueRandomizer sampling method carried over from CROSSBEAM.
2. LAMBDABEAM without random restarts: We use the LAMBDABEAM approach without random restarts as an ablation, also for 5 trials with different randomness for UniqueRandomizer sampling.
3. Enumeration: We run the same exhaustive bottom-up enumerative synthesis algorithm that was used to create the training data. We use 5 trials with different random orderings of DSL operations, which changes the enumeration order of programs with the same weight.
4. RobustFill [12]: This approach treats the synthesis problem as a sequence to sequence prediction task from the I/O examples to the program tokens. Specifically, we train a plain 3-layer LSTM-based encoder-decoder model on our synthetic training dataset, using approximately the same number of trainable parameters and training tasks as for the LAMBDABEAM model. We get model predictions via a single beam search of size 65536 which nearly exhausts the GPU memory and evaluate all resulting programs on the I/O examples. Since the beam search is deterministic, we perform 5 trials by re-training the model with different initializations.
5.  $\lambda^2$  [15]: This is a state-of-the-art symbolic program synthesizer that handles lambda functions and higher-order functions. We implemented our DSL within the  $\lambda^2$  framework (using the more extensible version provided by the  $\lambda^2$  authors).  $\lambda^2$  is deterministic so we only use 1 trial.
6. Python-Finetuned LLM: We try asking a pretrained large language model (LLM) to solve our evaluation tasks using Python code. Specifically, we use PaLM 62B that was trained for longer as described in Appendix F of Chowdhery et al. [9], with further fine-tuning on general Python code. The prompt contains natural language instructions and 2 examples of an I/O specification followed by Python code that solves the task (for 2 new tasks), and then the I/O specification of the evaluation task we wish to solve. We repeatedly draw batches of 16 independent samples with temperature sampling and run those programs on the I/O examples, until a solution is found or timeout is reached. We ran the LLM using 16 accelerators so this approach uses significantly more compute than the others over the same time limit. We repeat for 3 trials with different randomness for temperature sampling.
Figure 4 plots the synthesis performance of the various methods over time. Notably, LAMBDABEAM with restarts is the best approach for both handwritten and synthetic tasks. The gap is wider on the handwritten tasks where LAMBDABEAM with restarts solves 67.2 out of 100 tasks on average, which is  $24\%$  more tasks than the next best method  $\lambda^2$ . Figure 5 plots the various approaches' success rates for different task weights, which can be used as a measure of task difficulty. As expected, we observe that all methods perform worse on harder tasks with larger weight, but that LAMBDABEAM with restarts generally achieves higher success rates on the difficult tasks compared to other methods. This means that our approach scales better to larger programs compared to the other methods, except the LLM which has a more constant but lower success rate overall.
4 In preliminary experiments, we found that providing 3 few-shot examples in the prompt led to slower sampling without much change in program quality. On the other hand, using only 1 example led to noticeably worse program quality. We also tried asking for programs within our DSL via few-shot examples, but this was not as successful because the LLM was not trained on our DSL.
The LLM predicts Python code instead of using our DSL, so the weight according to our DSL is an inaccurate measure of the complexity of the corresponding Python code. Furthermore, difficulty for the LLM is more closely correlated with how "natural" the task is, i.e., its similarity to programs in the LLM's training data.
Figure 4: Synthesis results over time for various methods on the handwritten and synthetic evaluation tasks. Shaded areas represent the minimum and maximum across trials for that method.
Figure 5: Success rates of various methods broken down by the task weight (i.e., the smallest known weight of a solution). Task weights are bucketed such that each group contains at least 15 tasks.
Running LAMBDABEAM with random restarts helps overall but more so for the handwritten tasks. We believe this is because the synthetic evaluation tasks have the same distribution as the training tasks while the handwritten tasks are different. So, for the handwritten tasks, exploring the wrong part of the search space early on might cause further mistakes that lead the search astray, while the model may be better trained to stay on track for the synthetic tasks. This would also explain why more frequent restarts are effective for the handwritten tasks. We note that random restarts would not be possible for  $\lambda^2$ , enumeration, or RobustFill's beam search, and would not help the LLM where each sample is already independent.
We also identify false positives by running solutions on 2 held-out test cases per task, generated mostly synthetically with some manual editing. The results are in Figure 6, showing that LAMBDABEAM with restarts has the highest number of true positive solutions on handwritten tasks by a margin of nearly 8 tasks, while barely losing to Enumeration on synthetic tasks. While symbolic approaches (Enumeration and  $\lambda^2$ ) have fewer false positives due to focusing on small solutions, we observe that LAMBDABEAM has the fewest false positives among the neural approaches. The LLM produces many false positives on the synthetic tasks where the ground-truth solutions are less similar to programs seen during its training, and in fact many of its false positive solutions are if-elif-else chains that hardcode the examples in some way (which is feasible to implement in Python but not in our DSL). Finally, we note that some false positive solutions could be transformed into true positives with a postprocessing step, e.g., one that attempts to simplify or minimize subtrees of the solution. In this sense, false positive solutions may still be useful for synthesis, and LAMBDABEAM with restarts achieves the highest number of total positive solutions by a wide margin.
Appendix F contains analysis showing some of the differences in distributions between the handwritten and synthetic evaluation tasks, which helps to contextualize the experimental results. For example, lambda expressions are used in  $85\%$  of the handwritten tasks but only  $53\%$  of the synthetic tasks. The median task weight is 9 for handwritten tasks and only 7.5 for synthetic tasks. These comparisons suggest that the handwritten tasks are harder than the synthetic tasks on average, which
The synthetic tasks have randomly-generated I/O examples that are overall less informative than those in the handwritten tasks, and the "correct" solution is not chosen to be natural but rather is one with minimal weight by construction. Enumeration has an unfair advantage here, being the only method in our comparison that is guaranteed to find a minimal weight solution.
Figure 6: True positive versus false positive solutions measured using held-out test cases.
is also reflected in the overall performance in Figure 4. We observe that LAMBDABEAM achieves a greater performance gap over the other approaches on the handwritten tasks versus on the synthetic tasks, which is a promising trend because the handwritten tasks are both harder and more natural.
Although LAMBDABEAM resolves CROSSBEAM's limitations of not handling lambdas or looping computations, some other limitations are carried over. On-policy training is slow due to performing search during training, but this could be addressed with an initial training phase of off-policy teacher forcing. At evaluation time, even with UniqueRandomizer sampling to avoid duplicate argument lists within one sampling phase, our approach still encounters many duplicate values across sampling phases and across restarts. Finally, our DSL is small compared to general programming languages.
Machine learning for program synthesis has been an active area [17, 18, 1]. Within programming by example, deep learning architectures for sequences, like LSTMs and Transformers, have been particularly effective [12]. Our work builds upon CROSSBEAM [33], which itself combines three lines of research in program synthesis. The first are learned search strategies for program synthesis, that is, using a learned policy or value function to guide search [36, 20, 13], or multi-level strategies that combine the results of search over different spaces [25, 22, 19, 34]. The second are execution-guided neural synthesis methods, which guide the search over partial programs by evaluating them [37, 13, 7, 27, 8]. Finally, CROSSBEAM's use of imitation learning to train the policy is inspired by work in learning to search [11, 29, 5] and beam-aware training [23, 24].
In contrast, we are unaware of previous work that synthesizes helper functions, such as lambda functions, during neural program synthesis. The original DeepCoder DSL contains only a small set of predefined lambda functions. Even within symbolic program synthesis,  $\lambda^2$  is one of the few examples of work that synthesizes lambda expressions [15]. To control the size of the search space,  $\lambda^2$  employs type-directed synthesis, but we handle more general domains where the type system is not informative enough to reduce the search space sufficiently. DreamCoder [14] can also infer ad-hoc helper functions like  $\lambda^2$ , but its neural network provides no fine-grained guidance on how to compose those lambdas. Because DreamCoder is an algorithm for enriching an impoverished DSL to improve a neurally-guided program search, one could combine DreamCoder's DSL enrichment process with LAMBDABEAM's search strategy. Other work reuses fragments of code from partially-correct solutions [30, 34], but these are executable portions of straightline code, not lambda functions.
Our integer manipulation domain is inspired by DeepCoder [3] and subsequent work [37, 34].
We introduced the first neural search method for programming by example that is able to synthesize intermediate helper functions (lambdaes) by resolving two key difficulties. First, we algebraically represent lambda expressions in a canonical way and construct new lambdas with the MERGE operator that enforces desirable representational constraints. Second, we encode arbitrary lambda functions as inputs to a neural network by using property signatures to analyze the lambda's execution semantics. With these innovations, LAMBDABEAM learns a neural policy to drive a bottom-up search over programs. We experimentally show that LAMBDABEAM outperforms symbolic search, a sequence model, and a pretrained code LLM with 62 billion parameters.
The authors would like to thank Henryk Michalewski for his thoughtful ideas, and Christian Walder, Rif Saurous, and the anonymous reviewers for their helpful comments.
[29] Stephane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Conference on Artificial Intelligence and Statistics (AISTATS), 2011.
[30] Kensen Shi, Jacob Steinhardt, and Percy Liang. FrAngel: Component-based synthesis with control structures. Proceedings of the ACM on Programming Languages, 3(POPL), 2019.
[31] Kensen Shi, David Bieber, and Charles Sutton. Incremental sampling without replacement for sequence models. In International Conference on Machine Learning (ICML), 2020.
[32] Kensen Shi, David Bieber, and Rishabh Singh. TF-Coder: Program synthesis for tensor manipulations. ACM Transactions on Programming Languages and Systems (TOPLAS), 44(2): 1-36, 2022.
[33] Kensen Shi, Hanjun Dai, Kevin Ellis, and Charles Sutton. CrossBeam: Learning to search in bottom-up program synthesis. In International Conference on Learning Representations (ICLR), 2022.
[34] Disha Shrivastava, Hugo Larochelle, and Daniel Tarlow. Learning to combine per-example solutions for neural program synthesis. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
[35] Oriol Vinyals, Meire Fortunato, and Navdeep Jaitly. Pointer networks. In Advances in Neural Information Processing Systems (NIPS), 2015.
[36] Pengcheng Yin and Graham Neubig. A syntactic neural model for general-purpose code generation. In Association for Computational Linguistics (ACL), 2017.
[37] Amit Zohar and Lior Wolf. Automatic program synthesis of long programs with a learned garbage collector. In Advances in Neural Information Processing Systems (NeurIPS), 2018.
A Completeness of MERGE
The Merge operation is complete in the sense that it can generate all possible solutions in the domain-specific language (DSL) for a programming-by-example (PBE) problem.
We formalize our DSL in a subset of the lambda calculus. Let  $\mathcal{X} = \{x_1, \ldots, x_m\}$  be the set of input variables for the PBE task,  $\mathcal{V}$  be a countable set of variables that is disjoint from  $\mathcal{X}$ ,  $\mathcal{F}$  be the set of primitive functions in the DSL, and  $\mathcal{C}$  be a set of constants in the DSL. Then, our lambda calculus is:
$$
T := x \mid v \mid c \mid f (\mathbf {t} _ {1}, \dots , \mathbf {t} _ {k}) \mid \lambda v _ {1} \dots v _ {n}. \mathbf {t}
$$
$$
f o r \quad x \in \mathcal {X}, c \in \mathcal {C}, f \in \mathcal {F},
$$
$$
v, v _ {1}, \dots , v _ {n} \in \mathcal {V},
$$
$$
\mathbf {t}, \mathbf {t} _ {1}, \dots , \mathbf {t} _ {k} \in T.
$$
Let  $M$  be the set of terms obtainable by repeatedly applying MERGE (including the initial terms usable by MERGE):
$$
M := x \mid \lambda v. v \mid c \mid \operatorname {M E R G E} (f, \mathbf {a} _ {1}, i _ {1}, \dots , \mathbf {a} _ {k}, i _ {k})
$$
$$
\begin{array}{l} \text {f o r} x \in \mathcal {X}, v \in \mathcal {V}, c \in \mathcal {C}, f \in \mathcal {F}, \end{array}
$$
$$
\mathbf {a} _ {1}, \dots , \mathbf {a} _ {k} \in M,
$$
$$
i _ {1}, \dots , i _ {k} \in \mathcal {V} ^ {*}.
$$
We restate the definition
$$
\operatorname {M E R G E} (f, \mathbf {a} _ {1}, i _ {1}, \mathbf {a} _ {2}, i _ {2}, \dots) = \lambda v _ {1} v _ {2} \dots . f (\lambda u _ {1} u _ {2} \dots u _ {\ell_ {1}}. \mathbf {a} _ {1} (i _ {1}), \lambda u _ {1} u _ {2} \dots u _ {\ell_ {2}}. \mathbf {a} _ {2} (i _ {2}), \dots)
$$
$$
\text {w h e r e} \{v _ {1}, v _ {2}, \dots \} = \bigcup_ {k} i _ {k} - \left\{u _ {j} \right\} _ {j = 1} ^ {\max  \left\{\ell_ {1}, \ell_ {2}, \dots \right\}}
$$
From this definition, it is clear that  $M \subseteq T$ , i.e., MERGE is closed within the lambda calculus. However,  $M \neq T$  because MERGE imposes certain constraints, e.g.,  $\lambda v \cdot x$  is in  $T$  but cannot be constructed by MERGE. To precisely describe the constraints resulting from MERGE, we introduce the following definitions:
- Terms  $x \in \mathcal{X}, v \in \mathcal{V}$ , and  $c \in \mathcal{C}$  are atomic.
- A term  $\mathbf{s} = \lambda v_{1}\dots v_{n}$ .  $\mathbf{t}$  (possibly with  $n = 0$  such that  $\mathbf{s}$  is not a lambda expression) has exact lambda variables if  $\mathrm{FreeVars}(\mathbf{t}) - \mathcal{X} = \{v_1,\ldots ,v_n\}$ . Note that  $\mathbf{s}$  having exact lambda variables implies that  $\mathrm{FreeVars}(\mathbf{s})\subseteq \mathcal{X}$ .
- A term typechecks if every function application has the correct arity for every argument, e.g.,  $\operatorname{Map}(\mathbf{t}_1, \mathbf{t}_2)$  expects  $\mathbf{t}_1$  to have arity 1, while  $\mathbf{t}_2$  should have arity 0.
We now let  $S = \{\mathbf{s} \in T \mid \mathbf{s} \text{ has exact lambda variables and typechecks}\}$ .
The completeness of MERGE, in the sense that it can generate all solutions  $s$  to PBE problems (once the input variables  $x_{1},\ldots ,x_{m}$  are bound), follows from the two claims below.
Claim 1. If  $\mathbf{p} = \lambda x_1\ldots x_m$ .  $\mathbf{s}$  is a solution to a PBE problem, such that  $\mathbf{p}(x_1,\dots,x_m) = y$  for all I/O examples  $(x_{1},\ldots ,x_{m})\to y$  in the PBE specification, then  $\mathbf{s}\in S$ . That is,  $S$  is broad enough to cover all solutions to PBE problems.
Proof. To ensure that  $\mathbf{p} = \lambda x_1\ldots x_m$ ,  $\mathbf{s}$  is a valid solution program, we must have FreeVars  $(\mathbf{s}) - \mathcal{X} = \emptyset$  (so there are no unbound variables),  $\mathbf{s}$  must have arity 0 (since all inputs  $x_1,\dots,x_m$  are already bound), and  $\mathbf{s}$  must typecheck to avoid runtime errors. Together, these imply that  $\mathbf{s} \in S$ .
Claim 2.  $S \subseteq M$ . That is, MERGE can create any term in  $S$ , including all solutions to PBE problems.
Proof. Let  $s$  be any term in  $S$ . We will proceed by induction on the depth of  $s$ .
As the base case, if  $\mathbf{s}$  is atomic, then  $\mathbf{s} = x$  or  $\mathbf{s} = c$ , so  $\mathbf{s}$  is immediately in  $M$ . Note that  $\mathbf{s}$  cannot be  $v$  because  $v$  does not have exact lambda variables.