SlowGuess/NeurIPS_2023_OCR huggingface dataset:
LAMBDABEAM: Neural Program Search with Higher-Order Functions and Lambdas
Kensen Shi
Google DeepMind  
kshi@google.com
Google DeepMind hadai@google.com
Wen-Ding Li
Cornell University w1678@cornell.edu
Cornell University kellis@cornell.edu
Charles Sutton
Google DeepMind  
charlessbutton@google.com
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.
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. 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.
<sup>1</sup>DeepCoder [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.
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.
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,\dots ,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.
$\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
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.
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 (add, sort, map, 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. 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.
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
${}^{2}$  A 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.
(a) Term as a tree
Figure 3: Different ways of representing the lambda expression  $\lambda v_{1}$ . map  $(\lambda u.v_{1} + u^{2},\operatorname {sort}(x))$  where  $x$  is an input variable. The model predicts blue term pointers and variables, the search tries every red operation, and unshaded code tokens in (b) are known from the other shaded tokens.
(b) Term as straightline code
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
3.1 Building  $\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  $\mathfrak{t}_5 = (\lambda v_1.\mathrm{map}(\lambda u.v_1 + u^2,\mathrm{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  $\mathfrak{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  $\mathfrak{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 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 \ldots 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 \cdot v + 1)$  and  $a_2 = (\lambda v \cdot v - 1)$ , then we can share the variable by setting  $i_1 = i_2 = [v_1]$ , giving  $(\lambda v \cdot v + 1)(v_1) \times (\lambda v \cdot 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,  $\{i_k\}_{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  $\{a_k(i_k)\}_{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 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
$$
\operatorname {M E R G E} \left(f, a _ {1}, i _ {1}, a _ {2}, i _ {2}, \dots\right) = \lambda v _ {1} v _ {2} \dots f \left(\lambda u _ {1} u _ {2} \dots u _ {\ell_ {1}}. a _ {1} (i _ {1}), \lambda u _ {1} u _ {2} \dots u _ {\ell_ {2}}. a _ {2} (i _ {2}), \dots\right)
$$
$$
\text {w h e r e} \left\{v _ {1}, v _ {2}, \dots \right\} = \bigcup_ {k} i _ {k} - \left\{u _ {j} \right\} _ {j = 1} ^ {\max  \left\{\ell_ {1}, \ell_ {2}, \dots \right\}} \tag {1}
$$
where  $\ell_{k}$  is the arity of the  $k^{\mathrm{th}}$  argument to  $f$ . For example, the higher-order function 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., MERGE(square,  $v_{1}$ ,[]) returns  $\lambda v_{1}$ . square  $(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  $\mathfrak{t}_1$  through  $\mathfrak{t}_5$  have weights 1, 2, 5, 2, and 10 respectively.
3.2 Learning over  $\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 zipwith  $(\lambda u_1, u_2 \cdot u_1 + u_2 \cdot x, \text{sort}(x))$  and  $\text{scanl1}(\lambda u_1, u_2 \cdot u_1 + u_2 \cdot x)$  use the same lambda expression  $\lambda u_1, u_2 \cdot 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: