2403.16204

# SQL-Encoder: Improving NL2SQL In-Context Learning Through a Context-Aware Encoder Abstract Detecting structural similarity between queries is essential for selecting examples in in-context learning models. However, assessing structural similarity based solely on the natural language expressions of queries, without considering SQL queries, presents a significant challenge. This paper explores the significance of this similarity metric and proposes a model for accurately estimating it. To achieve this, we leverage a dataset comprising 170k question pairs, meticulously curated to train a similarity prediction model. Our comprehensive evaluation demonstrates that the proposed model adeptly captures the structural similarity between questions, as evidenced by improvements in Kendall-Tau distance and precision@k metrics. Notably, our model outperforms strong competitive embedding models from OpenAI and Cohere. Furthermore, compared to these competitive models, our proposed encoder enhances the downstream performance of NL2SQL models in 1-shot in-context learning scenarios by 1-2% for GPT-3.5-turbo, 4-8% for CodeLlama-7B, and 2-3% for CodeLlama-13B. 1 Introduction Scaling language models (LMs) has impressively enhanced sample efficiency and performance across numerous natural language processing (NLP) tasks (Brown et al., 2020). A notable area of improvement is in translating natural language questions into SQL (NL2SQL) queries. In this domain, models leveraging in-context learning, which involves providing a few example question and SQL pairs in the prompt, have achieved state-of-the-art performance (Pourreza and Rafiei, 2023; Gao et al., 2023) on NL2SQL benchmarks (Yu et al., 2018; Li et al., 2023). In-context learning, a capability that has emerged as a feature of large language models (LLMs) (Wei et al., 2022), enables LLMs to adapt to specific tasks with a few examples. This paper focuses on enhancing the in-context learning performance of LLMs for the NL2SQL task. We propose a novel method for effectively measuring the similarity between task-specific demonstrations, alongside a model trained to predict this similarity score. In the NL2SQL domain, the majority of prior research (Liu et al., 2021; Nan et al., 2023; Guo et al., 2023; Gao et al., 2023) has concentrated on cross-domain in-context learning. This approach typically involves selecting few-shot examples from question and SQL pairs derived from databases different from the one associated with the question at inference time. We contend that relying on cross-domain example selection fails to unlock the full capabilities of in-context learning in NL2SQL domain. This is because examples from other databases may not effectively inform the generation of SQL queries for a given question, especially when the underlying tables and columns differ significantly. Table 1 showcases the one-shot performance of CodeLlama 13B (Roziere et al., 2023) on the BIRD Li et al. (2023) dataset development set under two scenarios: one utilizing cross-domain examples and the other employing in-domain examples, where the in-context learning samples are from the exact same database as the question at inference time. Selection was based on embedding all examples and the given question with the OpenAI embedding model, then using cosine similarity to identify the most relevant example. The findings indicate a significant 12% enhancement in SQL generation performance when utilizing a single in-domain example, in contrast to cross-domain examples that appear to diminish the efficacy of one-shot in-context learning. This observation is in line with Gao et al. (2023), which noted a reduction in performance of the GPT-3.5-turbo model under one-shot cross-domain conditions It should be noted that in cross-domain scenarios, prior research (Liu et al., 2021; Nan et al., 2023; Guo et al., 2023; Gao et al., 2023) has demonstrated that employing a few-shot approach with more than one sample can ultimately enhance performance, outperforming the zero-shot scenario.. This paper aims to explore in-domain in-context learning scenarios, seeking to enhance performance in a way that truly reflects the strengths of in-context learning. The prevalent strategy for in-context learning with LLMs involves deriving question embeddings through high-performing embedding models, such as the top embedding models on MTEB (Muennighoff et al., 2022) benchmark. This technique entails measuring cosine similarity or dot product among the embeddings to identify the most closely related questions (Guo et al., 2023; Liu et al., 2021). However, recent studies (Nan et al., 2023; Gao et al., 2023) have shown that relying solely on question similarities may not always lead to the selection of the highest performing in-context learning samples, in terms of the accuracy of the generated SQL queries. Instead, these investigations have considered SQL query skeletons as an additional layer of information for finding samples for in-context learning. A key question investigated in this paper is if the structure of SQL queries and their similarity to other queries can be accurately predicted based solely on questions, without the need to generate SQL queries. This paper explores the significance of such similarity metrics and proposes novel measures, such as schema-linking similarity and SQL skeleton similarity with tree edit distance, to enhance in-context learning. Upon identifying the most effective similarity metric, we created a large dataset of 170K question pairs and their respective similarity scores. This dataset was used to train a cross-encoder language model (Reimers and Gurevych, 2019), specifically designed to predict the similarity score between a pair of questions model is accessible in huggingface: https://huggingface.co/MrezaPRZ/sql-encoder. This cross-encoder model, the first of its kind tailored for the NL2SQL task, has shown in our experiments to provide superior signals for selecting few-shot samples, outperforming both OpenAI (Neelakantan et al., 2022; OpenAI, 2023) and Cohere’s embedding (Cohere, 2023) models for in-context learning task. Its exceptional performance, further validated on out-of-domain questions, surpasses that of the OpenAI and Cohere models on such samples, demonstrating our model’s generalizability. | Method Zero-shot In-domain one-shot | Execution accuracy 26.66 38.14 | | --- | --- | | Cross-domain one-shot | 20.99 | Table 1: Performance comparison of the in-domain and cross-domain in-context learning on the development set of BIRD <details> <summary>2403.16204v1/extracted/5492121/Methodology.png Details</summary>  ### Visual Description # Technical Document Extraction: Flowchart Analysis ## Overview The image depicts a **two-section flowchart** illustrating a database schema and query processing logic for singer data. The diagram uses **textual components** (questions, buttons, schema links) and **mathematical notation** to represent data relationships. --- ## Top Section: Query Processing Flowchart ### 1. Database Schema - **Label**: `DATABASE` - **Schema**: `DB SCHEMA CREATE SINGER(...)` - Represents a database table/structure for singer data. ### 2. Query Components #### Question 1 (`N1`) - **Text**: `WHAT ARE THE DIFFERENT COUNTRIES WITH SINGERS ABOVE AGE 20?` - **Flow**: - Connects to `Q1` (Start button) via a black arrow. - Leads to `S1` (Schema Link). #### Start Button 1 (`Q1`) - **Label**: `Start` (green button with play icon) - **Action**: `SELECT DISTINCT COUNTRY ...` #### Schema Link 1 (`S1`) - **Parameters**: `[COUNTRY, SINGER, ...]` - Represents a database query or schema path. #### Question 2 (`N2`) - **Text**: `SHOW ALL COUNTRIES AND THE NUMBER OF SINGERS IN EACH COUNTRY.` - **Flow**: - Connects to `Q2` (Start button) via a black arrow. - Leads to `S2` (Schema Link). #### Start Button 2 (`Q2`) - **Label**: `Start` (green button with play icon) - **Action**: `SELECT COUNTRY, COUNT(*) ...` #### Schema Link 2 (`S2`) - **Parameters**: `[SINGER, COUNTRY, ...]` - Represents a database query or schema path. --- ## Bottom Section: Database Schema & Similarity Functions ### 1. Database Schema - **Label**: `DATABASE` - **Schema**: `DB SCHEMA CREATE SINGER(...)` - Repeats the top-section schema definition. ### 2. Similarity Functions #### Function 1 (`SIM'(N1, N2, D)`) - **Parameters**: - `N1`, `N2`: Query identifiers (from `N1` and `N2`). - `D`: Database schema reference. #### Function 2 (`SIM((N1, Q1, S1), (N2, Q2, S2))`) - **Parameters**: - `(N1, Q1, S1)`: Tuple of Question 1 components. - `(N2, Q2, S2)`: Tuple of Question 2 components. - Represents a similarity calculation between query results. --- ## Key Observations 1. **Flow Direction**: - Top flowchart progresses from `N1` → `Q1` → `S1` and `N2` → `Q2` → `S2`. - Bottom section links queries to similarity functions via `SIM'`. 2. **Schema Consistency**: - Both sections reference the same `DB SCHEMA CREATE SINGER(...)`. 3. **Mathematical Notation**: - Uses `SIM'` and `SIM(...)` to denote similarity operations between query results. --- ## Missing Elements - No numerical data, charts, or axis labels present. - No legends or color-coded data series. - No explicit spatial coordinates (e.g., [x, y] placements). --- ## Conclusion The diagram outlines a **query-driven database workflow** for singer data, emphasizing schema design and similarity analysis between query results. All textual components are explicitly labeled, with no omitted information. </details> Figure 1: Overview of SQL-encoder framework for predicting the similarity between two questions $N_{1}$ and $N_{2}$ on a given database schema D, where SIM’ serves as a proxy for the actual similarity SIM between the questions $N_{1}$ and $N_{2}$ , their respective SQL queries $Q_{1}$ and $Q_{2}$ and schema links $S_{1}$ and $S_{2}$ . 2 Methodology In this paper, our primary objective is to identify the most effective similarity metric to enhance the performance of in-context learning. Specifically, we aim to estimate the similarity between two queries $Q_{1}$ and $Q_{2}$ given their natural language expressions $N_{1}$ and $N_{2}$ and a Schema $S$ . Our initial step involves identifying the function that yields the highest performance in in-context learning. However, during testing, instead of having access to $Q_{1}$ and $Q_{2}$ , we work with natural language questions $N_{1}$ and $N_{2}$ alongside Schema $S$ . Here, $Q_{1}$ and $Q_{2}$ represent SQL queries on Schema S that respectively answer $N_{1}$ and $N_{2}$ . We have constructed a dataset comprising pairs of queries and have labelled the proxy function sim( $N_{1}$ , $N_{2}$ , $S$ ) with the actual similarity sim(( $N_{1}$ , $Q_{1}$ , $S_{1}$ ),( $N_{2}$ , $Q_{2}$ , $S_{2}$ )) that we want to approximate, where $S_{1}$ and $S_{2}$ are schema links of $N_{1}$ and $N_{2}$ . This approach allows us to train a model to predict this proxy similarity measure. Figure 1 outlines the methodology adopted in this study to train the model using the optimal similarity metric. 2.1 NL2SQL Similarity Metric In previous studies (Nan et al., 2023; Gao et al., 2023), it has been shown that relying exclusively on question similarity to select in-context samples leads to sub-optimal performance. In our research, we aim to compare various NL2SQL-specific similarity metrics and elucidate their contributions towards enhancing in-context learning performance. Figure 2 demonstrates an example of using different similarity metrics to find the most similar Question/SQL pair for a given Question/SQL pair from the development set of BIRD dataset. <details> <summary>2403.16204v1/extracted/5492121/sql_similarity.png Details</summary>  ### Visual Description # Technical Document Extraction: Flowchart Analysis ## Overview The image depicts a **decision flowchart** for identifying the most similar sample based on three similarity criteria: 1. **Question Similarity** 2. **Skeleton Similarity** 3. **Schema Linking Similarity** 4. **Mean of All Three** The flowchart begins with a **GIVEN SAMPLE** and progresses through conditional logic to determine the optimal similarity match. --- ## Component Breakdown ### 1. GIVEN SAMPLE **Text Block**: ```sql What is the unabbreviated mailing address of the school with the highest FRPM count for K-12 students? SELECT T2.MailStreet FROM frpm AS T1 INNER JOIN schools AS T2 ON T1.CDSCode = T2.CDSCode ORDER BY T1.FRPM Count (K-12) DESC LIMIT 1 ``` --- ### 2. FIND MOST SIMILAR SAMPLE **Flow Arrows**: - **Question Similarity** - **Skeleton Similarity** - **Schema Linking Similarity** - **Mean of All Three** --- ### 3. QUESTION SIMILARITY **Question**: ```text What is the number of SAT test takers of schools with the highest FRPM count for K-12 students? ``` **SQL Query**: ```sql SELECT NumTstTkr FROM satscores WHERE cds = (SELECT cds FROM frpm ORDER BY FRPM Count (K-12) DESC LIMIT 1) ``` --- ### 4. SKELETON SIMILARITY **Question**: ```text Please list the phone numbers of the schools with the top 3 SAT numbers with the highest free or reduced price meal count. ``` **SQL Query**: ```sql SELECT T1.Phone FROM schools AS T1 INNER JOIN satscores AS T2 ON T1.CDSCode = T2.CDSCode ORDER BY T2.NumGE1500 CAST(T2.NumTstTkr DESC) LIMIT 3 ``` --- ### 5. SCHEMA LINKING SIMILARITY **Question**: ```text What is the eligible free or reduced price meal rate for the top 5 schools in grades 1-12 with the highest free or reduced price meal count of the schools with the ownership code 66? ``` **SQL Query**: ```sql SELECT CAST(T1.FRPM Count (K-12) / T1.Enrollment (K-12) AS REAL) FROM frpm AS T1 INNER JOIN schools AS T2 ON T1.CDSCode = T2.CDSCode WHERE T2.CDSCode = '66' ORDER BY T1.FRPM Count (K-12) DESC LIMIT 5 ``` --- ### 6. MEAN OF ALL THREE **Question**: ```text List the top five schools by descending order, from the highest to the lowest, the most number of Enrollment (Ages 5-17). Please give their NCEs school identification number. ``` **SQL Query**: ```sql SELECT T1.NCESSchool FROM schools AS T1 INNER JOIN frpm AS T2 ON T1.CDSCode = T2.CDSCode ORDER BY T2.Enrollment (Ages 5-17) DESC LIMIT 5 ``` --- ## Key Observations 1. **Flow Logic**: - The flowchart evaluates similarity through three criteria, culminating in a weighted mean. - Each step includes a **question** and a **SQL query** to extract data. 2. **SQL Patterns**: - Common use of `INNER JOIN` between `frpm` and `schools` tables. - Frequent use of `ORDER BY` and `LIMIT` to prioritize results. - Aggregation functions like `CAST` and arithmetic operations (e.g., `/` for rate calculation). 3. **Data Focus**: - Emphasis on **K-12 student metrics** (FRPM count, SAT test takers, enrollment). - Ownership code (`66`) and age ranges (`5-17`) are critical filters. --- ## Language and Formatting - **Primary Language**: English (all text and SQL queries). - **No Non-English Elements Detected**. --- ## Final Output Structure The flowchart is a **decision tree** for data extraction, with each node representing a similarity criterion and its associated SQL query. The final output aggregates results from all three criteria to rank schools by relevance. </details> Figure 2: An example demonstrating the utilization of different similarity metrics to find the most similar Question/SQL pair. 2.1.1 Question Similarity Similar questions often reference the same table and column names, or employ identical keywords when ordering or aggregating results. Therefore, the most straightforward approach to identifying in-context learning samples is by leveraging question similarity. This method may embed all questions within the sample pool, and then, for any specific question, after its embedding, a similarity metric such as cosine similarity is applied to identify samples with the most similar question embeddings. However, relying solely on question similarity is inadequate because NL2SQL requires structured prediction, which necessitates more explicit information about the problem’s structure than what is provided by the input question alone (Nan et al., 2023). 2.1.2 Query Skeleton Similarity Recent studies by Nan et al. (2023); Gao et al. (2023) introduced the concept of utilizing SQL query skeleton similarity as a criterion for selecting in-context learning samples. This method aims to guide the model about the necessary syntax and structure for crafting SQL queries. Their methodology starts with generating a draft SQL query for the given question using a preliminary NL2SQL model, then, encoding both the draft query and all queries from the sample pool into binary syntax vectors according to SQL keyword presence, and then selecting the most similar queries as in-context learning samples. While promising, this approach, we argue, falls short by neglecting the structural and keyword order intricacies inherent to SQL queries, resembling the limitations found in the bag-of-words model (Hiz and Harris, 2012). Moreover, when the preliminary NL2SQL model produce a draft SQL that significantly diverges from a correct solution, this method fails to accurately assess similarity. <details> <summary>2403.16204v1/extracted/5492121/sql_tree.png Details</summary>  ### Visual Description # Technical Document Extraction: SQL Query Flowchart Analysis ## Diagram Overview The image depicts a flowchart representation of a SQL query structure. The diagram uses circular nodes connected by arrows to illustrate the logical flow of a database query. All text is in English. ## Key Components and Flow ### 1. Root Node: SELECT - **Label**: `SELECT` - **Position**: Top-center of diagram - **Function**: Initiates the query structure ### 2. Primary Branches from SELECT #### A. FROM Clause - **Path**: `SELECT → FROM` - **Subcomponents**: - `[MASK]` (placeholder) - `satscores` (table name) - `AS T1` (alias assignment) #### B. INNER JOIN Clause - **Path**: `SELECT → INNER JOIN` - **Subcomponents**: - `schools` (table name) - `AS T2` (alias assignment) #### C. ON Clause - **Path**: `SELECT → ON` - **Subcomponents**: - `T1.cds = T2.CDSCode` (join condition) #### D. WHERE Clause - **Path**: `SELECT → WHERE` - **Subcomponents**: - `T1.NumTstTkr > 500` (filter condition) ### 3. Secondary Branches #### A. FROM → satscores → AS T1 - **Flow**: `FROM → satscores → AS T1` - **Purpose**: Defines primary table and alias #### B. INNER JOIN → schools → AS T2 - **Flow**: `INNER JOIN → schools → AS T2` - **Purpose**: Specifies joined table and alias #### C. WHERE → T1.NumTstTkr > 500 - **Flow**: `WHERE → T1.NumTstTkr > 500` - **Purpose**: Applies row filter condition ### 4. Mask Placeholders - **Occurrences**: 5 instances of `[MASK]` - **Placement**: - 2 under `FROM` branch - 2 under `ON` branch - 1 under `WHERE` branch - **Purpose**: Represents omitted/placeholder elements in query structure ## Spatial Grounding - **Root Node Position**: [x=center, y=top] - **Branch Distribution**: - FROM/INNER JOIN branches: Left side - ON/WHERE branches: Right side - **Mask Distribution**: Evenly distributed across branches ## Trend Verification - **Query Structure Flow**: Top-down hierarchical structure - **Component Relationships**: - SELECT → FROM/INNER JOIN → ON → WHERE - JOIN conditions precede filter conditions - **Alias Assignment Pattern**: `AS [alias]` follows table names ## Data Extraction ### Original SQL Query (Transcribed) </details> Figure 3: An example of the process to construct a tree from SQL query after masking the schema mentions. To address these shortcomings, our work proposes the construction of abstract syntax trees from SQL query skeletons, as shown in Figure 3, employing tree edit distance (Zhang and Shasha, 1989) to measure the query dissimilarity. In order to turn this dissimilarity metric into a similarity metric between 0 and 1, we normalized the distance by dividing the distances by the maximum distance for all of the pairs and using 1 - distance as the similarity metric. This method maintains the structural integrity of SQL queries and surpasses the performance of the syntax vector approach, showcasing its efficacy in preserving the order and structure of SQL query keywords. 2.1.3 Schema-Linking Similarity Previous research in the NL2SQL field (Pourreza and Rafiei, 2023; Li et al., 2023; Pourreza and Rafiei, 2024; Domínguez et al., 2024) has highlighted that a significant number of failures in NL2SQL models stem from mismatches between the schema links (tables and columns) in generated SQL queries and those in reference SQL queries. Contrary to cross-domain in-context learning, in-domain settings allow the utilization of table and column names as additional criteria for selecting in-context samples. In our study, we introduce a novel approach tailored to in-domain in-context learning, which selects samples based on the Jaccard similarity of the sets of tables and columns used within the queries. 2.1.4 Mean Of All Three Each of the aforementioned similarity metrics captures a unique aspect of similarity between two question/SQL pairs. To holistically assess the similarity score across all these dimensions, we calculate the mean of all three metrics as a definitive measure for selecting in-context samples. This integrated approach yielded superior in-context learning performance in our experiments. From now on, when we mention a similarity score, we’re referring to the combined score from all three metrics. 2.2 Encoder Following an extensive evaluation of various similarity metrics to identify the most efficacious approach, we utilized the BIRD dataset (Li et al., 2023) to compile a comprehensive finetuning dataset. This dataset comprises 170,000 pairs of questions, each annotated with a corresponding similarity score, serving as our supervised finetuning corpus. To accommodate the requisite knowledge of the database schema, we constructed in the input sequence to the model as the concatenation of the database schema together with the first and second question. Given that our task is a regression task, where the objective is to predict a continuous number between 0 and 1, we have applied a sigmoid function to the logits output by the model. This ensures that the predictions are confined within the 0 to 1 range. Subsequently, we utilized the Mean Absolute Error (MAE) as the loss function: $$ \text{L}=\frac{1}{N}\sum_{i=1}^{N}\left|y_{i}-\frac{1}{1+e^{-z_{i}}}\right| \tag{1} $$ where $N$ is the number of samples, $y_{i}$ denotes the actual similarity score for the $i$ -th question pair, and $z_{i}$ represents the predicted similarity by the model for the $i$ -th question pair. Contrary to previous methodologies that employed cross-encoder or ranking models built on pretrained architectures such as BERT (Devlin et al., 2018) or ROBERTA (Liu et al., 2019), investigated by (Nogueira et al., 2019; Reimers and Gurevych, 2019; Chen et al., 2024), our approach necessitated the handling of larger input sequences exceeding 2096 tokens. This requirement stemmed from integrating a database schema with two distinct questions. Additionally, our procedure could be enhanced by a code pretraining phase, akin to the pretraining seen in models like CodeLlama (Roziere et al., 2023) or DeepSeek (Guo et al., 2024). Nevertheless, given the need for cross-encoder models to compute similarity scores between each pair efficiently, the model’s size and speed are critical considerations. Therefore, we opted for the smallest version of the DeepSeek models, which has 1.3 billion parameters, as our foundational model for similarity score prediction. In order to make a single number prediction, a single linear layer with an input feature size of 2048 and an output size of 1 has been attached to the output of the last token from the model. 3 Experiments 3.1 Models To evaluate the effectiveness of our trained encoder model in predicting similarity scores for in-domain in-context learning sample selection, we report the performance of the in-context learning on three models from distinct families and architectures serving as SQL generators. Specifically, we utilized the CodeLLama 7B and 13B models (Roziere et al., 2023) alongside the gpt-3.5-turbo model, to assess the performance of the in-context learning. 3.2 Datasets Our evaluation involved two comprehensive cross-domain datasets, Spider and BIRD (Yu et al., 2018; Li et al., 2023). Spider contains 10,181 questions linked to 5,693 SQL queries across 138 domains and 200 databases, divided into 8,659 training, 1,034 development, and 2,147 test examples from unique databases. BIRD comprises 12,751 question-SQL pairs from 95 databases across 37+ domains, including blockchain and healthcare, totaling 33.4 GB. The dataset breakdown for BIRD includes 1,534 development, 9,428 training, and 1,789 test queries. We used the development sets from both datasets for evaluation without performing hyperparameter tuning. To create the training dataset, we utilized BIRD’s training set. For each question, we sampled 20 (or the total number available questions per database if it was fewer than 20) other questions from the same database. We then calculated the mean of all three similarity metrics for the pair to serve as the prediction label. This approach resulted in a training dataset of roughly 170K question pairs, which was used to train our model. 3.3 Metrics In the NL2SQL field, model performance is mainly assessed using two metrics: exact set match accuracy and execution accuracy. Exact set match accuracy evaluates the alignment of components like SELECT, WHERE, and GROUP BY in both generated and reference SQL queries without considering their sequence. Execution accuracy, on the other hand, measures whether a query produced by the model and the reference query return the same results across various databases. Given the diversity in correct SQL query formulations for a single question, execution accuracy is considered a more robust metric. This metric has been used as the primary metric for both Spider and BIRD benchmarks. Our study also adopts execution accuracy as the primary evaluation metric. 3.4 Hyperparameters We fine-tuned the DeepSeek model with 1.3B parameter on 160K samples for three epochs, setting aside 10K samples for validation and finding the best checkpoint. A full parameter finetuning was executed at a learning rate of 5e-5, using the AdamW optimizer (Loshchilov and Hutter, 2017) and a batch size of 8, on a single H100 GPU. 3.5 Results 3.5.1 Encoder Similarity Prediction This section aims to assess the encoder model’s performance in predicting similarity scores. These scores are evaluated in terms of being aligned with the mean similarity scores obtained by averaging three similarity metrics, using the ground truth SQL query. To evaluate the similarity prediction score, we employed the Kendall Tau coefficient (Kendall, 1938) to compare the orderings generated by our model’s predicted scores against the orderings derived from the mean of the three similarity metrics. The Kendall Tau coefficient ranges between -1 and 1, where 1 indicates a perfect match between the orderings, and -1 indicates a completely opposite ordering. Additionally, given the significance of the highest-scoring samples in few-shot in-context learning—since these samples are selected for inclusion in the prompt, we also calculated the precision in identifying the top 1, 5, 10, 15, and 20 samples at the top 1, 5, 10, 15, and 20 results predicted by our encoder. Table 2 presents the mean Kendall Tau coefficients and precision@k for different values of $k$ on the Spider dataset’s development set, comparing the orderings from the best similarity metric (mean of three) with those obtained by our encoder, the orderings based on question similarity (using the OpenAI embedding model), the orderings based on SQL skeleton similarity, and the orderings based on schema-linking similarity together with the precision@10. It is important to note that for the last two similarity metrics, we utilize the ground truth SQL query for similarity prediction, hence they serve as upper bounds. However, our trained encoder model was designed to predict the mean of three similarities without access to the SQL query, relying solely on the questions. | Method | Mean KT | Mean P@1 | Mean P@5 | Mean P@10 | Mean P@15 | Mean P@20 | | --- | --- | --- | --- | --- | --- | --- | | Encoder | 0.126 | 0.52 | 0.62 | 0.69 | 0.73 | 0.76 | | Schema-linking | 0.115 | 0.26 | 0.63 | 0.69 | 0.71 | 0.72 | | SQL Skeleton | 0.060 | 0.16 | 0.38 | 0.46 | 0.54 | 0.6 | | Question-similarity | 0.100 | 0.41 | 0.58 | 0.65 | 0.69 | 0.73 | Table 2: Mean Kendall Tau coefficients and precision@k for different values of k, comparing the ordering of different similarity metrics to that of the mean of all three metric on the development set of Spider. The findings presented in Table 2 reveal that our encoder aligns more closely with the ordering derived from the mean of the three similarity metrics than any individual metric that contributes to the mean. Furthermore, in terms of precision, our encoder is generally outperforming or reaching the same performance as all of the other three metrics. We believe that both precision@k and the Kendall Tau coefficient could be enhanced by employing larger base models. This is because predicting similarity by inferring the resemblance between SQL queries is a complex task, and the model should infer the SQL query internally in order to predict the similarity score and even the largest NL2SQL models like GPT-4 still struggle to perfectly predict the SQL query (Li et al., 2023). However, a trade-off exists, as larger models tend to slow down the similarity prediction process. This trade-off is precisely why we opted for a smaller base model. 3.5.2 Down-Stream Performance of Different Similarity Metrics In this section, we evaluate different methods for measuring the similarity between in-context Question/SQL pairs, with the goal of finding the metric that achieves the best down-stream performance. This evaluation is done on the development set of BIRD. Having a pool of Question/SQL pairs from a database and a given question from the same database at inference time, in order to find the best similarity metric, we assume that we have access to the SQL query to measure the schema-linking and SQL skeleton similarity. This assumption was made to find an upper bound for the similarity metrics. Previous works Nan et al. (2023); Gao et al. (2023) used a preliminary NL2SQL model to predict a draft SQL query and used that as the SQL query to measure the SQL skeleton similarity. After finding the best metric, our trained model only relies on the input questions and the database schema and is supposed to infer the SQL query skeleton from the input questions and database schema. The results of the comparison between different similarity metrics are reported in Table 3 by using the gpt-3.5-turbo model as the SQL prediction model in a one-shot setting. For the question similarity method, we used OpenAI embedding model (OpenAI, 2023) to obtain the question embeddings. | Method | EX | | --- | --- | | Zero-shot | 42.96 | | Question similarity (1-shot) | 48.89 | | SQL-Skeleton syntax vector (1-shot) | 46.28 | | SQL-Skeleton tree edit distance (1-shot) | 48.57 | | Schema-linking similarity (1-shot) | 50.59 | | Mean of all three (1-shot) | 52.93 | Table 3: Performance of different similarity metrics in selecting one-shot in-context learning examples for NL2SQL with gpt-3.5-turbo-1106 used as the LLM to generate SQL queries on the development set of BIRD. The data in Table 3 underscores the superiority of our proposed approach, which averages question similarity, schema-linking similarity, and SQL query skeleton similarity using the tree edit distance method. The proposed similarity metric can outperforms the best single similarity metric by 2.34%. Furthermore, the tree edit distance showcases enhanced performance over the basic syntax vector similarity method mentioned in Nan et al. (2023), emphasizing the significance of acknowledging the structure and order of SQL keywords in the queries. 3.5.3 Encoder Performance on BIRD In this section, we compare the performance of our encoder model, which utilizes the database schema and two questions as inputs, with the question similarity approaches that rely on leading embedding models such as OpenAI and Cohere (OpenAI, 2023; Cohere, 2023). It is important to note that we are not comparing our encoder’s performance with the methods outlined in Nan et al. (2023); Gao et al. (2023), which assess SQL Skeleton similarity by employing a preliminary model to predict a draft SQL query. These studies have demonstrated that the best performance is achieved when question similarity is integrated with SQL skeleton similarity. The decision to exclude their performance from this comparison is twofold: firstly, our trained model is primarily a question similarity method, as it processes questions directly rather than draft SQL queries; and secondly, our method is orthogonal to these works and our trained model could be used in their proposed framework to predict the question similarity. Table 4 showcases the performance of our encoder model in comparison to other question similarity approaches that utilize OpenAI text-embedding-3-large and Cohere cohere-embed-english-v3 embedding models (OpenAI, 2023; Cohere, 2023). | Encoder Encoder OpenAI | Model gpt-3.5-turbo gpt-3.5-turbo | EX 50.46 48.89 | | --- | --- | --- | | Cohere | gpt-3.5-turbo | 49.93 | | Zero-shot | gpt-3.5-turbo | 42.96 | | Encoder | CodeLlama-7B | 38.59 | | OpenAI | CodeLlama-7B | 34.35 | | Cohere | CodeLlama-7B | 33.18 | | Zero-shot | CodeLlama-7B | 24.77 | | Encoder | CodeLlama-13B | 40.09 | | OpenAI | CodeLlama-13B | 38.14 | | Cohere | CodeLlama-13B | 37.16 | | Zero-shot | CodeLlama-13B | 26.66 | Table 4: Performance of different question-similarity prediction models in 1-shot in-conext learning for NL2SQL task. For Cohere and OpenAI embedding models an extra step of cosine similarity calculation is required. The results displayed in Table 4 highlight the superior performance of our question-similarity prediction model compared to conventional models designed for assessing the similarity between two questions. Unlike the other two embedding models, where the similarity score is derived using cosine similarity between the embedding vectors, our encoder directly predicts the similarity score. Our trained encoder has consistently outperformed the other two models across various models for SQL prediction in one-shot setting, demonstrating its effectiveness. 3.5.4 Encoder Performance Across BIRD Query Classes The BIRD dataset categorizes question/SQL pairs into three distinct classes: simple, moderate, and challenging, based on the complexity of the target SQL queries (Li et al., 2023). In this section, we compare the performance across different query difficulty levels to explore the performance gains for each class. Table 5 demonstrates the performance across classes using different question similarity methods on various models, including gpt-3.5-turbo-1106, codellama-7b, and codellama-13b for in-context learning on the development set of the BIRD dataset. | Encoder | Model | Sim | Mod | Cha | | --- | --- | --- | --- | --- | | Encoder | GPT | 59.26 | 39.78 | 28.47 | | OpenAI | GPT | 59.25 | 33.55 | 31.25 | | Cohere | GPT | 58.38 | 39.35 | 29.86 | | Encoder | CL-7B | 47.78 | 25.38 | 22.22 | | OpenAI | CL-7B | 44 | 21.29 | 14.58 | | Cohere | CL-7B | 41.41 | 22.58 | 14.58 | | Encoder | CL-13B | 49.62 | 28.17 | 17.36 | | OpenAI | CL13B | 46.92 | 26.88 | 18.06 | | Cohere | CL-13B | 45.08 | 27.31 | 18.06 | Table 5: Performance of various question-similarity prediction models in one-shot in-context learning for the NL2SQL task across simple (Sim), moderate (Mod), and challenging (Cha) difficulty levels on the BIRD development set. ’GPT’ refers to gpt-3.5-turbo, and ’CL’ denotes CodeLlama. Based on the results, it is evident that our encoder model consistently surpasses the in-context learning performance of two other models for question similarity, particularly for smaller models such as CodeLlama 7B, with the most significant performance improvement observed in challenging samples. In contrast, with larger models like gpt-3.5-turbo and CodeLlama 13B, our encoder model excels in the simple and moderate categories but does not perform as well as the two other embedding models in the challenging category. We believe that the performance enhancement of our encoder model is particularly noticeable in scenarios where in-context learning yields more substantial improvements over zero-shot performance, as can also be seen in Table 4. 3.5.5 Encoder Performance on Spider Given that our model was trained on the BIRD training set, concerns regarding the generalizability of our question-similarity prediction model may arise. To address these concerns, we evaluated our encoder model’s performance on the Spider development set, comparing it with other question-similarity prediction methods that utilize OpenAI’s text-embedding-3-large and Cohere’s cohere-embed-english-v3 embedding models (OpenAI, 2023; Cohere, 2023). For SQL prediction, we employed the CodeLlama 13B parameter model to assess its one-shot performance. The results are presented in Table 5, highlighting the superior performance of our encoder model with almost 4% gain for both exact set match and execution accuracy over the other methods. One interpretation of the observed improvement in exact set match over the other two methods could be that our encoder preferentially selects samples with SQL queries more similar to the target SQL query. | OpenAI | CodeLlama-13B | 69.1 | 46.3 | | --- | --- | --- | --- | | Cohere | CodeLlama-13B | 69.3 | 46.3 | | Zero-shot | CodeLlama-13B | 64.7 | 24.9 | Table 6: Performance of the in-domain in-context learning on Spider development set using question similarity. EX refers to execution accuracy and EM refers to exact set match accuracy. 3.5.6 Encoder Performance Across Spider Query Classes The Spider dataset categorizes its Question/SQL pairs into four complexity classes: easy, medium, hard, and extra hard, based on the complexity of the structured SQL queries. Mirroring our analysis in Section 3.5.4, we compare our question similarity predictor encoder model against two other embedding models across these complexity classes on the Spider development set, utilizing CodeLlama-13B as the SQL prediction model. The results are detailed in Table 7. | Encoder | CL-13B | 88.3 | 80.9 | 55.7 | 48.2 | | --- | --- | --- | --- | --- | --- | | OpenAI | CL13B | 84.7 | 75.3 | 55.2 | 44.0 | | Cohere | CL-13B | 85.5 | 75.8 | 55.2 | 42.8 | Table 7: Performance of various question-similarity prediction models in one-shot in-context learning for the NL2SQL task across easy, medium, hard, and extra hard difficulty levels on the Spider development set. ’CL’ denotes CodeLlama. Our encoder model consistently outperforms the other two embedding models in question similarity prediction across all complexity levels, achieving the most significant gains in the extra hard and medium classes. 4 Conclusion Embedding questions into vectors transcends term-based methods, enabling the detection of questions that may be similar but expressed differently. However, generic embedding vectors cannot capture all intricacies of questions and how they may be formulated in SQL. Our work is the first to develop a similarity function for identifying structural similarities between queries based on their natural language expressions. Our approach, utilizing a lightweight model with 1.3 billion parameters, surpasses the performance of the high performing embedding models across both Bird and Spider datasets. The pursuit of specialized functions to better discern structural similarities between queries and their integration into NL2SQL pipelines shows promise for enhancing model performance. Limitations Contrary to embedding models that yield a single embedding vector suitable for storage and subsequent reuse for computing similarities between question pairs—thereby only necessitating the embedding of new samples for distance calculations against pre-existing vectors—our methodology adopts a cross-encoder architecture. This architecture requires the processing of each question pair through the similarity prediction model, rendering it less efficient for contexts characterized by extensive sets of question pairs. This design decision was made due two primary reasons. Firstly, the accurate prediction of SQL query similarities from questions demands an analytical process that encompasses both inputs, a requirement effectively met by the cross-encoder architecture. Secondly, in-context learning is predominantly used for smaller sample sizes, finetuning methods are generally more advantageous for larger datasets. (Li et al., 2024; Pourreza and Rafiei, 2024) Ethics Statement In this paper, we emphasize the significance of ethical considerations in our research, ensuring compliance with the ACL Ethics Policy and adherence to ethical guidelines throughout. Efforts were made to minimize biases and discrimination in our methodology and analysis. We have committed to transparency, accuracy, and fairness in our findings, providing due citations and acknowledgments. This ethics statement underscores our dedication to integrity, respect for ethical standards, and our contribution to the responsible advancement of knowledge in our field. Acknowledgements References - Brown et al. (2020) Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2020. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901. - Chen et al. (2024) Jianlv Chen, Shitao Xiao, Peitian Zhang, Kun Luo, Defu Lian, and Zheng Liu. 2024. Bge m3-embedding: Multi-lingual, multi-functionality, multi-granularity text embeddings through self-knowledge distillation. arXiv preprint arXiv:2402.03216. - Cohere (2023) Cohere. 2023. Cohere’s embedding models. - Devlin et al. (2018) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2018. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805. - Domínguez et al. (2024) José Manuel Domínguez, Benjamín Errázuriz, and Patricio Daher. 2024. Blar-sql: Faster, stronger, smaller nl2sql. arXiv preprint arXiv:2401.02997. - Gao et al. (2023) Dawei Gao, Haibin Wang, Yaliang Li, Xiuyu Sun, Yichen Qian, Bolin Ding, and Jingren Zhou. 2023. Text-to-sql empowered by large language models: A benchmark evaluation. arXiv preprint arXiv:2308.15363. - Guo et al. (2023) Chunxi Guo, Zhiliang Tian, Jintao Tang, Pancheng Wang, Zhihua Wen, Kang Yang, and Ting Wang. 2023. A case-based reasoning framework for adaptive prompting in cross-domain text-to-sql. arXiv preprint arXiv:2304.13301. - Guo et al. (2024) Daya Guo, Qihao Zhu, Dejian Yang, Zhenda Xie, Kai Dong, Wentao Zhang, Guanting Chen, Xiao Bi, Y Wu, YK Li, et al. 2024. Deepseek-coder: When the large language model meets programming–the rise of code intelligence. arXiv preprint arXiv:2401.14196. - Hiz and Harris (2012) H. Hiz and Z. Harris. 2012. Papers on Syntax. Studies in Linguistics and Philosophy. Springer Netherlands. - Kendall (1938) Maurice G Kendall. 1938. A new measure of rank correlation. Biometrika, 30(1/2):81–93. - Li et al. (2024) Haoyang Li, Jing Zhang, Hanbing Liu, Ju Fan, Xiaokang Zhang, Jun Zhu, Renjie Wei, Hongyan Pan, Cuiping Li, and Hong Chen. 2024. Codes: Towards building open-source language models for text-to-sql. arXiv preprint arXiv:2402.16347. - Li et al. (2023) Jinyang Li, Binyuan Hui, Ge Qu, Binhua Li, Jiaxi Yang, Bowen Li, Bailin Wang, Bowen Qin, Ruiying Geng, Nan Huo, Xuanhe Zhou, Chenhao Ma, Guoliang Li, Kevin C. C. Chang, Fei Huang, Reynold Cheng, and Yongbin Li. 2023. Can llm already serve as a database interface? a big bench for large-scale database grounded text-to-sqls. - Liu et al. (2021) Jiachang Liu, Dinghan Shen, Yizhe Zhang, Bill Dolan, Lawrence Carin, and Weizhu Chen. 2021. What makes good in-context examples for gpt- $3$ ? arXiv preprint arXiv:2101.06804. - Liu et al. (2019) Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. 2019. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11692. - Loshchilov and Hutter (2017) Ilya Loshchilov and Frank Hutter. 2017. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101. - Muennighoff et al. (2022) Niklas Muennighoff, Nouamane Tazi, Loïc Magne, and Nils Reimers. 2022. Mteb: Massive text embedding benchmark. arXiv preprint arXiv:2210.07316. - Nan et al. (2023) Linyong Nan, Yilun Zhao, Weijin Zou, Narutatsu Ri, Jaesung Tae, Ellen Zhang, Arman Cohan, and Dragomir Radev. 2023. Enhancing few-shot text-to-sql capabilities of large language models: A study on prompt design strategies. arXiv preprint arXiv:2305.12586. - Neelakantan et al. (2022) Arvind Neelakantan, Tao Xu, Raul Puri, Alec Radford, Jesse Michael Han, Jerry Tworek, Qiming Yuan, Nikolas Tezak, Jong Wook Kim, Chris Hallacy, et al. 2022. Text and code embeddings by contrastive pre-training. arXiv preprint arXiv:2201.10005. - Nogueira et al. (2019) Rodrigo Nogueira, Wei Yang, Kyunghyun Cho, and Jimmy Lin. 2019. Multi-stage document ranking with bert. arXiv preprint arXiv:1910.14424. - OpenAI (2023) OpenAI. 2023. Openai’s embedding models. - Pourreza and Rafiei (2023) Mohammadreza Pourreza and Davood Rafiei. 2023. Din-sql: Decomposed in-context learning of text-to-sql with self-correction. arXiv preprint arXiv:2304.11015. - Pourreza and Rafiei (2024) Mohammadreza Pourreza and Davood Rafiei. 2024. Dts-sql: Decomposed text-to-sql with small large language models. arXiv preprint arXiv:2402.01117. - Reimers and Gurevych (2019) Nils Reimers and Iryna Gurevych. 2019. Sentence-bert: Sentence embeddings using siamese bert-networks. arXiv preprint arXiv:1908.10084. - Roziere et al. (2023) Baptiste Roziere, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Tal Remez, Jérémy Rapin, et al. 2023. Code llama: Open foundation models for code. arXiv preprint arXiv:2308.12950. - Wei et al. (2022) Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. 2022. Emergent abilities of large language models. arXiv preprint arXiv:2206.07682. - Yu et al. (2018) Tao Yu, Rui Zhang, Kai Yang, Michihiro Yasunaga, Dongxu Wang, Zifan Li, James Ma, Irene Li, Qingning Yao, Shanelle Roman, et al. 2018. Spider: A large-scale human-labeled dataset for complex and cross-domain semantic parsing and text-to-sql task. arXiv preprint arXiv:1809.08887. - Zhang and Shasha (1989) Kaizhong Zhang and Dennis Shasha. 1989. Simple fast algorithms for the editing distance between trees and related problems. SIAM journal on computing, 18(6):1245–1262.