LambdaBeam: Neural Program Search with Higher-Order Functions and Lambdas
Authors: Kensen Shi
Hanjun Dai
Wen-Ding Li
Kevin Ellis
Charles Sutton
Venue: NeurIPS 2023 poster (2023)
Keywords: Program Synthesis
Programming By Example
Lambdas
Functional Programming
TL;DR: We design, train, and evaluate a neural program synthesizer that can construct arbitrary lambda functions for use within higher-order functions to perform looping computations.

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.

LAMBDABEAM: Neural Program Search with Higher-Order Functions and Lambdas
Kensen Shi Google DeepMind kshi@google. com
Hanjun Dai Google DeepMind hadai@google.c com
Wen-Ding Li Cornell University w16788cornel1. edu
Kevin Ellis Cornell University kellis@cornell.co edu
Charles Sutton
Google DeepMind
tharleessutonn@goggee.ccm
Abstract
Search is an important technique in program synthesis that allows for adaptive strategies such as focusing on particular search directions based execution results. Several prior works have demonstrated that neural models are effective at guiding program synthesis searches. However, 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 search algorithm called LAMBDAABE. that can construct arbitrary lambda functions that compose operations within given DSL. We create semantic vector representations of the execution behavior of the lambda functions and train a neural policy network choose which lambdas to construct during search, and pass them as arguments to higher-order functions perform looping computations. Our experiments show that LAMBDABEAM outperforms neural, symbolic, and LLM-based techniques in an integer list manipulation domain.
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 algorithms can learn to adapt to information collected during the ongoing search, such as exccution results other analyses of candidate programs considered 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] neural models to define search space that is explored by traditional non-neural search, while BUSTLE [27], CROSSBEAM [33]. and Execution-Guidde Synthesis [7] neural models guide the search process itself. However, those prior works unable to generate programs with arbitrary looping computations, whether implemented via loop control structures 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 more systematic search strategy that adapts to valuable information like execution results during the search.
also play key role in PBE, because well-designed search
DeepCoder [3] supports higher-orde order functions but only a small set hardcoded lambda functions. Execution Guided Synthesis [7] supports variable- free while loops, but not loops with an iteration variable.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

The fundamental question explored in this paper is whether 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 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 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.
expression
In order to design 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 Az. x + 1, λy- y + 1, and λz, y. x + 1 in a canonical way to prune the search space. If Az. + 1 is their canonical representation, then how can we reuse that lambda to create a new lambda such as Az, y. (r + 1) x (y + 1) where the "o + 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 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 '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 LAMBDA BAMBDABEAM, 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-explore values, using neural search policy to choose the operation's '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 that LAMBDABI BEAM outperforms prior approaches including state-of-thee-rt symbolic search, neural sequence models trained from scratch, and 62 billion parameter large language model (LLM). We release our LAMBDABEAM code and trained model checkpoints at https: ://gi ://gtthub..../////// com/ ellisk44//LamdaBeen
tasks, we experimentally show
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 Nomain-speciffe language (DSL) L describing space of programs, and set of example inputs I = ..1..... IN and corresponding outputs O = {OI.... ON }. The goal is to find a program P E L such that P(I,) O for all í E {1...., .). The DSL L describes atomic values (constants and input variables) and operations that can be applied to arguments to produce new values. Programs in L are arbitrarily-nested compositions of operations applied atomic values other such compositions.
X-Calculus The lambda calculus [10, 28] is a formalism for universal computation. lambda calculus term is either a variable, function application, or lambda abstraction. Lambda abstractions
2

---

Search Context
Featurize lambdas
Constrained
Input / Output Examples
LAMIUABEAM
sequence prediction
Arguments Variable Tuples
First-order operations higher-order operations
Model
DSL
Operation
Main Synthesis Loop
Inpuls,
Explored Values
f(+)
Constants
(Inclading Lambdas)
Create new terms with Miseet:
Figure 1: Overview of LAMBDA BEAm BEAM which builds upon the prior work CROssBEAM (new elements in LAMBDABEAM shown in red).
Quanition
Argument
Seleclor Module
Take (int, limt)
Set of Eaploned Veluee
list
Expr
Woight
(vakee cmbeddings)
Checking
Checking
1/2
V/3
(art
Th Ay 71 12,
LSTM
LSTM
Domain
Output Values
specific
Embadiing Lookup
Embedding Lookup
Ex. [4, 11 (6, 2,
Search
Context
Input Values
Summary
(search
Modile
context summary)
Ex.
12,
Predicion
Search Context
Model
Take cy,
Sort (x))
Figure 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 term). We consider lambda abstractions that are allowed 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 X-caleulus 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 neural program synthesis search, many different expressions are encountered, possibly including lambda functions with arbitrary input and output types. How might a neural model Iearn to reason about these expressions? One option is to represent expressions 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-basee -baed testing in the programming languages literature [16]. Unary properties describe single value by mapping to a boolean, such as whether a list is sorted, or whether a string empty. Binary properties can describe the relationship between the input and output of PBE task, such whether the input and output lists have the same length. In either case, given a list property functions, we can evaluate all property functions form 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 more granular representation.
CROSsBEAM Our work LAMBDABEAM builds upon the prior work CROSsBEAM [33]. Both systems have a similar overall structure, illustrated in Figure with differences shown red. The core
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 emor.

---

Term
(b) Term straightline code
(c) Built using MERGE
X
ti
atomic term from DSL
map
02 f- Au1. square
MERGE( square, 01. [])
Au
sort
t3 + Au,, U2. add (U1 D2 (2)
MERGE add, U1 t2
[02]
add
æ
t + sort t1
MERGE sort, t1
square !.
t; Lty map Au.
13
(1 u), tq
MERGE map, L3 U1
t1 (])
Figure 3: Different ways of representing the lambda expression Au1. map Au. + u² sort(2)),
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.
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 powerful signal for the policy 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 Es.. Meanwhile, the I/O Module encodes the 1/O examples into an analogous vector ero. Then, the Search Context Summary Module combines the value encodings El and the 1/O encoding eso to produce latent "summary" e, of the search context so far. From this, the Argument Selector Module, which is recurrent pointer network [35], predicts an argument list for DSL operation via sequence of pointers into the table of explored values, thus generating a new expression to explore. By structuring the search this way, the neural policy is able take "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 1/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 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 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 CROSsB BEAM randomly samples different argument lists using UniqueRandomizer [31] avoid duplicate samples.
3
The LAMBDABEAM Approach
3.1 Building A-Terms
LAMBDABEAM constructs terms the lambda calculus during its search. natural design choice is to construct terms in a way that avoids considering semantically equivalent expressions. For example, the terms (Xx, y.. x - y), (Aa,
b). and (Ay, x. y x) are capable of expressing exactly the same computations, so there should be 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). Enforeing 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 onput-output behavior. However, the most straightforward way of building lambda terms bottom-up violates this important property. Consider the term t; (Ae1. map(Au. 01 + u² sort(x))) whose tree structure is illustrated in Figure 3(a), and where x refers to an input variable. This has subterms such as Au. ei + u², where the variable 2i occurs free. As vi is free, it is unclear what the semantics of this expression should be. Constructing the expression with en free would also introduce spurious redundancy with Xe,, U1 + u². During search, we would like to keep only canonical version of

---

those two terms to better prune the search space, which in this case would be t3 = Au1, U2 01 + 222 which as shown in Figure 3(b) can be used construct the larger term Tz.
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 and a list of terms 01 ... ak constructed earlier during search, and it returns a new term that applies I to ai a.... a.. MERGE does extra bookkeeping to ensure that every function (both f and any og 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 a1 a1...-4 a..- To evaluate MERGE on f and a1 .1...0k. we first apply each lambda argument dg tuple of variable names, denoted ig. This gives a name to each argument of each 0g- By reusing the same name across different og ss the same variable can be shared across arguments. For instance, the function f is multiplication (x), and we are merging with the arguments aD = (Au. v + 1) and a2 = (Av. F- - 1), then we can share the variable by setting in = i2 = [U]]. giving (Av. e + 1)(v1) x v = (u1 + x (V1 - 1). Alternatively, if we set 1 [v1] and in [02]. then merging gives different program, (v1 + 1) (02 1). These tuples of variable names, {ix) are also input to MERGE, because they determine how variable names are unified across arguments. Finally MERGE runs f on the arguments {au()s)) bind any new free variables, and lastly canonicalizes variable naming and variable ordering.
pads the expression with lambdas at the beginning to
Special care is needed for the arguments of higher-order functions like map, which themselves need to be functions. So far, each argument ag (fx) evaluates higher-order functions, MERGE automatically adds extra lambdas to each function-valued argument. These extra lambdas have variables denoted u1 122 .... All other variables used to unify variables across arguments are denoted U1, U2,.... Although this may seem ad-hoc at first, this convention actually corresponds to well-known canonicalization of typed lambda forms known as n--long normal form [28]. Putting everything together, MERGE is defined as
concrete value, not a function. To handle
MERGE (f, (1, à1, 02 i2....) Xey 02... ................................ f (Au1 U2. a1 (i(11, Aui az(t2)....)
where (21.V2,... = U {uy max (2
(1)
k
where l is the arity of the hthh argument to f. For example, the higher-order function map first takes function expecting one argument, so E 1, followed by a list (not a function) expecting no arguments, so l2 = O. We also enforce that | = arity(a;). The MERGE operation is complete, in the sense that we can use 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 da by pointing into the set of explored values, and variables in by emitting special tokens corresponding to vi, uI, U2, U2, etc. Figure 3b) and (c) illustrate how a nontrivial lambda expression 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, c.g., MERGE( square, Ui. []) returns Wun, square (v1) 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- fine-grained feedback to the model.
We define the weight of expression 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 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 t through ty have weights 1, 2, 2, and 10 respectively.
3.2 Learning over X- X-Expressions
One core technical challenge how to encode lambda expressions to allow neural models to reason about them. In LAMBDABEAM we solve this by constructing new generalization of property
5

---

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(Au,, ug. u1 + u2, I. sort(x)) and scanll (Auu, U2. #1 + U2, 2) use the same lambda expression Aui, U2. U1 + U2 but for different sets of [u1, (z) arguments. Thus, in order to describe the lambda expression 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, 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 inclusive, and other integers with varying combinations of magnitude, sign, even/odd parity, primality, and so on. There is also a tradcoff' 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 round-robin fashion.
To represent lambda First we evaluate f on t;, yielding a result r; = f(t,, Z). Then we use signature of unary properties to describe rı. Second, we use a signature of binary properties compare r; and O: intuitively, this helps the model learn about what work remains. Similarly, we use the binary properties to compare r, and laj. for each argument l2y E 44. By concatenating these, we obtain a single property vector for each tuple f.. Finally, we then reduce the property vectors across the runs of the lambda, i.c., computing for each property the fraction of runs where is applicable and the fraction of applicable runs where it is True. Encoding non lon-lambda expressions is similar, except that we use I/O examples (4) 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).
evaluate it each canonical argument tuple & with 1/O example (T, (T,
The properties used in our property signatures come from combinatorially combining hand-designed features as "building blocks" to create 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, 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 S which contains variable tokens for construct- ing lambdas, lambda expressions, and non-lambda expressions including constants and input variables. The Value Module embeds each element of S forming a matrix of embeddings Es E R (g(5)xx xd Ele- ments of S are embedded as follows. A variable token is embedded as vector of trainable parameters. Note that the set of such variable tokens is fixed and determined by the DSL. A lambda expression is embedded by + where is the property signature of this lambda function encoded by an MLP,, and is an embedding of the weight of this value. Non-lambda expressions are embedded like lambda expressions except using a different MIP to encode their property signatures.
Argument Selector Module Given an operator op, we use an operator-specific LSTMsp to select the arguments using a pointer mechanism [35] from the encoded value matrix Ess. in an autoregressive way. In addition to selecting (op)y arguments, for an argument dj that is lambda expression, we also need to predict the variables is as required for MERGE, where i& is a tuple of arity (ax) variable tokens. All of the ag and is predictions are done as single autoregressive sequence. Furthermore,
'The higher-ordet 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 U1, 12, M1, and u2.

---

for convenience we predict all of the ag arguments first, followed by the ig variable tuples which are constrained (via masking and padding) include exactly arity(a;) variable tokens.
Search with Restarts We also change the inference time search procedure. Recall that CROSS- BEAM random sampling during evaluation, making the search nondeterministic. In LAMBDA- BEAM, instead of performing one synthesis search until timeout, we restart the search from scratch whenever the search has run for a certain amount of time without finding a solution. Even though this discards work done in previous searches, in practice this helps 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 '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, 34j. However, DeepCoder's DSL only contains a hardcoded set of lambda functions and 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 (Ax. z 1). (Ax, y. - y), and (Az, y. max vH). 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 (Ax. max(2r, 0 100-)) xz}) 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 [- 1--25,,55 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 I/O examples and between I and input variables, and we sample tasks such that approximately 80% of them use lambdas in the ground-trum program. We used a time limit of I 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 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, I/O examples if the output is list or I/O examples the output an integer, and a handwritten ground-truthh 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
7

---

10 tasks of each weight between 3 and 12 inclusive. Appendix D contains example handwritten and synthetic tasks, along with LAMBDABEAM'S 's solutions for them.
Approaches Our experimeents compare several approaches:
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 V 100 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. LAMBDA BEAM without random restarts: We use the LAMBDA BEAM approach without random restarts as an ablation, also for 5 trials with different randomness for UniqueRandomizer sampling.
3. Enumeration: We run the same exhaustive bottom-up enumerative synthesis algorithm that was used to create the training data. We use 5 trials with different random orderings of DSL. operations, which changes the enumeration order of programs with the same weight.
4. RobustFill [12]: This approach treats the synthesis problem 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 LAMBDA BEAM 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. X2 [15]: This is wate-of-the-an art symbolic program synthesizer that handles lambda functions and higher-order functions. We implemented our DSL within the X2 framework (using the more extensible version provided by the X authors). X is deterministic so we only use 1 trial.
6. Python-Finetuned LLM: We try asking a pretrained large language model (LLM) to solve our evaluation tasks using Python code. Specifically, we use PaLM 62B that was trained for longer as described in Appendix F of Chowdhery et al. [9], with further fine-tuning on general Python code. The prompt contains natural language instructions and 2 examples of an I/O specification followed by Python code that solves the task (for 2 new tasks), and then the I/O specification of the evaluation task we wish to solve.* 4 We repeatedly draw batches of 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 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 X² 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 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 larger programs compared the other methods, except the LLM which has more constant but lower success rate overall."
'In preliminary experiments, we found that providing few-shot examples in the prompt led to slower sampling without much change in program quality. On the other hand, using only I 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 'The LLM predicts Python code instead of using our DSL, so the weight according to our DSL is an inaccurate measure of the complexity of the corresponding Python code. Furthermore, difficulty for the LLM is more closely correlated with how "natural" the task is, i... its similarity programs in the LLM's training data.
DSL.
8

---

100 Handwritten Tasks
100 Synthetic Tasks
100
100
80
76.8
Tooks Solved
57.2
71.6
60
60
LAMBDADEAM
Reslarts
LaMIaABaAM
40
54
40
33,0
Enumeration
20
50.8
20
A
RobustFill
47.8
Python-Finetuned LLM
0
0
100
200
300
400
500
600
0
100
200
300
400
500
600
Elapsed Time (sec)
Elapsed Time (sec)
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.
Handwritten lasks
Synthetic Tasks
100
100
Success Rate (%)
LAMBDABEAN Restarts
LAMBDABEAM
Enumeration
40
40
RobustFill
20
20
A²
-00- Python-Finetuned LLM
0
4-6
7-8
9-10
11-12
13-19
3-4
5-6
7-8
9-10
11-12
Task Weight
Task Weight
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 least 15 tasks.
Running LAMBDA BEAM with random restarts helps overall but more 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 stay 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 X², enumeration, or RobustFill's 's beam search, and would not help the LLM where each sample is already independent.
We also identify false positives by running solutions on 2 held-out test cases per task, generated mostly synthetically with some manual editing. The results are in Figure 6, showing that LAMBDABEAM with restarts has the highest number of true positive solutions on handwritten tasks by margin of nearly 8 tasks, while barely losing to Enumeration on synthetic tasks." 6 While symbolic approaches (Enumeration and 2 have fewer false positives due focusing on small solutions, we observe that LAMBda B BEAM has the fewest false positives among the neural approaches. The LLM produces many false positives 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 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. 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 hand- written and synthetic evaluation tasks, which helps to contextualize the experimental results. For example, lambda expressions are used in 85% of the handwritten tasks but only 53% of the synthetic tasks. The median task weight is 9 for handwritten tasks and only 7.5 for synthetic tasks. These comparisons suggest that the handwritten tasks are harder than the synthetic tasks on average, which
'The synthetic tasks have aandomly-gnerrtted handwritten tasks, and the "correct" solution is not chosen to be natural but rather one with minimal weight by construction. Enumeration has an unfair advantage here, being the only method in our comparison that is guaranteed find minimal weight solution.
examples that are overall less informative than those in the

---

Handwritten Tasks
Synthetic Tasks
100
100
03.6
Tasks Solved
80
80
82.6
75.0
60
47.6
60
67.4
02.0
00.4
40
42.3
40
22.0
Falsc Positives
20
20
True Positives
LB+R
Cnur
LLM
LD+R
Cnum
LLM
Figure 6: True positive versus false positive solutions measured using held-out test cases.
also reflected in the overall performance in Figure 4. 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.
We observe that LAMBDA BEAM achieves
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 UniqueRandomizee 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 general programming languages.
5 Related Work
Machine learning for program synthesis has been an active area [17, 18, 1J. 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 inspired by work in learning to search [11, 29, 5] and beam-aware aware training [23, 24].
In contrast, we unaware of previous work that synthesizes helper functions, such as lambda functions, during neural program synthesis. The original DeepCoder DSL contains only small of predefined lambda functions. Even within symbolic program synthesis, X² is one of the few examples of work that synthesizes lambda expressions [15]. To control the size of the search space, 2 employs type-directed synthesis, but we handle more general domains where the type system not informative enough to reduce the search space sufficiently. DreamCoder [14] can also infer ad hoc helper functions like X2, but its neural network provides no fine-graie fine-grained guidance on how to compose those lambdas. Because DreamCoder is an algorithm for enriching impoverished DSL to improve neurally-guides program search, one could combine DreamCoder's DSL. enrichment process with RAMBDAAEAM 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].
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 leams neural policy to drive a bottom-up search over programs. We experimentally show that LAMBDA BEAM outperforms symbolic search, a sequence model, and pretrained code LLM with 62 billion parameters.
10

---

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
[1] Miltiadis Allamanis, Earl Barr, Premkumar Devanbu, and Charles Sutton. A survey of machine learning for big code and naturalness. ACM Computing Surveys (CSUR), 51).. 2018.
[2] Rajeev Alur, Arjun Radhakrishna, and Abhishek Udupa. Scaling enumerative program syn- thesis via divide and conquer. In International Conference Tools and Algorithmss for the Construction and Analysis of Systems (TACAS), pages 319-336. Springer, 2017.
[3] Matej Balog, Alexander L Gaunt, Marc Brockschmidt, Sebastian Nowozin, and Daniel Tar- low. DeepCoder: Learning to write programs. In International Conference on Learning Representations (ICLR), 2017.
[4] Shraddha Barke, Hila Peleg, and Nadia Polikarpova. Just-in-time learning for bottom-up enumerative synthesis. In Object- (OOPSIA), 2020.
Object-Oriented Oriented Programming, Systems, Languages, and Applications
[5] 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 (ICMI), 2015.
[6] 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, 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.03377 2021. URL https: //arxiv. org/ abs 2207 03374.
[7] Xinyun Chen, Chang Liu, and Dawn Song. Execution-guided neural program synthesis. In International Conference on Learning Representations (ICLR), 2019.
[8] Xinyun Chen, Dawn Song, and Yuandong Tian. Latent execution for neural program synthesis. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
[9] 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 Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Pa Wei, Kathy modeling with pathways. Journal of Machine Learning Research (JMLR), 24: Scaling language
240:1-240:11: 113: 2023. URL http: / //jp:///mmmrr//r org/ happers///4/// 1144. html.
[10] Alonzo Church. The calculi of lambda--convesici Number Princeton University Press, 1985.
11

---

[11] Hal Daumé III. John Langford, and Daniel Marcu. Search-based structured prediction. Machine Learning. 75(3), 2009.
[12] 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.
[13] Kevin Ellis, Maxwell Nye, Yewen Pu, Felix Sosa, Josh Tenenbaum, and Armando Solar-Lezama. Write, execute, assess: Program synthesis with REPL. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
[14] Kevin Ellis, Catherine Wong, Maxwell Nye, Mathias Sablé-Meyer, Lucas Morales, Luke Hewitt, Luc Cary, Armando Solar- and Joshua B Tenenbaum. DreamCoder: Bootstrapping inductive program synthesis with wake-sleep library learning. In Programming Language Design and Implementation (PLDI), pages 835-858-0 850, 2021.
[15] John K Feser, Swarat Chaudhuri, and Isil Dillig. Synthesizing data structure transformations from imput-output examples. In Programming Language Design and Implementation (PLDI), pages 229-239, 2015.
[16] George Fink and Matt Bishop. Property-hasee testing: a new approach to testing for assurance. ACM SIGSOFT Software Engineering Notes, 22(4):74-80) 80, 1997.
[17] 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.
[18] Sumit Gulwani, Oleksandr Polozov, and Rishabh Singh. Program synthesis. Foundations and Trends in Programming Languages, 4(1-2), 2017.
[19] Joey Hong, David Dohan, Rishabh Singh, Charles Sutton, and Manzil Zaheer. Latent program- mer: Discrete latent codes for program synthesis. In International Conference on Machine Learning (ICML), 2021.
[20] Woosuk Lee, Kihong Hco, Rajeev Alur, and Mayur Naik. Accelerating search-based pro- gram synthesis using learned probabilistic models. In Programming Language Design and Implementation (PLDI), 2018.
[21] Zohar Manna and Richard J Waldinger. Toward automatic program synthesis. Communications of the ACM, 14(3):511165, 1971. URL https://doi. org/ 10. 1.0.114400 1145/362266..... 362568.
[22] Vijayaraghavan Murali, Letao Qi, Swarat Chaudhuri, and Chris Jermaine. Neural sketch learning for conditional program generation. In International Conference Learning Representations (ICLR), 2018.
[23] Renato Negrinho, Matthew Gormley, and Geoffrey Gordon. Learning beam search policies via imitation learning. In Advances in Neural Information Processing Systems (NeurIPS), 2018
[24] 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.
[25] Maxwell Nye, Luke Hewitt, Joshua Tenenbaum, and Armando Solar-Lezama, Learning to infer program sketches. In International Conference Machine Learning (ICML), 2019.
[26] Augustus Odena and Charles Sutton. Learning represent programs with property signatures. In International Conference on Learning Representations (ICLR), 2020.
[27] Augustus Odena, Kensen Shi, David Bieber, Rishabh Singh, Charles Sutton, and Hanjun Dai. BUSTLE: Bottom-up program synthesis through learning-guide exploration. In International Conference on Learning Representations (ICLR), 2021.
[28] Benjamin C Pierce. Types and programming languages. MIT press, 2002.
12

---

[29] Stephane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction no-regret online learning, Conference on Artificial Intelligncee and Statistics (AISTATS), 2011.
[30] Kensen Shi, Jacob Steinhardt, and Percy Liang. FrAngel: Component-based synthesis with control structures. Proceedings of the ACM on Programming Languages, 3(POPL), 2019.
[31] Kensen Shi, David Bieber, and Charles Sutton. Incremental sampling without replacement for sequence models. In International Conference Machine Learning (ICML), 2020.
[32] Kensen Shi, David Bieber, and Rishabh Singh. TF-Coder: Program synthesis for tensor manipulations. ACM Transactions on Programming Languages and Systems (TOPLAS), 44(2): 1-36, 2022.
[33] Kensen Shi, Hanjun Dai, Kevin Ellis,, and Charles Sutton. CrossBeam: Learning search in bottom-up program synthesis. In International Conference on Learning Representations (ICLR), 2022.
[34] Disha Shrivastava, Hugo Larochelle, and Daniel Tarlow. Learning to combine per-example solutions for neural program synthesis. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
[35] Oriol Vinyals, Meire Fortunato, and Navdeep Jaitly. Pointer networks. In Advances in Neural Information Processing Systems (NIPS), 2015.
[36] Pengeheng Yin and Graham Neubig. A syntactic neural model for general-purposs code generation. In Assocation for Computational Linguistics (ACL), 2017.
[37] Amit Zohar and Lior Wolf. Automatic program synthesis of long programs with a learned garbage collector. Advances in Neural Information Processing Systems (NeurIPS), 2018.
13

---

Appendices
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 grogramming-by-eeeee
(PBE) problem.
We formalize our DSL in a subset of the lambda calculus. Let X' [ ..t........... Fm be the set of imput variables for the PBE task, V be countable set of variables that is disjoint from X, F be the set of primitive functions the DSI,, and be set of constants in the DSL.. Then, our lambda calculus is:
T ::= x c f(t tx) 101
for E X, C, FF
v, 0I...... Un E V, t, t T.
Let M be the set of terms obtainable by repeatedly applying MERGE (including the initial terms usable by MERGE);
:: x Av. v e MERGE a1, "..,,,..., ik) for X, E V.C E C, f F, a1 a E M,
E V..
We restate the definition
MERGE (f, a, 21, a2, .2.... Ae D2... (Au U2... ai (fi) uu. a2(12)...)
where U2,.......) = U is [m)]martee max(ei, .2....
From this definition, it is clear that M C T, i.e., MERGE is closed within the lambda calculus. However, M * T because MERGE imposes certain constraints, egg., Av. a in but cannot be constructed by MERGE. To precisely describe the constraints resulting from MERGE, introduce the following definitions:
Terms x E v E and e are atomic. term s = A..... t (possibly with n = such that is not a lambda expression) has exact lambda variables if Free Vars(t) X' {v1, {uI,-... -0u.....,.. ), Note that having exact lambda variables implies that Free FreeVarst Vars(s) C X.
A term typechecks if every function application has the correct arity for every argument, e.g.. Map(t1, t2) expects t1 to have arity 1, while t2 should have arity 0.
We now let S {s E T s has exact lambda variables and typechecks].
The completeness MERGE, in the sense that can generate all solutions s to PBE problems (once the input variables ; Fin bound), follows from the two claims below.
Claim 1. p Am is a solution PBE problem, such that p(11...., rm) y for all I/ examples (æI..... .mm) + in the PBE specification, then S. That is, S is broad enough to cover all solutions to PBE problems.
Proof. To ensure that p Ar ......mm s a valid solution program, we must have Free FreeVars(s) 0 (so there are no unbound variables), s must have arity 0 (since all inputs æ1,..., Pm are already bound), and must typecheck to avoid runtime errors. Together, these imply that S.
Claim 2. S C M. That is, MERGE can create any term in S, including all solutions to PBE problems.
14

---

Proof. Let s be any term in S. We will proceed by induction on the depth of s.
As the base case, if s is atomic, then v because v does not have exact lambda variables.
c, so immediately in M. Note that cannot be
Then, assume the inductive hypothesis that any term in S, with depth less than that of is in M. There are two inductive cases to consider: either s = xn f(t],..., tq) where n might be 0, or u. where n 0 and t is atomic. Because s has exact lambda variables, the latter scenario is only possible if s = λv. v, which is immediately in M.
The remaining case is s = A f(t1..... aj and iy such that MERGE(F.. a1, .1....,.........................
..t.....ttt).. t²). Consider any t, for 1 S S k. We will construct
....................... aa, 2x) =.
t is atomic, then either , I, 1 u. or tú = c. tf = u. then a; Au. and ij [o]: otherwise, set a, = 1 and éj [], the empty tuple. In each case, aj E M. Because typechecks, we know that the j-th argument to f expects arity O. so in the MERGE definition, Ej = 0 and thus the j-th argument to f expands a (tn) tú in each case.?
ty is not atomic, then let E = Me. r, where f is the expected arity of the j-th argument (because typechecks). Let {uf..... f/ Free Vars(r) -X. Set a, >e and i [e'......' V'al. so that when applying MERGE, the j-th argument to f expands to We, (2,)) Aul...ue.. r = tj. Furthermore, a, has exact lambda variables by construction, and typechecks because r typechecks, so aj E S and a, E M by the inductive hypothesis.
With these choices of aj and ij. when expanding MERGE(J, a, iI, a in) according to the definition, the j-th argument becomes ty. and the lambda variables U2 in { jm f2.... are exactly Free FreeVarsi (a (,),,, = fvi,.... .......,,) since s has exact lambda variables. Therefore, un ftt......, t²) MERGE(f, a1, ii...., ak, it), E M.
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, (rpeproppetics(2) represents the type of x as boolean one-hot list. In our DSL, this returns a tuple of 5 booleans, representing whether is a lambda, boolean, int, list, or None (which used to indicate an error, e.g., signaling that an index is out of bounds).
Next, we define Basic Properties (a:) 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, N/A). Note that, if has type т, then all properties for type 7 evaluate True or False, while all properties for all other types r' * T evaluate to N/A. For our DSL, we use the following basic properties:
For boolean x: x itself.
For integer whether i equals -1, 0, 1, and 2; whether ir is positive and negative; whether is even; whether x is and modulo 3; and whether (x] is less than 5, 10, 20, 35, 50, 75, and 100.
For list whether is sorted, whether is sorted in reverse, and whether x contains all unique elements.
Then, (elevant(z) returns related objects that are relevant to understanding æ. For our DSL, this is only æ 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 ObjectSignature((a which takes single DSL object and returns fixed-length vector of property results, containing (ypeProperties(r) followed by Basic (nsiieroperties()) for each E Relevant(z:) For example, these properties include "assuming is int, is is even?" (a basic
'I practice for simplicity. when t, = v, we simply set ay ty and ig = [] as in the other cases, even though a, does not have exact lambda variables.
15

---

property applied to x)) as well as "assuming z is list, are there even number of elements in x²" (a basic property applied to an object relevant z). By applying basic properties to relevant objects in this compositional way, we reduce the effort needed to specify large number of properties.
We furthermore encode comparisons between two objects. We define Comparissonroopertisss Properties y) which evaluates hand-designed properties for comparing two objects z and y of the same type, for each different type in the DSL. This returns fixed-length vector of property results, where property for comparing type 7 evaluates to N/A if æ and y are not of type For our DSL, we use the following comparison properties:
For boolean x and y: whether x =
For integer x and y: whether æ = < y, and x > y, whether r is factor of y and vice versa; and whether [x y| is less than 2, 5, 10, and 20.
For list v and yy whether æ = y/ whether x is longer, shorter, or equal length compared to y; whether the lengths differ by at most 1; whether all x; yi for Ti, Vr E zip(æz, 3) and similarly for other comparisons S, >, 2, =, and =: whether x and y contain the same set of elements; and whether contains subset of elements compared to y and vice versa.
These properties are used in ComparisoonSignatureeaa y) which computes a fixed-length list of prop- erty results for any two DSL objects and y of any type, containing Comparison ropertteerr v) for all rz E Relevant(a (z) where type(rx) = type(y), and ComparisonProoerrtied type(z). Thus, "assuming is int and y is list, is æ 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.
ry) for all
Ty Relevant(y) where type(ry)
In the I/O Module of the neural policy (as in CROSS (33]), we use property signatures to encode set of I/O examples. For each example (L/),... In), O) we concatenate ObjectSignature( with ObjectSignature((.) and ComparisouSignaatuee
O) for all S i S 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 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 = 1fi, f using an I/O example (I.O) O) where the lambda evaluates to a result rt. we con- catenate ObjeetSignatuur(r]]]
CompartsonSSggaturr(r;; O), and ComparisonSignaturee.. for all 1 < j < m., and then reduce these across the runs of the lambda. To encode a non-lambda expression during search, for each I/O example (1,0) where the expression evaluates to result r, we concatenate ObjeclSignature(77 with ComparisonSignatui 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 outpur later, not how the value was created from the inputs.
r;)
O), and then reduce these across
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 property signature of length 359.
C Extension of the DeepCoder DSL
As mentioned Section 4. I, 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: (x: int, y: int) int: returr x ++
def Subtract (x: int, int) +> ints zeture x
def Multiply (X: int, yt int) +> int:
return x
16

---

def IntDivide (X: int, y: int) -> intr
def Square IX: int). -2
return
def (in(x: y: int) -> int: return min y)
def Max(x: int, int) -> int: return max (8, y)
def Greater int, y: int) +> bocl:
def int, y1 bool: return x x
def Equal int, int) bool: retuEr x ==
def IsEven Ix: int) -> bool: return
def IsOdd (x: int: bool:
8 != 0
def bool, int, ys intl int reture x if else y
def Head (xs: list) -> int:
return x6[0]
def Last (xG: 1st) +> ints xa[-1]
def (a: int, xs: list) list: return [s(:n]
def Drop (n: int, list) +> list: return xs [n:n]
def Access in: int, xE: list) -> int: return [n]
def Minimum (xS: list) +> intr (in(xs)
def (xs: 1ist] lnts returr max (xx(
def Reverse (XS: Histl +> List: roturr list (ravorsed
def Sort (xs: 1ist) >> List: return
(xs)
def Sum(xs: (xs: list: -> ints reture (xm)
higeer-order operations
dof Mapif:
callable| (int), int), xs: list) -> 1ist: return (E(x)
xs]
def Filterif' If: Callabl e [[nt]] bool], xat Hatt] >> Hat: zeturn
if f(x]
def Count callable (int), bool], xs: list) ints return [en([x for in xs if (())))
def ZipMith (f: bellablel[int intl, intl, list, list) is list: returr [E (x,x, y) for x, in zip (xs, ys)]
def Scanll (É: Callablel intl, intl, xs: list) >> listi ys [xs[0]] for in range (1, ien (xs)): ys appedd (f (ys [n-1], xs[n]))
17

---

D Example Tasks
This section contains selected example problems from 100 handwritten and 100 synthetic evaluation tasks. Each task is given a name for convenience purposes only, which not used by any method in experiments.
D.1 Handwritten Task
This task has inputs (x, f, and z). examples demonstrating the task ("in x, find instances of and replace them with r"), and handwritten ground solution using relatively complicated lambda function:
inputs_dis dict {
[[7,
6,
61.
[-6, -3, 4, 3, -, -3
2, 1,
5]
[18, 48 27, 26,
27,
28
27, 33]1,
1 [4, -3, 271
!= [-1, 7, 99]
outputa [[7, 2, -1, 6, -1, 2, 61
[-6,
-5,
51
[18, 48, 99, 26. 89 99 28
17. 99, 333]
aolution 'Nap (lambda u1: If (Equal (u1, f), 8, u1), x)
In inputs_dii dict, each of the entries for x, f, and r list of length 3, which contains the input for each of the 3 examples. Similarly, outputa is a list containing the output for each example. solution our handwritten solution.
CMMBDABEAM - Restarts finds the same solution of weight 10 in each of the trials, taking median time of 202 seconds:
Map (lambda u1: (lambda vi: If ((landda vi: Equal (f, (x1) r, (1)) x)
The solution looks complicated due to the MERGE operation causing lots of variable renames (i.e., ($x) in the MERGE definition). We have implemented algorithm to simplify the solution by statically resolving these renames. In this case, the solution simplifies to
Map (lambda u1: If (Bqual (f, u1), u1), x)
which is essentially identical the ground-trute solution.
D.2 Handwritten Task "multi :mmuliiippi :mlltppyy.odee
This task has input and uses multiple higher-order functions compute running product of only the odd elements:
Inputs_de dict = {
[[3, 5, 8, 2, 1]. 4]
[5
3,
[3, -4, -1, 8, 2, 0, 3, 0, 9, -1],
outputs [[3, 15, 16)
[8, 5, 15, 45,
[3, -3, 81, -811]
solution 'Scan11 (lambda ut, u2: Multiply u2), Filter (lambda ut: Iaddd (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 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_de diet { "x1": [[-9, -2, -10 -6, 0, 100 -, 3, 10.
[-1. -5, 8, 6]]
outputs= [o. o, 4, 41]
0,
11.
aolution "Nap(labba ui: (4, Max Max w1)), x1)
LAMBDA Restarts finds a correct solution all trials with a median time of 38 seconds, but the solutions slightly different (the simplified solutions are listed);
18

---

(landth u1, u2: (in (0, u2)), x1,
ZipWath (lambda ul, u2: (in Max Max u1)). x1, x1)
Reverse (Zipssth (1anbda ut, u2: Min Max (0, Reverse (Nap (lambda ul: (4, Mox u1)), (everse(x1))> # found in too trials
(everse(x1)))
Note that these are not the shortest solutions, but nevertheless all of these solutions are equivalent to the ground-truig truth solution. LAMBDABEAM'S solutions could benefit from a postprocessing simplification step, discussed in Section 4.3.
E More Details on LAMBDABEAS Architecture and Training
In our experiments, we used the following hyperparameters for the LAMBDABEAN BEAM model architecture and training procedure. Refer Figure for 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 's property signature using 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 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 projection layer (denoted as MLP.p 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., for different operators).
different sets of trainable parameters
Argument Selector Module: we use an operator-specific 3-layer LSTM with hidden size of 256.
The prediction head is 2-layer MLP with hidden size of 512.
During training, we generate on on-policy data with beam size 10, use an effective batch size of 32, and use the Adam optimizer with constant learning rate of 5 10
During evaluation, we use UniqueRandomizer with beam size 10.
F Analysis of Handwritten and Synthetic Tasks
Table I shows some differences in the distributions between our handwritten and synthetic evaluation tasks. This analysis may help contextualize the experimental results in Section
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 more, every one of the LLM's "solutions" are actually false positives using some form of "if the imput is (hardcoded) then return (hardcoded)" logic, which is easier implement when the output an integer as opposed a list. Table I shows that there are abnormally many synthetic tasks of weight 11-12 that have integer outputs, which helps explain the results.
19

---

Table I: 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 lambda expression in the ground truth solution, and the number of tasks where the desired output is integer (as opposed to 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
Output
Tasks
With Lambda
Int Output
3
0
10
3
30%
2
20%
4
3
2
%
2
67%
10
2
20%
6
60%
5
7
0
0%
3
43%
10
3
30%
5 50%
6
10
90%
4
40%
10
7
70%
3 30%
7
9
6
67%
7
78%
10
5 50%
5 50%
8
11
11
100%
4 36%
10
8 80%
2 20%
9
16
16
100%
3 19%
10
8 80%
3 30%
10
9
100%
11%
10
8
80%
3 30%
11
11
8
73%
5 45%
10
6
60%
7 70%
12
9
100%
3 33%
10
3 30%
70%
13
5
5
100%
2 40%
14
3
100%
33%
15
4
4
100%
1
25%
16
2
100%
0 0%
17
0
18
0
19
1
100%
0
0%
Total
100
85
85%
36
36%
100
53
53%
43 43%
20


Reviews:
[{"review_id": "AjjyJyoWaN", "reviewer": "Reviewer_ErYC", "summary": "This paper presents a method for training a neural module to guide a search-based program synthesis procedure that supports lambda functions. This is accomplished by leveraging the existing technique of property signatures, which essentially represent program constructs using a hand-designed vector of features. The authors design features for representing lambda functions, including evaluating the lambda function on a hardcoded set of inputs. This allows lambda functions to be incorporated in a prior bottom up program search technique called CrossBeam, and they introduce a new Merge operator to build lambda terms from the bottom up. The results on a modified version of the DeepCoder benchmark show their approach outperforms several strong baselines, including symbolic search and an LLM fine-tuned on Python.", "questions": "How many programs are there of weight at most 12 (i.e., that were sampled from for the training set)?\n\nWhat size is the property signature? (or how many properties are there)\n\nCan you elaborate on why you add an embedding of the weight of a value for embedding lambda expressions?\n\nIs Merge complete for lambda functions? This should be addressed in the paper.", "limitations": "The authors should also address the extent to which the property signatures are tailored to the DSL, and the implications for broader applicability / scalability.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Representing and synthesizing (lambda) functions is a significant step forward for neural program synthesis. Although the individual techniques are mostly prior art, I consider the combination of techniques to be novel.\n\nThe writing is generally clear, though I found the implementation details to be quite sparse in places (see Weaknesses)", "weaknesses": "The biggest weakness is that the approach is tested on only one synthetic dataset and also relies on hand designed features. Furthermore, the lambda functions only contain 2 variables. It remains to be seen if this approach could scale to more realistic datasets with more data types and more complex functions.\n\nAdditionally, there are almost no examples in the paper. At the very least, the appendix should include some examples of the synthesis dataset as well as programs synthesized by LambdaBeam.\n\nFinally, many of the implementation details are not listed out fully, which impedes reproducibility (e.g., the architecture of the models) and in severe cases, the reader's ability to contextualize the results (e.g., a list of the property signatures used and the evaluation tasks).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review!\n\n> there are almost no examples in the paper\n\nThis is a great point. Please see the global response for some example tasks and synthesized programs.\n\n> many of the implementation details are not listed out fully, which impedes reproducibility (e.g., the architecture of the models)\n\nWe provide more details below, and we will add them to the paper: \n\n* I/O encoder: it encodes the property signatures of I/O examples using a 2-layer relu-MLP with hidden size and output size of 512. \n* Value Module: it encodes the property signatures of a value using a 2-layer relu-MLP with hidden size of 512 and output (embedding) size of 256, with layer-norm applied after each linear projection. We have different MLPs for non-lambda and lambda expressions, as mentioned in Line 224 in the paper.\n* Argument Selector Module: we use an operator-specific 3-layer LSTM with hidden size 256. The prediction head is a 2-layer MLP of hidden size 512.\n\nDuring training, we use beam size 10 to generate on-policy data, where the effective batch size is 32 and a constant learning rate of 5e-4 is used with the Adam optimizer.\n\nWe will release our code and model checkpoints if accepted, to aid in reproducibility.\n\n> The biggest weakness is that the approach ... also relies on hand designed features\n\n> a list of the property signatures used\n\n> What size is the property signature?\n\nWe indeed use some hand-designed features. However, our approach actually requires less manual design than it might seem, because we devised a system of combinatorially combining property functions to greatly expand the richness of the property signatures. Appendix A (in the supplementary material) describes this in full detail.\n\nIn particular, in total across types in the DSL, we defined 20 “basic properties” of objects, 8 other objects “relevant” to understanding an object, and 24 properties for comparing two objects of the same type. These are all explicitly listed in Appendix A.\n\nWith only these building blocks, using the compositional approach in Appendix A, we encode lambda values with property signatures of length 558, and encode non-lambda values with property signatures of length 359.\n\n> How many programs are there of weight at most 12 (i.e., that were sampled from for the training set)?\n\nWhen generating training data, we first randomly generate some inputs to a synthesis problem, i.e., the input variables’ values across multiple examples (line 258). Then, the cardinality of the space of training programs varies depending on the random inputs (the number of inputs and their types), and because we exclude suboptimal solutions from the training set. That is, a program is excluded from the training set if there is a different program found earlier in our enumerative search that evaluates the same way on the input variables for all examples.\n\nThat said, we can still provide ballpark numbers. Consider the “map:replace” task from the global response. Then, starting from those inputs and ignoring the task’s outputs, our baseline enumeration finds the following programs (all with different behavior when run on the examples):\n* 313,842 programs with weight at most 8, in 5 minutes\n* 1,573,527 programs with weight at most 9, in 30 minutes\n* 8,390,593 programs with weight at most 10, in 170 minutes\n\nWhen generating training data, not all searches reach weight 12 within the 1 hour time limit (depending on the random inputs).\n\nOf course, there would be many more programs if we relax the constraint of solution optimality, and there are even more programs that are syntactically valid but fail to typecheck.\n\n> Can you elaborate on why you add an embedding of the weight of a value for embedding lambda expressions?\n\nThis is carried over from CrossBeam. Every value in CrossBeam and LambdaBeam (including lambda and non-lambda expressions) has an embedding of the value’s weight added to the embedding of the value. This helps the model understand the “cost” associated with using this value which may influence the model’s decisions, e.g., to not use values that have too high of weight. This may help the model avoid getting stuck exploring a “rabbit hole” of larger and larger expressions that ultimately do not lead to progress.\n\n> Is Merge complete for lambda functions?\n\nYes, Merge is complete, in the sense that we can use it to generate any function of the inputs in our DSL. Let $x_1 \\dots x_n$ represent the inputs to the programming-by-example (PBE) task. All solutions to the PBE task are a function $\\lambda x_1 \\dots x_n. t$, where $t$ has no unbound variables other than $x_1 \\dots x_n$. We can show that Merge generates the set of all terms $t$ in our DSL that have no unbound variables other than $x_1 \\dots x_n$.\n\nProof sketch: Use structural induction. The recursive case is where the expression has the form $t = \\lambda v_1 \\dots v_n. f(a_1 \\ldots a_K)$. For each of the $a_k$, we define a new term $b_k$ which binds all of the free variables in $a_k$. Inductively, $b_k$ can be generated using Merge. We then set variable tuples $i_k$ appropriately such that $Merge(f, b_1, i_1, b_2, i_2, \\dots)$ produces $t$. Intuitively, the $i_k$ can be seen as undoing any alpha-renaming when going from $a_k$ to $b_k$.\n\nWe will clarify this in the paper, and add a formal proof in an appendix.\n"}, {"review_id": "XSbnHgXD4X", "reviewer": "Reviewer_c5as", "summary": "In this work the authors introduce LambdaBeam a method crafted to explicitly handle lambda functions and higher-order functions for neurally guided program synthesis. Towards this goal the authors first introduce a method to represent lambda functions which enables variable order independent canonical representation, and eases creation of lambda functions by merging other lambda functions. Then, the authors adapt CrossBeam a pre-exiting method for program synthesis to synthesize lambda functions, while also employing property signatures to represent lambda functions. When deployed on integer list manipulation tasks, Lambda beam surpasses other competitive baselines, in terms of both speed and success rate.", "questions": "Why did the authors use DeepCoder benchmark? Is integer manipulation benchmark the right choice when solutions many involve a lot of hierarchical reasoning. I am especially concerned since the input contains only 3-5 examples, which might not be sufficient to triangulate the true function. Furthermore, the False Positive rate might be higher simply because of the benchmark used.\n\n\n", "limitations": "The paper does not discuss limitation or potential negative societal impact. Adding information regarding both these aspects in the appendix might further improve the paper.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "### Originality \nPrevious works do not explicitly model lambda functions, or learn how to compose them. This is an original contribution of the paper.\n\n### Quality & Clarity\nThe paper at a paragraph level is well written.\n\n### Significance\nThe approach towards representation of lambda functions is a interesting, useful and novel contribution that may be useful for many future program synthesis approaches. \nFurthermore, the approach is able to beat strong baselines such as lambda^2 an off-the-shelf program synthesis tool, and a 60B parameter large language model.", "weaknesses": "### Weaknesses\n\n1. I think the paper does not clearly justify *why* prior works cannot model lambda functions or higher-order functions. Particularly, it states Programming by examples (PBE) demands a more systematic search strategy but its unclear why that is the case. The paper also doubles down on the belief that other methods *cannot* synthesize programs with arbitrary looping computation, but its unclear why that is the case. For example, large scale models (GPT 3.5) can indeed produce programs with higher-order functions and lambda functions. Its unclear why the paper strongly posits that other methods cannot do this.\n\n2. The paper is not easy to understand and seems to depend on the reader being familiar with CrossBeam> On that note, LambdaBeam strongly depends on CrossBeam which is a small drawback as well (though the other contribution - representation of lambda functions is a general and useful contribution).\n\n ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review!\n\n> Weakness 1.\n\nWe were careful with our wording, but we will revise to make it more clear. When we discuss previous methods that cannot handle lambda functions (line 23), we are referring specifically to the types of prior works listed on lines 22-23, NOT referring to LLMs. We refer to this line of work, on deep learning for PBE, as “neural synthesis search”. Previous work in neural synthesis search (as cited on lines 22-23) indeed does not generate arbitrary looping computations. \n\nWhy not? We discuss a key difference in lines 32-26, and two difficulties in lines 41-49 and 50-56. In short: When explicitly searching over programs, the search space is very large. The best way that has been discovered in neural synthesis search is to evaluate partial programs, and use the result as input to the neural search policy. When building up a lambda function during search, the search method does not know yet what the inputs to the lambda will be, so we cannot directly evaluate the lambda. How might we represent lambdas efficiently (pruning the search space) and provide evaluation information to the neural search policy? These two difficulties are addressed by key contributions of our paper.\n\nAs for LLMs: Yes, LLMs *can* produce programs with loops, and quite impressively, but they perform synthesis from _natural language_ which is very different from PBE where we have only input-output examples. LLMs seem to depend strongly on natural language for their success. When given input-output examples alone (i.e., the PBE setting), we find in our results that LLMs perform poorly. This is why we say “PBE demands a more systematic search strategy”.\n\nTo our knowledge, LambdaBeam is the first neural synthesis search method that handles lambda functions or arbitrary looping computations. Note that $\\lambda^2$ handles lambdas but is not neural, and language models also handle loops and lambdas but do not perform any clever search (other than beam search or rejection sampling), and hence have trouble with PBE.\n\n> Weakness 2.\n\nWe will include more info about CrossBeam (see global response).\n\n> Why did the authors use DeepCoder benchmark?\n\nWe needed a benchmark that included arbitrary lambda functions, and the DeepCoder benchmark was the closest to that since it had hardcoded lambda functions. Thus, we extended DeepCoder to allow arbitrary lambda functions.\n\nFor the 100 handwritten evaluation tasks, we generally found that 3-5 examples are enough to sufficiently describe the task, and that false positive solutions are generally complicated unnatural programs that are accidentally correct on the examples. However, a higher proportion of the synthetic tasks are actually ambiguous, which can be seen from the significantly higher false positive rate on synthetic tasks compared to handwritten tasks.\n\n> The paper does not discuss limitation\n\nWe discuss limitations on lines 342 - 347.\n"}, {"review_id": "hsXK1zhZAO", "reviewer": "Reviewer_WhTE", "summary": "The paper presents LAMBDABEAM, a nn-based search method for program synthesis which is built upon CROSSBEAM and can handle lambda functions and higher-order functions. Specifically, to build lambda terms, LAMBDABEAM enforces that every term constructed during search has no free variables by introducing a novel operator called MERGE. Furthermore, to learn lambda expressions, LAMBDABEAM constructs a new generalization of property signatures to represent lambda expressions. The paper shows that LAMBDABEAM outperforms existing techniques in the integer list manipulation domain (a modified DeepCoder).", "questions": "1. Will CROSSBEAM be improved with restarts?\n2. Considering the elapsed time with quantitive computing resources used would be better if possible.\n3. How might this kind of nn-based search method be combined with LLMs?\n4. A experimental comparison with Dreamcoder is needed.", "limitations": "The authors have discussed the limitations in the results part.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This addresses a meaningful problem, that is how to search for lambdas and higher-order functions which would enable the synthesis of arbitrary looping computations and extend the boundary of neural program synthesis.\n2. The experiment result is promising.\n3. The writing is clear.\n4. In the current era dominated by LLMs in the field of program synthesis and code generation, this paper makes a good attempt towards small and meaningful works. I believe that this type of work and LLMs-related works will inspire and complement each other.", "weaknesses": "1. Placing a figure that illustrates the LAMBDABEAM Model architecture would be better. Although the design largely follows CROSSBEAM, a LAMBDABEAM figure is necessary for showing the differences and for readers unfamiliar with CROSSBEAM.\n2. Experimental settings are somewhat confusing. For example, what causes the differing number of I/O examples for list output and integer output in the hand-written?  What is the name of the pre-trained LLM since the performance gap among different LLMs on program synthesis is significant.  Can the authors further explain these settings?\n3. Can the authors fine-tune the LLM on the proposed DSL which might be a better comparison?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review!\n\n> Weakness 1.\n\nWe will include more info about CrossBeam (see global response).\n\n> what causes the differing number of I/O examples for list output and integer output in the hand-written?\n\nIn general, PBE tasks can be specified with fewer examples if the examples are more constraining. It would be easy for a program to coincidentally output the correct integer using the wrong approach, but it is less likely that an entire list matches by coincidence. For this reason, we used fewer examples (3) to specify the problem when the output is a list, and more examples (5) when the output is an integer. These numbers were chosen to avoid underspecifying the problem, while also being realistic in the number of examples a program synthesis user might want to provide.\n\n> Can the authors fine-tune the LLM on the proposed DSL\n\nIt would cost a large amount of compute resources to train the LLM (a very large sequence model) on our DSL. Instead, we have trained a smaller sequence model from scratch on our DSL, which is the RobustFill approach already included in our experiments.\n\n> Will CROSSBEAM be improved with restarts?\n\nWe hypothesize that CrossBeam will also improve with restarts for an appropriate restart frequency (a hyperparameter), but it is unclear how much improvement will be gained. The amount of improvement could vary by the domain or dataset used, as was the case in our experiments for handwritten vs synthetic tasks.\n\n(Also, we note that any such improvement to CrossBeam would not affect our experimental conclusions, since CrossBeam as a baseline would fail to solve a good majority of our evaluation problems which require using lambda expressions.)\n\n> Considering the elapsed time with quantitive computing resources used would be better if possible.\n\nWe agree it would be better, but this is very hard to do considering the different hardware used to run CPU-only approaches ($\\lambda^2$ and enumeration), CPU with GPU (LambdaBeam), mainly GPU (RobustFill), and multiple accelerators in parallel (the LLM).\n\n> How might this kind of nn-based search method be combined with LLMs?\n\nThis is an excellent question for future work! There are some very recent works that use LLMs to generate programs iteratively, e.g., to self-debug their predictions (https://arxiv.org/abs/2304.05128). In a similar vein, it would be very interesting to see whether LLMs can be made to perform program synthesis search guided by other sources of info such as program evaluations.\n\n> A experimental comparison with Dreamcoder is needed.\n\nDreamCoder is fundamentally an algorithm for enriching an impoverished DSL, and shows how that enrichment process can synergize with neurally-guided program search. Therefore, it does not make sense to compare *against* DreamCoder, but to experimentally consider *augmenting* DreamCoder with LambdaBeam (using our work as DreamCoder’s neurally-guided search strategy). While that would be a fascinating avenue to explore, we believe it is sufficiently involved to not be a reasonable piece of work to include within the scope of this paper. However, we will revise the paper to include this explanation of how DreamCoder and LambdaBeam could synergize in future systems.\n"}, {"review_id": "m0OkL1NHc7", "reviewer": "Reviewer_JwfZ", "summary": "The paper extends the CrossBeam method with synthesizing intermediate lambda\nfunctions to solve the programming by example task. The authors introduce the\nMerge operator to construct new lambda functions by choosing an operator from\nthe DSL, and its arguments from existing terms (variables or lambda\nfunctions). The inputs to the Merge operator are predicted by the neural model,\nso the lambda functions are built step-by-step bottom-up, similarly to the whole\nsynthesized program.\n\nLambda functions can not be executed on the input/ouput examples so they are\nexecuted on hardcoded canonical argument tuples instead and are encoded using\nproperty signatures computed on the results, the expected output of the program,\nand the arguments of the lambda function.\n", "questions": "There were 16 canonical argument tuples for each combination of tuple length and\nargument types. To me it seems somewhat low. Wouldn't more be helpful, or does\nthe slowdown of the search counteract the improvement?\n\nWhy do most of the methods perform worse on the handwritten tasks compared to\nthe synthetic tasks? Is it because the generation of the synthetic tasks is\nsimilar to their algorithm?\n\nFigure 4 (Synthetic tasks) shows significant improvement in the LLM results as\nthe task weight grows from 9-10 to 11-12. Is that an anomaly or is there an\nexplanation? It would be interesting to see what happens after task weight\n11-12.\n\n", "limitations": "Limitations are addressed at the end of the Results, maybe they could have their\nown section.\n", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "As the authors claim and also as far as I know this is the first neural search\nmethod to synthesize general-purpose (not hardcoded) lambda functions, which has\nbeen an open problem for years. DreamCoder (cited in the Related Work) can also\nsynthesize lambda functions but it does so by finding common program fragments\nin synthesized programs in a Bayesian framework.", "weaknesses": "The paper - understandably - refers to the CrossBeam paper many times. It\ncontains a summary of CrossBeam in lines 100-121, but reading that paper still\nhelped a lot to undertand this paper. Also, I think that Figure 2 from the\nCrossBeam paper should be included as Section 3.3 talks about parts of that\nFigure.\n\nIt would be good to include CrossBeam without lambda functions as a baseline for\nthe evaluations; currently it is hard to know how much of the improvement is due\nto the lambda functions.\n\nPart 3.1 could be clearer:\n- the example is at the end of the section, maybe an example-first approach\n  could be better\n- Merge ensures that there are no free variables and also unifies the arguments.\n  When reading through the section it seems from line 128 to 147 that the only\n  criterion we need is that we have no free variables and that alone ensures the\n  unification of the arguments.\n- I'm not sure these are correct:\n  - line 157 says that Merge runs the function $f$,\n  - line 160 says that $a_k(i_k)$ evaluates to a concrete value, I think it\n    should be an expression\n\nline 217 says that $S$ contains variable tokens, but I think it also contains\nlambda expressions.\n\nI'm not sure how to intepret line 223: \"an embedding of the weight of this value\".\n\nI couldn't find which LLM the authors used as a baseline, it should be cited.\n\nI couldn't find the range of the inputs and outputs, it would be good to include\nthem.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review!\n\n> include CrossBeam without lambda functions as a baseline\n\nCrossBeam would not perform well because 85/100 handwritten evaluation tasks and 53/100 synthetic evaluation tasks use a lambda function in the solution. CrossBeam would not be able to solve those problems and would end up being the worst in our comparison.\n\n> Part 3.1 could be clearer\n\nWe used “unify” loosely here. We mean that the arguments $a_k$ are alpha-renamed using $i_k$ so that we can choose which variables are used in which locations. For example, if we have a term $T = \\lambda v_1. (v_1 + 1)$, then this alpha-renaming step is needed to distinguish between $Merge(\\times, T, [v_1], T, [v_2]) = \\lambda v_1, v_2. (v_1 + 1) \\times (v_2 + 1)$ and $Merge(\\times, T, [v_1], T, [v_1]) = \\lambda v_1. (v_1 + 1) \\times (v_1 + 1)$. We will replace our usage of “unify” with “alpha-rename”.\n\nLine 157: Instead of “Merge runs $f$ on the arguments”, we should say “Merge creates an expression that calls $f$ on the arguments”.\n\nLine 160: The argument $a_k(i_k)$ is indeed an expression. We mean to point out that this expression is not immediately usable as a lambda in a higher-order function and needs to be wrapped with an explicit lambda. For example, suppose we have the term $T = \\lambda v_1. v_1 + 1$. We’d use this in a first-order function as follows: $Merge(\\times, T, [v_1], x, []) = \\lambda v_1. (v_1 + 1) \\times x$. If we use a higher-order function in the exact same way, we get $Merge(map, T, [v_1], x, []) = \\lambda v_1. map(v_1 + 1, x)$ which is not valid since the first argument to map must be a function. We instead need $Merge(map, T, [u_1], x, []) = map(\\lambda u_1. T(u_1), x) = map(\\lambda u_1. u_1 + 1, x)$, where the $\\lambda u_1.$ part is explicitly added by our definition of Merge because we know that $map$ requires its first argument to be a function with arity 1.\n\nLine 217: This sentence has ambiguous parsing which we will revise. $S$ contains variable tokens for constructing lambdas, and $S$ also contains lambda expressions and non-lambda expressions.\n\nLine 223: The weight of the value is an integer, and we look up the integer in a learned embedding table to obtain $z$. This is done in the same way as in CrossBeam (Section 3, “Value module” paragraph, in the CrossBeam paper), for both lambda and non-lambda values.\n\n> I couldn't find the range of the inputs and outputs\n\nFor inputs and outputs in our handwritten and synthetic tasks, all integers are in the range [-256, 255] as in DeepCoder, and lists have lengths in the range [0, 10] which we felt was reasonable for PBE users to specify.\n\n> There were 16 canonical argument tuples ... it seems somewhat low\n\nThrough profiling, we found that the majority of LambdaBeam’s time is spent computing property signatures (not running the model!), especially for lambda functions since they must be run many times. For this reason, we didn’t include too many canonical argument tuples.\n\n> Why do most of the methods perform worse on the handwritten tasks compared to the synthetic tasks?\n\nThere are two potential reasons. One is actually just visual: in Figure 4, the handwritten tasks have weight buckets shifted leftward compared to the synthetic tasks, because we have handwritten tasks of weight 13-19 but all synthetic tasks have weight at most 12. For tasks of weight <= 10, LambdaBeam+Restarts performs equally well on handwritten and synthetic tasks, and $\\lambda^2$ even performs better on the handwritten tasks.\n\nThe other reason is that the tasks are distributed differently in the handwritten vs synthetic datasets. For example:\n\n```\nThe handwritten evaluation tasks include:\n  * 20 tasks of weight  7 -  8: 17 with lambdas, 3 without\n  * 25 tasks of weight  9 - 10: 25 with lambdas, 0 without\n  * 20 tasks of weight 11 - 12: 17 with lambdas, 3 without\n\nThe synthetic evaluation tasks include:\n  * 20 tasks of weight  7 -  8: 13 with lambdas,  7 without\n  * 20 tasks of weight  9 - 10: 16 with lambdas,  4 without\n  * 20 tasks of weight 11 - 12:  9 with lambdas, 11 without\n```\n\nThe handwritten tasks include more tasks that use lambdas, and there are abnormally many synthetic tasks of weight 11-12 without lambdas which appear to be easier overall. There are likely other distribution shifts such as in the shape of the solution (deep vs wide expression trees), the operations used, or the I/O types. It is thus an impressive result that LambdaBeam+Restarts performs the best on both evaluation datasets.\n\n> Figure 4 (Synthetic tasks) shows significant improvement in the LLM results as the task weight grows from 9-10 to 11-12\n\nThis is an anomaly. In fact, for synthetic tasks of weight >= 8, all of the LLM’s “solutions” are false positives with some form of “if the input is <hardcoded> then return <hardcoded>” logic. Indeed, Figure 5 shows the LLM has a very high false positive rate on synthetic tasks. This solution pattern is easier to implement when the output is an integer as opposed to a list, and there are abnormally many synthetic tasks of weight 11-12 with integer outputs.\n\n```\nThe handwritten evaluation tasks include:\n  * 20 tasks of weight  7 -  8: 11 have int output,  9 have list output\n  * 25 tasks of weight  9 - 10:  4 have int output, 21 have list output\n  * 20 tasks of weight 11 - 12:  8 have int output, 12 have list output\n\nThe synthetic evaluation tasks include:\n  * 20 tasks of weight  7 -  8:  7 have int output, 13 have list output\n  * 20 tasks of weight  9 - 10:  6 have int output, 14 have list output\n  * 20 tasks of weight 11 - 12: 14 have int output,  6 have list output\n```\n\n"}]

Author rebuttal:
We appreciate all of the insightful reviews! We will revise our paper to incorporate our clarifications and new information wherever appropriate.

This global response includes information helpful for multiple reviewers, and we also respond to each reviewer individually.

**The paper could use more background on CrossBeam**
(reviewers c5as, WhTE, JwfZ)

If accepted, we will use the extra page to include more background on CrossBeam, and LambdaBeam’s relationship to CrossBeam (e.g., with a figure), to help alleviate this weakness.

**Which pre-trained LLM is used?**
(reviewers WhTE, JwfZ)

We omitted the name and citation for double-blind purposes. We will certainly include more details and citation if accepted.

**Minor correction**

Line 260 says each synthetic training task has between 3 and 5 I/O examples. Actually, synthetic tasks have between **2** and 5 I/O examples.

**The paper could use examples of evaluation tasks and synthesized programs**
(reviewer ErYC, but we think this context will be helpful to everyone)

We will add some evaluation tasks and synthesized programs to an appendix. We cannot upload new revisions during the review period, so we provide some interesting examples here. The task names are only for convenience and are not used in our comparisons.

**“map:replace”**
This handwritten 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.

LambdaBeam+Restarts finds the same solution (weight 10) in each of the 5 trials: `Map(lambda u1: (lambda v1: If((lambda v1: Equal(f, v1))(v1), r, v1))(u1), x)`, taking a median time of 202 seconds.

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.

**“multi:multiply_odds”**
This handwritten 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 (weight 11) that simplifies to the ground-truth solution, taking a median time of 75 seconds.

**“synthetic:weight_9_function_7”**
This synthetic 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)))` (this solution is 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 solutions could benefit from a postprocessing simplification step (line 339) but this is orthogonal to our contributions.


Meta-review: This paper presents a method for doing synthesis in the integer list manipulation domain (DeepCoder) with higher-order functions.  One new technique needed to do this deals with handling equivalent lambda representations for efficiency of pruning, which relies on a new MERGE operation. The paper also uses property signatures to analyze how lambdas are used and represent them appropriately in the neural model.

The reviewers found the paper to be tackling an interesting and novel topic, of how to model lambdas in program synthesis. They praised the execution of the method and the experimental results. The biggest weaknesses of the paper that were raised are (1) its dependence on CrossBeam, and its nature as an extension of that method, which primarily impacts the writing and the ability of this paper to stand alone in its current form. (2) The results on a single synthetic dataset, which makes the technique feel a bit unproven in broader or more realistic settings.
Decision: Accept (poster)