PET-SQL: A Prompt-Enhanced Two-Round Refinement of Text-to-SQL with Cross-consistency (2024)

3.1 Pre-processing: SQL-Tailored Prompting

When instructing LLMs to generate SQL queries, the usage of styles or templates for prompts significantly impacts LLMs’ performance. As [4] recommended, Code Representation ($CR_{p}$) and OpenAI Demonstration ($OD_{p}$) forms mixed with additional information are better choices (shown in Listing 1).To further exploit the potential of LLM, we enrich the prompt based on OpenAI Demonstration ($OD_{p}$), and we name our proposed as Reference-Enhanced representation ($RE_{p}$).

Building upon $OD_{p}$ consisting of the task instruction, database schema, and question components, our proposed $RE_{p}$ performs three modifications: optimization rule (OR), cell value references (CV), and foreign key declarations (FK), which all turn out effective.

Optimization Rule (OR): Concretely, in the task instruction, we raise a multi-task constraint rule for LLM, that is, highlighting the LLM to “minimize SQL execution time while ensuring correctness” (line 1 in Listing 2). The rationale is not only to focus on execution accuracy but also to make LLM aware of the efficiency of SQL statements. During the generation of efficient SQL statements by LLM, redundant characters and operators can be avoided, which often result in exceptions. As our experiment shows, such an optimization rule can boost the performance by as high as +2.3%.

Cell Value References (CV): As mentioned before, to avoid mistakes from wrong-formatted cell value, three rows in each table are randomly sampled and inserted in the prompt (lines 7-11 in Listing 2).The sampled cell values as references can help LLM understand the database formats and specifications. It relieves the dilemma of being unsure which element to query as the condition caused by the lack of standardization in filling out different tables. Without prior knowledge of cell values, organizing condition statements to query the number of men in a table can be confusing. Such a cell value reference can boost performance by +4.0%.

Foreign Key Declarations (FK): Thirdly, foreign key relations in the schema are added as suffixes in the prompt (lines 12-15 in Listing 2). It facilitates LLM in recognizing the connections between the tables involved in the database, thus improving the understanding of the user’s intent and automatically selecting the appropriate connection in the queries. Such foreign key can boost the performance by as high as +7.5%.

The prompt above is in a zero-shot setting and can be further enhanced by the few-shot approach. We will explain how to select demonstrations as shots, equip them with the prompt, and stimulate the LLM to yield the preliminary SQL (PreSQL) for a question. Such a PreSQL is a critical design for our schema linking with higher accuracy.

3.2 Pre-processing: To Generate A Preliminary SQL (PreSQL) for Schema Linking

To eschew the verbose information in the schema that may obstruct the LLMs’ performance,we further conduct the schema linking (SL), which identifies and narrows down the table and column references that are related to the database schema and condition values related to the natural language question.Schema link has been proven to boost the accuracy, enhance generalizability across domains, and facilitate the synthesis of complex queries [5; 3].

However, existing methods directly employed LLMs to “guess” the schema, and such a “guess” can be risky since LLMs, even GPT-4, have NOT been pre-trained or supervised and fine-tuned with such a domain-specific schema linking task. Though not being trained with schema linking, most LLMs, luckily, are trained with code generation tasks, such as SQL generation. So our intuitive yet very innovative solution is: can we solve the schema linking task (what LLM is not good at) by letting LLM generate SQL code (what LLM is good at)? Thus, rather than instructing the LLMs to link schema, we instruct LLMs to generate a preliminary SQL (PreSQL), and then the table and column entities mentioned in the PreSQL are parsed as the linking results. To ensure more accurate schema linking, we incorporate a few steps to generate the PreSQL:

Step 1 - Question De-semanticization[27]: the domain-related tokens (i.e., table names ‘singer’, column names ‘gender’, and values ‘male’) in questions are masked with a special token <mask> according to the database schema, and thus obtain the question skeletons that only exhibit the question’s intentions.

Step 2 - Demo Retrieval based on Similarity: Then, all the question skeletons and question-SQL pairs in the training set are constructed as a demonstration pool.Based on the semantic embeddings of question skeletons, we retrieve the top-$K$ similar samples from the pool with the target question. Like [4], the embedding model for question skeletons is a pre-trained sentence Transformer.

Step 3 - Few-shot demos + SQL-Tailored Prompt: After that, we organize the selected demonstrations with our prompt $RE_{p}$ following [4], that is, prefixing question-SQL pairs as the few-shot contexts in $RE_{p}$ (as shown in lines 1-10 in Listing 3). Finally, the few-shot $RE_{p}$ is used to prompt the LLM and generate PreSQL.

Step 4 - Schema Parsing: Then, the SL is implemented based on a simple principle:The schema information mentioned in the PreSQL is concise and relevant to the question, and thus, the table/column entities in the statement can be directly parsed as the linking results.

Rather than designing strategies such as sorting [5] or extraction [3] to make LLM output relevant database references, it is better to generate complete PreSQL statements directly. It is also proven in Table 4, via such a PreSQL pipeline, our execution accuracy can be boosted by as high as +12.6%, compared with the traditional methods that rely on LLM to directly link schema.

There are three advantages: (1) Most LLMs are pre-trained with Text2SQL or SQL-related corpus, so presumably, their ability on SQL generation is stronger than the schema linking output (e.g., a list with table.column format) following instructions. (2) Generating SQL is more generic and applicable for coding-specific LLMs such as CodeLlama [28]. (3) The execution accuracy is theoretically monotonic. For a correctly executed PreSQL, the linked results are always correct. Therefore, the resulting final SQL (FinSQL) should also be correct. Wrong SQL may be attributed to the LLM being confused by the verbose prompt, and there is room to be correct with a concise and simplified prompt. It should be noted that column linking has poor fault tolerance and may be inapplicable to weak LLMs, so we only perform table linking at this stage.

3.3 SQL Generation: To Generate the Final SQL (FinSQL)

The Final Prompt: After the schema linking, the mentioned tables and columns in PreSQL are utilized to reorganize and simplify the $RE_{p}$ prompt. Specifically, all the contexts that are irrelevant to the linked tables or columns in the prompt are removed adaptively, including the schema properties, database references, and foreign key declarations: e.g., in Listing 2, if only the table ‘singer’ is linked and the table ‘singer_in_concert’ not, then all the lines related to the table ‘singer_in_concert’ (Line 5, 10-14) can be pruned, as shown in Listing 4 in Appx. A.3. With the parsed schemas, the prompt is much simplified. According to our statistics, our tokens in prompt can be saved as much as 31% after the pruning of PreSQL-schema linking.

SQL Generation: To summarize, we first feed full schema information (all the tables and columns) following the Listing 3 into a strong LLM (e.g., GPT-4) to obtain a PreSQL, and parse the PreSQL for schema linking.Then, the prompt is simplified by pruning irrelevant schemas and used to instruct one LLM or multiple LLMs to generate the FinSQL. We will introduce how to orchestra multi-LLMs.

3.4 Post-calibration: Cross-Consistency Is Better than Self-Consistency

As a canonical trick, the self-consistency approach is widely employed for post-refining the executed results, which is conducted by increasing the randomness of an LLM output at high temperatures to generate diverse results and followed by majority voting among multiple executed results. However, it has been reported that high temperatures may hurt the performance of LLMs (increasing model hallucinations) [29], especially for deterministic tasks (e.g., coding). Besides, the diversity of a single model is usually limited.Instead of conducting self-consistency, we propose instructing several LLMs under lower temperatures to generate stable and high-quality SQLs and then casting votes across the SQLs’ executed results.Such a cross-consistency strategy can diversify the SQL queries and maintain LLMs’ performance as low temperatures are set up.

In terms of voting, we propose two feasible implementation strategies for cross-consistency: naive voting across several LLMs and difficulty-aware voting according to PreSQL complexity.

Naive voting across several LLMs: Then, the full schema prompt is simplified and used to instruct several LLMs, each of which generates a FinSQL with a low temperature.All the FinSQLs and the re-used PreSQL are executed with the SQL database engine, and the queried results are obtained.Finally, we check the results of each SQL execution and take the majority of the results as the final answer. The FinSQL will be selected randomly, whichever yields the final answer.

Difficulty-aware voting according to PreSQL complexity: In addition to the naive voting strategy, we further refine the voting rules based on the difficulty of PreSQL.As different LLMs specialize in questions with different complexity grades, putting them together in a voting pool can produce biased results.The complexity of questions is classified into four grades (i.e., easy, medium, hard, and extra) by the abstract syntax trees of PreSQL according to [30], and each grade of questions is solved by a distinct set of candidate LLMs for voting.With such fine-grained voting, we can maximize the potential of LLMs and significantly mitigate the voting bias.

4.1 Experimental Setup

Dataset and Metrics: We evaluate our PET-SQL in Spider [8] benchmark, which is a large-scale cross-domain Text2SQL dataset. It contains 8659 instances in training split and 1034 instances in development split over 200 databases, with non-overlapping databases in each set, and 2147 instances are holdout as the test set across another 34 databases. The evaluation of Text2SQL performance of methods is conducted by execution accuracy (EX), following the official test-suit¹¹1https://github.com/taoyds/test-suite-sql-eval defined in [30].Additionally, we also conduct experiments on the Bird-SQL [22] dataset to evaluate the effectiveness and generalization of our methods, with the results reported in Appendix.EX measures the proportion of questions in the evaluation set where the execution results of both the predicted and ground-truth inquiries are identical, relative to the total number of queries.

Evaluated LLMs: Five LLMs are used to validate the superiority of our PET-SQL, and we report the best results of PET-SQL on the test set, which can surpass rank 2 with about 1% improvement on the Spider leaderboard.The LLMs are: CodeLlama-34B [28], SQLCoder-34B [31], InternLM-70B [32], SenseChat-70B [33], and GPT-4 (version on 2023.06.13) [34].The first two are specific to coding, especially on SQL, while the others are generic LLMs. Implementation is detailed in Appx. A.1.

4.2 Overall Performance

4.3 A Deep Dive into Prompt Design

Our few-shot $RE_{p}$ prompt is better than code representation and OpenAI demonstration: We compare our few-shot $RE_{p}$ with $CR_{p}$ and $OD_{p}$, which are recommended by [4]. It should be noted that the selected question-SQL demonstrations in them are the same. They are evaluated on the Spider development and test datasets separately, with the same selected question-SQL demonstrations in them. As shown in Table 11 in Appx. A.6, we recap the conclusions here: our proposed prompt $RE_{p}$ demonstrates superior performance under the zero-shot setting across three different datasets and three different foundation LLMs, with roughly 1% to 7% improvements compared to $CR_{p}$. This indicates the versatility of $RE_{p}$, whether for coding-specific or generic LLMs.

Ablation of $RE_{p}$: We further scrutinize the effects of three modifications in $RE_{p}$ compared to $OD_{p}$: optimization rule (OR), cell values reference (CV), and foreign key declarations (FK) by ablation studies under the zero-shot setting. As shown in Table 3,(1) Our cell values reference is the most important: Discarding it in InternLM and SQLCoder leads to -2.4% and -6.4% decreases in EX. This indicates that ensuring the presence of sampled cell values is crucial for LLM generation, and with these references, LLM can better navigate the database and formulate accurate queries. This reduces ambiguity and improves the efficiency of the generated SQL queries. (More results can be referred to in Table 10 in Appx. A.5)(2) Optimization rule also always helps for all three models: Dropping OR also shows a consistent impact on all three models’ performances (-3%). (3) Foreign key instead has two-sided impact: In CodeLlama, it boosts the performance to the most extent by +7.5%, but in InternLM, not using FK is surprisingly better. To conclude, the CV has the greatest significance, followed by the OR and FK components.

4.4 A Deep Dive into Schema Linking

The linked schema is parsed from the PreSQL, which is generated by GPT-4 with full schema information. To validate its effect, we introduce the table recall metrics $R_{e}$ and $R_{s}$ (higher is better, defined in Appx. A.2), and compare the EX of FinSQLs generated by the LLMs based on simplified and unsimplified schemas separately.

To evaluate the superiority of the proposed PreSQL-based SL, in Table 4, we compare it with the SLunder demo-fixed prompt, which is used in DIN-SQL. It becomes evident that (1) the proposed SLcan significantly extract useful information. (2) It also significantly shortens the prompts. By using the linked schema, we can cut the average token length of prompt from 924.2 to 630.8, by more than 33%.

With the linked tables, we further simplify the prompt and then evaluate the performance of LLMs on the test set with the 9-shot setting. The results are demonstrated in Table 5.(1) The simplified prompt has improved performance across all LLMs achieving 1%$\sim$6% improvements, especially for CodeLlama, which validates the effectiveness of our PreSQL-based approach. (2) Smaller LLMs tend to benefit significantly from SL while larger LLMs gain slightly. This is likely due to the inherent differences in their capacities and parameters for retrieving useful information from raw text.In conclusion, SL is an effective technique that enhances the performance and versatility of LLMs by summarizing structured information.

4.5 A Deep Dive into Cross-consistency

In Table 6, we report the performance of cross-consistency (CC) on dev and test sets.Considering the token cost, we give up the SQL generation by GPT-4 on the dev set, and the best result from the remaining LLMs is chosen as the baseline (i.e., CodeLlama). For comparison, we also perform self-consistency with CodeLlama (i.e., calling it 5 times, $\textit{temperature}=0.5$).

It can be observed that,(1) Cross-consistency is much better than self-consistency due to diversity.On the dev set, the gain of self-consistency is trivial (only 0.8%), while the cross-consistency with naive voting achieves 6.9% improvements. On the test set, the naive voting of CC can already improve +1.2%, while In DAIL-SQL [4], self-consistency (five SQLs generated by GPT-4) only achieves +0.4%. Thus, bringing diverse LLMs can better boost the accuracy since different LLMs may make different errors, and letting them vote mitigates each other’s mistakes. (2) Difficulty-aware voting is a better implementation of CC. Compared to naive voting, the difficulty-aware voting strategy can further exploit the potential of LLMs and significantly mitigate the voting bias, improving the performance from 86.7% to 87.6% on the test set.

For questions in different difficulty levels, a tailored set of multi-LLMs can further optimize the accuracy. Table 7 gives the details of our recommended set of LLMs for each difficulty level and their EX on the test setusing the PreSQL only. (1) Smaller LLMs are sufficient to handle the easy questions: For easy questions, smaller LLMs can achieve 93.2% EX, demonstrating their efficacy in handling less complex questions. (2) For difficult questions, there is a clear need for more powerful LLM (i.e., GPT-4) to maintain high performance. This performance disparity underscores the importance of selecting the appropriate LLMs based on the difficulty of the questions.

CC as a plug-and-play can always boost performance: CC can serve as a plug-and-play module for other SOTA methods. We evaluate its performance with DAIL-SQL on the Spider-test set as shown in Table 8.With the introduction of CC, the DAIL-SQL can be further enhanced by +2% improvements. This suggests that the CC strategy is still effective on other methods. CodeLlama is used as the single LLM setting, and then SQLCoder and InternLM are added as the CC setting.

More Experiments in Appendix: We also provide more experimental results to evaluate the superiority or effectiveness of our proposed methods.Conducted on Bird-SQL dataset [22]: (1) the comparison between our PET-SQL and other SOTA methods (that is, DIN-SQL [3] and DAIL-SQL [4]) under different foundation LLMs are shown in Appx. A.4; (2) The ablation results of the proposed prompt $RE_{p}$ are further reported in Appx. A.5. Furthermore, we compare $RE_{p}$ with two well-defined prompts from DAIL-SQL [4] (i.e., $CR_{P}$ and $OD_{p}$) under the zero-shot setting, as detailed in Appx. A.6. In Appx. A.7, we also analyze the SL results by the PreSQL from GPT-4. Finally, the significance of GPT-4 in CC is underscored in Appx. A.8.

A.1 Implementation details

To decrease the randomness of LLMs output, the temperatures are set extremely low ($10^{-7}$ for all the non-OpenAI LLMs and 0 for GPT-4). The max lengths of input and output tokens are 4096 and 200, respectively.In the question skeleton-based demonstration retrieval, we traverse all the database schema and check if the table or column names appear in questions. The matched entities will be masked by <mask> token. Based on the similarities between the sentence embedding of the target and candidate question skeletons, top-9 similar question-SQL pairs in the demonstration pool are selected and then constructed as the in-context. For cross-consistency, the executed results from the above five LLMs and the PreSQL are voting with equal weights.

A.2 The definition of table recall metrics

A.3 The illustration of the simplified prompt with the schema linking results

As shown in Listing 4, the initial prompt includes several entities (e.g., “singer_in_concert”) that are not directly relevant to the primary question at hand, where the strikethrough (in lines 5, 10-14) indicates removing contexts in $RE_{p}$ via schema linking. By applying schema linking, we can streamline the prompt by eliminating unnecessary information and focusing on the critical elements that directly support the question “How many singers do we have?”.

A.4 The validation on the Bird-SQL dataset

Additionally, we have also conducted experiments on the bird develop (bird-dev) dataset [36], recognized for its heightened complexity. Table 9 presents the outcomes of SOTA methods with various foundational LLMs. The DAIL-SQL method stands out, exhibiting superior performance across all evaluated metrics on the Bird-dev dataset. Our proposed PET-SQL method performs admirably, particularly excelling in the SQLCoder with a score of 0.391.

A.5 The ablation of $RE_{p}$ on Bird-dev

To verify the effectiveness and generalization of the proposed $RE_{p}$, we also conduct the ablation study on the Bird-dev dataset. The results are shown in Table 10.The conclusion is the same as the one summarized in Table 3:The performance of all three LLMs (CodeLlama-34B, InternLM-70B, and SQLCoder-34B) decreases when any of the components (OR, CV, FK) are removed, which demonstrates that each component of $RE_{p}$ contributes to the LLMs’ ability to generalize and perform well on the Bird-dev dataset.

A.6 Further analysis of the proposed prompt $RE_{p}$

To further illustrate the effectiveness of the proposed prompt, we compared it with two well-defined prompts from DAIL-SQL [4], namely $CR_{p}$ and $OD_{p}$, under the zero-shot setting. The results are presented in Table 11. It is evident that even with different foundation LLMs, our proposed $RE_{p}$ remains effective, achieving abouta 7% to 10% relative performance boost across various datasets compared to the baseline prompts. This improvement signifies the robustness and generalizability of our proposed prompt $RE_{p}$ in enhancing model performance under the 0-shot setting.

A.7 Further analysis of the PreSQL-based schema linking

Table 12 shows the schema linking (SL for short) recall of GPT-4 (i.e., PreSQL).It can be seen that the recall of table linking reaches 94% while the EX of PreSQL is only 85.2%, and the limitations of LLMs may cause the gap. Prompt simplification with linked tables enables LLM to correct the wrong PreSQL cases.

In Table 13, we further summarize the average tables mentioned in the prompt and the ground truth (i.e., the average of GT tables). In the full schema (without SL), the average number of tables is 4.89, while in the simplified schema (with SL), it is 1.60 (close to GT). This means that SL significantly simplifies useless table information in prompts.

A.8 Further analysis of the cross consistency

GPT-4’s pivotal role is uncovered in enhancing the final outcomes. We conduct experiments to gauge its significance. Table 14 demonstrates the importance of GPT-4 within the Spider-test set:Specifically, utilizing GPT-4 alone yielded an accuracy of 85.50%, showcasing its substantial impact. Conversely, excluding GPT-4 from the other four models used for CC led to a noticeable decline in performance, resulting in an accuracy of 79.80%. Replacing GPT-4 with a weaker model (InternLM2-20B) within the CC framework further exacerbated this decrement, yielding an accuracy of 79.08%. However, employing all five models within the CC framework resulted in a notable enhancement, achieving an accuracy of 87.60%. These findings underscore GPT-4’s indispensable contribution to the overall performance, illustrating its pivotal role in enhancing CC.