[Paper Review] DTS-SQL: Decomposed Text-to-SQL with Small Large Language Models

Mohammadreza Pourreza, Davood Rafiei|arXiv (Cornell University)|Feb 2, 2024

Scientific Computing and Data Management4 citations

TL;DR

This paper proposes DTS-SQL, a two-stage fine-tuning method that decomposes text-to-SQL into schema linking and SQL generation, using small 7B LLMs. By separating these tasks, the approach improves execution accuracy by 3–7 percentage points over single-step fine-tuning, enabling open-source models to match the performance of proprietary models like GPT-4 on benchmark datasets.

ABSTRACT

Leading models for the text-to-SQL task heavily rely on proprietary Large Language Models (LLMs), posing concerns over data privacy. Closing the performance gap between small open-source models and large proprietary models is crucial to mitigate this reliance. To this end, we introduce a novel two-stage fine-tuning approach that decomposes the task into two simpler tasks. Through comprehensive evaluation on two large cross-domain datasets and two small LLMs, we show that this approach improves execution accuracy by 3 to 7 percent, effectively aligning the performance of open-source models with their proprietary counterparts.

Motivation & Objective

• To close the performance gap between small open-source LLMs and large proprietary models like GPT-4 in text-to-SQL tasks.
• To address privacy and cost concerns associated with using proprietary LLMs in enterprise settings.
• To improve the performance of small 7B-parameter LLMs on cross-domain text-to-SQL benchmarks through a novel two-stage fine-tuning approach.
• To demonstrate that decomposed training can enhance both schema linking and SQL generation accuracy in low-resource settings.
• To provide a replicable, open-source solution that outperforms existing fine-tuned open-source models on Spider and Spider-SYN datasets.

Proposed method

• The method employs a two-stage fine-tuning pipeline: first, a dedicated model is trained to predict relevant database tables from a question and schema, and second, a separate model generates SQL from the question and the predicted schema.
• Schema linking is performed using a fine-tuned LLM that takes the natural language question and database schema as input and outputs a set of relevant table names.
• SQL generation is performed using a second fine-tuned LLM that takes the question and the output of the schema linker as input to generate the final SQL query.
• The approach uses supervised fine-tuning (SFT) on question-SQL pairs, with the schema linker trained on table selection and the SQL generator on full query generation.
• The pipeline decouples reasoning for schema linking and query generation, reducing cognitive load on each model and improving overall accuracy.
• The method is evaluated on two large cross-domain benchmarks: Spider and Spider-SYN, using two 7B LLMs: Mistral and DeepSeek.

Figure 1: The prompt used for SQL generation. The database schema is where we put the tables representations.

Experimental results

Research questions

• RQ1 Can a two-stage decomposition of text-to-SQL into schema linking and SQL generation improve performance of small open-source LLMs?
• RQ2 To what extent can small 7B-parameter LLMs achieve performance comparable to GPT-4 when using this decomposition approach?
• RQ3 How does the performance of the schema-linking component affect the overall execution accuracy of the text-to-SQL pipeline?
• RQ4 Does separating the tasks lead to better generalization and reduced error propagation compared to end-to-end fine-tuning?
• RQ5 Can this method outperform existing fine-tuned open-source models on standard benchmarks like Spider and Spider-SYN?

Key findings

• The DTS-SQL method improves execution accuracy by 3 to 7 percentage points over single-step fine-tuning on both the Spider and Spider-SYN datasets.
• On the Spider development set, the method achieves 85.5% execution accuracy with DeepSeek 7B and 78.6% with Mistral 7B, surpassing all prior open-source fine-tuned models.
• The schema-linking model achieves 93.1% exact set match accuracy on Spider and 87.6% on Spider-SYN, demonstrating strong table selection performance.
• The method achieves performance comparable to GPT-4-based prompting techniques, such as DIN-SQL + GPT4 (74.2% EX) and DAIL-SQL + GPT4 (84.4% EX), despite using only small open-source LLMs.
• The upper bound performance of the SQL generator (86.6% EX with Mistral 7B) indicates that further improvements are possible with better schema linking.
• The results demonstrate that task decomposition significantly enhances the effectiveness of small LLMs in complex text-to-SQL generation tasks.

Figure 2: The prompt used for Schema linking. The database schema is where we put the tables representations.