Towards Neural Functional Program Evaluation

Neural models originally developed for natural language processing show promising performance for modeling computer programming languages (Chen et al., 2021; Austin et al., 2021). This has led to a number of interesting applications, and deep-learning models are now successfully and routinely applied in tools that assist developers in writing and understanding programs and code. For instance, neural language models can synthesize (Gulwani et al., 2017; Ellis et al., 2020), complete (Chen et al., 2021), and summarize programs (Elnaggar et al., 2021), whether they are written in mainstream languages like Python and Java, or in domain-specific languages like SQL and regex.

A less explored application of neural models is program evaluation (Reed and De Freitas, 2015; Zaremba and Sutskever, 2014). Here, the challenge is to predict the output of a program given its input. In software engineering, this problem is solved by program interpreters (Nystrom, 2021), which are implemented as rule-based, non-differentiable systems that take formal program code as input and produce outputs, effects, and/or errors. While these interpreters must reject programs that are syntactically incorrect or contain semantic errors, neural interpreters allow the evaluation of incomplete, informally-, or even incorrectly-specified programs. Developers could use them to predict output or effects of programs before they are fully written, which could aid in their very creation. Previous studies have focused mainly on simple imperative programming languages, e.g., FORTH (Bošnjak et al., 2017) and subsets of Python (Bieber et al., 2020). For these languages, evaluation is complicated by control flow, state, and mutability. State-of-the-art techniques solve these complications with specialized, complex model architectures, but results are not as promising as one might hope (Bieber et al. (2020) report an accuracy of $62.1\%$ on their task and dataset).

This paper studies neural program evaluation for functional programming languages, which share few of the complications of imperative languages (Feser et al., 2016). Our emphasis is on lambda calculi, which sit at the core of modern functional programming languages like Scheme and Haskell (Pierce and Benjamin, 2002). We show that program evaluation can be recast such that it is tractable for a language model based on the standard transformer architecture (Vaswani et al., 2017). Rather than predicting a program’s output given an input, we train the model to reduce a program to a form that cannot be reduced further. This generalizes the notion of program evaluation, as it can also be used to model partial application of functions and programs.

Unlike works that use neural programs to strengthen generalization of models (Chen et al., 2020; Nye et al., 2020) on another downstream task, we study the generalization of neural program evaluation explicitly. We propose three data splits to measure “type generalization”, “function compositional generalization”, and “reduction steps generalization” of our approach. Furthermore, we investigate the effect of syntax and abstraction on neural program evaluation. We define two lambda calculi that are equally capable of expressing semantically equivalent programs, yet are different in their syntax. We find that the neural program evaluation problem is slightly less tractable for the lambda calculus with the simpler syntax. Lastly, we compare two evaluation strategies for lambda calculus, lazy evaluation and eager evaluation, and find that they are similarly tractable.

Our experiments are based on synthetic data for two lambda calculi. The first lambda calculus (\nth1 LC) is untyped and has three syntactic constructs: variable, lambda, and application expressions, i.e.,

$\displaystyle e$

$\displaystyle=\mathrm{x}\,|\,{}\lambda\mathrm{x}.e\,|\,{}e\,e$

(1)

The second lambda calculus (\nth2 LC) adds types and additional syntactic constructs:

	$\displaystyle t$	$\displaystyle=\mathrm{Unit}{}\,|\,{}\mathrm{Bool}{}\,|\,{}\mathrm{List}\,t\,|\,{}t\to t$		(2)
	$\displaystyle e$	$\displaystyle=\mathrm{x}\,|\,{}\lambda\mathrm{x}.e\,|\,{}e\,e\,|\,{}(){}\,|\,{}\mathrm{True}{}\,|\,{}\mathrm{False}{}\,|\,{}\mathrm{if}\,e\,\mathrm{then}\,e\,\mathrm{else}\,e\,|\,{}\mathrm{Nil}{}\,|\,{}\mathrm{Cons}\,e\,e\,|\,{}\mathrm{Foldr}\,e\,e\,e$		(3)

Terms are generated exclusively from the \nth2 LC using types to guide the generation. First, the program type is chosen at random from the set of types $t$. Terms $e$ are then generated according to the type, beginning at the root of the term tree and proceeding down the branches to the leaves. Specifically, production is recursive and top to bottom (root to leaves). Generation starts with the outermost type constructor, and proceeds by sampling from all compatible term constructors. Production stops when none of the resulting leaf nodes contains hole terms, $e$. Care is taken to ensure that the generated program does not contain any free variables, is well-typed, and that it terminates without error. Terms from the \nth1 LC are obtained by transforming terms from the \nth2 LC. We therefore only consider the normalisable subset of the \nth1 LC. We use Church encoding to represent syntactic constructs from the \nth2 LC that do not appear in the \nth1 LC (Pierce and Benjamin, 2002). For example, lists are represented by their right fold:

$\displaystyle\mathrm{Nil}{}$

$\displaystyle=\lambda\mathrm{c}.\lambda\mathrm{n}.\mathrm{n}$

$\displaystyle\mathrm{Cons}\,h\,\tau$

$\displaystyle=\lambda\mathrm{c}.\lambda\mathrm{n}.\mathrm{c}\,h\,(\tau\,\mathrm{c}\,\mathrm{n})$

$\displaystyle\mathrm{Foldr}\,f\,e\,l$

$\displaystyle=l\,f\,e$

(4)

where $h$ and $\tau$ are head and tail terms, and where $f$, $e$, and $l$ refer to the combining function, the initial value, and the Church-encoded input list of elements, respectively.

We consider two reduction strategies: lazy and eager evaluation. Lazy evaluation reduces a term to weak head normal form (WHNF), which is a term that is not an application expression. For example, the term $(\lambda\mathrm{x}.\mathrm{x})(\lambda\mathrm{y}.\mathrm{y})$ is reduced to the WHNF $\lambda\mathrm{y}.\mathrm{y}$, while the term $\lambda\mathrm{y}.(\lambda\mathrm{x}.\mathrm{x})\,\mathrm{y}$ is already in WHNF and thus not reduced further. Eager evaluation not only reduces a term to WHNF, but also reduces all subterms of the term. We call this the “deep” reduction and refer to the result as the deep normal form (DNF). The terms $(\lambda\mathrm{x}.\mathrm{x})(\lambda\mathrm{y}.\mathrm{y})$ and $\lambda\mathrm{y}.(\lambda\mathrm{x}.\mathrm{x})\,\mathrm{y}$ are reduced to the same DNF $\lambda\mathrm{y}.\mathrm{y}$.

We define the number of reduction steps of a term to be the number of reductions performed on its root term and all its subterms. Reduction to WHNF is a special case of reduction to DNF, and the number of reduction steps to WHNF is always smaller or equal to the number of reduction steps to DNF. We use the number of reduction steps to denote the degree of evaluation complexity of a term, with more steps translating to higher complexity. This can then be used to evaluate the performance of our neural program evaluation models.

In our experiments, we train a language model based on the encoder-decoder transformer architecture to reduce programs to their normal forms. Depending on the experiment, the input to the network is a term in either the \nth1 LC or \nth2 LC, and the output is either the WHNF or the DNF of the input term.

Terms are encoded as strings, which are then converted to sequences of tokens with ids that are used as inputs to the network. The string representation of a term is obtained by pretty-printing the term in Haskell syntax (Jones, 2003). Variables are encoded as x0, x1, x2, etc. and lambda expressions are encoded as \x0 -> e. Application expressions are encoded as e1 e2. Where necessary, parentheses are used to denote the precedence of application. Thus, the \nth1 LC term $(\lambda\mathrm{x}.\mathrm{x})(\lambda\mathrm{y}.\mathrm{y})$ is encoded as (\x0 -> x0) (\x1 -> x1). In the \nth2 LC, $\mathrm{Nil}$ becomes [], and $\mathrm{Cons}\,e1\,e2$ becomes e1 : e2, while lists of one or more elements are pretty-printed using Haskell’s built-in list syntax: [e1, e2, ...]. We deviate from the usual Haskell syntax for if-then-else expressions: $\mathrm{if}\,e1\,\mathrm{then}\,e2\,\mathrm{else}\,e3$ is encoded as ite e1 e2 e3. The type of a \nth2 LC term is not encoded explicitly. Hence, both lambda calculi appear to be untyped to the models except that non-normalisable terms do not appear in the data. Our type-driven generation process avoids generating non-normalising terms. We distinguish evaluation with variable renaming (VR) and evaluation without variable renaming (NVR). In the former case, the variables in the reduced program are freshly generated in the order of their appearance, while in the latter case, the variable names are preserved during reduction, which we expect to be more tractable.

We use a pretrained T5 model (Raffel et al., 2020) as our base model. The number of parameters of the T5-Small and T5-Large models can be found in the appendix section A. We fine-tune the model using the Adafactor optimizer (Shazeer and Stern, 2018) with a maximum learning rate of $10^{-4}$ in a linear decay schedule and a batch size of $2048$. Predictions are made using greedy decoding.

A dataset of 1 million unique examples is generated by sampling from the distribution of terms in the \nth2 LC, which are subsequently converted to the \nth1 LC. We thus use the same dataset in experiments with the \nth1 LC and the \nth2 LC. We chose to generate from \nth2 LC so that we could exploit its types for type-driven generation of terms. Terms that are too long to fit within the model’s maximum input ($512$ tokens) and output lengths ($256$ tokens) are discarded.

Performance is evaluated using the average exact string match metric between the normal forms of the predicted terms and the normal forms of the ground-truth terms. For a given example, if all predicted terms match 100% with ground-truth terms, an exact match is found. The average is taken by dividing the number of exact matches over the total number of examples.

We compare neural execution of programs in two lambda calculi, \nth1 LC and \nth2 LC. Both express semantically equivalent programs, but \nth1 LC uses a simpler syntax at the cost of longer program length. Our experiments with T5 have shown near tractable performance for synthetic in-distribution \nth1 LC and \nth2 LC data, but stark performance differences are observed for out-of-distribution data. We analyze our models’ generalization capability on four splits of the data: uniformly random, by program type, by function composition, and by number of reduction steps. On average, models trained on \nth1 LC generalize less then models trained on \nth2 LC. We observe significant gains from using T5 with pretrained weights compared to randomly initialized models on in-distribution and unseen-type experiments. Input/output program lengths cannot completely account for differences in languages. Our experiments indicate that neural interpretation of functional programs is more tractable when expressed in more “sugared” syntax, but further evidence is needed to completely support this claim. Further, we show that T5-Large is consistently better than T5-Small. Pretraining on the T5 corpus and tasks yields better results than training from scratch. Our analysis of different evaluation strategies show that lazy reduction is slightly easier to learn than eager reduction, but may generalize less on nontrivial splits of the data. Results with and without variable renaming do not show a clear trend. This is unexpected since variable renaming is an additional syntactic operation that complicates the task and is not strictly necessary.

Future work on more difficult tasks and data splits will allow us to better understand the generalization properties of our models as well as the impact of syntactic complexity on performance.