Supplemental: Leveraging Language to Learn Program Search Heuristics and Abstractions

ions This contains the supplemental appendix to the 2021 ICML paper. It is organized sequentially in reference to the main text; S{N} refers back to section N in the main text. A complete release of code for our implementation, including command line scripts to replicate the experiments in the paper and links to the datasets, can be found at: https://bit.ly/3g9361W. S4. Base learning algorithm: DreamCoder The LAPS framework described in the main paper (Sec. 5) is a general one for extending Bayesian models of program learning to incorporate information from natural language (see (Liang et al., 2010; Lake et al., 2015; Dechter et al., 2013; Lake et al., 2013)). Our concrete implementation and experiments use the DreamCoder approach of (Ellis et al., 2021; 2018) as the base synthesis algorithm, which implements the hierarchical Bayesian formulation of program learning. It defines a modular interface with two primary learning components: a learned conditional inference model for search (as a neural search heuristic); and a learned abstraction algorithm for updating the program prior (based on program refactoring and compression) (Ellis et al., 2021). Each of these learning components has been additionally implemented in other work (such as (Devlin et al., 2017; Polosukhin & Skidanov, 2018; Nye et al., 2019; Parisotto et al., 2016; Balog et al., 2016) for neurally guided synthesis, and (Dechter et al., 2013; Zhang et al., 2017; Shin et al., 2019; Artzi et al., 2014; Dumancić & Cropper) for program abstraction learning). This supplementary section provides theoretical and implementation details on the DreamCoder algorithm we use in our experiments (summarized in Sec. 4). We match our implementation as closely as possible to the original work for comparison with published baselines. We provide key details relevant to the language-guided extension, but strongly recommend the original works which introduce the DreamCoder algorithm (Ellis et al., 2021; 2018) for further reference. S4.1 Program prior and MDL equivalence Hierarchical Bayesian program learning formulations require a prior over expressible programs. DreamCoder is learned iteratively: it is initialized with a base library L0 and returns a library Lf containing program abstractions learned from solving training tasks. Therefore, DreamCoder defines its program prior with respect to the current library Li maintained at each iteration. This is parameterized as a simple PCFG P[ρ|L, θL] whose productions are of the form li → lj ∈ L, each with a real-valued weight θLl, where the probability of a program ρ is given by P[ρ|L, θL] = ∏ l∈ρ P[l|L, θL] (Sec. 4.1). Minor complexity arises in order to support typing (Pierce, 2002): following (Ellis et al., 2018), the library Li is implemented as a set of polymorphically typed λ-calculus expressions. The only change this produces to the original prior definition is to restrict the set of possible productions under the PCFG: that is, permissible productions are of the form li → lj ∈ {L|li → lj is well typed}. The prior probabilities of programs are therefore calculated with respect to the set of well-typed productions. As discussed in the main paper, this prior definition is equivalent to a minimum description-length prior over programs under (L, θL) when all θL < 1.0, as the product of additional productions in an expression will strictly decrease as the number of productions in an expression increases. S4.2 Amortized conditional inference Figure 1. Architecture of the neural model Qi(ρ|t,Li). The model takes as input task examples t. These are encoded using a domainspecific encoder E(t). Task encodings feed to an MLP and activation layer and output a tensor Q. This parameterizes a distribution over program bigrams in the final DSL, which defines a conditional distribution from which to enumerate programs during search. To identify programs that solve tasks t while obtaining high probability under P[ρ|L, θL], DreamCoder trains a neural search heuristic Qi(ρ|t,Li) at each iteration i to approximate the inverse model. The training procedure in (Ellis et al., 2021) (summarized in Sec. 4.2) is a key contribution of the original work for learning in the distant supervision setting. The model is trained on samples from the generative prior (providing an endless