Latent Programmer: Discrete Latent Codes for Program Synthesis

1 INTRODUCTION . Our focus in this paper is program synthesis , one of the longstanding grand challenges of artificial intelligence research ( Manna & Waldinger , 1971 ; Summers , 1977 ) . The objective of program synthesis is to automatically write a program given a specification of its intended behavior , such as a natural language description or a small set of input-output examples . Search is an especially difficult challenge within program synthesis ( Alur et al. , 2013 ; Gulwani et al. , 2017 ) , and many different methods have been explored , including top-down search ( Lee et al. , 2018 ) , bottom up search ( Udupa et al. , 2013 ) , beam search ( Devlin et al. , 2017 ) , and many others ( see Section 2 ) . We take a different philosophy : Can we learn a representation of programs specifically to help search ? A natural way of representing a program is as a sequence of source code tokens , but the synthesis task requires searching over this representation , which can be difficult for longer , more complex programs . A programmer often starts by specifying high-level components of a program as a plan , then fills in the details of each component i.e . in string editing , a plan could be to extract the first name , then the last initial . We propose to use a sequence of latent variable tokens , called discrete latent codes , to represent such plans . Instead of having a fixed dictionary of codes , we let a model discover and learn what latent codes are useful and how to infer them from specification . Our hypothesis is that a discrete latent code – a sequence of discrete latent variables – can be a useful representation for search ( van den Oord et al. , 2017 ; Roy et al. , 2018 ; Kaiser et al. , 2018 ) . This is because we can employ standard methods from discrete search , such as beam search , over a compact space of high-level plans and then over programs conditioned on the plan , in a two-level procedure . We posit that the high-level search can help to organize the search over programs . In the string editing example earlier , a model could be confident that it needs to extract the last initial , but is less sure about whether it needs to extract a first name . By changing one token in the latent code , two-level search can explore alternative programs that do different things in the beginning . Whereas in traditional single-level search , the model would need to change multi-token prefixes of the alternatives , which is difficult to achieve in limited budget search . We propose the Latent Programmer , a program synthesis method that uses learned discrete representations to guide search via a two-level synthesis . The Latent Programmer is trained by a self- supervised learning principle . First a discrete autoencoder is trained on a set of programs to learn discrete latent codes , and then an encoder is trained to map the specification of the synthesis task to these latent codes . Finally , at inference time , Latent Programmer uses a two-level search . Given the specification , the model first produces a L-best list of latent codes from the latent predictor , and uses them to synthesize potential programs . On two different program synthesis domains , we find empirically that the Latent Programmer improves synthesis accuracy by over 10 % compared to standard sequence-to-sequence baselines as RobustFill ( Devlin et al. , 2017 ) . We also find that our method improves diversity of predictions , as well as accuracy on long programs . 2 BACKGROUND . Problem Setup The goal in program synthesis is to find a program in a given language that is consistent with a specification . Formally , we are given a domain specific language ( DSL ) which defines a space Y of programs . The task is described by a specification X ∈ X and is solved by some , possibly multiple , unknown program ( s ) Y ∈ Y . For example , each specification can be a set of input/output ( I/O ) examples denoted X = { ( I1 , O1 ) , . . . ( IN , ON ) } . Then , we say that we have solved specification X if we found a program Y which correctly solves all the examples : Y ( Ii ) = Oi , ∀i = 1 , . . . , N . As another example , each specification can be a natural language description of a task , and the corresponding program implements said task . An example string transformation synthesis task with four I/O examples together with a potential correct program in the string transformation DSL is shown in Figure 1 . Vector Quantization Traditionally , neural program synthesis techniques process the input specification as a set of sequences and predicts the output program token-by-token ( Devlin et al. , 2017 ) . In this work , we present a new approach for synthesis that performs structured planning in latent space using a discrete code . We conjecture that programs have an underlying discrete structure ; specifically , programs are compositional and modular with components that get reused across different problems . Our approach leverages this structure to guide the search over large program spaces . Following works in computer vision ( van den Oord et al. , 2017 ; Roy et al. , 2018 ) , we discover such discrete structure by using a Vector Quantized Variational Autoencoder ( VQ-VAE ) . VQ-VAEs work by feeding the intermediate representation of an autoencoder through a discretization bottleneck ( van den Oord et al. , 2017 ) . For completeness , we provide background on VQ-VAEs below . In a VQ-VAE , latent codes are drawn from a discrete set of learned vectors c ∈ RK×D , or codebook . Each element in the codebook can be viewed as either a token with id k ∈ [ K ] or as an embedding ck ∈ RD . To generate the discrete codes , the continuous autoencoder output e is quantized via nearest-neighbor lookup into the codebook . Formally , the token id qk ( e ) and quantized embedding qc ( e ) are defined as qc ( e ) = cqk ( e ) where qk ( e ) = arg min k∈ [ K ] ||e− ck||2 . ( 1 ) For input x , the training loss for a VQ-VAE consists of : a reconstruction loss for the encoder-decoder weights , a codebook loss that encourages codebook embeddings to be close to the continuous vectors which are quantized to them , and a commitment loss that encourages the encoded input ec ( x ) to `` commit '' to codes i.e . not switch which discrete code it is quantized to . The loss is given by , L ( c , θ , φ ) = log pθ ( x | qc ( ecφ ( x ) ) ) + ||sg ( ecφ ( x ) ) − c ) ||22 + β||sg ( c ) − ecφ ( x ) ||22 , ( 2 ) where θ , φ are the parameters of the decoder and encoder , respectively , sg ( · ) is the stop gradient operator that fixes the operand from being updated by gradients , and β controls the strength of the commitment loss . To stabilize training , van den Oord et al . ( 2017 ) also proposed removing the codebook loss and set the codebook to an exponential moving average ( EMA ) of encoded inputs . 3 SYNTHESIS WITH DISCRETE LATENT VARIABLES . We propose a two-level hierarchical approach to program synthesis that first performs high-level planning over an intermediate sequence , which is then used for fine-grained generation of the program . In our approach , a top-level module first infers a latent code , which gets used by a low-level module to generate the final program . 3.1 HIERARCHY OF TWO TRANSFORMERS . Our proposed Latent Programmer ( LP ) architecture consists of two Transformers in a two-level structure . The architecture comprises of two modules : a latent predictor which produces a latent code , which can be interpreted as a course sketch of the program , and a latent program decoder , which generates a program conditioned on the code . The latent code consists of discrete latent variables as tokens , which we arbitrarily denote TOK_1 , ... , TOK_K , whose meanings are assigned during training . Both components use a Transformer architecture due to their impressive performance on natural language tasks ( Vaswani et al. , 2017 ) . To help the model assign useful meanings to the latents , we also leverage a program encoder , which is only used during training . The program encoder ec ( Y ) encodes the true program Y = [ y1 , y2 , . . . , yT ] into a shorter sequence of discrete latent variables Z = [ z1 , z2 , . . . , zS ] , represented as codebook entries ; that is , each zi ∈ RD is one of K entries in a codebook c. The latent sequence serves as the ground-truth high-level plan for the task . The function ec ( Y ) is a Transformer encoder , followed by a stack of convolutions of stride 2 , each halving the size of the sequence . We apply the convolution ` times , which reduces a T -length program to a latent sequence of length dT/2 ` e . This provides temporal abstraction , since the high-level planning actions are made only every 2 ` steps . In summary , the program encoder is given by ec ( Y ) ← h ` ; hm ← Conv ( hm−1 ) for m ∈ 1 . . . ` ; h0 ← TransformerEncoder ( Y ) . ( 3 ) Here TransformerEncoder ( · ) applies a stack of self-attention and feed-forward units on input embeddings via a residual path , described in detail by Vaswani et al . ( 2017 ) . This will be used , along with the latent program decoder , as an autoencoder during training ( see Section 3.2 ) . The latent predictor lp ( X ) autoregressively predicts a coarse latent code lp ( X ) ∈ RS×K , conditioned on the program specification X . The latent predictor outputs a sequence of probabilities , which can be decoded using search algorithms such as beam search to generate a predicted latent code Z ′ . This is different than the program encoder , which outputs a single sequence Z , because we use the latent predictor to organize search over latent codes ; at test time , we will obtain a L-best list of latent token sequences from lp ( X ) . The latent predictor is given by a stack of Transformer blocks with the specification X as inputs . Similarly , the latent program decoder d ( Z , X ) defines an autoregressive distribution over program tokens given the specification X and the coarse plan Z ∈ RS×K , represented as codebook entries . The decoder is a Transformer that jointly attends to the latent sequence and program specification . This is performed via two separate attention modules , whose outputs are concatenated into the hidden unit . Formally , given a partially generated program Y ′ = [ y′1 , y ′ 2 , . . . , y ′ t−1 ] , and the encoded specification E = TransformerEncoder ( X ) , the latent program decoder performs ht = Concat ( TransformerDecoder ( Y ′ , E ) t−1 , TransformerDecoder ( Y ′ , Z ) t−1 ) , ( 4 ) where TransformerDecoder ( x , y ) denotes a Transformer decoder applied to outputs y while attending to inputs encoding x , and the subscript indexes an entry in the resulting output sequence . Finally , the distribution over output token k is given by dt ( Z , X ) = Softmax ( W ( ht ) ) , where W is a learned parameter matrix . Finally , the latent program decoder defines a distribution over programs autoregressively as p ( Y |Z , X ) = ∏ t p ( yt|y < t , Z , X ) , where p ( yt|y < t , Z , X ) = dt ( Z , X ) . When X is multiple I/O examples , each example is encoded asEi = TransformerDecoder ( Ii , Oi ) . Then , a separate hidden state per I/O is computed following equation 4 , followed by a late max-pool to get the final hidden state . Note that the program encoder and latent program decoder make up a VQ-VAE model of programs , with additional conditioning on the specification . The complete LP architecture is summarized in Figure 2 , and an end-to-end example run of our architecture is shown in Figure 4 .

Summary: This paper proposes a two-level hierarchical program synthesizer, Latent Programmer, which first predicts a sequence of latent codes from given input-output examples, and then decodes the latent codes into a program.  The sequence of latent codes can be viewed as a high-level synthesis plan, guiding the subsequent low-level synthesis. Latent Programer significantly outperforms RobustFill on string manipulation tasks and achieves state-of-the-art results on Python code generation tasks.