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: This paper proposes SYNCHROMESH, a framework for improving the reliability of pre-trained models for code generation. SYNCHROMESH first retrieves few-shot examples from a training set using Target Similarity Tuning. It then feeds the examples to a pre-trained language model and samples programs using Constrained Semantic Decoding, which can constraint the output to a set of valid programs in the target language. The authors evaluate the proposed approach by synthesizing code from natural language descriptions using GPT-3 and Codex in three real-world languages: SQL queries, Vega-Lite visualizations and SMCalFlow programs. The experimental results look promising.