SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Then, we assume the inductive hypothesis that any term in  $S$ , with depth less than that of  $\mathbf{s}$ , is in  $M$ . There are two inductive cases to consider: either  $\mathbf{s} = \lambda v_1 \dots v_n$ .  $f(\mathbf{t}_1, \ldots, \mathbf{t}_k)$  where  $n$  might be 0, or  $\mathbf{s} = \lambda v_1 \dots v_n$ .  $\mathbf{t}$  where  $n > 0$  and  $\mathbf{t}$  is atomic. Because  $\mathbf{s}$  has exact lambda variables, the latter scenario is only possible if  $\mathbf{s} = \lambda v$ .  $v$ , which is immediately in  $M$ .
The remaining case is  $\mathbf{s} = \lambda v_{1}\dots v_{n}$ .  $f(\mathbf{t}_1,\ldots ,\mathbf{t}_k)$ . Consider any  $\mathbf{t}_j$  for  $1\leq j\leq k$ . We will construct  $\mathbf{a}_j$  and  $i_j$  such that  $\mathrm{MERGE}(f,\mathbf{a}_1,i_1,\ldots ,\mathbf{a}_k,i_k) = \mathbf{s}$ .
- If  $\mathbf{t}_j$  is atomic, then either  $\mathbf{t}_j = x$ ,  $\mathbf{t}_j = v$ , or  $\mathbf{t}_j = c$ . If  $\mathbf{t}_j = v$ , then set  $\mathbf{a}_j = \lambda v$ .  $v$  and  $i_j = [v]$ ; otherwise, set  $\mathbf{a}_j = \mathbf{t}_j$  and  $i_j = [\cdot]$ , the empty tuple. In each case,  $\mathbf{a}_j \in M$ . Because s typechecks, we know that the  $j$ -th argument to  $f$  expects arity 0, so in the MERGE definition,  $\ell_j = 0$  and thus the  $j$ -th argument to  $f$  expands to  $\mathbf{a}_j(i_j) = \mathbf{t}_j$  in each case. $^7$
- If  $\mathbf{t}_j$  is not atomic, then let  $\mathbf{t}_j = \lambda u_1 \dots u_{\ell_j}$ .  $\mathbf{r}$ , where  $\ell_j$  is the expected arity of the  $j$ -th argument to  $f$  (because  $\mathbf{s}$  typechecks). Let  $\{v_1', \ldots, v_d'\} = \mathrm{FreeVars}(\mathbf{r}) - \mathcal{X}$ . Set  $\mathbf{a}_j = \lambda v_1' \dots v_d'$ .  $\mathbf{r}$  and  $i_j = [v_1', \ldots, v_d']$ , so that when applying MERGE, the  $j$ -th argument to  $f$  expands to  $\lambda u_1 \dots u_{\ell_j}$ .  $\mathbf{a}_j(i_j) = \lambda u_1 \dots u_{\ell_j}$ .  $\mathbf{r} = \mathbf{t}_j$ . Furthermore,  $\mathbf{a}_j$  has exact lambda variables by construction, and it typechecks because  $\mathbf{r}$  typechecks, so  $\mathbf{a}_j \in S$  and  $\mathbf{a}_j \in M$  by the inductive hypothesis.
With these choices of  $\mathbf{a}_j$  and  $i_j$ , when expanding  $\mathrm{MERGE}(f, \mathbf{a}_1, i_1, \ldots, \mathbf{a}_k, i_k)$  according to the definition, the  $j$ -th argument to  $f$  becomes  $\mathbf{t}_j$ , and the lambda variables  $\bigcup_{k} i_k - \{u_j\}_{j=1}^{\max \{\ell_1, \ell_2, \ldots\}}$  are exactly FreeVars  $(f(\mathbf{t}_1, \ldots, \mathbf{t}_k)) = \{v_1, \ldots, v_n\}$  since  $\mathbf{s}$  has exact lambda variables. Therefore,  $\mathbf{s} = \lambda v_1 \dots v_n$ .  $f(\mathbf{t}_1, \ldots, \mathbf{t}_k) = \mathrm{MERGE}(f, \mathbf{a}_1, i_1, \ldots, \mathbf{a}_k, i_k)$ , so  $\mathbf{s} \in 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, TypeProperties  $(x)$  represents the type of  $x$  as a boolean one-hot list. In our DSL, this returns a tuple of 5 booleans, representing whether  $x$  is a lambda, boolean, int, list, or None (which is used to indicate an error, e.g., signaling that an index is out of bounds).
Next, we define BasicProperties  $(x)$  to evaluate hand-designed "basic" properties that describe objects of each different type in the DSL. This returns a fixed-length vector of property results (each being True, False, or N/A). Note that, if  $x$  has type  $\tau$ , then all properties for type  $\tau$  evaluate to True or False, while all properties for all other types  $\tau' \neq \tau$  evaluate to N/A. For our DSL, we use the following basic properties:
- For boolean  $x$ :  $x$  itself.
- For integer  $x$ : whether  $x$  equals  $-1, 0, 1$ , and 2; whether  $x$  is positive and negative; whether  $x$  is even; whether  $x$  is 0 and 1 modulo 3; and whether  $|x|$  is less than 5, 10, 20, 35, 50, 75, and 100.
- For list  $x$ : whether  $x$  is sorted, whether  $x$  is sorted in reverse, and whether  $x$  contains all unique elements.
Then,  $\operatorname{Relevant}(x)$  returns related objects that are relevant to understanding  $x$ . For our DSL, this is only  $x$  itself for integer and boolean  $x$ , but for list  $x$ , the "relevant" objects are:  $x$  itself; the length of  $x$ ; the number of distinct elements in  $x$ ; the max, min, range, and sum of  $x$ ; and the first and last elements of  $x$  (defaulting to 0 if  $x$  is empty).
This culminates in ObjectSignature  $(x)$  which takes a single DSL object  $x$  and returns a fixed-length vector of property results, containing TypeProperties  $(x)$  followed by BasicProperties  $(r)$  for each  $r\in \mathrm{Relevant}(x)$ . For example, these properties include "assuming  $x$  is an int, is  $x$  is even?" (a basic
In practice for simplicity, when  $\mathbf{t}_j = v$ , we simply set  $\mathbf{a}_j = \mathbf{t}_j$  and  $i_j = []$  as in the other cases, even though  $\mathbf{a}_j = v$  does not have exact lambda variables.
property applied to  $x$ ) as well as "assuming  $x$  is a list, are there an even number of elements in  $x$ ?" (a basic property applied to an object relevant to  $x$ ). By applying basic properties to relevant objects in this compositional way, we reduce the effort needed to specify a large number of properties.
We furthermore encode comparisons between two objects. We define ComparisonProperties  $(x, y)$  which evaluates hand-designed properties for comparing two objects  $x$  and  $y$  of the same type, for each different type in the DSL. This returns a fixed-length vector of property results, where a property for comparing type  $\tau$  evaluates to N/A if  $x$  and  $y$  are not of type  $\tau$ . For our DSL, we use the following comparison properties:
- For boolean  $x$  and  $y$ : whether  $x = y$ .
- For integer  $x$  and  $y$ : whether  $x = y$ ,  $x < y$ , and  $x > y$ ; whether  $x$  is a factor of  $y$  and vice versa; and whether  $|x - y|$  is less than 2, 5, 10, and 20.
- For list  $x$  and  $y$ : whether  $x = y$ ; whether  $x$  is longer, shorter, or equal length compared to  $y$ ; whether the lengths differ by at most 1; whether all  $x_i < y_i$  for  $x_i, y_i \in \operatorname{zip}(x, y)$  and similarly for other comparisons  $\leq, >, \geq, =$ , and  $\neq$ ; whether  $x$  and  $y$  contain the same set of elements; and whether  $x$  contains a subset of elements compared to  $y$  and vice versa.
These properties are used in ComparisonSignature  $(x, y)$  which computes a fixed-length list of property results for any two DSL objects  $x$  and  $y$  of any type, containing ComparisonProperties  $(r_x, y)$  for all  $r_x \in \mathrm{Relevant}(x)$  where type  $(r_x) = \mathrm{type}(y)$ , and ComparisonProperties  $(x, r_y)$  for all  $r_y \in \mathrm{Relevant}(y)$  where type  $(r_y) = \mathrm{type}(x)$ . Thus, "assuming  $x$  is an int and  $y$  is a list, is  $x$  a factor of the length of  $y$ ?" is one resulting property. As usual, if  $x$  and  $y$  do not match the types assumed by the property, then the property evaluates to N/A.
In the I/O Module of the neural policy (as in CROSSBEAM [33]), we use property signatures to encode a set of I/O examples. For each example  $(\{I_1,\dots,I_n\},O)$  we concatenate ObjectSignature  $(O)$  with ObjectSignature  $(I_i)$  and ComparisonSignature  $(I_i,O)$  for all  $1\leq i\leq n$ . Then, we reduce these across I/O examples, computing for each property the fraction of examples where it is applicable (not N/A), and the fraction of examples where it is True among those where it is applicable (defaulting to 0.5 if it is N/A for all examples).
In the Value Module of the neural policy, we use property signatures to encode a value (lambda or non-lambda expression) that was found during search. To encode a lambda expression, we run it on canonical input tuples as described in Section 3.2. For each run of the lambda on canonical input tuple  $t_i = (t_{i,1},\dots ,t_{i,m})$  using an I/O example  $(I,O)$  where the lambda evaluates to a result  $r_i$ , we concatenate ObjectSignature  $(r_i)$ , ComparisonSignature  $(r_i,O)$ , and ComparisonSignature  $(t_{i,j},r_i)$  for all  $1\leq j\leq m$ , and then reduce these across the runs of the lambda. To encode a non-lambda expression during search, for each I/O example  $(I,O)$  where the expression evaluates to a result  $r$ , we concatenate ObjectSignature  $(r)$  with ComparisonSignature  $(r,O)$ , and then reduce these across I/O examples. Note that the signatures for values found during search do not contain comparisons to the I/O example inputs, because what ultimately matters is whether the value is useful for creating the output later, not how the value was created from the inputs.
In our implementation, encoding the set of I/O examples results in a property signature of length 1230, encoding a lambda expression results in a property signature of length 558, and encoding a non-lambda expression results in a property signature of length 359.
C Extension of the DeepCoder DSL
As mentioned in Section 4.1, we extended the DSL from DeepCoder [3]. Atomic terms in the DSL include variable names and the constant literals  $-1, 0, 1, 2, 3,$  and 4. The DSL contains 23 first-order and 5 higher-order operations, listed below with type annotations and Python implementations:
23 first-order operations
def Add(x: int, y: int) -> int:
    return x + y
def Subtract(x: int, y: int) -> int:
    return x - y
def Multiply(x: int, y: int) -> int:
    return x * y
def IntDivide(x: int, y: int) -> int: return x // y
def Square(x: int) -> int:
    return x ** 2
def Min(x: int, y: int) -> int: return min(x, y)
def Max(x: int, y: int) -> int:
    return max(x, y)
def Greater(x: int, y: int) -> bool: return x > y
def Less(x: int, y: int) -> bool:
    return x < y
def Equal(x: int, y: int) -> bool:
    return x == y
def IsEven(x: int) -> bool:
    return x % 2 == 0
def IsOdd(x: int) -> bool:
    return x % 2 != 0
def If(c: bool, x: int, y: int) -> int:
    return x if c else y
def Head(xs: list) -> int: return xs[0]
def Last(xs: list) -> int: return xs[-1]
def Take(n: int, xs: list) -> list:
    return xs[:n]
def Drop(n: int, xs: list) -> list:
    return xs[n:]
def Access(n: int, xs: list) -> int: return xs[n]
def Minimum(xs: list) -> int: return min(xs)
def Maximum(xs: list) -> int: return max(xs)
def Reverse(xs: list) -> list:
    return list(reversed(xs))
def Sort(xs: list) -> list: return sorted(xs)
def Sum(xs: list) -> int:
    return sum(xs)
5 higher-order operations
def Map(f: Callable[int], int], xs: list) -> list: return [f(x) for x in xs]
def Filter(f: Callable[[int], bool], xs: list) -> list: return [x for x in xs if f(x)]
def Count(f: Callable[[int], bool], xs: list) -> int: return len([x for x in xs if f(x)])
def ZipWith(f: Callable[[int, int], int], xs: list, ys: list) -> list: return [f(x, y) for x, y in zip(xs, ys)]
def Scan11(f: Callable[[int, int], int], xs: list) -> list:
    ys = [xs[0]]
    for n in range(1, len(xs)):
        ys.append(f(ys[n-1], xs[n]))
    return ys
D Example Tasks
This section contains selected example problems from our 100 handwritten and 100 synthetic evaluation tasks. Each task is given a name for convenience purposes only, which is not used by any method in our experiments.
D.1 Handwritten Task "map:replace"
This task has 3 inputs  $(\mathbf{x},\mathbf{f},$  and  $\mathbf{r})$  , 3 examples demonstrating the task ("in  $\mathbf{x}$  , find instances of f and replace them with r"), and a handwritten ground-truth solution using a relatively complicated lambda function:
inputs_dict = {  
'x': [[7, 2, 4, 6, 4, 2, 5],  
[-6, -3, 4, 3, -5, -3, 2, 1, 5],  
[18, 48, 27, 26, 27, 27, 28, 17, 27, 33]],  
'f': [4, -3, 27],  
'r': [-1, 7, 99],  
}  
outputs = [[7, 2, -1, 6, -1, 2, 5],  
[-6, 7, 4, 3, -5, 7, 2, 1, 5],  
[18, 48, 99, 26, 99, 99, 28, 17, 99, 33]]  
solution = 'Map( lambda u1: If(Equal(u1, f), r, u1), x)'
In inputs_dict, each of the entries for x, f, and r is a list of length 3, which contains the input for each of the 3 examples. Similarly, outputs is a list containing the output for each example. solution is our handwritten solution.
LAMBDABEAM + Restarts finds the same solution of weight 10 in each of the 5 trials, taking a median time of 202 seconds:
Map( lambda u1: (lambda v1: If((lambda v1: Equal(f, v1))(v1), r, v1)) (u1), x)
The solution looks complicated due to the MERGE operation causing lots of variable renames (i.e.,  $a_{k}(i_{k})$  in the MERGE definition). We have implemented an algorithm to simplify the solution by statically resolving these renames. In this case, the solution simplifies to
Map (lambda u1: If (Equal(f, u1), r, u1), x)
which is essentially identical to the ground-truth solution.
D.2 Handwritten Task "multi:multiply(odds)"
This task has 1 input and uses multiple higher-order functions to compute a running product of only the odd elements:
inputs_dict = {  
'x': [[3, 5, 8, 2, 1], [5, 2, 1, 3, 3, 1, 4], [3, -4, -1, 8, 2, 0, -3, 0, 9, -1]],  
}  
outputs = [[3, 15, 15], [5, 5, 15, 45, 45], [3, -3, 9, 81, -81]]  
solution = 'Scan11( lambda u1, u2: Multiply(u1, u2), Filter( lambda u1: IsOdd(u1), x))'
In each of the 5 trials, LAMBDABEAM + Restarts finds the same solution of weight 11 that simplifies to the ground-truth solution, taking a median time of 75 seconds.
D.3 Synthetic Task "synthetic:weight_9_function_7"
This task clips every element to the range  $[0,4]$ :
inputs_dict = {  
'x1': [[-9, -2, -10, -6, 0, -10, -6, 3, 1], [-1, -5, 8, 5]]  
}  
outputs = [[0, 0, 0, 0, 0, 0, 0, 3, 1], [0, 0, 4, 4]]  
solution = 'Map( lambda u1: Min(4, Max(0, u1)), x1)'
LAMBDABEAM + Restarts finds a correct solution in all 5 trials with a median time of 38 seconds, but the solutions are slightly different (the simplified solutions are listed):
ZipWith( lambda u1, u2: Min(4, Max(0, u2)), x1, x1)  
ZipWith( lambda u1, u2: Min(4, Max(0, u1)), x1, x1)  
Reverse(ZipWith( lambda u1, u2: Min(4, Max(0, u2)), x1, Reverse(x1)))  
Reverse(Map( lambda u1: Min(4, Max(0, u1), Reverse(x1)))) # found in two trials
Note that these are not the shortest solutions, but nevertheless all of these solutions are equivalent to the ground-truth solution. LAMBDABEAM's solutions could benefit from a postprocessing simplification step, as discussed in Section 4.3.
E More Details on LAMBDABEAM Architecture and Training
In our experiments, we used the following hyperparameters for the LAMBDABEAM model architecture and training procedure. Refer to Figure 2 for a diagram showing how the different modules interact.
- I/O Module: this encodes a property signature of the I/O examples using a 2-layer ReLU-MLP with hidden size and output size of 512.
- Value Module: this encodes each value's property signature using a 2-layer ReLU-MLP with hidden size of 512 and output (embedding) size of 256, with a layer-norm applied after each linear projection. We use different MLPs for lambda and non-lambda expressions.
- Search Context Summary Module: this module needs to represent the entire search state at the current stage, including the current operator to be expanded, the I/O specification, and the values explored so far. We compute the average of the set of value embeddings, concatenate it with the I/O embedding, and then apply a projection layer (denoted as  $MLP_{op}$  in Figure 2, which projects back to the embedding dimension) to get a vector representation. The model parameters used in the projection layers are indexed by the operator (i.e., we use different sets of trainable parameters for different operators).
- Argument Selector Module: we use an operator-specific 3-layer LSTM with hidden size of 256. The prediction head is a 2-layer MLP with hidden size of 512.
- During training, we generate on-policy data with beam size 10, use an effective batch size of 32, and use the Adam optimizer with a constant learning rate of  $5 \times 10^{-4}$ .
- During evaluation, we use UniqueRandomizer with beam size 10.
F Analysis of Handwritten and Synthetic Tasks
Table 1 shows some differences in the distributions between our handwritten and synthetic evaluation tasks. This analysis may help contextualize the experimental results in Section 4.
For example, Figure 5 shows that the LLM solved abnormally many synthetic tasks in the 11-12 weight bucket. In fact, for synthetic tasks of weight 8 or more, every one of the LLM's "solutions" are actually false positives using some form of "if the input is <hardcoded> then return <hardcoded>" logic, which is easier to implement when the output is an integer as opposed to a list. Table 1 shows that there are abnormally many synthetic tasks of weight 11-12 that have integer outputs, which helps to explain the results.
Table 1: Analysis of some differences in task distributions between the handwritten and synthetic evaluation tasks. We show the number of tasks with each weight, the number of tasks using a lambda expression in the ground truth solution, and the number of tasks where the desired output is an integer (as opposed to a list). There may also be other axes where the handwritten and synthetic tasks have different distributions.
Weight	100 Handwritten Tasks	100 Synthetic Tasks
# Tasks	With Lambda	Int Output	# Tasks	With Lambda	Int Output
3	0	-	-	10	3	30%	2 20%
4	3	2 67%	2	67%	10	2	20%	6 60%
5	7	0	0%	43%	10	3	30%	5 50%
6	10	9	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	9	100%	1 11%	10	8	80%	3 30%
11	11	8	73%	5 45%	10	6	60%	7 70%
12	9	9	100%	3 33%	10	3	30%	7 70%
13	5	5	100%	2 40%	0	-	-
14	3	3	100%	1 33%	0	-	-
15	4	4	100%	1 25%	0	-	-
16	2	2	100%	0 0%	0	-	-
17	0	-	-	0	-	-
18	0	-	-	0	-	-
19	1	1	100%	0 0%	0	-	-
Total	100	85	85%	36 36%	100	53	53%	43 43%