Learning to Combine Per-Example Solutions for Neural Program Synthesis

1 Introduction . Program synthesis from examples tackles the problem of coming up with a computer program that satisfies a given set of Input-Output ( IO ) examples . Since the space of possible programs is large , an exhaustive search can be extremely time-consuming . Therefore , development of systems for program synthesis that can come up with a solution ( program satisfying the given IO examples ) within a limited time , such that it is practical for real-world applications , is a challenging task . Neural-guided program synthesis systems [ 4 , 30 ] try to expedite the search by using a neural network conditioned on the IO examples as a learned heuristic for the search procedure . In these systems , a neural network outputs probabilities over programs or properties of programs ( e.g . functions ) . These probabilities are then utilized to guide a search like depth-first or beam search . These systems try to find a program that satisfies all IO examples simultaneously , which under most of the settings can be hard . What if instead , we try to find this program in parts ? To understand this motivation , imagine a process wherein a programmer is asked to write a program that satisfies a set of unit test cases . They may begin by figuring out a program that satisfies a subset of unit test cases first , and later modifying the program to incorporate other corner cases . Shi et al . [ 26 ] uses this intuition to iteratively refine a program by mining fragments of Java code from partial solutions , based on a set of rules and predefined heuristics . Gupta et al . [ 12 ] also uses the same intuition , but in a different application for program repair . In this work , we consider breaking the complex problem of finding a program that satisfies all N given IO examples ( called the global solution ) into N smaller , easy to solve sub-problems , where each sub-problem involves finding a program satisfying only one IO example ( called per-example ∗Correspondence to : < dishu.905 @ gmail.com > 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . solution ) . The cues present in these per-example ( PE ) solutions are then combined to provide useful signals that can help guide the search for the global solution effectively . As a motivating example , consider the left part of Figure 1 , where five IO examples are given as a specification ( green box ) and we need to find a global solution pg ( red box ) that satisfies these five examples . The first stage of our approach consists of performing per-example searches to find a program pi conditioned on the i-th IO example . In our example , we start from IO example # 1 and find program p1 . In addition , we also check if p1 satisfies any other examples ( # 3 in figure ) . Iterating through the examples in this way results in a set of programs ( p1 , p2 , p3 ) that , taken together , in the ideal scenario , would satisfy all five IO examples . Looking closely at the discovered PE solutions , we see that they contain fragments of the global solution . This brings us to the second stage of our approach that addresses the challenge of how best to aggregate these PE solutions to produce a global solution . Towards that goal , we propose a neural network based architecture , which we refer to as Cross Aggregator ( CA ) . It is designed to learn to combine the cues present in these PE solutions , in a way that helps guide the search for pg . We model this aggregation using a multi-head cross-attention mechanism , which leverages the state of step-wise execution of the PE solutions and the synthesized global solution so far ( see Section 3.2 for details ) . Our key contributions can be listed as follows : • We consider breaking the standard program synthesis pipeline into two stages : ( a ) discovering PE solutions , and ( b ) aggregating the PE solutions such that it leads to a global solution . We refer to our approach that uses neural networks at both these stages as Neural Per-Example Program Synthesis ( N-PEPS ) . • We propose a neural network based multi-head attention architecture called Cross Aggregator ( CA ) that makes use of step-wise execution information to learn to combine the PE cues such that it helps guide the search for the global solution . • We demonstrate via experiments with programs of different lengths and under two different evaluation settings that when given the same time budget , our formulation shows significant improvements in success rate when compared to PCCoder [ 30 ] ( one of the leading techniques for neural-guided program synthesis ) and other ablation baselines . 2 Background . Suppose we are given a set X = { ( xi , yi ) } Ni=1 = { ri } Ni=1 of N IO examples and our task is to come up with a program pg that satisfies these examples . The i-th IO example ri consists of a pair of input xi and output yi . The program consists of T lines ( excluding lines with input variable declarations ) , i.e . pg = [ ptg ] T t=1 . To be practically meaningful , we impose the constraint that pg has to be found within a given time budget , specified by a timeout value . The syntax and semantics of pg are governed by a domain-specific language ( DSL ) . We use the DSL provided by Balog et al . [ 4 ] , which contains first-order functions ( e.g . SORT , REVERSE ) and higher-order functions ( e.g . MAP , FILTER ) that can take lambda functions ( e.g . ( * 4 ) , ( < 0 ) ) as input . The inputs and outputs can be either an integer or a list of integers ( see Appendix F of Balog et al . [ 4 ] for more details about the DSL ) . The Predict and Collect Coder ( PCCoder ) [ 30 ] provides state-of-art results for this DSL and is illustrative of methods that directly solve for all available IO examples at once . We refer to these methods as Global Program Search ( GPS ) . We will be building on PCCoder to propose our per-example approach . 2.1 PCCoder . PCCoder synthesizes programs one line at a time , through a model based on the notion of a program state . The program state is a two-dimensional memory of size N × ( ν + 1 ) obtained during the execution of t lines ( steps ) of a program on a set of N inputs . This means that for each IO example ri , there are up to ν slots for storing the input and intermediate program variables , with an additional slot for storing the output ( see Appendix A.2 for more details ) . Note that the initial state at t = 0 consists of only the IO examples ( see left part of Figure 2 ) . PCCoder consists of two learnable components ( i.e . neural networks ) , Hθ and Wφ , with parameters θ and φ. Hθ obtains the embedding of the current program state by average-pooling the representation of the ν + 1 slots corresponding to individual examples ( white boxes inside the state in Figure 2 ) into a vector of fixed size in RZ , where Z denotes the embedding size ( see Appendix A.2 for details of how these representations of slots are obtained ) . Wφ maps this state embedding to predictions of three quantities of interest for the next line in the program : ( a ) the next operator ôt ( or function e.g . MAP ) ; ( b ) the next statement ŝt ( operator along with its arguments e.g . MAP ( /2 ) b ) ; and ( c ) next drop vector d̂t which represents positions of variables that can be dropped from the state . The dropping is desirable as it creates slots for storing new variables , which in turn allows for synthesizing longer programs . There is a module called DropExec which executes a given line of the program against an example ri and stores the resulting variable ci in the next available slot in the state . If all ν slots in the state are filled , a variable is dropped from one of the slots using the drop vector and ci is stored there . The updated state can then be used for predicting the next line ( see right part of Figure 2 ) . Next , we provide details of how training and inference is done in PCCoder . Training : For training Hθ and Wφ , several instances of a specification X and the ground-truth program pg are provided . Given an instance and line t of the program , training operates by obtaining the ground-truth values of statements ( st ) , operator ( ot ) and drop vector ( dt ) . The statement and operator values are represented as one-hot vectors of size equal to the number of statements ( ns ) and number of operators ( no ) , respectively in the DSL . The drop vector is a multi-hot vector of size ν with ones at positions corresponding to the variables in the program that can be dropped , i.e . variables that don ’ t appear in subsequent lines in the program . The step-wise loss L is the sum of cross-entropy losses between the actual and predicted statement and operator , and the binary cross entropy loss between each position in the actual and predicted drop vector . The task of predicting the operator is an auxiliary task , i.e . it is used only during training and not at inference time , and is found to improve the training performance . During training , to obtain the updated state , the DropExec module chooses the drop-index to be a random entry from those positions in the drop vector dt that are ones . The right part of Figure 2 illustrates the process of training at step 2 . Inference : Inference is done using complete anytime beam search ( CAB ) [ 29 ] where the time for search is upper bounded by the timeout value . The CAB algorithm operates by performing different beam searches repeatedly in an outer loop . The pruning conditions of the beam search ( i.e. , beam size , expansion size ) are weakened with each iteration of the outer loop , until a solution is found . The inner loop consists of different steps of a single beam search . At each step , the beam consists of the most promising program prefixes , with each prefix represented as a tuple of the current program state , synthesized program until now and the product of the probabilities of the statements in the synthesized program . To synthesize the next line of the program , prefixes are expanded by executing the statements in decreasing order of statement probabilities and taking the argmax of the drop vector probabilities . The statement and drop vector probabilities are obtained using the trained neural networks Hθ and Wφ . The search terminates if we find a candidate program prefix that satisfies all N IO examples . The corresponding program is the synthesized global solution pg . Note that the search may fail and not discover a global solution within the specified timeout . Appendix A gives details of training and inference procedures , and modules of PCCoder .

Summary: This paper proposes a new approach for neural program synthesis by first searching for per-example solutions and then learn to combine them into a global solution. The paper uses PCCoder for synthesizing the per-example solutions, and uses an aggregation model based on cross-attention mechanisms. The aggregation model attends to the partially executed program states by the per-example programs in order to choose the next program token. The proposed technique is evaluated on a DSL involving computation on integers and list of integers, and is shown to achieve better performance than PCCoder.