Synchromesh: Reliable Code Generation from Pre-trained Language Models

1 INTRODUCTION . Large language models ( LLMs ) trained on massive corpora of unsupervised data have been shown to perform a wide range of tasks , including natural language generation , semantic parsing and sentiment analysis ( Brown et al. , 2020 ; Devlin et al. , 2019 ; Raffel et al. , 2020 ) . This can be achieved without task-specific training , but rather by adapting the model to each task at test-time using textual prompts , which can contain examples and natural language descriptions . In many cases , this methodology was shown to provide good performance , reducing the need to annotate large datasets for each task of interest ( Brown et al. , 2020 ; Shin et al. , 2021 ) . An important application of LLMs is in synthesizing programs from natural language descriptions ( Austin et al. , 2021 ; Chen et al. , 2021 ) . But this task is still challenging for LLMs . First , they can commit conceptual errors , generating code that misses the intent behind the given description . For example , when asked to reverse an array , the model might generate code that simply swaps the first and last elements . Indeed , users of current natural language-to-code systems report that models often produce code that is unrelated to their query ( Xu et al. , 2021 ) . Even when they capture the right intent , LLMs can still make implementation errors : the generated code can fail to execute . For reversing an array , a model might generate a loop with the correct structure but with an off-by-one error , causing a runtime exception . These errors are common even with very large models . For example , Austin et al . ( 2021 ) tested models with up to 137B parameters on generating short Python programs from natural language . Still , 47 % of the failures were due to syntax , typing or run-time errors ( as opposed to running but producing incorrect output ) . This is in line with theoretical results in Merrill et al . ( 2021 ) showing that programming language semantics can not be fully inferred from ungrounded data . Together , both observations suggest that simply scaling up LLMs might be ineffective to obtain reliable performance , especially for longer programs . In this paper , we address both conceptual and implementation errors with SYNCHROMESH , a framework for reliable code generation from pre-trained models . Since LLMs are highly sensitive to which few-shot examples are given in their prompt , we propose Target Similarity Tuning ( TST ) : a method for dynamically selecting semantically relevant examples for a given description . TST mitigates conceptual errors by learning to select examples with similar intent , even when their natural language descriptions seem unrelated in form . Given relevant examples , we then generate programs with Constrained Semantic Decoding ( CSD ) , a novel method for enforcing rich syntactic and semantic constraints during code generation on top of a frozen language model . Rich language-specific constraints , ranging from syntax validity to scoping and type-checking , can be implemented under the simple abstraction of completion engines ( CE ) . CSD aligns these constraints with the language model ’ s token vocabulary by leveraging Brzozowski language derivatives ( Brzozowski , 1964 ) . This guarantees that all sampled programs satisfies the implemented constraints , preventing whole classes of implementation errors by construction . The pipeline is illustrated in Figure 1 . We demonstrate the generality of SYNCHROMESH in three real-world languages : SQL ( database queries ) , Vega-Lite ( data visualization ) and SMCalFlow ( calendar applications ) . In experiments with GPT-3 and Codex , we observe that SYNCHROMESH can eliminate whole classes of errors that make outputs from unconstrained models either fail to execute or produce trivial results ( e.g. , empty charts ) . Furthermore , eliminating invalid programs consistently improves prediction accuracy . In summary , we make the following contributions : • We propose Target Similarity Tuning for selecting few-shot examples based on the similarity of the programs they describe , improving relevance and downstream performance . • We introduce completion engines as an abstraction that can implement rich classes of semantic program constraints , as we demonstrate in SQL , Vega-Lite and SMCalFlow . • We introduce a general , constraint-observing decoding algorithm , which aligns programming language constraints with the language model ’ s token vocabulary . • We evaluate our method in three natural language-to-code tasks . CSD and TST both show strong complementary gains in output validity and prediction accuracy across domains . 2 TARGET SIMILARITY TUNING . In this section , we first overview the challenge posed by conceptual errors in programs synthesized by LLMs . We then introduce TST , which improves performance through more relevant example selection . Throughout , we will use a real example of synthesizing a SQL database query to answer a question posed in natural language . Suppose a data analyst has a relational database of airports and wants to answer the following question : “ Which city has the highest number of airports ? ” One procedure for turning this description into a SQL query is to use an LLM such as GPT-3 ( Brown et al. , 2020 ) or Codex ( Chen et al. , 2021 ) . To prompt the model for the task at hand , we would feed it with a natural language description of the task and a selection of input-output examples . Given the analyst ’ s question , how do we select the most relevant examples from a training pool ? Liu et al . ( 2021a ) proposed to retrieve examples with similar natural language descriptions using a pre-trained paraphrase detection model . Figure 2a shows the most similar example from the Spider natural language-to-SQL dataset ( Yu et al. , 2018 ) according to Sentence-BERT ( Reimers & Gurevych , 2019 ) . The query “ Which city has the highest elevation ? ” is similar on a surface level : it also asks “ Which city has the highest ? ” . This training query asks about “ elevation ” , a property that is readily available as a column in the Airports table . Figure 2b shows GPT-3 ’ s output when given this and a few other examples . The model attempts to mimic the top example , referring to a nonexistent column “ NumberOfAirports ” . The issue is that we picked the example in the prompt based on description similarity and not SQL query similarity . In fact , the SQL query in the chosen example had a simplistic structure that was significantly different from the structure of the desired SQL query , and this contributed to the failure at Point ( b ) in Figure 2 . We want to retrieve examples that have relevant program structures for the test query . We do so using our fine-tuning scheme called Target Similarity Tuning ( TST ) . Formally , suppose D is a dataset of programs and associated utterances , with Di = ( pi , ui ) . Let S ( pa , pb ) ∈ [ 0 , 1 ] denote a normalized similarity metric between programs . If fθ is a pre-trained similarity model for natural language sentences , TST consists in fine-tuning f to predict the similarity between target programs given by S from their descriptions . Precisely , we minimize the mean-squared error loss : LTST ( θ ) : = Ei , j∼D [ fθ ( ui , uj ) − S ( pi , pj ) ] 2 . We define S using the classical tree edit distance algorithm from Zhang & Shasha ( 1989 ) to compare Abstract Syntax Trees ( ASTs ) . Figure 2c shows GPT-3 ’ s output when given examples selected with TST . Now , the output query is correct : it performs a “ group by ” on the “ City ” column , and sorts by the count of records in each group . This structure was already present in the top example selected by TST , corresponding to “ Return the team with the most technicians ” . Even if the analyst ’ s question and this utterance are drastically different in natural language , they share similarity in the SQL query that they describe . The TST objective is able to properly capture this fact . As our experiments show in Section 4 , TST significantly boosts the performance of both GPT-3 and Codex . 3 CONSTRAINED SEMANTIC DECODING . We now present Constrained Semantic Decoding ( CSD ) as an approach to eliminate implementation errors from code generated by LLMs . We first illustrate CSD with an example , and then formalize it using the abstraction of CEs . The example in Figure 2 showed that TST can help LLMs generate the correct program . In general , however , TST only helps LLMs by guiding toward the correct structure , but the model still needs to fill all the specific implementation details correctly . Figure 3 shows a case where the model can not simply adapt one example from the prompt . Here , the user ’ s query is “ Which city has the highest number of departing flights ? ” This query is similar to the previous one – in fact , TST retrieves the same top-1 example as before . But now the correct SQL query needs to join the “ Airports ” and “ Flights ” tables . GPT-3 generates the join condition Flights.AirportCode = Airports.SourceAirport , but this condition has a subtle error : the column names of the two tables are swapped . Thus , this query fails to execute . In general , unconstrained language models often make such implementation errors : using undeclared variables , losing track of nesting levels when producing complex expressions , or calling functions using arguments of the wrong type . Even the smallest of such errors prevents generated code from executing . CSD prevents implementation errors by construction ( as opposed to repairing after-the-fact ) . Imagine we have access to an oracle , which we call a CE , that can take a partial program and return all tokens that can extend that partial program toward a complete correct program . When the LLM is generating the program token by token , CSD ensures that the next token is sampled from the set returned by the CE . In Figure 3 , after generating “ T1. ” inside the “ on ” clause , our SQL CE resolves the alias and constrains the model to output one of the columns from the “ Flights ” table . This fixes the error seen previously during generation and produces the correct SQL query . 3.1 COMPLETION ENGINES . We now formally define CEs . Let Σ be a base alphabet , and ΣL ⊆ Σ∗ be the ( potentially infinite ) set of tokens of the target language . Our goal is to sample programs from a language L ⊆ Σ∗L – the set of valid programs . A CE CL is a partial function from Σ∗L to a set of tokens . We use a regular expression over Σ to represent a set of tokens . The strings in the domain of CL are called completion points , and a CE satisfies the following axioms : ( A1 ) The empty string and every p ∈ L must be completion points . For every p ∈ L , CL ( p ) = r′ $ ′ , where r′ $ ′ is the regular expression that matches the stop token . ( A2 ) If s ∈ Σ∗L is a completion point and t fully matches CL ( s ) , then their concatenation st must also be a completion point . ( A3 ) The CE is exhaustive ; that is , if s is a completion point and s = tt0 , where t0 is a token , then t should be a completion point and CL ( t ) should match t0 . Furthermore , we assume that CEs are only called after maximal matches . For example , if a partial program ends in an identifier , the CE can assume that the identifier is complete . Our CEs are implemented in two layers : a context-free layer , which enforces syntactic validity , and a context-sensitive layer , which encodes semantic constraints that depend on language semantics and the user ’ s context ( e.g. , the database ) . Below , we describe an automatic method for constructing context-free CEs directly from the target language ’ s grammar . The context-sensitive layer of an engine is specific to the target language . Table 1 provides an overview of several constraints implemented by our CEs for SQL , Vega-Lite and SMCalFlow , three rich languages with different syntactic and semantic structures . A detailed description of the three CEs can be found in Appendix C. Deriving completions from grammars Computer language parsers are often automatically generated from a grammar . The grammar contains enough information to derive the context-free layer of CEs . To facilitate this process , we created a library that extends any parser generated by ANTLR ( Parr & Fisher , 2011 ) , a popular LL ( ∗ ) top-down parser generator , to provide token-level completions . Namely , we ( i ) let the ANTLR-generated parser process the given program prefix p , ( ii ) retrieve its state in the Augmented Transition Network ( ATN ) at the last program token , ( iii ) traverse the ATN from that state to enumerate all possible next token productions . This process yields ( a ) a list of productions and token types { τj } Kj=1 that are allowed to follow p and ( b ) a partial AST Tp . Each CE takes { τj } and Tp as input to generate semantic context-sensitive constraints .

Summary: The paper considers the problem of text to code translation, usually done in state-of-the-art models using a natural language input fed into a transformer model together with other similar input/output examples as a few-shot learning and code being directly output by the model.  The paper proposes two enhancements of this process.  * The first one called Target Similarity Tuning (TST) is a way to pick input/output examples similar to the input text. Previous works use a pretrained models for natural language similarity (Sentence-BERT is compared in the paper) that determine what examples to include in the few-shot prompt and the TST proposes to also fine-tune this model on code examples such that some semantic similarity in encoded - for example using the same structure of the queries [think of SQL queries to synthesize]. * The other improvement called Constrained Semantic Decoding (CSD) proposes to alter the decoder of the transformer model and to make it avoid generating programs that are impossible to complete to syntactically or semantically correct ones.