#  LambdaBeam: Neural Program Search with Higher-Order Functions and Lambdas

Kensen Shi
Google DeepMind
<EMAIL_ADDRESS>
&Hanjun Dai
Google DeepMind
<EMAIL_ADDRESS>
&Wen-Ding Li
Cornell University
<EMAIL_ADDRESS>
&&Kevin Ellis
Cornell University
<EMAIL_ADDRESS>
&Charles Sutton
Google DeepMind
<EMAIL_ADDRESS>
###### Abstract

Search is an important technique in program synthesis that allows for adaptive
strategies such as focusing on particular search directions based on execution
results. Several prior works have demonstrated that neural models are
effective at guiding program synthesis searches. However, a common drawback of
those approaches is the inability to handle iterative loops, higher-order
functions, or lambda functions, thus limiting prior neural searches from
synthesizing longer and more general programs. We address this gap by
designing a search algorithm called LambdaBeam that can construct arbitrary
lambda functions that compose operations within a given DSL. We create
semantic vector representations of the execution behavior of the lambda
functions and train a neural policy network to choose which lambdas to
construct during search, and pass them as arguments to higher-order functions
to perform looping computations. Our experiments show that LambdaBeam
outperforms neural, symbolic, and LLM-based techniques in an integer list
manipulation domain.

##  1 Introduction

Program synthesis involves finding a program meeting a given specification of
what the program should do [21, 18]. When the specification is in the form of
input/output examples, known as programming by example (PBE), combinatorial
search has been an especially popular technique [2, 30, 3, 32, 4, 27, 33].
Learning can also play a key role in PBE, because well-designed search
algorithms can learn to adapt to information collected during the ongoing
search, such as execution results or other analyses of candidate programs
considered so far. This information can be used to prune redundant parts of
the search space or focus on parts deemed more promising. For example,
DeepCoder [3] and TF-Coder [32] use neural models to define a search space
that is explored by traditional non-neural search, while Bustle [27],
CrossBeam [33], and Execution-Guided Synthesis [7] use neural models to guide
the search process itself. However, those prior works are unable to generate
programs with arbitrary looping computations, whether implemented via loop
control structures or through the use of higher-order functions with arbitrary
lambda functions.111DeepCoder [3] supports higher-order functions but only a
small set of hardcoded lambda functions. Execution-Guided Synthesis [7]
supports variable-free while loops, but not loops with an iteration variable.
Even though large language models (and other sequence models) can output
programs with loops and are very effective at synthesizing programs from
natural language [6], PBE demands a more systematic search strategy that
adapts to valuable information like execution results during the search.

The fundamental question explored in this paper is whether a neural program
synthesis search policy can learn to reason about lambdas and higher-order
functions, which would enable the synthesis of arbitrary looping computations
that were not previously possible with neural synthesis search techniques that
rely on intermediate expression evaluation. Previous work [27, 33] has shown
that neural models can effectively guide search when every candidate program
can be evaluated to produce a concrete value for the model to inspect,
allowing it to make decisions based on comparisons between explored values and
the desired output. Lambda functions however are extremely different: they
represent plans of functionality to be performed later, without specifying the
context in which this functionality will be used. As a result, reasoning about
lambdas requires a more abstract form of planning. Without knowing how the
lambda might be used later, the search policy must understand the different
behaviors of lambdas, predict whether a lambda will be useful for a given task
to prioritize search directions, and recognize when and how to use lambdas
within higher-order functions to actually perform useful computations.

In order to design a neural search algorithm that handles lambdas and higher-
order functions, we address some key difficulties. One challenge is in the
algebraic representation and subsequent manipulation of lambdas. We want to
represent “practically equivalent” lambdas like $\lambda x.\ x+1$, $\lambda
y.\ y+1$, and $\lambda x,y.\ x+1$ in a canonical way to prune the search
space. If $\lambda x.\ x+1$ is their canonical representation, then how can we
reuse that lambda to create a new lambda such as $\lambda x,y.\
(x+1)\times(y+1)$ where the “$\circ+1$” functionality is used in different
ways? We address this challenge by defining a new Merge operation that
combines lambda expressions into larger ones while allowing for variable
renaming and adhering to other representational constraints, therefore
enabling a bottom-up search algorithm to systematically build larger lambdas
from existing ones.

A second difficulty is in the encoding of lambdas when used as inputs to
neural models. A naive representation would be to encode the code tokens, but
slight changes in the code could lead to drastic differences in the lambda’s
behavior. Instead, we introduce a method of encoding lambdas that more
directly reflects their execution semantics. We do this using _property
signatures_ [26] in a way agnostic to how the lambda is used later (i.e., what
inputs are given to the lambda by a higher-order function), but still
analyzing the lambda’s behavior in the context of the current PBE task. One
conclusion of our work is that this encoding does enable neural models to
reason about lambdas effectively.

We present a new neural synthesis search method called LambdaBeam, which
combines our solutions to these challenges within the search framework of the
recent work CrossBeam [33]. CrossBeam performs a bottom-up search applying DSL
operations to previously-explored values, using a neural search policy to
choose the operation’s arguments with a pointer network. Thus, in LambdaBeam,
the neural policy is able to reason about lambdas by choosing which ones to
construct and when to use them, such that the search direction is tailored to
the synthesis task.

We demonstrate the effectiveness of LambdaBeam in the DeepCoder [3] domain of
integer list manipulation. We extend the DeepCoder DSL by adding many first-
order operations, keeping its higher-order functions, and replacing its
limited set of hardcoded lambda functions with arbitrary lambdas using
compositions of other DSL operations. Using a benchmark suite containing 100
natural hand-crafted evaluation tasks and 100 synthetically-generated tasks,
we experimentally show that LambdaBeam outperforms prior approaches including
state-of-the-art symbolic search, neural sequence models trained from scratch,
and a 62 billion parameter large language model (LLM). We release our
LambdaBeam code and trained model checkpoints at
https://github.com/ellisk42/LambdaBeam.

##  2 Background

Programming By Example  _Programming by Example_ (PBE) is the task of
synthesizing programs that satisfy a given set of input/output (I/O) examples.
In this task, we have a domain-specific language (DSL) $\mathcal{L}$
describing a space of programs, and a set of example inputs
$\mathcal{I}=\\{I_{1},\ldots,I_{N}\\}$ and corresponding outputs
$\mathcal{O}=\\{O_{1},\ldots,O_{N}\\}$. The goal is to find a program
$P\in\mathcal{L}$ such that $P(I_{i})=O_{i}$ for all $i\in\\{1,\ldots,N\\}$.
The DSL $\mathcal{L}$ describes atomic values (constants and input variables)
and operations that can be applied to arguments to produce new values.
Programs in $\mathcal{L}$ are arbitrarily-nested compositions of operations
applied to atomic values or other such compositions.

$\boldsymbol{\lambda}$-Calculus The lambda calculus [10, 28] is a formalism
for universal computation. A lambda calculus _term_ is either a variable,
function application, or a lambda abstraction. Lambda abstractions define new
functions by introducing new lexically-scoped variables, as well as a function
body (which is itself a term). We consider lambda abstractions that are
allowed to introduce multiple variables at once (we do not “Curry” our
lambdas). Terms in the lambda calculus are constructed recursively from
smaller terms, resulting in tree-like structures like in Figure 3(a). We
extend $\lambda$-calculus with domain-specific primitives
($\operatorname{add}$, $\operatorname{sort}$, $\operatorname{map}$,
$\operatorname{filter}$, etc.) as well as a basic type system known as simply-
typed lambda calculus; see [28] for details.

Property Signatures During a neural program synthesis search, many different
expressions are encountered, possibly including lambda functions with
arbitrary input and output types. How might a neural model learn to reason
about these expressions? One option is to represent expressions as source code
and apply standard text encoders like recurrent neural networks or
Transformers. However, two expressions with similar source code might have
very different execution behavior, or two syntactically distinct expressions
might have identical semantics. An alternative approach that is more
semantically aware is based on _property signatures_ [26, 27], inspired by
property-based testing in the programming languages literature [16]. Unary
properties describe a single value by mapping it to a boolean, such as whether
a list is sorted, or whether a string is empty. Binary properties can describe
the relationship between the input and output of a PBE task, such as whether
the input and output lists have the same length. In either case, given a list
of $k$ property functions, we can evaluate all property functions to form a
vector of length $k$ called a _property signature_. Each element of the
signature is either True, False, or N/A.222A property might be _not
applicable_ (N/A) if it does not apply to the types of objects currently under
consideration, or if the property is inconclusive because the code execution
resulted in an error. This vector may be used as input to a deep neural
network. Furthermore, given multiple executions of an expression (e.g., over
different examples), we can identify how often each property holds, leading to
a more granular representation.

Figure 1: Overview of LambdaBeam which builds upon the prior work CrossBeam
(new elements in LambdaBeam shown in red). Figure 2: Diagram of the LambdaBeam
model architecture, closely mirroring that of CrossBeam.

CrossBeam Our work LambdaBeam builds upon the prior work CrossBeam [33]. Both
systems have a similar overall structure, illustrated in Figure 1 with
differences shown in red. The core idea in CrossBeam is to use a neural policy
network to guide a bottom-up search over programs, where execution results of
explored expressions provide a powerful signal for the policy to guide the
search by exploring new expressions whose values are closer to the intended
output.

In CrossBeam, a table stores the expressions explored so far in search, along
with their execution values when run on the I/O examples. As diagrammed in
Figure 2, the Value Module encodes each explored value into a vector that is
stored in a matrix $E_{S}$. Meanwhile, the I/O Module encodes the I/O examples
into an analogous vector $\mathbf{e}_{IO}$. Then, the Search Context Summary
Module combines the value encodings $E_{S}$ and the I/O encoding
$\mathbf{e}_{IO}$ to produce a latent “summary” $\mathbf{e}_{c}$ of the search
context so far. From this, the Argument Selector Module, which is a recurrent
pointer network [35], predicts an argument list for a DSL operation via a
sequence of pointers into the table of explored values, thus generating a new
expression to explore. By structuring the search in this way, the neural
policy is able to take a “hands on” role during search, using information in
the search context to choose which programs should be explored next.

The network is trained on-policy using imitation learning. Each item in the
training set consists of a set of I/O examples and a target program. During
training, we run the search algorithm on the I/O examples where at each step
the policy network proposes values to explore by predicting argument lists for
DSL operations. We then apply a softmax loss that encourages the policy
network to make predictions that would lead to progress according to the
target program, instead of other argument lists that the policy network
actually proposed.

During evaluation, getting argument lists via beam search as in training can
lead to the search stalling if all new values are semantically equivalent to
values already seen, since beam search is deterministic if the value set is
unchanged. To address this, during evaluation CrossBeam randomly samples
different argument lists using UniqueRandomizer [31] to avoid duplicate
samples.

##  3 The LambdaBeam Approach

(a) Term as a tree | (b) Term as straightline code | (c) Built using Merge
---|---|---
$\lambda v_{1}$$\operatorname{map}$$\lambda u$$\operatorname{add}$$v_{1}$$\operatorname{square}$$u$$\lambda$$\operatorname{sort}$$x$ |  |
| $\texttt{t}_{1}\leftarrow x$ | atomic term from DSL
| $\texttt{t}_{2}\leftarrow\lambda v_{1}.\ \hbox{\pagecolor{lightred}\vphantom{$\operatorname{dq}${t}${}_{1}$}$\operatorname{square}$}(\hbox{\pagecolor{lightblue}\vphantom{$\operatorname{dq}${t}${}_{1}$}$v_{1}$})$ | $\textsc{Merge}{}(\hbox{\pagecolor{lightred}\vphantom{$\operatorname{dq}${t}${}_{1}$}$\operatorname{square}$},\hbox{\pagecolor{lightblue}\vphantom{$\operatorname{dq}${t}${}_{1}$}$v_{1}$},[\ ])$
| $\texttt{t}_{3}\leftarrow\lambda v_{1},v_{2}.\ \hbox{\pagecolor{lightred}\vphantom{$\operatorname{dq}${t}${}_{1}$}$\operatorname{add}$}(\hbox{\pagecolor{lightblue}\vphantom{$\operatorname{dq}${t}${}_{1}$}$v_{1}$},\hbox{\pagecolor{lightblue}\vphantom{$\operatorname{dq}${t}${}_{1}$}$\texttt{t}_{2}$}(\hbox{\pagecolor{lightblue}\vphantom{$\operatorname{dq}${t}${}_{1}$}$v_{2}$}))$ | $\textsc{Merge}{}(\hbox{\pagecolor{lightred}\vphantom{$\operatorname{dq}${t}${}_{1}$}$\operatorname{add}$},\hbox{\pagecolor{lightblue}\vphantom{$\operatorname{dq}${t}${}_{1}$}$v_{1}$},[\ ],\hbox{\pagecolor{lightblue}\vphantom{$\operatorname{dq}${t}${}_{1}$}$\texttt{t}_{2}$},[\hbox{\pagecolor{lightblue}\vphantom{$\operatorname{dq}${t}${}_{1}$}$v_{2}$}])$
| $\texttt{t}_{4}\leftarrow\hbox{\pagecolor{lightred}\vphantom{$\operatorname{dq}${t}${}_{1}$}$\operatorname{sort}$}(\hbox{\pagecolor{lightblue}\vphantom{$\operatorname{dq}${t}${}_{1}$}$\texttt{t}_{1}$})$ | $\textsc{Merge}{}(\hbox{\pagecolor{lightred}\vphantom{$\operatorname{dq}${t}${}_{1}$}$\operatorname{sort}$},\hbox{\pagecolor{lightblue}\vphantom{$\operatorname{dq}${t}${}_{1}$}$\texttt{t}_{1}$},[\ ])$
| $\texttt{t}_{5}\leftarrow\lambda v_{1}.\ \hbox{\pagecolor{lightred}\vphantom{$\operatorname{dq}${t}${}_{1}$}$\operatorname{map}$}(\lambda u.\ \hbox{\pagecolor{lightblue}\vphantom{$\operatorname{dq}${t}${}_{1}$}$\texttt{t}_{3}$}(\hbox{\pagecolor{lightblue}\vphantom{$\operatorname{dq}${t}${}_{1}$}$v_{1}$},\hbox{\pagecolor{lightblue}\vphantom{$\operatorname{dq}${t}${}_{1}$}$u$}),\hbox{\pagecolor{lightblue}\vphantom{$\operatorname{dq}${t}${}_{1}$}$\texttt{t}_{4}$})$ | $\textsc{Merge}{}(\hbox{\pagecolor{lightred}\vphantom{$\operatorname{dq}${t}${}_{1}$}$\operatorname{map}$},\hbox{\pagecolor{lightblue}\vphantom{$\operatorname{dq}${t}${}_{1}$}$\texttt{t}_{3}$},[\hbox{\pagecolor{lightblue}\vphantom{$\operatorname{dq}${t}${}_{1}$}$v_{1}$},\hbox{\pagecolor{lightblue}\vphantom{$\operatorname{dq}${t}${}_{1}$}$u$}],\hbox{\pagecolor{lightblue}\vphantom{$\operatorname{dq}${t}${}_{1}$}$\texttt{t}_{4}$},[\ ])$
Figure 3: Different ways of representing the lambda expression $\lambda
v_{1}.\ \operatorname{map}(\lambda u.\ v_{1}+u^{2},\operatorname{sort}(x))$,
where $x$ is an input variable. The model predicts $\operatorname{dq}$t1blue
term pointers and variables, the search tries every
$\operatorname{dq}$t1$\operatorname{red}$ operation, and unshaded code tokens
in (b) are known from the other shaded tokens.

###  3.1 Building $\boldsymbol{\lambda}$-Terms

LambdaBeam constructs terms in the lambda calculus during its search. A
natural design choice is to construct terms in a way that avoids considering
semantically equivalent expressions. For example, the terms $(\lambda x,y.\
x-y)$, $(\lambda a,b.\ a-b)$, and $(\lambda y,x.\ y-x)$ are all capable of
expressing exactly the same computations, so there should be a single
canonical way of building this family of terms.

An important aspect of our canonicalization of semantically equivalent
expressions is to enforce that every term constructed during search has no
free variables (but terms may include variables referring to the inputs given
in the I/O examples). Enforcing this property means that we can treat every
term we build during search as an ordinary program, and run it on different
inputs to probe its input-output behavior. However, the most straightforward
way of building lambda terms bottom-up violates this important property.
Consider the term $\texttt{t}_{5}=(\lambda v_{1}.\ \operatorname{map}(\lambda
u.\ v_{1}+u^{2},\operatorname{sort}(x)))$ whose tree structure is illustrated
in Figure 3(a), and where $x$ refers to an input variable. This has subterms
such as $\lambda u.\ v_{1}+u^{2}$, where the variable $v_{1}$ occurs free. As
$v_{1}$ is free, it is unclear what the semantics of this expression should
be. Constructing the expression with $v_{1}$ free would also introduce a
spurious redundancy with $\lambda v_{1},u.\ v_{1}+u^{2}$. During search, we
would like to keep only a canonical version of those two terms to better prune
the search space, which in this case would be $\texttt{t}_{3}=\lambda
v_{1},v_{2}.\ v_{1}+v_{2}^{2}$, which as shown in Figure 3(b) can be used to
construct the larger term $\texttt{t}_{5}$.

To build terms bottom-up while maintaining desired constraints such as that
every intermediate program has no free variables, we define an operator for
algebraically combining smaller terms to build larger programs. This operator,
which we call Merge, takes as input a primitive DSL function $f$ and a list of
terms $a_{1}\ldots a_{K}$ constructed earlier during search, and it returns a
new term that applies $f$ to $a_{1}\ldots a_{K}$. Merge does extra bookkeeping
to ensure that every function (both $f$ and any $a_{k}$ that are lambda terms)
is called with the correct number of arguments, and that the final output of
Merge has no free variables. The function $f$ can be a higher-order function,
such as $\operatorname{map}$, or a normal first-order function, such as
addition and subtraction.

Additional arguments to Merge specify how to unify lambda arguments (if any)
that appear in $a_{1}\ldots a_{K}$. To evaluate Merge on $f$ and $a_{1}\dots
a_{K}$, we first apply each lambda argument $a_{k}$ to a tuple of variable
names, denoted $i_{k}$. This gives a name to each argument of each $a_{k}$. By
reusing the same name across different $a_{k}$’s, the same variable can be
shared across arguments. For instance, if the function $f$ is multiplication
($\times$), and we are merging with the arguments $a_{1}=(\lambda v.\ v+1)$
and $a_{2}=(\lambda v.\ v-1)$, then we can share the variable by setting
$i_{1}=i_{2}=[v_{1}]$, giving $(\lambda v.\ v+1)(v_{1})\times(\lambda v.\
v-1)(v_{1})=(v_{1}+1)\times(v_{1}-1)$. Alternatively, if we set
$i_{1}=[v_{1}]$ and $i_{2}=[v_{2}]$, then merging gives a different program,
$(v_{1}+1)\times(v_{2}-1)$. These tuples of variable names,
$\left\\{i_{k}\right\\}_{k=1}^{K}$, are also input to Merge, because they
determine how variable names are unified across arguments. Finally Merge runs
$f$ on the arguments $\left\\{a_{k}(i_{k})\right\\}_{k=1}^{K}$, pads the
expression with lambdas at the beginning to bind any new free variables, and
lastly canonicalizes variable naming and variable ordering.

Special care is needed for the arguments of higher-order functions like
$\operatorname{map}$, which themselves need to be functions. So far, each
argument $a_{k}(i_{k})$ evaluates to a concrete value, not a function. To
handle higher-order functions, Merge automatically adds extra lambdas to each
function-valued argument. These extra lambdas have variables denoted
$u_{1},u_{2},\ldots$. All other variables used to unify variables across
arguments are denoted $v_{1},v_{2},\ldots$. Although this may seem ad-hoc at
first, this convention actually corresponds to a well-known canonicalization
of typed lambda forms known as $\eta$-long normal form [28]. Putting
everything together, Merge is defined as

$\displaystyle\textsc{Merge}{}\left(f,a_{1},i_{1},a_{2},i_{2},\dots\right)$
$\displaystyle=\lambda v_{1}v_{2}...\,.\ f\left(\lambda
u_{1}u_{2}...u_{\ell_{1}}.\ a_{1}(i_{1}),\lambda u_{1}u_{2}...u_{\ell_{2}}.\
a_{2}(i_{2}),\ldots\right)$ where
$\displaystyle\left\\{v_{1},v_{2},\ldots\right\\}=\bigcup_{k}i_{k}-\left\\{u_{j}\right\\}_{j=1}^{\max\\{\ell_{1},\ell_{2},\dots\\}}$
(1)

where $\ell_{k}$ is the arity of the $k^{\text{th}}$ argument to $f$. For
example, the higher-order function $\operatorname{map}$ first takes a function
expecting one argument, so $\ell_{1}=1$, followed by a list (not a function)
expecting no arguments, so $\ell_{2}=0$. We also enforce that
$|i_{k}|=\operatorname{arity}(a_{k})$. The Merge operation is complete, in the
sense that we can use it to generate any PBE solution within the DSL (see
Appendix A).

Fundamentally, our neural model predicts the inputs to Merge. It scans through
its different operators (functions $f$), and for each operator, generates
arguments $a_{k}$ by pointing into the set of explored values, and variables
$i_{k}$ by emitting special tokens corresponding to $v_{1}$, $u_{1}$, $v_{2}$,
$u_{2}$, etc. Figure 3(b) and (c) illustrate how a nontrivial lambda
expression is built step by step, with the tokens emitted by the neural model
highlighted in blue. In this figure, each call to Merge in the third column
returns the term that appears in the middle column, e.g.,
$\textsc{Merge}{}(\hbox{\pagecolor{lightred}\vphantom{$\operatorname{dq}${t}${}_{1}$}$\operatorname{square}$},\hbox{\pagecolor{lightblue}\vphantom{$\operatorname{dq}${t}${}_{1}$}$v_{1}$},[\
])$ returns $\lambda v_{1}.\
\hbox{\pagecolor{lightred}\vphantom{$\operatorname{dq}${t}${}_{1}$}$\operatorname{square}$}(\hbox{\pagecolor{lightblue}\vphantom{$\operatorname{dq}${t}${}_{1}$}$v_{1}$})$
and so on. Critically, Merge makes sure that (1) every intermediate term that
the model builds along the way can be evaluated so that the neural model can
inspect its progress, and also (2) intermediate terms are generated with
_every_ function application, giving fine-grained feedback to the model.

We define the _weight_ of an expression to be the number of nodes in its tree
representation using the Merge operator. More specifically, atomic terms like
variable names and literal constants have weight $1$, and a term constructed
with Merge has weight $1$ (for the operation) plus the sum of the weights of
all terms and variables in the Merge call. For example, in Figure 3(b), terms
$\texttt{t}_{1}$ through $\texttt{t}_{5}$ have weights $1$, $2$, $5$, $2$, and
$10$ respectively.

###  3.2 Learning over $\boldsymbol{\lambda}$-Expressions

One core technical challenge is how to encode lambda expressions to allow
neural models to reason about them. In LambdaBeam we solve this by
constructing a new generalization of property signatures which is designed to
represent lambda expressions and non-lambda expressions using similar sets of
properties.

Non-lambda expressions can be evaluated to produce a single result per I/O
example. However, we cannot evaluate lambda expressions in the same way,
because we do not know how the lambda expression will be used in the eventual
solution, so we do not know what its arguments will be. For instance, if $x$
is an input list, the expressions $\operatorname{zipwith}(\lambda
u_{1},u_{2}.\ u_{1}+u_{2},x,\operatorname{sort}(x))$ and
$\operatorname{scanl1}(\lambda u_{1},u_{2}.\ u_{1}+u_{2},x)$ use the same
lambda expression $\lambda u_{1},u_{2}.\ u_{1}+u_{2}$ but for different sets
of $(u_{1},u_{2})$ arguments. Thus, in order to describe the lambda
expression’s execution behavior, we run the lambda on a fixed set of
_canonical argument tuples_ based on the number of arguments and their types.
These argument tuples are held constant across training and evaluation so that
the model can learn from consistent signals.

In our experiments, we hardcoded the canonical argument tuples without
changing them afterward, trying to cover common values and a variety of
scenarios. For instance, our integer arguments include those between $-3$ and
$5$ inclusive, and other integers with varying combinations of magnitude,
sign, even/odd parity, primality, and so on. There is also a tradeoff in the
number of canonical argument tuples, where having more leads to finer-grained
execution information but more time spent running lambdas during search. In
our experiments, we use 16 canonical argument tuples for each combination of
tuple length and argument types in our DSL. Note that the lambda expression
can refer to inputs from the I/O examples. Instead of running the lambda on
each argument tuple for each example, for efficiency, we run on each argument
tuple once, under the context of one I/O example which is changed per argument
tuple in a round-robin fashion.

To represent a lambda $f$, we evaluate it on each canonical argument tuple
$t_{i}$ with I/O example $(I,O)$. First we evaluate $f$ on $t_{i}$, yielding a
result $r_{i}=f(t_{i},I)$. Then we use a signature of unary properties to
describe $r_{i}$. Second, we use a signature of binary properties to compare
$r_{i}$ and $O$; intuitively, this helps the model learn about what work
remains. Similarly, we use the binary properties to compare $r_{i}$ and
$t_{ij}$, for each argument $t_{ij}\in t_{i}$. By concatenating these, we
obtain a single property vector for each tuple $t_{i}$. Finally, we then
reduce the property vectors across the runs of the lambda, i.e., computing for
each property the fraction of runs where it is applicable and the fraction of
applicable runs where it is True. Encoding non-lambda expressions is similar,
except that we use I/O examples $(I_{i},O_{i})$ in place of the canonical
tuples. Note that the reduced property signatures for lambda and non-lambda
expressions have different formats and different lengths, and hence they are
embedded by separate parts of the neural model (Section 3.3).

The properties used in our property signatures come from combinatorially
combining hand-designed features as “building blocks” to create a rich space
of properties that describe individual objects as well as comparisons between
two objects. Appendix B contains more details.

###  3.3 LambdaBeam Model and Search

To guide the bottom-up search over lambda and non-lambda expressions, we
generally follow the design of the neural policy network in CrossBeam [33],
with the following major changes:

Value Module We maintain a set of explored values $\mathcal{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 $\mathcal{S}$ forming a matrix of embeddings
$E_{\mathcal{S}}\in\mathbb{R}^{|\mathcal{S}|\times d}$. Elements of
$\mathcal{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.333The 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}$. 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
$\operatorname{LSTM}_{op}$ to select the arguments using a pointer mechanism
[35] from the encoded value matrix $E_{\mathcal{S}}$, in an autoregressive
way. In addition to selecting $\operatorname{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
$\operatorname{arity}(a_{k})$ variable tokens. All of the $a_{k}$ and $i_{k}$
predictions are done as a single autoregressive sequence. Furthermore, 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 $\operatorname{arity}(a_{k})$ variable tokens.

Search with Restarts We also change the inference time search procedure.
Recall that CrossBeam uses random sampling during evaluation, making the
search nondeterministic. In LambdaBeam, 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
LambdaBeam solve more tasks because it may be otherwise difficult to recover
from exploring the wrong search direction.

##  4 Experiments

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. 1.

LambdaBeam with random restarts: We trained the LambdaBeam 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. 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. 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. 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. 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. 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.444In 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. 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.

###  4.3 Results

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.

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.555The 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.

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.

Figure 6: True positive versus false positive solutions measured using held-
out test cases.

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.666The 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. 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 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.

##  5 Related Work

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].

##  6 Conclusion

We introduced the first neural search method for programming by example that
is able to synthesize intermediate helper functions (lambdas) 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.

#### Acknowledgments

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.

## References

  * Allamanis et al. [2018] Miltiadis Allamanis, Earl T Barr, Premkumar Devanbu, and Charles Sutton.  A survey of machine learning for big code and naturalness.  _ACM Computing Surveys (CSUR)_ , 51(4), 2018.
  * Alur et al. [2017] Rajeev Alur, Arjun Radhakrishna, and Abhishek Udupa.  Scaling enumerative program synthesis via divide and conquer.  In _International Conference on Tools and Algorithms for the Construction and Analysis of Systems (TACAS)_ , pages 319–336. Springer, 2017\.
  * Balog et al. [2017] Matej Balog, Alexander L Gaunt, Marc Brockschmidt, Sebastian Nowozin, and Daniel Tarlow.  DeepCoder: Learning to write programs.  In _International Conference on Learning Representations (ICLR)_ , 2017.
  * Barke et al. [2020] Shraddha Barke, Hila Peleg, and Nadia Polikarpova.  Just-in-time learning for bottom-up enumerative synthesis.  In _Object-Oriented Programming, Systems, Languages, and Applications (OOPSLA)_ , 2020.
  * Chang et al. [2015] Kai-Wei Chang, Akshay Krishnamurthy, Alekh Agarwal, Hal Daumé III, and John Langford.  Learning to search better than your teacher.  In _International Conference on Machine Learning (ICML)_ , 2015\.
  * Chen et al. [2021a] Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde, Jared Kaplan, Harri Edwards, Yura Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Such, Dave Cummings, Matthias Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert-Voss, Will Guss, Alex Nichol, Igor Babuschkin, Suchir Balaji, Shantanu Jain, Andrew Carr, Jan Leike, Josh Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob McGrew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojciech Zaremba.  Evaluating large language models trained on code.  _arXiv preprint arXiv:2107.03374_ , 2021a.  URL https://arxiv.org/abs/2107.03374.
  * Chen et al. [2019] Xinyun Chen, Chang Liu, and Dawn Song.  Execution-guided neural program synthesis.  In _International Conference on Learning Representations (ICLR)_ , 2019.
  * Chen et al. [2021b] Xinyun Chen, Dawn Song, and Yuandong Tian.  Latent execution for neural program synthesis.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2021b.
  * Chowdhery et al. [2023] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel.  PaLM: Scaling language modeling with pathways.  _Journal of Machine Learning Research (JMLR)_ , 24:240:1–240:113, 2023.  URL http://jmlr.org/papers/v24/22-1144.html.
  * Church [1985] Alonzo Church.  _The calculi of lambda-conversion_.  Number 6. Princeton University Press, 1985.
  * Daumé III et al. [2009] Hal Daumé III, John Langford, and Daniel Marcu.  Search-based structured prediction.  _Machine Learning_ , 75(3), 2009.
  * Devlin et al. [2017] Jacob Devlin, Jonathan Uesato, Surya Bhupatiraju, Rishabh Singh, Abdel-Rahman Mohamed, and Pushmeet Kohli.  RobustFill: Neural program learning under noisy I/O.  In _International Conference on Machine Learning (ICML)_ , 2017\.
  * Ellis et al. [2019] Kevin Ellis, Maxwell Nye, Yewen Pu, Felix Sosa, Josh Tenenbaum, and Armando Solar-Lezama.  Write, execute, assess: Program synthesis with a REPL.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2019.
  * Ellis et al. [2021] Kevin Ellis, Catherine Wong, Maxwell Nye, Mathias Sablé-Meyer, Lucas Morales, Luke Hewitt, Luc Cary, Armando Solar-Lezama, and Joshua B Tenenbaum.  DreamCoder: Bootstrapping inductive program synthesis with wake-sleep library learning.  In _Programming Language Design and Implementation (PLDI)_ , pages 835–850, 2021.
  * Feser et al. [2015] John K Feser, Swarat Chaudhuri, and Isil Dillig.  Synthesizing data structure transformations from input-output examples.  In _Programming Language Design and Implementation (PLDI)_ , pages 229–239, 2015.
  * Fink and Bishop [1997] George Fink and Matt Bishop.  Property-based testing: a new approach to testing for assurance.  _ACM SIGSOFT Software Engineering Notes_ , 22(4):74–80, 1997.
  * Gottschlich et al. [2018] Justin Gottschlich, Armando Solar-Lezama, Nesime Tatbul, Michael Carbin, Martin Rinard, Regina Barzilay, Saman Amarasinghe, Joshua B Tenenbaum, and Tim Mattson.  The three pillars of machine programming.  In _International Workshop on Machine Learning and Programming Languages (MAPL at PLDI)_ , 2018.
  * Gulwani et al. [2017] Sumit Gulwani, Oleksandr Polozov, and Rishabh Singh.  Program synthesis.  _Foundations and Trends® in Programming Languages_ , 4(1-2), 2017.
  * Hong et al. [2021] Joey Hong, David Dohan, Rishabh Singh, Charles Sutton, and Manzil Zaheer.  Latent programmer: Discrete latent codes for program synthesis.  In _International Conference on Machine Learning (ICML)_ , 2021\.
  * Lee et al. [2018] Woosuk Lee, Kihong Heo, Rajeev Alur, and Mayur Naik.  Accelerating search-based program synthesis using learned probabilistic models.  In _Programming Language Design and Implementation (PLDI)_ , 2018\.
  * Manna and Waldinger [1971] Zohar Manna and Richard J Waldinger.  Toward automatic program synthesis.  _Communications of the ACM_ , 14(3):151–165, 1971\.  URL https://doi.org/10.1145/362566.362568.
  * Murali et al. [2018] Vijayaraghavan Murali, Letao Qi, Swarat Chaudhuri, and Chris Jermaine.  Neural sketch learning for conditional program generation.  In _International Conference on Learning Representations (ICLR)_ , 2018.
  * Negrinho et al. [2018] Renato Negrinho, Matthew Gormley, and Geoffrey Gordon.  Learning beam search policies via imitation learning.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2018.
  * Negrinho et al. [2020] Renato Negrinho, Matthew Gormley, and Geoffrey Gordon.  An empirical investigation of beam-aware training in supertagging.  In _Findings of the Association for Computational Linguistics: EMNLP_ , 2020.
  * Nye et al. [2019] Maxwell Nye, Luke Hewitt, Joshua Tenenbaum, and Armando Solar-Lezama.  Learning to infer program sketches.  In _International Conference on Machine Learning (ICML)_ , 2019\.
  * Odena and Sutton [2020] Augustus Odena and Charles Sutton.  Learning to represent programs with property signatures.  In _International Conference on Learning Representations (ICLR)_ , 2020.
  * Odena et al. [2021] Augustus Odena, Kensen Shi, David Bieber, Rishabh Singh, Charles Sutton, and Hanjun Dai.  BUSTLE: Bottom-up program synthesis through learning-guided exploration.  In _International Conference on Learning Representations (ICLR)_ , 2021.
  * Pierce [2002] Benjamin C Pierce.  _Types and programming languages_.  MIT press, 2002.
  * Ross et al. [2011] 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.
  * Shi et al. [2019] Kensen Shi, Jacob Steinhardt, and Percy Liang.  FrAngel: Component-based synthesis with control structures.  _Proceedings of the ACM on Programming Languages_ , 3(POPL), 2019.
  * Shi et al. [2020] Kensen Shi, David Bieber, and Charles Sutton.  Incremental sampling without replacement for sequence models.  In _International Conference on Machine Learning (ICML)_ , 2020.
  * Shi et al. [2022a] 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, 2022a.
  * Shi et al. [2022b] 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)_ , 2022b.
  * Shrivastava et al. [2021] 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.
  * Vinyals et al. [2015] Oriol Vinyals, Meire Fortunato, and Navdeep Jaitly.  Pointer networks.  In _Advances in Neural Information Processing Systems (NIPS)_ , 2015\.
  * Yin and Neubig [2017] Pengcheng Yin and Graham Neubig.  A syntactic neural model for general-purpose code generation.  In _Assocation for Computational Linguistics (ACL)_ , 2017.
  * Zohar and Wolf [2018] 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.

##  Appendix 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},\dots,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:

$\displaystyle T::=x\mid\ $ $\displaystyle v\mid c\mid
f(\mathbf{t}_{1},\dots,\mathbf{t}_{k})\mid\lambda v_{1}...v_{n}.\ \mathbf{t}$
for $\displaystyle x\in\mathcal{X},c\in\mathcal{C},f\in\mathcal{F},$
$\displaystyle v,v_{1},\dots,v_{n}\in\mathcal{V},$
$\displaystyle\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):

$\displaystyle M::=x\mid\ $ $\displaystyle\lambda v.\ v\mid
c\mid\textsc{Merge}(f,\mathbf{a}_{1},i_{1},\dots,\mathbf{a}_{k},i_{k})$ for
$\displaystyle
x\in\mathcal{X},v\in\mathcal{V},c\in\mathcal{C},f\in\mathcal{F},$
$\displaystyle\mathbf{a}_{1},\dots,\mathbf{a}_{k}\in M,$ $\displaystyle
i_{1},\dots,i_{k}\in\mathcal{V}^{*}.$

We restate the definition

$\displaystyle\textsc{Merge}{}\left(f,\mathbf{a}_{1},i_{1},\mathbf{a}_{2},i_{2},\dots\right)$
$\displaystyle=\lambda v_{1}v_{2}...\,.\ f\left(\lambda
u_{1}u_{2}...u_{\ell_{1}}.\ \mathbf{a}_{1}(i_{1}),\lambda
u_{1}u_{2}...u_{\ell_{2}}.\ \mathbf{a}_{2}(i_{2}),\ldots\right)$ where
$\displaystyle\left\\{v_{1},v_{2},\ldots\right\\}=\bigcup_{k}i_{k}-\left\\{u_{j}\right\\}_{j=1}^{\max\\{\ell_{1},\ell_{2},\dots\\}}$

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.\ 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}...v_{n}.\ \mathbf{t}$ (possibly with $n=0$
such that $\mathbf{s}$ is not a lambda expression) _has exact lambda
variables_ if
$\operatorname{FreeVars}(\mathbf{t})-\mathcal{X}=\\{v_{1},\dots,v_{n}\\}$.
Note that $\mathbf{s}$ having exact lambda variables implies that
$\operatorname{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
$\mathbf{s}$ to PBE problems (once the input variables $x_{1},\dots,x_{m}$ are
bound), follows from the two claims below.

Claim 1. If $\mathbf{p}=\lambda x_{1}...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},\dots,x_{m})\rightarrow 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}...x_{m}.\ \mathbf{s}$ is a valid
solution program, we must have
$\operatorname{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 $\mathbf{s}$ be any term in $S$. We will proceed by induction on the depth
of $\mathbf{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.

Then, we assume the inductive hypothesis that any term in $S$, with depth less
than that of $\mathbf{s}$, is in $M$. There are two inductive cases to
consider: either $\mathbf{s}=\lambda v_{1}...v_{n}.\
f(\mathbf{t}_{1},\dots,\mathbf{t}_{k})$ where $n$ might be $0$, or
$\mathbf{s}=\lambda v_{1}...v_{n}.\ \mathbf{t}$ where $n>0$ and $\mathbf{t}$
is atomic. Because $\mathbf{s}$ has exact lambda variables, the latter
scenario is only possible if $\mathbf{s}=\lambda v.\ v$, which is immediately
in $M$.

The remaining case is $\mathbf{s}=\lambda v_{1}...v_{n}.\
f(\mathbf{t}_{1},\dots,\mathbf{t}_{k})$. Consider any $\mathbf{t}_{j}$ for
$1\leq j\leq k$. We will construct $\mathbf{a}_{j}$ and $i_{j}$ such that
$\textsc{Merge}(f,\mathbf{a}_{1},i_{1},\dots,\mathbf{a}_{k},i_{k})=\mathbf{s}$.

  * •

If $\mathbf{t}_{j}$ is atomic, then either $\mathbf{t}_{j}=x$,
$\mathbf{t}_{j}=v$, or $\mathbf{t}_{j}=c$. If $\mathbf{t}_{j}=v$, then set
$\mathbf{a}_{j}=\lambda v.\ v$ and $i_{j}=[v]$; otherwise, set
$\mathbf{a}_{j}=\mathbf{t}_{j}$ and $i_{j}=[\ ]$, the empty tuple. In each
case, $\mathbf{a}_{j}\in M$. Because $\mathbf{s}$ typechecks, we know that the
$j$-th argument to $f$ expects arity 0, so in the Merge definition,
$\ell_{j}=0$ and thus the $j$-th argument to $f$ expands to
$\mathbf{a}_{j}(i_{j})=\mathbf{t}_{j}$ in each case.777In practice for
simplicity, when $\mathbf{t}_{j}=v$, we simply set
$\mathbf{a}_{j}=\mathbf{t}_{j}$ and $i_{j}=[\ ]$ as in the other cases, even
though $\mathbf{a}_{j}=v$ does not have exact lambda variables.

  * •

If $\mathbf{t}_{j}$ is not atomic, then let $\mathbf{t}_{j}=\lambda
u_{1}...u_{\ell_{j}}.\ \mathbf{r}$, where $\ell_{j}$ is the expected arity of
the $j$-th argument to $f$ (because $\mathbf{s}$ typechecks). Let
$\\{v^{\prime}_{1},\dots,v^{\prime}_{d}\\}=\operatorname{FreeVars}(\mathbf{r})-\mathcal{X}$.
Set $\mathbf{a}_{j}=\lambda v^{\prime}_{1}...v^{\prime}_{d}.\ \mathbf{r}$ and
$i_{j}=[v^{\prime}_{1},\dots,v^{\prime}_{d}]$, so that when applying Merge,
the $j$-th argument to $f$ expands to $\lambda u_{1}...u_{\ell_{j}}.\
\mathbf{a}_{j}(i_{j})=\lambda u_{1}...u_{\ell_{j}}.\
\mathbf{r}=\mathbf{t}_{j}$. Furthermore, $\mathbf{a}_{j}$ has exact lambda
variables by construction, and it typechecks because $\mathbf{r}$ typechecks,
so $\mathbf{a}_{j}\in S$ and $\mathbf{a}_{j}\in M$ by the inductive
hypothesis.

With these choices of $\mathbf{a}_{j}$ and $i_{j}$, when expanding
$\textsc{Merge}(f,\mathbf{a}_{1},i_{1},\dots,\mathbf{a}_{k},i_{k})$ according
to the definition, the $j$-th argument to $f$ becomes $\mathbf{t}_{j}$, and
the lambda variables
$\bigcup_{k}i_{k}-\left\\{u_{j}\right\\}_{j=1}^{\max\\{\ell_{1},\ell_{2},\dots\\}}$
are exactly
$\operatorname{FreeVars}(f(\mathbf{t}_{1},\dots,\mathbf{t}_{k}))=\\{v_{1},\dots,v_{n}\\}$
since $\mathbf{s}$ has exact lambda variables. Therefore, $\mathbf{s}=\lambda
v_{1}...v_{n}.\
f(\mathbf{t}_{1},\dots,\mathbf{t}_{k})=\textsc{Merge}(f,\mathbf{a}_{1},i_{1},\dots,\mathbf{a}_{k},i_{k})$,
so $\mathbf{s}\in M$. ∎

##  Appendix B More Details on Property Signatures

Here we describe in more detail the properties we use to encode lambda and
non-lambda values. We define the following helper functions to organize the
properties.

First, $\operatorname{TypeProperties}(x)$ represents the type of $x$ as a
boolean one-hot list. In our DSL, this returns a tuple of 5 booleans,
representing whether $x$ is a lambda, boolean, int, list, or None (which is
used to indicate an error, e.g., signaling that an index is out of bounds).

Next, we define $\operatorname{BasicProperties}(x)$ to evaluate hand-designed
“basic” properties that describe objects of each different type in the DSL.
This returns a fixed-length vector of property results (each being True,
False, or N/A). Note that, if $x$ has type $\tau$, then all properties for
type $\tau$ evaluate to True or False, while all properties for all other
types $\tau^{\prime}\neq\tau$ evaluate to N/A. For our DSL, we use the
following basic properties:

  * •

For boolean $x$: $x$ itself.

  * •

For integer $x$: whether $x$ equals $-1$, $0$, $1$, and $2$; whether $x$ is
positive and negative; whether $x$ is even; whether $x$ is $0$ and $1$ modulo
$3$; and whether $|x|$ is less than $5$, $10$, $20$, $35$, $50$, $75$, and
$100$.

  * •

For list $x$: whether $x$ is sorted, whether $x$ is sorted in reverse, and
whether $x$ contains all unique elements.

Then, $\operatorname{Relevant}(x)$ returns related objects that are relevant
to understanding $x$. For our DSL, this is only $x$ itself for integer and
boolean $x$, but for list $x$, the “relevant” objects are: $x$ itself; the
length of $x$; the number of distinct elements in $x$; the max, min, range,
and sum of $x$; and the first and last elements of $x$ (defaulting to 0 if $x$
is empty).

This culminates in $\operatorname{ObjectSignature}(x)$ which takes a single
DSL object $x$ and returns a fixed-length vector of property results,
containing $\operatorname{TypeProperties}(x)$ followed by
$\operatorname{BasicProperties}(r)$ for each $r\in\operatorname{Relevant}(x)$.
For example, these properties include “assuming $x$ is an int, is $x$ is
even?” (a basic property applied to $x$) as well as “assuming $x$ is a list,
are there an even number of elements in $x$?” (a basic property applied to an
object relevant to $x$). By applying basic properties to relevant objects in
this compositional way, we reduce the effort needed to specify a large number
of properties.

We furthermore encode comparisons between two objects. We define
$\operatorname{ComparisonProperties}(x,y)$ which evaluates hand-designed
properties for comparing two objects $x$ and $y$ of the same type, for each
different type in the DSL. This returns a fixed-length vector of property
results, where a property for comparing type $\tau$ evaluates to N/A if $x$
and $y$ are not of type $\tau$. For our DSL, we use the following comparison
properties:

  * •

For boolean $x$ and $y$: whether $x=y$.

  * •

For integer $x$ and $y$: whether $x=y$, $x<y$, and $x>y$; whether $x$ is a
factor of $y$ and vice versa; and whether $|x-y|$ is less than $2$, $5$, $10$,
and $20$.

  * •

For list $x$ and $y$: whether $x=y$; whether $x$ is longer, shorter, or equal
length compared to $y$; whether the lengths differ by at most $1$; whether all
$x_{i}<y_{i}$ for $x_{i},y_{i}\in\operatorname{zip}(x,y)$ and similarly for
other comparisons $\leq$, $>$, $\geq$, $=$, and $\neq$; whether $x$ and $y$
contain the same set of elements; and whether $x$ contains a subset of
elements compared to $y$ and vice versa.

These properties are used in $\operatorname{ComparisonSignature}(x,y)$ which
computes a fixed-length list of property results for any two DSL objects $x$
and $y$ of any type, containing $\operatorname{ComparisonProperties}(r_{x},y)$
for all $r_{x}\in\operatorname{Relevant}(x)$ where
$\operatorname{type}(r_{x})=\operatorname{type}(y)$, and
$\operatorname{ComparisonProperties}(x,r_{y})$ for all
$r_{y}\in\operatorname{Relevant}(y)$ where
$\operatorname{type}(r_{y})=\operatorname{type}(x)$. Thus, “assuming $x$ is an
int and $y$ is a list, is $x$ a factor of the length of $y$?” is one resulting
property. As usual, if $x$ and $y$ do not match the types assumed by the
property, then the property evaluates to N/A.

In the I/O Module of the neural policy (as in CrossBeam [33]), we use property
signatures to encode a set of I/O examples. For each example
$(\\{I_{1},\dots,I_{n}\\},O)$ we concatenate
$\operatorname{ObjectSignature}(O)$ with
$\operatorname{ObjectSignature}(I_{i})$ and
$\operatorname{ComparisonSignature}(I_{i},O)$ for all $1\leq i\leq n$. Then,
we reduce these across I/O examples, computing for each property the fraction
of examples where it is applicable (not N/A), and the fraction of examples
where it is True among those where it is applicable (defaulting to $0.5$ if it
is N/A for all examples).

In the Value Module of the neural policy, we use property signatures to encode
a value (lambda or non-lambda expression) that was found during search. To
encode a lambda expression, we run it on canonical input tuples as described
in Section 3.2. For each run of the lambda on canonical input tuple
$t_{i}=(t_{i,1},\dots,t_{i,m})$ using an I/O example $(I,O)$ where the lambda
evaluates to a result $r_{i}$, we concatenate
$\operatorname{ObjectSignature}(r_{i})$,
$\operatorname{ComparisonSignature}(r_{i},O)$, and
$\operatorname{ComparisonSignature}(t_{i,j},r_{i})$ for all $1\leq j\leq m$,
and then reduce these across the runs of the lambda. To encode a non-lambda
expression during search, for each I/O example $(I,O)$ where the expression
evaluates to a result $r$, we concatenate $\operatorname{ObjectSignature}(r)$
with $\operatorname{ComparisonSignature}(r,O)$, and then reduce these across
I/O examples. Note that the signatures for values found during search do not
contain comparisons to the I/O example inputs, because what ultimately matters
is whether the value is useful for creating the _output_ later, not how the
value was created from the inputs.

In our implementation, encoding the set of I/O examples results in a property
signature of length 1230, encoding a lambda expression results in a property
signature of length 558, and encoding a non-lambda expression results in a
property signature of length 359.

##  Appendix C Extension of the DeepCoder DSL

As mentioned in Section 4.1, we extended the DSL from DeepCoder [3]. Atomic
terms in the DSL include variable names and the constant literals $-1$, $0$,
$1$, $2$, $3$, and $4$. The DSL contains 23 first-order and 5 higher-order
operations, listed below with type annotations and Python implementations:

⬇

# 23 first-order operations

def Add(x: int, y: int) -¿ int:

return x + y

def Subtract(x: int, y: int) -¿ int:

return x - y

def Multiply(x: int, y: int) -¿ int:

return x * y

def IntDivide(x: int, y: int) -¿ int:

return x // y

def Square(x: int) -¿ int:

return x ** 2

def Min(x: int, y: int) -¿ int:

return min(x, y)

def Max(x: int, y: int) -¿ int:

return max(x, y)

def Greater(x: int, y: int) -¿ bool:

return x > y

def Less(x: int, y: int) -¿ bool:

return x < y

def Equal(x: int, y: int) -¿ bool:

return x == y

def IsEven(x: int) -¿ bool:

return x % 2 == 0

def IsOdd(x: int) -¿ bool:

return x % 2 != 0

def If(c: bool, x: int, y: int) -¿ int:

return x if c else y

def Head(xs: list) -¿ int:

return xs[0]

def Last(xs: list) -¿ int:

return xs[-1]

def Take(n: int, xs: list) -¿ list:

return xs[:n]

def Drop(n: int, xs: list) -¿ list:

return xs[n:]

def Access(n: int, xs: list) -¿ int:

return xs[n]

def Minimum(xs: list) -¿ int:

return min(xs)

def Maximum(xs: list) -¿ int:

return max(xs)

def Reverse(xs: list) -¿ list:

return list(reversed(xs))

def Sort(xs: list) -¿ list:

return sorted(xs)

def Sum(xs: list) -¿ int:

return sum(xs)

# 5 higher-order operations

def Map(f: Callable[[int], int], xs: list) -¿ list:

return [f(x) for x in xs]

def Filter(f: Callable[[int], bool], xs: list) -¿ list:

return [x for x in xs if f(x)]

def Count(f: Callable[[int], bool], xs: list) -¿ int:

return len([x for x in xs if f(x)])

def ZipWith(f: Callable[[int, int], int], xs: list, ys: list) -¿ list:

return [f(x, y) for x, y in zip(xs, ys)]

def Scanl1(f: Callable[[int, int], int], xs: list) -¿ list:

ys = [xs[0]]

for n in range(1, len(xs)):

ys.append(f(ys[n-1], xs[n]))

return ys

##  Appendix D Example Tasks

This section contains selected example problems from our 100 handwritten and
100 synthetic evaluation tasks. Each task is given a name for convenience
purposes only, which is not used by any method in our experiments.

###  D.1 Handwritten Task “map:replace”

This task has 3 inputs (x, f, and r), 3 examples demonstrating the task (“in
x, find instances of f and replace them with r”), and a handwritten ground-
truth solution using a relatively complicated lambda function:

⬇

inputs_dict = {

'x': [[7, 2, 4, 6, 4, 2, 5],

[-6, -3, 4, 3, -5, -3, 2, 1, 5],

[18, 48, 27, 26, 27, 27, 28, 17, 27, 33]],

'f': [4, -3, 27],

'r': [-1, 7, 99],

}

outputs = [[7, 2, -1, 6, -1, 2, 5],

[-6, 7, 4, 3, -5, 7, 2, 1, 5],

[18, 48, 99, 26, 99, 99, 28, 17, 99, 33]]

solution = 'Map(lambda u1: If(Equal(u1, f), r, u1), x)'

In inputs_dict, each of the entries for x, f, and r is a list of length 3,
which contains the input for each of the 3 examples. Similarly, outputs is a
list containing the output for each example. solution is our handwritten
solution.

LambdaBeam $+$ Restarts finds the same solution of weight 10 in each of the 5
trials, taking a median time of 202 seconds:

⬇

Map(lambda u1: (lambda v1: If((lambda v1: Equal(f, v1))(v1), r, v1))(u1), x)

The solution looks complicated due to the Merge operation causing lots of
variable renames (i.e., $a_{k}(i_{k})$ in the Merge definition). We have
implemented an algorithm to simplify the solution by statically resolving
these renames. In this case, the solution simplifies to

⬇

Map(lambda u1: If(Equal(f, u1), r, u1), x)

which is essentially identical to the ground-truth solution.

###  D.2 Handwritten Task “multi:multiply_odds”

This task has 1 input and uses multiple higher-order functions to compute a
running product of only the odd elements:

⬇

inputs_dict = {

'x': [[3, 5, 8, 2, 1],

[5, 2, 1, 3, 3, 1, 4],

[3, -4, -1, 8, 2, 0, -3, 0, 9, -1]],

}

outputs = [[3, 15, 15],

[5, 5, 15, 45, 45],

[3, -3, 9, 81, -81]]

solution = 'Scanl1(lambda u1, u2: Multiply(u1, u2), Filter(lambda u1:
IsOdd(u1), x))'

In each of the 5 trials, LambdaBeam $+$ Restarts finds the same solution of
weight 11 that simplifies to the ground-truth solution, taking a median time
of 75 seconds.

###  D.3 Synthetic Task “synthetic:weight_9_function_7”

This task clips every element to the range $[0,4]$:

⬇

inputs_dict = {

'x1': [[-9, -2, -10, -6, 0, -10, -6, 3, 1],

[-1, -5, 8, 5]]

}

outputs= [[0, 0, 0, 0, 0, 0, 0, 3, 1],

[0, 0, 4, 4]]

solution = 'Map(lambda u1: Min(4, Max(0, u1)), x1)'

LambdaBeam $+$ Restarts finds a correct solution in all 5 trials with a median
time of 38 seconds, but the solutions are slightly different (the simplified
solutions are listed):

⬇

ZipWith(lambda u1, u2: Min(4, Max(0, u2)), x1, x1)

ZipWith(lambda u1, u2: Min(4, Max(0, u1)), x1, x1)

Reverse(ZipWith(lambda u1, u2: Min(4, Max(0, u2)), x1, Reverse(x1)))

Reverse(Map(lambda u1: Min(4, Max(0, u1)), Reverse(x1))) # found in two trials

Note that these are not the shortest solutions, but nevertheless all of these
solutions are equivalent to the ground-truth solution. LambdaBeam’s solutions
could benefit from a postprocessing simplification step, as discussed in
Section 4.3.

##  Appendix E More Details on LambdaBeam Architecture and Training

In our experiments, we used the following hyperparameters for the LambdaBeam
model architecture and training procedure. Refer to Figure 2 for a diagram
showing how the different modules interact.

  * •

I/O Module: this encodes a property signature of the I/O examples using a
2-layer ReLU-MLP with hidden size and output size of 512.

  * •

Value Module: this encodes each value’s property signature using a 2-layer
ReLU-MLP with hidden size of 512 and output (embedding) size of 256, with a
layer-norm applied after each linear projection. We use different MLPs for
lambda and non-lambda expressions.

  * •

Search Context Summary Module: this module needs to represent the entire
search state at the current stage, including the current operator to be
expanded, the I/O specification, and the values explored so far. We compute
the average of the set of value embeddings, concatenate it with the I/O
embedding, and then apply a projection layer (denoted as $MLP_{op}$ in Figure
2, which projects back to the embedding dimension) to get a vector
representation. The model parameters used in the projection layers are indexed
by the operator (i.e., we use different sets of trainable parameters for
different operators).

  * •

Argument Selector Module: we use an operator-specific 3-layer LSTM with hidden
size of 256. The prediction head is a 2-layer MLP with hidden size of 512.

  * •

During training, we generate on-policy data with beam size 10, use an
effective batch size of 32, and use the Adam optimizer with a constant
learning rate of $5\times 10^{-4}$.

  * •

During evaluation, we use UniqueRandomizer with beam size 10.

##  Appendix F Analysis of Handwritten and Synthetic Tasks

Table 1 shows some differences in the distributions between our handwritten
and synthetic evaluation tasks. This analysis may help contextualize the
experimental results in Section 4.

For example, Figure 5 shows that the LLM solved abnormally many synthetic
tasks in the 11-12 weight bucket. In fact, for synthetic tasks of weight 8 or
more, every one of the LLM’s “solutions” are actually false positives using
some form of “if the input is $\langle$hardcoded$\rangle$ then return
$\langle$hardcoded$\rangle$” logic, which is easier to implement when the
output is an integer as opposed to a list. Table 1 shows that there are
abnormally many synthetic tasks of weight 11-12 that have integer outputs,
which helps to explain the results.

Table 1: Analysis of some differences in task distributions between the handwritten and synthetic evaluation tasks. We show the number of tasks with each weight, the number of tasks using a lambda expression in the ground truth solution, and the number of tasks where the desired output is an integer (as opposed to a list). There may also be other axes where the handwritten and synthetic tasks have different distributions. | 100 Handwritten Tasks | 100 Synthetic Tasks
---|---|---
Weight | # Tasks | With Lambda | Int Output | # Tasks | With Lambda | Int Output
03 |  000 | – | – |  010 |  03 30% |  02 20%
04 |  003 |  02 067% |  02 67% |  010 |  02 20% |  06 60%
05 |  007 |  00 000% |  03 43% |  010 |  03 30% |  05 50%
06 |  010 |  09 090% |  04 40% |  010 |  07 70% |  03 30%
07 |  009 |  06 067% |  07 78% |  010 |  05 50% |  05 50%
08 |  011 | 11 100% |  04 36% |  010 |  08 80% |  02 20%
09 |  016 | 16 100% |  03 19% |  010 |  08 80% |  03 30%
10 |  009 |  09 100% |  01 11% |  010 |  08 80% |  03 30%
11 |  011 |  08 073% |  05 45% |  010 |  06 60% |  07 70%
12 |  009 |  09 100% |  03 33% |  010 |  03 30% |  07 70%
13 |  005 |  05 100% |  02 40% |  000 | – | –
14 |  003 |  03 100% |  01 33% |  000 | – | –
15 |  004 |  04 100% |  01 25% |  000 | – | –
16 |  002 |  02 100% |  00 00% |  000 | – | –
17 |  000 | – | – |  000 | – | –
18 |  000 | – | – |  000 | – | –
19 |  001 |  01 100% |  00 00% |  000 | – | –
Total | 100 | 85 085% | 36 36% | 100 | 53 53% | 43 43%