PICARD: Parsing Incrementally for Constrained Auto-Regressive Decoding from Language Models

Torsten Scholak, Nathan Schucher, Dzmitry Bahdanau

Introduction

While there have been many successes in applying large pre-trained language models to downstream tasks, our ability to control and constrain the output of these models is still very limited. Many enterprise applications are out of reach because they require a degree of rigour and exactitude that language models are not able to deliver yet. If the target is a formal language like SQL, then we would like the model to adhere exactly and provably to the SQL specification with all its lexical, grammatical, logical, and semantical constraints. Unfortunately, with pre-training alone, language models may not satisfy these correctness requirements.

For text-to-SQL translation, the most widespread solution to constrained decoding is to make invalid SQL unrepresentable. For a while now it has been possible to restrict auto-regressive decoding to only those token sequences that correctly parse to SQL abstract syntax trees (Yin and Neubig, 2018; Lin et al., 2019; Wang et al., 2020). More recently, semi-auto-regressive improvements to this parsing paradigm have been proposed (Rubin and Berant, 2021). However, while effective, these approaches have in common that they are achieved at the expense of using a custom vocabulary of special control tokens or a custom model architecture, or both. Unfortunately, this makes them incompatible with generic pre-trained language model decoders. A less invasive and more compatible approach is to not constrain the generation process, but instead to filter finalized beam hypotheses by validity (Suhr et al., 2020; Lin et al., 2020). Yet, such filtering is at the expense of a very large beam size.

We address the expenses of these approaches with a novel incremental parsing method for constrained decoding called Picard, which stands for "Parsing Incrementally for Constrained Auto-Regressive Decoding." Picard is compatible with any existing auto-regressive language model decoder and vocabulary—including, but not limited to, those of large pre-trained transformers—and it does not require very large beam sizes. Picard is entirely absent from pre-training or fine-tuning of the model, and can be easily and optionally enabled at inference time. Picard operates directly on the output of the language model which, in the case of text-to-SQL translation, is the readable surface form of the SQL code.

In our experiments, we find that Picard can significantly improve the performance of a large pre-trained language model (Raffel et al., 2020) after it is fine-tuned on the text-to-SQL task. On the Spider text-to-SQL dataset (Yu et al., 2018), we find that a T5-Base model with Picard can outperform a T5-Large model without it, and likewise for a T5-Large and a T5-3B model. Significantly, with the help of Picard, a T5-3B model can be raised to state-of-the-art performance on the Spider and CoSQL datasets (Yu et al., 2019).

The Picard Method

Picard warps model prediction scores and integrates trivially with existing algorithms for greedy and beam search used in auto-regressive decoding from language models. Its arguments are the token ids of the current hypothesis and, for each vocabulary token, the log-softmax scores predicted by the model’s language modeling head. Picard also has access to SQL schema information, in particular, information about the names of tables and columns and about which column resides in which table.

At each generation step, Picard first restricts prediction to the top-kk highest probability tokens and then assigns a score of −∞-\infty to those that fail Picard’s numerous checks (see Figure 2). These checks are enabled by fast incremental parsing O’Sullivan and Gamari (2021) based on monadic combinators Leijen and Meijer (2001). There are four Picard mode settings that control their comprehensiveness: off (no checking), lexing, parsing without guards, and parsing with guards—the highest mode. A prediction that passes a higher mode will always pass a lower mode but not necessarily vice versa.

In lexing mode, Picard checks the output on a lexical level only. It attempts to convert the partial, detokenized model output to a white-space delimited sequence of individual SQL keywords like select, punctuation like (), operators like + and -, literals like string and number values in SQL conditions, and identifiers like aliases, tables, and columns—without being sensitive to the order in which these lexical items appear. By making it so, Picard can detect spelling errors in keywords or reject table and column names that are invalid for the given SQL schema. For instance, consider the question "What are the email, cell phone and home phone of each professional?" from Spider’s development set on the dog_kennels database. Our fine-tuned T5-Large model predicts select email_address, cell_phone, home_phone from professionals while the ground truth selects cell_number instead of the invalid cell_phone column. This mistake is caught and avoided by Picard in lexing mode.

2 Parsing without Guards

In the lowest parsing mode above lexing—referred to as parsing without guards—Picard checks the output on a grammatical level. Picard attempts to parse the detokenized model output to a data structure that represents the abstract syntax tree (AST) of the predicted SQL query. Contrary to lexing mode, the order in which keywords and clauses appear now matters. Picard can reject invalid query structures, e.g. find missing from clauses or incorrect orders of clauses and keywords. It can also detect a range of issues with compositions of SQL expressions: Number one, if Picard matches on a tid.cid pattern, but the table with the id tid does not contain a column with id cid, then that parse is rejected. Secondly, if Picard first matches on an alias.cid pattern and then later matches on the tid as alias pattern but tid does not contain cid, then that parse is also rejected. An equivalent rule also exists for sub-queries bound to table aliases. Lastly, Picard prohibits duplicate binding of a table alias in the same select scope, but permits shadowing of aliases defined in a surrounding scope. This can happen in nested SQL queries.

3 Parsing with Guards

In its highest parsing mode, Picard engages in additional analyses—called guards—while assembling the SQL AST. If Picard matches on tid.cid or alias.cid, then guards require that the table tid or the alias alias, respectively, is eventually brought into scope by adding it to the from clause. Moreover, the alias alias is constrained to resolve to a table or a sub-query that has the column cid in it. If Picard matches on the pattern cid, then another guard requires that exactly one table is eventually brought into scope that contains a column with that id. These guards are enforced eagerly in order to fail fast and to eject invalid hypotheses from the beam at the earliest possible time. The first time this is happening is after parsing the from clause.

Only with these guards, Picard is able to reject a wrong prediction from our fine-tuned T5-Large model like select maker, model from car_makers for the question "What are the makers and models?" Here, the correct table to use would have been model_list, since it is the only one in Spider’s car_1 schema that contains both a maker and a model column.

Additional checks and guards are conceivable, for instance, checking that only expressions of the same type are compared or that column types selected by union, except, or intersect queries match. We leave these additional checks to future work.

Experiments

Our experiments are mainly focused on Spider (Yu et al., 2018), a large multi-domain and cross-database dataset for text-to-SQL parsing. We train on the 7,000 examples in the Spider training set and evaluate on Spider’s development set and its hidden test set. We also report results on the CoSQL SQL-grounded dialog state tracking task (Yu et al., 2019), where we predict a SQL query for each question given previous questions in an interaction context. For this task, we train on both the Spider text-to-SQL training data and the CoSQL dialog state tracking training data, and evaluate on the CoSQL development and test sets.

Spider and CoSQL are both zero-shot settings. There is no overlap between questions or databases between the respective training, development, and test sets.

On Spider, we determine model performance based on three metrics: exact-set-match accuracy, execution accuracy, and test-suite execution accuracy (Zhong et al., 2020). Exact-set-match accuracy compares the predicted and the ground-truth SQL query by parsing both into a normalized data structure. This comparison is not sensitive to literal query values and can decrease under semantic-preserving SQL query rewriting. Execution accuracy compares the results of executing the predicted and ground-truth SQL queries on the database contents shipped with the Spider dataset. This metric is sensitive to literal query values, but suffers from a high false positive rate (Zhong et al., 2020). Lastly, test-suite execution accuracy extends execution to multiple database instances per SQL schema. The contents of these instances are optimized to lower the number of false positives and to provide the best approximation of semantic accuracy.

On CoSQL, we measure model performance in terms of the question match accuracy and the interaction match accuracy. Both metrics are based on exact-set-match accuracy. Interaction match accuracy is the joint accuracy over all questions in an interaction.

We are encouraged by results by Shaw et al. (2021), who showed that a pre-trained T5-Base or T5-3B model can not only learn the text-to-SQL task, but also generalize to unseen databases, and even that T5-3B can be competitive with the then-state-of-the-art (Choi et al., 2021; Wang et al., 2020)—all without modifications to the model. We therefore use T5 as the baseline for all our experiments.

In order to allow for generalization to unseen databases, we encode the schema together with the questions. We use the same serialization scheme used by Shaw et al. (2021). In experiments using database content, we detect and attach the database values to the column names in a fashion similar to the BRIDGE model by Lin et al. (2020). When fine-tuning for the CoSQL dialog state tracking task, we append the previous questions in the interaction in reverse chronological order to the input. Inputs exceeding the 512-token limit of T5 are truncated. The target is the SQL from the Spider and/or CoSQL training sets, unmodified except for a conversion of keywords and identifiers to lower case. We fine-tune T5 for up to 30723072 epochs using Adafactor (Shazeer and Stern, 2018), a batch size of 20482048, and a learning rate of 10−410^{-4}.

Our findings on the Spider dataset are summarized in Table 1 and Figure 1. Our reproductions of Shaw et al. (2021)’s results with T5 cannot compete with the current state of the art on Spider. The issue is that these models predict a lot of invalid SQL. For instance, 12% of the SQL queries generated by the T5-3B model on Spider’s development set result in an execution error. However, when these same models are augmented with Picard, we find substantial improvements. First, invalid SQL predictions become rare. For T5-3B with Picard, only 2% of the predictions are unusable. In these cases, beam search exited without finding a valid SQL prediction. Second, and most significantly, by using Picard, the T5-3B model is lifted to state-of-the-art performance. We measure an exact-set-match accuracy of 75.5% on the development set and 71.9% on the test set. The execution accuracy results are 79.3% and 75.1%, respectively. These numbers are on par or higher than those of the closest competitor, LGESQL + ELECTRA (Cao et al., 2021) (see Table 1). Furthermore, we achieve a test-suite execution accuracy of 71.9% on Spider’s development set.

Our findings on the CoSQL dialog state tracking dataset (see Table 2) are similar to those for Spider. Picard significantly improves the performance, and our fine-tuned T5-3B model achieves state-of-the-art performance.

Picard is not only improving performance, it is also fast. During evaluation of the T5-3B model on Spider, the decoding speed with beam size 44 on an NVIDIA A100-SXM4-40GB GPU was, on average, 2.5 seconds per sample without Picard and 3.1 seconds per sample with Picard.

Beam Size

Figure 1 shows results on Spider without and with Picard when parsing with guards for different beam sizes and sizes of T5. For each model size, Picard increases performance with increasing beam size. These increases are the strongest for the step from beam size 11 to 22, less pronounced from 22 to 44, and then saturating for beam sizes above 44. Even with greedy search (beam size 11), Picard allows for some modest improvements. Note that, without Picard, these models do not benefit from beam search. The number, kk, of highest-probability tokens that are processed by Picard at each decoding step has a modest to negligible impact on performance. It is the largest for T5-Base, smaller for T5-Large, and almost undetectable for T5-3B. We do not study the case k=1k=1, because it reduces the beam search to constrained greedy search.

Ablations

In Figure 3, we have condensed our ablation analysis for Picard. We show results for our T5-Large model in all four Picard checking modes and for four different beam sizes on the Spider development set. When checking incrementally at each decoding step, lexing shows a small improvement over the unconstrained T5 model. The results without Picard and with Picard in lexing mode are largely independent of the beam size. This is different when Picard is switched into the more sophisticated parsing modes. Both, with and without guards, improvements from Picard increase rapidly for increasing beam sizes, where parsing with guards clearly has a strong lead over parsing without them.

In order to compare Picard with the filtering-by-validity approach of Suhr et al. (2020) and Lin et al. (2020), we have studied also what happens when Picard is only checking hypotheses when the model predicts their finalization with the end-of-sequence token.This is not exactly equivalent to filtering a completely finalized beam, because the hypotheses rejected by Picard never enter it and never take up any space. In this restrained mode, Picard is still effective, but much less so compared to normal incremental operation. The gap between these two modes of operation only begins to shrink for large beam sizes. This is understandable since Lin et al. (2020) used beam sizes of at least 1616 and up to 6464 to reach optimal results with filtering while Suhr et al. (2020) used a beam of size 100100.

Conclusion

We propose and evaluate a new method, Picard, for simple and effective constrained decoding with large pre-trained language models. On both, the Spider cross-domain and cross-database text-to-SQL dataset and the CoSQL SQL-grounded dialog state tracking dataset, we find that the Picard decoding method not only significantly improves the performance of fine-tuned but otherwise unmodified T5 models, it also lifts a T5-3B model to state-of-the-art results on the established exact-match and execution accuracy metrics.

Acknowledgements

We thank Lee Zamparo for his contributions to the experiments on the CoSQL dataset. Further, we would like to thank Pete Shaw for his input on the reproduction of the T5 results on Spider. We would also like to extend our gratitude to Tao Yu and Yusen Zhang for their efforts in evaluating our model on the test split of the Spider and CoSQL datasets. Finally, we thank our anonymous reviewers for their time and valuable suggestions.

References