S YNCHROMESH : R ELIABLE C ODE G ENERATION FROM P RE - TRAINED L ANGUAGE M ODELS
A BSTRACT
Large pre-trained language models have been used to generate code, providing a ﬂexible interface for synthesizing programs from natural language speciﬁcations.
However, they often violate syntactic and semantic rules of their output language, limiting their practical usability.
In this paper, we propose S YNCHROMESH : a framework for substantially improving the reliability of pre-trained models for code generation.
S YNCHROMESH comprises two components.
First, it retrieves few-shot examples from a training bank using Target Similarity Tuning (TST), a novel method for semantic example selection.
TST learns to recognize utterances that describe similar target programs despite differences in surface natural language features.
Then, S YNCHROMESH feeds the examples to a pre-trained language model and samples programs using Constrained Semantic Decoding (CSD): a general framework for constraining the output to a set of valid programs in the target language.
CSD leverages constraints on partial outputs to sample complete correct programs, and needs neither re-training nor ﬁne-tuning of the language model.
We evaluate our methods 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.
These domains showcase rich constraints that CSD is able to enforce, including syntax, scope, typing rules and contextual logic.
We observe substantial complementary gains from CSD and
TST in prediction accuracy and in effectively preventing run-time errors
I NTRODUCTION
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-speciﬁc 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 ﬁrst 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 cannot 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 S YNCHROMESH , 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-speciﬁc 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 satisﬁes the implemented constraints, preventing whole classes of implementation errors by construction.
The pipeline is illustrated in Figure 1.
We demonstrate the generality of S YNCHROMESH 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 S YNCHROMESH 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.
T ARGET S IMILARITY T UNING
In this section, we ﬁrst 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 (cid:3) ?”.
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 signiﬁcantly 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 ﬁne-tuning scheme called Target Similarity Tuning (TST).
Formally, suppose D is a dataset of programs and associated utterances, with D i
= ( p i , u i ) .
Let S ( p a , p b ) ∈ [0 , 1] denote a normalized similarity metric between programs.
If f θ is a pre-trained similarity model for natural language sentences, TST consists in ﬁne-tuning f to predict the similarity between target programs given by
S from their descriptions.
Precisely, we minimize the mean-squared error loss:
[Equation]
We deﬁne 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 signiﬁcantly boosts the performance of both GPT-3 and Codex.
C ONSTRAINED S EMANTIC D ECODING
We now present Constrained Semantic Decoding (CSD) as an approach to eliminate implementation errors from code generated by LLMs.
We ﬁrst 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 ﬁll all the speciﬁc implementation details correctly.
Figure 3 shows a case where the model cannot simply adapt one example from the prompt.
Here, the user’s query is “Which city has the highest number of departing ﬂights?”
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 ﬁxes the error seen previously during generation and produces the correct SQL query.
3. C OMPLETION E NGINES
We now formally deﬁne CEs.
Let Σ be a base alphabet, and Σ L ⊆ Σ ∗ be the (potentially inﬁnite) 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 C L 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 C L are called completion points , and a CE satisﬁes the following axioms: (A1)
The empty string and every p ∈ L must be completion points.
For every p ∈ L , C L ( p ) =
r (cid:48) $ (cid:48) , where r (cid:48) $ (cid:48) is the regular expression that matches the stop token.
(A2)
If s ∈ Σ ∗ L is a completion point and t fully matches C L ( 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 = tt 0 , where t 0 is a token, then t should be a completion point and C L ( t ) should match t 0 .
Furthermore, we assume that CEs are only called after maximal matches.
For example, if a partial program ends in an identiﬁer, the CE can assume that the identiﬁer is complete.
[Table]
Given a partial program, CEs return a regular expression that matches the valid tokens that can follow.
Here, we show positive and negative token examples for each such regular expression.
This abstraction allows domain experts to encode a wide range of expressive code generation constraints.
Our CEs are implemented in two layers: a context-free layer, which enforces syntactic validity [Table] and a context-sensitive layer, which encodes semantic constraints that depend on language semantics and the user’s context (e [Table]
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 speciﬁc 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 preﬁx 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 T p .
Each CE takes { τ j } and T p as input to generate semantic context-sensitive constraints.
3. F ROM C OMPLETION E NGINES TO A D ECISION P ROCEDURE
We use CEs to guide sampling from an LLM.
A key component of our algorithm for constrained sampling is a decision procedure for membership in preﬁx-closure of the set L of all valid programs.
The preﬁx-closure L c of a language L contains all programs in L as well as all of their preﬁxes.
Intuitively, L c contains all partial programs that can be completed to a valid program.
Given a CE
C L , our ﬁrst goal is to build a decision procedure for L c : given a string s , does it belong to L c ?
We answer if s ∈ L c by repeatedly calling C L on certain preﬁxes p of s and matching the regular expression C L ( p ) with sufﬁxes of s .
We start with p being the empty string.
We ﬁnd the maximal preﬁx of s that matches the regular expression C L ( p ) and remove it from s and add it to p , and repeat until the match fails.
There are two cases: either s is empty now, which means the input string was a completion point and hence it is in L c , or s now is the remainder left after removing the largest preﬁx that was a completion point.
For the second case, we must check: does there exist a completion string c such that sc fully matches the regular expression C L ( p ) ?
This question can be efﬁciently answered by Brzozowski derivatives (Brzozowski, 1964).
Formally, the derivative of a formal language S with respect to a string u is another formal language u − 1 S
= { v : uv ∈ S } .
In other words, it is precisely the set of strings that can complete u to some string in
S .
If u − 1 S = ∅ , then no string in S starts with u .
Brzozowski derivatives are efﬁcient to compute for our regular languages (or regular expressions deﬁning them) – we describe a simple linear-time algorithm in the Appendix.
Given the derivative of C L ( p ) , answering whether s can be completed to belong to C L ( p ) reduces to performing a simple regular expression match.
This operation answers the case when the remainder is non-empty and completes our decision procedure for L c .
3. T HE C ONSTRAINED S EMANTIC D ECODING ALGORITHM
Using the decision procedure for L c , we can now describe the Constrained Semantic Decoding algorithm.
Suppose s
∈
L c is the language model’s output so far (we start with (cid:15) ).
If Σ M is the model’s vocabulary, we can compute the set of valid next tokens V M ( s )
= { t ∈ Σ M : st ∈ L c } by using our decision procedure for each token in the vocabulary Σ M .
In other words, we maintain the invariant that the model’s current partial output s is in L c , and make progress by using the model to sample from V M ( s ) , instead of the unconstrained Σ M .
Once we have a complete program, we are guaranteed that it will satisfy all constraints enforced by the CE.
One subtlety to note is that language models and programming languages have drastically different tokenizations; i.e., C L and LLM work with different tokens.
For instance, a long string literal is a single SQL token, but might span multiple tokens for the language model.
Similarly, a single token from the language model’s vocabulary might close multiple parentheses at once.
In general, token boundaries between the two can be arbitrarily misaligned.
Each decision of whether st belongs to
L c can potentially cross multiple completion points, or might not even ﬁnish a maximal match to the previous completion point (see the Appendix for an example prediction in Vega-Lite where this happens multiple times).
Nevertheless, our CSD algorithm described here naturally handles this alignment problem.
Hence, in S YNCHROMESH , CEs do not need to be aware of this issue – they can be fully implemented in terms of the target language’s tokens.
Our implementation applies substantial optimizations that leverage the structure of Byte-Pair Encoding vocabularies (namely, that many tokens are preﬁxes of longer tokens) and reuse computation.
We detail these optimizations in Appendix E.
In our experiments with GPT-3, CSD adds an average of 8% overhead to the sampling procedure – a relatively small impact to trade for output correctness.
E XPERIMENTS
We evaluate S YNCHROMESH in three tasks of synthesizing code from natural language descriptions.
For SQL, we use the Spider dataset (Yu et al., 2018).
For Vega-Lite, we use the NLV Corpus (Srinivasan et al., 2021).
For SMCalFlow, we use the dataset that introduced the language (Andreas et al., 2020).
In NLV, which has visualizations over 3 different datasets, we alternate using each dataset as a test-set by only using training examples from the other two datasets.
In Spider and
SMCalFlow, we use the training/validation set split given in each dataset.
Example selection model To select examples, we use Sentence-BERT (Reimers & Gurevych, 2019) to fetch the 5 closest examples by cosine similarity.
When using TST, we ﬁne-tuned the model with the TST objective in both the Spider and SMCalFlow training sets.
The NLV corpus is smaller and does not provide a clear train-test split to fairly evaluate TST.
Holding out one dataset and ﬁne-tuning on the remaining two yields S YNCHROMESH accuracies of over 90%.
However, we attribute that performance to the fact that NLV has only 10 distinct visualizations and the same participants labeled all three datasets.
For that reason, we omit Vega-Lite from the TST experiments.
[Table] Language models We used the two largest models from the GPT-3 family (Brown et al., 2020, with 13B and 175B parameters), as well as the largest Codex model (Chen et al., 2021).
Codex shares the same architecture with 175B GPT-3, but its training set contained a larger portion of source code in a variety of languages.
Our only access to the models was through the public OpenAI
HTTP API, which allowed us to apply constraints by adding a bias to logits.
We describe the necessary adaptations of CSD to this setting in Appendix F.
Metrics For Vega-Lite and SMCalFlow, we report the exact-match accuracy between predictions and ground-truth (ﬁeld order is disregarded in Vega-Lite).
In SQL, we instead measure execution accuracy , comparing query results instead.
For a more ﬁne-grained signal, we additionally measure the edit distance between the predicted and ground-truth ASTs using the normalized tree edit distance (Zhang & Shasha, 1989).
4. R ESULTS
Table 2 and Figure 4 summarize the results of evaluating S YNCHROMESH versus various baselines.
The key observations are as follows:
S YNCHROMESH improves reliability on top of all pre-trained LLMs
First it improves top-1 accuracy (exact or execution -measured) over any
pre- trained LLM in all domains. SMCalFlow beneﬁts the most, likely because this domain-speciﬁc language is absent in the LLM pre-training corpus.
For
SQL and SMCalFlow, the absolute gain is almost the same for equally-sized GPT-3 and Codex.
Second, S YNCHROMESH dramatically improves validity.
In SQL, it eliminates execution errors from 29% of the queries generated by GPT-3 13B (as validity improves from 43% to 72%).
Even
Codex beneﬁts, with 12% more queries executing successfully after S YNCHROMESH augmentation.
In Vega-Lite and SMCalFlow, S YNCHROMESH improves reliability even more substantially.
GPT-3 13B only produces valid charts for 55% of the queries in NLV; all errors are eliminated with
S YNCHROMESH .
This is nearly paralleled in SMCalFlow, in which all models produce well-typed programs 97% of the time or more with S YNCHROMESH . [Figure]
S YNCHROMESH brings the output closer to ground truth.
While error prevention alone is trivial (e.g., with a constant error-free prediction), that would improve neither accuracy nor the average edit distance to the ground-truth, unlike S YNCHROMESH results.
Again, we observe improvements in all domains and the most in SMCalFlow.
For GPT-3 175B, S YNCHROMESH reduced the average edit distance from 0.41 to 0.18.
TST and CSD bring complementary beneﬁts.
Our ablation studies reported in Table 2 show that their combination performs better than either one separately.
TST helps LLMs generate programs in the “vicinity” of the correct one, and CSD helps by “guiding” the models toward the correct one.
S YNCHROMESH adds more value for longer programs.
Program synthesis is hardest when the target program is complex.
Does S YNCHROMESH improve synthesis of longer programs, or are its beneﬁts coming from ﬁxes to small programs?
Figure 4(a) shows accuracies and (b) validity for SMCalFlow broken down by the length of the ground truth program (we show results for SQL in the Appendix).
Here, program lengths are shown as their percentile. With S YNCHROMESH , we see that accuracy decays at a slower pace, and validity remains high throughout, when compared to Codex alone. This indicates that S YNCHROMESH improves the ability of base models to generate longer programs.
LLMs augmented with S YNCHROMESH approach but underperform supervised models.
For context, we include state-of-the-art results at the time of writing for each task in Table 2.
We note that these methods ﬁne-tune or train the underlying language-to-code model on each task, thus are not directly comparable to LLMs with S YNCHROMESH .
That said, we observe that base LLMs—even Codex— substantially underperform supervised models (19% worse for SQL; 27% worse for SMCalFlow), and S YNCHROMESH helps narrow that gap (now 11% worse for SQL; 9% worse for SMCalFlow).
S YNCHROMESH outperforms “generate-then-test”.
CSD enforces program constraints during generation.
Instead, prior work has leveraged a “generate-then-test” approach: take multiple samples and ﬁlter out those that produce errors or violate constraints (Chen et al., 2021).
Figure 4(b) evaluates this approach with Codex, the highest performing base LLM.
We sample from Codex with a temperature τ = 0 .
7 to obtain diverse but high-quality samples.
We then compute the “Valid@K” metric by using the “Pass@K” estimator from Chen et al.
(2021) to calculate the probability of at least one valid sample among K , with K ≤ 5 .
In SQL, Codex needs 3 samples to match S YNCHROMESH (Valid@K = 85%).
In SMCalFlow and Vega-Lite, S YNCHROMESH is able to virtually eliminate errors with 1 sample, while “Valid@5” for Codex is still below 93%.
This provides evidence that even the best LLMs beneﬁt from incremental validation, especially in less popular languages.
4. D ISCUSSION
In our experiments, S YNCHROMESH was able to improve accuracies and program validity across all languages due to better examples to use as a reference (TST) and preventing errors during generation (CSD).
Yet, this approach has an important limitation.
While TST can reduce conceptual errors, and
CSD can guarantee that certain implementation errors never occur (e.g., type errors in SMCalFlow,
When those occur CSD is usually insufﬁcient to correct the prediction.
Figure 5 shows an example in Vega-Lite from the “Cars” dataset in NLV.
Here, the user asks for one scatter plot for each origin, indicating faceting (multiple charts).
GPT-3 alone produces an invalid Vega-Lite chart type, can eliminate this error, guiding GPT-3 to generate “point”
However conceptual error
Lite
Nonetheless, we argue that improving validity is especially important for user-facing applications.
Users of language-to-code systems might need to rephrase their request or to edit the system’s output.
But outputs that fail to even execute undermine user experience: ﬁxing an automatically generated program can be more cumbersome than writing it in the ﬁrst place.
In LLM-driven systems like
Github Copilot, implementation errors can remain unnoticed and introduce bugs or vulnerabilities.
R ELATED W ORK
Program synthesis is a long-standing AI challenge with the goal of generating computer programs from higher-level speciﬁcation (Gulwani et al., 2017).
In particular, synthesis from natural language descriptions has gained recent attention (Liu et al., 2016;
Yaghmazadeh et al., 2017), thanks to advances in natural language processing models such as Transformers (Vaswani et al., 2017).
Typically, LLMs such as GPT-3 (Brown et al., 2020) and Codex (Chen et al., 2021) output an unconstrained sequence of tokens, and still often make conceptual or implementation errors in generated programs (Austin et al., 2021).
Specialized training, e.g. to output an AST (Wang et al., 2020;
Yin & Neubig, 2017), can mitigate syntactic errors, but still does not guarantee accuracy or conformance to domain-speciﬁc semantic constraints.
Moreover, it requires a specialized architecture and a decoding procedure for each target language.
Instead, S YNCHROMESH applies such constraints at inference, neither using specialized architectures nor ﬁne-tuning the LLM.
The general idea of constraining LLMs when generating programs has been explored in recent work.
Shin et al.
(2021) applied syntactic constraints for semantic parsing.
However, their method requires enumerating all valid programs for determining valid next tokens for the LLM, and does not enforce semantic constraints.
In concurrent work, Scholak et al.
(2021) applied similar semantic constraints to synthesizing SQL queries.
The authors substantially improve the performance of an already ﬁnetuned model by leveraging an incremental parser.
We see CSD as a generalization of these efforts, as our completion engines can apply context-sensitive constraints by dynamically constructing regular expressions.
Aligning these constraints with the underlying model vocabulary does not require ﬁnetuning: S YNCHROMESH only trains the much smaller target similarity model (Section 2).
Since the emergence of LLMs, researchers have developed numerous techniques to adapt them to new domains (Liu et al., 2021b).
Many focus on prompting , i.e. generating a domain- and instancespeciﬁc input to an LLM to increase the likelihood of correctness.
In few-shot prompt augmentation ,
Gao et al.
(2020) use pre-trained sentence embeddings to select the closest prompt examples to the given input instance.
Liu et al.
(2021a) further ﬁne-tune sentence embeddings on the available training set of input utterances.
TST in S YNCHROMESH takes this approach a step further, and ﬁnetunes the embedding models based on output similarity.
It optimizes the amount of relevant output bits in the prompt, thereby reinforcing the necessary hints for the LLM.
C ONCLUSION
S YNCHROMESH augments program synthesis with pre-trained LLMs to prevent both conceptual and implementation errors in their generated programs.
We designed S YNCHROMESH to be easily usable for any domain with minimal NLP or LLM knowledge expected from a domain expert.
As such it (a) automatically generates the completion engine API from the language grammar, (b) does not require ﬁne-tuning the underlying LLM, drastically reducing the data/compute budget and (c) integrates into the decoding loop or inference API with minimum computational overhead.
Our method signiﬁcantly improves performance of both GPT-3 and Codex in three real-world programming languages, both by boosting prediction accuracy and consistently improving output validity.
While real-world and well-established, the domains we study are still not Turing-complete.
We envision extending S YNCHROMESH to a Turing-complete language like Python can vastly increase reliability of LLM-based systems like Github Copilot.
This requires further extension of CSD to integrate with the parser/interpreter of the target language and to study applicable classes of constraints.
The TST technique, however, can be used in any LLM-based language-to-code system.
A C ONSTRAINED S EMANTIC D ECODING A LGORITHM
In Section 3.
we described the Constrained Semantic Decoding algorithm in the text.
We provide the same algorithm in pseudo-code in Algorithms 1 and 2 in Figure 6 below.
ValidPrefix is our decision procedure for L c , while CSD samples a complete program from the model M by making calls to ValidPrefix .
We use s · t to denote concatenation of s and t , and s − 1 · t to denote the Input M Input token set Σ M while next token ← ; endreturn
Input : A string s Output: True iff s ∈ L c p
”” ; while preﬁx ⊥ do p next preﬁx ; s − end string obtained by removing the preﬁx s from t . The utility function startswith ( s, r ) returns the maximal preﬁx of s that matches the regular expression r , and returns ⊥ if there is no such match.
The function Sample ( Dist , S ) returns a token from the set S of tokens sampled from the distribution
Dist restricted to S .
The CSD procedure uses the model M on the partial program to generate a distribution on the next token, but constrains it to belong to the set of valid tokens determined by the completion engine C L .
B C OMPUTING B RZOZOWSKI DERIVATIVES
Regular expression derivatives can be expensive to compute in general.
The main challenge is that
Klenee stars fork the decision procedure: the algorithm must decide to repeat or skip the pattern inside a star.
However, in this work, all our regular expressions come from grammars of programming languages.
These languages have one important feature in common: tokens are deﬁned by greedy matching rules.
This means that Klenee stars consume as many characters as possible, and do not backtrack if a tokenization error occurs (the error simply propagates).
Under the assumption that
Klenee stars have greedy semantics, derivatives can be computed in linear time.
The algorithm for computing derivatives of a regular expression can be worked out by looking at all constructors of regular expressions.
Table 3 details this computation for regular expressions of a base alphabet Σ .
[Table]
If the derivative of R 1 w.r.t.
c (cid:48) is not ∅ , then it’s the concatenation of that with R 2 .
Otherwise, if R 1 matches (cid:15) , then it is simply the derivative of R 2 w.r.t.
c (cid:48) .
If not, then the result is ∅ .
[Table]
, simply
[Table]
the other
[Table]
If the derivative of R w.r.t.
c is not ∅ , then return the concatenation of that derivative with R ∗ .
Otherwise, return ∅ .
C C OMPLETION E NGINES
Here, we describe our completion engines for SQL, Vega-Lite and SMCalFlow in more detail.
C. SQL
SQL database queries are executed in the context of a particular database containing a set of tables.
Each table has a schema, which speciﬁes named columns, their data types and constraints such as foreign keys.
We refer to Yu et al.
for a more detailed description of the SQL language.
Our CE for SQL enforces that only columns that exist in the tables in the database are used.
The main challenge is that queries often specify aliases.
Thus, during parsing, we construct a symbol table mapping aliases to the tables they refer to.
We enforce that tables that already have an alias should only be referred to by their alias, not in their unqualiﬁed form.
Since aliases can be referred to in the SELECT clause before being deﬁned in the FROM clause, we also keep track of undeﬁned aliases to enforce that they will be assigned a table later.
Moreover, a condition in the WHERE clause might involve a nested query, which in its turn might redeﬁne or create new aliases.
As a result, our symbol table keeps a stack of scopes to properly resolve aliases in nested contexts.
We constrain numeric literals to either come from a set of common numbers (including 0 and 1 ) or from the user’s natural language question.
This prevents the model from copying arbitrary numbers from the few-shot examples in the prompt.
Finally, since SQL is case-insensitive, our CE returns case-insensitive regular expressions for keywords, table and column names.
C. V EGA -L ITE
Vega-Lite is a declarative language for specifying data visualizations given a data frame – a table where rows represent data points and columns represent attributes of various data types.
Its syntax is a subset of JSON.
Therefore, our Vega-Lite grammar accepts JSON objects that follow a subset of the Vega-Lite schema.
As in SQL, we use the user’s data frame to constrain valid ﬁeld names.
Additionally, in Vega-Lite, one must also specify a Vega-Lite type that is used to interpret the ﬁeld values.
We inspect the runtime values in the data frame to determine compatible Vega-Lite types.
For example, a string column is typically used as a categorical value ( nominal , in Vega-Lite).
However, if its entries have ISOformatted timestamps, it can be used as temporal , or quantitative if its values can be parsed as numbers.
Since JSON objects are unordered, it must handle two valid output orders: the model might output the ﬁeld name ﬁrst (we thus later constrain the type) or the data type ﬁrst (we then constrain the ﬁeld name to compatible columns).
Because Vega-Lite’s behavior is to silently ignore invalid data points, associating a column with an incompatible type simply produces an empty chart.
Our constraint completely prevents this class of errors.
We forbid repeated ﬁelds and limit the length of free-form string literals to 30 characters, which prevents the common failure mode of language models to enter repetition loops.
Similarly, we only allow an aggregation in one of the X-Y axes, but not both, since that collapses the chart to a single data point and is another common failure case.
Finally, we prevent the model from faceting (i.e., splitting into multiple charts) based on a column with too many ( > 50 ) distinct values: that typically crashes the Vega-Lite rendering engine since it allocates an overly large output image.
In summary, our constraints guarantee conformance to the Vega-Lite speciﬁcation and additionally avoid common mistakes that cause crashes or degenerate outputs.
C. SMC AL F LOW
SMCalFlow programs express responses to user queries about calendar events, weather, places, and people (Andreas et al., 2020).
It is a rich language with scoped variable declarations, generic types, polymorphic operators and a large API of over 400 functions, which can be composed to express complex actions like “cancel any meetings on the same day of my next doctor’s appointment”, or answer queries such as ”will it be raining during my next walking meeting with Frank?”
Our CE enforces that programs type-check 2 by construction.
An example 3 is given in Figure 1.
At the current point in inference, the model is producing an argument to size .
Since this function takes a list, its argument must be the return value of a function with return type List<T> .
Among all functions and methods in the SMCalFlow API, only 14 return lists, which severely limits the valid options for the callee.
Similarly, we keep track of declared variables inside let expressions, together with their types (inferred from their initialization expression), and use that data structure to limit options whenever a token of type identifier is syntactically allowed to follow.
Finally, we implemented heuristics based on user utterance patterns that avoid common failure cases we observed in GPT-3.
These tend to happen when the model blindly copies portions of the examples in the prompt without adaptation.
For instance, whenever the utterance contains exactly one of “a.m.” or “p.m”, this is usually represented in SMCalFlow by a call to a corresponding SMCalFlow function that constructs a Time object (e.g., NumberAM(5) ).
However, if the examples retrieved for the prompt tend to only have times in the opposite half of the day, GPT-3 might call the wrong function, and translate the time sub-expression in “Schedule it for 5pm” into NumberAM(5) .
To avoid this, if we detect exactly one of “a.m.” or “p.m.” in the utterance, we remove the time construction functions associated with the opposite pattern from the candidates.
We do the same ﬁltering with days of the week and months, which are also constructed by speciﬁc functions.
In all domains, the CE abstraction allows us to easily encode domain knowledge in a modular fashion.
Besides constraints coming from the language’s semantics, it further allows domain experts to analyze failure modes of the language model and to implement ﬁxes them in a modular and predictable manner.
D T HE TOKEN MISALIGNMENT CHALLENGE
The main challenge of CSD is in aligning constraints expressed in terms of programming language tokens with the model’s output, which happens in another token vocabulary (typically constructed with Byte-Pair Encoding).
Figure 7 shows an illustration of this challenge in Vega-Lite.
Arbitrary mismatches can occur: the ﬁrst BPE token includes the ﬁrst Vega-Lite token and the ﬁrst character of the second.
In the middle of the example, the "encoding" token in Vega-Lite spans 4 BPE tokens, with misaligned boundaries at the beginning and end.
Nonetheless, this issue is seamlessly handled by the CSD algorithm.
E O PTIMIZATIONS TO CSD
The algorithm described in Section 3.3 tests each token from the language model’s vocabulary individually.
However, many BPE tokens are preﬁxes of oen another, which lets us apply a signiﬁcant optimization.
If t 1 is a preﬁx of t 2 and t 1 is inadmissible after a partial program p , then t 2 is also inadmissible.
Thus, we test tokens by order of length, and keep rejected tokens in a Trie structure.
Before we test a token against the completion engine, we check whether one of its preﬁxes was already rejected.
If so, we can safely skip that token.
[Figure]
Another optimization consists in memoizing the enumerated completion points.
When testing a new partial program p , instead of starting from the empty string, we can start from the longest known completion point that is a preﬁx of p .
This, again, can be efﬁciently done by keeping completion points in a Trie.
F CSD WITHOUT DIRECT ACCESS TO THE LANGUAGE MODEL
We used the public OpenAI API to access GPT-3 and Codex.
Therefore, we did not have direct access to the underlying language models.
Even so, CSD can still be applied provided we can pass a bias to be added to the logits, which is available in the OpenAI API.
However, making one request at each token is too slow.
Instead, we apply a “rejection”-based sampling, as follows.
First, we request a complete program from the model.
Then, we iterate tokenby-token, validating it with the CSD algorithm against the completion engine.
If we ﬁnd a violation, we (a) use CSD to determine all valid next tokens, (b) make a request asking from just a single token, applying a logit bias to constrain it to the valid tokens, and then (c) continue generation after appending the new token.
Most (90%+) trajectories end after at most 3 corrections to the model.
In degenerate cases, we might need to correct the model after almost every token.
In our experiments, we capped CSD to apply at most 15 corrections, to control the time a request with CSD takes.
This only happened in less than .
5% of the cases, and could be completely avoided if we had direct access to the model (in which case CSD is efﬁcient enough to be applied at every token).
G A NALYSIS OF ACCURACY AND VALIDITY BY LENGTH IN SQL
Figure 8 shows the equivalent of Figure 4 for the SQL domain.
We notice that the largest gaps in validity happen for the longest queries, and the beneﬁts in accuracy are highest for queries around the 60% length percentile.
H TST F INE - TUNING DETAILS
In Section 2, we described TST, which ﬁne-tunes a sentence embedding model to attempt to capture program similarity.
Here, we give more details on our training procedure.
First, for both SQL and SMCalFlow, we selected a random set of 2000 examples, and computed the normalized tree edit distance between all pairs of examples.
We then ﬁne-tuned S-BERT for one epoch using the cosine similarity loss.
We used the AdamW optimizer with a learning rate of 2 × 10 − 5 – the default parameters in the S-BERT library 4 .
We did not perform any hyperparameter tuning.
Training took around 3 hours on a single GPU.
I C OMMENTED EXAMPLES OF TST AND CSD
Here, we provide three more detailed examples.
The ﬁrst examples shows the beneﬁt of using TST on top of Codex.
The remaining two show value of CSD on top of Codex.
Since the examples in
Sections 2 and 3 used SQL, we now give examples in SMCalFlow and Vega-Lite for diversity.
I. TST EXAMPLE IN SMC AL F LOW
Query: “Hi, I need to meet with Alex on this Sunday at 6pm at the CVS shopping center for an hour.”
This query asks the system to create an event.
Without TST, the top example (similar to the query) returned by vanilla S-BERT from the training set is:
User utterance: Hi, I need to meet with Alex tomorrow afternoon for an hour.
SMCalFlow: (Yield (CreateCommitEventWrapper (CreatePreflightEventWrapper (& (& (Event.start_?(DateTimeConstraint (Afternoon) (Tomorrow)))
(Event.duration_?
= (toHours 1))))
(Event.attendees_?
(AttendeeListHasRecipient (Execute (refer (extensionConstraint (RecipientWithNameLike (EmptyStructConstraint) (PersonName.apply "Alex")))))))))))
This example also involves creating an event, and also with a person named “Alex”.
However, the time structure of this example, “Tomorrow afternoon”, is not trivial to translate to “this Sunday at 6pm”, because these two speciﬁcations require different APIs from SMCalFlow.
Given this andmore examples, Codex generates: (Yield (CreateCommitEventWrapper (CreatePreflightEventWrapper (& (& (Event.start_?
(DateTimeConstraint (Afternoon) (ThisSunday)))
(Event.duration_?
=
(toHours Event
This prediction, while seeming reasonable at a surface level, violates the SMCalFlow APIs in multiple places.
First, the function “ThisSunday” does not exist.
The parameter 6pm was lost in this prediction, since the examples had no explicit time in it and Codex could not identify where to put it.
Besides that, the model is trying to specify four attributes of the created event:
start time, duration, location and attendees, but it would need one more conjunction (call to “&”) to do so.
This error likely came from the fact that the top example only speciﬁed three attributes.
The calls to the location APIs are also incorrect – “LocationWithNameLike” exists (and is present in other examples in the prompt), but is not really needed or appropriate in this context.
However, if we use TST to pick the examples for the prompt, Codex is able to solve the problem.
With TST, the top retrieved example is:
User utterance: Schedule 2 hours on Friday at 7:00 am for a Meeting with Kayla at the Sear’s Tower.
SMCalFlow: (Yield (CreateCommitEventWrapper (CreatePreflightEventWrapper (& (& (& (Event.start_?(?=(DateAtTimeWithDefaults (NextDOW(Friday)) (NumberAM7L))))
(Event.duration_? = (toHours 2)))) (Event.location_?
(?=(LocationKeyphrase.apply "Sear\’s Tower"))))
(Event.attendees_?
(AttendeeListHasRecipient (Execute(refer(extensionConstraint (RecipientWithNameLike(EmptyStructConstraint) (PersonName.apply"Kayla")))))))))))
Note that the utterance, from the perspective of its natural language structure, is less similar to the user’s request than the previous example.
However, it describes a query with the exact necessary structure: it schedules something for the next Friday (trivial to modify to Sunday) at a certain time of the day (“NumberAM” can be simply rewritten to “NumberPM”, which Codex does correctly).
The location attribute, which Codex had previously tried to adapt from a less related example, is now also speciﬁed here.
With this example, Codex is able to make a correct adaptation.
It predicts: (Yield (CreateCommitEventWrapper (CreatePreflightEventWrapper (& (& (& (Event.start_?(?=(DateAtTimeWithDefaults (NextDOW
This matches the reference program from the dataset.
I. CSD EXAMPLES IN V EGA -L ITE
Query: “average horsepower over years, colored by origin”
This query does not specify a chart type, but one could infer that a line chart is likely more appropriate since it is showing a quantity over time.
The dataset under consideration contains car models with their year of manufacture, horsepower, number of cylinders and many other attribute.
Thus, to obtain the “average horsepower” in a given year, we need an aggregation in Vega-Lite.
Codex alone predicts: {"mark": {"type": "line" }, "encoding":
[Equation]
[Equation]
[Equation]
[Equation]
"Year"
[Equation]
All the right components of the query are there.
However, even though the query asks for the “average horsepower”, the aggregation function in Vega-Lite is called “mean”.
Specifying an invalid aggregation function makes Vega-Lite raise an exception.
Thus, an actual user wouldn’t see any plot from this response.
CSD prevents Codex from generating “average” after it outputs the ﬁeld name “aggregate”, forcing it to choose one of the existing aggregation functions (max, min, mean, etc).
Constrained by CSD,
Codex generates: {"mark": {"type": "line" }
[Equation]
"encoding":
[Equation] [Equation]
[Equation] [Equation]
"Year
[Equation] [Equation]
[Equation]
yearmonth
[Equation]
[Equation] [Equation]
[Equation]
[Equation]
Origin [Equation] nominal
[Equation]
This is not only valid but also matches the ground-truth from the dataset. “Create a scatterplot of acceleration vs horsepower by model” This utterance asks for a chart that relates each car’s acceleration and horsepower. Codex generates:
[Equation]
[Equation]
[Equation] [Equation] [Equation] "point"
[Equation]
, "encoding":
[Equation] [Equation] [Equation]
"quantitative"
[Equation]
"y"
[Equation]
[Equation]
"Acceleration [Equation] [Equation] "quantitative"
[Equation] [Equation] [Equation] [Equation]
[Equation]
[Equation]
[Equation]
The produced plot, shown in Figure 9 (left) is in principle correct, and answers the user’s question.
The “by model” in the query might sound as if the user wants specialized plots by model.
Codex realizes this interpretation by coloring each point differently by the car’s model.
But in this dataset, each row represents a different car with a distinct model.
Thus, there are 303 unique values of the “Model” column, and Vega-Lite raises a run-time warning saying that the legend has overﬂowed for having too many values, so some values (the majority, in this case) were omitted.
When the model decides to facet on a column with too many values, the Vega-Lite run-time often runs out of memory for allocating an overly large image.
CSD avoids these errors by looking at the user’s dataset and determining which columns can be used for coloring or faceting, based on how many distinct values they have.
If none are available, then
CSD doesn’t let the model use these features.
In this example, CSD forces Codex to choose one of the columns with less than 30 distinct values.
Between those, Codex generates “Year”.
The resulting plot is shown in Figure 9 (right).
While in this example the “Year” column was not mentioned in the user’s request, the application of this constraint here steers Codex to avoid a run-time warning or error.
The ground-truth visualization for this description does not have any color speciﬁcation, so neither prediction counts as an exact match.
Nevertheless, from the point of view of the user’s experience, CSD can be helpful even when it cannot completely ﬁx the prediction.